from __future__ import annotations

"""Training loop for CVGL caption test on GTA-UAV-LR dataset.

Asymmetric DINOv3 encoders (drone LVD + satellite SAT) with LRSCLIP text fusion.
Single InfoNCE loss: query(drone+text) vs gallery(satellite).

Supports gin-config, W&B, TensorBoard, Grad-CAM, gradient monitoring,
PyTorch Profiler, and torchinfo model summary.
"""

import argparse
import json
import logging
import math
import time
import warnings
from dataclasses import dataclass, field
from pathlib import Path

import coloredlogs
import gin
import pandas as pd
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.amp import GradScaler, autocast
from torch.optim import AdamW
from torch.optim.lr_scheduler import LambdaLR
from torch.utils.data import DataLoader
from tqdm import tqdm

from src.datasets.gtauav_dataset import (
    GTAUAVDataset,
    GTAUAVDroneQuery,
    GTAUAVSatGallery,
    collate_drone_query,
    collate_gtauav_batch,
    collate_sat_gallery,
)
from src.datasets.dynamic_similarity_sampler import DynamicSimilaritySampler
from src.datasets.embedding_cache import EmbeddingCache
from src.datasets.mutually_exclusive_sampler import MutuallyExclusiveSampler
from src.losses.multi_infonce import InfoNCELoss
from src.losses.weighted_infonce import WeightedInfoNCELoss
from src.losses.hard_negatives import NegativeMemoryBank
from src.training.plot_metrics import generate_plots
from src.training.trackers import ExperimentTracker
from src.training.grad_monitor import compute_gradient_norms, log_gradient_summary
from src.training.profiling import TrainingProfiler, print_model_summary
from src.models.asymmetric_encoder import (
    AsymmetricEncoder,
    get_dino_transform,
    get_drone_train_transform,
    get_satellite_train_transform,
)
from src.models.sofia_fusion_encoder import SOFIAFusionEncoder
from src.models.sofia_v1 import SOFIAv1Config
from src.models.sofia_v1_fusion_encoder import SOFIAv1FusionEncoder
from src.models.sofia_v71 import (
    sofia_l_config,
    sofia_m_config,
    sofia_tiny_config,
)

LOGGER = logging.getLogger("caption_test.train_gtauav")

# Default paths.
_RGB_ROOT = "/home/servml/Документы/datasets/GTA-UAV-LR"
_CAPTION_ROOT = "/home/servml/Документы/datasets/GTA-UAV-LR-captions"
_TRAIN_JSON = "meta/train_80.json"
_TEST_JSON = "meta/test_20.json"

_DINO_WEB = "nn_models/DINO_WEB/dinov3-vitl16-pretrain-lvd1689m.pth"
_DINO_SAT = "nn_models/DINO_SAT/model.safetensors"
_LRSCLIP = "nn_models/LRSCLIP/DGTRS-CLIP-ViT-L-14.pt"


@gin.configurable(module="src.training.train_gtauav")
@dataclass
class TrainConfigGTAUAV:
    """Training configuration for GTA-UAV experiment."""

    # Data.
    train_json: str = _TRAIN_JSON
    test_json: str = _TEST_JSON
    rgb_root: str = _RGB_ROOT
    caption_root: str = _CAPTION_ROOT
    filter_meta: str | None = None

    # Model.
    dino_web_path: str = _DINO_WEB
    dino_sat_path: str = _DINO_SAT
    lrsclip_path: str = _LRSCLIP
    init_gate: float = 0.7
    baseline_mode: bool = False
    shared_encoder: bool = True  # single DINOv3 WEB for both branches (simpler, half the params)
    mona_bottleneck: int = 64
    mona_last_n_blocks: int = 12  # inject adapters only in last 12 of 24 ViT blocks
    gradient_checkpointing: bool = True  # trade compute for VRAM (allows larger batch)
    # StripNet backbone option (replaces DINOv3 when backbone="stripnet").
    backbone: str = "dinov3"  # "dinov3", "stripnet", or "sofia"
    stripnet_path: str = "nn_models/STRIPNET/stripnet_s.pth"
    stripnet_mona_last_n_stages: int = 2  # Conv-MONA in last N of 4 StripNet stages (0 = disable MONA)
    stripnet_freeze: bool = True  # If False, StripNet backbone is fully trainable (full fine-tune)
    stripnet_backbone_lr_factor: float = 0.1  # Backbone LR = learning_rate * factor (only when unfrozen)
    # SOFIA backbone options (used when backbone="sofia"). Trained from scratch — no pretrained checkpoint.
    sofia_preset: str = "Tiny"  # "Tiny" | "M" | "L"
    sofia_d_descriptor: int = 1024  # retrieval space (1024 = match TextFusionMLP out_dim)
    sofia_use_text_film_uav: bool = True  # mid-level text-FiLM in UAV head
    sofia_use_text_film_sat: bool = True  # mid-level text-FiLM in SAT head
    sofia_lora_rank: int = 4
    sofia_mamba_variant: str = "mamba2"  # "mamba1" | "mamba2" | "efficient_vmamba"
    sofia_mamba_backend: str = "auto"  # "auto" | "torch" | "mamba_ssm"
    # EVSSBridge (B6-inspired refinement between heterogeneous stages, opt-in).
    sofia_use_evss_bridge: bool = False
    sofia_evss_bridge_locations: list[str] = field(default_factory=lambda: ["pre_stage3"])
    # SOFIA v1 backbone options (used when backbone="sofia_v1"). StripNet+DCN, from scratch.
    sofia_v1_variant: str = "small"  # "tiny_tiny" | "tiny" | "small" | "small_v2"
    sofia_v1_dcn_variant: str = "v2"  # "v2" (torchvision DeformConv2d, stable) | "v4" (OpenGVLab, leaky)
    sofia_v1_d_descriptor: int = 1024
    sofia_v1_use_text_film_uav: bool = True
    sofia_v1_use_text_film_sat: bool = True
    sofia_v1_use_film_altitude: bool = True
    sofia_v1_lora_rank: int = 4

    # Training.
    resume_from: str | None = None  # path to checkpoint for resuming
    output_dir: str = "out/gtauav/with_text"
    epochs: int = 10
    batch_size: int = 8
    num_workers: int = 4
    learning_rate: float = 1e-4
    text_lr_factor: float = 0.1  # text encoder LR = learning_rate * factor
    weight_decay: float = 1e-4
    grad_clip: float = 1.0
    grad_accum_steps: int = 1  # gradient accumulation steps (effective_batch = batch_size * accum)
    use_amp: bool = True
    eval_every: int = 2
    warmup_epochs: int = 2
    seed: int = 42
    device: str = "cuda"

    # Loss.
    loss_type: str = "symmetric"  # "symmetric" (InfoNCE) or "weighted" (WeightedInfoNCE)
    tau_init: float = 0.07
    label_smoothing: float = 0.1
    learnable_temperature: bool = True
    weight_q2g: float = 0.6
    weight_g2q: float = 0.4
    neg_bank_size: int = 4096  # hard negative memory bank size (0 = disabled)

    # Sampling.
    sampler_type: str = "mutex"  # "mutex" (no false negatives) or "dss" (DSS + mutex)
    dss_reembed_every: int = 1  # Re-embed train queries every N epochs for DSS.
    dss_warmup_epochs: int = 1  # Use mutex-only for the first N epochs (fresh model embeddings aren't useful)
    dss_knn_device: str = "cuda"  # Device for similarity matmul in DSS sampler.
    dss_use_lsh: bool = False  # Approximate kNN via LSH (opt-in; exact is fast at 25K).
    dss_lsh_num_tables: int = 8
    dss_lsh_num_bits: int = 14
    dss_cache_dir: str | None = None  # Disk cache for embeddings; None = disabled.
    # Legacy alias kept for backward compatibility.
    use_mutex_sampler: bool = True

    # Tracking & diagnostics.
    use_wandb: bool = False
    use_tb: bool = True
    wandb_project: str = "caption-test-gtauav"
    wandb_run_name: str | None = None
    wandb_entity: str | None = None
    log_grad_norms: bool = True
    use_gradcam: bool = False
    gradcam_every: int = 5  # Grad-CAM every N epochs
    gradcam_samples: int = 8
    use_profiler: bool = False
    profiler_warmup: int = 3
    profiler_active: int = 5


def _set_seed(seed: int) -> None:
    import random as _random
    import numpy as _np
    _random.seed(seed)
    _np.random.seed(seed)
    torch.manual_seed(seed)
    torch.cuda.manual_seed_all(seed)


def _atomic_save(obj: dict, path: Path) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    tmp_path = path.with_suffix(path.suffix + ".tmp")
    torch.save(obj, tmp_path)
    tmp_path.replace(path)


def _build_param_groups(
    model: nn.Module,
    lr: float,
    text_lr_factor: float,
    stripnet_backbone_lr_factor: float = 0.1,
) -> list[dict]:
    """Build optimizer param groups with separate LR for text encoder and unfrozen StripNet backbone.

    Groups:
      - text_encoder.*       → lr * text_lr_factor (default 1e-5)
      - image_encoder.backbone.* (when StripNet unfrozen) → lr * stripnet_backbone_lr_factor (default 1e-5)
      - everything else (MONA, projection, TextFusionMLP, gates, tau, MONA on Conv,
                         SOFIA backbone+heads when backbone="sofia") → lr
    """
    text_params = []
    backbone_params = []
    other_params = []

    is_stripnet = isinstance(getattr(model, "image_encoder", None), nn.Module) and \
                  getattr(model, "backbone", "dinov3") == "stripnet"

    for name, param in model.named_parameters():
        if not param.requires_grad:
            continue
        if "text_encoder" in name:
            text_params.append(param)
        elif is_stripnet and name.startswith("image_encoder.backbone.") and "mona_" not in name:
            backbone_params.append(param)
        else:
            other_params.append(param)

    groups = [{"params": other_params, "lr": lr}]
    if text_params:
        groups.append({"params": text_params, "lr": lr * text_lr_factor})
    if backbone_params:
        groups.append({"params": backbone_params, "lr": lr * stripnet_backbone_lr_factor})

    return groups


def _cosine_warmup_schedule(
    warmup_steps: int,
    total_steps: int,
) -> callable:
    """Cosine annealing with linear warmup."""

    def lr_lambda(step: int) -> float:
        if step < warmup_steps:
            return step / max(warmup_steps, 1)
        progress = (step - warmup_steps) / max(total_steps - warmup_steps, 1)
        return 0.5 * (1.0 + math.cos(math.pi * progress))

    return lr_lambda


@torch.no_grad()
def _embed_drone_queries(
    model: AsymmetricEncoder,
    train_ds: GTAUAVDataset,
    device: str,
    batch_size: int,
    num_workers: int,
) -> torch.Tensor:
    """Forward all drone queries and return [N, D] embeddings on CPU.

    Used by DynamicSimilaritySampler to rank drones by visual similarity.
    Runs with model.eval() but restores original train state afterwards.
    """
    was_training = model.training
    model.eval()

    query_ds = GTAUAVDroneQuery(train_ds)
    loader = DataLoader(
        query_ds,
        batch_size=batch_size,
        shuffle=False,
        num_workers=num_workers,
        collate_fn=collate_drone_query,
        pin_memory=True,
    )

    embs: list[torch.Tensor] = []
    for batch in tqdm(loader, desc="  dss-embed-queries", unit="batch", leave=False):
        drone_img = batch["drone_img"].to(device, non_blocking=True)
        altitude = batch.get("altitude")
        if altitude is not None:
            altitude = altitude.to(device, non_blocking=True)
        q = model.encode_query(
            drone_img,
            batch["caption_l1"], batch["caption_l2"], batch["caption_l3"],
            altitude=altitude,
        )
        embs.append(q.cpu())

    if was_training:
        model.train()
    return torch.cat(embs, dim=0)


@torch.no_grad()
def _evaluate(
    model: AsymmetricEncoder,
    loader: DataLoader,
    device: str,
    loss_fn: nn.Module | None = None,
    epoch: int = 0,
    total_epochs: int = 1,
    k_values: tuple[int, ...] = (1, 5, 10),
    max_batches: int | None = None,
    desc: str = "eval",
) -> dict[str, float]:
    """Compute R@K and MRR on the full satellite gallery.

    Standard CVGL retrieval: forward every unique satellite in the dataset
    once (gallery), forward every drone query, then rank gallery by
    cosine similarity. A query counts as a hit@K if ANY of its valid
    satellite matches (pair_pos_sate_img_list ∪ pair_pos_semipos_sate_img_list)
    appears in the top-K.

    `max_batches` subsamples the drone queries (not the gallery) — useful
    for a quick train-side sanity check.
    """
    dataset = loader.dataset
    if not isinstance(dataset, GTAUAVDataset):
        raise TypeError(f"_evaluate expects GTAUAVDataset, got {type(dataset).__name__}")

    model.eval()

    batch_size = loader.batch_size or 32
    num_workers = getattr(loader, "num_workers", 0)
    pin_memory = getattr(loader, "pin_memory", False)

    gallery_ds = GTAUAVSatGallery(dataset)
    query_ds = GTAUAVDroneQuery(dataset)

    gallery_loader = DataLoader(
        gallery_ds,
        batch_size=batch_size,
        shuffle=False,
        num_workers=num_workers,
        pin_memory=pin_memory,
        collate_fn=collate_sat_gallery,
    )
    query_loader = DataLoader(
        query_ds,
        batch_size=batch_size,
        shuffle=False,
        num_workers=num_workers,
        pin_memory=pin_memory,
        collate_fn=collate_drone_query,
    )

    # --- Gallery forward (all unique sats) ---
    gallery_embs: list[torch.Tensor] = []
    gallery_names: list[str] = []
    for batch in tqdm(gallery_loader, desc=f"  {desc}-gallery", unit="batch", leave=False):
        sat_img = batch["sat_img"].to(device, non_blocking=True)
        g = model.encode_gallery(
            sat_img,
            batch["sat_caption_l1"], batch["sat_caption_l2"], batch["sat_caption_l3"],
        )
        gallery_embs.append(g.cpu())
        gallery_names.extend(batch["sat_names"])
    gallery = torch.cat(gallery_embs, dim=0)  # [N_sat, D]

    # --- Query forward (optionally subsampled via max_batches) ---
    query_embs: list[torch.Tensor] = []
    query_valid_names: list[list[str]] = []
    batch_losses: list[float] = []
    sat_name_to_idx: dict[str, int] = {name: i for i, name in enumerate(gallery_names)}

    for i, batch in enumerate(tqdm(query_loader, desc=f"  {desc}-query", unit="batch", leave=False)):
        if max_batches is not None and i >= max_batches:
            break
        drone_img = batch["drone_img"].to(device, non_blocking=True)
        altitude = batch.get("altitude")
        if altitude is not None:
            altitude = altitude.to(device, non_blocking=True)
        q = model.encode_query(
            drone_img,
            batch["caption_l1"], batch["caption_l2"], batch["caption_l3"],
            altitude=altitude,
        )
        query_embs.append(q.cpu())
        query_valid_names.extend(batch["valid_sat_names"])

        # Per-batch loss: use first valid sat per query as its paired gallery.
        if loss_fn is not None:
            pair_indices: list[int] = []
            for names in batch["valid_sat_names"]:
                for name in names:
                    if name in sat_name_to_idx:
                        pair_indices.append(sat_name_to_idx[name])
                        break
                else:
                    pair_indices.append(-1)
            if all(idx >= 0 for idx in pair_indices):
                paired_gallery = gallery[pair_indices].to(device)
                fake_embeddings = {
                    "query": q,
                    "gallery": paired_gallery,
                    "gate_q": model.fusion_query.gate_value,
                    "gate_g": model.fusion_gallery.gate_value,
                }
                loss_dict = loss_fn(fake_embeddings, epoch=epoch, total_epochs=total_epochs)
                batch_losses.append(float(loss_dict["total"].item()))

    query = torch.cat(query_embs, dim=0)  # [N_q, D]
    n_query = query.size(0)

    # --- Similarity + rankings ---
    sim = query @ gallery.t()  # [N_q, N_sat]
    sorted_idx = sim.argsort(dim=1, descending=True)

    metrics: dict[str, float] = {}
    if batch_losses:
        metrics["loss"] = sum(batch_losses) / len(batch_losses)

    # Precompute valid gallery index sets per query.
    valid_idx_per_query: list[set[int]] = []
    for names in query_valid_names:
        valid = {sat_name_to_idx[n] for n in names if n in sat_name_to_idx}
        valid_idx_per_query.append(valid)

    # R@K with multi-match.
    for k in k_values:
        hits = 0
        for i in range(n_query):
            top_k = set(sorted_idx[i, :k].tolist())
            if valid_idx_per_query[i] & top_k:
                hits += 1
        metrics[f"r@{k}_q2g"] = hits / max(n_query, 1)

    # MRR over valid matches (kept key `ap_q2g` for CSV/plot compatibility).
    mrr_sum = 0.0
    n_scored = 0
    for i in range(n_query):
        valid = valid_idx_per_query[i]
        if not valid:
            continue
        n_scored += 1
        for rank, gidx in enumerate(sorted_idx[i].tolist()):
            if gidx in valid:
                mrr_sum += 1.0 / (rank + 1)
                break
    metrics["ap_q2g"] = mrr_sum / max(n_scored, 1)

    # --- g2q (satellite → drone): invert ground-truth ---
    n_gallery = gallery.size(0)
    valid_q_per_sat: list[set[int]] = [set() for _ in range(n_gallery)]
    for q_idx, gset in enumerate(valid_idx_per_query):
        for g_idx in gset:
            valid_q_per_sat[g_idx].add(q_idx)

    sorted_idx_g2q = sim.t().argsort(dim=1, descending=True)  # [N_sat, n_query]
    n_scored_g2q = sum(1 for s in valid_q_per_sat if s)

    for k in k_values:
        hits_g2q = 0
        for i in range(n_gallery):
            valid = valid_q_per_sat[i]
            if not valid:
                continue
            top_k = set(sorted_idx_g2q[i, :k].tolist())
            if valid & top_k:
                hits_g2q += 1
        metrics[f"r@{k}_g2q"] = hits_g2q / max(n_scored_g2q, 1)

    mrr_sum_g2q = 0.0
    for i in range(n_gallery):
        valid = valid_q_per_sat[i]
        if not valid:
            continue
        for rank, qidx in enumerate(sorted_idx_g2q[i].tolist()):
            if qidx in valid:
                mrr_sum_g2q += 1.0 / (rank + 1)
                break
    metrics["ap_g2q"] = mrr_sum_g2q / max(n_scored_g2q, 1)

    metrics["n_query"] = float(n_query)
    metrics["n_gallery"] = float(n_gallery)
    metrics["n_scored_g2q"] = float(n_scored_g2q)

    metrics["gate_q"] = model.fusion_query.gate_value
    metrics["gate_g"] = model.fusion_gallery.gate_value
    return metrics


class CSVLogger:
    """Log train/val metrics to CSV files using pandas.

    Creates:
        {output_dir}/logs/train.csv               — epoch-level train averages
        {output_dir}/logs/val.csv                  — epoch-level val metrics
        {output_dir}/logs/train_batches.csv        — per-batch train metrics (all epochs)
        {output_dir}/logs/epoch_{N}_batches.csv    — per-batch for single epoch
    """

    def __init__(self, output_dir: Path) -> None:
        self.log_dir = output_dir / "logs"
        self.log_dir.mkdir(parents=True, exist_ok=True)
        self._current_epoch: int = -1
        self._batch_columns: list[str] | None = None
        self._cumulative_batch_path = self.log_dir / "train_batches.csv"
        self._epoch_batch_path: Path | None = None

        # Load existing CSV data on resume (so plots show full history).
        train_csv = self.log_dir / "train.csv"
        val_csv = self.log_dir / "val.csv"
        train_recall_csv = self.log_dir / "train_recall.csv"
        if train_csv.exists():
            self.train_rows = pd.read_csv(train_csv).to_dict("records")
            LOGGER.info("CSVLogger: loaded %d previous train epochs", len(self.train_rows))
        else:
            self.train_rows = []
        if val_csv.exists():
            self.val_rows = pd.read_csv(val_csv).to_dict("records")
            LOGGER.info("CSVLogger: loaded %d previous val epochs", len(self.val_rows))
        else:
            self.val_rows = []
        if train_recall_csv.exists():
            self.train_recall_rows = pd.read_csv(train_recall_csv).to_dict("records")
        else:
            self.train_recall_rows = []

    def log_batch(self, epoch: int, batch_idx: int, global_step: int, metrics: dict) -> None:
        """Log metrics for a single training batch. Writes to disk immediately."""
        row = {"epoch": epoch, "batch": batch_idx, "global_step": global_step, **metrics}

        # On new epoch, start a fresh per-epoch CSV.
        if epoch != self._current_epoch:
            self._current_epoch = epoch
            self._epoch_batch_path = self.log_dir / f"epoch_{epoch:03d}_batches.csv"

        # Determine columns on first call (consistent order).
        if self._batch_columns is None:
            self._batch_columns = list(row.keys())

        row_df = pd.DataFrame([row], columns=self._batch_columns)
        write_header = not self._cumulative_batch_path.exists()

        # Append to cumulative CSV.
        row_df.to_csv(
            self._cumulative_batch_path, mode="a", header=write_header, index=False,
        )
        # Append to per-epoch CSV.
        write_epoch_header = not self._epoch_batch_path.exists()
        row_df.to_csv(
            self._epoch_batch_path, mode="a", header=write_epoch_header, index=False,
        )

    def log_train(self, epoch: int, metrics: dict, lr: float, elapsed: float) -> None:
        """Log epoch-level train averages. Replaces existing entry for same epoch on resume."""
        row = {"epoch": epoch, "lr": lr, "elapsed_s": round(elapsed, 1), **metrics}
        # Remove previous entry for this epoch (resume may re-run it).
        self.train_rows = [r for r in self.train_rows if r.get("epoch") != epoch]
        self.train_rows.append(row)
        pd.DataFrame(self.train_rows).to_csv(self.log_dir / "train.csv", index=False)

    def log_val(self, epoch: int, metrics: dict) -> None:
        """Log val metrics. Replaces existing entry for same epoch on resume."""
        row = {"epoch": epoch, **metrics}
        self.val_rows = [r for r in self.val_rows if r.get("epoch") != epoch]
        self.val_rows.append(row)
        pd.DataFrame(self.val_rows).to_csv(self.log_dir / "val.csv", index=False)

    def log_train_recall(self, epoch: int, metrics: dict) -> None:
        """Log train recall metrics. Replaces existing entry for same epoch."""
        row = {"epoch": epoch, **metrics}
        self.train_recall_rows = [r for r in self.train_recall_rows if r.get("epoch") != epoch]
        self.train_recall_rows.append(row)
        pd.DataFrame(self.train_recall_rows).to_csv(self.log_dir / "train_recall.csv", index=False)


def _clear_vram() -> None:
    """Free VRAM from previous runs before starting."""
    import gc
    gc.collect()
    if torch.cuda.is_available():
        torch.cuda.empty_cache()
        torch.cuda.reset_peak_memory_stats()
        allocated = torch.cuda.memory_allocated() / 1e9
        LOGGER.info("VRAM cleared. Current usage: %.2f GB", allocated)


def train(cfg: TrainConfigGTAUAV) -> None:
    """Run full training loop."""
    coloredlogs.install(
        level="INFO",
        logger=LOGGER,
        fmt="%(asctime)s %(name)s %(levelname)s %(message)s",
    )
    _clear_vram()
    _set_seed(cfg.seed)
    output_dir = Path(cfg.output_dir)
    output_dir.mkdir(parents=True, exist_ok=True)

    # Save config.
    with (output_dir / "config.json").open("w") as f:
        json.dump(vars(cfg), f, indent=2)

    # --- Experiment tracker (W&B + TensorBoard) ---
    tracker = ExperimentTracker(
        output_dir=output_dir,
        config=vars(cfg),
        use_wandb=cfg.use_wandb,
        use_tb=cfg.use_tb,
        wandb_project=cfg.wandb_project,
        wandb_run_name=cfg.wandb_run_name,
        wandb_entity=cfg.wandb_entity,
    )

    # Model.
    start_epoch = 0
    resume_ckpt = None

    if cfg.resume_from is not None:
        LOGGER.info("Resuming from %s", cfg.resume_from)
        if cfg.backbone == "sofia":
            model, resume_ckpt = SOFIAFusionEncoder.load_checkpoint(
                cfg.resume_from,
                lrsclip_path=cfg.lrsclip_path,
                device=cfg.device,
            )
        elif cfg.backbone == "sofia_v1":
            model, resume_ckpt = SOFIAv1FusionEncoder.load_checkpoint(
                cfg.resume_from,
                lrsclip_path=cfg.lrsclip_path,
                device=cfg.device,
            )
        else:
            model, resume_ckpt = AsymmetricEncoder.load_checkpoint(
                cfg.resume_from,
                dino_web_path=cfg.dino_web_path,
                dino_sat_path=cfg.dino_sat_path,
                lrsclip_path=cfg.lrsclip_path,
                device=cfg.device,
            )
        start_epoch = resume_ckpt.get("epoch", -1) + 1
    else:
        mode_str = "baseline (no text)" if cfg.baseline_mode else "with text (L1/L2/L3)"
        if cfg.backbone == "sofia":
            enc_str = f"SOFIA-{cfg.sofia_preset} (text-FiLM uav={cfg.sofia_use_text_film_uav}, sat={cfg.sofia_use_text_film_sat})"
        elif cfg.backbone == "sofia_v1":
            enc_str = f"SOFIAv1-{cfg.sofia_v1_variant} (StripNet+DCNv4, text-FiLM uav={cfg.sofia_v1_use_text_film_uav}, sat={cfg.sofia_v1_use_text_film_sat})"
        elif cfg.backbone == "stripnet":
            enc_str = "StripNet-small (shared, 512→1024 proj)"
        else:
            enc_str = "shared DINOv3 WEB" if cfg.shared_encoder else "asymmetric (WEB + SAT)"
        LOGGER.info("Building model — %s, %s", mode_str, enc_str)
        if cfg.backbone == "sofia":
            preset_map = {"Tiny": sofia_tiny_config, "M": sofia_m_config, "L": sofia_l_config}
            if cfg.sofia_preset not in preset_map:
                raise ValueError(f"Unknown sofia_preset={cfg.sofia_preset!r}")
            sofia_cfg = preset_map[cfg.sofia_preset]()
            sofia_cfg.d_descriptor = cfg.sofia_d_descriptor
            sofia_cfg.text_film_dim = cfg.sofia_d_descriptor
            sofia_cfg.use_text_film_uav = cfg.sofia_use_text_film_uav and not cfg.baseline_mode
            sofia_cfg.use_text_film_sat = cfg.sofia_use_text_film_sat and not cfg.baseline_mode
            sofia_cfg.mamba_variant = cfg.sofia_mamba_variant
            sofia_cfg.mamba_backend = cfg.sofia_mamba_backend
            sofia_cfg.use_evss_bridge = cfg.sofia_use_evss_bridge
            sofia_cfg.evss_bridge_locations = list(cfg.sofia_evss_bridge_locations)
            model = SOFIAFusionEncoder(
                sofia_cfg=sofia_cfg,
                lrsclip_path=cfg.lrsclip_path,
                init_gate=cfg.init_gate,
                baseline_mode=cfg.baseline_mode,
                lora_rank=cfg.sofia_lora_rank,
                device=cfg.device,
            ).to(cfg.device)
        elif cfg.backbone == "sofia_v1":
            sofia_v1_cfg = SOFIAv1Config(
                variant=cfg.sofia_v1_variant,
                dcn_variant=cfg.sofia_v1_dcn_variant,
                d_descriptor=cfg.sofia_v1_d_descriptor,
                text_film_dim=cfg.sofia_v1_d_descriptor,
                use_text_film_uav=cfg.sofia_v1_use_text_film_uav and not cfg.baseline_mode,
                use_text_film_sat=cfg.sofia_v1_use_text_film_sat and not cfg.baseline_mode,
                use_film_altitude=cfg.sofia_v1_use_film_altitude,
            )
            model = SOFIAv1FusionEncoder(
                sofia_cfg=sofia_v1_cfg,
                lrsclip_path=cfg.lrsclip_path,
                init_gate=cfg.init_gate,
                baseline_mode=cfg.baseline_mode,
                lora_rank=cfg.sofia_v1_lora_rank,
                device=cfg.device,
            ).to(cfg.device)
        else:
            model = AsymmetricEncoder(
                dino_web_path=cfg.dino_web_path,
                dino_sat_path=cfg.dino_sat_path,
                lrsclip_path=cfg.lrsclip_path,
                init_gate=cfg.init_gate,
                baseline_mode=cfg.baseline_mode,
                shared_encoder=cfg.shared_encoder,
                mona_bottleneck=cfg.mona_bottleneck,
                mona_last_n_blocks=cfg.mona_last_n_blocks,
                device=cfg.device,
                backbone=cfg.backbone,
                stripnet_path=cfg.stripnet_path,
                stripnet_mona_last_n_stages=cfg.stripnet_mona_last_n_stages,
                stripnet_freeze=cfg.stripnet_freeze,
            ).to(cfg.device)
        LOGGER.info("embed_dim=%d", model.embed_dim)

    # --- Gradient checkpointing (trade compute for VRAM) ---
    # StripNet/SOFIA don't expose set_gradient_checkpointing — only DGTRS gets it.
    if cfg.gradient_checkpointing and cfg.backbone == "dinov3":
        if cfg.shared_encoder:
            model.image_encoder.set_gradient_checkpointing(True)
        else:
            model.drone_encoder.set_gradient_checkpointing(True)
            model.sat_encoder.set_gradient_checkpointing(True)
        if model.text_encoder is not None:
            model.text_encoder.transformer.gradient_checkpointing = True
        LOGGER.info("Gradient checkpointing enabled (DINOv3 + DGTRS)")
    elif cfg.gradient_checkpointing and cfg.backbone in ("stripnet", "sofia", "sofia_v1"):
        if model.text_encoder is not None:
            model.text_encoder.transformer.gradient_checkpointing = True
        LOGGER.info("Gradient checkpointing enabled (DGTRS only; %s doesn't support)", cfg.backbone)

    n_trainable = sum(p.numel() for p in model.trainable_parameters())
    n_total = sum(p.numel() for p in model.parameters())
    LOGGER.info(
        "trainable=%s (%.2f%%) total=%s",
        f"{n_trainable:,}", 100.0 * n_trainable / max(n_total, 1), f"{n_total:,}",
    )

    # --- Model summary (torchinfo) ---
    model_summary = print_model_summary(model, device=cfg.device)
    (output_dir / "model_summary.txt").write_text(model_summary)

    # --- W&B model watching (gradient + weight histograms) ---
    if tracker.has_wandb:
        tracker.watch_model(model, log_freq=50)

    # Loss.
    if cfg.loss_type == "symmetric":
        loss_fn = InfoNCELoss(
            temperature_init=cfg.tau_init,
            learnable_temperature=cfg.learnable_temperature,
            label_smoothing=cfg.label_smoothing,
            weight_q2g=cfg.weight_q2g,
            weight_g2q=cfg.weight_g2q,
        )
        loss_name = "SymmetricInfoNCE"
    elif cfg.loss_type == "weighted":
        loss_fn = WeightedInfoNCELoss(
            temperature_init=cfg.tau_init,
            learnable_temperature=cfg.learnable_temperature,
            label_smoothing=cfg.label_smoothing,
        )
        loss_name = "WeightedInfoNCE"
    else:
        raise ValueError(f"Unknown loss_type={cfg.loss_type!r} (expected 'symmetric' or 'weighted')")

    LOGGER.info(
        "Loss: %s  Temperature: %s (init=%.3f)  q2g=%.2f g2q=%.2f",
        loss_name,
        "learnable" if cfg.learnable_temperature else "fixed",
        cfg.tau_init, cfg.weight_q2g, cfg.weight_g2q,
    )

    # Hard negative memory bank.
    neg_bank = None
    if cfg.neg_bank_size > 0:
        neg_bank = NegativeMemoryBank(size=cfg.neg_bank_size, dim=model.embed_dim).to(cfg.device)
        LOGGER.info("Negative memory bank: size=%d, dim=%d", cfg.neg_bank_size, model.embed_dim)

    # Data — separate transforms for train (augmented) and eval (clean).
    drone_train_tf = get_drone_train_transform(image_size=256)
    sat_train_tf = get_satellite_train_transform(image_size=256)
    eval_tf = get_dino_transform(image_size=256)

    train_ds = GTAUAVDataset(
        pair_json=cfg.train_json,
        rgb_root=cfg.rgb_root,
        caption_root=cfg.caption_root,
        drone_transform=drone_train_tf,
        sat_transform=sat_train_tf,
        filter_meta=cfg.filter_meta,
    )
    test_ds = GTAUAVDataset(
        pair_json=cfg.test_json,
        rgb_root=cfg.rgb_root,
        caption_root=cfg.caption_root,
        image_transform=eval_tf,
        filter_meta=cfg.filter_meta,
    )

    sat_cand_list = [entry["sat_candidates"] for entry in train_ds.entries]
    # Backward compat: `use_mutex_sampler=False` overrides to plain shuffle.
    effective_sampler_type = cfg.sampler_type if cfg.use_mutex_sampler else "none"

    if effective_sampler_type == "dss":
        batch_sampler = DynamicSimilaritySampler(
            sat_cand_list, batch_size=cfg.batch_size, shuffle=True, seed=cfg.seed,
            knn_device=cfg.dss_knn_device,
            use_lsh=cfg.dss_use_lsh,
            lsh_num_tables=cfg.dss_lsh_num_tables,
            lsh_num_bits=cfg.dss_lsh_num_bits,
        )
        LOGGER.info(
            "Sampler: DynamicSimilarity — kNN on %s%s, warmup=%d, re-embed every %d epochs",
            cfg.dss_knn_device,
            " + LSH" if cfg.dss_use_lsh else "",
            cfg.dss_warmup_epochs, cfg.dss_reembed_every,
        )
    elif effective_sampler_type == "mutex":
        batch_sampler = MutuallyExclusiveSampler(
            sat_cand_list, batch_size=cfg.batch_size, shuffle=True, seed=cfg.seed,
        )
        LOGGER.info("Sampler: MutuallyExclusive — no false negatives within a batch")
    else:
        batch_sampler = None
        LOGGER.info("Sampler: default shuffle (no mutex / no DSS)")

    if batch_sampler is not None:
        train_loader = DataLoader(
            train_ds,
            batch_sampler=batch_sampler,
            num_workers=cfg.num_workers,
            collate_fn=collate_gtauav_batch,
            pin_memory=True,
        )
    else:
        train_loader = DataLoader(
            train_ds,
            batch_size=cfg.batch_size,
            shuffle=True,
            num_workers=cfg.num_workers,
            collate_fn=collate_gtauav_batch,
            pin_memory=True,
            drop_last=True,
        )

    emb_cache: EmbeddingCache | None = None
    if cfg.dss_cache_dir is not None:
        emb_cache = EmbeddingCache(cfg.dss_cache_dir)
        LOGGER.info("DSS embedding cache: %s", cfg.dss_cache_dir)
    test_loader = DataLoader(
        test_ds,
        batch_size=cfg.batch_size,
        shuffle=False,
        num_workers=cfg.num_workers,
        collate_fn=collate_gtauav_batch,
        pin_memory=True,
    )
    # Train eval loader: clean transforms (no augmentation), for R@K on train set.
    train_eval_ds = GTAUAVDataset(
        pair_json=cfg.train_json,
        rgb_root=cfg.rgb_root,
        caption_root=cfg.caption_root,
        image_transform=eval_tf,
        filter_meta=cfg.filter_meta,
    )
    train_eval_loader = DataLoader(
        train_eval_ds,
        batch_size=cfg.batch_size,
        shuffle=False,
        num_workers=cfg.num_workers,
        collate_fn=collate_gtauav_batch,
        pin_memory=True,
    )

    effective_batch = cfg.batch_size * cfg.grad_accum_steps
    LOGGER.info(
        "train=%d  test=%d  batch=%d  accum=%d  effective_batch=%d",
        len(train_ds), len(test_ds), cfg.batch_size, cfg.grad_accum_steps, effective_batch,
    )

    # Optimizer — per-group LR (text encoder gets lower LR).
    param_groups = _build_param_groups(
        model, cfg.learning_rate, cfg.text_lr_factor,
        stripnet_backbone_lr_factor=cfg.stripnet_backbone_lr_factor,
    )
    # Include loss temperature if learnable.
    if cfg.learnable_temperature and loss_fn.logit_scale is not None:
        param_groups[0]["params"].append(loss_fn.logit_scale)

    optimizer = AdamW(param_groups, weight_decay=cfg.weight_decay)

    lr_info = f"proj={cfg.learning_rate:.0e}"
    if not cfg.baseline_mode:
        lr_info += f"  text={cfg.learning_rate * cfg.text_lr_factor:.0e}"
    LOGGER.info("Optimizer: AdamW  LR: %s  warmup=%d epochs", lr_info, cfg.warmup_epochs)

    # Scheduler — cosine with linear warmup (counted in optimizer steps).
    steps_per_epoch = math.ceil(len(train_loader) / cfg.grad_accum_steps)
    total_steps = cfg.epochs * steps_per_epoch
    warmup_steps = cfg.warmup_epochs * steps_per_epoch
    with warnings.catch_warnings():
        warnings.filterwarnings("ignore", message=".*lr_scheduler.step.*optimizer.step.*")
        scheduler = LambdaLR(
            optimizer,
            lr_lambda=_cosine_warmup_schedule(warmup_steps, total_steps),
            last_epoch=-1,
        )
    scaler = GradScaler(enabled=cfg.use_amp)

    # Restore optimizer/scheduler/loss state on resume.
    if resume_ckpt is not None:
        if "optimizer_state" in resume_ckpt:
            optimizer.load_state_dict(resume_ckpt["optimizer_state"])
            LOGGER.info("Optimizer state restored")
        if "loss_state" in resume_ckpt:
            loss_fn.load_state_dict(resume_ckpt["loss_state"])
            LOGGER.info("Loss state restored (tau=%.4f)", loss_fn.current_temperature)
        # Set scheduler last_epoch so it resumes at the correct LR.
        scheduler.last_epoch = start_epoch * steps_per_epoch
        LOGGER.info("Resuming from epoch %d", start_epoch)

    history: list[dict] = []
    csv_logger = CSVLogger(output_dir)

    # --- Optional profiler (first epoch only) ---
    profiler = None
    if cfg.use_profiler and start_epoch == 0:
        profiler = TrainingProfiler(
            output_dir=output_dir,
            n_warmup=cfg.profiler_warmup,
            n_active=cfg.profiler_active,
        )
        profiler.start()

    LOGGER.info("Starting training for %d epochs (from epoch %d)", cfg.epochs, start_epoch)

    global_step = start_epoch * steps_per_epoch
    best_r1 = 0.0

    for epoch in range(start_epoch, cfg.epochs):
        model.train()
        if batch_sampler is not None:
            batch_sampler.set_epoch(epoch)

        # DSS re-embedding: refresh query embeddings before the epoch starts.
        if (
            isinstance(batch_sampler, DynamicSimilaritySampler)
            and epoch >= cfg.dss_warmup_epochs
            and (epoch - cfg.dss_warmup_epochs) % cfg.dss_reembed_every == 0
        ):
            query_embs: torch.Tensor | None = None
            if emb_cache is not None:
                query_embs = emb_cache.load(epoch)
            if query_embs is None:
                LOGGER.info("DSS: re-embedding %d train queries (epoch=%d)", len(train_ds), epoch)
                t_embed = time.time()
                query_embs = _embed_drone_queries(
                    model, train_ds, cfg.device,
                    batch_size=cfg.batch_size * cfg.grad_accum_steps,
                    num_workers=cfg.num_workers,
                )
                LOGGER.info("DSS: re-embed took %.1fs", time.time() - t_embed)
                if emb_cache is not None:
                    emb_cache.save(epoch, query_embs)
            t_sampler = time.time()
            batch_sampler.update_embeddings(query_embs)
            LOGGER.info("DSS: sampler update_embeddings took %.2fs", time.time() - t_sampler)

        epoch_start = time.time()
        agg: dict[str, float] = {}
        n_batches = 0

        pbar = tqdm(
            train_loader,
            desc=f"  Epoch {epoch + 1}/{cfg.epochs}",
            unit="batch",
            leave=False,
        )
        accum = cfg.grad_accum_steps
        for batch in pbar:
            # Zero gradients only at the start of each accumulation window.
            if n_batches % accum == 0:
                optimizer.zero_grad(set_to_none=True)

            drone_img = batch["drone_img"].to(cfg.device, non_blocking=True)
            sat_img = batch["sat_img"].to(cfg.device, non_blocking=True)
            altitude = batch.get("altitude")
            if altitude is not None:
                altitude = altitude.to(cfg.device, non_blocking=True)

            # Model forward in AMP (fp16 for DINOv3/DGTRS encoders).
            with autocast(device_type="cuda", enabled=cfg.use_amp):
                if cfg.baseline_mode:
                    embeddings = model(drone_img=drone_img, sat_img=sat_img, altitude=altitude)
                else:
                    embeddings = model(
                        drone_img=drone_img,
                        sat_img=sat_img,
                        caption_l1=batch["caption_l1"],
                        caption_l2=batch["caption_l2"],
                        caption_l3=batch["caption_l3"],
                        sat_caption_l1=batch["sat_caption_l1"],
                        sat_caption_l2=batch["sat_caption_l2"],
                        sat_caption_l3=batch["sat_caption_l3"],
                        altitude=altitude,
                    )
            # Loss — InfoNCE or WeightedInfoNCE. Only the latter uses positive_weights.
            queue_neg = neg_bank.get_queue() if neg_bank is not None else None
            loss_kwargs = {
                "embeddings": embeddings,
                "epoch": epoch,
                "total_epochs": cfg.epochs,
                "queue_negatives": queue_neg,
            }
            if isinstance(loss_fn, WeightedInfoNCELoss):
                loss_kwargs["positive_weights"] = batch["positive_weights"].to(
                    cfg.device, non_blocking=True,
                )
            loss_dict = loss_fn(**loss_kwargs)

            # Scale loss by accumulation steps so gradients average correctly.
            raw_loss = float(loss_dict["total"].item())  # save before backward
            total_loss = loss_dict["total"] / accum
            scaler.scale(total_loss).backward()

            # Enqueue current gallery AFTER backward. The queue buffer is aliased
            # into the autograd graph through `queue_neg` (a view returned by
            # `NegativeMemoryBank.get_queue`), so modifying it before backward
            # triggers "variable needed for gradient computation has been modified
            # by an inplace operation". Enqueueing here is semantically identical
            # — the next step's queue state is the same either way.
            if neg_bank is not None:
                neg_bank.enqueue(embeddings["gallery"].detach())

            # Optimizer step only after accumulating `accum` micro-batches.
            is_accum_step = (n_batches + 1) % accum == 0 or (n_batches + 1) == len(train_loader)
            if is_accum_step:
                if cfg.grad_clip > 0:
                    scaler.unscale_(optimizer)
                    nn.utils.clip_grad_norm_(
                        model.trainable_parameters(),
                        max_norm=cfg.grad_clip,
                    )

                # --- Gradient monitoring (after unscale, before step) ---
                if cfg.log_grad_norms and n_batches % (50 * accum) < accum:
                    grad_norms = compute_gradient_norms(model, loss_fn)
                    tracker.log_gradients(epoch, grad_norms, step=global_step)
                    if n_batches < accum:
                        log_gradient_summary(grad_norms)

                scaler.step(optimizer)
                scaler.update()
                with warnings.catch_warnings():
                    warnings.filterwarnings("ignore", message=".*lr_scheduler.step.*optimizer.step.*")
                    scheduler.step()
                global_step += 1

            # --- Per-batch tracking (log unscaled loss) ---
            step_metrics = {
                "loss": raw_loss,
                "temperature": float(loss_dict["temperature"].item()),
                "gate_q": float(loss_dict["gate_q"].item()),
                "gate_g": float(loss_dict["gate_g"].item()),
                "lr": optimizer.param_groups[0]["lr"],
            }
            tracker.log_train(epoch, step_metrics, step=global_step)
            csv_logger.log_batch(epoch, n_batches, global_step, step_metrics)

            for key, val in loss_dict.items():
                agg[key] = agg.get(key, 0.0) + float(val.item())
            n_batches += 1

            pbar.set_postfix(
                loss=f"{raw_loss:.3f}",
                tau=f"{step_metrics['temperature']:.4f}",
                gq=f"{step_metrics['gate_q']:.3f}",
                gg=f"{step_metrics['gate_g']:.3f}",
            )

            # --- Profiler step ---
            if profiler is not None:
                profiler.step()
                if profiler.is_done(n_batches):
                    profiler.export()
                    profiler = None

        elapsed = time.time() - epoch_start

        means = {k: v / max(n_batches, 1) for k, v in agg.items()}
        LOGGER.info(
            "epoch=%d  time=%.1fs  lr=%.2e  loss=%.4f  tau=%.4f  gate_q=%.4f  gate_g=%.4f",
            epoch, elapsed,
            optimizer.param_groups[0]["lr"],
            means.get("total", 0.0),
            means.get("temperature", 0.0),
            means.get("gate_q", 1.0),
            means.get("gate_g", 1.0),
        )

        epoch_record: dict = {
            "epoch": epoch,
            "elapsed_seconds": elapsed,
            "train": means,
        }

        # --- Log VRAM usage ---
        if torch.cuda.is_available():
            vram_gb = torch.cuda.max_memory_allocated() / 1e9
            tracker.log_scalar("system/vram_peak_gb", vram_gb, step=global_step)

        # Evaluation.
        train_recall = {}
        if (epoch + 1) % cfg.eval_every == 0 or epoch == cfg.epochs - 1:
            # Train R@K (subset — same size as test set for speed).
            train_eval_batches = len(test_loader)
            train_recall = _evaluate(
                model, train_eval_loader, cfg.device,
                loss_fn=loss_fn, epoch=epoch, total_epochs=cfg.epochs,
                max_batches=train_eval_batches, desc="eval-train",
            )
            epoch_record["train_recall"] = train_recall
            csv_logger.log_train_recall(epoch, train_recall)
            tracker.log_train(epoch, {f"recall/{k}": v for k, v in train_recall.items() if k.startswith("r@")}, step=global_step)

        # Log train metrics to CSV (includes recall/AP if eval ran this epoch).
        train_row = {**means}
        if "total" in train_row:
            train_row["train_loss"] = train_row.pop("total")
        if train_recall:
            train_row["r@1_q2g"] = train_recall.get("r@1_q2g", 0.0)
            train_row["r@5_q2g"] = train_recall.get("r@5_q2g", 0.0)
            train_row["r@10_q2g"] = train_recall.get("r@10_q2g", 0.0)
            train_row["ap_q2g"] = train_recall.get("ap_q2g", 0.0)
            train_row["r@1_g2q"] = train_recall.get("r@1_g2q", 0.0)
            train_row["r@5_g2q"] = train_recall.get("r@5_g2q", 0.0)
            train_row["r@10_g2q"] = train_recall.get("r@10_g2q", 0.0)
            train_row["ap_g2q"] = train_recall.get("ap_g2q", 0.0)
        csv_logger.log_train(epoch, train_row, optimizer.param_groups[0]["lr"], elapsed)
        generate_plots(csv_logger.log_dir)

        if train_recall:
            LOGGER.info(
                "train-recall epoch=%d  q2g R@1=%.4f R@5=%.4f R@10=%.4f AP=%.4f  "
                "g2q R@1=%.4f R@5=%.4f R@10=%.4f AP=%.4f  loss=%.4f",
                epoch,
                train_recall.get("r@1_q2g", 0.0),
                train_recall.get("r@5_q2g", 0.0),
                train_recall.get("r@10_q2g", 0.0),
                train_recall.get("ap_q2g", 0.0),
                train_recall.get("r@1_g2q", 0.0),
                train_recall.get("r@5_g2q", 0.0),
                train_recall.get("r@10_g2q", 0.0),
                train_recall.get("ap_g2q", 0.0),
                train_recall.get("loss", 0.0),
            )

            # Val R@K (full test set).
            val_metrics = _evaluate(
                model, test_loader, cfg.device,
                loss_fn=loss_fn, epoch=epoch, total_epochs=cfg.epochs,
                desc="eval-val",
            )
            epoch_record["val"] = val_metrics
            csv_logger.log_val(epoch, val_metrics)
            generate_plots(csv_logger.log_dir)
            tracker.log_val(epoch, val_metrics, step=global_step)

            # Track best R@1.
            r1 = val_metrics.get("r@1_q2g", 0.0)
            if r1 > best_r1:
                best_r1 = r1
                tracker.log_scalar("val/best_r@1_q2g", best_r1, step=global_step)

            LOGGER.info(
                "val epoch=%d  q2g R@1=%.4f R@5=%.4f R@10=%.4f AP=%.4f  "
                "g2q R@1=%.4f R@5=%.4f R@10=%.4f AP=%.4f  loss=%.4f  gate_q=%.4f",
                epoch,
                val_metrics.get("r@1_q2g", 0.0),
                val_metrics.get("r@5_q2g", 0.0),
                val_metrics.get("r@10_q2g", 0.0),
                val_metrics.get("ap_q2g", 0.0),
                val_metrics.get("r@1_g2q", 0.0),
                val_metrics.get("r@5_g2q", 0.0),
                val_metrics.get("r@10_g2q", 0.0),
                val_metrics.get("ap_g2q", 0.0),
                val_metrics.get("loss", 0.0),
                val_metrics.get("gate_q", 1.0),
            )

            # --- Grad-CAM visualization ---
            if cfg.use_gradcam and (epoch + 1) % cfg.gradcam_every == 0:
                from src.training.gradcam import generate_gradcam_samples
                overlays = generate_gradcam_samples(
                    model=model,
                    dataloader=test_loader,
                    device=cfg.device,
                    output_dir=str(output_dir),
                    n_samples=cfg.gradcam_samples,
                    epoch=epoch,
                )
                # Log first few overlays to tracker.
                for i, overlay in enumerate(overlays[:4]):
                    kind = "drone" if i % 2 == 0 else "sat"
                    tracker.log_image(
                        f"gradcam/{kind}_{i//2}",
                        overlay,
                        step=global_step,
                        caption=f"Epoch {epoch} {kind} Grad-CAM",
                    )

        history.append(epoch_record)

        # Save checkpoint. Model architecture flags go into the ckpt so
        # `AsymmetricEncoder.load_checkpoint` (or `SOFIAFusionEncoder.load_checkpoint`)
        # can rebuild the right shape.
        ckpt_obj = {
            "epoch": epoch,
            "model_state": model.state_dict(),
            "optimizer_state": optimizer.state_dict(),
            "loss_state": loss_fn.state_dict(),
            "baseline_mode": cfg.baseline_mode,
            "backbone": cfg.backbone,
        }
        if cfg.backbone in ("sofia", "sofia_v1"):
            ckpt_obj["sofia_cfg"] = model.sofia_cfg
        else:
            ckpt_obj["shared_encoder"] = cfg.shared_encoder
            ckpt_obj["mona_bottleneck"] = cfg.mona_bottleneck
            ckpt_obj["mona_last_n_blocks"] = cfg.mona_last_n_blocks
        _atomic_save(obj=ckpt_obj, path=output_dir / f"ckpt_epoch{epoch:03d}.pt")
        LOGGER.info("Checkpoint saved: ckpt_epoch%03d.pt", epoch)

    # Save history.
    history_path = output_dir / "history.json"
    with history_path.open("w", encoding="utf-8") as f:
        json.dump(history, f, indent=2)

    # Save final eval report.
    LOGGER.info("Running final evaluation...")
    final_metrics = _evaluate(
        model, test_loader, cfg.device,
        loss_fn=loss_fn, epoch=cfg.epochs - 1, total_epochs=cfg.epochs,
    )
    report = {
        "config": vars(cfg),
        "metrics": final_metrics,
        "history": history,
    }
    report_path = output_dir / "eval_report.json"
    with report_path.open("w", encoding="utf-8") as f:
        json.dump(report, f, indent=2)

    # --- Log final summary to W&B ---
    tracker.log_summary({
        "best_r@1_q2g": best_r1,
        "final_r@1_q2g": final_metrics.get("r@1_q2g", 0.0),
        "final_r@5_q2g": final_metrics.get("r@5_q2g", 0.0),
        "final_r@10_q2g": final_metrics.get("r@10_q2g", 0.0),
        "final_ap_q2g": final_metrics.get("ap_q2g", 0.0),
        "final_r@1_g2q": final_metrics.get("r@1_g2q", 0.0),
        "final_r@5_g2q": final_metrics.get("r@5_g2q", 0.0),
        "final_r@10_g2q": final_metrics.get("r@10_g2q", 0.0),
        "final_ap_g2q": final_metrics.get("ap_g2q", 0.0),
        "final_gate_q": final_metrics.get("gate_q", 1.0),
        "final_gate_g": final_metrics.get("gate_g", 1.0),
    })

    # --- Cleanup profiler if still running ---
    if profiler is not None:
        profiler.export()

    tracker.close()

    LOGGER.info("Training complete. Report: %s", report_path)
    LOGGER.info(
        "Final — q2g R@1=%.4f R@5=%.4f R@10=%.4f AP=%.4f  "
        "g2q R@1=%.4f R@5=%.4f R@10=%.4f AP=%.4f  gate_q=%.4f  gate_g=%.4f",
        final_metrics.get("r@1_q2g", 0.0),
        final_metrics.get("r@5_q2g", 0.0),
        final_metrics.get("r@10_q2g", 0.0),
        final_metrics.get("ap_q2g", 0.0),
        final_metrics.get("r@1_g2q", 0.0),
        final_metrics.get("r@5_g2q", 0.0),
        final_metrics.get("r@10_g2q", 0.0),
        final_metrics.get("ap_g2q", 0.0),
        final_metrics.get("gate_q", 1.0),
        final_metrics.get("gate_g", 1.0),
    )


def main() -> None:
    parser = argparse.ArgumentParser(description="GTA-UAV caption test training.")
    parser.add_argument(
        "--config", type=str, default=None,
        help="Path to gin config file (e.g. conf/gtauav_balanced.gin).",
    )
    parser.add_argument(
        "--baseline", action="store_true",
        help="Run baseline mode (no text).",
    )
    parser.add_argument(
        "--resume", type=str, default=None,
        help="Path to checkpoint to resume training from.",
    )
    parser.add_argument(
        "--output-dir", type=str, default=None,
        help="Override output directory.",
    )
    parser.add_argument(
        "--filter-meta", type=str, default=None,
        help="Path to seg_filter.json for excluding bad images.",
    )
    parser.add_argument(
        "--batch-size", type=int, default=None,
        help="Batch size.",
    )
    parser.add_argument(
        "--grad-accum", type=int, default=None,
        help="Gradient accumulation steps (effective_batch = batch_size * accum).",
    )
    parser.add_argument(
        "--epochs", type=int, default=None,
        help="Number of epochs.",
    )
    parser.add_argument(
        "--lr", type=float, default=None,
        help="Learning rate for projections.",
    )
    parser.add_argument(
        "--text-lr-factor", type=float, default=None,
        help="Text encoder LR = lr * factor (default 0.1 = 10x lower).",
    )
    parser.add_argument(
        "--warmup-epochs", type=int, default=None,
        help="Linear warmup epochs.",
    )
    parser.add_argument(
        "--init-gate", type=float, default=None,
        help="Initial gate value (image weight).",
    )
    # Tracking flags.
    parser.add_argument("--wandb", action="store_true", help="Enable W&B tracking.")
    parser.add_argument("--no-tb", action="store_true", help="Disable TensorBoard.")
    parser.add_argument("--gradcam", action="store_true", help="Enable Grad-CAM visualization.")
    parser.add_argument("--profile", action="store_true", help="Enable PyTorch profiler (first epoch).")
    parser.add_argument("--no-grad-norms", action="store_true", help="Disable gradient norm logging.")
    # Gin overrides.
    parser.add_argument(
        "--gin-param", type=str, nargs="*", default=[],
        help="Gin parameter overrides (e.g. 'TrainConfigGTAUAV.epochs=30').",
    )
    args = parser.parse_args()

    # Parse gin config if provided.
    if args.config is not None:
        gin.parse_config_file(args.config)
    if args.gin_param:
        gin.parse_config(args.gin_param)

    # Create config (gin bindings apply via @gin.configurable).
    cfg = TrainConfigGTAUAV()

    # CLI overrides take priority over gin.
    if args.baseline:
        cfg.baseline_mode = True
    if args.resume is not None:
        cfg.resume_from = args.resume
    if args.batch_size is not None:
        cfg.batch_size = args.batch_size
    if args.grad_accum is not None:
        cfg.grad_accum_steps = args.grad_accum
    if args.epochs is not None:
        cfg.epochs = args.epochs
    if args.lr is not None:
        cfg.learning_rate = args.lr
    if args.text_lr_factor is not None:
        cfg.text_lr_factor = args.text_lr_factor
    if args.warmup_epochs is not None:
        cfg.warmup_epochs = args.warmup_epochs
    if args.init_gate is not None:
        cfg.init_gate = args.init_gate
    if args.filter_meta is not None:
        cfg.filter_meta = args.filter_meta

    # Tracking overrides.
    if args.wandb:
        cfg.use_wandb = True
    if args.no_tb:
        cfg.use_tb = False
    if args.gradcam:
        cfg.use_gradcam = True
    if args.profile:
        cfg.use_profiler = True
    if args.no_grad_norms:
        cfg.log_grad_norms = False

    if args.output_dir is not None:
        cfg.output_dir = args.output_dir
    elif args.baseline and args.output_dir is None:
        cfg.output_dir = "out/gtauav/baseline"

    train(cfg)


if __name__ == "__main__":
    main()
