# Balanced configuration — primary test setup.
# L = 1.0 * L_img_img + 0.3 * L_sat_cap + 0.3 * L_drone_cap + 0.1 * L_cap_cap

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.shared_image_head = True
DualEncoderCaptionTest.device = "cuda"

ProjectionHead.in_dim = 512
ProjectionHead.out_dim = 512
ProjectionHead.use_mlp = False

# ---- Loss ----
MultiTermInfoNCE.temperature_init = 0.1
MultiTermInfoNCE.temperature_final = 0.01
MultiTermInfoNCE.label_smoothing = 0.1
MultiTermInfoNCE.asym_drone_to_sat = 0.6
MultiTermInfoNCE.asym_sat_to_drone = 0.4
MultiTermInfoNCE.warmup_epochs = 3
MultiTermInfoNCE.text_ramp_epochs = 10
MultiTermInfoNCE.lambda_ii = 1.0
MultiTermInfoNCE.lambda_sc_max = 0.3
MultiTermInfoNCE.lambda_dc_max = 0.3
MultiTermInfoNCE.lambda_cc_max = 0.1

# ---- Dataset ----
VisLocCaptionDataset.caption_strategy = "hybrid"
VisLocCaptionDataset.drop_caption_prob = 0.0
VisLocCaptionDataset.seed = 42

# ---- Training ----
TrainConfig.train_manifest = "data/visloc_train.json"
TrainConfig.val_manifest = "data/visloc_val.json"
TrainConfig.image_root = "data/visloc/images"
TrainConfig.output_dir = "out/caption_test/balanced"
TrainConfig.epochs = 30
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 = 1
TrainConfig.seed = 42
TrainConfig.device = "cuda"