# pyright: reportCallIssue=none
# Mitsuba 3 は variant ベースの動的属性解決を行うため Pylance/Pyright が
# mi.* を object 型と推論し reportCallIssue を誤検知する（GUIDE_06 の方針）．
"""各メッシュ face 中心での照度ベイクモジュール．

2 つのモードを提供する：

- mode="direct": shadow ray + コサイン項 + 距離減衰による直接照度（W/m²）．
  多重反射を無視する．高速（face ×光源数の積分）．

- mode="path": コサイン重点サンプリング + Mitsuba path tracer による多重反射込み照度（W/m²）．
  E = ∫_hemisphere L_i(ω)·cosθ dω を Monte Carlo 積分で評価する．
  コサイン重点サンプリングの場合 E ≈ (π/N)·Σ L_i（PDF=cosθ/π のため cosθ キャンセル）．
  完全シーン（emitter + BSDF を含む）を使うため多重反射経路もすべて積分される．
  scalar_rgb variant を使用（LLVM 不要）．per-ray スカラ評価のため遅いが正確に完走する．
  計算コストは face 数 × samples 倍（直接より大幅に重い）．

{mesh_id: irradiance_per_face (N_f,) float32} の形式で返す．
"""

from __future__ import annotations

import logging
import shutil
from pathlib import Path
from typing import Literal

import numpy as np

from sm_tias_lightsim.geometry.loader import load_stl
from sm_tias_lightsim.geometry.transforms import scale_vertices
from sm_tias_lightsim.runtime.llvm_init import ensure_drjit_libllvm_path, log_variant_status
from sm_tias_lightsim.scene.schema import (
    AreaLight,
    DirectionalLight,
    PointLight,
    SceneConfig,
)

_logger = logging.getLogger(__name__)

# YAML 座標 (mm) → m 変換係数
_MM_TO_M = 0.001

# shadow ray の原点オフセット（face 法線方向に少しずらして自己交差を防ぐ）
_SHADOW_RAY_EPSILON = 1e-4

# DirectionalLight の shadow ray の最大距離（無限遠をシミュレート）
_DIRECTIONAL_MAXT = 1e9

# 方向ベクトルの長さがこれ以下なら縮退とみなしてスキップする
_DIR_NORM_EPSILON = 1e-9


def _to_scalar_intensity(value: float | tuple[float, float, float]) -> float:
    """スカラまたは RGB タプルを輝度スカラに変換する．

    RGB の場合は Rec.709 係数で輝度に変換する．
    スカラの場合はそのまま返す．

    Args:
        value: スカラまたは RGB タプル．

    Returns:
        輝度スカラ値（float）．
    """
    if isinstance(value, (int, float)):
        return float(value)
    r, g, b = value
    return 0.2126 * r + 0.7152 * g + 0.0722 * b


def _shadow_ray_blocked(
    scene: object,
    origin: np.ndarray,
    direction: np.ndarray,
    maxt: float,
) -> bool:
    """shadow ray が遮蔽されているかを判定する．

    origin → direction 方向に距離 maxt まで ray を飛ばし，
    途中に交差があれば True（遮蔽あり）を返す．

    Args:
        scene: Mitsuba の Scene オブジェクト（load_dict 済み）．
        origin: ray の起点（m 単位）．
        direction: ray の方向（単位ベクトル）．
        maxt: 最大ヒット距離（m 単位）．

    Returns:
        遮蔽あり（途中に交差あり）なら True．
    """
    import mitsuba as mi

    o = mi.Point3f(float(origin[0]), float(origin[1]), float(origin[2]))
    d = mi.Vector3f(float(direction[0]), float(direction[1]), float(direction[2]))
    ray = mi.Ray3f(o, d)  # type: ignore[call-arg]
    ray.maxt = maxt  # type: ignore[assignment]
    si = scene.ray_intersect_preliminary(ray)  # type: ignore[attr-defined]
    return bool(si.is_valid())


def _compute_face_centroids_and_normals(
    vertices: np.ndarray, faces: np.ndarray
) -> tuple[np.ndarray, np.ndarray]:
    """face の重心と法線ベクトルを計算する．

    Args:
        vertices: 頂点配列 (N_v, 3)，m 単位．
        faces: 面配列 (N_f, 3)，頂点インデックス．

    Returns:
        (centroids, normals) のタプル．
        - centroids: face 重心 (N_f, 3)．
        - normals: 単位法線ベクトル (N_f, 3)．長さ 0 の面はゼロベクトルになる．
    """
    v0 = vertices[faces[:, 0]]
    v1 = vertices[faces[:, 1]]
    v2 = vertices[faces[:, 2]]

    centroids = (v0 + v1 + v2) / 3.0

    edge1 = v1 - v0
    edge2 = v2 - v0
    raw_normals = np.cross(edge1, edge2)  # (N_f, 3)

    norms = np.linalg.norm(raw_normals, axis=1, keepdims=True)  # (N_f, 1)
    norms = np.where(norms == 0.0, 1.0, norms)
    normals = raw_normals / norms

    return centroids, normals


def _accumulate_finite_light(
    irradiance: np.ndarray,
    centroids: np.ndarray,
    normals: np.ndarray,
    scene: object,
    light_pos_m: np.ndarray,
    intensity: float,
    emit_dir: np.ndarray | None = None,
) -> None:
    """距離減衰のある光源（PointLight / AreaLight）の照度を face ごとに累積する．

    PointLight と AreaLight は「光源位置 → face 距離」「cos_theta」「shadow ray」の
    計算構造が共通のため，このヘルパーに集約する．
    AreaLight に固有の「裏面スキップ」は emit_dir が指定されたときのみ評価する．

    Args:
        irradiance: 累積先の照度配列 (N_f,)，in-place で更新する．
        centroids: face 重心 (N_f, 3)．
        normals: face 単位法線 (N_f, 3)．
        scene: Mitsuba Scene オブジェクト．
        light_pos_m: 光源位置（m 単位）．
        intensity: 光源の等価スカラ強度（W/sr 相当）．
        emit_dir: AreaLight の放射方向単位ベクトル（正規化済み）．
            None のとき裏面チェックをスキップする（PointLight 用）．
    """
    for fi in range(len(centroids)):
        centroid = centroids[fi]
        normal = normals[fi]

        to_light = light_pos_m - centroid
        distance = float(np.linalg.norm(to_light))
        if distance < _SHADOW_RAY_EPSILON:
            continue

        dir_n = to_light / distance
        cos_theta = float(np.dot(normal, dir_n))
        if cos_theta <= 0.0:
            continue

        # AreaLight の裏面には放射しない（放射方向との角度チェック）
        # emit_dir が face から光源への方向と逆向きの場合はスキップ
        if emit_dir is not None:
            emit_cos = float(np.dot(-dir_n, emit_dir))
            if emit_cos <= 0.0:
                continue

        # shadow ray: origin を法線方向にわずかにオフセットし，
        # maxt は光源手前までに制限して自己遮蔽と光源貫通を防ぐ
        origin = centroid + _SHADOW_RAY_EPSILON * normal
        if _shadow_ray_blocked(scene, origin, dir_n, distance - _SHADOW_RAY_EPSILON):
            continue

        irradiance[fi] += float(intensity * cos_theta / (distance**2))


def bake_irradiance(
    scene_config: SceneConfig,
    project_root: Path,
    *,
    mode: Literal["direct", "path"] = "direct",
    samples: int = 64,
    exclude_mesh_ids: set[str] | None = None,
    max_depth: int | None = None,
) -> dict[str, np.ndarray]:
    """各メッシュの face 中心で照度（W/m²）を計算する．

    2 つのモードを提供する：

    - mode="direct"（既定）: shadow ray + コサイン項 + 距離減衰による直接照度のみ．
      多重反射を無視する．高速．scalar_rgb variant を使用．

    - mode="path": コサイン重点サンプリング + Mitsuba path tracer による多重反射込み照度．
      完全シーン（emitter + BSDF を含む）に対してパス積分を評価するため，
      内壁の多重反射で均一化された照度分布が得られる（MiniTIAS の設計目標と整合）．
      **評価板は直接照度ベイクでは真っ暗になるが，path ベイクでは多重反射で非ゼロになる．**
      scalar_rgb variant を使用（LLVM 不要・DRJIT_LIBLLVM_PATH 不要）．
      計算コストは face 数 × samples 倍（直接ベイクより大幅に重い）．
      目安: 100k face × samples=64 ≈ 数分（face 数依存）．

    Args:
        scene_config: 検証済み SceneConfig．
        project_root: STL 相対パスの基点ディレクトリ．
        mode: 計算モード．"direct"（直接照度, 後方互換）または "path"（多重反射込み）．
        samples: mode="path" 時の face あたりサンプル数．大きいほど低ノイズだが重い．
            既定 64．高速確認は 16 程度，本番は 256 以上推奨．mode="direct" では無視．
        exclude_mesh_ids: ベイク計算からスキップする mesh_id の集合（任意）．
            ハイブリッドモードで評価板を投影で着色する場合に path ベイクから除外して
            計算量を節約するために使う．除外された mesh_id はキーとして返さない．
        max_depth: mode="path" 時の path integrator の最大反射回数（任意）．
            None のとき scene_config.render.max_depth を使用（後方互換）．
            指定時はその値を使う（render 設定から独立して制御できる）．
            mode="direct" では無視される（direct ベイクは path integrator を使わない）．
            SPEC_03: depth≥12 推奨（depth=8 は後段多重反射を取りこぼし勾配を過大評価する）．

    Returns:
        {mesh_id: irradiance_per_face (N_f,) float32} の辞書．
        irradiance の単位は W/m²（light の radiance / intensity の単位に依存）．
        exclude_mesh_ids に指定した id はキーに含まれない．

    Raises:
        ValueError: mode が不正な値の場合．
        FileNotFoundError: STL ファイルが見つからない場合．
    """
    if mode not in ("direct", "path"):
        msg = f"mode={mode!r} は不正です．'direct' または 'path' を指定してください．"
        raise ValueError(msg)

    if exclude_mesh_ids is None:
        exclude_mesh_ids = set()

    import mitsuba as mi

    from sm_tias_lightsim.scene.builder import build_scene

    if mode == "direct":
        # Mitsuba variant を初期化（scalar_rgb を使う：shadow ray 評価に最適）
        # scalar_rgb は LLVM 不要なので ensure_drjit_libllvm_path() は呼ばない
        # runner.py と同じ variant が既に設定されている場合はそのまま続行する
        desired_variant = "scalar_rgb"
        _variant_set_succeeded = False
        try:
            mi.set_variant(desired_variant)
            _variant_set_succeeded = True
        except Exception:
            _logger.warning(
                "Mitsuba variant の切り替えに失敗しました．現在の variant で続行します．"
            )
        if _variant_set_succeeded:
            log_variant_status(_logger, desired_variant, mi.variant())

        # shadow ray 評価専用シーン（AreaLight の rectangle 形状を除外）を構築する．
        # AreaLight を含む通常シーンでは，面光源自身の rectangle にヒットして
        # 正面メッシュへの shadow ray が常に遮蔽されてしまう問題を回避するため，
        # exclude_emitters=True で AreaLight をシーンから除外する．
        scene_dict, tempdir = build_scene(
            scene_config, project_root=project_root, exclude_emitters=True
        )
        try:
            scene = mi.load_dict(scene_dict)
            return _bake_direct(
                scene_config, project_root, scene, exclude_mesh_ids=exclude_mesh_ids
            )
        finally:
            shutil.rmtree(tempdir, ignore_errors=True)

    else:  # mode == "path"
        return _bake_path(
            scene_config,
            project_root,
            samples=samples,
            exclude_mesh_ids=exclude_mesh_ids,
            max_depth=max_depth,
        )


def _set_variant_once(desired_variant: str) -> None:
    """Mitsuba variant をプロセス内で 1 度だけ設定する．

    既に同じ variant が設定されている場合は何もしない．
    異なる variant が設定されている場合は警告を出して続行する（変更しない）．

    llvm_* variant 要求時は DRJIT_LIBLLVM_PATH を自動探索して設定する
    （Speed-2 wavefront 化で llvm_ad_rgb 等を要求するケースに備えた前準備）．
    scalar_* / cuda_* は LLVM-C.dll 不要なので探索しない．

    Args:
        desired_variant: 設定したい variant 文字列．
    """
    import mitsuba as mi

    current = mi.variant()
    if current == desired_variant:
        return
    if current is None:
        # Speed-2 wavefront 化用の前準備: llvm_* variant の場合は DRJIT_LIBLLVM_PATH を自動探索する
        if desired_variant.startswith("llvm_"):
            ensure_drjit_libllvm_path()
        mi.set_variant(desired_variant)
        log_variant_status(_logger, desired_variant, mi.variant())
    else:
        _logger.warning(
            "Mitsuba variant は既に %r に設定されています．"
            "要求された variant %r は無視されます．"
            "variant を変更するにはプロセスを再起動してください．",
            current,
            desired_variant,
        )


def _bake_path(
    scene_config: SceneConfig,
    project_root: Path,
    *,
    samples: int,
    exclude_mesh_ids: set[str] | None = None,
    max_depth: int | None = None,
) -> dict[str, np.ndarray]:
    """多重反射込み path ベイクの内部実装．

    コサイン重点サンプリングで hemisphere 方向を N 本生成し，
    各レイを Mitsuba path integrator で評価して入射照度を積分する．

    積分式（コサイン重点サンプリング，PDF = cosθ/π）：

        E = ∫_H² L_i(ω) cosθ dω
          ≈ (1/N) Σ L_i(ωk) cosθk / pdf(ωk)
          = (π/N) Σ L_i(ωk)      ← cosθ と PDF のコサインがキャンセル

    完全シーン（emitter + BSDF を含む）を使うため内壁の多重反射も積分される．
    scalar_rgb variant を使用（LLVM 不要）．per-ray スカラ評価のため遅いが正確に完走する．

    Args:
        scene_config: 検証済み SceneConfig．
        project_root: STL 相対パスの基点ディレクトリ．
        samples: face あたりサンプル数．大きいほど低ノイズ・高コスト．
        exclude_mesh_ids: ベイク計算からスキップする mesh_id の集合（任意）．
        max_depth: path integrator の最大反射回数（任意）．
            None のとき scene_config.render.max_depth を使用（後方互換）．
            指定時はその値で scene_dict["integrator"]["max_depth"] を上書きしてから
            mi.load_dict を呼ぶ（render 設定から独立して制御できる）．

    Returns:
        {mesh_id: irradiance_per_face (N_f,) float32}．
    """
    if exclude_mesh_ids is None:
        exclude_mesh_ids = set()
    import mitsuba as mi

    from sm_tias_lightsim.scene.builder import build_scene

    # path ベイクは scalar_rgb variant を使う（LLVM 不要）．
    # llvm_ad_rgb 等の LLVM variant を per-ray スカラループで使うと
    # LLVM JIT の megakernel シンボル解決が破綻してクラッシュするため，
    # scalar_rgb（JIT 無し）で per-ray 評価を行う（_bake_direct と同方式）．
    _set_variant_once("scalar_rgb")

    # 完全シーン（emitter と BSDF を含む）を構築する．
    # exclude_emitters=False にすることで AreaLight の rectangle も含まれ，
    # path tracer が emitter を光源として正しくサンプリングできる．
    scene_dict, tempdir = build_scene(
        scene_config, project_root=project_root, exclude_emitters=False
    )

    # max_depth を render 設定から独立して上書きする．
    # build_scene は scene_dict["integrator"] = {"type": "path", "max_depth": ...} を構築する．
    # None のとき scene_config.render.max_depth（YAML 値）をそのまま使う（後方互換）．
    # 指定時は mi.load_dict 前に dict を上書きして integrator に反映させる．
    effective_max_depth = max_depth if max_depth is not None else scene_config.render.max_depth
    scene_dict["integrator"]["max_depth"] = effective_max_depth
    _logger.info(
        "path ベイク: max_depth=%d (render.max_depth=%d, 引数指定=%s)",
        effective_max_depth,
        scene_config.render.max_depth,
        str(max_depth) if max_depth is not None else "なし（render 設定を継承）",
    )

    try:
        scene = mi.load_dict(scene_dict)
        integrator = scene.integrator()
        sampler = mi.load_dict({"type": "independent"})
        # wavefront_size=1: 1 レイずつスカラで評価する（face ループ内で毎回呼ぶため）
        sampler.seed(0, wavefront_size=1)  # type: ignore[call-arg]

        result: dict[str, np.ndarray] = {}
        n_meshes = len(scene_config.geometry.meshes)

        for mesh_idx, mesh_cfg in enumerate(scene_config.geometry.meshes):
            if mesh_cfg.id in exclude_mesh_ids:
                _logger.info(
                    "path ベイク: メッシュ %d/%d '%s' はスキップします（exclude_mesh_ids）",
                    mesh_idx + 1,
                    n_meshes,
                    mesh_cfg.id,
                )
                continue

            raw_path = Path(mesh_cfg.path)
            stl_path = raw_path if raw_path.is_absolute() else project_root / mesh_cfg.path

            loaded = load_stl(stl_path, mesh_cfg.unit)
            scaled = scale_vertices(loaded, target_unit="m")

            centroids, normals = _compute_face_centroids_and_normals(scaled.vertices, scaled.faces)
            n_faces = len(centroids)
            irradiance = np.zeros(n_faces, dtype=np.float64)

            _logger.info(
                "path ベイク: メッシュ %d/%d '%s'  faces=%d  samples=%d",
                mesh_idx + 1,
                n_meshes,
                mesh_cfg.id,
                n_faces,
                samples,
            )

            rng = np.random.default_rng(seed=mesh_idx)

            for fi in range(n_faces):
                centroid = centroids[fi]
                normal = normals[fi]

                # face 法線座標系を構築（コサイン重点サンプリングのローカル→ワールド変換用）
                n_mi = mi.Vector3f(float(normal[0]), float(normal[1]), float(normal[2]))
                frame = mi.Frame3f(n_mi)

                # origin を法線方向にオフセット（自己交差回避）
                origin = centroid + _SHADOW_RAY_EPSILON * normal

                total_radiance = np.zeros(3, dtype=np.float64)

                for _ in range(samples):
                    # コサイン重点サンプリング: [0,1)² → ローカル hemisphere 方向
                    u0, u1 = float(rng.random()), float(rng.random())
                    local_dir = mi.warp.square_to_cosine_hemisphere(mi.Point2f(u0, u1))
                    world_dir = frame.to_world(local_dir)
                    # scalar_rgb では Vector3f は通常の Mitsuba 型なので numpy 変換で取り出す
                    wd = np.array(world_dir).flatten()

                    ray = mi.Ray3f(  # type: ignore[call-arg]
                        mi.Point3f(float(origin[0]), float(origin[1]), float(origin[2])),
                        mi.Vector3f(float(wd[0]), float(wd[1]), float(wd[2])),
                    )

                    # path integrator で多重反射込みの入射放射輝度を評価する
                    # 返り値（Color3f, Bool）を numpy 変換して取り出す
                    radiance_rgb, valid, _aovs = integrator.sample(  # type: ignore[call-arg]
                        scene, sampler, ray
                    )
                    radiance_arr = np.array(radiance_rgb).flatten()  # (3,) float32
                    valid_arr = np.array(valid).flatten()  # (1,) bool

                    if bool(valid_arr[0]):
                        total_radiance[0] += float(radiance_arr[0])
                        total_radiance[1] += float(radiance_arr[1])
                        total_radiance[2] += float(radiance_arr[2])

                # E = (π/N) Σ L_i  (コサイン重点サンプリングの正規化)
                # Σ は valid サンプルのみ．N はサンプル総数（valid 比率で bias なし）
                if samples > 0:
                    radiance_mean = total_radiance / samples
                    irradiance_rgb = np.pi * radiance_mean
                    # スカラ照度として Rec.709 輝度に変換
                    irradiance[fi] = (
                        0.2126 * irradiance_rgb[0]
                        + 0.7152 * irradiance_rgb[1]
                        + 0.0722 * irradiance_rgb[2]
                    )

            result[mesh_cfg.id] = irradiance.astype(np.float32)
            _logger.info(
                "  '%s': min=%.4e  max=%.4e  std=%.4e  (valid_faces=%d/%d)",
                mesh_cfg.id,
                float(irradiance.min()),
                float(irradiance.max()),
                float(irradiance.std()),
                int((irradiance > 0).sum()),
                n_faces,
            )

    finally:
        shutil.rmtree(tempdir, ignore_errors=True)

    return result


def _bake_direct(
    scene_config: SceneConfig,
    project_root: Path,
    scene: object,
    *,
    exclude_mesh_ids: set[str] | None = None,
) -> dict[str, np.ndarray]:
    """直接照度ベイクの内部実装．

    各メッシュの face 中心に対して，全光源からの直接照度を shadow ray で計算する．

    Args:
        scene_config: 検証済み SceneConfig．
        project_root: STL 相対パスの基点ディレクトリ．
        scene: Mitsuba Scene オブジェクト（load_dict 済み）．
        exclude_mesh_ids: ベイク計算からスキップする mesh_id の集合（任意）．

    Returns:
        {mesh_id: irradiance_per_face (N_f,) float32}．
    """
    if exclude_mesh_ids is None:
        exclude_mesh_ids = set()

    result: dict[str, np.ndarray] = {}

    for mesh_cfg in scene_config.geometry.meshes:
        if mesh_cfg.id in exclude_mesh_ids:
            _logger.info(
                "direct ベイク: メッシュ '%s' はスキップします（exclude_mesh_ids）", mesh_cfg.id
            )
            continue
        raw_path = Path(mesh_cfg.path)
        stl_path = raw_path if raw_path.is_absolute() else project_root / mesh_cfg.path

        loaded = load_stl(stl_path, mesh_cfg.unit)
        scaled = scale_vertices(loaded, target_unit="m")

        centroids, normals = _compute_face_centroids_and_normals(scaled.vertices, scaled.faces)
        n_faces = len(centroids)
        irradiance = np.zeros(n_faces, dtype=np.float32)

        for light in scene_config.lights:
            if isinstance(light, PointLight):
                intensity = _to_scalar_intensity(light.intensity)
                light_pos_m = np.array(light.position, dtype=np.float64) * _MM_TO_M
                _accumulate_finite_light(
                    irradiance, centroids, normals, scene, light_pos_m, intensity
                )

            elif isinstance(light, AreaLight):
                # 面光源: 中心点を点光源として近似する（TODO: 面サンプリングで精度向上可能）
                # 近似の精度を上げるには center を N 点サンプリングして平均する方法がある．
                # 現時点ではシンプルに中心点のみで計算する．
                radiance = _to_scalar_intensity(light.radiance)
                area = (light.size[0] * _MM_TO_M) * (light.size[1] * _MM_TO_M)
                # 点光源近似の等価強度: radiance × area（Lambertian 放射源の近似）
                equiv_intensity = radiance * area

                center_m = np.array(light.center, dtype=np.float64) * _MM_TO_M
                dir_vec = np.array(light.direction, dtype=np.float64)
                dir_norm = np.linalg.norm(dir_vec)
                if dir_norm < _DIR_NORM_EPSILON:
                    continue
                emit_dir = dir_vec / dir_norm

                _accumulate_finite_light(
                    irradiance, centroids, normals, scene, center_m, equiv_intensity, emit_dir
                )

            elif isinstance(light, DirectionalLight):
                # YAML の direction は「光の進行方向」なので，face への入射方向は逆転する
                irradiance_val = _to_scalar_intensity(light.irradiance)
                dir_vec = np.array(light.direction, dtype=np.float64)
                dir_norm = np.linalg.norm(dir_vec)
                if dir_norm < _DIR_NORM_EPSILON:
                    continue
                incident_dir = -dir_vec / dir_norm

                for fi in range(n_faces):
                    centroid = centroids[fi]
                    normal = normals[fi]

                    cos_theta = float(np.dot(normal, incident_dir))
                    if cos_theta <= 0.0:
                        continue

                    origin = centroid + _SHADOW_RAY_EPSILON * normal
                    if _shadow_ray_blocked(scene, origin, incident_dir, float(_DIRECTIONAL_MAXT)):
                        continue

                    irradiance[fi] += float(irradiance_val * cos_theta)

        result[mesh_cfg.id] = irradiance

    return result