"""PyVista 3D インタラクティブビューアモジュール．

シーンのメッシュ・LED・センサを PyVista で可視化する．
照度ベイク結果 (irradiance) があればメッシュを照度で着色する．

## 投影モード（--render 引数）

luminance_map（2D 輝度マップ, np.ndarray）が指定された場合、
評価板メッシュ（id="evaluation_board"）の各頂点を orthographic センサ規約で
2D 輝度マップへ射影し、per-vertex scalar で着色する．
投影規約の詳細は visualize/projection.py の docstring を参照．

投影モードでは:
  - 評価板が主役：viridis カラーマップ + scalar bar 1 本（単位: Luminance）
  - 他メッシュは半透明ニュートラル（グレー, opacity=0.2）
  - 配色は percentile クリップで高ダイナミックレンジ潰れを防止

## 直接ベイクモード（既存動作）

irradiance dict が指定された場合は各メッシュの per-face 照度で着色する（後方互換）．
irradiance の scalar bar は 1 本に統合（重複 scalar bar のバグを解消）．
"""

from __future__ import annotations

import logging
from pathlib import Path
from typing import TYPE_CHECKING

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

if TYPE_CHECKING:
    import pyvista as pv

_logger = logging.getLogger(__name__)

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

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

# 光源・センサを示す球の半径（m 単位）
_MARKER_SPHERE_RADIUS = 0.005

# AreaLight の放射方向矢印のスケール（m 単位）
_AREA_LIGHT_ARROW_SCALE = 0.02

# DirectionalLight の矢印をシーン対角線長に対して拡大する比率
_DIRECTIONAL_ARROW_SCALE_RATIO = 0.3

# センサの視線矢印をシーン対角線長に対してクランプする比率
_SENSOR_ARROW_SCALE_RATIO = 0.2

# 投影モードで評価板以外のメッシュの不透明度
_BACKGROUND_MESH_OPACITY = 0.2

# 直接ベイクモードで着色メッシュ以外の不透明度
_UNBAKED_MESH_OPACITY = 0.5


def _build_pyvista_polydata(vertices: np.ndarray, faces: np.ndarray) -> pv.PolyData:
    """頂点・面配列から PyVista PolyData を構築する．

    PyVista の faces 形式は各行が [n_vertices, v0, v1, v2, ...] の 1D フラット配列．

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

    Returns:
        PyVista PolyData オブジェクト．
    """
    import pyvista as pv

    n_faces = len(faces)
    # PyVista 形式: [3, v0, v1, v2, 3, v0, v1, v2, ...]
    pv_faces = np.hstack([np.full((n_faces, 1), 3, dtype=np.int64), faces]).flatten()
    return pv.PolyData(vertices.astype(np.float32), pv_faces)


def _get_scene_bbox_center_and_diag(
    scene_config: SceneConfig, project_root: Path
) -> tuple[np.ndarray, float]:
    """シーン全メッシュのバウンディングボックス中心と対角線長を返す．

    DirectionalLight の矢印位置を決めるために使う．
    メッシュが空なら原点と 1.0 を返す．

    Args:
        scene_config: SceneConfig．
        project_root: STL 相対パスの基点ディレクトリ．

    Returns:
        (center (3,), diagonal_length) のタプル．
    """
    all_vertices: list[np.ndarray] = []
    for mesh_cfg in scene_config.geometry.meshes:
        try:
            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")
            all_vertices.append(scaled.vertices)
        except Exception as exc:
            _logger.warning("バウンディングボックス計算でメッシュを読み込めませんでした: %s", exc)

    if not all_vertices:
        return np.zeros(3), 1.0

    combined = np.vstack(all_vertices)
    mins = combined.min(axis=0)
    maxs = combined.max(axis=0)
    center = (mins + maxs) / 2.0
    diagonal = float(np.linalg.norm(maxs - mins))
    return center, max(diagonal, 1e-6)


def _apply_percentile_clip(
    values: np.ndarray, clip_low: float, clip_high: float
) -> np.ndarray:
    """percentile クリップを適用して高 DR 潰れを防ぐ．

    Args:
        values: スカラ配列．
        clip_low: 下限 percentile（%）．
        clip_high: 上限 percentile（%）．

    Returns:
        クリップ済みの配列（元と同 dtype）．
    """
    p_low = float(np.percentile(values, clip_low))
    p_high = float(np.percentile(values, clip_high))
    return np.clip(values, p_low, p_high)


def _prepare_projection_scalars(
    scene_config: SceneConfig,
    project_root: Path,
    luminance_map: np.ndarray,
    clip_low: float,
    clip_high: float,
) -> tuple[str | None, np.ndarray | None]:
    """投影モード用に評価板 id とper-vertex 輝度スカラを準備する．

    評価板メッシュを特定して STL をロードし、orthographic 投影で
    per-vertex 輝度スカラ配列を計算して返す．

    Args:
        scene_config: 検証済み SceneConfig．
        project_root: STL 相対パスの基点ディレクトリ．
        luminance_map: sm-render 出力の 2D 輝度マップ (H, W) または (H, W, C)．
        clip_low: percentile クリップの下限（%）．
        clip_high: percentile クリップの上限（%）．

    Returns:
        (eval_board_id, projected_scalars) のタプル．
        失敗時は (None, None) を返す（例外は投げない）．
    """
    from sm_tias_lightsim.visualize.projection import (
        find_eval_board_mesh_id,
        project_luminance_to_vertices,
    )

    eval_board_id = find_eval_board_mesh_id(scene_config)
    _logger.info("投影モード: 評価板メッシュ id = '%s'", eval_board_id)

    eval_mesh_cfg = next(m for m in scene_config.geometry.meshes if m.id == eval_board_id)
    try:
        raw_path = Path(eval_mesh_cfg.path)
        stl_path = raw_path if raw_path.is_absolute() else project_root / eval_mesh_cfg.path
        eval_loaded = load_stl(stl_path, eval_mesh_cfg.unit)
        eval_scaled = scale_vertices(eval_loaded, target_unit="m")
        projected_scalars = project_luminance_to_vertices(
            eval_scaled.vertices,
            luminance_map,
            scene_config.sensor,
            clip_low=clip_low,
            clip_high=clip_high,
        )
        _logger.info(
            "投影スカラ: min=%.4f, max=%.4f, std=%.4f",
            float(projected_scalars.min()),
            float(projected_scalars.max()),
            float(projected_scalars.std()),
        )
    except Exception as exc:
        _logger.warning("投影計算に失敗しました: %s．シーンレイアウトのみ表示します．", exc)
        return None, None

    return eval_board_id, projected_scalars


def show_scene(
    scene_config: SceneConfig,
    project_root: Path,
    irradiance: dict[str, np.ndarray] | None = None,
    luminance_map: np.ndarray | None = None,
    *,
    off_screen: bool = False,
    screenshot: Path | None = None,
    clip_low: float = 1.0,
    clip_high: float = 99.0,
) -> None:
    """PyVista でシーンとベイク結果を表示する．

    メッシュ・LED・センサ・座標軸を含む 3D インタラクティブビューアを開く．

    表示モードの優先順位:
      1. luminance_map が指定された場合（投影モード）:
           評価板メッシュに 2D 輝度マップを orthographic 投影して per-vertex 着色．
           他メッシュは半透明ニュートラル（opacity=0.2）．
           scalar bar は評価板の 1 本のみ（単位: Luminance）．
      2. irradiance が指定された場合（直接ベイクモード）:
           各メッシュを per-face 照度で着色（viridis）．
           scalar bar は 1 本に統合（旧実装の重複バグを解消）．
      3. 両方 None の場合: シーンレイアウトのみ（--no-bake 相当）．

    off_screen=True のときは GUI ウィンドウを開かず，screenshot が指定されていれば保存する．
    screenshot=None かつ off_screen=True のときは何も出力しない（no-op）．

    Args:
        scene_config: 検証済み SceneConfig．
        project_root: STL 相対パスの基点ディレクトリ．
        irradiance: {mesh_id: irradiance_per_face (N_f,)} の辞書．
            None かつ luminance_map も None なら着色なし（後方互換）．
        luminance_map: sm-render 出力の 2D 輝度マップ (H, W) または (H, W, C)．
            指定時は投影モードで評価板に投影する（irradiance より優先）．
        off_screen: True のとき GUI ウィンドウを開かない（screenshot 専用）．
        screenshot: 保存先 PNG パス．指定時は off_screen で screenshot を保存する．
        clip_low: percentile クリップの下限（%）．既定 1．
        clip_high: percentile クリップの上限（%）．既定 99．
    """
    import pyvista as pv

    # screenshot 指定があれば off_screen を強制する
    if screenshot is not None:
        off_screen = True

    plotter = pv.Plotter(off_screen=off_screen)

    # 投影モードか否かを決定
    use_projection = luminance_map is not None

    # 投影モードの場合：評価板メッシュ id と per-vertex スカラを準備する
    projected_scalars: np.ndarray | None = None
    eval_board_id: str | None = None
    if use_projection:
        eval_board_id, projected_scalars = _prepare_projection_scalars(
            scene_config, project_root, luminance_map, clip_low, clip_high  # type: ignore[arg-type]
        )
        if eval_board_id is None:
            # 投影計算失敗 → レイアウト表示のみにフォールバック
            use_projection = False

    # ---- メッシュ追加 ----
    scalar_bar_added = False  # scalar bar を 1 本だけ追加するフラグ

    for mesh_cfg in scene_config.geometry.meshes:
        try:
            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")
        except Exception as exc:
            _logger.warning("メッシュ '%s' を読み込めませんでした: %s", mesh_cfg.id, exc)
            continue

        polydata = _build_pyvista_polydata(scaled.vertices, scaled.faces)

        if use_projection and mesh_cfg.id == eval_board_id and projected_scalars is not None:
            # 投影モード: 評価板を per-vertex 輝度で着色（viridis, 知覚均等）
            # point_data に格納して per-vertex スカラとして扱う
            polydata.point_data["luminance"] = projected_scalars

            show_bar = not scalar_bar_added
            plotter.add_mesh(
                polydata,
                scalars="luminance",
                cmap="viridis",
                show_scalar_bar=show_bar,
                scalar_bar_args={"title": "Luminance (a.u.)"},
                label=mesh_cfg.id,
            )
            if show_bar:
                scalar_bar_added = True

        elif use_projection:
            # 投影モード: 評価板以外は半透明ニュートラルグレー（背景扱い）
            plotter.add_mesh(
                polydata,
                color="lightgray",
                opacity=_BACKGROUND_MESH_OPACITY,
                label=mesh_cfg.id,
            )

        elif irradiance is not None and mesh_cfg.id in irradiance:
            # 直接ベイクモード: per-face 照度で着色
            irr = irradiance[mesh_cfg.id]
            # percentile クリップを適用
            irr_clipped = _apply_percentile_clip(irr.astype(np.float32), clip_low, clip_high)
            polydata.cell_data["irradiance"] = irr_clipped

            show_bar = not scalar_bar_added
            plotter.add_mesh(
                polydata,
                scalars="irradiance",
                cmap="viridis",
                show_scalar_bar=show_bar,
                scalar_bar_args={"title": "Irradiance (W/m²)"},
                label=mesh_cfg.id,
            )
            if show_bar:
                scalar_bar_added = True

        else:
            # 着色なし: グレー半透明
            plotter.add_mesh(
                polydata,
                color="lightgray",
                opacity=_UNBAKED_MESH_OPACITY,
                label=mesh_cfg.id,
            )

    # ---- シーン全体のバウンディングボックス（DirectionalLight 用）----
    bbox_center, bbox_diag = _get_scene_bbox_center_and_diag(scene_config, project_root)

    # ---- LED / 光源表示 ----
    for light in scene_config.lights:
        if isinstance(light, PointLight):
            pos_m = [c * _MM_TO_M for c in light.position]
            sphere = pv.Sphere(radius=_MARKER_SPHERE_RADIUS, center=pos_m)
            plotter.add_mesh(sphere, color="yellow", label=f"LED:{light.id}")

        elif isinstance(light, AreaLight):
            center_m = [c * _MM_TO_M for c in light.center]
            size_m = [light.size[0] * _MM_TO_M, light.size[1] * _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:
                dir_vec = dir_vec / dir_norm

            plane = pv.Plane(
                center=center_m,
                direction=dir_vec.tolist(),
                i_size=size_m[0],
                j_size=size_m[1],
            )
            plotter.add_mesh(plane, color="orange", opacity=0.7, label=f"AreaLight:{light.id}")

            arrow = pv.Arrow(
                start=center_m,
                direction=dir_vec.tolist(),
                scale=_AREA_LIGHT_ARROW_SCALE,
            )
            plotter.add_mesh(arrow, color="orange", label=f"AreaLight:{light.id}:dir")

        elif isinstance(light, DirectionalLight):
            # バウンディングボックス中心から矢印 1 本で進行方向を示す
            dir_vec = np.array(light.direction, dtype=np.float64)
            dir_norm = np.linalg.norm(dir_vec)
            if dir_norm < _DIR_NORM_EPSILON:
                continue
            dir_vec = dir_vec / dir_norm

            arrow_scale = bbox_diag * _DIRECTIONAL_ARROW_SCALE_RATIO
            arrow = pv.Arrow(
                start=bbox_center.tolist(),
                direction=dir_vec.tolist(),
                scale=arrow_scale,
            )
            plotter.add_mesh(arrow, color="lightyellow", label=f"DirLight:{light.id}")

    # ---- センサ表示 ----
    sensor_pos_m = [c * _MM_TO_M for c in scene_config.sensor.position]
    sensor_target_m = [c * _MM_TO_M for c in scene_config.sensor.target]

    sensor_sphere = pv.Sphere(radius=_MARKER_SPHERE_RADIUS, center=sensor_pos_m)
    plotter.add_mesh(sensor_sphere, color="blue", label="sensor")

    sensor_pos_arr = np.array(sensor_pos_m)
    sensor_target_arr = np.array(sensor_target_m)
    view_dir = sensor_target_arr - sensor_pos_arr
    view_len = np.linalg.norm(view_dir)
    if view_len > _DIR_NORM_EPSILON:
        view_dir_n = view_dir / view_len
        sensor_arrow = pv.Arrow(
            start=sensor_pos_m,
            direction=view_dir_n.tolist(),
            scale=min(float(view_len), bbox_diag * _SENSOR_ARROW_SCALE_RATIO),
        )
        plotter.add_mesh(sensor_arrow, color="blue", label="sensor:view")

    # ---- 座標軸と境界表示 ----
    plotter.add_axes()
    plotter.show_bounds(grid="back", location="outer")

    # ---- 投影モード: 評価板正面が見えるカメラ視点を設定 ----
    if use_projection:
        # センサ位置から評価板を見る視点に設定
        # センサは -Z 側から +Z 方向を向いているので、
        # ビューワは少し引いて斜め上から俯瞰する位置に設定する
        sensor_pos = np.array(sensor_pos_m)
        sensor_tgt = np.array(sensor_target_m)
        # センサ位置からさらに -Z 方向（センサの裏側）に引いた位置
        cam_offset = sensor_pos - sensor_tgt
        cam_dist = np.linalg.norm(cam_offset)
        if cam_dist < 1e-9:
            cam_dist = 0.1
        # カメラを評価板正面かつ斜め上から
        cam_pos = sensor_pos + cam_offset / cam_dist * 0.25
        cam_pos[1] += 0.05  # 少し上から
        plotter.camera.position = cam_pos.tolist()
        plotter.camera.focal_point = sensor_tgt.tolist()
        plotter.camera.up = (0.0, 1.0, 0.0)
        plotter.reset_camera()

    # ---- 表示 / 保存 ----
    if screenshot is not None:
        plotter.show(screenshot=str(screenshot))
    elif not off_screen:
        plotter.show()
    # off_screen=True かつ screenshot=None の場合は no-op