#!/usr/bin/env python3
"""Plot the validated V5 and Legacy three-role reproduction summaries."""
from __future__ import annotations

import argparse
import hashlib
import json
import os
import shutil
import tempfile
from pathlib import Path

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

L0 = 20.0
VORTICITY_LIMIT = 0.001
SENSOR_PAIRS = ((0, 1, "upper"), (2, 3, "center"), (4, 5, "lower"))
SENSOR_COLORS = ("#0072B2", "#D55E00", "#009E73")
V5_GROUPS = {
    "ill_075L_seed43": "ill_075L", "ill_1L_seed43": "ill_1L",
    "ill_15L_seed43": "ill_15L", "ill_2L_seed43": "ill_2L",
    "kar_d075_seed44": "kar_d075", "kar_d15_seed45": "kar_d15",
    "kar_d2_seed45": "kar_d2", "kar_re60_seed43": "kar_re60",
    "kar_re100_seed41": "kar_re100", "kar_re100_seed42": "kar_re100",
    "kar_re100_seed43": "kar_re100", "kar_re100_seed44": "kar_re100",
    "kar_re100_seed45": "kar_re100", "kar_re200_seed43": "kar_re200",
    "kar_re400_seed43": "kar_re400",
}
LEGACY_PERIODIC = (
    "illusion_075L", "illusion_1L", "illusion_15L", "karman_re50",
    "karman_re100", "karman_re200", "karman_re400",
)
LEGACY_VORTEX = (
    "vortex_lamb_y000", "vortex_taylor_ym2L", "vortex_taylor_ym1L",
    "vortex_taylor_y000", "vortex_taylor_yp1L", "vortex_taylor_yp2L",
)
LEGACY_GROUPS = LEGACY_PERIODIC + ("steady",) + LEGACY_VORTEX + ("erase",)
SR_TRAINING = LEGACY_PERIODIC
SR_GENERALIZATION = (
    "karman_re25", "karman_re70", "karman_re150", "karman_re300",
    "illusion_05L", "illusion_06L", "illusion_08L", "illusion_12L", "illusion_2L",
)
SR_VARIANT_CASE = {
    "k_front0": "karman_re100", "k_rear0": "karman_re100", "k_rear1": "karman_re100",
    "i_front0": "illusion_15L", "i_front1": "illusion_15L", "i_rear0": "illusion_15L", "i_rear1": "illusion_15L",
}


def fail(message: str) -> None:
    raise RuntimeError(message)


def require_exact_groups(root: Path) -> None:
    v5_controlled = {p.parent.name for p in (root / "v5").glob("*/controlled") if p.is_dir()}
    if v5_controlled != set(V5_GROUPS):
        fail(f"V5 controlled groups differ: missing={set(V5_GROUPS)-v5_controlled}, extra={v5_controlled-set(V5_GROUPS)}")
    legacy = {p.name for p in (root / "legacy").iterdir() if p.is_dir()}
    if legacy != set(LEGACY_GROUPS):
        fail(f"Legacy groups differ: missing={set(LEGACY_GROUPS)-legacy}, extra={legacy-set(LEGACY_GROUPS)}")


def roles_for(root: Path, pipeline: str, group: str):
    if pipeline == "v5":
        case = V5_GROUPS[group]
        return (("target", root / "v5" / case / "target"),
                ("controlled", root / "v5" / group / "controlled"),
                ("physical-zero", root / "v5" / case / "zero"))
    middle = "constant" if group == "steady" else "controlled"
    return (("target", root / "legacy" / group / "target"),
            (middle, root / "legacy" / group / middle),
            ("physical-zero", root / "legacy" / group / "zero"))


def field_spec(pipeline: str, group: str, role_dir: Path):
    if pipeline == "v5" or group in LEGACY_PERIODIC + SR_GENERALIZATION or (group == "erase" and role_dir.name != "target"):
        return "phase_fields.npz", "target_phase", 0.0
    if group in LEGACY_VORTEX:
        return "event_fields.npz", "relative_offsets", -10
    return "late_field.npz", "field_indices", None


def load_role(role_dir: Path, pipeline: str, group: str):
    if not role_dir.is_dir():
        fail(f"missing role directory: {role_dir}")
    required = {"metadata.json", "timeseries.csv"}
    names = {p.name for p in role_dir.iterdir() if p.is_file()}
    if not required <= names:
        fail(f"missing required files in {role_dir}: {required-names}")
    field_name, selector_name, selector_value = field_spec(pipeline, group, role_dir)
    field_candidates = names & {"phase_fields.npz", "event_fields.npz", "late_field.npz"}
    if field_candidates != {field_name}:
        fail(f"field files differ in {role_dir}: expected {field_name}, found {sorted(field_candidates)}")
    metadata = json.loads((role_dir / "metadata.json").read_text())
    units = metadata.get("units", {}).get("sensors", "raw sensor units")
    data = np.genfromtxt(role_dir / "timeseries.csv", delimiter=",", names=True)
    if data.ndim != 1 or data.size == 0:
        fail(f"invalid timeseries shape in {role_dir}: {data.shape}")
    expected_sensors = {f"sensors_{i}" for i in range(6)}
    actual_sensors = {n for n in (data.dtype.names or ()) if n.startswith("sensors_")}
    if actual_sensors != expected_sensors:
        fail(f"sensor columns differ in {role_dir}: {actual_sensors}")
    sensors = np.column_stack([data[f"sensors_{i}"] for i in range(6)])
    if sensors.shape != (data.size, 6) or not np.isfinite(sensors).all():
        fail(f"invalid sensor values in {role_dir}: {sensors.shape}")
    field_path = role_dir / field_name
    with np.load(field_path) as z:
        required_keys = {"ux", "uy", selector_name}
        if not required_keys <= set(z.files):
            fail(f"missing field keys in {field_path}: {required_keys-set(z.files)}")
        ux, uy = np.array(z["ux"]), np.array(z["uy"])
        selector = np.array(z[selector_name])
    if ux.ndim != 3 or ux.shape != uy.shape or ux.shape[0] != selector.shape[0]:
        fail(f"invalid canonical field shapes in {field_path}: ux={ux.shape}, uy={uy.shape}, selector={selector.shape}")
    expected_count = 8 if field_name == "phase_fields.npz" else 5 if field_name == "event_fields.npz" else 1
    if ux.shape[0] != expected_count or selector.shape != (expected_count,):
        fail(f"unexpected retained field count in {field_path}: {ux.shape[0]}")
    if selector_value is not None and not np.isclose(selector[0], selector_value, atol=1e-12, rtol=0):
        fail(f"slot 0 selector in {field_path} is {selector[0]}, expected {selector_value}")
    if not np.isfinite(ux[0]).all() or not np.isfinite(uy[0]).all():
        fail(f"nonfinite field in {field_path} slot 0")
    omega = np.gradient(uy[0], axis=1) - np.gradient(ux[0], axis=0)
    return omega, sensors, units, field_path, selector_name, selector[0]


def padded_limits(values):
    lo, hi = float(np.min(values)), float(np.max(values))
    if not np.isfinite([lo, hi]).all():
        fail("nonfinite sensor limits")
    span = hi - lo
    pad = 0.05 * span if span > 0 else max(abs(lo) * 0.05, 1e-12)
    return [lo - pad, hi + pad]


def sha256_file(path: Path):
    digest = hashlib.sha256()
    with path.open("rb") as stream:
        for chunk in iter(lambda: stream.read(1024 * 1024), b""):
            digest.update(chunk)
    return digest.hexdigest()


def plot_role_set(stage: Path, data_root: Path, group: str, roles, output_name: str):
    loaded = [load_role(path, "legacy", group) for _, path in roles]
    shapes = {item[0].shape for item in loaded}
    if len(shapes) != 1:
        fail(f"role field shapes differ for SR/{group}: {shapes}")
    all_sensors = np.concatenate([item[1] for item in loaded], axis=0)
    xlim = padded_limits(all_sensors[:, [0, 2, 4]])
    ylim = padded_limits(all_sensors[:, [1, 3, 5]])
    ny, nx = next(iter(shapes)); extent = (0, (nx - 1) / L0, 0, (ny - 1) / L0)
    fig, axes = plt.subplots(2, len(roles), figsize=(6 * len(roles), 8), constrained_layout=True, squeeze=False)
    images = []
    for col, ((label, _), (omega, sensors, units, _, _, _)) in enumerate(zip(roles, loaded)):
        images.append(axes[0, col].imshow(omega, origin="lower", extent=extent, cmap="RdBu_r", vmin=-VORTICITY_LIMIT, vmax=VORTICITY_LIMIT, aspect="equal"))
        axes[0, col].set_title(f"LEGACY / {group} — {label}")
        axes[0, col].set_xlabel("x/L0"); axes[0, col].set_ylabel("y/L0")
        axes[0, col].set_xlim(extent[:2]); axes[0, col].set_ylim(extent[2:])
        for (u, v, sensor_label), color in zip(SENSOR_PAIRS, SENSOR_COLORS):
            axes[1, col].plot(sensors[:, u], sensors[:, v], color=color, lw=1.1, alpha=0.85, label=sensor_label)
        axes[1, col].axhline(0, color="0.35", lw=0.7); axes[1, col].axvline(0, color="0.35", lw=0.7)
        axes[1, col].grid(True, alpha=0.25); axes[1, col].set_xlim(xlim); axes[1, col].set_ylim(ylim)
        axes[1, col].set_xlabel(f"sensor u ({units})"); axes[1, col].set_ylabel(f"sensor v ({units})")
        axes[1, col].set_box_aspect(0.8); axes[1, col].legend(loc="best", frameon=False)
    cbar = fig.colorbar(images[0], ax=axes[0, :], orientation="vertical", shrink=0.92, pad=0.015)
    cbar.set_label(r"$\omega_z$ (lattice$^{-1}$)")
    output = stage / "legacy" / f"{output_name}.png"; output.parent.mkdir(parents=True, exist_ok=True)
    fig.savefig(output, dpi=160); plt.close(fig)
    role_entries = []
    for (label, path), item in zip(roles, loaded):
        sources = [path / "metadata.json", path / "timeseries.csv", item[3]]
        role_entries.append({"role": label, "source": str(path), "field_file": str(item[3]), "field_slot": 0,
                             "selector": item[4], "selector_value": float(item[5]),
                             "source_files": [{"path": str(source), "sha256": sha256_file(source), "bytes": source.stat().st_size} for source in sources]})
    return {"pipeline": "legacy-sr", "group": group, "plot": f"legacy/{output_name}.png", "roles": role_entries,
            "vorticity_limit": VORTICITY_LIMIT, "sensor_limits": {"u": xlim, "v": ylim}}


def sr_plot_specs(data_root: Path):
    legacy = data_root / "legacy"; specs = []
    for case in SR_TRAINING:
        base = legacy / case
        specs.append((case, (("Target", base / "target"), ("PPO", base / "controlled"), ("SR", base / "sr"), ("Zero (physical, no control)", base / "zero")), case))
    for case in SR_GENERALIZATION:
        base = legacy / case
        specs.append((case, (("Target", base / "target"), ("SR", base / "sr"), ("Zero (physical, no control)", base / "zero")), case))
    for variant, case in SR_VARIANT_CASE.items():
        base = legacy / case
        specs.append((case, (("Target", base / "target"), ("Parent SR", base / "sr"), (f"Variant SR ({variant})", base / f"sr_{variant}"), ("Zero (physical, no control)", base / "zero")), f"variants/{case}_{variant}"))
    return specs


def render_sr_package(stage: Path, data_root: Path, allow_partial=False):
    entries = []
    for group, roles, output_name in sr_plot_specs(data_root):
        missing = [str(path) for _, path in roles if not path.is_dir()]
        if missing:
            if allow_partial:
                continue
            fail(f"missing SR plot roles for {group}: {missing}")
        entries.append(plot_role_set(stage, data_root, group, roles, output_name))
    if not allow_partial and len(entries) != 23:
        fail(f"expected 23 SR plots, got {len(entries)}")
    return entries


def plot_group(stage: Path, data_root: Path, pipeline: str, group: str):
    roles = roles_for(data_root, pipeline, group)
    loaded = [load_role(path, pipeline, group) for _, path in roles]
    shapes = {x[0].shape for x in loaded}
    if len(shapes) != 1:
        fail(f"role field shapes differ for {pipeline}/{group}: {shapes}")
    limit = VORTICITY_LIMIT
    all_sensors = np.concatenate([x[1] for x in loaded], axis=0)
    xlim = padded_limits(all_sensors[:, [0, 2, 4]])
    ylim = padded_limits(all_sensors[:, [1, 3, 5]])
    ny, nx = next(iter(shapes))
    extent = (0, (nx - 1) / L0, 0, (ny - 1) / L0)
    fig, axes = plt.subplots(2, 3, figsize=(18, 8), constrained_layout=True)
    images = []
    for col, ((role, _), (omega, sensors, units, _, _, _)) in enumerate(zip(roles, loaded)):
        images.append(axes[0, col].imshow(omega, origin="lower", extent=extent, cmap="RdBu_r", vmin=-limit, vmax=limit, aspect="equal"))
        axes[0, col].set_title(f"{pipeline.upper()} / {group} — {role}")
        axes[0, col].set_xlabel("x/L0")
        axes[0, col].set_ylabel("y/L0")
        axes[0, col].set_xlim(extent[:2]); axes[0, col].set_ylim(extent[2:])
        for (u, v, label), color in zip(SENSOR_PAIRS, SENSOR_COLORS):
            axes[1, col].plot(sensors[:, u], sensors[:, v], color=color, lw=1.1, alpha=0.85, label=label)
        axes[1, col].axhline(0, color="0.35", lw=0.7); axes[1, col].axvline(0, color="0.35", lw=0.7)
        axes[1, col].grid(True, alpha=0.25)
        axes[1, col].set_xlim(xlim); axes[1, col].set_ylim(ylim)
        axes[1, col].set_xlabel(f"sensor u ({units})"); axes[1, col].set_ylabel(f"sensor v ({units})")
        axes[1, col].set_box_aspect(0.8)
        axes[1, col].legend(loc="best", frameon=False)
    cbar = fig.colorbar(images[0], ax=axes[0, :], orientation="vertical", shrink=0.92, pad=0.015)
    cbar.set_label(r"$\omega_z$ (lattice$^{-1}$)")
    output = stage / pipeline / f"{group}.png"
    output.parent.mkdir(parents=True, exist_ok=True)
    fig.savefig(output, dpi=160)
    plt.close(fig)
    return {
        "pipeline": pipeline, "group": group, "plot": f"{pipeline}/{group}.png",
        "roles": [{"role": role, "source": str(path), "field_file": str(item[3]),
                   "field_slot": 0, "selector": item[4], "selector_value": float(item[5])}
                  for (role, path), item in zip(roles, loaded)],
        "vorticity_limit": limit, "sensor_limits": {"u": xlim, "v": ylim},
    }


def main() -> None:
    here = Path(__file__).resolve().parent
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--data-root", type=Path, default=here / "data" / "reproduction")
    parser.add_argument("--output-root", type=Path)
    parser.add_argument("--campaign", choices=("ppo", "sr"), default="ppo")
    parser.add_argument("--allow-partial", action="store_true")
    args = parser.parse_args()
    data_root = args.data_root.resolve()
    output_root = (args.output_root or here / "data" / ("reproduction_plots_sr" if args.campaign == "sr" else "reproduction_plots")).absolute()
    if args.campaign == "ppo": require_exact_groups(data_root)
    if output_root.is_symlink():
        fail(f"refusing symlink output root: {output_root}")
    output_root.parent.mkdir(parents=True, exist_ok=True)
    stage = Path(tempfile.mkdtemp(prefix=f".{output_root.name}.staging-", dir=output_root.parent))
    backup = output_root.with_name(f".{output_root.name}.old-{os.getpid()}")
    try:
        if args.campaign == "sr":
            entries = render_sr_package(stage, data_root, allow_partial=args.allow_partial)
            manifest = {"schema": "drl-pinball-sr-reproduction-plots-v1", "plot_count": len(entries),
                        "expected_plot_count": 23, "allow_partial": args.allow_partial,
                        "no_ppo_cases": list(SR_GENERALIZATION),
                        "no_ppo_contract": "Generalization conditions intentionally contain Target/SR/Zero only; PPO is absent by acquisition design.",
                        "variant_contract": "Each diagnostic binds the named variant to its same-case parent SR.",
                        "vorticity_limit": VORTICITY_LIMIT,
                        "vorticity_contract": "fixed symmetric [-0.001, +0.001] for every vorticity panel",
                        "plots": entries}
        else:
            entries = [plot_group(stage, data_root, "v5", group) for group in sorted(V5_GROUPS)]
            entries += [plot_group(stage, data_root, "legacy", group) for group in LEGACY_GROUPS]
            if len(entries) != 30:
                fail(f"expected 30 plots, got {len(entries)}")
            manifest = {"schema": "drl-pinball-reproduction-plots-v1", "plot_count": 30,
                        "vorticity_limit": VORTICITY_LIMIT,
                        "vorticity_contract": "fixed symmetric [-0.001, +0.001] for every vorticity panel", "plots": entries}
        (stage / "manifest.json").write_text(json.dumps(manifest, indent=2) + "\n")
        if output_root.exists():
            os.replace(output_root, backup)
        os.replace(stage, output_root)
        if backup.exists():
            shutil.rmtree(backup)
    except Exception:
        shutil.rmtree(stage, ignore_errors=True)
        if backup.exists() and not output_root.exists():
            os.replace(backup, output_root)
        raise
    print(f"wrote {len(entries)} plots and manifest to {output_root}")


if __name__ == "__main__":
    main()