"""シーン設定・シーン幾何の物理的健全性を監査するモジュール．

「エラーも警告も出ないまま結果だけが静かに歪む」タイプの不具合を検出することを目的とする．
TEST_12（シェーディング法線アーティファクト）では，レンダリング規約の取り違えが
CoV の 6 割超を占める系統誤差を生みながら 4 ステップにわたり気づかれなかった．
同種の故障モードを自動検出できるようにするのが本モジュールの役割．

現在の監査項目:

- **多重反射の打ち切りバイアス** (`estimate_truncation_residual`)
  SmTIAS の照明は LED 直接光が評価板に届かない間接照明のみの構成（SPEC_03）であり，
  輝度は白筐体内壁の多重反射の累積で決まる．`render.max_depth` が小さいと
  高次の反射が切り捨てられ，絶対輝度が系統的に過小評価される．
- **シェーディング法線の同一性・面の向き** (`audit_scene_normals`)
  読み込み済み Mitsuba シーンにレイを飛ばし，シェーディング法線が幾何法線と一致するか
  （＝`face_normals` が実効か）と，空洞内部から見て裏向きの面が無いかを実測する．

設定に対する検査は副作用を持たない純関数．`audit_scene_normals` は
読み込み済みシーンへの読み取り専用クエリで，seed 固定により決定論的．
"""

from __future__ import annotations

import math
from dataclasses import dataclass
from typing import Any

import numpy as np

from sm_tias_lightsim.scene.schema import (
    BlendBSDFConfig,
    DiffuseMaterial,
    MaterialConfig,
    RoughConductorBSDFConfig,
    SceneConfig,
)

# 打ち切り残差の既定許容値（1%）．これを超えると絶対輝度が無視できない過小評価になる．
DEFAULT_TRUNCATION_TOLERANCE = 0.01

# material="none" の roughconductor は完全反射体として扱う（schema の仕様）
_PERFECT_REFLECTOR = 1.0


def material_max_reflectance(material: MaterialConfig) -> float | None:
    """1 材料の最大反射率を返す（不透明反射材のみ）．

    Args:
        material: SceneConfig.materials の値．

    Returns:
        反射率の最大成分．反射率を定義できない材料（誘電体等）は None．
    """
    if isinstance(material, DiffuseMaterial):
        r = material.reflectance
        return float(r) if isinstance(r, (int, float)) else float(max(r))

    if isinstance(material, RoughConductorBSDFConfig):
        # material="none" 固定＝完全反射体（schema の仕様）
        return _PERFECT_REFLECTOR

    if isinstance(material, BlendBSDFConfig):
        values = [material_max_reflectance(c) for c in material.components]
        known = [v for v in values if v is not None]
        return max(known) if known else None

    # DielectricMaterial / ThinDielectricMaterial は吸収を持たず反射率で表せない
    return None


def scene_max_reflectance(config: SceneConfig) -> float | None:
    """シーン中の最大反射率を返す．

    多重反射の減衰率を支配するのは最も明るい材料なので，最大値を使う．

    Args:
        config: 検証済み SceneConfig．

    Returns:
        最大反射率．反射率を定義できる材料が 1 つも無ければ None．
    """
    values = [material_max_reflectance(m) for m in config.materials.values()]
    known = [v for v in values if v is not None]
    return max(known) if known else None


def estimate_truncation_residual(reflectance: float, max_depth: int) -> float:
    """多重反射を max_depth で打ち切ることによる相対エネルギー欠損を見積もる．

    反射率 rho の閉じた空洞では，n 回反射後に残るエネルギーが rho^n なので，
    深度 max_depth で打ち切ると収束値に対して概ね rho^(max_depth - 1) を取りこぼす．

    実測との対応（本モジュール導入時の検証）:

    | シーン | rho | max_depth | 本式の予測欠損 | 実測欠損 |
    | --- | --- | --- | --- | --- |
    | 積分球 | 0.80 | 16 | 3.5% | 4.6% |
    | 基準設計シーン | 0.85 | 16 | 8.7% | 9.7% |

    Args:
        reflectance: シーン中の最大反射率（0 以上）．
        max_depth: `render.max_depth`．

    Returns:
        収束値に対する相対エネルギー欠損の見積もり（0〜1．1 に近いほど未収束）．
        max_depth が 1 以下なら 1.0（ほぼ何も拾えていない）を返す．
    """
    if max_depth <= 1:
        return 1.0
    return float(reflectance ** (max_depth - 1))


def check_truncation_bias(
    config: SceneConfig, tolerance: float = DEFAULT_TRUNCATION_TOLERANCE
) -> str | None:
    """max_depth が多重反射を十分に拾えているか検査し，不足なら警告文を返す．

    Args:
        config: 検証済み SceneConfig．
        tolerance: 許容する相対エネルギー欠損（既定 1%）．

    Returns:
        欠損が tolerance を超える場合は警告メッセージ．十分なら None．
        反射率を定義できる材料が無い場合も None．
    """
    reflectance = scene_max_reflectance(config)
    if reflectance is None:
        return None

    residual = estimate_truncation_residual(reflectance, config.render.max_depth)
    if residual <= tolerance:
        return None

    return (
        f"render.max_depth={config.render.max_depth} は最大反射率 {reflectance:.3f} の"
        f"多重反射を拾いきれていません（相対エネルギー欠損の見積もり {residual:.1%}）．"
        f"絶対輝度が系統的に過小評価されます．"
        f"max_depth >= {required_max_depth(reflectance, tolerance)} を推奨します．"
    )


def required_max_depth(reflectance: float, tolerance: float = DEFAULT_TRUNCATION_TOLERANCE) -> int:
    """相対エネルギー欠損を tolerance 以下に収めるのに必要な max_depth を返す．

    Args:
        reflectance: シーン中の最大反射率．
        tolerance: 許容する相対エネルギー欠損．

    Returns:
        必要な max_depth（1 以上の整数）．reflectance が 0 以下または 1 以上で
        解が定まらない場合はそれぞれ 1 / 大きな既定値を返す．
    """
    if reflectance <= 0.0:
        return 1
    if reflectance >= 1.0:
        # 吸収の無い完全反射体では有限深度で収束しない．実用上の上限を返す．
        return 1024
    return 1 + math.ceil(math.log(tolerance) / math.log(reflectance))


@dataclass(frozen=True)
class NormalAuditResult:
    """`audit_scene_normals` の測定結果．

    Attributes:
        hits: シーンに命中したレイの本数．
        backface_hits: 法線がレイの発射点と反対を向いていた命中数．
            空洞内部から飛ばした場合，これが 0 でないメッシュは裏返っている．
        max_shading_normal_deviation_deg: シェーディング法線と幾何法線の
            最大なす角 [deg]．`face_normals` が実効なら浮動小数点誤差程度になる．
    """

    hits: int
    backface_hits: int
    max_shading_normal_deviation_deg: float


def _sample_directions(rng: np.random.Generator, count: int) -> np.ndarray:
    """単位球面上の一様方向を count 本サンプリングする．"""
    directions = rng.normal(size=(count, 3))
    return directions / np.linalg.norm(directions, axis=1, keepdims=True)


def _intersect_batched(scene: Any, origins: np.ndarray, directions: np.ndarray):
    """ベクトル化 variant（llvm_* / cuda_*）で全レイを 1 回の交差計算にまとめる．

    Returns:
        (valid, geometric, shading) のタプル．いずれも (N,) / (N, 3) の numpy 配列．
    """
    import mitsuba as mi

    ray = mi.Ray3f(
        mi.Point3f(origins[:, 0], origins[:, 1], origins[:, 2]),
        mi.Vector3f(directions[:, 0], directions[:, 1], directions[:, 2]),
    )
    si = scene.ray_intersect(ray)
    valid = np.asarray(si.is_valid()).reshape(-1)
    # Dr.Jit の 3 成分ベクトルは (3, N) で numpy 化されるので転置して (N, 3) に揃える
    geometric = np.asarray(si.n).reshape(3, -1).T
    shading = np.asarray(si.sh_frame.n).reshape(3, -1).T
    return valid, geometric, shading


def _intersect_per_ray(scene: Any, origins: np.ndarray, directions: np.ndarray):
    """scalar_* variant 用に 1 本ずつ交差計算する（ベクトル化非対応のため）．

    Returns:
        `_intersect_batched` と同じ形式のタプル．
    """
    import mitsuba as mi

    count = len(origins)
    valid = np.zeros(count, dtype=bool)
    geometric = np.zeros((count, 3), dtype=float)
    shading = np.zeros((count, 3), dtype=float)

    for i in range(count):
        ray = mi.Ray3f(
            mi.Point3f(float(origins[i, 0]), float(origins[i, 1]), float(origins[i, 2])),
            mi.Vector3f(float(directions[i, 0]), float(directions[i, 1]), float(directions[i, 2])),
        )
        si = scene.ray_intersect(ray)
        if not bool(si.is_valid()):
            continue
        valid[i] = True
        geometric[i] = np.asarray(si.n).reshape(3)
        shading[i] = np.asarray(si.sh_frame.n).reshape(3)

    return valid, geometric, shading


def audit_scene_normals(
    scene: Any,
    origins: np.ndarray,
    rays_per_origin: int = 500,
    seed: int = 0,
) -> NormalAuditResult:
    """空洞内部からレイを飛ばし，法線の健全性を実測する．

    2 つの「気づけない故障モード」を同時に検査する．

    1. **シェーディング法線の汚染**（TEST_12 の原因）．
       PLY に法線が無いと Mitsuba は既定でスムーズ頂点法線を自動生成し，
       巨大三角形の平面がエッジ頂点の平均化法線で描画されて偽の異方性を生む．
       `face_normals` が実効なら幾何法線と一致する．
    2. **面の裏返り**．片面 BSDF では裏向きの面が無警告で真っ黒になる
       （`materials.build_scene_bsdf_dict` の twosided ラップで救済済みだが，
       裏返り自体はモデリング誤りの兆候なので検出する）．

    レイはベクトル化 variant では 1 回の交差計算にまとめる．llvm_* で 1 本ずつ
    呼ぶと JIT カーネルが大量に生成され，実行時に不安定になるため．

    Args:
        scene: `mi.load_dict` で読み込み済みの Mitsuba シーン．
        origins: レイ発射点 (N, 3)．シーンと同じ単位（m）で空洞内部の点を与える．
        rays_per_origin: 発射点あたりのレイ本数．
        seed: 方向サンプリングの乱数シード（決定論性のため固定する）．

    Returns:
        NormalAuditResult．
    """
    import mitsuba as mi

    origins = np.asarray(origins, dtype=float)
    rng = np.random.default_rng(seed)

    ray_origins = np.repeat(origins, rays_per_origin, axis=0)
    ray_directions = _sample_directions(rng, len(ray_origins))

    intersect = _intersect_per_ray if mi.variant().startswith("scalar") else _intersect_batched
    valid, geometric, shading = intersect(scene, ray_origins, ray_directions)

    if not valid.any():
        return NormalAuditResult(hits=0, backface_hits=0, max_shading_normal_deviation_deg=0.0)

    geometric = geometric[valid]
    shading = shading[valid]
    directions = ray_directions[valid]

    cos_deviation = np.abs(np.einsum("ij,ij->i", geometric, shading))
    deviation_deg = np.degrees(np.arccos(np.clip(cos_deviation, 0.0, 1.0)))
    backface = np.einsum("ij,ij->i", geometric, -directions) < 0.0

    return NormalAuditResult(
        hits=int(valid.sum()),
        backface_hits=int(backface.sum()),
        max_shading_normal_deviation_deg=float(deviation_deg.max()),
    )