#!/usr/bin/env python3
"""
关节模组一键验证入口（编排 urdf_to_mjcf → generate_schema → simulate_report）
====================================================================

把「URDF → MJCF → schema → 仿真报告」整条流水线串起来，一条命令跑完。
这相当于后端处理「用户上传一份 URDF + mesh」时干的事：你拿任意一份关节模组 URDF
丢进来，只要显式给出输入/输出关节，就能得到自包含 MJCF、schema JSON 和仿真报告。

三个步骤各自都能单独跑（见 docs/architecture.md 的「分步运行」），本脚本只是把它们按顺序编排，
并把产物统一放到同一个工作目录（默认 = URDF 所在目录）。

用法（在 scripts/ 目录下运行）：
  python3 validate_module.py \
      --urdf /path/to/joint_module.urdf \
      --input sun_input_joint --output carrier_output_joint

常用可选参数：
  --work-dir DIR     产物输出目录（默认 URDF 所在目录）
  --ratio N          减速比（URDF 输出关节无 <mimic> 时手动指定）
  --torque-limit T   输入力矩限位 [N·m]（默认 ±10，占位）
  --position-limit P 输入关节无 <limit> 时的位置限位 ±rad（默认 ±12π）
  --show             弹可视化窗口（默认无头模式）
  --plot             追加 report_curves.png 曲线图

产物（都在 work-dir 下）：
  <module_id>.xml   自包含 MJCF（meshdir="meshes"，网格已拷到 meshes/）
  <module_id>.json  schema JSON
  report.txt        仿真报告
  timeseries.csv    逐时间步观测
  report_curves.png 曲线图（--plot）
"""
import argparse
import os
import re
import subprocess
import sys


SCRIPTS_DIR = os.path.dirname(os.path.abspath(__file__))


def robot_name_of(urdf):
    """从 URDF 根元素 <robot name="..."> 取机器名，作为默认 module_id。"""
    with open(urdf, encoding="utf-8") as f:
        head = f.read(2048)
    m = re.search(r'<robot[^>]*\bname="([^"]+)"', head)
    return m.group(1) if m else os.path.splitext(os.path.basename(urdf))[0]


def run(cmd, what):
    print(f"\n== {what} ==")
    print("  " + " ".join(cmd))
    r = subprocess.run(cmd, cwd=SCRIPTS_DIR)
    if r.returncode != 0:
        sys.exit(f"[失败] {what}（退出码 {r.returncode}）")


def main():
    ap = argparse.ArgumentParser(description="关节模组一键验证（URDF→MJCF→schema→报告）")
    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("--work-dir", default=None, help="产物输出目录（默认 URDF 所在目录）")
    ap.add_argument("--module-id", default=None, help="模组唯一标识（默认 URDF 的 robot 名）")
    ap.add_argument("--ratio", type=float, default=None, help="手动指定减速比")
    ap.add_argument("--torque-limit", type=float, default=10.0, help="输入力矩限位 [N·m]")
    ap.add_argument("--position-limit", type=float, default=None, help="位置限位 ±rad")
    ap.add_argument("--load-torque", type=float, default=-3.0,
                    help="输出端恒值负载 [N·m]（负=阻力，默认 -3.0）")
    ap.add_argument("--damping", type=float, default=0.01, help="关节粘性阻尼（默认 0.01）")
    ap.add_argument("--mode", choices=["normal", "overload"], default="normal",
                    help="仿真模式（normal / overload，默认 normal）")
    ap.add_argument("--show", action="store_true", help="弹可视化窗口（默认无头）")
    ap.add_argument("--plot", action="store_true", help="导出曲线 PNG")
    args = ap.parse_args()

    urdf = os.path.abspath(args.urdf)
    if not os.path.isfile(urdf):
        sys.exit(f"找不到 URDF：{urdf}")

    work_dir = os.path.abspath(args.work_dir) if args.work_dir else os.path.dirname(urdf)
    os.makedirs(work_dir, exist_ok=True)

    module_id = args.module_id or robot_name_of(urdf)
    xml_path = os.path.join(work_dir, module_id + ".xml")
    json_path = os.path.join(work_dir, module_id + ".json")
    mesh_dir = os.path.join(work_dir, "meshes")

    py = sys.executable

    # 1) URDF → MJCF
    cmd1 = [py, os.path.join(SCRIPTS_DIR, "urdf_to_mjcf.py"),
            "--urdf", urdf, "--out", xml_path, "--meshdir", mesh_dir,
            "--input", args.input, "--output", args.output,
            "--damping", str(args.damping),
            "--torque-limit", str(args.torque_limit)]
    run(cmd1, "① URDF → MJCF")

    # 2) URDF → schema
    cmd2 = [py, os.path.join(SCRIPTS_DIR, "generate_schema.py"),
            "--urdf", urdf, "--input", args.input, "--output", args.output,
            "--model", os.path.basename(xml_path), "--out", json_path,
            "--module-id", module_id,
            "--torque-limit", str(args.torque_limit),
            "--load-torque", str(args.load_torque),
            "--damping", str(args.damping),
            "--mode", args.mode]
    if args.ratio is not None:
        cmd2 += ["--ratio", str(args.ratio)]
    if args.position_limit is not None:
        cmd2 += ["--position-limit", str(args.position_limit)]
    run(cmd2, "② URDF → schema")

    # 3) schema → 仿真报告
    cmd3 = [py, os.path.join(SCRIPTS_DIR, "simulate_report.py"),
            "--schema", json_path, "--headless"]
    if args.plot:
        cmd3 += ["--plot"]
    if args.show:
        cmd3.remove("--headless")
    run(cmd3, "③ 仿真 + 报告")

    print("\n全部完成。产物：")
    print(f"  MJCF   : {xml_path}")
    print(f"  schema : {json_path}")
    print(f"  报告   : {os.path.join(work_dir, 'report.txt')}")
    print(f"  观测   : {os.path.join(work_dir, 'timeseries.csv')}")
    if args.plot:
        print(f"  曲线   : {os.path.join(work_dir, 'report_curves.png')}")


if __name__ == "__main__":
    main()