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
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, collate_gtauav_batch
from src.losses.multi_infonce import InfoNCELoss
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,
)

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

    # 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
    use_amp: bool = True
    eval_every: int = 2
    warmup_epochs: int = 2
    seed: int = 42
    device: str = "cuda"

    # Loss.
    tau_init: float = 0.07
    label_smoothing: float = 0.1
    weight_q2g: float = 0.6
    weight_g2q: float = 0.4
    learnable_temperature: 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: AsymmetricEncoder,
    lr: float,
    text_lr_factor: float,
) -> list[dict]:
    """Build optimizer param groups with separate LR for text encoder."""
    text_params = []
    other_params = []

    for name, param in model.named_parameters():
        if not param.requires_grad:
            continue
        if "text_encoder" in name:
            text_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})

    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 _evaluate(
    model: AsymmetricEncoder,
    loader: DataLoader,
    device: str,
    k_values: tuple[int, ...] = (1, 5, 10),
) -> dict[str, float]:
    """Compute R@K on validation set."""
    model.eval()
    all_query: list[torch.Tensor] = []
    all_gallery: list[torch.Tensor] = []

    for batch in tqdm(loader, desc="  eval", unit="batch", leave=False):
        drone_img = batch["drone_img"].to(device, non_blocking=True)
        sat_img = batch["sat_img"].to(device, non_blocking=True)

        if model.baseline_mode:
            embeddings = model(drone_img=drone_img, sat_img=sat_img)
        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"],
            )
        all_query.append(embeddings["query"].cpu())
        all_gallery.append(embeddings["gallery"].cpu())

    query = torch.cat(all_query, dim=0)
    gallery = torch.cat(all_gallery, dim=0)

    sim = query @ gallery.t()
    n = sim.size(0)
    targets = torch.arange(n)

    metrics: dict[str, float] = {}
    sorted_idx = sim.argsort(dim=1, descending=True)
    for k in k_values:
        top_k = sorted_idx[:, :k]
        hit = (top_k == targets.unsqueeze(1)).any(dim=1).float()
        metrics[f"r@{k}_q2g"] = float(hit.mean().item())

    sorted_idx_g2q = sim.t().argsort(dim=1, descending=True)
    for k in k_values:
        top_k = sorted_idx_g2q[:, :k]
        hit = (top_k == targets.unsqueeze(1)).any(dim=1).float()
        metrics[f"r@{k}_g2q"] = float(hit.mean().item())

    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}_train.csv      — per-epoch summary
        {output_dir}/logs/epoch_{N}_val.csv        — per-epoch val
        {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.train_rows: list[dict] = []
        self.val_rows: list[dict] = []
        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

    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."""
        row = {"epoch": epoch, "lr": lr, "elapsed_s": round(elapsed, 1), **metrics}
        self.train_rows.append(row)
        pd.DataFrame(self.train_rows).to_csv(self.log_dir / "train.csv", index=False)
        pd.DataFrame([row]).to_csv(self.log_dir / f"epoch_{epoch:03d}_train.csv", index=False)

    def log_val(self, epoch: int, metrics: dict) -> None:
        row = {"epoch": epoch, **metrics}
        self.val_rows.append(row)
        pd.DataFrame(self.val_rows).to_csv(self.log_dir / "val.csv", index=False)
        pd.DataFrame([row]).to_csv(self.log_dir / f"epoch_{epoch:03d}_val.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)
        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)"
        LOGGER.info("Building model — %s", mode_str)
        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,
            device=cfg.device,
        ).to(cfg.device)

    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.
    loss_fn = InfoNCELoss(
        temperature_init=cfg.tau_init,
        label_smoothing=cfg.label_smoothing,
        weight_q2g=cfg.weight_q2g,
        weight_g2q=cfg.weight_g2q,
        learnable_temperature=cfg.learnable_temperature,
    )
    LOGGER.info(
        "Temperature: %s (init=%.3f)",
        "learnable" if cfg.learnable_temperature else "cosine schedule",
        cfg.tau_init,
    )

    # 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,
    )

    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,
    )
    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,
    )

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

    # Optimizer — per-group LR (text encoder gets lower LR).
    param_groups = _build_param_groups(model, cfg.learning_rate, cfg.text_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.
    steps_per_epoch = len(train_loader)
    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()
        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,
        )
        for batch in pbar:
            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)

            # 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)
                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"],
                    )
            # Loss in fp32 (learnable temperature gradient overflows in fp16).
            loss_dict = loss_fn(
                embeddings=embeddings,
                epoch=epoch,
                total_epochs=cfg.epochs,
            )

            total_loss = loss_dict["total"]
            scaler.scale(total_loss).backward()

            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 == 0:
                grad_norms = compute_gradient_norms(model, loss_fn)
                tracker.log_gradients(epoch, grad_norms, step=global_step)
                if n_batches == 0:
                    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()

            # --- Per-step tracking ---
            step_metrics = {
                "loss": float(total_loss.item()),
                "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
            global_step += 1

            pbar.set_postfix(
                loss=f"{total_loss.item():.3f}",
                tau=f"{loss_dict['temperature'].item():.4f}",
                gq=f"{loss_dict['gate_q'].item():.3f}",
                gg=f"{loss_dict['gate_g'].item():.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 train metrics to CSV + generate plots every epoch.
        csv_logger.log_train(epoch, means, optimizer.param_groups[0]["lr"], elapsed)
        generate_plots(csv_logger.log_dir)

        # --- 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.
        if (epoch + 1) % cfg.eval_every == 0 or epoch == cfg.epochs - 1:
            val_metrics = _evaluate(model, test_loader, cfg.device)
            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  R@1=%.4f  R@5=%.4f  R@10=%.4f  gate_q=%.4f  gate_g=%.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("gate_q", 1.0),
                val_metrics.get("gate_g", 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.
        _atomic_save(
            obj={
                "epoch": epoch,
                "model_state": model.state_dict(),
                "optimizer_state": optimizer.state_dict(),
                "loss_state": loss_fn.state_dict(),
            },
            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)
    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_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 — R@1=%.4f  R@5=%.4f  R@10=%.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("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(
        "--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.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()
