"""Тест разрушения текста — диагностика shortcut / устойчивости к missing-modality.

Берёт ОБУЧЕННЫЙ чекпоинт и прогоняет тот же multi-positive retrieval-eval
(протокол §6.2), но портит текстовый вход, измеряя падение R@1. Отвечает на
вопрос «модель реально использует текст или это shortcut/бутафория» (§3.5, §8.2).

Режимы порчи (текст меняется СИММЕТРИЧНО на обеих сторонах — query и gallery,
иначе слитые векторы несопоставимы):
    full    — без изменений (референс);
    zero    — текст убран (пустые токены) → чистый image-retrieval.
              Вклад текста = R@1(full) − R@1(zero). Для H-TA-2: малая разница
              при обучении с text-dropout = модель устойчива к отсутствию текста.
    shuffle — каждому объекту подставлен ЧУЖОЙ валидный caption (глобальная
              перестановка). Если R@1 не падает — текст не несёт геометрию
              (модель цепляется за лексический шаблон = косвенный shortcut).

Интерпретация (§8.2): малое падение при высоком train R@1 → shortcut / текст не
используется; большое падение → текст реально нужен.

Скрипт НЕ трогает даталоадер/обучение — только читает чекпоинт и гоняет eval.

Запуск (проще всего — по каталогу прогона с config.json + best_model.pth):
    python eval_text_destruction.py --run_dir outputs/text_aug/B0/seed42/exp_level1-level2-level3_ep10_bs64

Либо явно:
    python eval_text_destruction.py \
        --checkpoint outputs/.../best_model.pth \
        --data_root /media/.../GTA-UAV-LR \
        --descriptions_path /media/.../GTA-UAV-LR-captions_v2 \
        --text_levels level1 level2 level3
"""
from __future__ import annotations

import argparse
import json
import logging
import sys
from contextlib import nullcontext
from pathlib import Path

import torch

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

from src.models.dual_encoder import build_dual_encoder
from src.models.dgtrs.model import tokenize_dgtrs
from src.metrics import compute_bidirectional_metrics
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.text_destruction")

MODES = ("full", "zero", "shuffle")


# ---------------------------------------------------------------------------
# Кодирование loader'а с подменой токенов (pool задаёт токены в порядке loader)
# ---------------------------------------------------------------------------

@torch.no_grad()
def collect_tokens(loader) -> torch.Tensor:
    """Собрать все токены loader'а в один тензор [N, L] (порядок loader'а)."""
    return torch.cat([batch["tokens"] for batch in loader], dim=0)


@torch.no_grad()
def encode(model, loader, encode_fn, device, amp_ctx,
           pool: torch.Tensor | None = None) -> torch.Tensor:
    """Прогнать loader через encode_fn; если pool задан — берём токены из него.

    pool[N, L] — токены в порядке loader'а (для zero/shuffle). None — исходные
    токены из батча (режим full).
    """
    embs = []
    off = 0
    for batch in loader:
        images = batch["image"].to(device)
        bs = images.shape[0]
        tokens = batch["tokens"] if pool is None else pool[off:off + bs]
        tokens = tokens.to(device)
        with amp_ctx:
            embs.append(encode_fn(images, tokens).float().cpu())
        off += bs
    return torch.cat(embs, dim=0)


def make_pool(mode: str, orig: torch.Tensor, empty_row: torch.Tensor,
              generator: torch.Generator) -> torch.Tensor | None:
    """Построить pool токенов для режима порчи.

    Args:
        mode: full / zero / shuffle.
        orig: исходные токены [N, L] (в порядке loader'а).
        empty_row: токены пустого текста [L] (SOT+EOT+паддинг).
        generator: для воспроизводимой перестановки (shuffle).
    """
    n = orig.shape[0]
    if mode == "full":
        return None
    if mode == "zero":
        return empty_row.unsqueeze(0).expand(n, -1).clone()
    if mode == "shuffle":
        perm = torch.randperm(n, generator=generator)
        return orig[perm].clone()
    raise ValueError(f"Unknown mode: {mode}")


# ---------------------------------------------------------------------------
# Загрузка обученной модели из чекпоинта + config
# ---------------------------------------------------------------------------

def resolve_paths(args) -> tuple[Path, dict]:
    """Найти чекпоинт и config: либо из --run_dir, либо из явных аргументов."""
    if args.run_dir:
        run_dir = Path(args.run_dir)
        ckpt_path = Path(args.checkpoint) if args.checkpoint else run_dir / "best_model.pth"
        cfg_path = run_dir / "config.json"
        config = json.loads(cfg_path.read_text()) if cfg_path.exists() else {}
    else:
        if not args.checkpoint:
            raise SystemExit("Нужен либо --run_dir, либо --checkpoint")
        ckpt_path = Path(args.checkpoint)
        config = {}
    # аргументы командной строки перекрывают config
    for key in ("data_root", "descriptions_path", "test_meta", "image_size",
                "fused_dim", "shared_dim", "mona_bottleneck",
                "dgtrs_checkpoint", "stripnet_checkpoint"):
        val = getattr(args, key, None)
        if val is not None:
            config[key] = val
    if args.text_levels is not None:
        config["text_levels"] = args.text_levels
    return ckpt_path, config


def load_model(ckpt_path: Path, config: dict, device: torch.device):
    """Собрать DualEncoder по config и загрузить веса из чекпоинта."""
    model = build_dual_encoder(
        dgtrs_checkpoint=config["dgtrs_checkpoint"],
        stripnet_checkpoint=config["stripnet_checkpoint"],
        fused_dim=config.get("fused_dim", 512),
        shared_dim=config.get("shared_dim", 512),
        freeze_text=True,
        freeze_image_backbone=True,
        inject_mona=config.get("inject_mona", True),
        mona_bottleneck=config.get("mona_bottleneck", 64),
        device=str(device),
    )
    ckpt = torch.load(ckpt_path, map_location=device)
    state = ckpt.get("model_state_dict", ckpt)
    missing, unexpected = model.load_state_dict(state, strict=False)
    if missing:
        LOGGER.warning("Отсутствуют ключи при загрузке: %d (напр. %s)", len(missing), missing[:3])
    if unexpected:
        LOGGER.warning("Лишние ключи при загрузке: %d (напр. %s)", len(unexpected), unexpected[:3])
    model.eval()
    LOGGER.info("✅ Модель загружена: %s (epoch=%s)", ckpt_path, ckpt.get("epoch", "?"))
    return model


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

def main():
    args = parse_args()
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    ckpt_path, config = resolve_paths(args)

    context_length = config.get("context_length", 248)
    text_levels = config.get("text_levels", ["level1"])
    LOGGER.info("🔎 text_levels=%s | test_meta=%s", text_levels,
                config.get("test_meta", "cross-area-drone2sate-test.json"))

    model = load_model(ckpt_path, config, device)

    amp_ctx = (torch.autocast(device_type="cuda", dtype=torch.bfloat16)
               if device.type == "cuda" else nullcontext())

    eval_data = build_multipos_eval(
        data_root=config["data_root"],
        test_meta=config.get("test_meta", "cross-area-drone2sate-test.json"),
        descriptions_path=config["descriptions_path"],
        text_levels=text_levels,
        image_size=config.get("image_size", 256),
        batch_size=args.batch_size,
        num_workers=args.num_workers,
        context_length=context_length,
    )

    # Пустой текст (токены SOT+EOT+паддинг) — общий для zero-режима.
    empty_row = tokenize_dgtrs("", context_length=context_length, truncate=True).squeeze(0)

    # Исходные токены обеих сторон (в порядке loader'ов) — нужны для shuffle.
    drone_orig = collect_tokens(eval_data["drone_loader"])
    gallery_orig = collect_tokens(eval_data["gallery_loader"])
    LOGGER.info("📊 queries=%d, gallery=%d", drone_orig.shape[0], gallery_orig.shape[0])

    gen = torch.Generator().manual_seed(args.seed)

    results: dict[str, dict] = {}
    for mode in MODES:
        drone_pool = make_pool(mode, drone_orig, empty_row, gen)
        gallery_pool = make_pool(mode, gallery_orig, empty_row, gen)

        q = encode(model, eval_data["drone_loader"], model.encode_drone,
                   device, amp_ctx, pool=drone_pool)
        g = encode(model, eval_data["gallery_loader"], model.encode_satellite,
                   device, amp_ctx, pool=gallery_pool)

        m = compute_bidirectional_metrics(
            q, g, ks=(1, 5, 10),
            positives_q2g=eval_data["positives_q2g"],
            positives_g2q=eval_data["positives_g2q"],
        )
        results[mode] = {
            "q2g_recall@1": m["q2g_recall@1"],
            "q2g_recall@5": m["q2g_recall@5"],
            "q2g_recall@10": m["q2g_recall@10"],
            "q2g_mAP": m["q2g_mAP"],
            "g2q_recall@1": m["g2q_recall@1"],
        }
        LOGGER.info("[%s] q2g R@1=%.4f R@5=%.4f R@10=%.4f | mAP=%.4f",
                    mode, m["q2g_recall@1"], m["q2g_recall@5"],
                    m["q2g_recall@10"], m["q2g_mAP"])

    _print_report(results)

    out_path = Path(args.output) if args.output else (
        (Path(args.run_dir) / "text_destruction.json") if args.run_dir
        else Path("text_destruction.json"))
    out_path.parent.mkdir(parents=True, exist_ok=True)
    out_path.write_text(json.dumps(results, indent=2, ensure_ascii=False))
    LOGGER.info("💾 Сохранено: %s", out_path)


def _print_report(results: dict[str, dict]) -> None:
    """Таблица режимов + падение R@1 относительно full."""
    full_r1 = results["full"]["q2g_recall@1"]
    print("\n" + "=" * 60)
    print("  ТЕСТ РАЗРУШЕНИЯ ТЕКСТА (q2g, cross-area)")
    print("=" * 60)
    print(f"{'Режим':<10} {'R@1':>8} {'R@5':>8} {'R@10':>8} {'ΔR@1 vs full':>14}")
    print("-" * 60)
    for mode in MODES:
        r = results[mode]
        drop = r["q2g_recall@1"] - full_r1
        drop_str = "—" if mode == "full" else f"{drop*100:+.2f} п.п."
        print(f"{mode:<10} {r['q2g_recall@1']*100:>7.2f}% {r['q2g_recall@5']*100:>7.2f}% "
              f"{r['q2g_recall@10']*100:>7.2f}% {drop_str:>14}")
    print("-" * 60)
    zero_drop = (results["zero"]["q2g_recall@1"] - full_r1) * 100
    shuf_drop = (results["shuffle"]["q2g_recall@1"] - full_r1) * 100
    print(f"Вклад текста (full−zero):    {-zero_drop:+.2f} п.п.")
    print(f"Падение при подмене (shuffle): {shuf_drop:+.2f} п.п.")
    print("Интерпретация: малое падение при высоком train R@1 → shortcut/текст "
          "не используется;\n               большое падение → текст реально нужен (§8.2).")
    print("=" * 60)


def parse_args():
    p = argparse.ArgumentParser(description="Тест разрушения текста (H-TA-2 / shortcut-диагностика)")
    p.add_argument("--run_dir", type=str, default=None,
                   help="Каталог прогона с config.json и best_model.pth")
    p.add_argument("--checkpoint", type=str, default=None,
                   help="Путь к .pth (по умолчанию run_dir/best_model.pth)")
    # переопределения config (если run_dir не задан или нужно поменять)
    p.add_argument("--data_root", type=str, default=None)
    p.add_argument("--descriptions_path", type=str, default=None)
    p.add_argument("--test_meta", type=str, default=None)
    p.add_argument("--text_levels", nargs="+", default=None)
    p.add_argument("--image_size", type=int, default=None)
    p.add_argument("--fused_dim", type=int, default=None)
    p.add_argument("--shared_dim", type=int, default=None)
    p.add_argument("--mona_bottleneck", type=int, default=None)
    p.add_argument("--dgtrs_checkpoint", type=str, default=None)
    p.add_argument("--stripnet_checkpoint", type=str, default=None)
    # прогон
    p.add_argument("--batch_size", type=int, default=64)
    p.add_argument("--num_workers", type=int, default=8)
    p.add_argument("--seed", type=int, default=42)
    p.add_argument("--output", type=str, default=None,
                   help="Куда сохранить JSON (по умолчанию run_dir/text_destruction.json)")
    return p.parse_args()


if __name__ == "__main__":
    main()
