#!/usr/bin/env bash
# 同时开两个 MuJoCo：O12 右手标定 vs 未标定 URDF
#
# 标定:   /home/lxqs/Linkerbot_URDF/012/right/linkerhand_o12_right_biaoding.urdf
# 未标定: /home/lxqs/Linkerbot_URDF/012/right/linkerhand_o12_right.urdf
#
# 启动前各自同步 MJCF（限位/轴/body 零位）；两窗各自用对应 URDF 的限位+mimic 映射。
# 共用 /o12/right/joint_cmd（12 路 SDK 弧度）。
#
# 用法:
#   ./hands/O12/scripts/run_o12_right_dual_compare.sh
#   ./hands/O12/scripts/run_o12_right_dual_compare.sh --no-pd-drive
#   ./hands/O12/scripts/run_o12_right_dual_compare.sh --pd-drive --slew-rad-s 4
set -eo pipefail
ROOT="$(cd "$(dirname "${BASH_SOURCE[0]}")/../../.." && pwd)"
cd "$ROOT"

URDF_DIR="${O12_RIGHT_URDF_DIR:-$ROOT/assets/vendor/Linkerbot_URDF/012/right}"
URDF_CALIBRATED="${O12_RIGHT_URDF_CALIBRATED:-$URDF_DIR/linkerhand_o12_right_biaoding.urdf}"
URDF_UNCALIBRATED="${O12_RIGHT_URDF_UNCALIBRATED:-$URDF_DIR/linkerhand_o12_right.urdf}"
MESHDIR="${O12_RIGHT_MESHDIR:-$URDF_DIR/meshes}"

PKG_SRC="$ROOT/src/linkerhand-sim/linker_hand_mujoco_ros2/linker_hand_mujoco_ros2"
O12_DIR="$PKG_SRC/urdf/O12/linker_hand_o12_right"
TEMPLATE_XML="$O12_DIR/linker_hand_o12_right.xml"
CALIB_XML="$O12_DIR/linker_hand_o12_right.xml"
STOCK_XML="$O12_DIR/linker_hand_o12_right_stock.xml"
CALIB_URDF_PKG="$O12_DIR/linkerhand_o12_right.urdf"
STOCK_URDF_PKG="$O12_DIR/linkerhand_o12_right_stock.urdf"

if [[ ! -f "$URDF_CALIBRATED" ]]; then
  echo "[error] missing calibrated URDF: $URDF_CALIBRATED" >&2
  exit 1
fi
if [[ ! -f "$URDF_UNCALIBRATED" ]]; then
  echo "[error] missing uncalibrated URDF: $URDF_UNCALIBRATED" >&2
  exit 1
fi
if [[ ! -f "$TEMPLATE_XML" ]]; then
  echo "[error] missing O12 MJCF template: $TEMPLATE_XML" >&2
  exit 1
fi
if [[ ! -d "$MESHDIR" ]]; then
  echo "[error] missing meshes: $MESHDIR" >&2
  exit 1
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

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 "[o12-right-dual] sync calibrated + uncalibrated MJCF from Linkerbot_URDF"
python3 - "$URDF_CALIBRATED" "$URDF_UNCALIBRATED" "$MESHDIR" "$TEMPLATE_XML" \
  "$CALIB_XML" "$STOCK_XML" "$CALIB_URDF_PKG" "$STOCK_URDF_PKG" <<'PY'
import math
import re
import shutil
import sys
from pathlib import Path

import mujoco
import numpy as np

urdf_cal = Path(sys.argv[1])
urdf_unc = Path(sys.argv[2])
meshdir = Path(sys.argv[3])
template_xml = Path(sys.argv[4])
calib_xml = Path(sys.argv[5])
stock_xml = Path(sys.argv[6])
calib_urdf_pkg = Path(sys.argv[7])
stock_urdf_pkg = Path(sys.argv[8])

BODY_FROM_JOINT = {
    "thumb_cmc_roll": "thumb_metacarpals_base1",
    "thumb_cmc_yaw": "thumb_metacarpals_base2",
    "thumb_cmc_pitch": "thumb_metacarpals",
    "thumb_mcp": "thumb_proximal",
    "thumb_dip": "thumb_distal",
    "index_mcp_roll": "index_metacarpals",
    "index_mcp_pitch": "index_proximal",
    "index_pip": "index_middle",
    "index_dip": "index_distal",
    "middle_mcp_roll": "middle_metacarpals",
    "middle_mcp_pitch": "middle_proximal",
    "middle_pip": "middle_middle",
    "middle_dip": "middle_distal",
    "ring_mcp_pitch": "ring_proximal",
    "ring_pip": "ring_middle",
    "ring_dip": "ring_distal",
    "pinky_mcp_pitch": "pinky_proximal",
    "pinky_pip": "pinky_middle",
    "pinky_dip": "pinky_distal",
}


def rpy_to_quat(roll, pitch, yaw):
    cr, sr = np.cos(roll / 2), np.sin(roll / 2)
    cp, sp = np.cos(pitch / 2), np.sin(pitch / 2)
    cy, sy = np.cos(yaw / 2), np.sin(yaw / 2)
    return np.array(
        [
            cr * cp * cy + sr * sp * sy,
            sr * cp * cy - cr * sp * sy,
            cr * sp * cy + sr * cp * sy,
            cr * cp * sy - sr * sp * cy,
        ]
    )


def localize_urdf(src: Path, dst: Path) -> None:
    text = src.read_text()
    text = re.sub(r"package://[^/]+/meshes/", "meshes/", text)
    dst.write_text(text)


def parse_urdf(path: Path) -> dict[str, dict]:
    text = path.read_text()
    joints: dict[str, dict] = {}
    for m in re.finditer(r"<joint\b([^>]*)>(.*?)</joint>", text, re.S):
        attrs, body = m.group(1), m.group(2)
        name_m = re.search(r'name="([^"]+)"', attrs)
        if not name_m:
            continue
        name = name_m.group(1)
        if "revolute" not in attrs:
            continue
        origin = re.search(r'<origin[^>]*xyz="([^"]+)"[^>]*rpy="([^"]+)"', body)
        lim = re.search(r'lower="([^"]+)"[^>]*upper="([^"]+)"', body)
        axis = re.search(r'<axis[^>]*xyz="([^"]+)"', body)
        xyz = [float(x) for x in origin.group(1).split()] if origin else [0.0, 0.0, 0.0]
        rpy = [float(x) for x in origin.group(2).split()] if origin else [0.0, 0.0, 0.0]
        joints[name] = {
            "lo": float(lim.group(1)),
            "hi": float(lim.group(2)),
            "xyz": xyz,
            "rpy": rpy,
            "axis": axis.group(1).strip() if axis else "0 0 1",
        }
    return joints


def apply_urdf_to_mjcf(template: str, joints: dict[str, dict], model_name: str) -> str:
    xml = re.sub(
        r'model="[^"]+"',
        f'model="{model_name}"',
        template,
        count=1,
    )
    for jname, info in joints.items():
        lo, hi = info["lo"], info["hi"]
        ax = info["axis"]
        jpat = re.compile(rf'<joint name="{re.escape(jname)}"[^>]*/?>')
        m = jpat.search(xml)
        if not m:
            raise SystemExit(f"missing joint in MJCF: {jname}")
        attrs = m.group(0)
        # keep damping/frictionloss if present
        damp = re.search(r'\bdamping="[^"]+"', attrs)
        fric = re.search(r'\bfrictionloss="[^"]+"', attrs)
        extra = ""
        if damp:
            extra += f" {damp.group(0)}"
        if fric:
            extra += f" {fric.group(0)}"
        new = (
            f'<joint name="{jname}" pos="0 0 0" axis="{ax}" '
            f'range="{lo:g} {hi:g}"{extra}/>'
        )
        xml = xml[: m.start()] + new + xml[m.end() :]

        act_pat = re.compile(
            rf'(<position name="{re.escape(jname)}_pos"[^>]*\bctrlrange=")[^"]+(")'
        )
        if act_pat.search(xml):
            xml = act_pat.sub(rf"\g<1>{lo:g} {hi:g}\2", xml, count=1)

    for jname, bname in BODY_FROM_JOINT.items():
        if jname not in joints:
            continue
        j = joints[jname]
        xyz, rpy = j["xyz"], j["rpy"]
        q = rpy_to_quat(*rpy)
        pos = f"{xyz[0]:g} {xyz[1]:g} {xyz[2]:g}"
        m = re.search(rf'<body name="{re.escape(bname)}"[^>]*>', xml)
        if not m:
            raise SystemExit(f"missing body in MJCF: {bname}")
        if np.allclose(q, [1, 0, 0, 0], atol=1e-9):
            new = f'<body name="{bname}" pos="{pos}">'
        else:
            quat = "{:.8g} {:.8g} {:.8g} {:.8g}".format(*q)
            new = f'<body name="{bname}" pos="{pos}" quat="{quat}">'
        xml = xml[: m.start()] + new + xml[m.end() :]
    return xml


# meshes → 包内（两窗共用）
mesh_out = calib_xml.parent / "meshes"
mesh_out.mkdir(exist_ok=True)
for f in meshdir.glob("*"):
    if f.is_file():
        shutil.copy2(f, mesh_out / f.name)

localize_urdf(urdf_cal, calib_urdf_pkg)
localize_urdf(urdf_unc, stock_urdf_pkg)
shutil.copy2(urdf_cal, calib_urdf_pkg.parent / "linkerhand_o12_right_biaoding.urdf")

template = template_xml.read_text()
cal_joints = parse_urdf(urdf_cal)
unc_joints = parse_urdf(urdf_unc)

calib_xml.write_text(
    apply_urdf_to_mjcf(template, cal_joints, "linkerhand_o12_right_calibrated")
)
stock_xml.write_text(
    apply_urdf_to_mjcf(template, unc_joints, "linkerhand_o12_right_uncalibrated")
)
mujoco.MjModel.from_xml_path(str(calib_xml))
mujoco.MjModel.from_xml_path(str(stock_xml))

print(f"[o12-right-dual] calibrated URDF: {urdf_cal}")
print(f"[o12-right-dual] calibrated MJCF: {calib_xml}")
print(f"[o12-right-dual] uncalibrated URDF: {urdf_unc}")
print(f"[o12-right-dual] uncalibrated MJCF: {stock_xml}")
for j in ("thumb_cmc_roll", "thumb_cmc_yaw", "index_mcp_roll", "index_mcp_pitch"):
    c, u = cal_joints[j], unc_joints[j]
    print(
        f"[o12-right-dual] {j}: cal lim=[{c['lo']:g},{c['hi']:g}] rpy={c['rpy']} | "
        f"unc lim=[{u['lo']:g},{u['hi']:g}] rpy={u['rpy']}"
    )
PY

EXTRA=("$@")

cleanup() {
  [[ -n "${PID_CALIB:-}" ]] && kill "$PID_CALIB" 2>/dev/null || true
  [[ -n "${PID_UNCAL:-}" ]] && kill "$PID_UNCAL" 2>/dev/null || true
}
trap cleanup EXIT INT TERM

echo "[o12-right-dual] launch calibrated window"
"$ROOT/hands/O12/scripts/run_o12_right_mujoco_plan.sh" \
  --xml "$CALIB_XML" \
  --mapping-urdf "$URDF_CALIBRATED" \
  --o12-urdf "$URDF_CALIBRATED" \
  --node-name o12_right_calibrated \
  --window-title "O12 right 标定 (biaoding)" \
  --cmd-topic /o12/right/joint_cmd \
  --sim-state-topic /sim/mujoco/o12/right/calibrated/joint_states \
  "${EXTRA[@]}" &
PID_CALIB=$!

sleep 1

echo "[o12-right-dual] launch uncalibrated window"
"$ROOT/hands/O12/scripts/run_o12_right_mujoco_plan.sh" \
  --xml "$STOCK_XML" \
  --mapping-urdf "$URDF_UNCALIBRATED" \
  --o12-urdf "$URDF_UNCALIBRATED" \
  --node-name o12_right_uncalibrated \
  --window-title "O12 right 未标定 (stock)" \
  --cmd-topic /o12/right/joint_cmd \
  --sim-state-topic /sim/mujoco/o12/right/uncalibrated/joint_states \
  "${EXTRA[@]}" &
PID_UNCAL=$!

echo "[o12-right-dual] both started. shared cmd: /o12/right/joint_cmd (12 rad)"
echo "[o12-right-dual] calibrated:   $URDF_CALIBRATED"
echo "[o12-right-dual] uncalibrated: $URDF_UNCALIBRATED"
echo "[o12-right-dual] Ctrl+C stops both."
wait "$PID_CALIB" "$PID_UNCAL"