"""Цикл обучения dual-encoder модели на GTA-UAV (symmetric fusion-архитектура).

И дрон, и спутник имеют собственное текстовое описание; слияние
(картинка+текст) происходит СИММЕТРИЧНО на обеих сторонах, каждая —
со своим экземпляром GatedFusion (веса не общие, т.к. визуальные
домены различаются).

Использование:
    python train.py \\
        --data_root /path/to/GTA-UAV \\
        --descriptions_path /path/to/descriptions.json \\
        --text_levels level1 \\
        --dgtrs_checkpoint /path/to/DGTRS-CLIP-ViT-B-16 \\
        --stripnet_checkpoint /path/to/stripnet_small.pth \\
        --epochs 50 --batch_size 64 --micro_batch_size 64 --bf16
"""
from __future__ import annotations

import argparse
import json
import logging
import os
import random
import sys
import time
from contextlib import nullcontext
from pathlib import Path

import numpy as np
import torch
import torch.nn as nn
from torch.optim import AdamW
from torch.optim.lr_scheduler import CosineAnnealingLR, LinearLR, SequentialLR

sys.path.insert(0, str(Path(__file__).resolve().parent))

from src.models.dual_encoder import build_dual_encoder, get_trainable_params
from src.losses import InfoNCELoss
from src.metrics import compute_bidirectional_metrics
from src.data.gta_uav import build_dataloaders
from src.data.gta_uav_eval import build_multipos_eval

logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s %(name)s %(levelname)s %(message)s",
)
LOGGER = logging.getLogger("cvgl.train")


def set_seed(seed: int) -> None:
    """Зафиксировать все источники случайности (протокол §0.4: seed=42 dev).

    Покрывает python random, numpy, torch (CPU/CUDA) и порядок сэмплирования
    в DataLoader/semi-positive выборе — для воспроизводимого equal-budget.
    """
    random.seed(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)
    torch.cuda.manual_seed_all(seed)
    os.environ["PYTHONHASHSEED"] = str(seed)
    LOGGER.info("🎲 Seed fixed: %d", seed)


# ---------------------------------------------------------------------------
# GPU info
# ---------------------------------------------------------------------------

def log_gpu_info(device: torch.device) -> None:
    if device.type != "cuda":
        return
    name = torch.cuda.get_device_name(device)
    total = torch.cuda.get_device_properties(device).total_memory / 1024**3
    LOGGER.info("🖥️  GPU: %s (%.1f GB VRAM)", name, total)
    cap = torch.cuda.get_device_capability(device)
    bf16_ok = cap[0] >= 8
    LOGGER.info(
        "   Compute capability: %d.%d | BF16: %s",
        cap[0], cap[1], "✅ supported" if bf16_ok else "❌ use FP16 instead",
    )


def log_vram_usage(prefix: str = "") -> None:
    if not torch.cuda.is_available():
        return
    allocated = torch.cuda.memory_allocated() / 1024**3
    reserved = torch.cuda.memory_reserved() / 1024**3
    LOGGER.info("   %sVRAM: %.2f GB allocated / %.2f GB reserved", prefix, allocated, reserved)


def get_amp_context(use_bf16: bool, use_fp16: bool, device: torch.device):
    if use_bf16:
        return torch.autocast(device_type=device.type, dtype=torch.bfloat16)
    elif use_fp16:
        return torch.autocast(device_type=device.type, dtype=torch.float16)
    else:
        return nullcontext()


# ---------------------------------------------------------------------------
# Evaluation
# ---------------------------------------------------------------------------

@torch.no_grad()
@torch.no_grad()
def _encode_loader(model, loader, encode_fn, device, amp_ctx) -> torch.Tensor:
    """Прогнать loader через encode_fn (encode_drone/encode_satellite)."""
    embs = []
    for batch in loader:
        images = batch["image"].to(device)
        tokens = batch["tokens"].to(device)
        with amp_ctx:
            emb = encode_fn(images, tokens)
        embs.append(emb.float().cpu())
    return torch.cat(embs, dim=0)


def evaluate(
    model: nn.Module,
    eval_data: dict,
    device: torch.device,
    amp_ctx,
) -> dict[str, float]:
    """Multi-positive retrieval eval (протокол §6.2).

    Дрон-запросы и УНИКАЛЬНАЯ спутниковая галерея кодируются раздельно;
    для каждого дрона учитываются ВСЕ его positive/semi-positive тайлы
    (hit-if-any). Считаются оба направления: q2g (primary) и g2q.

    Args:
        eval_data: dict из build_multipos_eval (drone_loader, gallery_loader,
            positives_q2g, positives_g2q).
    """
    model.eval()

    query_emb = _encode_loader(
        model, eval_data["drone_loader"], model.encode_drone, device, amp_ctx,
    )
    gallery_emb = _encode_loader(
        model, eval_data["gallery_loader"], model.encode_satellite, device, amp_ctx,
    )

    metrics = compute_bidirectional_metrics(
        query_emb, gallery_emb, ks=(1, 5, 10),
        positives_q2g=eval_data["positives_q2g"],
        positives_g2q=eval_data["positives_g2q"],
        return_ranks=True,
    )
    # Отделяем per-query массивы (0-indexed ранги первого попадания) от
    # числовых метрик: они пойдут в .npy для bootstrap CI / paired-теста,
    # а не в history.json (иначе JSON-дамп упадёт на np.ndarray).
    per_query = {
        "q2g_best_ranks": metrics.pop("q2g_best_ranks"),
        "g2q_best_ranks": metrics.pop("g2q_best_ranks"),
    }
    # Primary-алиасы: recall@k / mAP = q2g (для best-model и collect_results).
    for k in (1, 5, 10):
        metrics[f"recall@{k}"] = metrics[f"q2g_recall@{k}"]
    metrics["mAP"] = metrics["q2g_mAP"]

    model.train()
    return metrics, per_query


# ---------------------------------------------------------------------------
# Training loop
# ---------------------------------------------------------------------------

def train_one_epoch(
    model: nn.Module,
    train_loader,
    criterion: nn.Module,
    optimizer: torch.optim.Optimizer,
    device: torch.device,
    epoch: int,
    grad_accumulate_steps: int = 1,
    max_grad_norm: float = 1.0,
    amp_ctx=None,
    scaler: torch.amp.GradScaler | None = None,
) -> dict[str, float]:
    model.train()
    if amp_ctx is None:
        amp_ctx = nullcontext()

    total_loss = 0.0
    total_acc_d2s = 0.0
    total_acc_s2d = 0.0
    n_steps = 0

    optimizer.zero_grad()

    for batch_idx, batch in enumerate(train_loader):
        drone_images = batch["drone_image"].to(device)
        drone_tokens = batch["drone_tokens"].to(device)
        satellite_images = batch["satellite_image"].to(device)
        satellite_tokens = batch["satellite_tokens"].to(device)

        with amp_ctx:
            outputs = model(drone_images, drone_tokens, satellite_images, satellite_tokens)
            logits = outputs["logits"]
            loss_dict = criterion(logits)
            loss = loss_dict["loss"]
            scaled_loss = loss / grad_accumulate_steps

        if scaler is not None:
            scaler.scale(scaled_loss).backward()
        else:
            scaled_loss.backward()

        total_loss += loss.item()
        total_acc_d2s += loss_dict["acc_d2s"].item()
        total_acc_s2d += loss_dict["acc_s2d"].item()
        n_steps += 1

        if (batch_idx + 1) % grad_accumulate_steps == 0:
            if scaler is not None:
                scaler.unscale_(optimizer)
                torch.nn.utils.clip_grad_norm_(
                    [p for p in model.parameters() if p.requires_grad],
                    max_grad_norm,
                )
                scaler.step(optimizer)
                scaler.update()
            else:
                torch.nn.utils.clip_grad_norm_(
                    [p for p in model.parameters() if p.requires_grad],
                    max_grad_norm,
                )
                optimizer.step()
            optimizer.zero_grad()

        if (batch_idx + 1) % 50 == 0:
            LOGGER.info(
                "  [Epoch %d] Step %d/%d | loss=%.4f | acc_d2s=%.3f | "
                "acc_s2d=%.3f | τ=%.4f",
                epoch, batch_idx + 1, len(train_loader),
                loss.item(),
                loss_dict["acc_d2s"].item(),
                loss_dict["acc_s2d"].item(),
                outputs["temperature"].item(),
            )

    if len(train_loader) % grad_accumulate_steps != 0:
        if scaler is not None:
            scaler.unscale_(optimizer)
            torch.nn.utils.clip_grad_norm_(
                [p for p in model.parameters() if p.requires_grad],
                max_grad_norm,
            )
            scaler.step(optimizer)
            scaler.update()
        else:
            torch.nn.utils.clip_grad_norm_(
                [p for p in model.parameters() if p.requires_grad],
                max_grad_norm,
            )
            optimizer.step()
        optimizer.zero_grad()

    return {
        "loss": total_loss / max(n_steps, 1),
        "acc_d2s": total_acc_d2s / max(n_steps, 1),
        "acc_s2d": total_acc_s2d / max(n_steps, 1),
    }


# ---------------------------------------------------------------------------
# Main
# ---------------------------------------------------------------------------

def main(args):
    set_seed(args.seed)
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    LOGGER.info("🚀 Device: %s", device)
    log_gpu_info(device)

    exp_name = f"exp_{'-'.join(args.text_levels)}_ep{args.epochs}_bs{args.batch_size}"
    output_dir = Path(args.output_dir) / exp_name
    output_dir.mkdir(parents=True, exist_ok=True)
    LOGGER.info("📁 Output: %s", output_dir)

    config = vars(args)
    with open(output_dir / "config.json", "w") as f:
        json.dump(config, f, indent=2, default=str)

    train_loader, _ = build_dataloaders(
        data_root=args.data_root,
        descriptions_path=args.descriptions_path,
        text_levels=args.text_levels,
        train_meta=args.train_meta,
        test_meta=args.test_meta,
        batch_size=args.micro_batch_size,
        num_workers=args.num_workers,
        image_size=args.image_size,
        build_test=False,  # оценка идёт через multi-positive eval ниже
        mutually_exclusive=args.mutually_exclusive,
        seed=args.seed,
    )

    # Multi-positive eval (протокол §6.2): уникальная галерея + positive-карты.
    eval_data = build_multipos_eval(
        data_root=args.data_root,
        test_meta=args.test_meta,
        descriptions_path=args.descriptions_path,
        text_levels=args.text_levels,
        image_size=args.image_size,
        batch_size=args.micro_batch_size,
        num_workers=args.num_workers,
    )

    model = build_dual_encoder(
        dgtrs_checkpoint=args.dgtrs_checkpoint,
        stripnet_checkpoint=args.stripnet_checkpoint,
        fused_dim=args.fused_dim,
        shared_dim=args.shared_dim,
        freeze_text=True,
        freeze_image_backbone=True,
        inject_mona=args.inject_mona,
        mona_bottleneck=args.mona_bottleneck,
        device=str(device),
    )
    log_vram_usage("After model load: ")

    if args.compile and hasattr(torch, "compile"):
        LOGGER.info("⚡ Compiling model with torch.compile (mode=%s)", args.compile_mode)
        model = torch.compile(model, mode=args.compile_mode)

    use_bf16 = args.bf16 and device.type == "cuda"
    use_fp16 = args.fp16 and device.type == "cuda" and not use_bf16
    amp_ctx = get_amp_context(use_bf16, use_fp16, device)
    scaler = torch.amp.GradScaler("cuda") if use_fp16 else None

    precision_str = "BF16" if use_bf16 else ("FP16" if use_fp16 else "FP32")
    LOGGER.info("🔢 Precision: %s", precision_str)

    trainable_params = [p for p in model.parameters() if p.requires_grad]
    optimizer = AdamW(
        trainable_params,
        lr=args.lr,
        weight_decay=args.weight_decay,
        betas=(0.9, 0.98),
    )

    # Linear-warmup + cosine (протокол §8.1 / §0.1). Warmup зажимается до
    # epochs-1, чтобы на коротких прогонах (смок epochs=1) деградировать
    # к чистому косинусу без ошибок.
    warmup_epochs = max(0, min(args.warmup_epochs, args.epochs - 1))
    cosine = CosineAnnealingLR(
        optimizer,
        T_max=args.epochs - warmup_epochs,
        eta_min=args.lr * 0.01,
    )
    if warmup_epochs > 0:
        warmup = LinearLR(
            optimizer, start_factor=0.01, end_factor=1.0, total_iters=warmup_epochs,
        )
        scheduler = SequentialLR(
            optimizer, schedulers=[warmup, cosine], milestones=[warmup_epochs],
        )
        LOGGER.info("📉 LR schedule: linear-warmup(%d) + cosine(%d)", warmup_epochs, args.epochs - warmup_epochs)
    else:
        scheduler = cosine
        LOGGER.info("📉 LR schedule: cosine (%d epochs, no warmup)", args.epochs)

    criterion = InfoNCELoss(label_smoothing=args.label_smoothing)

    grad_accumulate_steps = max(1, args.batch_size // args.micro_batch_size)
    LOGGER.info(
        "⚙️  Effective batch=%d (micro=%d × accumulate=%d)",
        args.batch_size, args.micro_batch_size, grad_accumulate_steps,
    )

    start_epoch = 1
    if args.resume and (output_dir / "latest_model.pth").exists():
        ckpt = torch.load(output_dir / "latest_model.pth", map_location=device)
        model.load_state_dict(ckpt["model_state_dict"])
        optimizer.load_state_dict(ckpt["optimizer_state_dict"])
        if "scheduler_state_dict" in ckpt:
            scheduler.load_state_dict(ckpt["scheduler_state_dict"])
        start_epoch = ckpt["epoch"] + 1
        LOGGER.info("🔄 Resumed from epoch %d", start_epoch - 1)

    best_recall1 = 0.0
    history = []

    history_path = output_dir / "history.json"
    if args.resume and history_path.exists():
        with open(history_path) as f:
            history = json.load(f)

    log_vram_usage("Before training: ")

    for epoch in range(start_epoch, args.epochs + 1):
        t0 = time.time()

        train_metrics = train_one_epoch(
            model, train_loader, criterion, optimizer,
            device, epoch,
            grad_accumulate_steps=grad_accumulate_steps,
            max_grad_norm=args.max_grad_norm,
            amp_ctx=amp_ctx,
            scaler=scaler,
        )

        scheduler.step()

        if epoch % args.eval_every == 0 or epoch == args.epochs:
            eval_metrics, per_query = evaluate(model, eval_data, device, amp_ctx)
            # Per-query ранги последнего eval (перезапись) — источник для
            # bootstrap CI / paired-теста между вариантами (протокол §8.3).
            np.save(output_dir / "last_ranks_q2g.npy", per_query["q2g_best_ranks"])
            np.save(output_dir / "last_ranks_g2q.npy", per_query["g2q_best_ranks"])
        else:
            eval_metrics, per_query = {}, None

        elapsed = time.time() - t0

        LOGGER.info(
            "📈 Epoch %d/%d (%.0fs) | loss=%.4f | "
            "R@1=%.3f R@5=%.3f R@10=%.3f | mAP=%.3f",
            epoch, args.epochs, elapsed,
            train_metrics["loss"],
            eval_metrics.get("recall@1", 0),
            eval_metrics.get("recall@5", 0),
            eval_metrics.get("recall@10", 0),
            eval_metrics.get("mAP", 0),
        )

        if epoch == 1:
            log_vram_usage("After first epoch: ")

        record = {
            "epoch": epoch,
            "lr": scheduler.get_last_lr()[0],
            **{f"train_{k}": v for k, v in train_metrics.items()},
            **{f"eval_{k}": v for k, v in eval_metrics.items()},
            "elapsed_s": elapsed,
        }
        history.append(record)

        recall1 = eval_metrics.get("recall@1", 0)
        if recall1 > best_recall1:
            best_recall1 = recall1
            torch.save(
                {
                    "epoch": epoch,
                    "model_state_dict": model.state_dict(),
                    "eval_metrics": eval_metrics,
                    "config": config,
                },
                output_dir / "best_model.pth",
            )
            # Ранги best-эпохи — для честного сравнения вариантов по лучшей
            # точке (bootstrap CI / McNemar в compare-скрипте после обучения).
            if per_query is not None:
                np.save(output_dir / "best_ranks_q2g.npy", per_query["q2g_best_ranks"])
                np.save(output_dir / "best_ranks_g2q.npy", per_query["g2q_best_ranks"])
            LOGGER.info("💾 New best model (R@1=%.4f)", recall1)

        torch.save(
            {
                "epoch": epoch,
                "model_state_dict": model.state_dict(),
                "optimizer_state_dict": optimizer.state_dict(),
                "scheduler_state_dict": scheduler.state_dict(),
            },
            output_dir / "latest_model.pth",
        )

        with open(history_path, "w") as f:
            json.dump(history, f, indent=2)

    LOGGER.info("=" * 60)
    LOGGER.info("🏁 Training complete. Best R@1: %.4f", best_recall1)
    LOGGER.info("📁 Results: %s", output_dir)


# ---------------------------------------------------------------------------
# CLI
# ---------------------------------------------------------------------------

def parse_args():
    p = argparse.ArgumentParser(description="Train CVGL fusion dual-encoder on GTA-UAV")

    # Data
    p.add_argument("--data_root", type=str, required=True)
    p.add_argument("--descriptions_path", type=str, required=True)
    p.add_argument("--text_levels", nargs="+", default=["level1"])
    p.add_argument("--train_meta", default="cross-area-drone2sate-train.json")
    p.add_argument("--test_meta", default="cross-area-drone2sate-test.json")
    p.add_argument("--image_size", type=int, default=256)
    p.add_argument("--num_workers", type=int, default=8)

    # Model
    p.add_argument("--dgtrs_checkpoint", type=str, required=True)
    p.add_argument("--stripnet_checkpoint", type=str, required=True)
    p.add_argument("--fused_dim", type=int, default=512,
                   help="Размерность вектора после слияния картинки и текста")
    p.add_argument("--shared_dim", type=int, default=512)
    p.add_argument("--inject_mona", action="store_true", default=True)
    p.add_argument("--mona_bottleneck", type=int, default=64)
    p.add_argument(
        "--mutually_exclusive", action=argparse.BooleanOptionalAction, default=True,
        help="Батчи без общих позитивных тайлов (§6.4); --no-mutually_exclusive "
             "→ обычный random shuffle.",
    )

    # Training
    p.add_argument("--epochs", type=int, default=50)
    p.add_argument("--batch_size", type=int, default=64)
    p.add_argument("--micro_batch_size", type=int, default=64)
    p.add_argument("--lr", type=float, default=1e-4)
    p.add_argument("--weight_decay", type=float, default=0.01)
    p.add_argument("--max_grad_norm", type=float, default=1.0)
    p.add_argument("--label_smoothing", type=float, default=0.1)
    p.add_argument("--eval_every", type=int, default=1)

    # Performance
    p.add_argument("--bf16", action="store_true", default=True)
    p.add_argument("--fp16", action="store_true", default=False)
    p.add_argument("--compile", action="store_true", default=False)
    p.add_argument("--compile_mode", default="reduce-overhead",
                   choices=["default", "reduce-overhead", "max-autotune"])

    # Resume / output
    p.add_argument("--resume", action="store_true", default=False)
    p.add_argument("--output_dir", type=str, default="outputs")
    p.add_argument("--seed", type=int, default=42)
    p.add_argument("--warmup_epochs", type=int, default=3,
                   help="Эпох линейного warmup перед косинусом (§8.1).")

    return p.parse_args()


if __name__ == "__main__":
    main(parse_args())
