plot_module_compare.py

#!/usr/bin/env python3
"""
Visualize Local vs PROCESS module comparison results in output/module_compare/.

Usage:
    python plot_module_compare.py
    python plot_module_compare.py --input-dir output/module_compare --no-show
"""

from __future__ import annotations

import argparse
import re
from pathlib import Path

import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
from matplotlib.patches import Patch

ROOT = Path(__file__).resolve().parent
DEFAULT_INPUT_DIR = ROOT / "output" / "module_compare"

IMPLEMENTATION_COLORS = {
    "implemented": "#2563eb",
    "partial": "#f59e0b",
    "stub": "#94a3b8",
}
MATCH_OK = "#16a34a"
MATCH_FAIL = "#dc2626"
LOCAL_COLOR = "#2563eb"
PROCESS_COLOR = "#f97316"
REL_ERR_CLIP = 200.0


def load_data(input_dir: Path) -> tuple[pd.DataFrame, pd.DataFrame]:
    comparison_path = input_dir / "module_comparison.csv"
    inventory_path = input_dir / "module_inventory.csv"
    if not comparison_path.is_file():
        raise FileNotFoundError(f"Missing {comparison_path}; run compare_modules_process.py first")
    if not inventory_path.is_file():
        raise FileNotFoundError(f"Missing {inventory_path}; run compare_modules_process.py first")
    comparison = pd.read_csv(comparison_path)
    inventory = pd.read_csv(inventory_path)
    return comparison, inventory


def metric_label(row: pd.Series) -> str:
    variant = row["variant"]
    if variant and variant != "default":
        return f"{row['metric']}\n({variant})"
    return str(row["metric"])


def clip_rel_err(value: float, clip: float = REL_ERR_CLIP) -> tuple[float, bool]:
    if abs(value) > clip:
        return np.sign(value) * clip, True
    return value, False


def plot_implementation_overview(inventory: pd.DataFrame, output_dir: Path) -> Path:
    fig, axes = plt.subplots(1, 2, figsize=(11, 4.5))

    impl_counts = inventory["implementation"].value_counts().reindex(
        ["implemented", "partial", "stub"], fill_value=0
    )
    axes[0].bar(
        impl_counts.index,
        impl_counts.values,
        color=[IMPLEMENTATION_COLORS[k] for k in impl_counts.index],
        edgecolor="white",
    )
    axes[0].set_title("Local module implementation status")
    axes[0].set_ylabel("Count")
    axes[0].grid(axis="y", alpha=0.25)

    auto_counts = inventory["auto_compare"].value_counts()
    auto_labels = ["Compared", "Not compared"]
    auto_values = [auto_counts.get(True, 0), auto_counts.get(False, 0)]
    axes[1].bar(
        auto_labels,
        auto_values,
        color=["#16a34a", "#cbd5e1"],
        edgecolor="white",
    )
    axes[1].set_title("Automated Local vs PROCESS compare")
    axes[1].set_ylabel("Count")
    axes[1].grid(axis="y", alpha=0.25)

    fig.suptitle("Module inventory overview", fontsize=12, y=1.02)
    fig.tight_layout()
    path = output_dir / "module_status_overview.png"
    fig.savefig(path, dpi=160, bbox_inches="tight")
    plt.close(fig)
    return path


def plot_match_summary(comparison: pd.DataFrame, output_dir: Path) -> Path:
    summary_rows = []
    for module_id, group in comparison.groupby("module_id", sort=False):
        summary_rows.append(
            {
                "module_id": module_id,
                "matched": int(group["match"].sum()),
                "total": len(group),
                "match_rate_pct": 100.0 * group["match"].mean(),
            }
        )
    summary = pd.DataFrame(summary_rows).sort_values("match_rate_pct")

    fig, ax = plt.subplots(figsize=(9, max(4, 0.45 * len(summary) + 1.5)))
    y = np.arange(len(summary))
    ax.barh(y, summary["match_rate_pct"], color=MATCH_OK, alpha=0.85)
    for idx, row in summary.iterrows():
        pos = summary.index.get_loc(idx)
        ax.text(
            min(row["match_rate_pct"] + 1.5, 102),
            pos,
            f"{row['matched']}/{row['total']}",
            va="center",
            fontsize=9,
        )
    ax.set_yticks(y)
    ax.set_yticklabels(summary["module_id"])
    ax.set_xlim(0, 105)
    ax.set_xlabel("Metrics within 1% relative error (%)")
    ax.set_title("Per-module match rate")
    ax.grid(axis="x", alpha=0.25)
    fig.tight_layout()
    path = output_dir / "module_match_summary.png"
    fig.savefig(path, dpi=160, bbox_inches="tight")
    plt.close(fig)
    return path


def plot_relative_errors(comparison: pd.DataFrame, output_dir: Path) -> Path:
    plot_df = comparison.copy()
    plot_df["label"] = plot_df.apply(metric_label, axis=1)
    plot_df["module_metric"] = plot_df["module_id"] + " / " + plot_df["label"].str.replace("\n", " ")

    clip_vals = []
    clipped = []
    colors = []
    for row in plot_df.itertuples():
        val, is_clip = clip_rel_err(row.rel_err_pct)
        clip_vals.append(val)
        clipped.append(is_clip)
        colors.append(MATCH_OK if row.match else MATCH_FAIL)

    fig, ax = plt.subplots(figsize=(max(12, len(plot_df) * 0.55), 5.5))
    bars = ax.bar(plot_df["module_metric"], clip_vals, color=colors, alpha=0.9)
    ax.axhline(0.0, color="black", linewidth=0.8)
    ax.axhline(1.0, color="#64748b", linestyle="--", linewidth=0.8)
    ax.axhline(-1.0, color="#64748b", linestyle="--", linewidth=0.8)
    ax.set_ylabel("Relative error (%)")
    ax.set_title(f"Relative error vs PROCESS (clipped to ±{REL_ERR_CLIP:g}% for display)")
    ax.tick_params(axis="x", rotation=35, labelsize=8)
    ax.grid(axis="y", alpha=0.25)
    ax.legend(
        handles=[
            Patch(facecolor=MATCH_OK, label="Within 1%"),
            Patch(facecolor=MATCH_FAIL, label="Outside 1%"),
        ],
        loc="upper right",
    )

    for bar, row, is_clip in zip(bars, plot_df.itertuples(), clipped):
        if is_clip:
            ax.annotate(
                f"{row.rel_err_pct:+.2g}%",
                (bar.get_x() + bar.get_width() / 2, bar.get_height()),
                ha="center",
                va="bottom" if row.rel_err_pct >= 0 else "top",
                fontsize=7,
                color="#7f1d1d",
            )

    fig.tight_layout()
    path = output_dir / "module_relative_errors.png"
    fig.savefig(path, dpi=160, bbox_inches="tight")
    plt.close(fig)
    return path


def plot_local_vs_process(comparison: pd.DataFrame, output_dir: Path) -> Path:
    modules = list(comparison["module_id"].drop_duplicates())
    ncols = 2
    nrows = int(np.ceil(len(modules) / ncols))
    fig, axes = plt.subplots(nrows, ncols, figsize=(6.5 * ncols, 4.2 * nrows), squeeze=False)

    for idx, module_id in enumerate(modules):
        ax = axes[idx // ncols][idx % ncols]
        sub = comparison[comparison["module_id"] == module_id].copy()
        labels = [metric_label(row) for _, row in sub.iterrows()]
        x = np.arange(len(labels))
        width = 0.36

        local_vals = sub["local"].to_numpy()
        process_vals = sub["process"].to_numpy()
        pos_local = np.abs(local_vals[np.abs(local_vals) > 0])
        pos_process = np.abs(process_vals[np.abs(process_vals) > 0])
        dynamic_range = 1.0
        if len(pos_local) > 0 and len(pos_process) > 0:
            dynamic_range = max(
                np.max(pos_local) / np.min(pos_local),
                np.max(pos_process) / np.min(pos_process),
            )
        use_log = dynamic_range > 1e4 or np.any(local_vals < 0) or np.any(process_vals < 0)

        if use_log:
            local_plot = np.clip(np.abs(local_vals), np.finfo(float).tiny, None)
            process_plot = np.clip(np.abs(process_vals), np.finfo(float).tiny, None)
            ax.set_yscale("log")
        else:
            local_plot = local_vals
            process_plot = process_vals

        ax.bar(x - width / 2, local_plot, width, label="Local", color=LOCAL_COLOR, alpha=0.92)
        ax.bar(x + width / 2, process_plot, width, label="PROCESS", color=PROCESS_COLOR, alpha=0.92)
        ax.set_xticks(x)
        ax.set_xticklabels(labels, rotation=25, ha="right", fontsize=8)
        ax.set_title(module_id, fontsize=10, fontweight="bold")
        ax.grid(axis="y", alpha=0.25)
        if idx == 0:
            ax.legend(fontsize=8)

    for idx in range(len(modules), nrows * ncols):
        axes[idx // ncols][idx % ncols].axis("off")

    fig.suptitle("Local vs PROCESS by module (baseline case)", y=1.01, fontsize=12)
    fig.tight_layout()
    path = output_dir / "module_local_vs_process.png"
    fig.savefig(path, dpi=160, bbox_inches="tight")
    plt.close(fig)
    return path


def plot_fusion_reactivity(comparison: pd.DataFrame, output_dir: Path) -> Path | None:
    sub = comparison[
        (comparison["module_id"] == "fusion_reactivity") & comparison["metric"].str.startswith("sigmv_dt_t")
    ].copy()
    if sub.empty:
        return None

    temps = []
    for metric in sub["metric"]:
        match = re.search(r"t([\d.]+)kev", metric)
        if match:
            temps.append(float(match.group(1)))
        else:
            temps.append(np.nan)
    sub["temp_kev"] = temps
    sub = sub.sort_values("temp_kev")

    fig, ax = plt.subplots(figsize=(7, 4.5))
    ax.plot(sub["temp_kev"], sub["local"], "o-", color=LOCAL_COLOR, label="Local")
    ax.plot(sub["temp_kev"], sub["process"], "s--", color=PROCESS_COLOR, label="PROCESS")
    ax.set_xlabel("T_i (keV)")
    ax.set_ylabel("<σv>_DT (m³/s)")
    ax.set_title("D-T fusion reactivity\n(fus_rate vs bosch_hale_reactivity)")
    ax.set_yscale("log")
    ax.grid(alpha=0.25)
    ax.legend()
    fig.tight_layout()
    path = output_dir / "module_fusion_reactivity_curve.png"
    fig.savefig(path, dpi=160, bbox_inches="tight")
    plt.close(fig)
    return path


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description="Plot module comparison CSV results.")
    parser.add_argument(
        "--input-dir",
        type=Path,
        default=DEFAULT_INPUT_DIR,
        help=f"Directory containing module_comparison.csv (default: {DEFAULT_INPUT_DIR})",
    )
    parser.add_argument("--show", action="store_true", help="Open the overview figure after saving")
    return parser.parse_args()


def main() -> int:
    args = parse_args()
    input_dir = args.input_dir
    comparison, inventory = load_data(input_dir)

    paths = [
        plot_implementation_overview(inventory, input_dir),
        plot_match_summary(comparison, input_dir),
        plot_relative_errors(comparison, input_dir),
        plot_local_vs_process(comparison, input_dir),
    ]
    reactivity_path = plot_fusion_reactivity(comparison, input_dir)
    if reactivity_path is not None:
        paths.append(reactivity_path)

    print(f"Read: {input_dir / 'module_comparison.csv'}")
    print(f"Read: {input_dir / 'module_inventory.csv'}")
    for path in paths:
        print(f"Saved: {path}")

    if args.show:
        fig, ax = plt.subplots()
        img = plt.imread(paths[0])
        ax.imshow(img)
        ax.axis("off")
        plt.show()

    return 0


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