"""Fusion 360 から出力した LED ユニット STL から LED 位置を抽出する.

STL ファイルに含まれる連結成分を解析し，カバー（最大の連結成分）と LED 群
（小さい連結成分）を識別する．各 LED の中心座標・サイズ・法線方向（発光方向）を
抽出して scene_baseline.yaml の lights セクションに組み込める形式で出力する.

使い方:
    # 1 ファイル（複数ユニット含む）
    uv run python scripts/extract_led_positions_from_stl.py \\
        configs/meshes/LEDユニット_all.stl

    # 複数 STL ファイル（各ユニット別ファイル）
    uv run python scripts/extract_led_positions_from_stl.py \\
        configs/meshes/LEDユニット下.stl \\
        configs/meshes/LEDユニット右.stl \\
        configs/meshes/LEDユニット左.stl

    # 出力を YAML ファイルに保存
    uv run python scripts/extract_led_positions_from_stl.py \\
        configs/meshes/LEDユニット_all.stl \\
        --output configs/leds_extracted.yaml

仕様:
- 連結成分のうち面数最大をカバー（光らない），残りを LED（発光面）と識別
- 各 LED の法線方向は extents 最小軸（板の厚み方向）から推定
- 法線の符号は「カバー中心 → LED 中心」のベクトルに基づき外向きに決定
- LED 位置は (Y, X, Z) 順で sort（下基板は X 順，右・左基板は Y 順になる想定）
"""

from __future__ import annotations

import argparse
from collections import defaultdict
from pathlib import Path
from typing import Any

import numpy as np
import trimesh
import yaml

# 既存 scene_baseline.yaml の LED radiance（Step 8-6 で校正済み）．新規 YAML 生成時に使う．
DEFAULT_RADIANCE = 2807.2469840551203


def _split_connected_components(mesh: trimesh.Trimesh) -> list[list[int]]:
    """頂点共有で連結する面群をグループ化（union-find）．

    Returns:
        各グループは face index のリスト．
    """
    n_verts = len(mesh.vertices)
    parent = list(range(n_verts))

    def find(x: int) -> int:
        while parent[x] != x:
            parent[x] = parent[parent[x]]
            x = parent[x]
        return x

    def union(a: int, b: int) -> None:
        ra, rb = find(a), find(b)
        if ra != rb:
            parent[ra] = rb

    for face in mesh.faces:
        union(int(face[0]), int(face[1]))
        union(int(face[1]), int(face[2]))

    groups: dict[int, list[int]] = defaultdict(list)
    for fi, face in enumerate(mesh.faces):
        groups[find(int(face[0]))].append(fi)
    return list(groups.values())


def _component_stats(mesh: trimesh.Trimesh, face_indices: list[int]) -> dict[str, Any]:
    """連結成分の統計（頂点・面数・bbox 中心・extents）を計算する．

    `face_indices` は STL 内の元 face index リストで，後で独立 STL として
    書き出す際に頂点・面の再マッピングに使う．
    """
    face_verts_idx = np.unique(mesh.faces[face_indices].flatten())
    verts = mesh.vertices[face_verts_idx]
    bmin = verts.min(axis=0)
    bmax = verts.max(axis=0)
    return {
        "n_verts": len(face_verts_idx),
        "n_faces": len(face_indices),
        "center": (bmin + bmax) / 2.0,
        "extents": bmax - bmin,
        "bmin": bmin,
        "bmax": bmax,
        "face_indices": face_indices,
    }


def _export_component_as_stl(mesh: trimesh.Trimesh, face_indices: list[int],
                              output_path: Path) -> None:
    """連結成分を独立 STL ファイルとして書き出す．

    元メッシュの face_indices に対応する頂点・面を抽出し，頂点 index を再マッピング
    してから新規 Trimesh として書き出す．元の座標系（ワールド）は保持する．
    """
    face_verts_idx = np.unique(mesh.faces[face_indices].flatten())
    old_to_new = {int(old): new for new, old in enumerate(face_verts_idx)}
    new_vertices = mesh.vertices[face_verts_idx]
    new_faces = np.array(
        [[old_to_new[int(v)] for v in mesh.faces[fi]] for fi in face_indices]
    )
    sub_mesh = trimesh.Trimesh(vertices=new_vertices, faces=new_faces, process=False)
    output_path.parent.mkdir(parents=True, exist_ok=True)
    sub_mesh.export(str(output_path), file_type="stl")


def _infer_led_direction(led_center: np.ndarray, led_extents: np.ndarray,
                        cover_center: np.ndarray) -> np.ndarray:
    """LED 発光面の法線方向を推定する．

    LED は薄板状（extents の最小軸が板の厚み）．発光方向はその最小軸方向で，
    カバー中心から LED 中心への差分の符号で外向きを決める．
    """
    min_axis = int(np.argmin(led_extents))
    direction = np.zeros(3)
    diff = led_center - cover_center
    sign = 1.0 if diff[min_axis] >= 0 else -1.0
    direction[min_axis] = sign
    return direction


def _led_size_xy(led_extents: np.ndarray) -> tuple[float, float]:
    """LED 発光面のサイズ（厚み以外の 2 軸）を取得する．

    パッケージ命名規則どおり (2.8, 3.5) mm 等が想定される．
    返り値は (短辺, 長辺) ではなく (X 方向幅, Y 方向幅) を意図しない．
    AreaLight の size は (width, height) の 2 値で，向きは direction で別途指定する．
    """
    min_axis = int(np.argmin(led_extents))
    other = [i for i in range(3) if i != min_axis]
    return (float(led_extents[other[0]]), float(led_extents[other[1]]))


def _classify_components(comps: list[dict]) -> tuple[list[dict], list[dict]]:
    """連結成分を「LED 候補」「カバー候補」に分類する．

    分類基準（SMD 2835 LED と LED 基板カバーの想定形状に基づく）:
      - LED 候補: extents の最大 ≤ 4mm かつ最小 ≤ 1.5mm（板状で 2.8×3.5×1mm 想定）
      - カバー候補: extents の最大 > 30mm かつ最小 ≤ 3mm（細長い板状、LED 基板カバー想定）

    上記いずれにも該当しない大きな成分（筐体、評価板、スマホ等）は除外する．
    """
    led_candidates = []
    cover_candidates = []
    for c in comps:
        e = c["extents"]
        if max(e) <= 4.0 and min(e) <= 1.5:
            led_candidates.append(c)
        elif max(e) > 30.0 and min(e) <= 3.0:
            cover_candidates.append(c)
    return led_candidates, cover_candidates


def extract_leds_from_mesh(mesh: trimesh.Trimesh) -> tuple[list[dict], list[dict]]:
    """STL メッシュから LED 群とカバー群を抽出する．

    複数の LED 基板カバー（下・右・左）を含む複合 STL に対応．
    各 LED について「最寄りカバー」を探し，そのカバー中心からの差分で
    発光方向（direction）の符号を決定する．

    Returns:
        (leds, covers) のタプル．
        - leds: 各 LED の {"center", "size", "direction", "extents", "nearest_cover_idx"} のリスト
        - covers: 各カバーの統計（center, extents 等）のリスト
    """
    groups = _split_connected_components(mesh)
    if len(groups) < 2:
        msg = f"STL に連結成分が {len(groups)} 個しかない．カバー + LED 群が必要"
        raise ValueError(msg)

    comps = [_component_stats(mesh, g) for g in groups]
    led_components, cover_components = _classify_components(comps)

    if len(cover_components) == 0:
        msg = "カバー候補が見つからない（細長い板状の連結成分が必要）"
        raise ValueError(msg)
    if len(led_components) == 0:
        msg = "LED 候補が見つからない（2.8x3.5x1mm 想定の小さな板状成分が必要）"
        raise ValueError(msg)

    cover_centers = np.array([c["center"] for c in cover_components])
    leds = []
    for c in led_components:
        # 最寄りカバー（中心間距離が最小）を選ぶ
        distances = np.linalg.norm(cover_centers - c["center"], axis=1)
        nearest_idx = int(np.argmin(distances))
        nearest_cover = cover_components[nearest_idx]

        size = _led_size_xy(c["extents"])
        direction = _infer_led_direction(c["center"], c["extents"], nearest_cover["center"])
        leds.append({
            "center": c["center"].tolist(),
            "size": list(size),
            "direction": direction.tolist(),
            "extents": c["extents"].tolist(),
            "nearest_cover_idx": nearest_idx,
        })
    return leds, cover_components


def _classify_cover_position(cover_center: np.ndarray) -> str:
    """カバー中心位置から `bottom` / `right` / `left` / `top` を判定する．

    判定基準: |Y| > |X| なら上下軸（Y 符号で bottom/top），そうでなければ左右軸（X 符号で left/right）．
    MiniTIAS の 3 基板配置 (下基板 Y<0, 右基板 X>0, 左基板 X<0) に対応．
    """
    x, y, _ = cover_center
    if abs(y) > abs(x):
        return "bottom" if y < 0 else "top"
    return "right" if x > 0 else "left"


def _assign_led_ids(leds: list[dict], covers: list[dict]) -> list[dict]:
    """各 LED に id を割り当てて並べ替える．

    - 各 LED の所属基板を cover の位置（bottom/right/left/top）で判定
    - 同基板内では「下基板: X 昇順」「右/左基板: Y 昇順」「上基板: X 昇順」で並べる
    - id 命名: `led_{label}_{index}` 例: `led_bottom_0`, `led_right_5`
    """
    cover_labels = [_classify_cover_position(np.array(c["center"])) for c in covers]
    grouped: dict[str, list[dict]] = defaultdict(list)
    for led in leds:
        label = cover_labels[led["nearest_cover_idx"]]
        grouped[label].append(led)

    # 並べ替えキー: bottom/top は X 昇順，right/left は Y 昇順
    sort_keys = {
        "bottom": lambda l: l["center"][0],
        "top": lambda l: l["center"][0],
        "right": lambda l: l["center"][1],
        "left": lambda l: l["center"][1],
    }
    # 出力順: bottom → right → left → top（読みやすい固定順）
    label_order = ["bottom", "right", "left", "top"]
    result = []
    for label in label_order:
        if label not in grouped:
            continue
        group = sorted(grouped[label], key=sort_keys[label])
        for i, led in enumerate(group):
            led_with_id = dict(led)
            led_with_id["assigned_id"] = f"led_{label}_{i}"
            result.append(led_with_id)
    return result


def _format_led_yaml_entries(leds: list[dict], radiance: float = DEFAULT_RADIANCE) -> list[dict]:
    """LED リストを scene_baseline.yaml の lights セクション形式に整形する．

    `assigned_id` フィールド（_assign_led_ids で付与）が無い場合は単純な番号で代用．
    """
    entries = []
    for i, led in enumerate(leds):
        led_id = led.get("assigned_id", f"led_{i:02d}")
        entries.append({
            "id": led_id,
            "type": "area",
            "shape": "rectangle",
            "center": [round(c, 4) for c in led["center"]],
            "size": [round(s, 4) for s in led["size"]],
            "direction": [round(d, 4) for d in led["direction"]],
            "radiance": [radiance, radiance, radiance],
        })
    return entries


def main() -> None:
    parser = argparse.ArgumentParser(description=__doc__.splitlines()[0])
    parser.add_argument("stl_files", nargs="+", type=Path,
                        help="解析する STL ファイル（複数可）")
    parser.add_argument("--output", type=Path, default=None,
                        help="出力 YAML パス（省略時は標準出力）")
    parser.add_argument("--radiance", type=float, default=DEFAULT_RADIANCE,
                        help=f"各 LED の radiance 値（既定 {DEFAULT_RADIANCE}）")
    parser.add_argument("--export-covers", type=Path, default=None, metavar="DIR",
                        help=(
                            "LED 基板カバーを bottom/right/left ごとに独立 STL として書き出す．"
                            " 指定したディレクトリに LED基板カバー{bottom,right,left}.stl を保存．"
                        ))
    parser.add_argument("--summary", action="store_true",
                        help="LED 位置サマリーを stderr に表示")
    args = parser.parse_args()

    all_leds: list[dict] = []
    for stl in args.stl_files:
        if not stl.is_file():
            msg = f"STL ファイルが見つかりません: {stl}"
            raise FileNotFoundError(msg)
        print(f"# Processing {stl}", flush=True)
        mesh = trimesh.load_mesh(str(stl))
        leds, covers = extract_leds_from_mesh(mesh)
        print(f"  detected {len(covers)} cover(s):", flush=True)
        for i, cover in enumerate(covers):
            label = _classify_cover_position(cover["center"])
            print(f"    cover[{i}] label={label:<6} center=({cover['center'][0]:+.2f}, "
                  f"{cover['center'][1]:+.2f}, {cover['center'][2]:+.2f}) "
                  f"extents=({cover['extents'][0]:.2f}, {cover['extents'][1]:.2f}, "
                  f"{cover['extents'][2]:.2f})", flush=True)

        # ID 付与（bottom/right/left/top + 並び順）
        leds_with_ids = _assign_led_ids(leds, covers)
        print(f"  detected {len(leds_with_ids)} LED(s) (id assigned):", flush=True)
        for led in leds_with_ids:
            print(f"    {led['assigned_id']:<16} "
                  f"center=({led['center'][0]:+.2f}, {led['center'][1]:+.2f}, "
                  f"{led['center'][2]:+.2f}) size={led['size']} dir={led['direction']}",
                  flush=True)
        all_leds.extend(leds_with_ids)

        # カバー書き出し（--export-covers 指定時のみ）．
        # ファイル名は既存 `LED基板カバー.stl` と整合させて日本語化（下/右/左）．
        jp_labels = {"bottom": "下", "right": "右", "left": "左", "top": "上"}
        if args.export_covers is not None:
            for cover in covers:
                label = _classify_cover_position(cover["center"])
                jp = jp_labels.get(label, label)
                out_path = args.export_covers / f"LED基板カバー{jp}.stl"
                _export_component_as_stl(mesh, cover["face_indices"], out_path)
                print(f"  exported cover [{label}] → {out_path}", flush=True)

    print(f"\n# Total: {len(all_leds)} LED(s) from {len(args.stl_files)} STL file(s)", flush=True)

    # YAML 形式で出力
    entries = _format_led_yaml_entries(all_leds, radiance=args.radiance)
    yaml_str = yaml.safe_dump(entries, default_flow_style=False,
                              allow_unicode=True, sort_keys=False)

    if args.output is None:
        print("\n# --- lights セクション（scene_baseline.yaml にコピペ可） ---")
        print(yaml_str)
    else:
        args.output.write_text(yaml_str, encoding="utf-8")
        print(f"\n# Wrote {len(entries)} LED entries to {args.output}", flush=True)


if __name__ == "__main__":
    main()