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

shadow ray + コサイン項 + 距離減衰による直接照度（W/m²）を計算し，
{mesh_id: irradiance_per_face (N_f,) float32} の形式で返す．

多重反射（インターリフレクション）は Phase A では無視する．
将来的に mode="path" を追加することで path tracing ベースの計算に切り替え可能．
"""

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.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"] = "direct",
) -> dict[str, np.ndarray]:
    """各メッシュの face 中心で照度（W/m²）を計算する（直接照度のみ）．

    Phase A: shadow ray + コサイン項 + 距離減衰による直接照度のみを計算する．
    多重反射は無視する（Phase B で mode="path" を追加予定）．

    shadow ray の評価には Mitsuba の `scene.ray_intersect_preliminary` を利用する．
    遮蔽された face は寄与をスキップする．

    Args:
        scene_config: 検証済み SceneConfig．
        project_root: STL 相対パスの基点ディレクトリ．
        mode: 計算モード．現在は "direct" のみサポート．

    Returns:
        {mesh_id: irradiance_per_face (N_f,) float32} の辞書．
        irradiance の単位は W/m²（intensity の単位に依存）．

    Raises:
        NotImplementedError: mode が "direct" 以外の場合．
        FileNotFoundError: STL ファイルが見つからない場合．
    """
    if mode != "direct":
        msg = f"mode={mode!r} は未サポートです．現在は 'direct' のみ実装しています．"
        raise NotImplementedError(msg)

    import mitsuba as mi

    from sm_tias_lightsim.scene.builder import build_scene

    # Mitsuba variant を初期化（scalar_rgb を使う：shadow ray 評価に最適）
    # runner.py と同じ variant が既に設定されている場合はそのまま続行する
    desired_variant = "scalar_rgb"
    try:
        mi.set_variant(desired_variant)
    except Exception:
        _logger.warning(
            "Mitsuba variant の切り替えに失敗しました．現在の 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)
    finally:
        shutil.rmtree(tempdir, ignore_errors=True)


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

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

    Args:
        scene_config: 検証済み SceneConfig．
        project_root: STL 相対パスの基点ディレクトリ．
        scene: Mitsuba Scene オブジェクト（load_dict 済み）．

    Returns:
        {mesh_id: irradiance_per_face (N_f,) float32}．
    """
    result: dict[str, np.ndarray] = {}

    for mesh_cfg in scene_config.geometry.meshes:
        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