#!/usr/bin/env bash
# 同时开两个 MuJoCo：L6 v3.1 左手标定 vs 未标定 URDF
#
# 未标定: /home/lxqs/linkerhand-urdf/l6/left
# 标定:   /home/lxqs/Linkerbot_URDF/L6/LEFT（直接用 linker_hand_l6v31_left.xml，与单开相同）
#
# 两个窗口共用同一条 6 维 SDK 控制话题（/cb_left_hand_control_cmd）。
#
# 用法:
#   ./hands/L6/scripts/run_l6_left_dual_compare.sh
#   ./hands/L6/scripts/run_l6_left_dual_compare.sh --pd-drive
#   ./hands/L6/scripts/run_l6_left_dual_compare.sh --cmd-topic /cb_left_hand_control_cmd --speed 1.0
set -eo pipefail
ROOT="$(cd "$(dirname "${BASH_SOURCE[0]}")/../../.." && pwd)"
cd "$ROOT"

URDF_UNCALIBRATED="$ROOT/hands/L6/reference/left_uncalibrated.urdf"
URDF_CALIBRATED="$ROOT/hands/L6/reference/left_calibrated.urdf"
PKG_L6="$ROOT/src/linkerhand-sim/linker_hand_mujoco_ros2/linker_hand_mujoco_ros2/urdf/L6/linker_hand_l6v31_left"
CALIB_XML="$PKG_L6/linker_hand_l6v31_left.xml"
STOCK_XML="$PKG_L6/linker_hand_l6v31_left_stock.xml"

if [[ ! -f "$URDF_UNCALIBRATED" ]]; then
  echo "[error] missing uncalibrated URDF: $URDF_UNCALIBRATED" >&2
  exit 1
fi
if [[ ! -f "$URDF_CALIBRATED" ]]; then
  echo "[error] missing calibrated URDF: $URDF_CALIBRATED" >&2
  exit 1
fi
if [[ ! -f "$CALIB_XML" ]]; then
  echo "[error] missing calibrated MJCF: $CALIB_XML" >&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
PKG_SRC="$ROOT/src/linkerhand-sim/linker_hand_mujoco_ros2"
export PYTHONPATH="$PKG_SRC${PYTHONPATH:+:$PYTHONPATH}"

echo "[l6-left-dual] calibrated  xml=$CALIB_XML (same as single-run)"
echo "[l6-left-dual] build uncalibrated MJCF from $URDF_UNCALIBRATED"
python3 - "$CALIB_XML" "$URDF_UNCALIBRATED" "$STOCK_XML" <<'PY'
import math
import re
import sys
from pathlib import Path

template_xml = Path(sys.argv[1])
urdf_uncal = Path(sys.argv[2])
out_uncal = Path(sys.argv[3])

joint_body = {
    "lh_thumb_cmc_roll": "lh_thumb_metacarpals_base1",
    "lh_thumb_cmc_pitch": "lh_thumb_metacarpals",
}
MESHDIR_UNCAL = "/home/lxqs/linkerhand-urdf/l6/left/meshes"


def parse_urdf_limits(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)
        jtype_m = re.search(r'type="([^"]+)"', attrs)
        if not jtype_m or jtype_m.group(1) != "revolute":
            continue
        origin = re.search(r'<origin[^>]*xyz="([^"]+)"[^>]*rpy="([^"]+)"', body)
        lim = re.search(r'lower="([^"]+)"[^>]*upper="([^"]+)"', body)
        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)),
            "rpy": rpy,
        }
    return joints


def rpy_to_quat(r: float, p: float, y: float) -> tuple[float, float, float, float]:
    cr, sr = math.cos(r * 0.5), math.sin(r * 0.5)
    cp, sp = math.cos(p * 0.5), math.sin(p * 0.5)
    cy, sy = math.cos(y * 0.5), math.sin(y * 0.5)
    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 w, x, yq, z


def apply_urdf_to_mjcf(template: str, urdf_path: Path, model_name: str) -> str:
    joints = parse_urdf_limits(urdf_path)
    xml = template.replace(
        'model="linkerhand_l6v31_left"',
        f'model="{model_name}"',
        1,
    )
    xml = re.sub(
        r'(<compiler[^>]*meshdir=")[^"]+(")',
        rf"\1{MESHDIR_UNCAL}\2",
        xml,
        count=1,
    )

    for jname, info in joints.items():
        lo, hi = info["lo"], info["hi"]
        pat = rf'(<joint name="{re.escape(jname)}"[^>]*\srange=")([^"]+)(")'
        xml = re.sub(
            pat,
            lambda m, lo=lo, hi=hi: f"{m.group(1)}{lo} {hi}{m.group(3)}",
            xml,
            count=1,
        )
        r, p, y = info["rpy"]
        bname = joint_body.get(jname)
        if bname and abs(r) + abs(p) + abs(y) > 1e-9:
            w, x, yq, z = rpy_to_quat(r, p, y)
            quat = f"{w:.9f} {x:.9f} {yq:.9f} {z:.9f}"
            bpat = rf'(<body name="{re.escape(bname)}"[^>]*)(>)'

            def add_quat(m, quat=quat):
                tag = m.group(1)
                if 'quat="' in tag:
                    tag = re.sub(r'quat="[^"]+"', f'quat="{quat}"', tag, count=1)
                else:
                    tag = tag + f' quat="{quat}"'
                return tag + m.group(2)

            xml = re.sub(bpat, add_quat, xml, count=1)
        elif bname:
            bpat = rf'(<body name="{re.escape(bname)}"[^>]*)\squat="[^"]+"([^>]*>)'
            xml = re.sub(bpat, r"\1\2", xml, count=1)

    act_map = {
        "lh_thumb_cmc_roll_pos": "lh_thumb_cmc_roll",
        "lh_thumb_cmc_pitch_pos": "lh_thumb_cmc_pitch",
        "lh_index_mcp_pitch_pos": "lh_index_mcp_pitch",
        "lh_middle_mcp_pitch_pos": "lh_middle_mcp_pitch",
        "lh_ring_mcp_pitch_pos": "lh_ring_mcp_pitch",
        "lh_pinky_mcp_pitch_pos": "lh_pinky_mcp_pitch",
    }
    for act, jname in act_map.items():
        lo, hi = joints[jname]["lo"], joints[jname]["hi"]
        pat = rf'(<position name="{act}"[^>]*\sctrlrange=")([^"]+)(")'
        xml = re.sub(
            pat,
            lambda m, lo=lo, hi=hi: f"{m.group(1)}{lo} {hi}{m.group(3)}",
            xml,
            count=1,
        )
    return xml


import mujoco

out_uncal.write_text(
    apply_urdf_to_mjcf(template_xml.read_text(), urdf_uncal, "linkerhand_l6v31_left_uncalibrated")
)
mujoco.MjModel.from_xml_path(str(out_uncal))
print(f"[l6-left-dual] uncalibrated URDF: {urdf_uncal}")
print(f"[l6-left-dual] uncalibrated meshes: {MESHDIR_UNCAL}")
print(f"[l6-left-dual] uncalibrated MJCF -> {out_uncal}")
PY

echo "[l6-left-dual] calibrated  URDF: $URDF_CALIBRATED"
echo "[l6-left-dual] calibrated  meshes: $PKG_L6/meshes (package, 标定安装)"

EXTRA=("$@")

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

echo "[l6-left-dual] launch uncalibrated window"
"$ROOT/hands/L6/scripts/run_l6_left_mujoco_plan.sh" \
  --xml "$STOCK_XML" \
  --mapping-urdf "$URDF_UNCALIBRATED" \
  --node-name l6_left_uncalibrated \
  --window-title "L6 left 未标定 (linkerhand-urdf)" \
  --cmd-topic /cb_left_hand_control_cmd \
  --sim-state-topic /sim/mujoco/l6/left/uncalibrated/joint_state \
  "${EXTRA[@]}" &
PID_UNCAL=$!

sleep 1

echo "[l6-left-dual] launch calibrated window"
"$ROOT/hands/L6/scripts/run_l6_left_mujoco_plan.sh" \
  --xml "$CALIB_XML" \
  --mapping-urdf "$URDF_CALIBRATED" \
  --node-name l6_left_calibrated \
  --window-title "L6 left 标定 (Linkerbot_URDF)" \
  --cmd-topic /cb_left_hand_control_cmd \
  --sim-state-topic /sim/mujoco/l6/left/calibrated/joint_state \
  "${EXTRA[@]}" &
PID_CALIB=$!

echo "[l6-left-dual] both started. cmd topic: /cb_left_hand_control_cmd"
echo "[l6-left-dual] Ctrl+C stops both."
wait "$PID_UNCAL" "$PID_CALIB"