from __future__ import annotations

"""Training loop for caption quality test on cross-view geo-localization.

GeoRSCLIP dual encoder with GatedFusion on query branch.
Single InfoNCE loss: query(drone+text) vs gallery(satellite).
"""

import argparse
import json
import logging
import time
from pathlib import Path

import gin
import torch
import torch.nn as nn
from torch.amp import GradScaler, autocast
from torch.optim import AdamW
from torch.optim.lr_scheduler import CosineAnnealingLR
from torch.utils.data import DataLoader

from src.datasets.visloc_with_captions import (
    GeoLocCaptionDataset,
    collate_caption_batch,
)
from src.eval.evaluate import evaluate_retrieval
from src.losses.multi_infonce import InfoNCELoss
from src.models.dual_encoder import DualEncoderCaptionTest

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


@gin.configurable
class TrainConfig:
    """Top-level training configuration.

    Args:
        train_query_file: Path to train_query.txt.
        val_query_file: Path to test_query.txt (used as val).
        data_root: Root of UAV-GeoLoc dataset.
        output_dir: Checkpoint and log output directory.
        epochs: Number of training epochs.
        batch_size: Mini-batch size.
        num_workers: DataLoader workers.
        learning_rate: AdamW initial LR.
        weight_decay: AdamW weight decay.
        grad_clip: Max gradient norm (0 disables).
        use_amp: Enable fp16 mixed-precision.
        eval_every: Run validation every N epochs.
        seed: Random seed.
        device: torch device.
    """

    def __init__(
        self,
        train_query_file: str = "Index/train_query.txt",
        val_query_file: str = "Index/test_query.txt",
        data_root: str = "/mnt/data1tb/cvgl_datasets/UAV-GeoLoc",
        output_dir: str = "out/caption_test",
        epochs: int = 10,
        batch_size: int = 128,
        num_workers: int = 4,
        learning_rate: float = 1e-4,
        weight_decay: float = 1e-4,
        grad_clip: float = 1.0,
        use_amp: bool = True,
        eval_every: int = 2,
        seed: int = 42,
        device: str = "cuda",
    ) -> None:
        self.train_query_file = train_query_file
        self.val_query_file = val_query_file
        self.data_root = data_root
        self.output_dir = Path(output_dir)
        self.epochs = epochs
        self.batch_size = batch_size
        self.num_workers = num_workers
        self.learning_rate = learning_rate
        self.weight_decay = weight_decay
        self.grad_clip = grad_clip
        self.use_amp = use_amp
        self.eval_every = eval_every
        self.seed = seed
        self.device = device


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 train(config_path: str) -> None:
    """Run full training loop from gin config."""
    gin.parse_config_file(config_path)
    cfg = TrainConfig()

    logging.basicConfig(
        level=logging.INFO,
        format="%(asctime)s %(name)s %(levelname)s %(message)s",
    )
    _set_seed(cfg.seed)
    cfg.output_dir.mkdir(parents=True, exist_ok=True)

    # Model + loss.
    model = DualEncoderCaptionTest().to(cfg.device)
    loss_fn = InfoNCELoss().to(cfg.device)

    preprocess = model.preprocess

    train_ds = GeoLocCaptionDataset(
        query_file=cfg.train_query_file,
        data_root=cfg.data_root,
        image_transform=preprocess,
    )
    val_ds = GeoLocCaptionDataset(
        query_file=cfg.val_query_file,
        data_root=cfg.data_root,
        image_transform=preprocess,
    )

    train_loader = DataLoader(
        train_ds,
        batch_size=cfg.batch_size,
        shuffle=True,
        num_workers=cfg.num_workers,
        collate_fn=collate_caption_batch,
        pin_memory=True,
        drop_last=True,
    )
    val_loader = DataLoader(
        val_ds,
        batch_size=cfg.batch_size,
        shuffle=False,
        num_workers=cfg.num_workers,
        collate_fn=collate_caption_batch,
        pin_memory=True,
    )

    optimizer = AdamW(
        model.trainable_parameters(),
        lr=cfg.learning_rate,
        weight_decay=cfg.weight_decay,
    )
    scheduler = CosineAnnealingLR(optimizer, T_max=cfg.epochs)
    scaler = GradScaler(enabled=cfg.use_amp)

    n_trainable = sum(p.numel() for p in model.trainable_parameters())
    n_total = sum(p.numel() for p in model.parameters())
    LOGGER.info(
        "trainable=%d (%.2f%%) total=%d train=%d val=%d",
        n_trainable, 100.0 * n_trainable / n_total,
        n_total, len(train_ds), len(val_ds),
    )

    history: list[dict] = []

    for epoch in range(cfg.epochs):
        model.train()
        epoch_start = time.time()
        agg: dict[str, float] = {}
        n_batches = 0

        for batch in train_loader:
            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)
            caption_drone = batch["caption_drone"]

            with autocast(device_type="cuda", enabled=cfg.use_amp):
                embeddings = model(
                    drone_img=drone_img,
                    sat_img=sat_img,
                    caption_drone=caption_drone,
                )
                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,
                )
            scaler.step(optimizer)
            scaler.update()

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

        scheduler.step()
        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=%.4f",
            epoch, elapsed,
            optimizer.param_groups[0]["lr"],
            means.get("total", 0.0),
            means.get("temperature", 0.0),
            means.get("gate", 1.0),
        )

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

        # Validation.
        if (epoch + 1) % cfg.eval_every == 0 or epoch == cfg.epochs - 1:
            model.eval()
            val_metrics = evaluate_retrieval(
                model=model,
                loader=val_loader,
                device=cfg.device,
            )
            epoch_record["val"] = val_metrics
            LOGGER.info(
                "val epoch=%d R@1_q2g=%.4f R@5_q2g=%.4f R@10_q2g=%.4f",
                epoch,
                val_metrics.get("r@1_query_to_gallery", 0.0),
                val_metrics.get("r@5_query_to_gallery", 0.0),
                val_metrics.get("r@10_query_to_gallery", 0.0),
            )

        history.append(epoch_record)

        _atomic_save(
            obj={
                "epoch": epoch,
                "model_state": model.state_dict(),
                "optimizer_state": optimizer.state_dict(),
                "config_path": config_path,
            },
            path=cfg.output_dir / f"ckpt_epoch{epoch:03d}.pt",
        )

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

    LOGGER.info("training complete, history saved to %s", history_path)


def main() -> None:
    parser = argparse.ArgumentParser(description="Caption quality test training.")
    parser.add_argument("--config", type=str, required=True, help="Gin config file.")
    args = parser.parse_args()
    train(config_path=args.config)


if __name__ == "__main__":
    main()
