#!/usr/bin/env python3
"""Render publication SR figures and tables from frozen article CSV/JSON artifacts."""
from __future__ import annotations

import argparse
import csv
import hashlib
import json
from pathlib import Path
from typing import Any, Mapping, Sequence

import matplotlib.pyplot as plt
import numpy as np

TRAINING_SCENES = (
    "karman_re50", "karman_re100", "karman_re200", "karman_re400",
    "illusion_0.75L", "illusion_1L", "illusion_1.5L",
)
LABELS = ("K50", "K100", "K200", "K400", "I0.75L", "I1L", "I1.5L")
ACTIONS = ("front", "upper", "lower")
PACKAGE = "article2-plotting-package-20260721"
STANDARD = "article2-timeseries-csv-20260720"
LONG_SR = "article2-long-timeseries-csv-20260720"
COLORS = {"Target": "#222222", "PPO": "#2878B5", "SR": "#D95319"}


def read_rows(path: Path) -> list[dict[str, str]]:
    with path.open(newline="", encoding="utf-8") as handle:
        return list(csv.DictReader(handle))


def write_rows(path: Path, fieldnames: Sequence[str], rows: Sequence[Mapping[str, Any]]) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    with path.open("w", newline="", encoding="utf-8") as handle:
        writer = csv.DictWriter(handle, fieldnames=fieldnames, lineterminator="\n")
        writer.writeheader()
        writer.writerows(rows)


def final_similarity(root: Path, package: str, scene: str, source: str) -> float:
    return float(read_rows(root / package / scene / f"{source}_dtw_convergence.csv")[-1]["similarity"])


def save(fig: plt.Figure, output: Path, stem: str) -> None:
    output.mkdir(parents=True, exist_ok=True)
    fig.savefig(output / f"{stem}.png", dpi=300, bbox_inches="tight")
    fig.savefig(output / f"{stem}.pdf", bbox_inches="tight")
    plt.close(fig)


def plot_performance(root: Path, output: Path) -> None:
    ppo200 = np.array([final_similarity(root, STANDARD, s, "ppo") for s in TRAINING_SCENES])
    sr200 = np.array([final_similarity(root, STANDARD, s, "sr") for s in TRAINING_SCENES])
    ppo400 = np.array([final_similarity(root / PACKAGE / "long_ppo", "", s, "ppo") for s in TRAINING_SCENES])
    sr400 = np.array([final_similarity(root, LONG_SR, s, "sr") for s in TRAINING_SCENES])
    x = np.arange(len(LABELS)); width = 0.36
    fig, axes = plt.subplots(1, 2, figsize=(12.2, 4.7), gridspec_kw={"width_ratios": (1.15, 1)})
    axes[0].bar(x - width / 2, ppo400, width, label="PPO — 400 control steps", color=COLORS["PPO"])
    axes[0].bar(x + width / 2, sr400, width, label="SR — 400 control steps", color=COLORS["SR"])
    axes[0].set_ylabel("Legacy DTW similarity"); axes[0].set_ylim(0.75, 1.005)
    axes[0].set_title("(a) Closed-loop performance at 400 control steps")
    axes[0].legend(frameon=False, fontsize=9)
    axes[1].axhline(0, color="0.25", linewidth=1)
    axes[1].scatter(x - 0.08, ppo400 - ppo200, label="PPO: 400 − 200 control steps", marker="o", s=48, color=COLORS["PPO"])
    axes[1].scatter(x + 0.08, sr400 - sr200, label="SR: 400 − 200 control steps", marker="s", s=48, color=COLORS["SR"])
    axes[1].set_ylabel("Change in legacy DTW similarity")
    axes[1].set_title("(b) Duration stability from 200 to 400 steps")
    axes[1].legend(frameon=False, fontsize=9)
    for ax in axes:
        ax.set_xticks(x, LABELS); ax.set_xlabel("Training case"); ax.grid(axis="y", alpha=0.25)
    fig.tight_layout(); save(fig, output, "01_training_case_performance")


def generalization_points(repo_root: Path) -> list[dict[str, Any]]:
    path = repo_root / "src/SR_analysis/results/runs/article2-generalization-summary-20260720/generalization.csv"
    points: list[dict[str, Any]] = [dict(row) for row in read_rows(path)]
    run_root = repo_root / "src/SR_analysis/results/runs"
    for scene in TRAINING_SCENES:
        objective = "karman" if scene.startswith("karman") else "illusion"
        parameter = float(scene.removeprefix("karman_re").removeprefix("illusion_").removesuffix("L"))
        points.append({"objective": objective, "scene": scene, "parameter": parameter,
                       "classification": "training", "legacy_dtw": final_similarity(run_root, STANDARD, scene, "sr")})
    return points


def plot_generalization(repo_root: Path, output: Path) -> None:
    rows = generalization_points(repo_root)
    fig, axes = plt.subplots(1, 2, figsize=(11, 4.4))
    styles = {"training": ("o", "#222222", "Training condition"),
              "interpolation": ("^", "#2E8B57", "Unseen interpolation"),
              "extrapolation": ("s", "#A23B72", "Unseen extrapolation")}
    for ax, objective, xlabel in zip(axes, ("karman", "illusion"), ("Reynolds-number label", "Target-size label, L")):
        subset = [r for r in rows if r["objective"] == objective]
        for category, (marker, color, label) in styles.items():
            selected = sorted((r for r in subset if r["classification"] == category), key=lambda r: float(r["parameter"]))
            ax.scatter([float(r["parameter"]) for r in selected], [float(r["legacy_dtw"]) for r in selected],
                       marker=marker, color=color, edgecolor="white", linewidth=0.6, s=65, label=label, zorder=3)
        ax.set_xlabel(xlabel); ax.set_ylabel("Legacy DTW similarity"); ax.set_ylim(0.7, 1.01)
        ax.grid(alpha=0.25); ax.legend(frameon=False, fontsize=8.5)
        ax.set_title("Kármán conditions" if objective == "karman" else "Illusion conditions")
    fig.suptitle("Frozen SR controller: training and pointwise unseen 200-step realizations", y=1.01)
    fig.tight_layout(); save(fig, output, "02_pointwise_generalization")


def plot_steady(repo_root: Path, output: Path) -> None:
    rows = read_rows(repo_root / "src/SR_analysis/results/runs/article2-steady-analysis-20260720/steady_rotation_sweep.csv")
    magnitude = np.asarray([float(r["rear_alpha_magnitude"]) for r in rows])
    similarity = np.asarray([float(r["legacy_dtw"]) for r in rows])
    error = np.asarray([float(r["mean_abs_sensor_error"]) for r in rows])
    fig, ax = plt.subplots(figsize=(7.2, 4.5)); secondary = ax.twinx()
    ax.plot(magnitude, similarity, marker="o", label="DTW similarity")
    secondary.plot(magnitude, error, marker="s", linestyle="--", label="Mean sensor error", color="tab:orange")
    ax.axvline(3.4312086805, linestyle=":", color="0.35", label="SR rear constant 3.4312")
    ax.set_xlabel("Rear counter-rotation magnitude |α|"); ax.set_ylabel("Legacy DTW similarity")
    secondary.set_ylabel("Mean absolute sensor error"); ax.set_ylim(0.6, 1.01); ax.grid(alpha=0.25)
    lines = ax.get_lines() + secondary.get_lines(); ax.legend(lines, [line.get_label() for line in lines], frameon=False, loc="center right")
    ax.set_title("Steady-cloak constant-rotation calibration"); fig.tight_layout()
    save(fig, output, "03_steady_rotation_calibration")


def build_offline_table(root: Path, tables: Path) -> None:
    source = root / PACKAGE / "offline_predictions.csv"; rows = read_rows(source); output = []
    for scene in TRAINING_SCENES:
        subset = [r for r in rows if r["scene"] == scene]
        values = {a: float(np.sqrt(np.mean([float(r[f"residual_alpha_{a}"]) ** 2 for r in subset]))) for a in ACTIONS}
        aggregate = float(np.sqrt(np.mean([float(r[f"residual_alpha_{a}"]) ** 2 for r in subset for a in ACTIONS])))
        output.append({"objective": subset[0]["objective"], "training_case": scene, "n_causally_aligned_states": len(subset),
                       **{f"rmse_alpha_{a}": f"{values[a]:.8f}" for a in ACTIONS}, "aggregate_rmse_alpha": f"{aggregate:.8f}"})
    fields = ("objective", "training_case", "n_causally_aligned_states", *(f"rmse_alpha_{a}" for a in ACTIONS), "aggregate_rmse_alpha")
    write_rows(tables / "offline_action_rmse.csv", fields, output)
    lines = ["# Offline SR-vs-PPO next-action RMSE", "",
             "RMSE is evaluated on causally aligned PPO-visited states: each post-action state predicts the next PPO action. This is an offline imitation diagnostic, not closed-loop performance.", "",
             "| Objective | Training case | States | Front RMSE | Upper RMSE | Lower RMSE | Aggregate RMSE |",
             "|---|---|---:|---:|---:|---:|---:|"]
    for r in output:
        lines.append(f"| {r['objective']} | {r['training_case']} | {r['n_causally_aligned_states']} | {r['rmse_alpha_front']} | {r['rmse_alpha_upper']} | {r['rmse_alpha_lower']} | {r['aggregate_rmse_alpha']} |")
    lines += ["", "All action errors are in dimensionless surface-speed α. Aggregate RMSE pools the three action channels and all aligned states within a case."]
    (tables / "offline_action_rmse.md").write_text("\n".join(lines) + "\n", encoding="utf-8")


def _load_json(path: Path) -> dict[str, Any]:
    return json.loads(path.read_text(encoding="utf-8"))


def build_term_deletion_table(repo_root: Path, root: Path, tables: Path) -> None:
    formula_root = repo_root / "src/SR_analysis/results/runs/article-ablation-formulas-v2-20260718/formulas"
    parent_root = repo_root / "src/SR_analysis/results/runs"
    specs = {
        "k_front0": ("karman", "front odd projection", "karman_front__delete_t0.json", "article-refit-karman-topology-a-20260718/formulas/joint_front.json", "article-L2-karman-20260718"),
        "k_rear0": ("karman", "upper/lower shared-symmetry rear", "karman_rear__delete_t0.json", "article-refit-karman-topology-a-20260718/formulas/joint_rear_shared_upper.json", "article-L2-karman-20260718"),
        "k_rear1": ("karman", "upper/lower shared-symmetry rear", "karman_rear__delete_t1.json", "article-refit-karman-topology-a-20260718/formulas/joint_rear_shared_upper.json", "article-L2-karman-20260718"),
        "i_front0": ("illusion", "front odd projection", "illusion_front__delete_t0.json", "article-refit-illusion-topology-a-20260718/formulas/joint_front.json", "article-L2-illusionA-20260718"),
        "i_front1": ("illusion", "front odd projection", "illusion_front__delete_t1.json", "article-refit-illusion-topology-a-20260718/formulas/joint_front.json", "article-L2-illusionA-20260718"),
        "i_rear0": ("illusion", "upper/lower shared-symmetry rear", "illusion_rear__delete_t0.json", "article-refit-illusion-topology-a-20260718/formulas/joint_rear_shared_upper.json", "article-L2-illusionA-20260718"),
        "i_rear1": ("illusion", "upper/lower shared-symmetry rear", "illusion_rear__delete_t1.json", "article-refit-illusion-topology-a-20260718/formulas/joint_rear_shared_upper.json", "article-L2-illusionA-20260718"),
    }
    ablations = read_rows(root / PACKAGE / "ablation_summary.csv"); output = []
    for variant, (objective, head, variant_file, parent_file, parent_run) in specs.items():
        variant_json = _load_json(formula_root / variant_file); parent_json = _load_json(parent_root / parent_file)
        deleted = variant_json["variant_metadata"]["term_identity"]
        variant_rows = [r for r in ablations if r["variant"] == variant]
        similarities = [float(r["similarity"]) for r in variant_rows]
        parent_metrics = {}
        for path in (parent_root / parent_run / "validations").glob("*.json"):
            record = _load_json(path); parent_metrics[record["scene"]] = float(record["metrics"]["legacy_reference_cycle_vs_last_recorded_cycle"]["similarity"])
        for r in sorted(variant_rows, key=lambda x: TRAINING_SCENES.index(x["scene"])):
            parent_value = parent_metrics[r["scene"]]; value = float(r["similarity"])
            output.append({"objective": objective, "head": head, "variant": variant,
                           "parent_formula": parent_json["deployment_expression"], "deleted_term": deleted,
                           "retained_formula": variant_json["deployment_expression"], "case": r["scene"], "steps": r["steps"],
                           "legacy_dtw": f"{value:.8f}", "parent_legacy_dtw_40_step": f"{parent_value:.8f}",
                           "delta_vs_parent_40_step": f"{value-parent_value:+.8f}", "variant_mean": f"{np.mean(similarities):.8f}",
                           "variant_min": f"{np.min(similarities):.8f}"})
    fields = tuple(output[0]); write_rows(tables / "term_deletion.csv", fields, output)
    lines = ["# Forty-step closed-loop term deletion", "",
             "Every comparison uses the same 40-control-step window and the legacy DTW metric. Parent values come from the scientifically comparable parent L2 runs; deltas are deletion minus parent. Formula identities are read from the canonical parent and deletion JSON artifacts.", "",
             "Front formulas are scalar generators deployed through the odd projection `α_F(x) = ½[f(x) − f(Gx)]`. Rear formulas define the upper action; the lower action is mapped by `α_L(x) = −α_U(Gx)`.", "",
             "| Objective | Head | Variant | Parent formula | Deleted term | Retained formula | Case | DTW | Parent DTW | Δ | Mean | Min |",
             "|---|---|---|---|---|---|---|---:|---:|---:|---:|---:|"]
    for r in output:
        fmt = lambda s: f"`{s.replace('*', '·')}`"
        lines.append(f"| {r['objective']} | {r['head']} | {r['variant']} | {fmt(r['parent_formula'])} | {fmt(r['deleted_term'])} | {fmt(r['retained_formula'])} | {r['case']} | {r['legacy_dtw']} | {r['parent_legacy_dtw_40_step']} | {r['delta_vs_parent_40_step']} | {r['variant_mean']} | {r['variant_min']} |")
    lines += ["", "Interpretation: Kármán deletion effects are term-dependent, with deleting the rear constant producing the largest degradation; the tested front term is weak over this short window. Illusion deletions remain stable and comparatively close to their parent, so these runs do not establish uniqueness of every term. Means and minima summarize the listed training cases only."]
    (tables / "term_deletion.md").write_text("\n".join(lines) + "\n", encoding="utf-8")


def circular_lag(reference: np.ndarray, signal: np.ndarray) -> int:
    """Return deterministic roll applied to signal to maximize normalized circular correlation."""
    ref = np.asarray(reference, dtype=float) - np.mean(reference); sig = np.asarray(signal, dtype=float) - np.mean(signal)
    if len(ref) != len(sig) or len(ref) == 0: raise ValueError("signals must have equal non-zero length")
    scores = np.array([np.dot(ref, np.roll(sig, lag)) for lag in range(len(ref))]) / (np.linalg.norm(ref) * np.linalg.norm(sig))
    lag = int(np.flatnonzero(np.isclose(scores, scores.max(), rtol=0, atol=1e-12))[0])
    return lag if lag <= len(ref) // 2 else lag - len(ref)


def plot_example_timeseries(root: Path, output: Path, package_dir: Path) -> dict[str, Any]:
    source = root / STANDARD / "karman_re100"
    data = {name: read_rows(source / f"{name.lower()}_wide.csv") for name in ("Target", "PPO", "SR")}
    start, stop = 96, 146  # 50 samples, approximately three target cycles, late in the 150-row target record.
    phase = np.linspace(0, 6 * np.pi, stop - start, endpoint=False)
    signal_name = "sensors_center_uy"
    reference = np.array([float(r[signal_name]) for r in data["Target"][start:stop]])
    lags = {"Target": 0}
    aligned: dict[str, list[dict[str, str]]] = {"Target": data["Target"][start:stop]}
    for name in ("PPO", "SR"):
        segment = data[name][start:stop]
        lag = circular_lag(reference, np.array([float(r[signal_name]) for r in segment]))
        lags[name] = lag; aligned[name] = list(np.roll(np.asarray(segment, dtype=object), lag))
    fig, axes = plt.subplots(2, 3, figsize=(12, 7.2))
    sensors = (("upper", "Upper sensor"), ("center", "Center sensor"), ("lower", "Lower sensor"))
    styles = {"Target": "-", "PPO": "--", "SR": "-."}
    for col, (sensor, title) in enumerate(sensors):
        for name in ("Target", "PPO", "SR"):
            rows = aligned[name]
            axes[0, col].plot([float(r[f"sensors_{sensor}_ux"]) for r in rows], [float(r[f"sensors_{sensor}_uy"]) for r in rows],
                              color=COLORS[name], linestyle=styles[name], linewidth=1.7, label=name)
        axes[0, col].set_xlabel(r"$u/U_0$"); axes[0, col].set_ylabel(r"$v/U_0$"); axes[0, col].set_title(title); axes[0, col].grid(alpha=0.22)
        axes[0, col].scatter(float(aligned["Target"][0][f"sensors_{sensor}_ux"]), float(aligned["Target"][0][f"sensors_{sensor}_uy"]), color=COLORS["Target"], marker="o", s=20, zorder=4)
    axes[0, 0].legend(frameon=False, ncol=3, fontsize=9)
    for col, action in enumerate(ACTIONS):
        for name in ("PPO", "SR"):
            axes[1, col].plot(phase / (2 * np.pi), [float(r[f"actions_alpha_{action}"]) for r in aligned[name]], color=COLORS[name], linestyle=styles[name], linewidth=1.6, label=name)
        axes[1, col].set_xlabel("Target phase / 2π"); axes[1, col].set_ylabel("Action α"); axes[1, col].set_title(f"{action.capitalize()} action"); axes[1, col].set_xticks((0, 1, 2, 3)); axes[1, col].grid(alpha=0.22)
    axes[1, 0].legend(frameon=False, ncol=2, fontsize=9)
    fig.suptitle("Kármán Re100: three stable target cycles with phase-aligned controllers", y=0.995)
    fig.tight_layout(); save(fig, output, "06_example_timeseries_karman_re100")
    metadata = {"scene": "karman_re100", "source_package": STANDARD, "target_index_start_inclusive": start,
                "target_index_stop_exclusive": stop, "target_t_D_start": float(data["Target"][start]["t_D"]),
                "target_t_D_end": float(data["Target"][stop-1]["t_D"]), "samples": stop-start, "displayed_target_cycles": 3,
                "alignment_signal": signal_name, "alignment_method": "mean-centered normalized circular cross-correlation; smallest maximizing roll selected",
                "roll_lags_samples": lags, "control_dt_D_over_U0": 0.4, "target_actions_plotted": False}
    (package_dir / "phase_alignment.json").write_text(json.dumps(metadata, indent=2) + "\n", encoding="utf-8")
    return metadata


def update_manifest(repo_root: Path, package_dir: Path) -> None:
    old = _load_json(package_dir / "manifest.json"); summary = dict(old.get("summary", {}))
    summary.pop("diagnostic_figures", None); summary.update({"publication_figures": 5, "publication_tables": 2, "presentation_pages": 2})
    suffixes = {".csv", ".png", ".pdf", ".md", ".json", ".npz"}; artifacts = []
    for path in sorted(p for p in package_dir.rglob("*") if p.is_file() and p.suffix in suffixes and p.name != "manifest.json"):
        artifacts.append({"path": str(path.relative_to(repo_root)), "sha256": hashlib.sha256(path.read_bytes()).hexdigest()})
    record = {"schema_version": "sr-plotting-package-v3", "source_policy": old["source_policy"], "summary": summary, "artifacts": artifacts}
    (package_dir / "manifest.json").write_text(json.dumps(record, indent=2) + "\n", encoding="utf-8")


def write_package_readme(package_dir: Path, metadata: Mapping[str, Any]) -> None:
    text = f"""# SR publication plotting package

This package contains five generated publication figures, two replacement tables, and two PPT-ready 16:9 summary pages. All values are derived from frozen canonical artifacts without formula refitting.

- `01_training_case_performance`: panel (a) compares PPO and SR absolute closed-loop legacy DTW similarity at **400 control steps** for all seven training cases; panel (b) reports each controller's 400-minus-200 change. Here 200 and 400 are run durations in control steps, not Reynolds numbers.
- `02_pointwise_generalization`: 200-step SR results at training conditions and sampled unseen interpolation/extrapolation conditions. Each unseen point is one realization. Categories are unconnected because the evidence does not establish a continuous parameter law or statistical robustness.
- `03_steady_rotation_calibration`: the retained disturbance-free constant rear-rotation sweep diagnostic.
- `06_example_timeseries_karman_re100`: late-window phase portraits and PPO/SR actions over approximately three target cycles. Target indices {metadata['target_index_start_inclusive']}–{metadata['target_index_stop_exclusive'] - 1} (`t_D={metadata['target_t_D_start']:.1f}`–`{metadata['target_t_D_end']:.1f}`) are used. PPO and SR segments at the same indices are circularly aligned to target `sensors_center_uy` by mean-centered normalized cross-correlation; deterministic roll lags are PPO={metadata['roll_lags_samples']['PPO']} and SR={metadata['roll_lags_samples']['SR']} samples. No target actions exist or are plotted.
- `07_flow_field_comparison_karman_re100`: a deterministic stable-cycle snapshot trio. Over [96,146), downstream center-sensor `(ux,uy)` is standardized separately per trajectory; an exhaustive joint search minimizes wrapped phase errors plus a 0.001 rad/sample temporal-separation penalty with a same-direction branch check. This matches center-sensor limit-cycle phase, not six-sensor state equality or exact full-field identity.
- `tables/offline_action_rmse.*`: offline SR-vs-PPO next-action RMSE on causally aligned PPO-visited states. This is not closed-loop performance.
- `tables/term_deletion.*`: actual 40-step closed-loop deletion results, canonical formula identities, same-window parent comparisons, and aggregate mean/min summaries.

The `presentation/` directory contains two white-background 16:9 PNG/PDF pages and a source-artifact README. PNG and PDF files share each retained figure stem. `phase_alignment.json` records the exact Figure 06 selection/alignment contract. `manifest.json` hashes every package CSV, Markdown, JSON/NPZ metadata or field artifact, and publication figure.
"""
    (package_dir / "README.md").write_text(text, encoding="utf-8")


def main(argv: Sequence[str] | None = None) -> int:
    parser = argparse.ArgumentParser(description=__doc__); parser.add_argument("--repo-root", type=Path, default=Path(__file__).resolve().parents[3]); parser.add_argument("--output-dir", type=Path, required=True)
    args = parser.parse_args(argv); repo_root = args.repo_root.resolve(); root = repo_root / "src/SR_analysis/results/runs"
    output = args.output_dir if args.output_dir.is_absolute() else repo_root / args.output_dir; package_dir = root / PACKAGE
    plt.style.use("seaborn-v0_8-whitegrid"); plt.rcParams.update({"font.size": 10, "axes.titlesize": 11, "axes.labelsize": 10, "legend.fontsize": 9, "pdf.fonttype": 42})
    for obsolete in ("04_offline_action_rmse", "05_term_deletion"):
        for suffix in (".png", ".pdf"): (output / f"{obsolete}{suffix}").unlink(missing_ok=True)
    plot_performance(root, output); plot_generalization(repo_root, output); plot_steady(repo_root, output)
    tables = package_dir / "tables"; build_offline_table(root, tables); build_term_deletion_table(repo_root, root, tables)
    metadata = plot_example_timeseries(root, output, package_dir); write_package_readme(package_dir, metadata); update_manifest(repo_root, package_dir)
    return 0


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