"""バッフルが評価板への直接光を実際に塞いでいるかをレイキャストで直接検証する.

「指標が良くなったから効いている」という論法を取らないための検証スクリプト．
評価板上の多数の点から LED 発光面上の多数の点へレイを飛ばし，**LED に到達する前に
別の物体（バッフル）に遮られること**を Mitsuba の交差計算で実測する．

判定:
    board_baffle 構成 : 全レイが遮蔽されていること（blocked == total）が合格条件．
                        1 本でも通っていればバッフル寸法が不足している．
    axial_wall  構成  : LED が板と反対向きなので「到達はするが発光面の裏側」．
                        本スクリプトは遮蔽本数だけを報告し，合否判定は行わない．

実行方法:
    uv run python scripts/verify_baffle_occlusion.py \
        --scene configs/_tmp_sphere_radius/scene_board_baffle_r100.yaml
    uv run python scripts/verify_baffle_occlusion.py --all-radii
"""

from __future__ import annotations

import argparse
import shutil
import sys
from pathlib import Path

import numpy as np

sys.path.insert(0, str(Path(__file__).resolve().parent))
sys.path.insert(0, str(Path(__file__).resolve().parents[1] / "src"))

from sweep_sphere_radius import (  # noqa: E402
    BOARD_SURFACE_Z,
    CAVITIES,
    RADII,
    TMP_DIR,
    job_name,
)

from sm_tias_lightsim.runtime.llvm_init import ensure_drjit_libllvm_path  # noqa: E402
from sm_tias_lightsim.scene.builder import build_scene  # noqa: E402
from sm_tias_lightsim.scene.schema import AreaLight, load_scene_yaml  # noqa: E402

ROOT = Path(__file__).resolve().parents[1]
MM_TO_M = 0.001

# 評価板の実体半幅（configs/meshes/照明ムラ補正ボード.stl と一致）
BOARD_HALF = (49.5, 51.0)
# レイ起点を板面から少しキャビティ側へ浮かせる量（mm）．板自身への自己交差を避ける
SURFACE_OFFSET = 0.05
# レイが LED 発光面そのものに当たったと誤判定しないための許容量（m）
HIT_EPSILON = 1.0e-5
# LED 外周サンプルの位相オフセット（rad）．対称面上にサンプルを置かないための任意値
LED_RING_PHASE = 0.1234


def board_samples(grid: int) -> np.ndarray:
    """評価板のキャビティ側の面を格子状にサンプリングした点（mm）を返す.

    z は板面（BOARD_SURFACE_Z）から SURFACE_OFFSET だけキャビティ側へ浮かせる．
    面上に厳密に置くと自分自身の面に交差して「遮蔽された」と誤判定するため．
    四隅は必ず含める（見込み角が最大＝最悪条件になる）．
    """
    gx = np.linspace(-BOARD_HALF[0], BOARD_HALF[0], grid)
    gy = np.linspace(-BOARD_HALF[1], BOARD_HALF[1], grid)
    mesh_x, mesh_y = np.meshgrid(gx, gy)
    z = np.full(mesh_x.size, BOARD_SURFACE_Z - SURFACE_OFFSET)
    return np.stack([mesh_x.ravel(), mesh_y.ravel(), z], axis=1)


def led_samples(light: AreaLight, count: int) -> np.ndarray:
    """LED 円板上のサンプル点（mm）を返す（中心 + 外周リング．外周が最悪条件）.

    リングの位相を LED_RING_PHASE だけずらす．0 始まりだとサンプル点が XZ / YZ 対称面に
    乗り，メッシュのエッジ上を走るレイになって交差判定が数値的に不安定になるため．
    """
    center = np.asarray(light.center, dtype=float)
    radius = float(getattr(light, "radius", 0.0) or 0.0)
    angles = LED_RING_PHASE + np.linspace(0.0, 2.0 * np.pi, count, endpoint=False)
    ring = np.stack(
        [center[0] + radius * np.cos(angles), center[1] + radius * np.sin(angles),
         np.full(count, center[2])],
        axis=1,
    )
    return np.vstack([center[None, :], ring])


def check_occlusion(scene_yaml: Path, board_grid: int, led_count: int) -> dict[str, float]:
    """板 × LED の全組合せについて，LED 到達前に遮蔽されるレイの割合を測る.

    Returns:
        total / blocked / clearance_min（遮蔽点から LED までの最短余裕・mm）を含む dict．
    """
    ensure_drjit_libllvm_path()
    import mitsuba as mi

    config = load_scene_yaml(scene_yaml)
    if mi.variant() is None:
        mi.set_variant(config.render.variant)

    scene_dict, tempdir = build_scene(config, project_root=scene_yaml.parent)
    try:
        scene = mi.load_dict(scene_dict)
        lights = [light for light in config.lights if isinstance(light, AreaLight)]
        if not lights:
            msg = f"{scene_yaml}: AreaLight が見つかりません"
            raise ValueError(msg)

        origins_mm = board_samples(board_grid)
        targets_mm = np.vstack([led_samples(light, led_count) for light in lights])
        # 全組合せを 1 回の交差計算にまとめる（llvm_* で 1 本ずつ呼ぶと JIT が不安定）
        origins = np.repeat(origins_mm, len(targets_mm), axis=0) * MM_TO_M
        targets = np.tile(targets_mm, (len(origins_mm), 1)) * MM_TO_M
        delta = targets - origins
        distance = np.linalg.norm(delta, axis=1)
        direction = delta / distance[:, None]

        ray = mi.Ray3f(
            mi.Point3f(origins[:, 0], origins[:, 1], origins[:, 2]),
            mi.Vector3f(direction[:, 0], direction[:, 1], direction[:, 2]),
        )
        si = scene.ray_intersect(ray)
        hit_t = np.asarray(si.t).reshape(-1)
        valid = np.asarray(si.is_valid()).reshape(-1)
        # LED 発光面そのものへの命中は「遮蔽なし」．それより手前で何かに当たれば遮蔽
        blocked = valid & (hit_t < distance - HIT_EPSILON)
        clearance = (distance - hit_t)[blocked]
        return {
            "total": float(len(distance)),
            "blocked": float(blocked.sum()),
            "clearance_min_mm": float(clearance.min() / MM_TO_M) if blocked.any() else float("nan"),
        }
    finally:
        shutil.rmtree(tempdir, ignore_errors=True)


def main() -> int:
    """指定シーン（または全半径）のバッフル遮蔽を検証する."""
    parser = argparse.ArgumentParser(description="バッフルの遮蔽をレイキャストで検証する")
    parser.add_argument("--scene", type=Path, default=None, help="検証するシーン YAML")
    parser.add_argument("--all-radii", action="store_true",
                        help="_tmp_sphere_radius の board_baffle 全半径を検証する")
    parser.add_argument("--cavity", choices=CAVITIES, default="similar_g2",
                        help="--all-radii のときに検証するキャビティ形状")
    parser.add_argument("--board-grid", type=int, default=41,
                        help="評価板の 1 辺あたりサンプル数（既定 41 = 1681 点）")
    parser.add_argument("--led-samples", type=int, default=32,
                        help="LED 円板の外周サンプル数（既定 32）")
    args = parser.parse_args()

    if args.all_radii:
        scenes = [TMP_DIR / f"scene_board_baffle_{args.cavity}_{job_name(r)}.yaml"
                  for r in RADII]
    elif args.scene is not None:
        scenes = [args.scene]
    else:
        parser.error("--scene か --all-radii のいずれかを指定してください")

    failures = 0
    print(f"{'シーン':>44} {'レイ本数':>10} {'遮蔽':>10} {'判定':>8} {'最小余裕mm':>11}")
    for scene_yaml in scenes:
        if not scene_yaml.exists():
            print(f"{scene_yaml.name:>44}  (未生成のためスキップ)")
            continue
        result = check_occlusion(scene_yaml, args.board_grid, args.led_samples)
        ok = result["blocked"] == result["total"]
        failures += 0 if ok else 1
        print(f"{scene_yaml.name:>44} {result['total']:10.0f} {result['blocked']:10.0f} "
              f"{'OK' if ok else '不合格':>8} {result['clearance_min_mm']:11.2f}")

    if failures:
        print(f"\n!! {failures} 件で直接光が通っています．バッフル寸法を見直してください")
        return 1
    print("\n全シーンで評価板から LED への直線は完全に遮蔽されています")
    return 0


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