#!/usr/bin/env python3
"""Generate master comparison report from eval outputs.

Reads metrics.json from all output/{train,reproduce}/ directories,
produces:
  1. Master comparison table (markdown)
  2. Summary bar chart (train vs reproduce DTW similarity)
  3. Action statistics summary

Usage:
    conda run -n pycuda_3_10 python generate_report.py
"""
from __future__ import annotations

import json
import sys
from pathlib import Path
from typing import Any, Dict, List

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

_OUT_BASE = Path(__file__).resolve().parent / "output"
_REPORT_DIR = _OUT_BASE / "reports"


def log(msg: str) -> None:
    print(f"[report] {msg}", flush=True)


def load_metrics(scene_dir: Path) -> Dict[str, Any]:
    p = scene_dir / "metrics.json"
    if p.exists():
        with open(p) as f:
            return json.load(f)
    return {}


def load_all_seeds(scene_dir: Path) -> List[Dict]:
    p = scene_dir / "all_seeds.json"
    if p.exists():
        with open(p) as f:
            return json.load(f)
    return []


# ---------------------------------------------------------------------------
# Tables
# ---------------------------------------------------------------------------
def generate_comparison_table() -> str:
    """Generate markdown comparison table for overlapping scenes."""
    overlap = [
        ("re100_karman", "karman_cloak", "Karman Cloak Re100"),
        ("illusion_075L", "illusion_075L", "Illusion 0.75L"),
        ("illusion_1L", "illusion_1L", "Illusion 1.0L"),
        ("illusion_15L", "illusion_15L", "Illusion 1.5L"),
    ]

    lines = []
    lines.append("# Master Comparison Table")
    lines.append("")
    lines.append("| Scene | Side | Best Seed | DTW sim | Reward | "
                 "aF mean | aB mean | aT mean | r_cd | r_cl | r_sim |")
    lines.append("|-------|------|-----------|:---:|:---:|:---:|:---:|:---:|:---:|:---:|:---:|")

    for trn_key, rep_key, name in overlap:
        m_trn = load_metrics(_OUT_BASE / "train" / trn_key)
        m_rep = load_metrics(_OUT_BASE / "reproduce" / rep_key)

        dtw_t = _fmt(m_trn.get("dtw_sim_v5") or m_trn.get("sim_raw_mean"))
        dtw_r = _fmt(m_rep.get("dtw_sim_v5") or m_rep.get("sim_raw_mean"))

        # Train row
        lines.append(
            f"| {name} | Train | {m_trn.get('best_seed', '-')} | "
            f"{dtw_t} | {_fmt(m_trn.get('reward_mean'))} | "
            f"{_fmt(m_trn.get('aF_mean'))} | {_fmt(m_trn.get('aB_mean'))} | "
            f"{_fmt(m_trn.get('aT_mean'))} | "
            f"{_fmt(m_trn.get('r_cd_mean'))} | {_fmt(m_trn.get('r_cl_mean'))} | "
            f"{_fmt(m_trn.get('r_sim_mean'))} |"
        )
        # Reproduce row
        lines.append(
            f"| {name} | Reproduce | - | "
            f"{dtw_r} | - | "
            f"- | - | - | - | - | - |"
        )

    # Train-only scenes
    train_only = ["transfer_re60", "transfer_re200", "transfer_re400"]
    lines.append("")
    lines.append("## Train Only (Cross-Re Transfer)")
    lines.append("| Scene | Best Seed | Reward | "
                 "r_cd | r_cl | r_sim | sim_raw |")
    lines.append("|-------|-----------|---:|-----:|-----:|-----:|---:|")
    for trn_key in train_only:
        m = load_metrics(_OUT_BASE / "train" / trn_key)
        lines.append(
            f"| {trn_key} | {m.get('best_seed', '-')} | "
            f"{_fmt(m.get('reward_mean'))} | "
            f"{_fmt(m.get('r_cd_mean'))} | {_fmt(m.get('r_cl_mean'))} | "
            f"{_fmt(m.get('r_sim_mean'))} | {_fmt(m.get('sim_raw_mean'))} |"
        )

    # Reproduce-only
    repro_only = ["steady_cloak", "vortex_lamb", "vortex_taylor"]
    lines.append("")
    lines.append("## Reproduce Only")
    lines.append("| Scene | Notes |")
    lines.append("|-------|-------|")
    for r_key in repro_only:
        lines.append(f"| {r_key} | open-loop / vortex |")

    # Per-scene all-seeds detail
    lines.append("")
    lines.append("## Per-Scene Seed Details (Train)")
    for trn_key in ["re100_karman", "transfer_re60", "transfer_re200",
                     "transfer_re400", "illusion_075L", "illusion_1L", "illusion_15L"]:
        seeds = load_all_seeds(_OUT_BASE / "train" / trn_key)
        if not seeds:
            continue
        lines.append(f"\n### {trn_key}")
        lines.append("| Seed | Reward | r_cd | r_cl | r_sim | sim_raw | Time (s) |")
        lines.append("|------|--------|------|------|-------|---------|---------:|")
        for s in sorted(seeds, key=lambda x: -x.get("reward", -float("inf"))):
            lines.append(
                f"| {s['seed']} | {_fmt(s.get('reward'))} | "
                f"{_fmt(s.get('r_cd'))} | {_fmt(s.get('r_cl'))} | "
                f"{_fmt(s.get('r_sim'))} | {_fmt(s.get('sim_raw'))} | "
                f"{_fmt_int(s.get('dt_sec'))} |"
            )

    return "\n".join(lines)


def _fmt(v, precision=4):
    if v is None or v == -1.0:
        return "-"
    if isinstance(v, float):
        return f"{v:.{precision}f}"
    return str(v)


def _fmt_int(v):
    if v is None:
        return "-"
    return f"{int(v)}"


# ---------------------------------------------------------------------------
# Bar chart
# ---------------------------------------------------------------------------
def generate_dtw_barchart() -> None:
    overlap = [
        ("re100_karman", "karman_cloak", "Karman Cloak"),
        ("illusion_075L", "illusion_075L", "Illusion 0.75L"),
        ("illusion_1L", "illusion_1L", "Illusion 1L"),
        ("illusion_15L", "illusion_15L", "Illusion 1.5L"),
    ]

    names = []
    trn_vals = []
    rep_vals = []
    for trn_key, rep_key, name in overlap:
        m_trn = load_metrics(_OUT_BASE / "train" / trn_key)
        m_rep = load_metrics(_OUT_BASE / "reproduce" / rep_key)
        t = m_trn.get("dtw_sim_v5") or m_trn.get("sim_raw_mean", 0)
        r = m_rep.get("dtw_sim_v5") or m_rep.get("sim_raw_mean", 0)
        names.append(name)
        trn_vals.append(t if t and t > 0 else 0)
        rep_vals.append(r if r and r > 0 else 0)

    fig, ax = plt.subplots(figsize=(10, 5))
    x = np.arange(len(names))
    w = 0.35
    ax.bar(x - w / 2, rep_vals, w, label="Reproduce (legacy)", color="#d62728")
    ax.bar(x + w / 2, trn_vals, w, label="Train (V5)", color="#1f77b4")
    ax.set_xticks(x); ax.set_xticklabels(names, rotation=15, ha="right")
    ax.set_ylabel("DTW Similarity"); ax.set_ylim(0, 1.05)
    ax.set_title("Train vs Reproduce - DTW Similarity")
    ax.legend(); ax.grid(axis="y", alpha=0.3)
    for i, (tv, rv) in enumerate(zip(trn_vals, rep_vals)):
        if tv > 0: ax.text(i + w / 2, tv + 0.01, f"{tv:.3f}", ha="center", fontsize=9)
        if rv > 0: ax.text(i - w / 2, rv + 0.01, f"{rv:.3f}", ha="center", fontsize=9)
    fig.tight_layout()
    fig.savefig(_REPORT_DIR / "summary_dtw_barchart.png", dpi=150,
                bbox_inches="tight")
    plt.close(fig)
    log("  -> summary_dtw_barchart.png")


# ---------------------------------------------------------------------------
# Main
# ---------------------------------------------------------------------------
def main() -> int:
    _REPORT_DIR.mkdir(parents=True, exist_ok=True)
    log(f"Train dir: {_OUT_BASE / 'train'}")
    log(f"Reproduce dir: {_OUT_BASE / 'reproduce'}")
    log(f"Reports dir: {_REPORT_DIR}")

    table = generate_comparison_table()
    (_REPORT_DIR / "comparison_table.md").write_text(table)
    log("  -> comparison_table.md")

    generate_dtw_barchart()

    log("All reports generated.")
    return 0


if __name__ == "__main__":
    raise SystemExit(main())
