"""spp とノイズの関係を輝度マップで見せるパネル（TEST_19 の補助図・スライド用）.

TEST_19 の spp スイープ（output/spp_noise_sweep/）から，条件ごとに spp 256〜8192 の輝度マップを
横に並べる．各マップは σ=20 平滑化の 99 パーセンタイル（ロバスト最大）を 1 とし，同じ
カラースケールで描くので，spp が上がるにつれてノイズだけが減っていく様子が見える．
各セルの下にノイズ CoV と GPU 時間を書く．

実行方法:
    uv run python scripts/build_spp_noise_maps.py
    uv run python scripts/build_spp_noise_maps.py --conditions sphere_d340_th000 sphere_d500_th090
    uv run python scripts/build_spp_noise_maps.py --normalize minmax       # 各マップの min=0・max=1
    uv run python scripts/build_spp_noise_maps.py --normalize percentile   # 1% 点=0・99% 点=1
    uv run python scripts/build_spp_noise_maps.py --normalize mean --lang en   # 平均=1・英語版
    uv run python scripts/build_spp_noise_maps.py --normalize mean --lang en --slide \
        --conditions sphere_d340_th000   # スライド用 3×2（題・指標なし，spp だけ）
    uv run python scripts/build_spp_noise_maps.py --lang en --slide --min-metric raw \
        --conditions sphere_d340_th000   # 最大=1 で正規化し，spp の横に画像全体 min を併記
"""

from __future__ import annotations

import argparse
import sys
from pathlib import Path

import matplotlib
import numpy as np
import pandas as pd
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().parents[1] / "src"))

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

ROOT = Path(__file__).resolve().parents[1]
SWEEP_DIR = ROOT / "output" / "spp_noise_sweep"
SPPS = [256, 512, 1024, 2048, 4096, 8192]
SEED = 0
LABELS = {
    "ja": {
        "sphere_d340_th000": "球 φ340 軸上",
        "sphere_d340_th090": "球 φ340 θ=90°",
        "sphere_d500_th090": "球 φ500 θ=90°",
        "cube_L234_th000": "立方体 L=234 軸上",
        "cube_L234_th090": "立方体 L=234 θ=90°",
        "cube_L340_th090": "立方体 L=340 θ=90°",
    },
    "en": {
        "sphere_d340_th000": "Sphere φ340, on-axis",
        "sphere_d340_th090": "Sphere φ340, θ=90°",
        "sphere_d500_th090": "Sphere φ500, θ=90°",
        "cube_L234_th000": "Cube L=234, on-axis",
        "cube_L234_th090": "Cube L=234, θ=90°",
        "cube_L340_th090": "Cube L=340, θ=90°",
    },
}
TEXT = {
    "ja": {
        "cell": "ノイズ CoV {noise:.3f}\nGPU {t:.1f} s",
        "cbar": "相対輝度",
        "cbar_slide": {"mean": "輝度\n（ROI 平均で正規化）",
                       "robust_peak": "輝度\n（各マップの最大で正規化）",
                       "max": "輝度\n（各マップの最大で正規化）",
                       "minmax": "輝度\n（各マップの min/max で正規化）",
                       "percentile": "輝度\n（各マップの 1〜99% 点で正規化）"},
        "title": "spp と MC ノイズ（{norm}・seed 0）",
        "norm": {"minmax": "各マップの画素の最小 = 0・最大 = 1",
                 "percentile": "各マップの画素の 1% 点 = 0・99% 点 = 1",
                 "mean": "ROI 平均 = 1・共通スケール",
                 "robust_peak": "ロバスト最大 = 1 で正規化・共通スケール"},
    },
    "en": {
        "cell": "noise CoV {noise:.3f}\nGPU {t:.1f} s",
        "cbar": "relative luminance",
        "cbar_slide": {"mean": "Luminance\n(normalized by mean)",
                       "robust_peak": "Luminance\n(normalized by max of each)",
                       "max": "Luminance\n(normalized by max of each)",
                       "minmax": "Luminance\n(normalized by min/max of each)",
                       "percentile": "Luminance\n(normalized by 1-99% of each)"},
        "title": "Luminance maps vs spp ({norm}, seed 0)",
        "norm": {"minmax": "each map: pixel min = 0, max = 1",
                 "percentile": "each map: 1st percentile = 0, 99th percentile = 1",
                 "mean": "ROI mean = 1, common scale",
                 "robust_peak": "robust peak = 1, common scale"},
    },
}

plt.rcParams["font.family"] = ["Meiryo", "Yu Gothic", "MS Gothic", "sans-serif"]


def load_roi(condition: str, spp: int) -> np.ndarray:
    """条件・spp の luminance.npy を読んで ROI を返す（seed 0）."""
    path = SWEEP_DIR / f"{condition}_spp{spp}_s{SEED}" / "luminance.npy"
    if not path.exists():
        raise FileNotFoundError(f"レンダ結果がありません: {path}")
    return extract_roi_xywh(np.load(path).astype(np.float64), SIM_ROI)


def light_smooth(roi: np.ndarray, sigma: float = 4.0) -> np.ndarray:
    """TEST_18 のマップパネルと同じ軽い平滑化（メディアン 5 px → ガウス σ px）.

    σ=12 にすると画像全体 min/max 比（``fullimage_raw_min_max``）と同じ平滑化になり，
    マップの見た目とラベルの min が一致する．
    """
    return gaussian_filter(median_filter(roi, size=5), sigma=sigma)


def normalized(roi: np.ndarray, mode: str) -> np.ndarray:
    """ROI を正規化して返す.

    Args:
        roi: 輝度 ROI（生，または平滑化後）．
        mode: ``robust_peak`` は σ=20 平滑化の 99 パーセンタイルを 1 に，``minmax`` は各マップの
            画素の最小を 0・最大を 1 に（ノイズの裾まで含めた値域．低 spp はファイアフライで
            ほぼ真っ黒になる），``percentile`` は画素の 1% 点を 0・99% 点を 1，``mean`` は
            ROI 平均を 1 にする（共通スケール）．
    """
    if mode == "minmax":
        return (roi - roi.min()) / (roi.max() - roi.min())
    if mode == "percentile":
        lo, hi = np.percentile(roi, [1.0, 99.0])
        return (roi - lo) / (hi - lo)
    if mode == "mean":
        return roi / roi.mean()
    if mode == "max":
        return roi / roi.max()  # 平滑化後のマップに使う（生の画素だとファイアフライで決まる）
    return roi / np.percentile(gaussian_filter(roi, 20.0), 99.0)


def summary_rows(condition: str) -> pd.DataFrame:
    """summary.csv から条件・seed 0 の行を spp 順で返す."""
    df = pd.read_csv(SWEEP_DIR / "summary.csv")
    rows = df[(df["condition"] == condition) & (df["seed"] == SEED)]
    return rows.set_index("spp").sort_index()


def draw(conditions: list[str], output: Path, vmin: float, vmax: float, mode: str,
         lang: str = "ja") -> None:
    """条件を行，spp を列にして輝度マップを描く."""
    text = TEXT[lang]
    n_rows = len(conditions)
    fig, axes = plt.subplots(n_rows, len(SPPS), figsize=(2.6 * len(SPPS), 2.9 * n_rows))
    axes = np.atleast_2d(axes)
    for ci, condition in enumerate(conditions):
        rows = summary_rows(condition)
        for si, spp in enumerate(SPPS):
            ax = axes[ci, si]
            img = normalized(load_roi(condition, spp), mode)
            im = ax.imshow(img, cmap="hot", vmin=vmin, vmax=vmax)
            noise = rows.loc[spp, "noise_cov"] if spp in rows.index else np.nan
            t_render = rows.loc[spp, "t_render"] if spp in rows.index else np.nan
            if ci == 0:
                ax.set_title(f"spp {spp}", fontsize=11)
            ax.set_xlabel(text["cell"].format(noise=noise, t=t_render), fontsize=9)
            ax.set_xticks([])
            ax.set_yticks([])
        axes[ci, 0].set_ylabel(LABELS[lang].get(condition, condition), fontsize=10)
    fig.suptitle(text["title"].format(norm=text["norm"][mode]), fontsize=12)
    fig.tight_layout(rect=(0, 0, 0.96, 0.96), h_pad=3.0)
    cbar_ax = fig.add_axes((0.965, 0.15, 0.012, 0.7))
    fig.colorbar(im, cax=cbar_ax, label=text["cbar"])
    output.parent.mkdir(parents=True, exist_ok=True)
    fig.savefig(output, dpi=150)
    plt.close(fig)
    print(f"保存: {output}")


def draw_slide(condition: str, output: Path, vmin: float, vmax: float, mode: str,
               lang: str = "en", n_cols: int = 3, min_col: str | None = None,
               smooth: bool = False, smooth_sigma: float = 4.0) -> None:
    """スライド貼り付け用: 1 筐体の spp 6 段階を n_cols 列 × 2 行に並べる（題・枠・指標なし）.

    各マップの下に ``spp=N`` を書き（``min_col`` を指定すると summary.csv のその列を
    ``(min=0.93)`` として併記），右にカラーバーを置く（既存スライドの
    「Comparison by cube edge lengths」と同じ体裁）．
    """
    text = TEXT[lang]
    rows = summary_rows(condition)
    n_rows = int(np.ceil(len(SPPS) / n_cols))
    fig, axes = plt.subplots(n_rows, n_cols, figsize=(2.9 * n_cols + 1.3, 2.9 * n_rows))
    axes = np.atleast_2d(axes)
    im = None
    for i, spp in enumerate(SPPS):
        ax = axes[i // n_cols, i % n_cols]
        roi = load_roi(condition, spp)
        img = normalized(light_smooth(roi, smooth_sigma) if smooth else roi, mode)
        if smooth and mode == "robust_peak":
            img = img / img.max()  # 表示するマップ自身の最大を 1 に（max of each を文字どおりに）
        im = ax.imshow(img, cmap="hot", vmin=vmin, vmax=vmax)
        label = f"spp={spp}"
        if min_col == "display":
            label += f" (min={img.min():.3f})"  # 表示しているマップそのものの最小（max=1）
        elif min_col is not None and spp in rows.index:
            label += f" (min={rows.loc[spp, min_col]:.3f})"
        ax.set_xlabel(label, fontsize=14, labelpad=6)
        ax.set_xticks([])
        ax.set_yticks([])
        for spine in ax.spines.values():
            spine.set_visible(False)
    for ax in axes.flat[len(SPPS):]:
        ax.axis("off")
    fig.tight_layout(rect=(0, 0, 0.87, 1.0), h_pad=1.0, w_pad=1.0)
    cbar_ax = fig.add_axes((0.875, 0.25, 0.016, 0.5))
    cbar = fig.colorbar(im, cax=cbar_ax, format="%.2f")
    cbar.set_label(text["cbar_slide"][mode], fontsize=12)
    cbar.ax.tick_params(labelsize=11)
    output.parent.mkdir(parents=True, exist_ok=True)
    fig.savefig(output, dpi=200, facecolor="white")
    plt.close(fig)
    print(f"保存: {output}")


def main() -> int:
    """パネルを描く."""
    parser = argparse.ArgumentParser(description="spp とノイズの輝度マップパネル")
    parser.add_argument("--conditions", nargs="+",
                        default=["sphere_d340_th000", "sphere_d500_th090"])
    parser.add_argument("--normalize",
                        choices=("robust_peak", "minmax", "percentile", "mean", "max"),
                        default="robust_peak")
    parser.add_argument("--vmin", type=float, default=None,
                        help="カラースケール下限（既定: robust_peak 0.5 / mean 0.5 / 他 0）")
    parser.add_argument("--vmax", type=float, default=None,
                        help="カラースケール上限（既定: robust_peak 1.3 / mean 1.5 / 他 1）")
    parser.add_argument("--lang", choices=("ja", "en"), default="ja", help="図中の文字の言語")
    parser.add_argument("--slide", action="store_true",
                        help="スライド貼り付け用（--conditions の先頭 1 条件を 3×2 に並べる）")
    parser.add_argument("--smooth", action="store_true",
                        help="--slide でマップに TEST_18 と同じ軽い平滑化を掛ける")
    parser.add_argument("--smooth-sigma", type=float, default=4.0,
                        help="--smooth のガウス σ [px]（12 にすると min ラベルの定義と一致）")
    parser.add_argument("--min-metric", choices=("none", "corrected", "raw", "ring", "display"),
                        default="none",
                        help="--slide で spp の横に書く min/max 比: corrected=画像全体"
                             "（ノイズ補正あり）/ raw=画像全体（補正なし）/ ring=輪帯平均"
                             "/ display=表示マップ自身の最小（図と一致する）")
    parser.add_argument("--output", type=Path, default=None,
                        help="出力 PNG（既定: spp_noise_maps.png / spp_noise_maps_minmax.png）")
    args = parser.parse_args()
    # スライド用は「最大 = 1」の意味を保つため上端を 1.0 に固定（超えた画素は白に飽和）
    slide_lo = 0.85 if args.smooth else 0.5
    defaults = {"robust_peak": (slide_lo, 1.0) if args.slide else (0.5, 1.3),
                "mean": (0.5, 1.5)}.get(args.normalize, (0.0, 1.0))
    vmin = args.vmin if args.vmin is not None else defaults[0]
    vmax = args.vmax if args.vmax is not None else defaults[1]
    suffix = ("" if args.normalize == "robust_peak" else f"_{args.normalize}")
    suffix += "_en" if args.lang == "en" else ""
    if args.slide:
        condition = args.conditions[0]
        min_col = {"corrected": "fullimage_min_max", "raw": "fullimage_raw_min_max",
                   "ring": "ring_min_max", "display": "display"}.get(args.min_metric)
        if min_col is not None:
            suffix += f"_min-{args.min_metric}"
        if args.smooth:
            suffix += f"_smooth{args.smooth_sigma:g}"
        output = args.output or SWEEP_DIR / f"spp_noise_slide_{condition}{suffix}.png"
        draw_slide(condition, output, vmin, vmax, args.normalize, args.lang, min_col=min_col,
                   smooth=args.smooth, smooth_sigma=args.smooth_sigma)
        return 0
    output = args.output or SWEEP_DIR / f"spp_noise_maps{suffix}.png"
    draw(args.conditions, output, vmin, vmax, args.normalize, args.lang)
    return 0


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