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).
"""

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

import coloredlogs
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.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"


@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
    proj_dim: int = 512
    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 = 64
    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


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="  🔎 Evaluating", 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"],
            )
        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"] = model.fusion.gate_value
    return metrics


def train(cfg: TrainConfigGTAUAV) -> None:
    """Run full training loop."""
    coloredlogs.install(
        level="INFO",
        logger=LOGGER,
        fmt="%(asctime)s %(name)s %(levelname)s %(message)s",
    )
    _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)

    # 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,
            proj_dim=cfg.proj_dim,
            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:,}",
    )

    # 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] = []

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

    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}/{cfg.epochs - 1}",
            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)

            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"],
                    )
                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()
            with warnings.catch_warnings():
                warnings.filterwarnings("ignore", message=".*lr_scheduler.step.*optimizer.step.*")
                scheduler.step()

            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"{total_loss.item():.3f}",
                tau=f"{loss_dict['temperature'].item():.4f}",
                gate=f"{loss_dict['gate'].item():.3f}",
            )

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

        # 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
            LOGGER.info(
                "🎯 val epoch=%d  R@1=%.4f  R@5=%.4f  R@10=%.4f  gate=%.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", 1.0),
            )

        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)

    LOGGER.info("✅ Training complete. Report: %s", report_path)
    LOGGER.info(
        "📊 Final — R@1=%.4f  R@5=%.4f  R@10=%.4f  gate=%.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", 1.0),
    )


def main() -> None:
    parser = argparse.ArgumentParser(description="GTA-UAV caption test training.")
    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=64,
        help="Batch size.",
    )
    parser.add_argument(
        "--epochs", type=int, default=10,
        help="Number of epochs.",
    )
    parser.add_argument(
        "--lr", type=float, default=1e-4,
        help="Learning rate for projections.",
    )
    parser.add_argument(
        "--text-lr-factor", type=float, default=0.1,
        help="Text encoder LR = lr * factor (default 0.1 = 10x lower).",
    )
    parser.add_argument(
        "--warmup-epochs", type=int, default=2,
        help="Linear warmup epochs.",
    )
    parser.add_argument(
        "--init-gate", type=float, default=0.7,
        help="Initial gate value (image weight).",
    )
    args = parser.parse_args()

    cfg = TrainConfigGTAUAV()
    cfg.baseline_mode = args.baseline
    cfg.resume_from = args.resume
    cfg.batch_size = args.batch_size
    cfg.epochs = args.epochs
    cfg.learning_rate = args.lr
    cfg.text_lr_factor = args.text_lr_factor
    cfg.warmup_epochs = args.warmup_epochs
    cfg.init_gate = args.init_gate

    if args.filter_meta is not None:
        cfg.filter_meta = args.filter_meta

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

    train(cfg)


if __name__ == "__main__":
    main()
