"""Fidex ExoskeletonConfig for the parametric smith300_modular robot.

Bridges the CAD output (`fidex_config.json` emitted by param_arm.py) to the
FiducialExoskeletons schema (LinkConfig + GridBoards), mirroring exo_configs/so100_adhesive.py.

To run the fidex estimator on this robot:
  1. cp this file -> FiducialExoskeletons/exo_configs/smith300_modular.py
  2. provide an MJCF of the arm at base_xml_path (load smith300_modular.urdf in MuJoCo,
     name bodies L0..L6 to match mujoco_name, save via mj_saveLastXML).
  3. import SMITH300_MODULAR_CONFIG and feed it to the exo_utils pipeline.
Note dict is DICT_4X4_250 (set exo_utils.ARUCO_DICT accordingly — it defaults to 6X6).
"""
import json
import numpy as np
from cv2 import aruco

try:
    from .exoskeleton import ExoskeletonConfig, LinkConfig          # when placed in exo_configs/
except ImportError:                                                 # standalone (CAD dir) fallback
    from dataclasses import dataclass

    @dataclass
    class LinkConfig:
        mujoco_name: str; pybullet_name: str; robot_mesh_path: str; exo_mesh_path: str
        aruco_offset_pos: np.ndarray; aruco_offset_rot: np.ndarray
        aruco_board_name: str; board_length: float

    class ExoskeletonConfig:
        pass

CAD_DIR = "/data/cameron/repos/smith300_para_stuff"
_cfg = json.load(open(f"{CAD_DIR}/fidex_config.json"))
_dict = aruco.getPredefinedDictionary(aruco.DICT_4X4_250)


def _gridboard(c):
    gx, gy = c["grid"]
    n = gx * gy
    ids = c["ids"]
    if len(ids) < n:
        ids = list(range(ids[0], ids[0] + n))
    B = c["board_length"]
    unit = gx * 0.8 + (gx - 1) * 0.2          # marker:separation = 0.8:0.2
    mlen, sep = 0.8 * B / unit, 0.2 * B / unit
    return aruco.GridBoard((gx, gy), mlen, sep, _dict, np.array(ids, dtype=np.int32).reshape(-1, 1))


class Smith300ModularConfig(ExoskeletonConfig):
    name = "smith300_modular"
    base_xml_path = f"{CAD_DIR}/smith300_modular.mjcf"   # generate from the URDF (see header)
    compiler_meshdir = f"{CAD_DIR}/"
    aruco_boards = {c["aruco_board_name"]: f"{CAD_DIR}/{c['marker_image']}" for c in _cfg}
    aruco_board_objects = {c["aruco_board_name"]: _gridboard(c) for c in _cfg}
    links = {
        c["link"]: LinkConfig(
            mujoco_name=c["mujoco_name"], pybullet_name=c["link"],
            robot_mesh_path=c["robot_mesh_path"], exo_mesh_path=c["exo_mesh_path"],
            aruco_offset_pos=np.array(c["aruco_offset_pos"], dtype=float),
            aruco_offset_rot=np.array(c["aruco_offset_rot"], dtype=float),
            aruco_board_name=c["aruco_board_name"], board_length=c["board_length"],
        )
        for c in _cfg
    }


SMITH300_MODULAR_CONFIG = Smith300ModularConfig()

if __name__ == "__main__":
    cfg = SMITH300_MODULAR_CONFIG
    print(f"{cfg.name}: {len(cfg.links)} links, {len(cfg.aruco_board_objects)} boards")
    for k, lc in cfg.links.items():
        print(f"  {k}: board={lc.aruco_board_name} {lc.board_length*1000:.0f}mm pos(mm)={lc.aruco_offset_pos.tolist()}")