"""シミュ評価板 STL の上に ROI 矩形を 3D で重ねて表示するスクリプト．

使い方:
    uv run python scripts/visualize_roi_3d.py [--white-board configs/white_board_bbox.json] \
        [--real-roi-x 539 --real-roi-y 1015 --real-roi-w 1400 --real-roi-h 1409] \
        [--screenshot output/roi_3d.png]

意図:
    実測 ROI が白色板内で占める「相対位置」を、シミュの評価板（同じ形状の白色板を想定）の
    物理サイズに変換し、評価板 STL の表面に 3D 矩形として重ねる．
    人間が見て「シミュの ROI が評価板のここに対応する」と確認するため．

備考:
    白色板の物理サイズは「評価板の bbox（X=99, Y=102 mm）」を使う
    （白色板 = 評価板の有効測定面と仮定）．
"""

from __future__ import annotations

import argparse
import json
from pathlib import Path

import numpy as np
import pyvista as pv
import trimesh


def main() -> None:
    parser = argparse.ArgumentParser(description="Visualize ROI on eval board (3D)")
    parser.add_argument("--white-board", type=Path, default=Path("configs/white_board_bbox.json"))
    parser.add_argument("--real-roi-x", type=int, default=539)
    parser.add_argument("--real-roi-y", type=int, default=1015)
    parser.add_argument("--real-roi-w", type=int, default=1400)
    parser.add_argument("--real-roi-h", type=int, default=1409)
    parser.add_argument(
        "--screenshot", type=Path, default=None, help="出力 PNG．未指定なら GUI を開く"
    )
    parser.add_argument("--off-screen", action="store_true")
    parser.add_argument(
        "--flip-y",
        action="store_true",
        default=True,
        # 実機の物理上下と画像 y の向きが逆の場合に反転する（デフォルト ON）
        help="実測 PNG の y 軸を反転して物理 Y にマップ．デフォルト ON",
    )
    parser.add_argument(
        "--no-flip-y",
        dest="flip_y",
        action="store_false",
        help="y 軸を反転しない",
    )
    parser.add_argument(
        "--flip-x",
        action="store_true",
        default=False,
        help="実測 PNG の x 軸を反転して物理 X にマップ（実機の物理左右と画像 x の向きが逆の場合）",
    )
    args = parser.parse_args()

    if not args.white_board.exists():
        msg = (
            f"白色板 bbox JSON が見つかりません: {args.white_board}\n"
            "先に `uv run python scripts/select_white_board.py` を実行してください．"
        )
        raise FileNotFoundError(msg)

    wb = json.loads(args.white_board.read_text(encoding="utf-8"))
    wb_x, wb_y, wb_w, wb_h = wb["x"], wb["y"], wb["width"], wb["height"]
    print(f"白色板 bbox (real PNG): x={wb_x}, y={wb_y}, w={wb_w}, h={wb_h}")

    rx, ry, rw, rh = args.real_roi_x, args.real_roi_y, args.real_roi_w, args.real_roi_h
    print(f"実測 ROI: x={rx}, y={ry}, w={rw}, h={rh}")

    # ROI の相対位置（白色板内 0〜1）
    rel_x = (rx - wb_x) / wb_w
    rel_y = (ry - wb_y) / wb_h
    rel_w = rw / wb_w
    rel_h = rh / wb_h
    print(f"  相対位置: x={rel_x:.3f}, y={rel_y:.3f}, w={rel_w:.3f}, h={rel_h:.3f}")

    # 評価板 STL を読み込み
    eval_stl = Path("configs/meshes/照明ムラ補正ボード.stl")
    m = trimesh.load(eval_stl, force="mesh")
    print(f"\n評価板 STL bbox (mm): min={m.bounds[0]}, max={m.bounds[1]}")
    board_xmin, board_xmax = m.bounds[0, 0], m.bounds[1, 0]
    board_ymin, board_ymax = m.bounds[0, 1], m.bounds[1, 1]
    board_w = board_xmax - board_xmin
    board_h = board_ymax - board_ymin

    # ROI の評価板内 mm 範囲
    if args.flip_x:
        flipped_x_start = 1.0 - (rel_x + rel_w)
        roi_xmin_mm = board_xmin + flipped_x_start * board_w
        roi_xmax_mm = roi_xmin_mm + rel_w * board_w
        print(f"  X 反転適用: rel_x_orig={rel_x:.3f} → flipped={flipped_x_start:.3f}")
    else:
        roi_xmin_mm = board_xmin + rel_x * board_w
        roi_xmax_mm = roi_xmin_mm + rel_w * board_w
    if args.flip_y:
        # 実測 PNG の y 増加（画像下方向）= 物理 Y 減少と仮定して反転
        # 白色板内 y 相対位置 [rel_y, rel_y+rel_h] を 1.0 から引いて反転
        flipped_y_start = 1.0 - (rel_y + rel_h)
        roi_ymin_mm = board_ymin + flipped_y_start * board_h
        roi_ymax_mm = roi_ymin_mm + rel_h * board_h
        print(f"  Y 反転適用: rel_y_orig={rel_y:.3f} → flipped={flipped_y_start:.3f}")
    else:
        roi_ymin_mm = board_ymin + rel_y * board_h
        roi_ymax_mm = roi_ymin_mm + rel_h * board_h

    # 評価板の裏面（平坦面、-Z 法線、Z=最小）に ROI を重ねる
    # 表面は凹凸が多くて視認しにくいため平坦な裏面を採用
    z_bottom = m.bounds[0, 2]
    z_offset = -1.0  # 裏面のさらに -Z 方向に矩形を浮かせて視認性確保

    print("\nROI の評価板上 mm 範囲:")
    print(f"  X: [{roi_xmin_mm:.2f}, {roi_xmax_mm:.2f}]")
    print(f"  Y: [{roi_ymin_mm:.2f}, {roi_ymax_mm:.2f}]")
    print(f"  Z: {z_bottom + z_offset:.2f}（裏面 Z={z_bottom:.2f} のさらに {z_offset}mm 下）")

    # PyVista で表示
    plotter = pv.Plotter(
        off_screen=args.off_screen or args.screenshot is not None,
        window_size=(1400, 1000),
    )

    # ----- シーン全体のメッシュを追加 -----
    scene_meshes = [
        ("照明ムラ補正ボード.stl", "lightgray", 0.85, "評価板"),
        ("前側.stl", "white", 0.20, "前側筐体"),
        ("後側.stl", "white", 0.20, "後側筐体"),
        ("スマホホルダ.stl", "tan", 0.30, "スマホホルダ"),
        ("LED基板カバー.stl", "lightblue", 0.40, "LED 基板カバー"),
        ("SH-02M.stl", "dimgray", 0.30, "スマホ"),
    ]
    for stl_name, color, opacity, label in scene_meshes:
        stl_path = Path(f"configs/meshes/{stl_name}")
        if not stl_path.exists():
            continue
        mesh = trimesh.load(stl_path, force="mesh")
        v = mesh.vertices * 0.001
        f = np.hstack([np.full((len(mesh.faces), 1), 3), mesh.faces]).flatten()
        plotter.add_mesh(pv.PolyData(v, f), color=color, opacity=opacity, label=label)

    # LED 18 個を黄色い小球で表示
    import yaml
    cfg = yaml.safe_load(Path("configs/scene_baseline.yaml").read_text(encoding="utf-8"))
    for light in cfg.get("lights", []):
        if light.get("type") == "area":
            cx, cy, cz = light["center"]
            pos = np.array([cx, cy, cz]) * 0.001
            plotter.add_mesh(pv.Sphere(radius=0.003, center=pos), color="yellow")

    # ROI 矩形（評価板の裏面 -Z 側に浮かせて表示）
    roi_corners_mm = np.array([
        [roi_xmin_mm, roi_ymin_mm, z_bottom + z_offset],
        [roi_xmax_mm, roi_ymin_mm, z_bottom + z_offset],
        [roi_xmax_mm, roi_ymax_mm, z_bottom + z_offset],
        [roi_xmin_mm, roi_ymax_mm, z_bottom + z_offset],
    ])
    roi_corners_m = roi_corners_mm * 0.001
    # 矩形を 2 三角形で
    faces = np.array([[3, 0, 1, 2], [3, 0, 2, 3]]).flatten()
    roi_poly = pv.PolyData(roi_corners_m, faces)
    plotter.add_mesh(roi_poly, color="red", opacity=0.5, label="ROI（実測対応）")
    # 矩形の枠線も
    for i in range(4):
        a = roi_corners_m[i]
        b = roi_corners_m[(i + 1) % 4]
        line = pv.Line(a, b)
        plotter.add_mesh(line, color="red", line_width=5)

    plotter.add_axes()
    annotation = (
        f"Red = ROI on eval board BACK surface (flat, -Z side)\n"
        f"Range (mm): X=[{roi_xmin_mm:.1f},{roi_xmax_mm:.1f}]"
        f" Y=[{roi_ymin_mm:.1f},{roi_ymax_mm:.1f}]"
    )
    plotter.add_text(annotation, position="upper_left", font_size=10, color="black")

    # 座標軸（X+/Y+/Z+ の向きが見える）
    plotter.add_axes(interactive=True)
    plotter.show_bounds(
        grid="back",
        location="outer",
        xtitle="X [m]",
        ytitle="Y [m]",
        ztitle="Z [m]",
        all_edges=True,
    )

    # デフォルトの視点（screenshot 時のみ．GUI 時はマウスでぐりぐり回せる）
    if args.screenshot is not None:
        plotter.camera_position = [(0.15, -0.20, 0.05), (0, 0, -0.05), (0, 0, 1)]

    if args.screenshot is not None:
        args.screenshot.parent.mkdir(parents=True, exist_ok=True)
        plotter.show(screenshot=str(args.screenshot))
        print(f"\n保存: {args.screenshot}")
    else:
        plotter.show()


if __name__ == "__main__":
    main()