#!/usr/bin/env bash
# MuJoCo O30i 右手单窗仿真（V2_0819）
# 默认使用 V10 纯 URDF：保留对指锚点，调整新零位，保持线性换算。
# O30_RIGHT_URDF 可覆盖；启动姿态对应第五张照片的零位数组。
# 订阅 /cb_right_hand_control_cmd（20 维 SDK 0~255，顺序见 linkerhand-o30-ros2 JOINT_MAP）
# 发布 /sim/mujoco/o30/right/joint_state
#
# 用法:
#   ./hands/O30/scripts/run_o30_right_mujoco_plan.sh
#   ./hands/O30/scripts/run_o30_right_mujoco_plan.sh --pd-drive
#   O30_RIGHT_PREPARE_ONLY=1 ./hands/O30/scripts/run_o30_right_mujoco_plan.sh  # 只生成/校验模型
set -eo pipefail
ROOT="$(cd "$(dirname "${BASH_SOURCE[0]}")/../../.." && pwd)"
cd "$ROOT"

PKG_SRC="$ROOT/src/linkerhand-sim/linker_hand_mujoco_ros2/linker_hand_mujoco_ros2"
O30_DIR="$PKG_SRC/urdf/O30/linker_hand_o30_right"
URDF_SRC="${O30_RIGHT_URDF:-$O30_DIR/linkerhand_o30_right_zero_v10_calibrated.urdf}"
MESH_SRC="${O30_RIGHT_MESHDIR:-$O30_DIR/meshes}"
XML="$O30_DIR/linker_hand_o30_right.xml"
URDF_PKG="$O30_DIR/linkerhand_o30_right.urdf"

if [[ ! -f "$URDF_SRC" ]]; then
  echo "[error] missing O30 right URDF: $URDF_SRC" >&2
  exit 1
fi
if [[ ! -d "$MESH_SRC" ]]; then
  echo "[error] missing O30 right meshes: $MESH_SRC" >&2
  exit 1
fi

VENV_SITE="$ROOT/venv/lib/python3.12/site-packages"
if [[ -d "$VENV_SITE" ]]; then
  export PYTHONPATH="$VENV_SITE${PYTHONPATH:+:$PYTHONPATH}"
fi
export PYTHONPATH="$ROOT/src/linkerhand-sim/linker_hand_mujoco_ros2${PYTHONPATH:+:$PYTHONPATH}"

echo "[o30-right] sync MJCF from $URDF_SRC"
python3 - "$URDF_SRC" "$MESH_SRC" "$O30_DIR" "$XML" "$URDF_PKG" <<'PY'
import re
import shutil
import sys
import xml.etree.ElementTree as ET
from pathlib import Path

import mujoco

urdf_src = Path(sys.argv[1])
mesh_src = Path(sys.argv[2])
out_dir = Path(sys.argv[3])
xml_path = Path(sys.argv[4])
urdf_pkg = Path(sys.argv[5])

out_dir.mkdir(parents=True, exist_ok=True)
mesh_out = out_dir / "meshes"
mesh_out.mkdir(exist_ok=True)
for f in mesh_src.iterdir():
    if f.is_file() and f.resolve() != (mesh_out / f.name).resolve():
        shutil.copy2(f, mesh_out / f.name)

text = urdf_src.read_text()
text = re.sub(r"package://[^/]+/meshes/", "meshes/", text)
urdf_pkg.write_text(text)


def attrs_of(chunk: str) -> dict[str, str]:
    return dict(re.findall(r'(\w+)="([^"]*)"', chunk))


def parse_links(urdf: str) -> dict[str, dict]:
    links: dict[str, dict] = {}
    for m in re.finditer(r'<link\s+name="([^"]+)"[^>]*>(.*?)</link>', urdf, re.S):
        name, body = m.group(1), m.group(2)
        inertial = re.search(r"<inertial>(.*?)</inertial>", body, re.S)
        pos = "0 0 0"
        mass = "0.01"
        fullinertia = "1e-6 1e-6 1e-6 0 0 0"
        if inertial:
            ib = inertial.group(1)
            origin = re.search(r"<origin\b([^/]*)/>", ib)
            mass_m = re.search(r"<mass\b([^/]*)/>", ib)
            inertia = re.search(r"<inertia\b([^/]*)/>", ib)
            if origin:
                pos = attrs_of(origin.group(1)).get("xyz", pos)
            if mass_m:
                mass = attrs_of(mass_m.group(1)).get("value", mass)
            if inertia:
                ia = attrs_of(inertia.group(1))
                fullinertia = " ".join(
                    ia.get(k, "0") for k in ("ixx", "iyy", "izz", "ixy", "ixz", "iyz")
                )
        color = "0.7 0.7 0.7 1"
        vis = re.search(r"<visual>(.*?)</visual>", body, re.S)
        if vis:
            col = re.search(r'rgba="([^"]+)"', vis.group(1))
            if col:
                color = col.group(1)
        mesh = None
        mesh_scale = None
        mesh_m = re.search(r"<mesh\b([^>]*)/>", body)
        if mesh_m:
            ma = attrs_of(mesh_m.group(1))
            fn = ma.get("filename", "")
            mesh = Path(fn).name
            mesh_scale = ma.get("scale")
        links[name] = {
            "pos": pos,
            "mass": mass,
            "fullinertia": fullinertia,
            "color": color,
            "mesh": mesh,
            "mesh_scale": mesh_scale,
        }
    return links


def parse_joints(urdf: str) -> list[dict]:
    joints = []
    for m in re.finditer(r"<joint\b([^>]*)>(.*?)</joint>", urdf, re.S):
        attrs, body = m.group(1), m.group(2)
        name = re.search(r'name="([^"]+)"', attrs)
        jtype = re.search(r'type="([^"]+)"', attrs)
        if not name or not jtype or jtype.group(1) != "revolute":
            continue
        origin = re.search(r"<origin\b([^/]*)/>", body)
        axis = re.search(r"<axis\b([^/]*)/>", body)
        lim = re.search(r"<limit\b([^/]*)/>", body)
        parent = re.search(r'<parent[^>]*link="([^"]+)"', body)
        child = re.search(r'<child[^>]*link="([^"]+)"', body)
        oa = attrs_of(origin.group(1)) if origin else {}
        la = attrs_of(lim.group(1)) if lim else {}
        joints.append(
            {
                "name": name.group(1),
                "xyz": oa.get("xyz", "0 0 0"),
                "rpy": oa.get("rpy", "0 0 0"),
                "axis": attrs_of(axis.group(1)).get("xyz", "0 0 1") if axis else "0 0 1",
                "lo": float(la.get("lower", "0")),
                "hi": float(la.get("upper", "0")),
                "parent": parent.group(1) if parent else "",
                "child": child.group(1) if child else "",
            }
        )
    return joints


def rpy_to_quat(rpy: str) -> str | None:
    import math

    r, p, y = [float(x) for x in rpy.split()]
    if abs(r) + abs(p) + abs(y) < 1e-12:
        return None
    cr, sr = math.cos(r / 2), math.sin(r / 2)
    cp, sp = math.cos(p / 2), math.sin(p / 2)
    cy, sy = math.cos(y / 2), math.sin(y / 2)
    w = cr * cp * cy + sr * sp * sy
    x = sr * cp * cy - cr * sp * sy
    yq = cr * sp * cy + sr * cp * sy
    z = cr * cp * sy - sr * sp * cy
    return f"{w:.12g} {x:.12g} {yq:.12g} {z:.12g}"


links = parse_links(text)
joints = parse_joints(text)
by_parent: dict[str, list[dict]] = {}
for j in joints:
    by_parent.setdefault(j["parent"], []).append(j)

root = "hand_base_link"
# mesh name -> scale (from first link that uses it)
mesh_scale: dict[str, str | None] = {}
for info in links.values():
    if info["mesh"] and info["mesh"] not in mesh_scale:
        mesh_scale[info["mesh"]] = info["mesh_scale"]
meshes = sorted(mesh_scale.keys())

lines = [
    '<mujoco model="linkerhand_o30_right">',
    '  <compiler angle="radian" meshdir="meshes"/>',
    '  <option timestep="0.002" gravity="0 0 -9.81" integrator="implicitfast"/>',
    '  <statistic center="0.01 0 0.08" extent="0.16"/>',
    "  <default>",
    '    <joint damping="0.15" armature="0.01"/>',
    '    <position kp="80" kv="8" forcelimited="true" forcerange="-8 8"/>',
    '    <default class="visual">',
    '      <geom type="mesh" contype="0" conaffinity="0" group="2"/>',
    "    </default>",
    '    <default class="collision">',
    '      <geom type="mesh" contype="1" conaffinity="1" condim="3" group="3"',
    '            rgba="0 1 0 0.15" friction="1 0.5 0.01"',
    '            solref="0.004 1" solimp="0.95 0.99 0.001" margin="0"/>',
    "    </default>",
    "  </default>",
    "  <asset>",
]
for mesh in meshes:
    stem = Path(mesh).stem
    sc = mesh_scale.get(mesh)
    if sc:
        lines.append(f'    <mesh name="{stem}" file="{mesh}" scale="{sc}"/>')
    else:
        lines.append(f'    <mesh name="{stem}" file="{mesh}"/>')
lines.append("  </asset>")
lines.append("  <worldbody>")


def emit_body(link: str, indent: int, joint: dict | None) -> None:
    pad = "  " * indent
    info = links[link]
    if joint is None:
        lines.append(f'{pad}<body name="{link}">')
    else:
        quat = rpy_to_quat(joint["rpy"])
        pose = f'pos="{joint["xyz"]}"'
        if quat:
            pose += f' quat="{quat}"'
        lines.append(f'{pad}<body name="{link}" {pose}>')
    lines.append(
        f'{pad}  <inertial pos="{info["pos"]}" mass="{info["mass"]}" '
        f'fullinertia="{info["fullinertia"]}"/>'
    )
    if joint is not None:
        lines.append(
            f'{pad}  <joint name="{joint["name"]}" pos="0 0 0" axis="{joint["axis"]}" '
            f'range="{joint["lo"]:.12g} {joint["hi"]:.12g}"/>'
        )
    mesh = Path(info["mesh"]).stem if info["mesh"] else None
    if mesh:
        lines.append(
            f'{pad}  <geom class="visual" rgba="{info["color"]}" mesh="{mesh}"/>'
        )
        lines.append(f'{pad}  <geom class="collision" mesh="{mesh}"/>')
    for child_j in by_parent.get(link, []):
        emit_body(child_j["child"], indent + 1, child_j)
    lines.append(f"{pad}</body>")


emit_body(root, 2, None)
lines.append("  </worldbody>")
lines.append("  <contact>")
for j in joints:
    lines.append(f'    <exclude body1="{j["parent"]}" body2="{j["child"]}"/>')
for other in (
    "thumb_metacarpals_base2",
    "index_proximal",
    "middle_proximal",
    "ring_proximal",
    "pinky_proximal",
):
    lines.append(f'    <exclude body1="hand_base_link" body2="{other}"/>')
lines.append("  </contact>")
lines.append("  <actuator>")
for j in joints:
    lines.append(
        f'    <position name="{j["name"]}_pos" joint="{j["name"]}" '
        f'ctrlrange="{j["lo"]:.12g} {j["hi"]:.12g}"/>'
    )
lines.append("  </actuator>")
sdk_mapping = ET.fromstring(text).find('o30_sdk_mapping')
if sdk_mapping is not None:
    if sdk_mapping.get('version') != '1' or sdk_mapping.get('interpolation') != 'linear':
        raise ValueError('Unsupported O30 SDK mapping format')
    lines.append('  <custom>')
    for item in sdk_mapping.findall('map'):
        name = item.get('joint', '')
        if name not in {j['name'] for j in joints}:
            raise ValueError(f'Unknown mapped joint: {name}')
        sdk = item.get('sdk', '').split()
        radians = item.get('radians', '').split()
        if len(sdk) != len(radians):
            raise ValueError(f'O30 mapping knot count mismatch: {name}')
        data = ' '.join(format(float(v), '.16g') for pair in zip(sdk, radians) for v in pair)
        lines.append(f'    <numeric name="o30_sdk_map_{name}" data="{data}"/>')
    lines.append('  </custom>')
lines.append("</mujoco>")
xml_path.write_text("\n".join(lines) + "\n")
model = mujoco.MjModel.from_xml_path(str(xml_path))
from linker_hand_mujoco_ros2.utils.o30_right_joints import configure_o30_right_limits
configure_o30_right_limits(xml_path)
print(f"[o30-right] wrote {xml_path} njnt={model.njnt} nbody={model.nbody} ngeom={model.ngeom}")
PY

if [[ "${O30_RIGHT_PREPARE_ONLY:-0}" == "1" ]]; then
  exit 0
fi

set +u
# shellcheck disable=SC1091
source /opt/ros/jazzy/setup.bash
if [[ -f "$ROOT/install/setup.bash" ]]; then
  # shellcheck disable=SC1091
  source "$ROOT/install/setup.bash"
fi
set -u
# Keep this checkout ahead of installed copies after sourcing ROS workspaces.
export PYTHONPATH="$ROOT/src/linkerhand-sim/linker_hand_mujoco_ros2${PYTHONPATH:+:$PYTHONPATH}"

echo "[o30-right] urdf=$URDF_PKG"
echo "[o30-right] xml=$XML"
echo "[o30-right] cmd-topic=/cb_right_hand_control_cmd (20ch SDK 0~255, JOINT_MAP 顺序)"
echo "[o30-right] state-topic=/sim/mujoco/o30/right/joint_state"
echo "[o30-right] order: thumb_roll,yaw, idx/mid/ring/little_yaw, thumb/idx/mid/ring/little_root1, idx/mid/ring/little_root2, thumb/idx/mid/ring/little_tip"
echo "[o30-right] default mode=kinematic; 加 --pd-drive 看接触"

exec python3 -m linker_hand_mujoco_ros2.o30_right_plan_bridge \
  --xml "$XML" \
  --node-name o30_mujoco_plan_bridge_right \
  --window-title "O30i right V10 zero calibrated" \
  --cmd-topic /cb_right_hand_control_cmd \
  --sim-state-topic /sim/mujoco/o30/right/joint_state \
  --mapping g20_urdf \
  "$@"