"""Сборка результатов экспериментов: таблица + CSV + графики.

Использование:
    python collect_results.py          # v1 + v2
    python collect_results.py v1       # только v1
    python collect_results.py v2       # только v2
"""
from __future__ import annotations

import csv
import json
import sys
from pathlib import Path

import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt
import numpy as np


SCRIPT_DIR = Path(__file__).resolve().parent
OUTPUTS_DIR = SCRIPT_DIR / "outputs"
RESULTS_DIR = SCRIPT_DIR / "results"


# ---------------------------------------------------------------------------
# Загрузка
# ---------------------------------------------------------------------------

def load_experiment(exp_dir: Path) -> dict | None:
    history_path = exp_dir / "history.json"
    config_path = exp_dir / "config.json"
    if not history_path.exists():
        return None

    with open(history_path) as f:
        history = json.load(f)
    config = {}
    if config_path.exists():
        with open(config_path) as f:
            config = json.load(f)
    if not history:
        return None

    # primary метрика — q2g R@1 (drone→satellite), см. протокол §6.2;
    # eval_recall@* дублируют q2g, eval_mAP дублирует q2g_mAP.
    best = max(history, key=lambda r: r.get("eval_recall@1", 0))
    latest = history[-1]

    # Протокол сплита определяем из train_meta (§5.2): нужен для cross−same gap.
    train_meta = config.get("train_meta", "")
    if "same-area" in train_meta:
        protocol = "same-area"
    elif "cross-area" in train_meta:
        protocol = "cross-area"
    else:
        protocol = "?"

    return {
        "dir": exp_dir.name,
        "protocol": protocol,
        "seed": config.get("seed", "?"),
        "text_levels": " + ".join(config.get("text_levels", ["?"])),
        "epochs_done": latest["epoch"],
        "epochs_total": config.get("epochs", "?"),
        "best_epoch": best["epoch"],
        # q2g (primary). gap = R@1(same-area) − R@1(cross-area) считается в
        # compute_gaps(), если для варианта есть прогоны ОБОИХ протоколов.
        "best_R@1": best.get("eval_recall@1", 0),
        "best_R@5": best.get("eval_recall@5", 0),
        "best_R@10": best.get("eval_recall@10", 0),
        "best_R@1%": best.get("eval_q2g_recall@1%", 0),
        "best_mAP": best.get("eval_mAP", 0),          # настоящий AP, не MRR
        "best_median_rank": best.get("eval_q2g_median_rank", 0),
        "best_mean_rank": best.get("eval_q2g_mean_rank", 0),
        # g2q (satellite→drone), та же best-эпоха
        "best_g2q_R@1": best.get("eval_g2q_recall@1", 0),
        "best_g2q_R@5": best.get("eval_g2q_recall@5", 0),
        "best_g2q_R@10": best.get("eval_g2q_recall@10", 0),
        "best_g2q_R@1%": best.get("eval_g2q_recall@1%", 0),
        "best_g2q_mAP": best.get("eval_g2q_mAP", 0),
        "latest_loss": latest.get("train_loss", 0),
        "latest_R@1": latest.get("eval_recall@1", 0),
        "avg_epoch_time": sum(r.get("elapsed_s", 0) for r in history) / len(history),
        "_history": history,
    }


def collect_version(version: str) -> list[dict]:
    """Собрать результаты одной версии (v1 или v2).

    Ищет history.json рекурсивно под outputs/<version>/, поэтому находит
    прогоны обоих протоколов (same/cross), даже если они лежат на разной
    глубине (напр. seed/exp или seed/protocol/exp).
    """
    version_dir = OUTPUTS_DIR / version
    if not version_dir.exists():
        print(f"⚠️  Папка не найдена: {version_dir}")
        return []

    results = []
    for hist_path in sorted(version_dir.rglob("history.json")):
        data = load_experiment(hist_path.parent)
        if data:
            data["version"] = version
            results.append(data)
    return results


# ---------------------------------------------------------------------------
# Таблица в консоль
# ---------------------------------------------------------------------------

def print_table(results: list[dict], version: str) -> None:
    if not results:
        print(f"  {version}: нет данных")
        return

    results.sort(key=lambda r: -r["best_R@1"])

    header = (
        f"{'Levels':<24} {'Proto':>10} {'Seed':>5} {'Prog':<8} {'BestEp':>6} "
        f"{'R@1':>7} {'R@5':>7} {'R@10':>7} {'R@1%':>7} {'mAP':>7} "
        f"{'medR':>6} {'g2q@1':>7} {'Loss':>8} {'Time':>6}"
    )
    sep = "─" * len(header)

    print(sep)
    print(f"  {version.upper()} Results  (q2g = primary; g2q@1 = satellite→drone)")
    print(sep)
    print(header)
    print(sep)

    for r in results:
        prog = f"{r['epochs_done']}/{r['epochs_total']}"
        print(
            f"{r['text_levels']:<24} {r.get('protocol','?'):>10} {str(r['seed']):>5} "
            f"{prog:<8} {r['best_epoch']:>6} "
            f"{r['best_R@1']:>7.4f} {r['best_R@5']:>7.4f} {r['best_R@10']:>7.4f} "
            f"{r['best_R@1%']:>7.4f} {r['best_mAP']:>7.4f} {r['best_median_rank']:>6.0f} "
            f"{r['best_g2q_R@1']:>7.4f} "
            f"{r['latest_loss']:>8.4f} {r['avg_epoch_time']:>5.0f}s"
        )

    print(sep)
    print(f"  {len(results)} experiments")
    print()


# ---------------------------------------------------------------------------
# cross−same gap (§3.5, §5.3): диагностика text-shortcut
# ---------------------------------------------------------------------------

def compute_gaps(results: list[dict]) -> list[dict]:
    """Спарить same-area и cross-area прогоны и посчитать gap по R@1.

    gap = R@1(same-area) − R@1(cross-area). Большой gap = модель опирается на
    area-specific shortcut (в т.ч. текстовый шаблон), не переносящийся на
    невиданную область. Пары ищутся по (version, seed, text_levels): это ОДИН
    вариант, обученный под двумя протоколами. same-area R@1 честен только при
    отдельном обучении на same-area-train (§5.3), поэтому его нельзя заменить
    доп. eval'ом cross-модели.
    """
    from collections import defaultdict

    groups: dict[tuple, dict[str, dict]] = defaultdict(dict)
    for r in results:
        key = (r["version"], r["seed"], r["text_levels"])
        groups[key][r["protocol"]] = r

    gaps = []
    for (version, seed, levels), by_proto in groups.items():
        same = by_proto.get("same-area")
        cross = by_proto.get("cross-area")
        if not (same and cross):
            continue
        gaps.append({
            "version": version,
            "seed": seed,
            "text_levels": levels,
            "same_R@1": same["best_R@1"],
            "cross_R@1": cross["best_R@1"],
            "gap_R@1": same["best_R@1"] - cross["best_R@1"],
            "same_mAP": same["best_mAP"],
            "cross_mAP": cross["best_mAP"],
        })
    gaps.sort(key=lambda g: -g["gap_R@1"])
    return gaps


def print_gap_table(gaps: list[dict], version: str) -> None:
    if not gaps:
        print(f"  {version}: gap не посчитан — нужны прогоны ОБОИХ протоколов "
              f"(same-area И cross-area) для одного варианта (§5.3).")
        print()
        return

    header = (f"{'Levels':<24} {'Seed':>5} {'R@1 same':>9} {'R@1 cross':>10} "
              f"{'gap (same−cross)':>17}")
    sep = "─" * len(header)
    print(sep)
    print(f"  {version.upper()} — cross−same gap (диагностика text-shortcut; "
          f"меньше gap = лучше генерализация)")
    print(sep)
    print(header)
    print(sep)
    for g in gaps:
        print(f"{g['text_levels']:<24} {str(g['seed']):>5} "
              f"{g['same_R@1']:>9.4f} {g['cross_R@1']:>10.4f} "
              f"{g['gap_R@1']*100:>15.2f} п.п.")
    print(sep)
    print()


def save_gap_csv(gaps: list[dict], version: str) -> None:
    if not gaps:
        return
    fields = ["version", "seed", "text_levels", "same_R@1", "cross_R@1",
              "gap_R@1", "same_mAP", "cross_mAP"]
    path = RESULTS_DIR / f"gap_{version}.csv"
    with open(path, "w", newline="") as f:
        writer = csv.DictWriter(f, fieldnames=fields, extrasaction="ignore")
        writer.writeheader()
        for g in gaps:
            writer.writerow(g)
    print(f"  CSV: {path}")


# ---------------------------------------------------------------------------
# CSV
# ---------------------------------------------------------------------------

def save_csv(results: list[dict], version: str) -> None:
    fields = [
        "version", "protocol", "seed", "text_levels", "epochs_done", "epochs_total",
        "best_epoch",
        "best_R@1", "best_R@5", "best_R@10", "best_R@1%", "best_mAP",
        "best_median_rank", "best_mean_rank",
        "best_g2q_R@1", "best_g2q_R@5", "best_g2q_R@10", "best_g2q_R@1%", "best_g2q_mAP",
        "latest_loss", "latest_R@1", "avg_epoch_time", "dir",
    ]
    path = RESULTS_DIR / f"results_{version}.csv"
    with open(path, "w", newline="") as f:
        writer = csv.DictWriter(f, fieldnames=fields, extrasaction="ignore")
        writer.writeheader()
        for r in results:
            writer.writerow(r)
    print(f"  CSV: {path}")


# ---------------------------------------------------------------------------
# Графики
# ---------------------------------------------------------------------------

def save_plots(results: list[dict], version: str) -> None:
    if not results:
        return

    colors = plt.cm.tab10(np.linspace(0, 1, max(len(results), 1)))

    # --- 1. Loss ---
    fig, ax = plt.subplots(figsize=(8, 5))
    for r, c in zip(results, colors):
        h = r["_history"]
        ax.plot([e["epoch"] for e in h],
                [e.get("train_loss", 0) for e in h],
                label=r["text_levels"], color=c, linewidth=1.5)
    ax.set_xlabel("Epoch")
    ax.set_ylabel("Train Loss")
    ax.set_title(f"{version.upper()} — Training Loss")
    ax.legend(fontsize=8)
    ax.grid(True, alpha=0.3)
    plt.tight_layout()
    path = RESULTS_DIR / f"loss_{version}.png"
    plt.savefig(path, dpi=150)
    plt.close()
    print(f"  PNG: {path}")

    # --- 2. Recall@1 ---
    fig, ax = plt.subplots(figsize=(8, 5))
    for r, c in zip(results, colors):
        h = r["_history"]
        eps = [e["epoch"] for e in h if e.get("eval_recall@1") is not None]
        r1s = [e["eval_recall@1"] for e in h if e.get("eval_recall@1") is not None]
        if eps:
            ax.plot(eps, r1s, label=r["text_levels"], color=c,
                    linewidth=1.5, marker=".", markersize=3)
    ax.set_xlabel("Epoch")
    ax.set_ylabel("Recall@1")
    ax.set_title(f"{version.upper()} — Recall@1")
    ax.legend(fontsize=8)
    ax.grid(True, alpha=0.3)
    plt.tight_layout()
    path = RESULTS_DIR / f"recall1_{version}.png"
    plt.savefig(path, dpi=150)
    plt.close()
    print(f"  PNG: {path}")

    # --- 3. Recall@5 ---
    fig, ax = plt.subplots(figsize=(8, 5))
    for r, c in zip(results, colors):
        h = r["_history"]
        eps = [e["epoch"] for e in h if e.get("eval_recall@5") is not None]
        r5s = [e["eval_recall@5"] for e in h if e.get("eval_recall@5") is not None]
        if eps:
            ax.plot(eps, r5s, label=r["text_levels"], color=c,
                    linewidth=1.5, marker=".", markersize=3)
    ax.set_xlabel("Epoch")
    ax.set_ylabel("Recall@5")
    ax.set_title(f"{version.upper()} — Recall@5")
    ax.legend(fontsize=8)
    ax.grid(True, alpha=0.3)
    plt.tight_layout()
    path = RESULTS_DIR / f"recall5_{version}.png"
    plt.savefig(path, dpi=150)
    plt.close()
    print(f"  PNG: {path}")
 
    # --- 4. Recall@10 ---
    fig, ax = plt.subplots(figsize=(8, 5))
    for r, c in zip(results, colors):
        h = r["_history"]
        eps = [e["epoch"] for e in h if e.get("eval_recall@10") is not None]
        r10s = [e["eval_recall@10"] for e in h if e.get("eval_recall@10") is not None]
        if eps:
            ax.plot(eps, r10s, label=r["text_levels"], color=c,
                    linewidth=1.5, marker=".", markersize=3)
    ax.set_xlabel("Epoch")
    ax.set_ylabel("Recall@10")
    ax.set_title(f"{version.upper()} — Recall@10")
    ax.legend(fontsize=8)
    ax.grid(True, alpha=0.3)
    plt.tight_layout()
    path = RESULTS_DIR / f"recall10_{version}.png"
    plt.savefig(path, dpi=150)
    plt.close()
    print(f"  PNG: {path}")
 
    # --- 3. Bar chart лучших ---
    fig, ax = plt.subplots(figsize=(8, 5))
    labels = [r["text_levels"] for r in results]
    r1 = [r["best_R@1"] for r in results]
    r5 = [r["best_R@5"] for r in results]
    r10 = [r["best_R@10"] for r in results]
 
    x = np.arange(len(labels))
    w = 0.25
    ax.bar(x - w, r1, w, label="R@1", color="#4C78A8")
    ax.bar(x, r5, w, label="R@5", color="#54A24B")
    ax.bar(x + w, r10, w, label="R@10", color="#E45756")
    ax.set_xticks(x)
    ax.set_xticklabels(labels, fontsize=8, rotation=20, ha="right")
    ax.set_ylabel("Score")
    ax.set_title(f"{version.upper()} — Best Recall")
    ax.legend()
    ax.grid(True, alpha=0.3, axis="y")
    for i, v in enumerate(r1):
        ax.text(i - w, v + 0.005, f"{v:.3f}", ha="center", fontsize=7)
    plt.tight_layout()
    path = RESULTS_DIR / f"best_recall_{version}.png"
    plt.savefig(path, dpi=150)
    plt.close()
    print(f"  PNG: {path}")

    # --- 6. gate σ(α)-траектории (§9.2 modality-collapse) ---
    # Есть только у новых прогонов (train_drone_gate/train_satellite_gate).
    has_gate = any(
        any("train_drone_gate" in e for e in r["_history"]) for r in results
    )
    if has_gate:
        fig, ax = plt.subplots(figsize=(8, 5))
        for r, c in zip(results, colors):
            h = r["_history"]
            eps = [e["epoch"] for e in h if "train_drone_gate" in e]
            gd = [e["train_drone_gate"] for e in h if "train_drone_gate" in e]
            gs = [e.get("train_satellite_gate") for e in h if "train_drone_gate" in e]
            if not eps:
                continue
            ax.plot(eps, gd, color=c, linewidth=1.5, label=f"{r['text_levels']} (drone)")
            ax.plot(eps, gs, color=c, linewidth=1.0, linestyle="--", alpha=0.7)
        ax.axhline(0.5, color="gray", linewidth=0.8, alpha=0.5)
        ax.set_ylim(0, 1)
        ax.set_xlabel("Epoch")
        ax.set_ylabel("gate σ(α) — вес картинки")
        ax.set_title(f"{version.upper()} — gate σ(α) (0/1 = modality collapse)")
        ax.legend(fontsize=7)
        ax.grid(True, alpha=0.3)
        plt.tight_layout()
        path = RESULTS_DIR / f"gate_{version}.png"
        plt.savefig(path, dpi=150)
        plt.close()
        print(f"  PNG: {path}")

    # --- 7. Pareto: R@1 vs стоимость (avg epoch time) ---
    fig, ax = plt.subplots(figsize=(8, 5))
    for r, c in zip(results, colors):
        ax.scatter(r["avg_epoch_time"], r["best_R@1"], color=c, s=60,
                   label=f"{r['text_levels']} ({r.get('protocol','?')})")
        ax.annotate(r["text_levels"], (r["avg_epoch_time"], r["best_R@1"]),
                    fontsize=7, xytext=(4, 4), textcoords="offset points")
    ax.set_xlabel("Средн. время эпохи, с (прокси стоимости)")
    ax.set_ylabel("Best R@1")
    ax.set_title(f"{version.upper()} — Pareto: R@1 vs cost")
    ax.grid(True, alpha=0.3)
    plt.tight_layout()
    path = RESULTS_DIR / f"pareto_{version}.png"
    plt.savefig(path, dpi=150)
    plt.close()
    print(f"  PNG: {path}")


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

def main():
    # Парсинг: python collect_results.py [v1|v2]
    if len(sys.argv) > 1 and sys.argv[1] in ("v1", "v2"):
        versions = [sys.argv[1]]
    else:
        versions = ["v1", "v2"]

    if not OUTPUTS_DIR.exists():
        print(f"❌ Папка outputs не найдена: {OUTPUTS_DIR}")
        return

    RESULTS_DIR.mkdir(exist_ok=True)

    for version in versions:
        results = collect_version(version)

        # Таблица в консоль
        print_table(results, version)

        # cross−same gap (диагностика text-shortcut) — если есть пары протоколов
        gaps = compute_gaps(results)
        print_gap_table(gaps, version)

        if results:
            # CSV и графики в results/
            save_csv(results, version)
            save_gap_csv(gaps, version)
            save_plots(results, version)

    print(f"\n📁 Все результаты: {RESULTS_DIR}")


if __name__ == "__main__":
    main()