"""Сборка результатов экспериментов: таблица + 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]

    return {
        "dir": exp_dir.name,
        "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); split — только official cross-area (§5.3),
        # same-area прогонов нет, поэтому cross−same gap не считается.
        "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).

    Структура: outputs/<version>/<seed>/<exp_dir>/history.json.
    """
    version_dir = OUTPUTS_DIR / version
    if not version_dir.exists():
        print(f"⚠️  Папка не найдена: {version_dir}")
        return []

    results = []
    for exp_dir in sorted(version_dir.glob("*/*")):
        if exp_dir.is_dir() and (exp_dir / "history.json").exists():
            data = load_experiment(exp_dir)
            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} {'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} {str(r['seed']):>5} {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()


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

def save_csv(results: list[dict], version: str) -> None:
    fields = [
        "version", "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}")


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

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

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


if __name__ == "__main__":
    main()