"""半径スイープの幾何確認図（軸断面 + 3D スクリーンショット一覧）を生成する.

CLAUDE.md の「幾何を変えるシーンは本番レンダの前に必ず 3D で配置確認する」に対応する
確認用の図．sweep_sphere_radius.py が生成したシーン YAML と同じ式で軸断面を描き，
評価板・ロフト・球・LED・一次照射キャップの位置関係を半径ごとに重ね描きする．

3D スクリーンショットは事前に次で生成しておく（無い半径は自動でスキップ）:
    uv run sm-visualize --scene configs/_tmp_sphere_radius/scene_r<R>.yaml \
        --no-bake --screenshot <shots_dir>/r<R>.png

実行方法:
    uv run python scripts/build_radius_geometry_figure.py --shots-dir <dir>
"""

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))

from generate_sphere_loft_housing import _blend_radius  # noqa: E402
from sweep_sphere_radius import (  # noqa: E402
    BAFFLE_GAP,
    BOARD_HALF,
    BOARD_SURFACE_Z,
    CAVITIES,
    RADII,
    RIM_INSET,
    RIM_Z,
    VARIANTS,
    baffle_radius,
    center_z,
    cut_radius,
    cut_z,
    job_name,
    led_z,
    loft_blend,
)

ROOT = Path(__file__).resolve().parents[1]
DEFAULT_OUT = ROOT / "output" / "sphere_radius_sweep" / "geometry_check.png"

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


def loft_profile(radius: float, cavity: str, n: int = 200) -> tuple[np.ndarray, np.ndarray]:
    """ロフト部（リム矩形 → 球断面円）の X-Z 断面プロファイルを返す.

    メッシュ生成本体（generate_sphere_loft_housing）と同じブレンド関数を
    theta=0（+X 方向）について評価するので，図と STL が食い違わない．
    """
    rect_r = np.asarray(BOARD_HALF[0] - RIM_INSET)
    z_cut, r_cut = cut_z(radius, cavity), cut_radius(radius, cavity)
    span = z_cut - RIM_Z
    m1 = (-(z_cut - center_z(radius, cavity)) / r_cut) * span
    m2 = (-(radius**2) / r_cut**3) * span**2
    u = np.linspace(0.0, 1.0, n)
    rho = np.array([
        float(_blend_radius(float(t), rect_r, r_cut, m1, m2, loft_blend(cavity))) for t in u
    ])
    return rho, RIM_Z + u * span


def sphere_profile(radius: float, cavity: str, n: int = 240) -> tuple[np.ndarray, np.ndarray]:
    """球お椀部（接続面より下）の X-Z 断面プロファイルを返す."""
    cz = center_z(radius, cavity)
    phi = np.linspace(float(np.arccos((cut_z(radius, cavity) - cz) / radius)), np.pi, n)
    return radius * np.sin(phi), cz + radius * np.cos(phi)


def draw_cross_sections(ax, radii: list[float], variant: str, cavity: str) -> None:
    """全半径の軸断面を重ね描きする（LED の向きとバッフルは構成に応じて描き分ける）."""
    colors = plt.cm.viridis(np.linspace(0.0, 0.9, len(radii)))
    to_board = variant == "board_baffle"
    for radius, color in zip(radii, colors, strict=True):
        lx, lz = loft_profile(radius, cavity)
        sx, sz = sphere_profile(radius, cavity)
        x = np.concatenate([lx, sx])
        z = np.concatenate([lz, sz])
        ax.plot(x, z, color=color, lw=1.4, label=f"r={radius:g}")
        ax.plot(-x, z, color=color, lw=1.4)
        zl = led_z(radius, variant, cavity)
        ax.plot(0.0, zl, marker="^" if to_board else "v", color=color, ms=7)
        if to_board:
            # バッフル円板（実寸）．評価板への直接光を完全に塞ぐ最小径
            rb = baffle_radius(radius, variant, cavity)
            ax.plot([-rb, rb], [zl + BAFFLE_GAP, zl + BAFFLE_GAP], color="black", lw=2.5)
            # 板の四隅へ向かう直接光の視線（バッフルで切られることを示す）
            for sign in (-1.0, 1.0):
                ax.plot([0.0, sign * np.hypot(*BOARD_HALF)], [zl, BOARD_SURFACE_Z],
                        color=color, lw=0.5, ls=":", alpha=0.6)
        else:
            cap_r = float(np.sqrt(radius**2 - (zl - center_z(radius, cavity)) ** 2))
            ax.plot([-cap_r, cap_r], [zl, zl], color=color, lw=0.7, ls=":")
    ax.plot([-BOARD_HALF[0], BOARD_HALF[0]], [BOARD_SURFACE_Z] * 2, color="crimson", lw=3.0)
    ax.text(0, 8, "評価板 99×102mm（不変）", color="crimson", ha="center", fontsize=9)
    ax.plot([cut_radius(r, cavity) for r in radii], [cut_z(r, cavity) for r in radii],
            color="black", lw=1.0, ls=":", marker="x", ms=4, label="接合点")
    ax.set_aspect("equal")
    ax.set_xlabel("X (mm)")
    ax.set_ylabel("Z (mm)")
    title = ("軸断面（▲=LED 評価板向き，黒線=バッフル，点線=遮断される直接光）" if to_board
             else "軸断面（▼=LED 壁向き，点線=一次照射キャップの端）")
    ax.set_title(title, fontsize=11)
    ax.legend(fontsize=7, loc="lower right", ncol=2)
    ax.grid(alpha=0.3)


def draw_shots(fig, gs, shots_dir: Path, radii: list[float]) -> None:
    """3D スクリーンショットをタイル表示する（存在するものだけ）."""
    available = [r for r in radii if (shots_dir / f"{job_name(r)}.png").exists()]
    for i, radius in enumerate(available[:4]):
        ax = fig.add_subplot(gs[i // 2, 1 + i % 2])
        ax.imshow(plt.imread(shots_dir / f"{job_name(radius)}.png"))
        ax.set_title(f"r={radius:g}mm", fontsize=10)
        ax.axis("off")


def main() -> int:
    """幾何確認図を生成して保存する."""
    parser = argparse.ArgumentParser(description="半径スイープの幾何確認図を生成")
    parser.add_argument("--shots-dir", type=Path, required=True,
                        help="sm-visualize が出力した r<R>.png を置いたディレクトリ")
    parser.add_argument("--output", type=Path, default=None, help="出力 PNG パス")
    parser.add_argument("--variant", choices=VARIANTS, default="board_baffle",
                        help="描画する LED 構成")
    parser.add_argument("--cavity", choices=CAVITIES, default="similar_g2",
                        help="描画するキャビティ形状")
    args = parser.parse_args()

    output = args.output or DEFAULT_OUT.with_name(
        f"geometry_check_{args.variant}_{args.cavity}.png")
    fig = plt.figure(figsize=(16, 9))
    gs = fig.add_gridspec(2, 3, width_ratios=[1.35, 1, 1])
    ax = fig.add_subplot(gs[:, 0])
    draw_cross_sections(ax, RADII, args.variant, args.cavity)
    draw_shots(fig, gs, args.shots_dir, [88.0, 100.0, 170.0, 250.0])
    fig.suptitle(
        f"球体筐体 半径スイープ 幾何確認 [{args.variant} / {args.cavity}]", fontsize=13
    )
    fig.tight_layout()
    output.parent.mkdir(parents=True, exist_ok=True)
    fig.savefig(output, dpi=110)
    print(f"出力: {output}")
    return 0


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