#!/usr/bin/env python3
"""把 Flexiv Elements 的 .traj（文本）转换成 arm_application 的 YAML 路点。"""

from __future__ import annotations

import argparse
import math
import re
import sys
from pathlib import Path
from typing import Any


def parse_flexiv_traj(text: str) -> list[dict[str, Any]]:
    """解析 Elements Studio 导出的散点 .traj 文件。"""
    blocks = re.split(r'pattern_params\s*\{', text)[1:]
    waypoints: list[dict[str, Any]] = []
    for block in blocks:
        name_m = re.search(r'pattern_name:\s*"([^"]+)"', block)
        vel_m = re.search(r'velocity:\s*"([^"]+)"', block)
        joints: dict[str, float] = {}
        for j in range(1, 8):
            m = re.search(rf'joint{j}:\s*"([^"]+)"', block)
            if m:
                joints[f'joint{j}'] = math.radians(float(m.group(1)))
        if len(joints) != 7:
            continue
        positions = [joints[f'joint{i}'] for i in range(1, 8)]
        waypoints.append({
            'name': name_m.group(1) if name_m else f'WPT{len(waypoints)}',
            'positions': positions,
            'velocity': float(vel_m.group(1)) if vel_m else 0.2,
        })
    return waypoints


def assign_times(waypoints: list[dict[str, Any]], min_seg: float = 3.0) -> None:
    """按关节空间的段长累加出 time_from_start。"""
    t = 0.0
    prev = None
    for wp in waypoints:
        pos = wp['positions']
        if prev is None:
            t += min_seg
        else:
            delta = max(abs(a - b) for a, b in zip(pos, prev))
            # 关节最大运动量每弧度约 4 秒，且不低于 min_seg
            t += max(min_seg, delta * 4.0)
        wp['time'] = round(t, 2)
        prev = pos


def to_yaml(
    waypoints: list[dict[str, Any]],
    *,
    joint_prefix: str = 'Rizon4-063989_',
    velocity_scale: float | None = None,
) -> str:
    if velocity_scale is None:
        velocity_scale = waypoints[0].get('velocity', 0.2) if waypoints else 0.2
    times = [wp['time'] for wp in waypoints]
    flat: list[float] = []
    for wp in waypoints:
        flat.extend(wp['positions'])
    pos_line = ', '.join(f'{v:.6f}' for v in flat)
    time_line = ', '.join(str(t) for t in times)
    names = ', '.join(f'  # {wp.get("name", "")}' if i == 0 else '' for i, wp in enumerate(waypoints))
    _ = names  # 名称保留在下面的注释里
    lines = [
        '# 由 Flexiv Elements .traj 自动转换而来（关节角：度 -> 弧度）',
        '# ROS 参数：扁平的 waypoint_positions + waypoint_times（不使用嵌套字典）',
        '# 按需修改 velocity_scale / loop_count（无需重新编译）。',
        'run_joint_trajectory:',
        '  ros__parameters:',
        f'    joint_prefix: "{joint_prefix}"',
        '    joint_names: [joint1, joint2, joint3, joint4, joint5, joint6, joint7]',
        f'    velocity_scale: {velocity_scale}',
        '    loop_count: 1                # 1=执行一次，0=无限循环，N=重复 N 次',
        '    pause_between_loops: 0.5',
        f'    waypoint_times: [{time_line}]',
        '    waypoint_positions: [',
    ]
    for i, wp in enumerate(waypoints):
        pos = ', '.join(f'{v:.6f}' for v in wp['positions'])
        lines.append(f'      {pos},  # {wp.get("name", f"WPT{i}")}')
    lines.append('    ]')
    return '\n'.join(lines) + '\n'


def main() -> int:
    ap = argparse.ArgumentParser(description=__doc__)
    ap.add_argument('input', type=Path, help='Flexiv Elements 的 .traj 文件')
    ap.add_argument('-o', '--output', type=Path, help='输出的 YAML 路径')
    ap.add_argument('--prefix', default='Rizon4-063989_', help='关节名前缀（现场 SN + _）')
    ap.add_argument('--min-seg', type=float, default=3.0, help='每段的最小秒数')
    args = ap.parse_args()

    text = args.input.read_text(encoding='utf-8')
    wps = parse_flexiv_traj(text)
    if not wps:
        print('错误：没有解析到任何路点', file=sys.stderr)
        return 1
    assign_times(wps, min_seg=args.min_seg)
    yaml_text = to_yaml(wps, joint_prefix=args.prefix)

    out = args.output
    if out:
        out.parent.mkdir(parents=True, exist_ok=True)
        out.write_text(yaml_text, encoding='utf-8')
        print(f'已写入 {len(wps)} 个路点 -> {out}')
    else:
        print(yaml_text)
    return 0


if __name__ == '__main__':
    raise SystemExit(main())