# Balanced: GatedFusion with text captions enabled.
# query = sigma(alpha) * drone + (1-sigma(alpha)) * text -> InfoNCE vs satellite

import src.datasets.visloc_with_captions
import src.losses.multi_infonce
import src.models.dual_encoder
import src.training.train

# ---- Dual encoder ----
DualEncoderCaptionTest.variant = "ViT-B-32"
DualEncoderCaptionTest.pretrained_path = "checkpoints/RS5M_ViT-B-32.pt"
DualEncoderCaptionTest.unfreeze_mode = "last_block"
DualEncoderCaptionTest.embed_dim = 512
DualEncoderCaptionTest.use_mlp_heads = False
DualEncoderCaptionTest.baseline_mode = False
DualEncoderCaptionTest.init_gate = 0.7
DualEncoderCaptionTest.device = "cuda"

# ---- Fusion ----
GatedFusion.init_gate = 0.7
GatedFusion.baseline_mode = False

# ---- Loss ----
InfoNCELoss.temperature_init = 0.1
InfoNCELoss.temperature_final = 0.01
InfoNCELoss.label_smoothing = 0.1
InfoNCELoss.weight_q2g = 0.6
InfoNCELoss.weight_g2q = 0.4

# ---- Dataset ----
GeoLocCaptionDataset.drop_caption_prob = 0.0
GeoLocCaptionDataset.seed = 42

# ---- Training ----
TrainConfig.train_query_file = "/mnt/data1tb/cvgl_datasets/UAV-GeoLoc/Index/train_query.txt"
TrainConfig.val_query_file = "/mnt/data1tb/cvgl_datasets/UAV-GeoLoc/Index/val_query.txt"
TrainConfig.data_root = "/mnt/data1tb/cvgl_datasets/UAV-GeoLoc"
TrainConfig.output_dir = "out/caption_test/balanced"
TrainConfig.epochs = 10
TrainConfig.batch_size = 128
TrainConfig.num_workers = 4
TrainConfig.learning_rate = 1e-4
TrainConfig.weight_decay = 1e-4
TrainConfig.grad_clip = 1.0
TrainConfig.use_amp = True
TrainConfig.eval_every = 2
TrainConfig.seed = 42
TrainConfig.device = "cuda"