"""実機とシミュの輝度マップの「形」を，全灯と単灯 3 灯（下・右・左）で突き合わせる.

評価の決め手は集約指標（min 等）ではなく輝度マップの形そのものである（ユーザー方針）．
本スクリプトは 4 構成それぞれについて，実機（DNG 由来の白板 ROI）とシミュを同じ処理
（シミュ ROI 解像度へリサイズ → メディアン 5 → ガウス σ=4 → ROI 平均で正規化）で並べ，

- 実機 / シミュ / 差分（シミュ − 実機）のマップ
- 形の一致度: 差分の RMS・最大，SSIM，NCC，方位フーリエ（a1・b1・a2・a4），
  上下比・左右比，傾き（1 次平面フィットの振幅と方位），お椀（2 次成分）の大きさ

を出力する．正規化は ROI 平均（=1）にする．単灯は絶対輝度が実機とシミュで揃っていない
（露出・LED 出力差）ので，形だけを比べるにはこれが自然．

実行方法:
    uv run python scripts/compare_real_vs_sim_shapes.py
    uv run python scripts/compare_real_vs_sim_shapes.py --sim-set legacy   # ρ=0.85 blend の旧レンダ
"""

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
from skimage.transform import resize

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 analyze_fourfold_component import aggregate, ring_harmonics  # noqa: E402
from build_radial_comparison import MAX_DISTANCE  # noqa: E402

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

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

ROOT = Path(__file__).resolve().parents[1]
OUT_DIR = ROOT / "output" / "real_vs_sim"
REAL_DIR = ROOT / "lightsim" / "luminance_maps"

CONFIGS = ("full", "bottom", "left", "right")
CONFIG_LABEL = {"full": "全灯 18 灯", "bottom": "下 単灯", "left": "左 単灯", "right": "右 単灯"}
REAL = {
    "full": REAL_DIR / "full_MiniTIAS_QM_20260601_124500.npy",
    "bottom": REAL_DIR / "bottom_SmTIAS_QM_20260615_141420.npy",
    "left": REAL_DIR / "left_SmTIAS_QM_20260615_150300.npy",
    "right": REAL_DIR / "right_SmTIAS_QM_20260615_170948.npy",
}
SIM_SETS = {
    # TEST_14 で確定した内壁（純 diffuse ρ=0.81）・spp 4096（本スクリプト用にサーバでレンダ）
    "rho081": {
        "label": "ρ=0.81 純 diffuse・spp 4096",
        "full": OUT_DIR / "sim_rho081_spp4096" / "luminance.npy",
        "bottom": OUT_DIR / "sim_single_bottom_rho081_spp4096" / "luminance.npy",
        "left": OUT_DIR / "sim_single_left_rho081_spp4096" / "luminance.npy",
        "right": OUT_DIR / "sim_single_right_rho081_spp4096" / "luminance.npy",
    },
    # TEST_12/14 の旧レンダ（blendbsdf ρ=0.85・ω=0.05・α=0.1・spp 1024）
    "legacy": {
        "label": "ρ=0.85 光沢 5%（旧材質）・spp 1024",
        "full": ROOT / "output" / "full_base_fix" / "luminance.npy",
        "bottom": ROOT / "output" / "singles_fix" / "bottom" / "luminance.npy",
        "left": ROOT / "output" / "singles_fix" / "left" / "luminance.npy",
        "right": ROOT / "output" / "singles_fix" / "right" / "luminance.npy",
    },
}
MEDIAN = 5
SIGMA = 4.0
DISPLAY_RANGE = (0.85, 1.15)  # ROI 平均 = 1 に対する表示レンジ
DIFF_RANGE = 0.08


def load_sim(path: Path) -> np.ndarray:
    """シミュ luminance.npy → 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]
    return extract_roi_xywh(arr, SIM_ROI)


def load_real(path: Path, shape: tuple[int, int]) -> np.ndarray:
    """実機 ROI（DNG 由来・LSC 適用済み）をシミュ ROI の解像度へリサイズして返す（生）."""
    arr = np.load(path).astype(np.float64)
    return resize(arr, shape, anti_aliasing=True, preserve_range=True)


def normalized_shape(roi: np.ndarray) -> np.ndarray:
    """平滑化（メディアン 5 → ガウス σ=4）して ROI 平均で正規化した形."""
    smooth = gaussian_filter(median_filter(roi, size=MEDIAN), sigma=SIGMA)
    return smooth / float(smooth.mean())


def plane_fit(shape_map: np.ndarray) -> tuple[float, float, float]:
    """1 次平面（傾き）を当てはめ，(振幅 %, 方位 deg, 2 次成分の振幅 %) を返す.

    振幅は ROI の半幅あたりの変化量（TEST_14 の「傾き」と同じ定義）．方位は画像座標で
    +X（右）を 0 deg，+Y（画像上）を 90 deg とする．2 次成分は傾きを除いた後の
    x²+y² 項（お椀）の振幅．
    """
    h, w = shape_map.shape
    y, x = np.mgrid[0:h, 0:w]
    x = (x - w / 2.0) / (w / 2.0)
    y = -(y - h / 2.0) / (h / 2.0)  # 画像上を +Y にする
    design = np.stack([np.ones(x.size), x.ravel(), y.ravel(),
                       (x**2 + y**2).ravel(), (x**2 - y**2).ravel(), (x * y).ravel()], axis=1)
    coeffs, *_ = np.linalg.lstsq(design, shape_map.ravel(), rcond=None)
    tilt = float(np.hypot(coeffs[1], coeffs[2])) * 100.0
    azimuth = float(np.degrees(np.arctan2(coeffs[2], coeffs[1])))
    bowl = float(coeffs[3]) * 100.0
    return tilt, azimuth, bowl


def shape_metrics(sim_map: np.ndarray, real_map: np.ndarray, sim_roi: np.ndarray,
                  real_roi: np.ndarray) -> dict[str, float]:
    """形の一致度と，各マップの形の記述子を返す."""
    diff = sim_map - real_map
    inner = compute_normalized_distance_map(*sim_map.shape) <= MAX_DISTANCE
    out = {
        "rms_diff": float(np.sqrt(np.mean(diff[inner] ** 2))),
        "max_abs_diff": float(np.abs(diff[inner]).max()),
        "ssim": float(calc_ssim(sim_map, real_map,
                                data_range=float(DISPLAY_RANGE[1] - DISPLAY_RANGE[0]))),
        "ncc": float(calc_ncc(sim_map, real_map)),
    }
    for tag, roi, smap in (("sim", sim_roi, sim_map), ("real", real_roi, real_map)):
        rings = ring_harmonics(roi, (1, 2, 4))
        for m in (1, 2, 4):
            out[f"{tag}_a{m}"], _ = aggregate(rings, f"a{m}")
            out[f"{tag}_b{m}"], _ = aggregate(rings, f"b{m}")
        h, w = smap.shape
        out[f"{tag}_top_bottom"] = float(smap[: h // 3].mean() / smap[-(h // 3):].mean())
        out[f"{tag}_left_right"] = float(smap[:, : w // 3].mean() / smap[:, -(w // 3):].mean())
        out[f"{tag}_peak_mean"] = float(np.percentile(smap, 99.0))
        out[f"{tag}_tilt"], out[f"{tag}_tilt_az"], out[f"{tag}_bowl"] = plane_fit(smap)
    return out


def cell_title(who: str, cfg: str, m: dict, tag: str) -> str:
    """マップの見出し（上下比・左右比・傾き・お椀）."""
    return (f"{who} {CONFIG_LABEL[cfg]}\n上下比 {m[f'{tag}_top_bottom']:.3f}・"
            f"左右比 {m[f'{tag}_left_right']:.3f}\n傾き {m[f'{tag}_tilt']:.1f}%@"
            f"{m[f'{tag}_tilt_az']:.0f}°・お椀 {m[f'{tag}_bowl']:+.1f}%")


def draw(results: dict[str, dict], maps: dict[str, tuple[np.ndarray, np.ndarray]],
         sim_label: str, output: Path) -> None:
    """3 行（実機 / シミュ / 差分）× 構成数の列のパネル（results にある構成だけ描く）."""
    configs = [c for c in CONFIGS if c in results]
    fig, axes = plt.subplots(3, len(configs), figsize=(4.6 * len(configs), 13.5), squeeze=False)
    for c, cfg in enumerate(configs):
        real_map, sim_map = maps[cfg]
        m = results[cfg]
        for r, (data, title) in enumerate((
            (real_map, cell_title("実機", cfg, m, "real")),
            (sim_map, cell_title("シミュ", cfg, m, "sim")),
        )):
            ax = axes[r, c]
            im = ax.imshow(data, cmap="hot", vmin=DISPLAY_RANGE[0], vmax=DISPLAY_RANGE[1])
            ax.set_title(title, fontsize=8.5)
            ax.axis("off")
            if c == len(configs) - 1:
                fig.colorbar(im, ax=ax, fraction=0.046, label="輝度 / ROI 平均")
        ax = axes[2, c]
        im = ax.imshow(sim_map - real_map, cmap="coolwarm", vmin=-DIFF_RANGE, vmax=DIFF_RANGE)
        ax.set_title(f"差分（シミュ − 実機）\nRMS {m['rms_diff']:.3f}・最大 "
                     f"{m['max_abs_diff']:.3f}\nSSIM {m['ssim']:.3f}・NCC {m['ncc']:.3f}",
                     fontsize=8.5)
        ax.axis("off")
        if c == len(configs) - 1:
            fig.colorbar(im, ax=ax, fraction=0.046)
    fig.suptitle(f"輝度マップの形の一致 — 実機 vs シミュ（{sim_label}）．"
                 "全セル同一処理（シミュ解像度へリサイズ・メディアン 5・ガウス σ=4・ROI 平均 = 1）",
                 fontsize=12)
    fig.tight_layout()
    output.parent.mkdir(parents=True, exist_ok=True)
    fig.savefig(output, dpi=110)
    plt.close(fig)
    print(f"出力: {output}")


def main() -> int:
    """4 構成の形の比較を図と表にする."""
    parser = argparse.ArgumentParser(description="実機 vs シミュの輝度マップの形の比較")
    parser.add_argument("--sim-set", choices=tuple(SIM_SETS), default="rho081")
    args = parser.parse_args()
    sim_set = SIM_SETS[args.sim_set]
    results: dict[str, dict] = {}
    maps: dict[str, tuple[np.ndarray, np.ndarray]] = {}
    for cfg in CONFIGS:
        sim_path, real_path = sim_set[cfg], REAL[cfg]
        if not sim_path.exists() or not real_path.exists():
            print(f"スキップ {cfg}: {sim_path if not sim_path.exists() else real_path} が無い")
            continue
        sim_roi = load_sim(sim_path)
        real_roi = load_real(real_path, sim_roi.shape)
        sim_map, real_map = normalized_shape(sim_roi), normalized_shape(real_roi)
        results[cfg] = shape_metrics(sim_map, real_map, sim_roi, real_roi)
        maps[cfg] = (real_map, sim_map)
    if not results:
        print("エラー: 比較できる構成がありません")
        return 1
    draw(results, maps, sim_set["label"], OUT_DIR / f"shapes_{args.sim_set}.png")
    with (OUT_DIR / f"shapes_{args.sim_set}.csv").open("w", newline="", encoding="utf-8") as fh:
        keys = ["config"] + list(next(iter(results.values())).keys())
        writer = csv.DictWriter(fh, fieldnames=keys)
        writer.writeheader()
        for cfg, m in results.items():
            writer.writerow({"config": cfg, **m})
    print(f"\n=== 形の一致（{sim_set['label']}）===")
    print(f"{'構成':>8} {'RMS差':>6} {'最大差':>6} {'SSIM':>6} {'NCC':>6} | "
          f"{'上下比 実/シ':>13} {'左右比 実/シ':>13} {'傾き% 実/シ':>13} {'方位 実/シ':>12} "
          f"{'お椀% 実/シ':>12} {'b1% 実/シ':>11}")
    for cfg, m in results.items():
        print(f"{CONFIG_LABEL[cfg]:>8} {m['rms_diff']:6.3f} {m['max_abs_diff']:6.3f} "
              f"{m['ssim']:6.3f} {m['ncc']:6.3f} | "
              f"{m['real_top_bottom']:6.3f}/{m['sim_top_bottom']:6.3f} "
              f"{m['real_left_right']:6.3f}/{m['sim_left_right']:6.3f} "
              f"{m['real_tilt']:6.1f}/{m['sim_tilt']:6.1f} "
              f"{m['real_tilt_az']:5.0f}/{m['sim_tilt_az']:5.0f} "
              f"{m['real_bowl']:+5.1f}/{m['sim_bowl']:+5.1f} "
              f"{m['real_b1'] * 100:+5.1f}/{m['sim_b1'] * 100:+5.1f}")
    return 0


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