#!/usr/bin/env python3
"""
关节模组 schema 生成器（接口 2 的后端实现）
====================================================================

从一份关节模组 URDF 生成前端要的配置 JSON（schema），即 [schema.md](schema.md)
定义的观测契约：输入/输出关节名、减速比、限位、默认仿真参数。

输入/输出关节由用户**显式指定**（--input / --output）。因为自动识别在真实 URDF 上
不可靠（见 docs/architecture.md 说明：多级级联有多个正 multiplier 的 mimic，且 fixed 关节也没有
mimic，程序分不清哪级是末端输出）。

减速比来源优先级（同 schema.md 第 5 节）：
  1. mimic  —— 输出关节 <mimic multiplier> 的倒数（默认）
  2. manual —— --ratio 手动指定（URDF 无 mimic 时）

用法（在 scripts/ 目录下运行）：
  python3 generate_schema.py \
      --urdf ../urdf/planetary_joint_split_motor_demo.urdf \
      --input sun_input_joint --output carrier_output_joint \
      --model planetary_joint_split_motor_demo.xml \
      --out planetary_joint_split_motor_demo.json

依赖：仅标准库（xml.etree / json / argparse），不需要 trimesh / mujoco。
"""

import argparse
import json
import math
import os
import xml.etree.ElementTree as ET


# 默认仿真参数（与 schema.md / report_spec.md 保持一致）
DEFAULTS = {
    "timestep": 0.001,
    "duration": 4.0,
    "kp": 20.0,
    "kd": 0.3,
    "amplitude": 2.0 * math.pi,
    "frequency": 0.25,
    "load_torque": -3.0,
    "damping": 0.01,
}
# 输入关节无 <limit>（continuous）时的位置限位占位：±6 圈 = ±12π
DEFAULT_POSITION_LIMIT = 12.0 * math.pi
# 输入力矩限位占位（应填真实电机额定值，见 schema.md limits.torque）
DEFAULT_TORQUE_LIMIT = 10.0


def _float(attr):
    try:
        return float(attr)
    except (TypeError, ValueError):
        return None


def parse_joints(root):
    """返回 {关节名: <joint> 元素}。"""
    return {j.get("name"): j for j in root.findall("joint")}


def get_mimic(joint_el):
    """返回 (joint, multiplier, offset) 或 None。"""
    m = joint_el.find("mimic")
    if m is None:
        return None
    return {
        "joint": m.get("joint"),
        "multiplier": _float(m.get("multiplier")),
        "offset": _float(m.get("offset", 0.0)),
    }


def get_limit(joint_el):
    """返回 (lower, upper, velocity) 或 None。关节无 <limit>（continuous）返回 None。"""
    lim = joint_el.find("limit")
    if lim is None:
        return None
    return {
        "lower": _float(lim.get("lower")),
        "upper": _float(lim.get("upper")),
        "velocity": _float(lim.get("velocity")),
    }


def main():
    ap = argparse.ArgumentParser(description="关节模组 URDF → schema JSON")
    ap.add_argument("--urdf", required=True, help="输入 URDF 路径")
    ap.add_argument("--input", required=True, help="输入关节名（电机端，显式指定）")
    ap.add_argument("--output", required=True, help="输出关节名（模组末端，显式指定）")
    ap.add_argument("--model", default=None,
                    help="要引用进 schema.model 的 MJCF 文件名（默认 <robot名>.xml）")
    ap.add_argument("--out", default=None, help="schema 输出 JSON 路径（默认 <robot名>.json）")
    ap.add_argument("--module-id", default=None, help="module_id（默认 robot 名）")
    ap.add_argument("--name", default=None, help="显示名（默认 robot 名）")
    ap.add_argument("--ratio", type=float, default=None,
                    help="手动指定减速比（URDF 无 mimic 时用，gear_ratio_source=manual）")
    ap.add_argument("--torque-limit", type=float, default=DEFAULT_TORQUE_LIMIT,
                    help="输入力矩限位 [N·m]，默认 ±10（占位）")
    ap.add_argument("--position-limit", type=float, default=DEFAULT_POSITION_LIMIT,
                    help="输入关节无 <limit> 时的位置限位 ±rad，默认 ±12π")
    for k, v in DEFAULTS.items():
        ap.add_argument(f"--{k.replace('_', '-')}", type=float, default=v,
                        help=f"仿真参数 {k}（默认 {v}）")
    ap.add_argument("--mode", choices=["normal", "overload"], default="normal",
                    help="仿真模式（normal / overload，默认 normal）")
    args = ap.parse_args()

    tree = ET.parse(args.urdf)
    robot = tree.getroot()
    robot_name = robot.get("name")
    joints = parse_joints(robot)

    if args.input not in joints:
        ap.error(f"输入关节 '{args.input}' 在 URDF 中不存在；可用关节：{sorted(joints)}")
    if args.output not in joints:
        ap.error(f"输出关节 '{args.output}' 在 URDF 中不存在；可用关节：{sorted(joints)}")

    # ---- 减速比 ----
    mimic = get_mimic(joints[args.output])
    gear_ratio = None
    gear_ratio_source = None
    if mimic is not None and mimic["multiplier"] not in (None, 0.0):
        gear_ratio = 1.0 / abs(mimic["multiplier"])
        gear_ratio_source = "mimic"
    elif args.ratio is not None:
        gear_ratio = args.ratio
        gear_ratio_source = "manual"
    else:
        ap.error("输出关节没有 <mimic>，请用 --ratio 手动指定减速比")

    # ---- 限位 ----
    in_limit = get_limit(joints[args.input])
    if in_limit is not None and in_limit["lower"] is not None and in_limit["upper"] is not None:
        pos_lim = [in_limit["lower"], in_limit["upper"]]
    else:
        pos_lim = [-args.position_limit, args.position_limit]
    velocity = in_limit["velocity"] if in_limit is not None else None

    module_id = args.module_id or robot_name
    model = args.model or (module_id + ".xml")
    out_path = args.out or (module_id + ".json")

    schema = {
        "module_id": module_id,
        "name": args.name or robot_name,
        "model": model,
        "input_joint": args.input,
        "output_joint": args.output,
        "gear_ratio": gear_ratio,
        "gear_ratio_source": gear_ratio_source,
        "limits": {
            "position": pos_lim,
            "torque": [-args.torque_limit, args.torque_limit],
        },
        "simulation": {
            "timestep": args.timestep,
            "duration": args.duration,
            "kp": args.kp,
            "kd": args.kd,
            "amplitude": args.amplitude,
            "frequency": args.frequency,
            "load_torque": args.load_torque,
            "damping": args.damping,
            "mode": args.mode,
        },
    }
    if velocity is not None:
        schema["limits"]["velocity"] = [-velocity, velocity]

    with open(out_path, "w", encoding="utf-8") as f:
        json.dump(schema, f, ensure_ascii=False, indent=2)
        f.write("\n")

    print(f"已生成 schema: {os.path.abspath(out_path)}")
    print(f"  module_id   : {module_id}")
    print(f"  input_joint : {args.input}")
    print(f"  output_joint: {args.output}")
    print(f"  gear_ratio  : {gear_ratio:.6f}  (source={gear_ratio_source})")
    print(f"  position_lim: {pos_lim[0]:.4f} ~ {pos_lim[1]:.4f} rad")
    print(f"  torque_lim  : ±{args.torque_limit} N·m")
    print(f"  mode        : {args.mode}")
    print(f"  load_torque : {args.load_torque} N·m")
    print(f"  damping     : {args.damping}")


if __name__ == "__main__":
    main()