#!/usr/bin/env bash
# 同时开两个 MuJoCo：O6 左手标定 vs 未标定 URDF
#
# 未标定: /home/lxqs/linkerhand-urdf/o6/left
# 标定:   /home/lxqs/Linkerbot_URDF/06/LEFT（9/3 版，直接用 linker_hand_o6_left.xml）
#
# 两个窗口共用同一条 6 维 SDK 控制话题（默认 /cb_left_hand_control_cmd）。
#
# 用法:
#   ./hands/O6/scripts/run_o6_left_dual_compare.sh
#   ./hands/O6/scripts/run_o6_left_dual_compare.sh --pd-drive
#   ./hands/O6/scripts/run_o6_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/O6/reference/left_uncalibrated.urdf"
URDF_CALIBRATED="$ROOT/hands/O6/reference/left_calibrated.urdf"
PKG_O6="$ROOT/src/linkerhand-sim/linker_hand_mujoco_ros2/linker_hand_mujoco_ros2/urdf/O6/linker_hand_o6_left"
# 9/3 标定 MJCF：与 ./hands/O6/scripts/run_o6_left_mujoco_plan.sh 单开完全相同
CALIB_XML="$PKG_O6/linker_hand_o6_left.xml"
STOCK_XML="$PKG_O6/linker_hand_o6_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 (9/3): $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 "[o6-dual] calibrated  xml=$CALIB_XML (9/3, same as single-run)"
echo "[o6-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_yaw": "lh_thumb_metacarpals_base2"}


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 apply_urdf_to_mjcf(template: str, urdf_path: Path, model_name: str) -> str:
    joints = parse_urdf_limits(urdf_path)
    xml = template.replace(
        'model="linkerhand_o6_left"',
        f'model="{model_name}"',
        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 = math.cos(r * 0.5) * math.cos(p * 0.5) * math.cos(y * 0.5) + math.sin(
                r * 0.5
            ) * math.sin(p * 0.5) * math.sin(y * 0.5)
            x = math.sin(r * 0.5) * math.cos(p * 0.5) * math.cos(y * 0.5) - math.cos(
                r * 0.5
            ) * math.sin(p * 0.5) * math.sin(y * 0.5)
            yq = math.cos(r * 0.5) * math.sin(p * 0.5) * math.cos(y * 0.5) + math.sin(
                r * 0.5
            ) * math.cos(p * 0.5) * math.sin(y * 0.5)
            z = math.cos(r * 0.5) * math.cos(p * 0.5) * math.sin(y * 0.5) - math.sin(
                r * 0.5
            ) * math.sin(p * 0.5) * math.cos(y * 0.5)
            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_yaw_pos": "lh_thumb_cmc_yaw",
        "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_o6_left_uncalibrated")
)
mujoco.MjModel.from_xml_path(str(out_uncal))
print(f"[o6-dual] uncalibrated MJCF -> {out_uncal}")
PY

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 "[o6-dual] launch uncalibrated window"
"$ROOT/hands/O6/scripts/run_o6_left_mujoco_plan.sh" \
  --xml "$STOCK_XML" \
  --mapping-urdf "$URDF_UNCALIBRATED" \
  --node-name o6_left_uncalibrated \
  --window-title "O6 left 未标定 (linkerhand-urdf)" \
  --sim-state-topic /sim/mujoco/o6/left/uncalibrated/joint_state \
  "${EXTRA[@]}" &
PID_UNCAL=$!

sleep 1

echo "[o6-dual] launch calibrated window (9/3 MJCF)"
"$ROOT/hands/O6/scripts/run_o6_left_mujoco_plan.sh" \
  --xml "$CALIB_XML" \
  --mapping-urdf "$URDF_CALIBRATED" \
  --node-name o6_left_calibrated \
  --window-title "O6 left 标定 (Linkerbot_URDF 9/3)" \
  --sim-state-topic /sim/mujoco/o6/left/calibrated/joint_state \
  "${EXTRA[@]}" &
PID_CALIB=$!

echo "[o6-dual] both started. same cmd topic (default /cb_left_hand_control_cmd)."
echo "[o6-dual] Ctrl+C stops both."
wait "$PID_UNCAL" "$PID_CALIB"