"""半径スイープの結果を同心円（輪帯）プロファイルで比較する.

CoV は MC ノイズに支配されて半径間の比較ができない（本スイープでは絶対輝度が
1 桁変わるため）．一方**同心円リングの方位平均はリング内の画素数だけノイズを
1/sqrt(N) に落とす**ので，ノイズ水準の違う条件どうしでも「中心から周辺への
落ち方（ドームの形）」を公平に比較できる．

輪帯の定義は評価パイプラインと同一（`evaluation/radial.py` の正規化楕円距離ビン）．
輝度は条件ごとに**最大の輪帯を 1** として正規化するので，絶対輝度の 1/r^2 の違いは
相殺され，曲線は「ピークから何割落ちるか」を直接表す（曲線の最小値がそのまま
輪帯 min/max 比になる）．各リングには標準誤差（std/sqrt(N)）の帯を重ねて，
ノイズが平均化で潰れていることを図の上で確認できるようにする．

分割数の既定は 20 で，評価パイプライン（`calc_radial_profile` の既定）および
`build_radius_sweep_visual.py` の輝度マップ正規化と同一にしてある．図どうしで
「ピーク＝どの輪帯か」がずれないようにするため．

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

from __future__ import annotations

import argparse
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))
sys.path.insert(0, str(Path(__file__).resolve().parents[1] / "src"))

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

from sm_tias_lightsim.evaluation.roi_presets import SIM_ROI, extract_roi_xywh  # noqa: E402
from sm_tias_lightsim.evaluation.zones import compute_normalized_distance_map  # noqa: E402

ROOT = Path(__file__).resolve().parents[1]
# 旧構成（LED 壁向き・バッフルなし）の r=100．バッフルの効果を同じ図で比較するための参照
LEGACY_R100 = ROOT / "output" / "sphere_r100_loft_disk_fix" / "luminance.npy"
# 最外ビンは ROI の四隅しか含まず画素数が激減し統計が荒れるため，代表値の算出から外す
MAX_DISTANCE = 0.95

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


def radial_stats(luminance_npy: Path, n_bins: int) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
    """輪帯ごとの (ビン中心距離, 平均輝度/ROI平均, 標準誤差/ROI平均) を返す.

    平滑化は掛けない．方位平均そのものがノイズ低減になっており，
    平滑化を重ねると本来見たい径方向の構造まで鈍らせるため．

    ビン割りは `evaluation/radial.py:calc_radial_profile` と同一規約
    （正規化楕円距離 [0, 1] を n_bins 等分・`digitize(right=True)`）．
    """
    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()

    distance = compute_normalized_distance_map(*roi.shape)
    edges = np.linspace(0.0, 1.0, n_bins + 1)
    centers = (edges[:-1] + edges[1:]) / 2.0
    indices = np.digitize(distance.ravel(), edges, right=True)
    flat = roi.ravel()

    means = np.full(n_bins, np.nan)
    sems = np.full(n_bins, np.nan)
    for i in range(1, n_bins + 1):
        values = flat[indices == i]
        if values.size:
            means[i - 1] = values.mean() / mean
            sems[i - 1] = values.std(ddof=1) / np.sqrt(values.size) / mean
    return centers, means, sems


def normalize_to_peak(
    centers: np.ndarray, means: np.ndarray, sems: np.ndarray
) -> tuple[np.ndarray, np.ndarray]:
    """最大の輪帯平均を 1 とするスケールに直した (平均, 標準誤差) を返す.

    ROI 平均で割る正規化だと曲線ごとに基準の高さが違い「どれだけ落ちるか」が
    読みにくい．ピークを 1 に揃えると縦軸がそのまま「ピーク比」になる．
    四隅しか含まない最外ビンは統計が荒いのでピーク探索から外す．
    """
    valid = (~np.isnan(means)) & (centers <= MAX_DISTANCE)
    peak = means[valid].max()
    return means / peak, sems / peak


def ring_min_max_ratio(centers: np.ndarray, means: np.ndarray) -> float:
    """輪帯平均の min/max 比を返す（1 が理想＝ドームが無い）.

    中心リング基準の比は，中心が 1 リングだけで統計が荒く，かつプロファイルの
    ピークが中心そのものではなく d≈0.05 にあるため代表値として不安定だった．
    全輪帯の min/max なら非単調なプロファイルにも頑健に効く
    （`evaluation/radial.py:calc_radial_min_max_ratio` と同じ考え方）．
    四隅しか含まない最外ビン（d > MAX_DISTANCE）は除外する．
    """
    valid = (~np.isnan(means)) & (centers <= MAX_DISTANCE)
    return float(means[valid].min() / means[valid].max())


def draw_all_radii(ax, variant: str, cavity: str, n_bins: int) -> list[tuple[float, float]]:
    """全半径の輪帯プロファイルを重ね描きし，(半径, 周辺/中心比) のリストを返す."""
    colors = plt.cm.viridis(np.linspace(0.0, 0.9, len(RADII)))
    ratios: list[tuple[float, float]] = []
    for radius, color in zip(RADII, colors, strict=True):
        path = OUTPUT_BASE / run_dir_name(variant, cavity) / job_name(radius) / "luminance.npy"
        if not path.exists():
            continue
        centers, means, sems = radial_stats(path, n_bins)
        ratios.append((radius, ring_min_max_ratio(centers, means)))
        means, sems = normalize_to_peak(centers, means, sems)
        ax.plot(centers, means, color=color, lw=1.6, label=f"r={radius:g}")
        ax.fill_between(centers, means - sems, means + sems, color=color, alpha=0.35, lw=0)
    ax.axhline(1.0, color="gray", lw=0.8, ls="--")
    ax.set_xlabel("正規化楕円距離（0=中心, 1=ROI 四隅）")
    ax.set_ylabel("輪帯平均輝度（ピーク=1）")
    ax.set_title("同心円プロファイルの半径依存（帯は ±1 標準誤差）", fontsize=11)
    ax.legend(fontsize=8, ncol=2)
    ax.grid(alpha=0.3)
    return ratios


def draw_dome_depth(ax, ratios: list[tuple[float, float]], legacy: float | None) -> None:
    """輪帯 min/max の半径依存を描く（1 に近いほどドームが浅い＝均一）."""
    radii = [r for r, _ in ratios]
    values = [v for _, v in ratios]
    ax.plot(radii, values, marker="o", color="tab:red", label="新構成（LED 壁面 + バッフル）")
    if legacy is not None:
        ax.axhline(legacy, color="tab:blue", lw=1.4, ls="-.",
                   label=f"旧 axial_wall r=100 ({legacy:.3f})")
    ax.axhline(1.0, color="gray", lw=0.8, ls="--")
    ax.set_xlabel("球内半径 (mm)")
    ax.set_ylabel(f"輪帯 min/max（d<={MAX_DISTANCE}）")
    ax.set_title("ドームの深さ（1 が理想）", fontsize=11)
    ax.legend(fontsize=8)
    ax.grid(alpha=0.3)


def draw_config_compare(ax, variant: str, cavity: str, n_bins: int) -> None:
    """同じ r=100 で旧構成（バッフルなし）と新構成の輪帯プロファイルを比較する."""
    run_dir = OUTPUT_BASE / run_dir_name(variant, cavity)
    targets = [
        (run_dir / "r100" / "luminance.npy", f"新 {variant}/{cavity} r=100", "tab:red"),
        (LEGACY_R100, "旧 axial_wall r=100（バッフルなし）", "tab:blue"),
    ]
    for path, label, color in targets:
        if not path.exists():
            continue
        centers, means, sems = radial_stats(path, n_bins)
        means, sems = normalize_to_peak(centers, means, sems)
        ax.plot(centers, means, color=color, lw=2.0, label=label)
        ax.fill_between(centers, means - sems, means + sems, color=color, alpha=0.3, lw=0)
    ax.axhline(1.0, color="gray", lw=0.8, ls="--")
    ax.set_xlabel("正規化楕円距離")
    ax.set_ylabel("輪帯平均輝度（ピーク=1）")
    ax.set_title("同じ r=100 での構成比較", fontsize=11)
    ax.legend(fontsize=9)
    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("--bins", type=int, default=20,
                        help="輪帯の分割数（既定 20＝評価パイプラインと同一）")
    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)
    output = args.output or run_dir / "radial_comparison.png"
    fig, axes = plt.subplots(1, 3, figsize=(19, 6))
    ratios = draw_all_radii(axes[0], args.variant, args.cavity, args.bins)
    legacy = None
    if LEGACY_R100.exists():
        legacy = ring_min_max_ratio(*radial_stats(LEGACY_R100, args.bins)[:2])
    draw_dome_depth(axes[1], ratios, legacy)  # 旧構成は水平線で参照値として重ねる
    draw_config_compare(axes[2], args.variant, args.cavity, args.bins)
    fig.suptitle(
        f"同心円（輪帯）プロファイル比較 [{args.variant} / {args.cavity}]"
        f"　輪帯 {args.bins} 分割・方位平均でノイズを低減・最大輪帯を 1 に正規化", fontsize=14
    )
    fig.tight_layout()
    output.parent.mkdir(parents=True, exist_ok=True)
    fig.savefig(output, dpi=110)
    print(f"出力: {output}")
    for radius, ratio in ratios:
        print(f"  r={radius:6.1f}  輪帯 min/max = {ratio:.4f}")
    if legacy is not None:
        print(f"  旧 axial_wall r=100  輪帯 min/max = {legacy:.4f}")
    return 0


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