"""半径スイープの結果を一覧図にする（輝度マップ一覧 + 指標の半径依存カーブ）.

数値だけで形を判断しないための図．TEST_14 では総合スコアと残差 RMS が別の最良条件を
指す食い違いが起き，輝度マップそのものを見て初めて解釈できた．本スクリプトは
summary.csv の指標カーブと，各半径の輝度マップを必ず並べて出す．

輝度マップは平滑化（median 5px → gaussian sigma=4）した上で **20 分割の輪帯平均の最大値を
1.0** として正規化する．絶対輝度は半径で 1 桁変わるため何かで割らないと形が比較できないが，
画素単位の最大値は MC ノイズの外れ値（firefly）に張り付いて使えない（本スイープでは
max/min が 2.9〜15.5 まで暴れる）．輪帯平均ならリング内の画素数だけノイズが 1/sqrt(N) に
落ちるので，最大値そのものが安定する．分割は評価パイプラインと同一
（`evaluation/radial.py:calc_radial_profile` の既定 20 ビン）．

実行方法:
    uv run python scripts/build_radius_sweep_visual.py --cavity similar_radial
"""

from __future__ import annotations

import argparse
import csv
import sys
from pathlib import Path

import matplotlib
import numpy as np
from scipy.ndimage import gaussian_filter, median_filter

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

sys.path.insert(0, str(Path(__file__).resolve().parent))
sys.path.insert(0, str(Path(__file__).resolve().parents[1] / "src"))

from sweep_sphere_radius import (  # noqa: E402
    CAVITIES,
    OUTPUT_BASE,
    VARIANTS,
    job_name,
    run_dir_name,
)

from sm_tias_lightsim.evaluation.radial import calc_radial_profile  # noqa: E402
from sm_tias_lightsim.evaluation.roi_presets import SIM_ROI, extract_roi_xywh  # noqa: E402

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

# 指標カーブに出す列と表示名（小さいほど均一 / 1 に近いほど均一を注記する）
METRIC_PANELS = [
    ("gradient_magnitude", "勾配 [%]（小さいほど均一）", False),
    ("center_periphery_ratio", "中心/周辺比（1 が理想）", True),
    ("radial_min_max_ratio", "動径 min/max（1 が理想）", True),
    ("angular_ratio", "angular_ratio（0=同心円）", False),
    ("raw_mean_roi", "生の ROI 平均輝度（絶対値・対数）", False),
]

# CoV 分解に使う平滑化強度．sigma=4 では MC ノイズが残るため構造の推定には 12 を使う
# （sigma を 4→20 と変えても構造値は 4% しか動かないことを実測して確認済み）
STRUCTURE_SIGMA = 12

# 輝度マップの正規化に使う輪帯の分割数（評価パイプラインの既定と同一）
NORMALIZE_BINS = 20


def read_summary(path: Path) -> list[dict[str, float]]:
    """summary.csv を読み，半径昇順の辞書リストで返す."""
    with path.open(encoding="utf-8") as fh:
        rows = [{k: (float(v) if v else float("nan")) for k, v in row.items()}
                for row in csv.DictReader(fh)]
    return sorted(rows, key=lambda r: r["radius_mm"])


def normalized_map(luminance_npy: Path) -> np.ndarray:
    """輪帯平均の最大値を 1.0 に正規化した平滑化輝度マップを返す.

    正規化の基準は**生の ROI から取った 20 分割の輪帯平均の最大値**とする．
    画素単位の最大値は MC ノイズの外れ値に張り付いて条件間で比較できないが，
    輪帯平均なら方位方向の平均でノイズが 1/sqrt(N) に落ちて安定する．
    表示するマップ自体は形を見るために平滑化する（median 5px → gaussian sigma=4）．
    """
    arr = np.load(luminance_npy).astype(np.float64)
    if arr.ndim == 3:
        arr = 0.2126 * arr[..., 0] + 0.7152 * arr[..., 1] + 0.0722 * arr[..., 2]
    roi = extract_roi_xywh(arr, SIM_ROI)
    profile = calc_radial_profile(roi, n_bins=NORMALIZE_BINS)
    peak = float(np.nanmax(profile[:, 1]))
    return gaussian_filter(median_filter(roi, size=5), sigma=4) / peak


def decompose_cov(luminance_npy: Path, seed: int = 0) -> tuple[float, float, float]:
    """CoV を「構造」と「MC ノイズ」に分解して (生, 構造, ノイズ) を返す.

    生の CoV は絶対輝度が半径で 1 桁変わる本スイープでは MC ノイズに支配され，
    そのまま半径依存を読むと「大きい球ほど不均一」という誤った結論になる．

    構造は STRUCTURE_SIGMA で平滑化した像の CoV から，ノイズの漏れ込み分を
    二乗で差し引いて求める．漏れ込み量は残差（生 - 平滑）を空間シャッフルして
    同じ平滑化を掛けた std として実測する（構造を壊しノイズの統計だけ残す）．
    """
    arr = np.load(luminance_npy).astype(np.float64)
    if arr.ndim == 3:
        arr = 0.2126 * arr[..., 0] + 0.7152 * arr[..., 1] + 0.0722 * arr[..., 2]
    roi = extract_roi_xywh(arr, SIM_ROI)
    mean = roi.mean()
    residual = roi - gaussian_filter(median_filter(roi, size=5), sigma=4)
    smoothed = gaussian_filter(median_filter(roi, size=5), sigma=STRUCTURE_SIGMA)
    flat = residual.ravel().copy()
    np.random.default_rng(seed).shuffle(flat)
    leak = gaussian_filter(
        median_filter(flat.reshape(residual.shape), size=5), sigma=STRUCTURE_SIGMA
    ).std() / mean
    structure = float(np.sqrt(max((smoothed.std() / smoothed.mean()) ** 2 - leak**2, 0.0)))
    return float(roi.std() / mean), structure, float(residual.std() / mean)


def draw_cov_panel(ax, rows: list[dict[str, float]], variant: str, cavity: str) -> None:
    """CoV を生・構造・ノイズに分けて描く（生だけ見ると誤読するため）."""
    radii, raws, structures, noises = [], [], [], []
    for row in rows:
        path = (OUTPUT_BASE / run_dir_name(variant, cavity)
                / job_name(row["radius_mm"]) / "luminance.npy")
        if not path.exists():
            continue
        raw, structure, noise = decompose_cov(path)
        radii.append(row["radius_mm"])
        raws.append(raw)
        structures.append(structure)
        noises.append(noise)
    ax.plot(radii, raws, marker="o", color="gray", label="生 CoV（= metrics.csv）")
    ax.plot(radii, noises, marker="^", color="tab:orange", label="MC ノイズ成分")
    ax.plot(radii, structures, marker="s", color="tab:red", lw=2.2, label="構造成分（真の不均一）")
    ax.set_xlabel("球内半径 (mm)")
    ax.set_title("CoV の分解（構造 vs MC ノイズ）", fontsize=11)
    ax.legend(fontsize=8)
    ax.grid(alpha=0.3)


def draw_maps(fig, gs, rows: list[dict[str, float]], variant: str, cavity: str) -> None:
    """各半径の輝度マップを「輪帯平均の最大値=1.0」で正規化して並べる（共通スケール）."""
    maps: list[tuple[float, np.ndarray]] = []
    for row in rows:
        path = (OUTPUT_BASE / run_dir_name(variant, cavity)
                / job_name(row["radius_mm"]) / "luminance.npy")
        if path.exists():
            maps.append((row["radius_mm"], normalized_map(path)))
    if not maps:
        return
    stacked = np.concatenate([m.ravel() for _, m in maps])
    # 上端は 1.0（＝最大輪帯）に固定し，下端だけデータに合わせる
    vmin = float(np.percentile(stacked, 0.5))
    for i, (radius, data) in enumerate(maps[:10]):
        ax = fig.add_subplot(gs[i // 5, i % 5])
        image = ax.imshow(data, cmap="hot", vmin=vmin, vmax=1.0)
        ax.set_title(f"r={radius:g}mm", fontsize=10)
        ax.axis("off")
        if i == len(maps[:10]) - 1:
            fig.colorbar(image, ax=ax, fraction=0.046, label="最大輪帯（20 分割）=1.0")


def draw_metrics(fig, gs, rows: list[dict[str, float]]) -> None:
    """指標の半径依存カーブを描く（理想値がある指標には水平線を引く）."""
    radii = [row["radius_mm"] for row in rows]
    for i, (key, label, has_ideal) in enumerate(METRIC_PANELS, start=1):
        ax = fig.add_subplot(gs[2 + i // 3, (i % 3) * 2 : (i % 3) * 2 + 2])
        ax.plot(radii, [row[key] for row in rows], marker="o", color="tab:blue")
        if has_ideal:
            ax.axhline(1.0, color="gray", lw=0.8, ls="--")
        if key == "raw_mean_roi":
            ax.set_yscale("log")
        ax.set_xlabel("球内半径 (mm)")
        ax.set_title(label, fontsize=10)
        ax.grid(alpha=0.3)


def main() -> int:
    """半径スイープの一覧図を生成する."""
    parser = argparse.ArgumentParser(description="半径スイープの結果一覧図を生成")
    parser.add_argument("--variant", choices=VARIANTS, default="board_baffle")
    parser.add_argument("--cavity", choices=CAVITIES, default="similar_radial")
    parser.add_argument("--output", type=Path, default=None, help="出力 PNG パス")
    args = parser.parse_args()

    run_dir = OUTPUT_BASE / run_dir_name(args.variant, args.cavity)
    summary = run_dir / "summary.csv"
    if not summary.exists():
        print(f"エラー: {summary} がありません（先にスイープを実行してください）")
        return 1
    rows = read_summary(summary)
    output = args.output or run_dir / "sweep_overview.png"

    fig = plt.figure(figsize=(20, 13))
    gs = fig.add_gridspec(4, 6, height_ratios=[1.1, 1.1, 1, 1])
    draw_maps(fig, gs, rows, args.variant, args.cavity)
    draw_cov_panel(fig.add_subplot(gs[2, 0:2]), rows, args.variant, args.cavity)
    draw_metrics(fig, gs, rows)
    fig.suptitle(
        f"球体筐体 半径スイープ 結果一覧 [{args.variant} / {args.cavity}]"
        "　輝度マップは 20 分割の輪帯平均の最大値を 1.0 に正規化", fontsize=15
    )
    fig.tight_layout()
    output.parent.mkdir(parents=True, exist_ok=True)
    fig.savefig(output, dpi=100)
    print(f"出力: {output}")
    return 0


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