"""輪帯ごとの方位プロファイルから cos 4φ 成分を抽出し「四隅の痕跡」を定量する（TEST_16）.

立方体キャビティの壁の四角さが評価板に写り込むなら，対称性から必ず 90° 回転対称
（m=4）の方位成分として現れる．球は四回対称の壁を持たないため，球側の m=4 は
評価系（長方形 ROI・楕円ビン・MC ノイズ）が作る背景値を与える．
**立方体 − 球（同内面積）の m=4 が「立方体化が持ち込む四隅ムラ」の正味の量**になる．

設計上の要点:
    - 板は左右・上下に鏡映対称なので，四隅の構造は板軸に揃った **cos 4φ の符号付き
      係数 a4 だけ**に出る（a4 > 0 = 四隅が暗い，a4 < 0 = 辺の中央が暗い）．直交成分
      b4 は対称性からゼロのはずで，**b4 が生きたノイズ対照**になる（a4 と同じ統計で
      計算されるが真の信号を含まない）
    - 輪帯は評価パイプラインと同じ正規化楕円距離 20 ビン．ただし d が辺中点
      （d≈0.707）を超える輪帯は ROI の四隅方向にしか画素が無く，方位カバレッジの
      欠けが m=0（平均）を m=4 に漏れ込ませる偽信号を作るため，**全 72 セクタに
      画素がある輪帯だけ**を使う
    - セクタ平均を先に取ってからフーリエ係数を出す（セクタ間の画素数ムラで方位の
      重みが偏るのを防ぐ）．誤差はセクタ平均の標準誤差を伝播する

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

from __future__ import annotations

import argparse
import sys
from pathlib import Path

import matplotlib
import numpy as np

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_cube_sphere_comparison import (  # noqa: E402
    SPHERE_RUN_DIR,
    nearest_by_area,
    read_summary,
)
from sweep_cube_size import OUTPUT_BASE as CUBE_OUTPUT_BASE  # noqa: E402
from sweep_cube_size import RUN_DIR as CUBE_RUN_DIR  # noqa: E402
from sweep_cube_size import job_name as cube_job_name  # noqa: E402
from sweep_sphere_radius import OUTPUT_BASE as SPHERE_OUTPUT_BASE  # noqa: E402
from sweep_sphere_radius import job_name as sphere_job_name  # noqa: E402

from sm_tias_lightsim.evaluation.roi_presets import SIM_ROI, extract_roi_xywh  # 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"]

N_RINGS = 20  # 評価パイプラインの輪帯分割と同一
N_SECTORS = 72  # 方位セクタ数（m=4 の周期 90° あたり 18 セクタ）
MIN_PIXELS_PER_SECTOR = 5  # 全セクタがこれ以上の画素を持つ輪帯だけを「全方位カバー」とする
HARMONICS = (1, 2, 4)
# 集約に使う輪帯の距離範囲．内側は信号が原理的に小さく（どの方向も壁が遠い），
# 外側はカバレッジ限界（辺中点 d≈0.707）まで
AGGREGATE_RANGE = (0.35, 0.75)


def load_roi(luminance_npy: Path) -> np.ndarray:
    """luminance.npy を読み ROI（輝度 2 次元配列）を返す."""
    arr = np.load(luminance_npy).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 ring_harmonics(roi: np.ndarray) -> list[dict[str, float]]:
    """全方位カバーの各輪帯について方位フーリエ係数（ring mean 比）を返す.

    Returns:
        輪帯ごとの dict のリスト．キーは
        d_center / n_pixels と，各 m について am / bm / se_m
        （am, bm はセクタ平均のフーリエ係数を輪帯平均で割った相対値，
        se_m は両係数共通の 1σ 伝播誤差）．
    """
    height, width = roi.shape
    distance = compute_normalized_distance_map(height, width)
    ys = np.arange(height, dtype=np.float64).reshape(-1, 1) - height / 2.0
    xs = np.arange(width, dtype=np.float64).reshape(1, -1) - width / 2.0
    phi = np.arctan2(np.broadcast_to(ys, roi.shape), np.broadcast_to(xs, roi.shape))

    ring_edges = np.linspace(0.0, 1.0, N_RINGS + 1)
    ring_index = np.digitize(distance.ravel(), ring_edges, right=True)
    sector_index = ((phi.ravel() + np.pi) / (2.0 * np.pi) * N_SECTORS).astype(int) % N_SECTORS
    sector_phi = -np.pi + (np.arange(N_SECTORS) + 0.5) * (2.0 * np.pi / N_SECTORS)
    flat = roi.ravel()

    results: list[dict[str, float]] = []
    for ring in range(1, N_RINGS + 1):
        mask = ring_index == ring
        if not mask.any():
            continue
        values = flat[mask]
        sectors = sector_index[mask]
        counts = np.bincount(sectors, minlength=N_SECTORS)
        if counts.min() < MIN_PIXELS_PER_SECTOR:
            continue  # 方位カバレッジが欠けた輪帯は偽の m=4 を作るため除外
        sums = np.bincount(sectors, weights=values, minlength=N_SECTORS)
        sq_sums = np.bincount(sectors, weights=values**2, minlength=N_SECTORS)
        means = sums / counts
        # セクタ内標準偏差 → セクタ平均の標準誤差（MC ノイズは画素間独立とみなす）
        variances = np.maximum(sq_sums / counts - means**2, 0.0)
        sems = np.sqrt(variances / counts)
        ring_mean = float(np.mean(means))

        row: dict[str, float] = {
            "d_center": float((ring_edges[ring - 1] + ring_edges[ring]) / 2.0),
            "n_pixels": float(mask.sum()),
        }
        for m in HARMONICS:
            cos_m, sin_m = np.cos(m * sector_phi), np.sin(m * sector_phi)
            row[f"a{m}"] = float(2.0 / N_SECTORS * np.sum(means * cos_m) / ring_mean)
            row[f"b{m}"] = float(2.0 / N_SECTORS * np.sum(means * sin_m) / ring_mean)
            # Var(a) = (2/K)^2 Σ sem^2 cos^2 ≈ (2/K^2) Σ sem^2（cos^2 の平均 1/2）
            row[f"se_{m}"] = float(
                np.sqrt(4.0 / N_SECTORS**2 * np.sum(sems**2 * cos_m**2)) / ring_mean
            )
        results.append(row)
    return results


def aggregate(rings: list[dict[str, float]], key: str) -> tuple[float, float]:
    """AGGREGATE_RANGE の輪帯を逆分散重みで集約し (推定値, 標準誤差) を返す."""
    m = key.lstrip("ab")
    selected = [r for r in rings if AGGREGATE_RANGE[0] <= r["d_center"] <= AGGREGATE_RANGE[1]]
    weights = np.array([1.0 / r[f"se_{m}"] ** 2 for r in selected])
    values = np.array([r[key] for r in selected])
    return float(np.sum(weights * values) / np.sum(weights)), float(1.0 / np.sqrt(np.sum(weights)))


def analyze(label: str, luminance_npy: Path) -> dict[str, object]:
    """1 条件を解析して集約値と輪帯プロファイルを返す."""
    rings = ring_harmonics(load_roi(luminance_npy))
    result: dict[str, object] = {"label": label, "rings": rings}
    for key in ("a4", "b4", "a2", "a1", "b1"):
        result[key], result[f"{key}_se"] = aggregate(rings, key)
    return result


def print_report(cube_results: list[dict], sphere_results: list[dict]) -> None:
    """集約表と立方体−球の差を標準出力に書く."""
    print(f"\n=== 方位フーリエ成分（輪帯 d∈[{AGGREGATE_RANGE[0]}, {AGGREGATE_RANGE[1]}] "
          "の逆分散重み集約・輪帯平均比 %）===")
    print("a4 > 0 が「四隅が暗い」．b4 は対称性からゼロのはずのノイズ対照．")
    print(f"{'条件':>14} {'a4 [%]':>14} {'b4 [%]（対照）':>16} {'a2 [%]':>14} {'|m1| [%]':>10}")
    for res in cube_results + sphere_results:
        m1 = float(np.hypot(res["a1"], res["b1"]))
        print(f"{res['label']:>14} {res['a4'] * 100:7.3f}±{res['a4_se'] * 100:5.3f} "
              f"{res['b4'] * 100:8.3f}±{res['b4_se'] * 100:5.3f}   "
              f"{res['a2'] * 100:7.3f}±{res['a2_se'] * 100:5.3f} {m1 * 100:10.3f}")

    print("\n=== 立方体 − 球（同内面積ペア）の a4 差 ===")
    print(f"{'ペア':>22} {'Δa4 [%]':>14} {'有意性':>8}")
    for cube, sphere in zip(cube_results, sphere_results, strict=True):
        delta = cube["a4"] - sphere["a4"]
        se = float(np.hypot(cube["a4_se"], sphere["a4_se"]))
        sigma = abs(delta) / se
        print(f"{cube['label'] + ' vs ' + sphere['label']:>22} "
              f"{delta * 100:7.3f}±{se * 100:5.3f} {sigma:6.1f}σ")


def draw_figure(cube_results: list[dict], sphere_results: list[dict], output: Path) -> None:
    """a4 の輪帯プロファイルと集約値の比較図を描く."""
    fig, axes = plt.subplots(1, 3, figsize=(19, 5.5))

    colors = plt.cm.viridis(np.linspace(0.0, 0.9, len(cube_results)))
    for results, ax, kind in ((cube_results, axes[0], "立方体"), (sphere_results, axes[1], "球")):
        for res, color in zip(results, colors, strict=True):
            d = [r["d_center"] for r in res["rings"]]
            a4 = [r["a4"] * 100 for r in res["rings"]]
            se = [r["se_4"] * 100 for r in res["rings"]]
            ax.errorbar(d, a4, yerr=se, marker="o", ms=3, lw=1.4, color=color,
                        label=res["label"])
        ax.axhline(0.0, color="gray", lw=0.8, ls="--")
        ax.set_xlabel("正規化楕円距離（輪帯中心）")
        ax.set_ylabel("a4（輪帯平均比 %・正=四隅が暗い）")
        ax.set_title(f"{kind}: cos 4φ 係数の輪帯プロファイル（±1σ）", fontsize=11)
        ax.legend(fontsize=8)
        ax.grid(alpha=0.3)
    y_all = [abs(r["a4"] * 100) + r["se_4"] * 100
             for res in cube_results + sphere_results for r in res["rings"]]
    for ax in axes[:2]:
        ax.set_ylim(-max(y_all) * 1.05, max(y_all) * 1.05)

    ax = axes[2]
    x = np.arange(len(cube_results))
    for results, offset, color, label in (
        (cube_results, -0.12, "tab:red", "立方体 a4"),
        (sphere_results, 0.12, "tab:blue", "球 a4"),
    ):
        ax.errorbar(x + offset, [r["a4"] * 100 for r in results],
                    yerr=[r["a4_se"] * 100 for r in results],
                    fmt="s", capsize=4, color=color, label=label)
    ax.errorbar(x, [r["b4"] * 100 for r in cube_results],
                yerr=[r["b4_se"] * 100 for r in cube_results],
                fmt="x", capsize=3, color="gray", label="立方体 b4（ノイズ対照）")
    ax.axhline(0.0, color="gray", lw=0.8, ls="--")
    ax.set_xticks(x)
    ax.set_xticklabels([f"{c['label']}\nvs {s['label']}"
                        for c, s in zip(cube_results, sphere_results, strict=True)], fontsize=8)
    ax.set_ylabel("集約 a4（輪帯平均比 %）")
    ax.set_title(f"集約値の比較（d∈[{AGGREGATE_RANGE[0]}, {AGGREGATE_RANGE[1]}]・±1σ）",
                 fontsize=11)
    ax.legend(fontsize=8)
    ax.grid(alpha=0.3)

    fig.suptitle(
        "四隅の痕跡の定量: 輪帯ごとの方位プロファイルの cos 4φ 成分（板軸フレーム・符号付き）",
        fontsize=14,
    )
    fig.tight_layout()
    output.parent.mkdir(parents=True, exist_ok=True)
    fig.savefig(output, dpi=110)
    print(f"\n出力: {output}")


def main() -> int:
    """立方体全条件と同内面積の球について m=4 成分を抽出する."""
    parser = argparse.ArgumentParser(description="cos 4φ 成分の抽出（四隅の痕跡の定量）")
    parser.add_argument("--output", type=Path, default=None, help="出力 PNG パス")
    args = parser.parse_args()

    cube_run = CUBE_OUTPUT_BASE / CUBE_RUN_DIR
    sphere_run = SPHERE_OUTPUT_BASE / SPHERE_RUN_DIR
    cube_rows = read_summary(cube_run / "summary.csv", "size_mm")
    sphere_rows = read_summary(sphere_run / "summary.csv", "radius_mm")

    cube_results, sphere_results = [], []
    for row in cube_rows:
        lum = cube_run / cube_job_name(row["size_mm"]) / "luminance.npy"
        if not lum.exists():
            continue
        cube_results.append(analyze(f"L={row['size_mm']:g}", lum))
        sphere = nearest_by_area(row["cavity_area_mm2"], sphere_rows)
        sphere_results.append(analyze(
            f"r={sphere['radius_mm']:g}",
            sphere_run / sphere_job_name(sphere["radius_mm"]) / "luminance.npy",
        ))

    print_report(cube_results, sphere_results)
    draw_figure(cube_results, sphere_results,
                args.output or cube_run / "fourfold_component.png")
    return 0


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