"""照明配置スイープ（TEST_18）の輝度マップを条件別に並べたパネルを生成する.

数値指標（build_led_placement_comparison.py）とは別に，輝度分布の「形」を目で比べるための
画像パネル．各パネルは球（上段）と同内面積の立方体（下段）を同じ条件の並びで置き，
平滑化後の 99 パーセンタイルを 1.0 とするロバスト最大値正規化（`--normalize` で輪帯ピーク・
ROI 平均にも切替可）・共通カラースケールで描く．LED は +X＝画像右側．

正規化の注意: 従来の輪帯ピーク正規化（20 輪帯平均の最大 = 1.0）は同心円状の分布を前提にしており，
非対称な分布では同じ輪帯に明るい側と暗い側が混ざって基準が下がり，明るい側が 1.0 を超えて飛ぶ．
本スクリプトの既定はこれを避けるため方位に依らないロバスト最大値を使う．

出力（output/led_placement_sweep/maps/）:
    maps_position_<ペア>.png : θ を横に並べる（13 段階）．4 ペア分
    maps_direction.png       : θ × 向き（板向き / 壁法線向き / 真下向き）．代表ペア
    maps_count.png           : 灯数 N=1/2/3/4（θ=90°）と N=4 の θ 違い．全ペア
    maps_azimuth.png         : 立方体の方位角 φ（球は対照）
    maps_ring.png            : 真下向き N=1/2/4/8 と板向き N=8
    maps_material.png        : 旧材質 vs 校正材質（θ=0/30/90/120・多灯）

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

from __future__ import annotations

import argparse
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"))

import sweep_led_placement as sweep  # noqa: E402
from build_radius_sweep_visual import normalized_map  # noqa: E402
from sweep_led_placement import (  # noqa: E402
    AIMS,
    AZIMUTH_CUBE,
    AZIMUTH_SPHERE,
    COUNT_THETA,
    CUBE_N8_AZIMUTH,
    DIRECTION_THETAS,
    MATERIAL_SUFFIX,
    SHAPES,
    SIZE_PAIRS,
    THETA_GRID,
    Condition,
)

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

SHAPE_LABEL = {"sphere": "球", "cube": "立方体"}
AIM_LABEL = {"board": "板向き+バッフル", "normal": "壁法線向き+バッフル", "down": "真下向き"}
CMAP = "hot"
# 共通スケールの下限（最大=1.0 に対する比）．全パネルで固定して見比べられるようにする
VMIN = 0.80
REP_PAIR = 2  # 代表ペア（r=170 / L=234）

# 正規化の方式（--normalize）:
#   robust_peak: 平滑化後の 99 パーセンタイルを 1.0 にする（既定）．方位に依らず「最も明るい所
#                からの落ち込み」を表す．画素の最大値は MC ノイズの外れ値に張り付くので
#                パーセンタイルで避ける
#   ring_peak  : 20 輪帯平均の最大値を 1.0（従来）．同心円状の分布向けで，非対称な分布では同じ輪帯に
#                明るい側と暗い側が混ざって基準が下がり，明るい側が 1.0 を超えて飛ぶ
#   mean       : ROI 平均を 1.0．CoV・mean anchor 正規化と同じ基準（明るい側が 1 を超える）
NORMALIZE_MODES = ("robust_peak", "ring_peak", "mean")
ROBUST_PERCENTILE = 99.0
SMOOTH_MEDIAN = 5
SMOOTH_SIGMA = 4.0
_normalize = "robust_peak"


def output_root(material: str) -> Path:
    """材質系列に対応する出力ディレクトリ."""
    return Path(sweep.ROOT) / "output" / f"led_placement_sweep{MATERIAL_SUFFIX[material]}"


def cond_name(shape: str, size: float, theta: float, aim: str, n: int,
              azimuth: float = 0.0, rep: int = 1) -> str:
    """条件名（出力ディレクトリ名）を組み立てる."""
    return Condition("", shape, size, theta, aim, n, rep, azimuth).name


def smoothed_roi(path: Path) -> np.ndarray:
    """luminance.npy の ROI を表示用に平滑化して返す（normalized_map と同じ平滑化）."""
    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)
    return gaussian_filter(median_filter(roi, size=SMOOTH_MEDIAN), sigma=SMOOTH_SIGMA)


def normalize_map(path: Path, mode: str = "robust_peak") -> np.ndarray:
    """輝度マップを指定方式で正規化して返す（表示用に平滑化済み）."""
    if mode == "ring_peak":
        return normalized_map(path)
    smooth = smoothed_roi(path)
    if mode == "mean":
        return smooth / float(smooth.mean())
    if mode == "robust_peak":
        return smooth / float(np.percentile(smooth, ROBUST_PERCENTILE))
    raise ValueError(f"未知の正規化方式: {mode}（{', '.join(NORMALIZE_MODES)}）")


def load_map(name: str, material: str = "legacy") -> np.ndarray | None:
    """条件名から正規化・平滑化済みの輝度マップを返す（無ければ None）."""
    path = output_root(material) / name / "luminance.npy"
    return normalize_map(path, _normalize) if path.exists() else None


def colorbar_label() -> str:
    """正規化方式に応じたカラーバーの説明."""
    return {
        "robust_peak": f"輝度（平滑化後 {ROBUST_PERCENTILE:g} パーセンタイル = 1.0）",
        "ring_peak": "輝度（最大の輪帯平均 = 1.0）",
        "mean": "輝度（ROI 平均 = 1.0）",
    }[_normalize]


def row_label(shape: str, size: float) -> str:
    """行見出し（形状とサイズ）."""
    return f"{SHAPE_LABEL[shape]} {'r' if shape == 'sphere' else 'L'}={size:g}"


def pair_sizes(pair: int) -> dict[str, float]:
    """ペア番号 → {形状: サイズ}."""
    return {shape: SIZE_PAIRS[pair][SHAPES.index(shape)] for shape in SHAPES}


def draw_grid(
    cells: list[list[tuple[str, np.ndarray | None]]],
    row_labels: list[str],
    title: str,
    output: Path,
    cell_size: float = 2.3,
) -> None:
    """(行 × 列) のマップ格子を共通スケールで描く．cells[row][col] = (見出し, マップ)."""
    n_rows, n_cols = len(cells), max(len(r) for r in cells)
    figsize = (cell_size * n_cols + 1.2, cell_size * n_rows + 1.0)
    fig, axes = plt.subplots(n_rows, n_cols, figsize=figsize, squeeze=False)
    image = None
    for r, row in enumerate(cells):
        for c in range(n_cols):
            ax = axes[r, c]
            ax.set_xticks([])
            ax.set_yticks([])
            if c >= len(row) or row[c][1] is None:
                ax.axis("off")
                continue
            label, data = row[c]
            vmax = 1.0 if _normalize != "mean" else 1.0 + (1.0 - VMIN)
            image = ax.imshow(data, cmap=CMAP, vmin=VMIN, vmax=vmax)
            ax.set_title(label, fontsize=8)
            if c == 0:
                ax.set_ylabel(row_labels[r], fontsize=10, fontweight="bold")
    if image is not None:
        fig.colorbar(image, ax=axes.ravel().tolist(), fraction=0.015, pad=0.01,
                     label=colorbar_label())
    fig.suptitle(title, fontsize=13)
    output.parent.mkdir(parents=True, exist_ok=True)
    fig.savefig(output, dpi=110, bbox_inches="tight")
    plt.close(fig)
    print(f"出力: {output}")


def panel_position(out_dir: Path) -> None:
    """θ を横に並べた位置スイープのパネル（ペアごとに 1 枚）."""
    for pair in range(len(SIZE_PAIRS)):
        sizes = pair_sizes(pair)
        cells, labels = [], []
        for shape in SHAPES:
            row = []
            for theta in THETA_GRID:
                data = load_map(cond_name(shape, sizes[shape], theta, "board", 1))
                row.append((f"θ={theta:g}°", data))
            cells.append(row)
            labels.append(row_label(shape, sizes[shape]))
        draw_grid(cells, labels,
                  f"位置スイープ ペア{pair + 1}（1 灯・板向き・バッフル付き）: "
                  "LED 極角 θ（0=底の中心・90=側壁中段・120=板の直下）．LED は画像右側",
                  out_dir / f"maps_position_pair{pair + 1}.png", cell_size=1.9)


def panel_direction(out_dir: Path) -> None:
    """θ × 向きのパネル（代表ペア）．行 = 形状 × 向き，列 = θ."""
    sizes = pair_sizes(REP_PAIR)
    cells, labels = [], []
    for shape in SHAPES:
        for aim in AIMS:
            row = []
            for theta in DIRECTION_THETAS:
                data = load_map(cond_name(shape, sizes[shape], theta, aim, 1))
                row.append((f"θ={theta:g}°", data))
            cells.append(row)
            labels.append(f"{SHAPE_LABEL[shape]}\n{AIM_LABEL[aim]}")
    draw_grid(cells, labels,
              f"向きの比較（1 灯・ペア{REP_PAIR + 1}: r={sizes['sphere']:g} / "
              f"L={sizes['cube']:g}）．LED は画像右側",
              out_dir / "maps_direction.png", cell_size=2.6)


def panel_count(out_dir: Path) -> None:
    """灯数のパネル（全ペア）．列 = N=1/2/3/4（θ=90°）+ N=4 の θ=60/120."""
    cells, labels = [], []
    for pair in range(len(SIZE_PAIRS)):
        sizes = pair_sizes(pair)
        for shape in SHAPES:
            size = sizes[shape]
            row = [(f"N={n} θ=90°", load_map(cond_name(shape, size, COUNT_THETA, "board", n)))
                   for n in (1, 2, 3, 4)]
            row += [(f"N=4 θ={theta:g}°", load_map(cond_name(shape, size, theta, "board", 4)))
                    for theta in (60.0, 120.0)]
            cells.append(row)
            labels.append(row_label(shape, size))
    draw_grid(cells, labels,
              "灯数の比較（板向き・各灯にバッフル・総光束一定）．1 灯目は画像右側，以降等方位",
              out_dir / "maps_count.png", cell_size=2.2)


def panel_azimuth(out_dir: Path) -> None:
    """方位角のパネル（θ=90°・1 灯・板向き）．立方体は φ=0〜40，球は 0 と 45."""
    cells, labels = [], []
    for pair in range(len(SIZE_PAIRS)):
        sizes = pair_sizes(pair)
        for shape in SHAPES:
            phis = [0.0] + (AZIMUTH_CUBE if shape == "cube" else AZIMUTH_SPHERE)
            row = [(f"φ={phi:g}°", load_map(cond_name(shape, sizes[shape], 90.0, "board", 1, phi)))
                   for phi in phis]
            cells.append(row)
            labels.append(row_label(shape, sizes[shape]))
    draw_grid(cells, labels,
              "方位角 φ の比較（θ=90°・1 灯・板向き）．立方体: φ=0 が面の中央・45 が稜．"
              "LED は画像右から φ だけ反時計回り",
              out_dir / "maps_azimuth.png", cell_size=2.2)


def panel_ring(out_dir: Path) -> None:
    """真下向き多灯のパネル（全ペア）．列 = 真下 N=1/2/4/8 + 板向き N=8."""
    cells, labels = [], []
    for pair in range(len(SIZE_PAIRS)):
        sizes = pair_sizes(pair)
        for shape in SHAPES:
            row = []
            for n in (1, 2, 4, 8):
                az = CUBE_N8_AZIMUTH if (shape == "cube" and n == 8) else 0.0
                name = cond_name(shape, sizes[shape], 90.0, "down", n, az)
                row.append((f"真下 N={n}", load_map(name)))
            az = CUBE_N8_AZIMUTH if shape == "cube" else 0.0
            name = cond_name(shape, sizes[shape], 90.0, "board", 8, az)
            row.append(("板向き N=8", load_map(name)))
            cells.append(row)
            labels.append(row_label(shape, sizes[shape]))
    draw_grid(cells, labels,
              "真下向き多灯（バッフル無し）と板向き多灯（θ=90° の輪・総光束一定）",
              out_dir / "maps_ring.png", cell_size=2.2)


def panel_material(out_dir: Path) -> None:
    """旧材質 vs 校正材質のパネル（校正材質で回した主要条件）."""
    specs = [(0.0, "board", 1), (30.0, "board", 1), (90.0, "board", 1), (120.0, "board", 1),
             (90.0, "board", 2), (90.0, "board", 4), (90.0, "down", 1), (90.0, "down", 4)]
    spec_label = {"board": "板向き", "down": "真下"}
    cells, labels = [], []
    # 校正材質は多灯・真下向きをペア 3 だけ回しているので，ペア 1 は位置の 4 条件だけ並べる
    for pair, spec_subset in ((0, specs[:4]), (2, specs)):
        sizes = pair_sizes(pair)
        for shape in SHAPES:
            for material, mat_label in (("legacy", "旧材質 ρ=0.85 光沢5%"),
                                        ("calibrated", "校正 ρ=0.81 diffuse")):
                row = []
                for theta, aim, n in spec_subset:
                    data = load_map(cond_name(shape, sizes[shape], theta, aim, n), material)
                    row.append((f"θ={theta:g}° {spec_label[aim]} N={n}", data))
                cells.append(row)
                labels.append(f"{row_label(shape, sizes[shape])}\n{mat_label}")
    draw_grid(cells, labels,
              "内壁材質の比較（旧材質 blendbsdf vs TEST_14 校正材質 純 diffuse ρ=0.81）",
              out_dir / "maps_material.png", cell_size=2.2)


def main() -> int:
    """全パネルを生成する."""
    parser = argparse.ArgumentParser(description="照明配置スイープの輝度マップパネルを生成")
    parser.add_argument("--output", type=Path, default=None, help="出力ディレクトリ")
    parser.add_argument("--normalize", choices=NORMALIZE_MODES, default="robust_peak",
                        help="正規化方式（既定 robust_peak＝平滑化後 99 パーセンタイル=1.0）")
    args = parser.parse_args()
    global _normalize  # noqa: PLW0603
    _normalize = args.normalize
    out_dir = args.output or output_root("legacy") / "maps"
    panel_position(out_dir)
    panel_direction(out_dir)
    panel_count(out_dir)
    panel_azimuth(out_dir)
    panel_ring(out_dir)
    panel_material(out_dir)
    return 0


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