#!/usr/bin/env python3
"""Generate publication-quality figures for scratch-only DRL training results.

The script discovers completed canonical ``*_seed*`` runs, reads their TensorBoard
metrics, and deliberately excludes transfer-learning runs. It produces PNG and
PDF figures plus CSV/JSON summaries without rerunning CFD.

Usage:
    python src/drl_pinball/train/plot_scratch_results.py
    python src/drl_pinball/train/plot_scratch_results.py --window 20
"""
from __future__ import annotations

import argparse
import csv
import json
import re
import sys
from dataclasses import dataclass
from pathlib import Path
from typing import Dict, Iterable, List, Sequence

from drl_pinball.train.case_registry import CASE_IDS

import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt
import numpy as np
from matplotlib.lines import Line2D
from tensorboard.backend.event_processing.event_accumulator import EventAccumulator

_SRC = Path(__file__).resolve().parents[2]
if str(_SRC) not in sys.path:
    sys.path.insert(0, str(_SRC))

TAGS = {
    "reward": "eval/avg_reward",
    "r_cd": "eval/r_cd",
    "r_cl": "eval/r_cl",
    "r_sim": "eval/r_sim",
}
SEED_RE = re.compile(r"^(?P<case>.+)_seed(?P<seed>\d+)$")
COLORS = {
    "blue": "#2474B5",
    "orange": "#D67A1F",
    "green": "#278A62",
    "red": "#C8463A",
    "purple": "#7356A5",
    "gray": "#68717A",
    "black": "#202428",
}


@dataclass
class RunData:
    name: str
    case: str
    seed: int
    path: Path
    steps: np.ndarray
    metrics: Dict[str, np.ndarray]

    @property
    def reward(self) -> np.ndarray:
        return self.metrics["reward"]

    def summary(self) -> dict:
        best_idx = int(np.argmax(self.reward))
        tail = self.reward[-50:]
        return {
            "run": self.name,
            "case": self.case,
            "seed": self.seed,
            "episodes": len(self.reward),
            "best_reward": float(self.reward[best_idx]),
            "best_episode": int(self.steps[best_idx]),
            "final_reward": float(self.reward[-1]),
            "last50_mean": float(np.mean(tail)),
            "last50_median": float(np.median(tail)),
            "last50_std": float(np.std(tail)),
            "retention_ratio": float(np.median(tail) / self.reward[best_idx]),
            "collapse_severity": float(self.reward[best_idx] - np.median(tail)),
            "best_r_cd": float(self.metrics["r_cd"][best_idx]),
            "best_r_cl": float(self.metrics["r_cl"][best_idx]),
            "best_r_sim": float(self.metrics["r_sim"][best_idx]),
        }


def configure_style() -> None:
    plt.rcParams.update({
        "font.family": "DejaVu Sans",
        "font.size": 9,
        "axes.titlesize": 10,
        "axes.labelsize": 9,
        "legend.fontsize": 8,
        "xtick.labelsize": 8,
        "ytick.labelsize": 8,
        "axes.spines.top": False,
        "axes.spines.right": False,
        "axes.linewidth": 0.8,
        "grid.alpha": 0.18,
        "grid.linewidth": 0.6,
        "figure.dpi": 150,
        "savefig.dpi": 300,
        "pdf.fonttype": 42,
        "ps.fonttype": 42,
    })


def smooth(values: np.ndarray, window: int) -> np.ndarray:
    if window <= 1 or len(values) < window:
        return values.copy()
    kernel = np.ones(window, dtype=float) / window
    core = np.convolve(values, kernel, mode="valid")
    left = window // 2
    right = len(values) - len(core) - left
    return np.pad(core, (left, right), mode="edge")


def incumbent_indices(reward: np.ndarray) -> np.ndarray:
    """Index of the best checkpoint discovered up to each episode."""
    indices = np.empty(len(reward), dtype=int)
    best_idx = 0
    for idx, value in enumerate(reward):
        if value > reward[best_idx]:
            best_idx = idx
        indices[idx] = best_idx
    return indices


def running_best(reward: np.ndarray) -> np.ndarray:
    """Monotone best-so-far reward envelope."""
    return np.maximum.accumulate(reward)


def load_tag(tb_dir: Path, tag: str) -> tuple[np.ndarray, np.ndarray]:
    accumulator = EventAccumulator(str(tb_dir), size_guidance={"scalars": 0})
    accumulator.Reload()
    if tag not in accumulator.Tags().get("scalars", []):
        raise KeyError(f"Missing TensorBoard tag {tag!r} in {tb_dir}")
    by_step = {int(event.step): float(event.value) for event in accumulator.Scalars(tag)}
    steps = np.array(sorted(by_step), dtype=int)
    return steps, np.array([by_step[step] for step in steps], dtype=float)


def discover_runs(output_dir: Path) -> List[RunData]:
    runs: List[RunData] = []
    for run_dir in sorted(output_dir.glob("*_seed*")):
        if not run_dir.is_dir():
            continue
        match = SEED_RE.match(run_dir.name)
        if match is None or match.group("case") not in CASE_IDS:
            continue
        metrics: Dict[str, np.ndarray] = {}
        canonical_steps = None
        for short, tag in TAGS.items():
            steps, values = load_tag(run_dir / "tb", tag)
            if canonical_steps is None:
                canonical_steps = steps
            elif not np.array_equal(canonical_steps, steps):
                raise ValueError(f"Metric step mismatch in {run_dir}")
            metrics[short] = values
        runs.append(RunData(
            name=run_dir.name,
            case=match.group("case"),
            seed=int(match.group("seed")),
            path=run_dir,
            steps=canonical_steps,
            metrics=metrics,
        ))
    if not runs:
        raise FileNotFoundError(f"No scratch runs found under {output_dir}")
    return runs


def save_figure(fig: plt.Figure, out_dir: Path, stem: str) -> None:
    fig.savefig(out_dir / f"{stem}.png", bbox_inches="tight", facecolor="white")
    fig.savefig(out_dir / f"{stem}.pdf", bbox_inches="tight", facecolor="white")
    plt.close(fig)


def panel_label(ax: plt.Axes, label: str) -> None:
    ax.text(-0.11, 1.04, label, transform=ax.transAxes, fontsize=11,
            fontweight="bold", va="bottom")


def baseline_runs(runs: Sequence[RunData]) -> List[RunData]:
    result = sorted((run for run in runs if run.case == "kar_re100"), key=lambda run: run.seed)
    if len(result) < 2:
        raise ValueError("Re100 multi-seed plot requires at least two scratch seeds")
    return result


def plot_multiseed(runs: Sequence[RunData], out_dir: Path, window: int) -> None:
    base = baseline_runs(runs)
    steps = base[0].steps
    best_matrix = np.vstack([running_best(run.reward) for run in base])
    best_median = np.median(best_matrix, axis=0)
    best_q25, best_q75 = np.percentile(best_matrix, [25, 75], axis=0)
    best_minimum, best_maximum = best_matrix.min(axis=0), best_matrix.max(axis=0)

    fig, axes = plt.subplots(2, 1, figsize=(7.2, 7.0), sharex=True,
                             gridspec_kw={"height_ratios": [1.15, 1.0]})
    ax = axes[0]
    seed_colors = [COLORS["blue"], COLORS["orange"], COLORS["green"],
                   COLORS["red"], COLORS["purple"]]
    for run, color in zip(base, seed_colors):
        ax.plot(run.steps, run.reward, color=color, alpha=0.10, lw=0.45)
        ax.step(run.steps, running_best(run.reward), where="post", color=color, lw=1.45,
                label=f"Seed {run.seed}")
        best = int(np.argmax(run.reward))
        ax.scatter(run.steps[best], run.reward[best], s=23, color=color,
                   edgecolor="white", linewidth=0.5, zorder=4)
    ax.set_ylabel("Evaluation reward")
    ax.set_title("Raw exploration outcomes and best-so-far policy envelopes")
    ax.set_ylim(0, 1.02)
    ax.grid(True)
    ax.legend(ncol=5, loc="lower center", frameon=False)
    panel_label(ax, "a")

    ax = axes[1]
    ax.fill_between(steps, best_minimum, best_maximum, color=COLORS["blue"], alpha=0.10,
                    step="post", label="Seed range")
    ax.fill_between(steps, best_q25, best_q75, color=COLORS["blue"], alpha=0.24,
                    step="post", label="Interquartile range")
    ax.step(steps, best_median, where="post", color=COLORS["black"], lw=2.0,
            label="Median best-so-far")
    ax.set_xlabel("Training iteration")
    ax.set_ylabel("Best reward discovered")
    ax.set_title("Across-seed progress of the incumbent policy")
    ax.set_ylim(0, 1.02)
    ax.grid(True)
    ax.legend(frameon=False, ncol=3, loc="lower center")
    panel_label(ax, "b")
    fig.suptitle("Re100 exploration is volatile while policy discovery progresses monotonically",
                 fontsize=12, y=0.995)
    fig.text(0.5, 0.005,
             "Thin lines are per-iteration evaluations; stair-step curves retain the best checkpoint found up to each training iteration.",
             ha="center", fontsize=8, color=COLORS["gray"])
    fig.tight_layout(rect=(0, 0.03, 1, 0.97))
    save_figure(fig, out_dir, "fig01_re100_multiseed_learning")


def plot_exploration_drawdown(runs: Sequence[RunData], out_dir: Path) -> None:
    """Show discovery progress separately from exploratory policy drawdown."""
    base = baseline_runs(runs)
    steps = base[0].steps
    best_matrix = np.vstack([running_best(run.reward) for run in base])
    drawdown_matrix = np.vstack([
        running_best(run.reward) - run.reward for run in base
    ])

    fig, axes = plt.subplots(2, 1, figsize=(7.2, 6.2), sharex=True)
    ax = axes[0]
    best_med = np.median(best_matrix, axis=0)
    best_q25, best_q75 = np.percentile(best_matrix, [25, 75], axis=0)
    ax.fill_between(steps, best_q25, best_q75, step="post",
                    color=COLORS["green"], alpha=0.24)
    ax.step(steps, best_med, where="post", color=COLORS["green"], lw=2.0)
    ax.set_ylabel("Best reward discovered")
    ax.set_title("Incumbent policy improves as new checkpoints are discovered")
    ax.set_ylim(0, 1.02)
    ax.grid(True)
    panel_label(ax, "a")

    ax = axes[1]
    draw_med = np.median(drawdown_matrix, axis=0)
    draw_q25, draw_q75 = np.percentile(drawdown_matrix, [25, 75], axis=0)
    draw_max = np.max(drawdown_matrix, axis=0)
    ax.fill_between(steps, 0, draw_max, color=COLORS["red"], alpha=0.08,
                    label="Maximum seed drawdown")
    ax.fill_between(steps, draw_q25, draw_q75, color=COLORS["red"], alpha=0.22,
                    label="Interquartile range")
    ax.plot(steps, draw_med, color=COLORS["red"], lw=1.6,
            label="Median drawdown")
    ax.set_xlabel("Training iteration")
    ax.set_ylabel("Incumbent − current reward")
    ax.set_title("Aggressive exploration causes temporary performance drawdown")
    ax.set_ylim(bottom=0)
    ax.grid(True)
    ax.legend(frameon=False, ncol=3, loc="upper left")
    panel_label(ax, "b")
    fig.suptitle("Policy discovery and exploratory volatility are distinct training outcomes",
                 fontsize=12)
    fig.text(0.5, 0.005,
             "Drawdown measures the loss of the current exploratory policy relative to the best checkpoint already retained.",
             ha="center", fontsize=8, color=COLORS["gray"])
    fig.tight_layout(rect=(0, 0.035, 1, 0.95))
    save_figure(fig, out_dir, "fig02_re100_discovery_drawdown")


def plot_retention(runs: Sequence[RunData], out_dir: Path) -> None:
    base = baseline_runs(runs)
    summaries = [run.summary() for run in base]
    seeds = [str(item["seed"]) for item in summaries]
    best = np.array([item["best_reward"] for item in summaries])
    stable = np.array([item["last50_median"] for item in summaries])
    final = np.array([item["final_reward"] for item in summaries])
    retention = stable / best

    fig, axes = plt.subplots(1, 2, figsize=(7.2, 3.5), gridspec_kw={"width_ratios": [1.2, 1]})
    ax = axes[0]
    x = np.arange(len(seeds))
    for idx in range(len(seeds)):
        ax.plot([x[idx]] * 3, [best[idx], stable[idx], final[idx]], color="#B8BEC4", lw=1.2, zorder=1)
    ax.scatter(x, best, color=COLORS["green"], s=40, label="Best checkpoint", zorder=3)
    ax.scatter(x, stable, color=COLORS["blue"], s=40, marker="s", label="Last-50 median", zorder=3)
    ax.scatter(x, final, color=COLORS["red"], s=42, marker="^", label="Iteration 500", zorder=3)
    ax.set_xticks(x, seeds)
    ax.set_xlabel("Random seed")
    ax.set_ylabel("Evaluation reward")
    ax.set_ylim(0, 1.02)
    ax.set_title("Attained, retained, and final performance")
    ax.grid(True, axis="y")
    ax.legend(frameon=False, loc="lower left")
    panel_label(ax, "a")

    ax = axes[1]
    bars = ax.bar(x, retention, width=0.62, color=COLORS["blue"], alpha=0.82)
    ax.axhline(1.0, color=COLORS["black"], lw=0.8, ls="--")
    for bar, value in zip(bars, retention):
        ax.text(bar.get_x() + bar.get_width() / 2, value + 0.025, f"{value:.2f}",
                ha="center", va="bottom", fontsize=8)
    ax.set_xticks(x, seeds)
    ax.set_xlabel("Random seed")
    ax.set_ylabel("Retention ratio")
    ax.set_ylim(0, 1.08)
    ax.set_title("Last-50 median / best reward")
    ax.grid(True, axis="y")
    panel_label(ax, "b")
    fig.suptitle("High-reward policies are repeatedly found but not always preserved", fontsize=12)
    fig.tight_layout(rect=(0, 0, 1, 0.94))
    save_figure(fig, out_dir, "fig02_re100_policy_retention")


def plot_components(runs: Sequence[RunData], out_dir: Path, window: int) -> None:
    base = baseline_runs(runs)
    steps = base[0].steps
    labels = [("reward", "Total reward"), ("r_cd", r"Drag component $r_{CD}$"),
              ("r_cl", r"Lift component $r_{CL}$"), ("r_sim", r"DTW component $r_{sim}$")]
    fig, axes = plt.subplots(2, 2, figsize=(7.2, 5.8), sharex=True, sharey=True)
    for letter, ax, (metric, title) in zip("abcd", axes.flat, labels):
        matrix = np.vstack([
            run.metrics[metric][incumbent_indices(run.reward)] for run in base
        ])
        median = np.median(matrix, axis=0)
        q25, q75 = np.percentile(matrix, [25, 75], axis=0)
        minimum, maximum = matrix.min(axis=0), matrix.max(axis=0)
        ax.fill_between(steps, minimum, maximum, color=COLORS["purple"], alpha=0.09,
                        step="post")
        ax.fill_between(steps, q25, q75, color=COLORS["purple"], alpha=0.23,
                        step="post")
        ax.step(steps, median, where="post", color=COLORS["purple"], lw=1.8)
        ax.set_title(title)
        ax.set_ylim(0, 1.02)
        ax.grid(True)
        panel_label(ax, letter)
    for ax in axes[-1, :]:
        ax.set_xlabel("Training iteration")
    for ax in axes[:, 0]:
        ax.set_ylabel("Component value")
    legend = [
        Line2D([0], [0], color=COLORS["purple"], lw=1.8, label="Median incumbent policy"),
        Line2D([0], [0], color=COLORS["purple"], lw=7, alpha=0.23, label="Interquartile range"),
        Line2D([0], [0], color=COLORS["purple"], lw=7, alpha=0.09, label="Seed range"),
    ]
    fig.legend(handles=legend, frameon=False, ncol=3, loc="lower center", bbox_to_anchor=(0.5, 0.005))
    fig.suptitle("Reward composition of the best policy discovered so far", fontsize=12)
    fig.tight_layout(rect=(0, 0.06, 1, 0.95))
    save_figure(fig, out_dir, "fig03_re100_reward_components")


def select_one(runs: Sequence[RunData], case: str) -> RunData:
    matches = [run for run in runs if run.case == case]
    if len(matches) != 1:
        raise ValueError(f"Expected one scratch run for {case}, found {len(matches)}")
    return matches[0]


def plot_case_overview(runs: Sequence[RunData], out_dir: Path, window: int) -> None:
    base = baseline_runs(runs)
    base_steps = base[0].steps
    base_envelopes = np.vstack([running_best(run.reward) for run in base])
    base_median = np.median(base_envelopes, axis=0)
    base_q25, base_q75 = np.percentile(base_envelopes, [25, 75], axis=0)

    groups = [
        ("Cross-Reynolds-number scratch training",
         [("Re60", "kar_re60"), ("Re100 (5-seed median)", None),
          ("Re200", "kar_re200"), ("Re400", "kar_re400")]),
        ("Variable disturbance-cylinder size",
         [(r"$d/L_0=0.75$", "kar_d075"), (r"$d/L_0=1.0$ (baseline)", None),
          (r"$d/L_0=1.5$", "kar_d15"), (r"$d/L_0=2.0$", "kar_d2")]),
        ("Hydrodynamic illusion target size",
         [(r"$D_t/L_0=0.75$", "ill_075L"), (r"$D_t/L_0=1.0$", "ill_1L"),
          (r"$D_t/L_0=1.5$", "ill_15L"), (r"$D_t/L_0=2.0$", "ill_2L")]),
    ]
    palette = [COLORS["blue"], COLORS["orange"], COLORS["green"], COLORS["red"]]
    fig, axes = plt.subplots(1, 3, figsize=(11.3, 3.6), sharex=True, sharey=True)
    for letter, ax, (title, entries) in zip("abc", axes, groups):
        for (label, case), color in zip(entries, palette):
            if case is None:
                ax.fill_between(base_steps, base_q25, base_q75, color=color, alpha=0.13,
                                step="post")
                ax.step(base_steps, base_median, where="post", color=color, lw=1.8,
                        label=label)
            else:
                run = select_one(runs, case)
                ax.plot(run.steps, run.reward, color=color, alpha=0.10, lw=0.45)
                ax.step(run.steps, running_best(run.reward), where="post", color=color,
                        lw=1.7, label=label)
                best = int(np.argmax(run.reward))
                ax.scatter(run.steps[best], run.reward[best], color=color, s=18,
                           edgecolor="white", linewidth=0.4, zorder=3)
        ax.set_title(title)
        ax.set_xlabel("Training iteration")
        ax.set_ylim(0, 1.02)
        ax.grid(True)
        ax.legend(frameon=False, loc="lower right")
        panel_label(ax, letter)
    axes[0].set_ylabel("Evaluation reward")
    fig.suptitle("Best-so-far policy discovery across the expanded control matrix", fontsize=12)
    fig.text(0.5, 0.005,
             "Stair-step curves are incumbent best checkpoints; faint lines show exploratory evaluations. Re100 reports five-seed median and IQR.",
             ha="center", fontsize=8, color=COLORS["gray"])
    fig.tight_layout(rect=(0, 0.04, 1, 0.94))
    save_figure(fig, out_dir, "fig04_scratch_case_overview")


def plot_case_stability(runs: Sequence[RunData], out_dir: Path) -> None:
    ordered_cases = [
        "kar_re60", "kar_re100", "kar_re200", "kar_re400",
        "kar_d075", "kar_d15", "kar_d2",
        "ill_075L", "ill_1L", "ill_15L", "ill_2L",
    ]
    labels = ["Re60", "Re100", "Re200", "Re400", "d0.75", "d1.5", "d2.0",
              "Ill 0.75", "Ill 1.0", "Ill 1.5", "Ill 2.0"]
    values = []
    for case in ordered_cases:
        case_runs = [run for run in runs if run.case == case]
        per_run = [run.summary() for run in case_runs]
        values.append({
            "best": float(np.median([item["best_reward"] for item in per_run])),
            "stable": float(np.median([item["last50_median"] for item in per_run])),
            "final": float(np.median([item["final_reward"] for item in per_run])),
            "best_min": float(np.min([item["best_reward"] for item in per_run])),
            "best_max": float(np.max([item["best_reward"] for item in per_run])),
        })
    x = np.arange(len(labels))
    width = 0.24
    fig, ax = plt.subplots(figsize=(10.2, 4.1))
    ax.bar(x - width, [v["best"] for v in values], width, color=COLORS["green"], label="Best checkpoint")
    ax.bar(x, [v["stable"] for v in values], width, color=COLORS["blue"], label="Last-50 median")
    ax.bar(x + width, [v["final"] for v in values], width, color=COLORS["red"], label="Iteration 500")
    re100_idx = ordered_cases.index("kar_re100")
    low = values[re100_idx]["best_min"]
    high = values[re100_idx]["best_max"]
    center = values[re100_idx]["best"]
    ax.errorbar(re100_idx - width, center, yerr=[[center - low], [high - center]], fmt="none",
                color=COLORS["black"], capsize=3, lw=1.0, label="Re100 seed range")
    ax.axvline(3.5, color="#BBC1C6", lw=0.8)
    ax.axvline(6.5, color="#BBC1C6", lw=0.8)
    ax.text(1.5, 1.015, "Cross-Re", ha="center", va="bottom", fontsize=8)
    ax.text(5.0, 1.015, "Variable disturbance", ha="center", va="bottom", fontsize=8)
    ax.text(8.5, 1.015, "Illusion", ha="center", va="bottom", fontsize=8)
    ax.set_xticks(x, labels, rotation=32, ha="right")
    ax.set_ylabel("Evaluation reward")
    ax.set_ylim(0, 1.08)
    ax.set_title("Attained performance and late-training stability across scratch cases")
    ax.grid(True, axis="y")
    ax.legend(frameon=False, ncol=4, loc="lower center")
    fig.text(0.5, 0.005,
             "Re100 bars are medians across five seeds; all other cases contain one seed and do not represent seed uncertainty.",
             ha="center", fontsize=8, color=COLORS["gray"])
    fig.tight_layout(rect=(0, 0.04, 1, 1))
    save_figure(fig, out_dir, "fig05_case_performance_stability")


def write_summaries(runs: Sequence[RunData], out_dir: Path) -> None:
    rows = [run.summary() for run in runs]
    fields = list(rows[0])
    with (out_dir / "scratch_training_summary.csv").open("w", newline="") as handle:
        writer = csv.DictWriter(handle, fieldnames=fields)
        writer.writeheader()
        writer.writerows(rows)
    base_rows = [row for row in rows if row["case"] == "kar_re100"]
    report = {
        "scope": "Scratch-only DRL training; transfer runs excluded",
        "run_count": len(rows),
        "case_count": len(set(row["case"] for row in rows)),
        "re100_multiseed": {
            "seed_count": len(base_rows),
            "best_reward_mean": float(np.mean([row["best_reward"] for row in base_rows])),
            "best_reward_std": float(np.std([row["best_reward"] for row in base_rows])),
            "best_reward_min": float(np.min([row["best_reward"] for row in base_rows])),
            "best_reward_max": float(np.max([row["best_reward"] for row in base_rows])),
            "retention_ratio_mean": float(np.mean([row["retention_ratio"] for row in base_rows])),
        },
        "runs": rows,
    }
    with (out_dir / "scratch_training_summary.json").open("w") as handle:
        json.dump(report, handle, indent=2)


def main() -> int:
    parser = argparse.ArgumentParser(description=__doc__)
    script_dir = Path(__file__).resolve().parent
    parser.add_argument("--output-root", type=Path, default=script_dir / "output")
    parser.add_argument("--figure-dir", type=Path, default=script_dir / "results" / "figures")
    parser.add_argument("--window", type=int, default=20, help="Moving-average window in episodes")
    args = parser.parse_args()
    args.figure_dir.mkdir(parents=True, exist_ok=True)
    configure_style()
    runs = discover_runs(args.output_root)
    write_summaries(runs, args.figure_dir)
    plot_multiseed(runs, args.figure_dir, args.window)
    plot_exploration_drawdown(runs, args.figure_dir)
    plot_retention(runs, args.figure_dir)
    plot_components(runs, args.figure_dir, args.window)
    plot_case_overview(runs, args.figure_dir, args.window)
    plot_case_stability(runs, args.figure_dir)
    print(f"Generated 6 figure sets from {len(runs)} scratch runs in {args.figure_dir}")
    return 0


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