"""軸上スイープ（TEST_15 球 10 半径・TEST_16 立方体 10 サイズ）の輝度マップを正規化方式で比較する.

同じ luminance.npy を 3 方式で正規化して並べ，正規化の違いだけで見え方がどう変わるかを確かめる．
軸上配置は同心円状の分布なので輪帯ピーク正規化が本来の想定どおり働く条件であり，
非対称な分布（TEST_18）で採用したロバスト最大値との差が小さければ，過去の図の読みは正規化に
依存していなかったと言える．

方式（build_led_placement_maps と同じ定義）:
    ring_peak   : 20 輪帯平均の最大値 = 1.0（TEST_15/16 の図で使った従来方式）
    robust_peak : 平滑化後の 99 パーセンタイル = 1.0（TEST_18 の既定）
    mean        : ROI 平均 = 1.0

出力（output/_normalization_comparison/）:
    norm_sphere.png / norm_cube.png : 行 = 方式，列 = サイズ（10）．共通スケール
    norm_summary.csv                : 各マップの基準値（ピーク・パーセンタイル・平均）と比

実行方法:
    uv run python scripts/build_normalization_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))
sys.path.insert(0, str(Path(__file__).resolve().parents[1] / "src"))

from build_led_placement_maps import (  # noqa: E402
    CMAP,
    ROBUST_PERCENTILE,
    normalize_map,
    smoothed_roi,
)
from sweep_cube_size import OUTPUT_BASE as CUBE_OUTPUT_BASE  # noqa: E402
from sweep_cube_size import RUN_DIR as CUBE_RUN_DIR  # noqa: E402
from sweep_cube_size import SIZES as CUBE_SIZES  # noqa: E402
from sweep_cube_size import job_name as cube_job_name  # noqa: E402
from sweep_sphere_radius import OUTPUT_BASE as SPHERE_OUTPUT_BASE  # noqa: E402
from sweep_sphere_radius import RADII as SPHERE_RADII  # noqa: E402
from sweep_sphere_radius import job_name as sphere_job_name  # noqa: E402

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"]

ROOT = Path(__file__).resolve().parents[1]
OUT_DIR = ROOT / "output" / "_normalization_comparison"
SPHERE_RUN_DIR = "board_baffle_similar_radial"

MODES = ("ring_peak", "robust_peak", "mean")
MODE_LABEL = {
    "ring_peak": "輪帯ピーク\n（20 輪帯平均の最大=1）",
    "robust_peak": f"ロバスト最大\n（平滑化後 {ROBUST_PERCENTILE:g}%ile=1）",
    "mean": "ROI 平均\n（mean=1）",
}
# 共通スケール．mean は 1 を中心に上下対称，他は上限 1.0
VMIN = 0.80
VMAX = {"ring_peak": 1.0, "robust_peak": 1.0, "mean": 1.0 + (1.0 - VMIN)}


def conditions() -> dict[str, list[tuple[str, Path]]]:
    """形状 → [(ラベル, luminance.npy)] を返す（存在するものだけ）."""
    sphere_dir = SPHERE_OUTPUT_BASE / SPHERE_RUN_DIR
    sphere = [(f"φ{2 * r:g}", sphere_dir / sphere_job_name(r) / "luminance.npy")
              for r in SPHERE_RADII]
    cube = [(f"L={size:g}", CUBE_OUTPUT_BASE / CUBE_RUN_DIR / cube_job_name(size) / "luminance.npy")
            for size in CUBE_SIZES]
    return {"sphere": [c for c in sphere if c[1].exists()],
            "cube": [c for c in cube if c[1].exists()]}


def reference_values(path: Path) -> dict[str, float]:
    """3 方式の基準値（正規化に使う分母）と生の ROI 統計を返す."""
    arr = np.load(path).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=20)
    smooth = smoothed_roi(path)
    return {
        "raw_mean": float(roi.mean()),
        "ring_peak": float(np.nanmax(profile[:, 1])),
        "robust_peak": float(np.percentile(smooth, ROBUST_PERCENTILE)),
        "mean": float(smooth.mean()),
        "smooth_max": float(smooth.max()),
    }


def draw(shape: str, items: list[tuple[str, Path]], output: Path) -> None:
    """行 = 正規化方式，列 = サイズのパネルを描く."""
    n_cols = len(items)
    figsize = (1.9 * n_cols + 1.5, 2.1 * len(MODES) + 1.0)
    fig, axes = plt.subplots(len(MODES), n_cols, figsize=figsize, squeeze=False)
    images = {}
    for r, mode in enumerate(MODES):
        for c, (label, path) in enumerate(items):
            ax = axes[r, c]
            data = normalize_map(path, mode)
            images[mode] = ax.imshow(data, cmap=CMAP, vmin=VMIN, vmax=VMAX[mode])
            ax.set_xticks([])
            ax.set_yticks([])
            if r == 0:
                ax.set_title(label, fontsize=9)
            if c == 0:
                ax.set_ylabel(MODE_LABEL[mode], fontsize=9, fontweight="bold")
        fig.colorbar(images[mode], ax=axes[r, :].tolist(), fraction=0.012, pad=0.01)
    title = "球 similar_radial（TEST_15）" if shape == "sphere" else "立方体（TEST_16）"
    fig.suptitle(f"正規化方式の比較 — {title} 軸上 1 灯・板向き・バッフル付き．"
                 f"上 2 行は {VMIN:g}〜1.0，mean 行は {VMIN:g}〜{VMAX['mean']:g}", fontsize=12)
    output.parent.mkdir(parents=True, exist_ok=True)
    fig.savefig(output, dpi=110, bbox_inches="tight")
    plt.close(fig)
    print(f"出力: {output}")


def write_summary(conds: dict[str, list[tuple[str, Path]]]) -> None:
    """各マップの基準値と比を CSV と標準出力に書く."""
    rows = []
    for shape, items in conds.items():
        for label, path in items:
            ref = reference_values(path)
            rows.append({
                "shape": shape, "size": label,
                **{k: round(v, 6) for k, v in ref.items()},
                "ring_over_robust": round(ref["ring_peak"] / ref["robust_peak"], 4),
                "ring_over_mean": round(ref["ring_peak"] / ref["mean"], 4),
                "robust_over_mean": round(ref["robust_peak"] / ref["mean"], 4),
            })
    OUT_DIR.mkdir(parents=True, exist_ok=True)
    with (OUT_DIR / "norm_summary.csv").open("w", newline="", encoding="utf-8") as fh:
        writer = csv.DictWriter(fh, fieldnames=list(rows[0].keys()))
        writer.writeheader()
        writer.writerows(rows)
    print(f"\n{'形状':>7} {'サイズ':>6} {'輪帯ピーク':>10} {'ロバスト最大':>12} {'平均':>8} "
          f"{'輪帯/ロバスト':>12} {'輪帯/平均':>9} {'ロバスト/平均':>12}")
    for r in rows:
        print(f"{r['shape']:>7} {r['size']:>6} {r['ring_peak']:10.4f} {r['robust_peak']:12.4f} "
              f"{r['mean']:8.4f} {r['ring_over_robust']:12.4f} {r['ring_over_mean']:9.4f} "
              f"{r['robust_over_mean']:12.4f}")


def main() -> int:
    """正規化方式の比較パネルと集約表を生成する."""
    parser = argparse.ArgumentParser(description="軸上スイープの輝度マップを正規化方式で比較")
    parser.parse_args()
    conds = conditions()
    if not conds["sphere"] and not conds["cube"]:
        print("エラー: TEST_15/16 の結果が見つかりません")
        return 1
    for shape, items in conds.items():
        if items:
            draw(shape, items, OUT_DIR / f"norm_{shape}.png")
    write_summary(conds)
    return 0


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