"""Training recipe: loss + optimizer + sampler.

Three concerns kept together because they form one coherent recipe — they
co-vary across experiments. Splitting Loss vs Optimizer vs Sampler can be
done later if a need emerges.
"""

from __future__ import annotations

import gin


@gin.configurable
class TrainingConfig:
    """Loss + optimizer + sampler.

    Selects between InfoNCELoss and WeightedInfoNCELoss via `loss_type`.
    Selects between DSS / mutex / plain shuffle via `sampler_type`.
    """

    def __init__(
        self,
        # ---- Loss: shared between InfoNCELoss and WeightedInfoNCELoss ----
        loss_type: str = "symmetric",            # 'symmetric' | 'weighted'
        tau_init: float = 0.07,
        tau_min: float = 0.01,
        tau_max: float = 0.1,
        learnable_temperature: bool = True,
        label_smoothing: float = 0.1,
        # ---- Loss: InfoNCELoss-only ----
        tau_final: float = 0.01,                 # cosine-schedule final tau (when not learnable)
        weight_q2g: float = 0.6,
        weight_g2q: float = 0.4,
        hard_mining_k: int = 0,
        neg_bank_size: int = 0,
        # ---- Loss: WeightedInfoNCELoss-only ----
        weighted_loss_k: float = 5.0,            # sigmoid steepness for weight→eps mapping
        # ---- Optimizer ----
        learning_rate: float = 1e-4,
        text_lr_factor: float = 0.1,             # lr * factor for DGTRS-CLIP/LoRA params
        weight_decay: float = 1e-4,
        grad_clip: float = 1.0,
        # ---- Sampler ----
        sampler_type: str = "mutex",             # 'mutex' | 'dss' | 'none'
        dss_warmup_epochs: int = 1,
        dss_reembed_every: int = 1,
        dss_knn_device: str = "cuda",
        dss_use_lsh: bool = False,
        dss_lsh_num_tables: int = 8,
        dss_lsh_num_bits: int = 14,
        dss_cache_dir: str | None = None,
        # Legacy alias (kept until train_gtauav.py is rewritten in step 4).
        use_mutex_sampler: bool = True,
    ) -> None:
        # Loss (shared).
        self.loss_type = loss_type
        self.tau_init = tau_init
        self.tau_min = tau_min
        self.tau_max = tau_max
        self.learnable_temperature = learnable_temperature
        self.label_smoothing = label_smoothing
        # Loss (InfoNCE-specific).
        self.tau_final = tau_final
        self.weight_q2g = weight_q2g
        self.weight_g2q = weight_g2q
        self.hard_mining_k = hard_mining_k
        self.neg_bank_size = neg_bank_size
        # Loss (WeightedInfoNCE-specific).
        self.weighted_loss_k = weighted_loss_k
        # Optimizer.
        self.learning_rate = learning_rate
        self.text_lr_factor = text_lr_factor
        self.weight_decay = weight_decay
        self.grad_clip = grad_clip
        # Sampler.
        self.sampler_type = sampler_type
        self.dss_warmup_epochs = dss_warmup_epochs
        self.dss_reembed_every = dss_reembed_every
        self.dss_knn_device = dss_knn_device
        self.dss_use_lsh = dss_use_lsh
        self.dss_lsh_num_tables = dss_lsh_num_tables
        self.dss_lsh_num_bits = dss_lsh_num_bits
        self.dss_cache_dir = dss_cache_dir
        self.use_mutex_sampler = use_mutex_sampler


def get_training_cfg(path2cfg: str) -> TrainingConfig:
    """Load ONLY training config (TESTING ONLY — production uses load_all_configs)."""
    gin.clear_config()
    gin.parse_config_file(f"{path2cfg}training.gin")
    return TrainingConfig()