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

ВАЖНОЕ ИЗМЕНЕНИЕ: модель теперь принимает ЧЕТЫРЕ входа вместо трёх:
    drone_images, drone_tokens         → сливаются в drone_emb (TextFusionMLP)
    satellite_images, satellite_tokens → сливаются в satellite_emb (TextFusionMLP)
    drone_emb и satellite_emb сравниваются между собой → cosine similarity → InfoNCE

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

Оптимизировано под RTX 4090 (24 GB VRAM, Ada Lovelace): BF16 AMP,
micro_batch=64 по умолчанию (effective batch = micro_batch при отсутствии
gradient accumulation).

Использование:
    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 sys
import time
from contextlib import nullcontext
from pathlib import Path

import torch
import torch.nn as nn
from torch.optim import AdamW
from torch.optim.lr_scheduler import CosineAnnealingLR

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_retrieval_metrics, format_metrics
from src.data.gta_uav import build_dataloaders

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


# ---------------------------------------------------------------------------
# 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()
def evaluate(
    model: nn.Module,
    test_loader,
    device: torch.device,
    amp_ctx,
) -> dict[str, float]:
    """Прогнать test set: drone (fused) vs satellite (fused) — симметрично."""
    model.eval()

    all_drone_emb = []
    all_satellite_emb = []

    for batch in test_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:
            drone_emb = model.encode_drone(drone_images, drone_tokens)
            satellite_emb = model.encode_satellite(satellite_images, satellite_tokens)

        all_drone_emb.append(drone_emb.float().cpu())
        all_satellite_emb.append(satellite_emb.float().cpu())

    all_drone_emb = torch.cat(all_drone_emb, dim=0)
    all_satellite_emb = torch.cat(all_satellite_emb, dim=0)

    metrics = compute_retrieval_metrics(
        all_drone_emb, all_satellite_emb, ks=[1, 5, 10],
    )

    model.train()
    return metrics


# ---------------------------------------------------------------------------
# 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_t2i = 0.0
    total_acc_i2t = 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_t2i += loss_dict["acc_t2i"].item()
        total_acc_i2t += loss_dict["acc_i2t"].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_t2i"].item(),
                loss_dict["acc_i2t"].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_t2i": total_acc_t2i / max(n_steps, 1),
        "acc_i2t": total_acc_i2t / max(n_steps, 1),
    }


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

def main(args):
    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, test_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,
    )

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

    scheduler = CosineAnnealingLR(
        optimizer,
        T_max=args.epochs,
        eta_min=args.lr * 0.01,
    )

    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 = evaluate(model, test_loader, device, amp_ctx)
        else:
            eval_metrics = {}

        elapsed = time.time() - t0

        LOGGER.info(
            "📈 Epoch %d/%d (%.0fs) | loss=%.4f | "
            "R@1=%.3f R@5=%.3f R@10=%.3f | AP=%.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("AP", 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",
            )
            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=384)
    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)

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

    return p.parse_args()


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