"""後側筐体ロフト 0.5（従来）と 1.4 の白板照明を直接比較するスクリプト.

後側筐体のロフトパラメータ違い（従来=0.5 / 再印刷=1.4）について，
全灯点灯と単灯（下・左・右）の白板 DNG を筐体間で直接比較する．

出力（2 図）:
- 全灯比較図: 輝度マップ 2 枚（共通絶対カラースケール）＋ 比マップ（w1.4/w0.5）
  ＋ 動径プロファイル重ね描き ＋ 指標表
- 単灯比較図: 輝度マップ 2×3（行=筐体，列=下/左/右，共通絶対カラースケール）
  ＋ 動径プロファイル重ね描き（実線=w1.4，破線=w0.5）＋ 指標表

すべて linear + LSC 補正・ROI=smtias.whiteboard で揃えて算出する．

使い方:
    python scripts/compare_loft_housings.py
    python scripts/compare_loft_housings.py --w14-all-dng <dng> --suffix _shot1 --only all
"""

import argparse
import sys
from pathlib import Path

import matplotlib.pyplot as plt
import numpy as np
from scipy.ndimage import binary_dilation, median_filter

PROJECT_ROOT = Path(__file__).resolve().parent.parent
sys.path.insert(0, str(PROJECT_ROOT))

from src.analysis.spatial import (  # noqa: E402
    calc_distance_map,
    calc_radial_min_max_ratio,
    calc_radial_profile,
    calc_spatial_uniformity,
)
from src.analysis.uniformity import calc_uniformity  # noqa: E402
from src.config import load_roi_config  # noqa: E402
from src.io.dng_pipeline import dng_luminance  # noqa: E402
from src.io.loader import extract_roi  # noqa: E402

QUANT_DIR = PROJECT_ROOT / "data" / "smtias" / "quantitative"
ROI_CONFIG = PROJECT_ROOT / "config" / "roi_config.json"
ROI_KEY = "smtias.whiteboard"
FIGURES_DIR = PROJECT_ROOT / "output" / "smtias" / "figures"

# 比較対象（筐体ロフト別．全灯 1 枚＋単灯 3 枚）
HOUSINGS = {
    "Loft w=0.5 (original)": {
        "all": QUANT_DIR / "SmTIAS_QM_20260601_124500.dng",
        "singles": [
            ("Bottom LED", QUANT_DIR / "SmTIAS_QM_20260615_141420.dng"),
            ("Left LED", QUANT_DIR / "SmTIAS_QM_20260615_150300.dng"),
            ("Right LED", QUANT_DIR / "SmTIAS_QM_20260615_170948.dng"),
        ],
    },
    "Loft w=1.4 (reprinted)": {
        "all": QUANT_DIR / "SmTIAS_QM_20260727_124257.dng",
        "singles": [
            ("Bottom LED", QUANT_DIR / "SmTIAS_QM_20260727_134105.dng"),
            ("Left LED", QUANT_DIR / "SmTIAS_QM_20260727_134853.dng"),
            ("Right LED", QUANT_DIR / "SmTIAS_QM_20260727_140129.dng"),
        ],
    },
}
LED_COLORS = {"Bottom LED": "steelblue", "Left LED": "seagreen", "Right LED": "indianred"}

# 共通絶対カラースケールの設定（外れ値に強いパーセンタイルで vmin/vmax を決める）
SCALE_PERCENTILES = (1.0, 99.0)

# 暗点インペイントの設定
DESPOT_SMOOTH_KERNEL = 5  # 画素ノイズ抑制用の小メディアン
DESPOT_BG_KERNEL = 61  # 照明背景推定用の大メディアン（汚れより十分大きく取る）
DESPOT_RATIO = 0.96  # 背景比でこの割合より暗い領域を汚れとみなす
DESPOT_DILATE = 7  # 汚れの薄い縁まで含めるためのマスク膨張回数

METRIC_ROWS = [
    ("Mean luminance", "mean", "{:.2f}"),
    ("CoV", "cov", "{:.3f}"),
    ("Max/min ratio", "max_min_ratio", "{:.3f}"),
    ("Center/periphery", "center_periphery_ratio", "{:.3f}"),
    ("Gradient (%)", "gradient_magnitude", "{:.2f}"),
    ("Radial min/max", "radial_min_max_ratio", "{:.3f}"),
]


def _luminance(dng_path: Path, roi: dict) -> np.ndarray:
    """DNG を linear + LSC 補正で輝度化し ROI を切り出す（portrait）."""
    luma_full = dng_luminance(str(dng_path), apply_lsc=True)
    return extract_roi(luma_full, roi["x"], roi["y"], roi["width"], roi["height"])


def _inpaint_dark_spots(luminance: np.ndarray, label: str) -> np.ndarray:
    """白板の汚れ等による局所的な暗点を背景メディアン値で埋める.

    画素ノイズに埋もれないよう小メディアンで平滑化した上で，汚れより十分大きい
    大メディアン（照明の緩やかな勾配は追従して保存される）と比較し，
    DESPOT_RATIO 未満に暗い領域を汚れとみなす．薄い縁まで含めるため
    マスクを膨張させてから背景値で置換する．
    """
    smoothed = median_filter(luminance, size=DESPOT_SMOOTH_KERNEL)
    background = median_filter(luminance, size=DESPOT_BG_KERNEL)
    mask = smoothed < background * DESPOT_RATIO
    mask = binary_dilation(mask, iterations=DESPOT_DILATE)
    if mask.any():
        print(f"  暗点インペイント（{label}）: {int(mask.sum())} px "
              f"({mask.mean() * 100:.3f}%) を背景メディアンで置換")
    return np.where(mask, background, luminance)


def _metrics(luminance: np.ndarray) -> dict:
    """均一性・空間・動径の指標をまとめて算出する."""
    distance_map = calc_distance_map(*luminance.shape)
    profile = calc_radial_profile(luminance, distance_map)
    radial = calc_radial_min_max_ratio(profile)
    return {
        **calc_uniformity(luminance),
        **calc_spatial_uniformity(luminance)["zone_stats"],
        "radial_min_max_ratio": radial["radial_min_max_ratio"],
        "profile": profile,
    }


def _metrics_table(ax: plt.Axes, col_labels: list[str], metrics: list[dict]) -> None:
    """指標表を描画する."""
    ax.axis("off")
    rows = [
        [name] + [fmt.format(m[key]) for m in metrics]
        for name, key, fmt in METRIC_ROWS
    ]
    table = ax.table(
        cellText=rows,
        colLabels=["Metric"] + col_labels,
        cellLoc="center",
        loc="center",
    )
    table.auto_set_font_size(False)
    table.set_fontsize(8)
    table.scale(1.0, 1.6)
    table.auto_set_column_width(list(range(len(col_labels) + 1)))
    for j in range(len(col_labels) + 1):
        table[0, j].set_facecolor("#dfe6ee")
        table[0, j].set_text_props(weight="bold")


def plot_all_on(roi: dict, suffix: str = "", despot: bool = False) -> None:
    """全灯点灯のロフト 0.5 vs 1.4 比較図を生成して保存する."""
    labels = list(HOUSINGS.keys())
    lumas = [_luminance(HOUSINGS[label]["all"], roi) for label in labels]
    if despot:
        lumas = [_inpaint_dark_spots(luma, label) for label, luma in zip(labels, lumas)]
    metrics = [_metrics(luma) for luma in lumas]

    ratio = lumas[1] / np.maximum(lumas[0], 1e-6)
    ratio_stats = {
        "mean": float(np.mean(ratio)),
        "median": float(np.median(ratio)),
        "p5": float(np.percentile(ratio, 5)),
        "p95": float(np.percentile(ratio, 95)),
    }

    flat_all = np.concatenate([luma.ravel() for luma in lumas])
    vmin, vmax = np.percentile(flat_all, SCALE_PERCENTILES)

    fig = plt.figure(figsize=(15, 9))
    gs = fig.add_gridspec(2, 3, height_ratios=[1.0, 0.85], hspace=0.28, wspace=0.22)

    map_axes = []
    im = None
    for col, (label, luma) in enumerate(zip(labels, lumas)):
        ax = fig.add_subplot(gs[0, col])
        im = ax.imshow(luma, cmap="hot", interpolation="nearest", vmin=vmin, vmax=vmax)
        ax.set_title(f"{label}  (mean={metrics[col]['mean']:.1f})", fontsize=11)
        ax.set_xlabel("X (px)")
        if col == 0:
            ax.set_ylabel("Y (px)")
        map_axes.append(ax)
    cbar = fig.colorbar(im, ax=map_axes, fraction=0.020, pad=0.02)
    cbar.set_label(
        f"Luminance (common scale, p{SCALE_PERCENTILES[0]:g}-p{SCALE_PERCENTILES[1]:g})"
    )

    # 比マップ（w1.4 / w0.5．中央値中心の発散スケール）
    ax_ratio = fig.add_subplot(gs[0, 2])
    center = ratio_stats["median"]
    half_range = max(ratio_stats["p95"] - center, center - ratio_stats["p5"])
    im_r = ax_ratio.imshow(
        ratio, cmap="RdBu_r", interpolation="nearest",
        vmin=center - half_range, vmax=center + half_range,
    )
    ax_ratio.set_title(f"Ratio w1.4/w0.5  (median={ratio_stats['median']:.3f})", fontsize=11)
    ax_ratio.set_xlabel("X (px)")
    fig.colorbar(im_r, ax=ax_ratio, fraction=0.040, pad=0.02).set_label("Ratio")

    ax_prof = fig.add_subplot(gs[1, 0:2])
    for label, color, marker, m in zip(labels, ["gray", "darkorange"], ["o", "s"], metrics):
        profile = m["profile"]
        ax_prof.plot(
            profile[:, 0], profile[:, 1], marker=marker, ms=4, color=color,
            label=f"{label}  radial min/max={m['radial_min_max_ratio']:.3f}",
        )
    ax_prof.set_title("Radial profile (absolute luminance)", fontsize=11)
    ax_prof.set_xlabel("Normalized distance from center")
    ax_prof.set_ylabel("Mean luminance")
    ax_prof.grid(True, linestyle="--", alpha=0.5)
    ax_prof.legend(loc="lower left", fontsize=9)

    ax_tbl = fig.add_subplot(gs[1, 2])
    _metrics_table(ax_tbl, ["w=0.5", "w=1.4"], metrics)

    fig.suptitle("All-LEDs comparison: rear housing loft w=0.5 vs w=1.4", fontsize=14, y=0.98)

    output_path = FIGURES_DIR / f"compare_loft_housings_all{suffix}.png"
    output_path.parent.mkdir(parents=True, exist_ok=True)
    fig.savefig(str(output_path), dpi=150, bbox_inches="tight")
    plt.close(fig)

    print(f"全灯の筐体比較図を保存しました: {output_path}")
    for label, m in zip(labels, metrics):
        print(f"  {label}: mean={m['mean']:.2f}, CoV={m['cov']:.4f}")
    print(f"  Ratio w1.4/w0.5: mean={ratio_stats['mean']:.4f}, "
          f"median={ratio_stats['median']:.4f}, "
          f"p5={ratio_stats['p5']:.4f}, p95={ratio_stats['p95']:.4f}")


def plot_singles(roi: dict) -> None:
    """単灯 3 種のロフト 0.5 vs 1.4 比較図（2×3）を生成して保存する."""
    housing_labels = list(HOUSINGS.keys())
    led_labels = [led for led, _ in HOUSINGS[housing_labels[0]]["singles"]]
    lumas = {
        housing: [_luminance(path, roi) for _, path in HOUSINGS[housing]["singles"]]
        for housing in housing_labels
    }
    metrics = {
        housing: [_metrics(luma) for luma in lumas[housing]]
        for housing in housing_labels
    }

    flat_all = np.concatenate(
        [luma.ravel() for housing in housing_labels for luma in lumas[housing]]
    )
    vmin, vmax = np.percentile(flat_all, SCALE_PERCENTILES)

    fig = plt.figure(figsize=(15, 13))
    gs = fig.add_gridspec(
        3, 3, height_ratios=[1.0, 1.0, 0.8], hspace=0.30, wspace=0.22
    )

    map_axes = []
    im = None
    for row, housing in enumerate(housing_labels):
        for col, led in enumerate(led_labels):
            ax = fig.add_subplot(gs[row, col])
            im = ax.imshow(
                lumas[housing][col], cmap="hot", interpolation="nearest",
                vmin=vmin, vmax=vmax,
            )
            mean = metrics[housing][col]["mean"]
            ax.set_title(f"{led} / {housing}  (mean={mean:.1f})", fontsize=10)
            if row == 1:
                ax.set_xlabel("X (px)")
            if col == 0:
                ax.set_ylabel("Y (px)")
            map_axes.append(ax)
    cbar = fig.colorbar(im, ax=map_axes, fraction=0.020, pad=0.02)
    cbar.set_label(
        f"Luminance (common scale, p{SCALE_PERCENTILES[0]:g}-p{SCALE_PERCENTILES[1]:g})"
    )

    # 動径プロファイル重ね描き（色=LED，実線=w1.4 / 破線=w0.5）
    ax_prof = fig.add_subplot(gs[2, 0:2])
    for housing, style in zip(housing_labels, ["--", "-"]):
        for led, m in zip(led_labels, metrics[housing]):
            profile = m["profile"]
            ax_prof.plot(
                profile[:, 0], profile[:, 1], style, color=LED_COLORS[led],
                label=f"{led} ({housing.split()[1]})",
            )
    ax_prof.set_title("Radial profile (absolute luminance)", fontsize=11)
    ax_prof.set_xlabel("Normalized distance from center")
    ax_prof.set_ylabel("Mean luminance")
    ax_prof.grid(True, linestyle="--", alpha=0.5)
    ax_prof.legend(loc="upper right", fontsize=8, ncol=2)

    # 指標表（列=LED×筐体）
    ax_tbl = fig.add_subplot(gs[2, 2])
    col_labels = []
    tbl_metrics = []
    for led_idx, led in enumerate(led_labels):
        for housing in housing_labels:
            # 例: "B 0.5"（LED 頭文字＋ロフト値）
            col_labels.append(f"{led[0]} {housing.split()[1].removeprefix('w=')}")
            tbl_metrics.append(metrics[housing][led_idx])
    _metrics_table(ax_tbl, col_labels, tbl_metrics)

    fig.suptitle(
        "Single-LED comparison: rear housing loft w=0.5 vs w=1.4", fontsize=14, y=0.98
    )

    output_path = FIGURES_DIR / "compare_loft_housings_singles.png"
    output_path.parent.mkdir(parents=True, exist_ok=True)
    fig.savefig(str(output_path), dpi=150, bbox_inches="tight")
    plt.close(fig)

    print(f"単灯の筐体比較図を保存しました: {output_path}")
    for housing in housing_labels:
        for led, m in zip(led_labels, metrics[housing]):
            print(f"  {led} / {housing}: mean={m['mean']:.2f}, CoV={m['cov']:.4f}")


def main() -> None:
    """ロフト 0.5 vs 1.4 の全灯・単灯比較図を生成して保存する."""
    parser = argparse.ArgumentParser(
        description="後側筐体ロフト 0.5 vs 1.4 の全灯・単灯比較図を生成する."
    )
    parser.add_argument(
        "--w14-all-dng",
        default=None,
        help="w=1.4 側の全灯 DNG を差し替える（既定: 07-27 清掃後撮影）",
    )
    parser.add_argument(
        "--suffix",
        default="",
        help="出力ファイル名の末尾に付ける文字列（例: _shot1．既定図の上書きを避ける）",
    )
    parser.add_argument(
        "--only",
        choices=["all", "singles", "both"],
        default="both",
        help="生成する図の選択（既定: both）",
    )
    parser.add_argument(
        "--despot",
        action="store_true",
        help="全灯図で白板汚れ等の局所暗点を局所メディアンでインペイントする",
    )
    args = parser.parse_args()

    if args.w14_all_dng:
        HOUSINGS["Loft w=1.4 (reprinted)"]["all"] = Path(args.w14_all_dng)

    roi = load_roi_config(str(ROI_CONFIG), ROI_KEY)
    if roi is None:
        raise SystemExit(f"ROI '{ROI_KEY}' が config に見つかりません")
    if args.only in ("all", "both"):
        plot_all_on(roi, args.suffix, args.despot)
    if args.only in ("singles", "both"):
        plot_singles(roi)


if __name__ == "__main__":
    main()