"""画像全体で評価する，ノイズに強い最小/最大比の定義を軸上スイープで検証する.

輪帯平均の min/max は方位方向に平均するので同心円状でない分布（TEST_18）を正しく評価できない．
一方，画素や平滑化後のパーセンタイルは MC ノイズの裾を拾って大きいキャビティで頭打ちになる
（build_minmax_definition_comparison.py）．本スクリプトは「画像全体を使い，かつノイズに
引きずられない」定義を 3 つ用意し，対称な分布（軸上スイープ）では輪帯平均と一致することを
確かめる（対称ならどの定義も同じ値になるはずなので，輪帯平均が基準になる）．

定義:
    block_<px>   : ROI を約 px 角のブロックに分けて平均し，その min / max．
                   ノイズは 1/sqrt(ブロック画素数) に落ちる．方位には依らない
    smooth_corr  : σ=12 で平滑化した min / max を，ノイズ場の極値で補正する．
                   ノイズ場 = 残差（生 − σ=4 平滑化）を空間シャッフルして同じ σ=12 平滑化を
                   掛けたもの（構造を壊しノイズの統計だけ残す．decompose_cov と同じ発想）．
                   min には +E[ノイズ場の最小]，max には −E[ノイズ場の最大] を掛ける
    poly<n>      : n 次の 2 次元多項式を最小二乗で当てはめ，フィット面の min / max
    ring         : 輪帯平均（基準．同心円状でのみ正しい）

各定義について「ノイズ床」もシャッフル法で見積もる（一様な像に同じノイズを載せたときに
定義が返す min/max 比．1 からのずれが小さいほどノイズに強い）．

出力（output/_normalization_comparison/）:
    fullimage_minmax.csv / fullimage_minmax.png

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

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

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_normalization_comparison import OUT_DIR, conditions  # noqa: E402
from build_radial_comparison import MAX_DISTANCE  # 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"]

BLOCK_SIZES = (20, 40, 80)
SMOOTH_SIGMA = 12.0
RESIDUAL_SIGMA = 4.0
POLY_DEGREES = (2, 4)
N_SHUFFLE = 4
SEED = 0

DEFINITIONS = ["ring"] + [f"block_{b}" for b in BLOCK_SIZES] + ["smooth_raw", "smooth_corr"] \
    + [f"poly{d}" for d in POLY_DEGREES]
DEF_LABEL = {
    "ring": "輪帯平均（基準）",
    "smooth_raw": f"平滑化 σ={SMOOTH_SIGMA:g} 補正なし",
    "smooth_corr": f"平滑化 σ={SMOOTH_SIGMA:g} ノイズ補正",
    **{f"block_{b}": f"ブロック平均 {b}px" for b in BLOCK_SIZES},
    **{f"poly{d}": f"多項式 {d} 次フィット" for d in POLY_DEGREES},
}
DEF_STYLE = {
    "ring": dict(marker="o", ls="-", color="black"),
    "block_20": dict(marker="s", ls="--", color="tab:blue", alpha=0.6),
    "block_40": dict(marker="s", ls="-", color="tab:blue"),
    "block_80": dict(marker="s", ls=":", color="tab:blue", alpha=0.6),
    "smooth_raw": dict(marker="D", ls=":", color="tab:red", alpha=0.6),
    "smooth_corr": dict(marker="D", ls="-", color="tab:red"),
    "poly2": dict(marker="^", ls=":", color="tab:green", alpha=0.6),
    "poly4": dict(marker="^", ls="-", color="tab:green"),
}


def load_roi(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 block_means(roi: np.ndarray, block: int) -> np.ndarray:
    """ROI をおよそ block 画素角のブロックに等分して平均した 2 次元配列を返す（端は切り捨て）."""
    h, w = roi.shape
    nh, nw = max(h // block, 1), max(w // block, 1)
    bh, bw = h // nh, w // nw
    trimmed = roi[: nh * bh, : nw * bw]
    return trimmed.reshape(nh, bh, nw, bw).mean(axis=(1, 3))


def poly_fit(roi: np.ndarray, degree: int) -> np.ndarray:
    """2 次元多項式（全次数 ≤ degree）を最小二乗で当てはめたフィット面を返す."""
    h, w = roi.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)
    columns = [x**i * y**j for i in range(degree + 1) for j in range(degree + 1 - i)]
    design = np.stack([c.ravel() for c in columns], axis=1)
    coeffs, *_ = np.linalg.lstsq(design, roi.ravel(), rcond=None)
    return (design @ coeffs).reshape(h, w)


def noise_field_extremes(roi: np.ndarray, rng: np.random.Generator) -> tuple[float, float]:
    """ノイズだけの場に σ=SMOOTH_SIGMA の平滑化を掛けたときの (min, max) の期待値を返す."""
    residual = roi - gaussian_filter(median_filter(roi, size=5), sigma=RESIDUAL_SIGMA)
    mins, maxs = [], []
    for _ in range(N_SHUFFLE):
        flat = residual.ravel().copy()
        rng.shuffle(flat)
        field = gaussian_filter(flat.reshape(roi.shape), sigma=SMOOTH_SIGMA)
        mins.append(field.min())
        maxs.append(field.max())
    return float(np.mean(mins)), float(np.mean(maxs))


def min_max_all(roi: np.ndarray, rng: np.random.Generator) -> dict[str, float]:
    """全定義の最小/最大比を返す."""
    result: dict[str, float] = {}
    profile = calc_radial_profile(roi, n_bins=20)
    valid = (~np.isnan(profile[:, 1])) & (profile[:, 0] <= MAX_DISTANCE)
    result["ring"] = float(profile[valid, 1].min() / profile[valid, 1].max())
    for b in BLOCK_SIZES:
        means = block_means(roi, b)
        result[f"block_{b}"] = float(means.min() / means.max())
    smooth = gaussian_filter(roi, sigma=SMOOTH_SIGMA)
    result["smooth_raw"] = float(smooth.min() / smooth.max())
    n_min, n_max = noise_field_extremes(roi, rng)
    result["smooth_corr"] = float((smooth.min() - n_min) / (smooth.max() - n_max))
    for d in POLY_DEGREES:
        fit = poly_fit(roi, d)
        result[f"poly{d}"] = float(fit.min() / fit.max())
    return result


def noise_floor(roi: np.ndarray, rng: np.random.Generator) -> dict[str, float]:
    """一様な像に同じノイズを載せたときの各定義の min/max 比（1 からのずれ = ノイズの影響）."""
    residual = roi - gaussian_filter(median_filter(roi, size=5), sigma=RESIDUAL_SIGMA)
    flat = residual.ravel().copy()
    rng.shuffle(flat)
    uniform = roi.mean() + flat.reshape(roi.shape)
    floor = min_max_all(uniform, rng)
    return {f"floor_{k}": v for k, v in floor.items()}


def draw(rows: list[dict], output: Path) -> None:
    """定義ごとの min/max 比とノイズ床をサイズに対して描く."""
    fig, axes = plt.subplots(2, 2, figsize=(15, 10))
    for col, (shape, xlabel) in enumerate((("sphere", "球 内径 φ [mm]"),
                                           ("cube", "立方体 内寸 L [mm]"))):
        sel = [r for r in rows if r["shape"] == shape]
        scale = 2.0 if shape == "sphere" else 1.0  # 球は内径で描く
        sizes = [r["size_mm"] * scale for r in sel]
        ax = axes[0, col]
        for d in DEFINITIONS:
            ax.plot(sizes, [r[d] for r in sel], label=DEF_LABEL[d], ms=5, lw=1.4, **DEF_STYLE[d])
        ax.set_xlabel(xlabel)
        ax.set_ylabel("最小/最大比")
        ax.set_title(("球 similar_radial（TEST_15）" if shape == "sphere" else "立方体（TEST_16）")
                     + "：各定義の値", fontsize=11)
        ax.grid(alpha=0.3)
        ax.legend(fontsize=7, ncol=2)
        ax = axes[1, col]
        for d in DEFINITIONS:
            ax.plot(sizes, [1.0 - r[f"floor_{d}"] for r in sel], label=DEF_LABEL[d], ms=5,
                    lw=1.4, **DEF_STYLE[d])
        ax.set_yscale("log")
        ax.set_xlabel(xlabel)
        ax.set_ylabel("ノイズ床 = 1 − (一様像 + ノイズ の min/max)（対数）")
        ax.set_title("ノイズだけで生じる見かけの不均一（小さいほどノイズに強い）", fontsize=11)
        ax.grid(alpha=0.3, which="both")
        ax.legend(fontsize=7, ncol=2)
    fig.suptitle("画像全体で評価する最小/最大比の定義 — 軸上スイープ（対称な分布）での検証．"
                 "対称なら輪帯平均（黒）と一致するのが正しい", fontsize=13)
    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:
    """全定義の min/max 比とノイズ床を計算し，表と図を出す."""
    parser = argparse.ArgumentParser(description="画像全体で評価する min/max 比の定義を検証")
    parser.parse_args()
    rng = np.random.default_rng(SEED)
    rows = []
    for shape, items in conditions().items():
        for label, path in items:
            roi = load_roi(path)
            size = float(label[1:]) / 2.0 if label.startswith("φ") else float(label.split("=")[1])
            row = {"shape": shape, "size_mm": size,
                   "raw_mean": float(roi.mean())}
            row.update(min_max_all(roi, rng))
            row.update(noise_floor(roi, rng))
            rows.append(row)
    if not rows:
        print("エラー: TEST_15/16 の結果が見つかりません")
        return 1
    OUT_DIR.mkdir(parents=True, exist_ok=True)
    with (OUT_DIR / "fullimage_minmax.csv").open("w", newline="", encoding="utf-8") as fh:
        writer = csv.DictWriter(fh, fieldnames=list(rows[0].keys()))
        writer.writeheader()
        writer.writerows(rows)
    draw(rows, OUT_DIR / "fullimage_minmax.png")

    print("\n=== 最小/最大比（上段）とノイズ床 1−floor（下段）===")
    print(f"{'形状':>7} {'サイズ':>6} " + " ".join(f"{d:>11}" for d in DEFINITIONS))
    for r in rows:
        print(f"{r['shape']:>7} {r['size_mm']:6g} "
              + " ".join(f"{r[d]:11.4f}" for d in DEFINITIONS))
        print(f"{'':>7} {'床':>6} "
              + " ".join(f"{1.0 - r['floor_' + d]:11.4f}" for d in DEFINITIONS))
    return 0


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