"""spp スイープの解析: 計算時間・GPU 率・ノイズ・指標の安定性を spp に対して描く.

sweep_spp_noise.py の summary.csv（と併走した nvidia-smi のサンプル gpu_util.csv）を読み，

1. spp vs 計算時間（段階別の積み上げ）と GPU 率（mi.render / 全体）
2. spp vs ノイズ CoV（1/sqrt(spp) の理論線付き）と生の輝度依存
3. spp vs 指標（構造 CoV・輪帯 min/max・画像全体 min/max の補正あり/なし）．seed 2 つの差も

を出力する（output/spp_noise_sweep/spp_noise.png・spp_summary_table.csv）．

実行方法:
    uv run python scripts/build_spp_noise_comparison.py
"""

from __future__ import annotations

import argparse
import csv
import sys
from pathlib import Path

import matplotlib
import numpy as np

matplotlib.use("Agg")
import matplotlib.pyplot as plt  # noqa: E402

sys.path.insert(0, str(Path(__file__).resolve().parent))

from sweep_spp_noise import CONDITIONS, OUTPUT_BASE, SPPS  # noqa: E402

plt.rcParams["font.family"] = ["Yu Gothic", "MS Gothic", "DejaVu Sans"]

COND_LABEL = {
    "sphere_d340_th000": "球 φ340 軸上",
    "sphere_d340_th090": "球 φ340 θ=90°",
    "cube_L234_th000": "立方体 L=234 軸上",
    "cube_L234_th090": "立方体 L=234 θ=90°",
    "sphere_d500_th090": "球 φ500 θ=90°",
    "cube_L340_th090": "立方体 L=340 θ=90°",
}
STAGES = [("t_load_yaml", "YAML 読込"), ("t_build_scene", "PLY 書き出し"),
          ("t_load_dict", "シーン読込"), ("t_render", "mi.render（GPU）"),
          ("t_save", "保存"), ("t_evaluate", "評価")]


def load_rows() -> list[dict]:
    """summary.csv を数値化して読む."""
    rows = []
    with (OUTPUT_BASE / "summary.csv").open(encoding="utf-8") as fh:
        for r in csv.DictReader(fh):
            row = {"name": r["name"], "condition": r["condition"],
                   "render_variant": r["render_variant"]}
            for k, v in r.items():
                if k not in row:
                    row[k] = float(v) if v not in ("", None) else float("nan")
            rows.append(row)
    return rows


def by_condition_spp(rows: list[dict], key: str) -> dict[str, dict[int, tuple[float, float]]]:
    """条件 → spp → (seed 平均, seed 差の半分) を返す."""
    out: dict[str, dict[int, tuple[float, float]]] = {}
    for cond, _ in CONDITIONS:
        out[cond] = {}
        for spp in SPPS:
            vals = [r[key] for r in rows if r["condition"] == cond and int(r["spp"]) == spp]
            if vals:
                spread = float(np.ptp(vals)) / 2.0 if len(vals) > 1 else 0.0
                out[cond][spp] = (float(np.mean(vals)), spread)
    return out


def gpu_util_summary() -> str:
    """nvidia-smi サンプルの要約（あれば）."""
    path = OUTPUT_BASE / "gpu_util.csv"
    if not path.exists():
        return ""
    utils = []
    with path.open(encoding="utf-8") as fh:
        for r in csv.DictReader(fh):
            try:
                utils.append(float(r["gpu_util"]))
            except (ValueError, KeyError):
                continue
    if not utils:
        return ""
    arr = np.array(utils)
    return (f"nvidia-smi サンプル {len(arr)} 点: 使用率 平均 {arr.mean():.0f}%・"
            f"稼働（>0%）の割合 {np.mean(arr > 0) * 100:.0f}%")


def draw(rows: list[dict], output: Path) -> None:
    """3 行 × 3 列の図を描く."""
    fig, axes = plt.subplots(3, 3, figsize=(20, 15))
    colors = plt.cm.tab10(np.linspace(0, 1, 10))[: len(CONDITIONS)]

    # 1-1 段階別の時間（代表条件 球 φ340 θ=90 の seed 平均を積み上げ）
    ax = axes[0, 0]
    rep = "sphere_d340_th090"
    bottom = np.zeros(len(SPPS))
    for key, label in STAGES:
        series = by_condition_spp(rows, key)[rep]
        vals = np.array([series.get(s, (np.nan, 0))[0] for s in SPPS])
        ax.bar(range(len(SPPS)), vals, bottom=bottom, label=label)
        bottom += np.nan_to_num(vals)
    ax.set_xticks(range(len(SPPS)))
    ax.set_xticklabels([str(s) for s in SPPS])
    ax.set_xlabel("spp")
    ax.set_ylabel("時間 [s]")
    ax.set_title(f"段階別の時間（{COND_LABEL[rep]}）", fontsize=11)
    ax.legend(fontsize=8)
    ax.grid(alpha=0.3, axis="y")

    # 1-2 全体時間 vs spp（全条件）と render 時間
    ax = axes[0, 1]
    total = by_condition_spp(rows, "t_total")
    render = by_condition_spp(rows, "t_render")
    for (cond, _), c in zip(CONDITIONS, colors, strict=True):
        xs = [s for s in SPPS if s in total[cond]]
        ax.plot(xs, [total[cond][s][0] for s in xs], marker="o", color=c, label=COND_LABEL[cond])
        ax.plot(xs, [render[cond][s][0] for s in xs], marker="x", ls=":", color=c, alpha=0.7)
    ax.set_xscale("log", base=2)
    ax.set_yscale("log")
    ax.set_xlabel("spp")
    ax.set_ylabel("時間 [s]（実線=全体，点線=mi.render）")
    ax.set_title("計算時間 vs spp", fontsize=11)
    ax.legend(fontsize=7)
    ax.grid(alpha=0.3, which="both")

    # 1-3 GPU 率
    ax = axes[0, 2]
    frac = by_condition_spp(rows, "gpu_fraction")
    for (cond, _), c in zip(CONDITIONS, colors, strict=True):
        xs = [s for s in SPPS if s in frac[cond]]
        ax.plot(xs, [frac[cond][s][0] * 100 for s in xs], marker="o", color=c,
                label=COND_LABEL[cond])
    ax.set_xscale("log", base=2)
    ax.set_xlabel("spp")
    ax.set_ylabel("GPU 率 [%] = mi.render / 全体")
    ax.set_title("処理時間に占める GPU（レンダ本体）の割合", fontsize=11)
    ax.set_ylim(0, 100)
    ax.legend(fontsize=7)
    ax.grid(alpha=0.3, which="both")

    # 2-1 ノイズ CoV vs spp + 理論線
    ax = axes[1, 0]
    noise = by_condition_spp(rows, "noise_cov")
    for (cond, _), c in zip(CONDITIONS, colors, strict=True):
        xs = [s for s in SPPS if s in noise[cond]]
        ys = [noise[cond][s][0] for s in xs]
        ax.plot(xs, ys, marker="o", color=c, label=COND_LABEL[cond])
        if 1024 in noise[cond]:
            ref = noise[cond][1024][0]
            ax.plot(xs, [ref * np.sqrt(1024 / s) for s in xs], ls=":", color=c, alpha=0.6)
    ax.set_xscale("log", base=2)
    ax.set_yscale("log")
    ax.set_xlabel("spp")
    ax.set_ylabel("残差ノイズ CoV（点線 = 1/√spp の理論線，1024 基準）")
    ax.set_title("ノイズ vs spp", fontsize=11)
    ax.legend(fontsize=7)
    ax.grid(alpha=0.3, which="both")

    # 2-2 ノイズ CoV × sqrt(spp) の一定性（生の輝度に対して）
    ax = axes[1, 1]
    mean = by_condition_spp(rows, "raw_mean_roi")
    for (cond, _), c in zip(CONDITIONS, colors, strict=True):
        xs = [s for s in SPPS if s in noise[cond]]
        ax.plot(xs, [noise[cond][s][0] * np.sqrt(s) for s in xs], marker="o", color=c,
                label=f"{COND_LABEL[cond]}（生 mean {mean[cond].get(1024, (np.nan, 0))[0]:.2f}）")
    ax.set_xscale("log", base=2)
    ax.set_xlabel("spp")
    ax.set_ylabel("ノイズ CoV × √spp（一定なら 1/√spp 則）")
    ax.set_title("1/√spp 則の検証（生の輝度が低いほど大きい）", fontsize=11)
    ax.legend(fontsize=7)
    ax.grid(alpha=0.3, which="both")

    # 2-3 ノイズ場の極値（補正量）
    ax = axes[1, 2]
    nmin = by_condition_spp(rows, "fullimage_noise_min")
    nmax = by_condition_spp(rows, "fullimage_noise_max")
    for (cond, _), c in zip(CONDITIONS, colors, strict=True):
        xs = [s for s in SPPS if s in nmin[cond]]
        m = [mean[cond][s][0] for s in xs]
        spans = [(nmax[cond][s][0] - nmin[cond][s][0]) / mm for s, mm in zip(xs, m, strict=True)]
        ax.plot(xs, spans, marker="o", color=c, label=COND_LABEL[cond])
    ax.set_xscale("log", base=2)
    ax.set_yscale("log")
    ax.set_xlabel("spp")
    ax.set_ylabel("ノイズ場の (max − min) / 平均")
    ax.set_title("画像全体 min/max のノイズ補正量", fontsize=11)
    ax.legend(fontsize=7)
    ax.grid(alpha=0.3, which="both")

    # 3-x 指標の安定性
    for ax, key, title in ((axes[2, 0], "structure_cov", "構造 CoV"),
                           (axes[2, 1], "ring_min_max", "輪帯平均 min/max"),
                           (axes[2, 2], "fullimage_min_max", "画像全体 min/max（ノイズ補正）")):
        series = by_condition_spp(rows, key)
        for (cond, _), c in zip(CONDITIONS, colors, strict=True):
            xs = [s for s in SPPS if s in series[cond]]
            ax.errorbar(xs, [series[cond][s][0] for s in xs],
                        yerr=[series[cond][s][1] for s in xs], marker="o", capsize=3,
                        color=c, label=COND_LABEL[cond])
            if key == "fullimage_min_max":
                raw = by_condition_spp(rows, "fullimage_raw_min_max")
                ax.plot(xs, [raw[cond][s][0] for s in xs], ls=":", color=c, alpha=0.6)
        ax.set_xscale("log", base=2)
        ax.set_xlabel("spp")
        ax.set_ylabel(title + ("（点線 = 補正なし）" if key == "fullimage_min_max" else ""))
        ax.set_title(f"{title} vs spp（誤差棒 = seed 2 つの差の半分）", fontsize=11)
        ax.legend(fontsize=7)
        ax.grid(alpha=0.3, which="both")

    note = gpu_util_summary()
    fig.suptitle("spp スイープ: 計算時間・GPU 率・ノイズ・指標の安定性"
                 "（Newton GPU 0・cuda_ad_rgb・512²）" + (f"\n{note}" if note else ""),
                 fontsize=13)
    fig.tight_layout()
    output.parent.mkdir(parents=True, exist_ok=True)
    fig.savefig(output, dpi=100)
    plt.close(fig)
    print(f"出力: {output}")


def print_table(rows: list[dict]) -> None:
    """spp ごとの要約表（条件平均）を出す."""
    print(f"\n{'spp':>6} {'全体[s]':>8} {'render[s]':>10} {'固定費[s]':>9} {'GPU率':>6} "
          f"{'noise CoV':>10} {'noise×√spp':>11} {'構造CoV seed差':>14} {'min seed差':>11}")
    out_rows = []
    for spp in SPPS:
        sel = [r for r in rows if int(r["spp"]) == spp]
        if not sel:
            continue
        total = np.mean([r["t_total"] for r in sel])
        render = np.mean([r["t_render"] for r in sel])
        noise = np.mean([r["noise_cov"] for r in sel])
        struct_diff, min_diff = [], []
        for cond, _ in CONDITIONS:
            c = [r for r in sel if r["condition"] == cond]
            if len(c) >= 2:
                struct_diff.append(abs(c[0]["structure_cov"] - c[1]["structure_cov"]))
                min_diff.append(abs(c[0]["fullimage_min_max"] - c[1]["fullimage_min_max"]))
        row = {"spp": spp, "t_total": total, "t_render": render, "t_fixed": total - render,
               "gpu_fraction": render / total, "noise_cov": noise,
               "noise_x_sqrt": noise * np.sqrt(spp),
               "structure_seed_diff": float(np.mean(struct_diff)) if struct_diff else np.nan,
               "min_seed_diff": float(np.mean(min_diff)) if min_diff else np.nan}
        out_rows.append(row)
        print(f"{spp:>6} {total:8.1f} {render:10.1f} {total - render:9.1f} "
              f"{render / total * 100:5.0f}% {noise:10.4f} {noise * np.sqrt(spp):11.3f} "
              f"{row['structure_seed_diff']:14.4f} {row['min_seed_diff']:11.4f}")
    with (OUTPUT_BASE / "spp_summary_table.csv").open("w", newline="", encoding="utf-8") as fh:
        writer = csv.DictWriter(fh, fieldnames=list(out_rows[0].keys()))
        writer.writeheader()
        writer.writerows(out_rows)
    note = gpu_util_summary()
    if note:
        print(note)


def main() -> int:
    """spp スイープの図と表を出す."""
    parser = argparse.ArgumentParser(description="spp スイープの解析")
    parser.parse_args()
    rows = load_rows()
    if not rows:
        print("エラー: summary.csv が空です")
        return 1
    draw(rows, OUTPUT_BASE / "spp_noise.png")
    print_table(rows)
    return 0


if __name__ == "__main__":
    sys.exit(main())