# GTA-UAV Balanced: Asymmetric DINOv3 (WEB+SAT) with L1/L2/L3 captions.
# query = sigma(alpha) * drone + (1-sigma(alpha)) * text -> InfoNCE vs gallery
# 10 epochs, MONA all 24 blocks, 1024-dim retrieval, hard negative bank.
#
# NOTE: TrainConfigGTAUAV is registered by train_gtauav.py before gin parsing.
# InfoNCELoss is registered via import below.

import src.losses.multi_infonce

# ---- Training ----
TrainConfigGTAUAV.epochs = 10
TrainConfigGTAUAV.batch_size = 8
TrainConfigGTAUAV.num_workers = 4
TrainConfigGTAUAV.learning_rate = 1e-4
TrainConfigGTAUAV.text_lr_factor = 0.1
TrainConfigGTAUAV.weight_decay = 1e-4
TrainConfigGTAUAV.grad_clip = 1.0
TrainConfigGTAUAV.grad_accum_steps = 8
TrainConfigGTAUAV.use_amp = True
TrainConfigGTAUAV.eval_every = 1
TrainConfigGTAUAV.warmup_epochs = 2
TrainConfigGTAUAV.seed = 42
TrainConfigGTAUAV.device = "cuda"

# ---- Model ----
TrainConfigGTAUAV.init_gate = 0.7
TrainConfigGTAUAV.baseline_mode = False
TrainConfigGTAUAV.shared_encoder = False
TrainConfigGTAUAV.gradient_checkpointing = True

# ---- Loss ----
TrainConfigGTAUAV.tau_init = 0.07
TrainConfigGTAUAV.label_smoothing = 0.1
TrainConfigGTAUAV.weight_q2g = 0.6
TrainConfigGTAUAV.weight_g2q = 0.4
TrainConfigGTAUAV.learnable_temperature = True
TrainConfigGTAUAV.neg_bank_size = 4096

# ---- Output ----
TrainConfigGTAUAV.output_dir = "out/gtauav/with_text"

# ---- Tracking ----
TrainConfigGTAUAV.use_wandb = False
TrainConfigGTAUAV.use_tb = True
TrainConfigGTAUAV.use_gradcam = True
TrainConfigGTAUAV.gradcam_every = 5
TrainConfigGTAUAV.use_profiler = False
TrainConfigGTAUAV.log_grad_norms = True

# ---- InfoNCE Loss (gin-configurable) ----
InfoNCELoss.temperature_init = 0.07
InfoNCELoss.temperature_final = 0.01
InfoNCELoss.label_smoothing = 0.1
InfoNCELoss.weight_q2g = 0.6
InfoNCELoss.weight_g2q = 0.4
InfoNCELoss.learnable_temperature = True