"""XR 手柄 UDP 数据模拟发送器。

用于在没有真实 PICO/XR 手柄时向 udp_controller_receiver 发送测试数据。
默认轨迹按 +X/-X、+Y/-Y、+Z/-Z 成对扫轴，便于用肉眼检查机械臂方向映射。
"""

import argparse
import json
import math
import socket
import time


BASE_POS = [0.0, 1.0, 0.0]
# 方向名称会打印到终端，方便现场观察机械臂运动时对照坐标轴。
AXIS_DIRECTIONS = (
    ("XR +X", (1.0, 0.0, 0.0)),
    ("XR -X", (-1.0, 0.0, 0.0)),
    ("XR +Y", (0.0, 1.0, 0.0)),
    ("XR -Y", (0.0, -1.0, 0.0)),
    ("XR +Z", (0.0, 0.0, 1.0)),
    ("XR -Z", (0.0, 0.0, -1.0)),
)
AXIS_SWEEP_SEGMENTS = (
    ("XR +X  right", (1.0, 0.0, 0.0), "hold"),
    ("XR -X  left", (-1.0, 0.0, 0.0), "hold"),
    ("CENTER / X done", (0.0, 0.0, 0.0), "center"),
    ("XR +Y  up", (0.0, 1.0, 0.0), "hold"),
    ("XR -Y  down", (0.0, -1.0, 0.0), "hold"),
    ("CENTER / Y done", (0.0, 0.0, 0.0), "center"),
    ("XR +Z  back", (0.0, 0.0, 1.0), "hold"),
    ("XR -Z  front", (0.0, 0.0, -1.0), "hold"),
    ("CENTER / Z done", (0.0, 0.0, 0.0), "center"),
)
RPY_DIRECTIONS = (
    ("XR roll +", (1.0, 0.0, 0.0)),
    ("XR roll -", (-1.0, 0.0, 0.0)),
    ("XR pitch +", (0.0, 1.0, 0.0)),
    ("XR pitch -", (0.0, -1.0, 0.0)),
    ("XR yaw +", (0.0, 0.0, 1.0)),
    ("XR yaw -", (0.0, 0.0, -1.0)),
)


def _validate_positive(name: str, value: float) -> float:
    if value <= 0.0:
        raise argparse.ArgumentTypeError(f"{name} must be > 0")
    return value


def _positive_float(name: str):
    def parser(value: str) -> float:
        return _validate_positive(name, float(value))

    return parser


def _position_from_direction(direction: tuple[float, float, float], amplitude: float) -> list[float]:
    return [
        BASE_POS[0] + amplitude * direction[0],
        BASE_POS[1] + amplitude * direction[1],
        BASE_POS[2] + amplitude * direction[2],
    ]


def _pattern_period(pattern: str, hold_seconds: float, center_seconds: float) -> float:
    if pattern == "axis_sweep":
        return sum(
            hold_seconds if segment_type == "hold" else center_seconds
            for _, _, segment_type in AXIS_SWEEP_SEGMENTS
        )
    if pattern == "axis_steps":
        return (hold_seconds + center_seconds) * len(AXIS_DIRECTIONS)
    return 5.0


def _rotation_pattern_period(pattern: str, hold_seconds: float, center_seconds: float) -> float:
    if pattern == "rpy_steps":
        return (hold_seconds + center_seconds) * len(RPY_DIRECTIONS)
    return 0.0


def _axis_sweep_position(
    t: float,
    amplitude: float,
    hold_seconds: float,
    center_seconds: float,
    initial_center_seconds: float,
) -> tuple[str, list[float]]:
    if t < initial_center_seconds:
        return "CENTER / lock origin", BASE_POS.copy()

    # +方向和-方向连续切换，比每次都很快回中心更容易看清实际运动方向。
    cycle_t = (t - initial_center_seconds) % _pattern_period(
        "axis_sweep",
        hold_seconds,
        center_seconds,
    )
    for label, direction, segment_type in AXIS_SWEEP_SEGMENTS:
        duration = hold_seconds if segment_type == "hold" else center_seconds
        if cycle_t < duration:
            return label, _position_from_direction(direction, amplitude)
        cycle_t -= duration
    return "CENTER", BASE_POS.copy()


def _axis_step_position(
    t: float,
    amplitude: float,
    hold_seconds: float,
    center_seconds: float,
    initial_center_seconds: float,
) -> tuple[str, list[float]]:
    if t < initial_center_seconds:
        return "CENTER / lock origin", BASE_POS.copy()

    # 每个方向由“保持目标点”和“回中心”两段组成，循环遍历六个方向。
    segment_seconds = hold_seconds + center_seconds
    cycle_t = (t - initial_center_seconds) % (segment_seconds * len(AXIS_DIRECTIONS))
    direction_index = int(cycle_t // segment_seconds)
    segment_t = cycle_t - direction_index * segment_seconds
    label, direction = AXIS_DIRECTIONS[direction_index]

    if segment_t >= hold_seconds:
        return f"CENTER after {label}", BASE_POS.copy()

    return label, _position_from_direction(direction, amplitude)


def _sine_position(t: float, amplitude: float, phase: float) -> tuple[str, list[float]]:
    return "SINE X", [
        amplitude * math.sin(2.0 * math.pi * 0.2 * t + phase),
        BASE_POS[1],
        BASE_POS[2],
    ]


def _euler_to_quaternion(roll: float, pitch: float, yaw: float) -> list[float]:
    cy = math.cos(yaw * 0.5)
    sy = math.sin(yaw * 0.5)
    cp = math.cos(pitch * 0.5)
    sp = math.sin(pitch * 0.5)
    cr = math.cos(roll * 0.5)
    sr = math.sin(roll * 0.5)
    return [
        sr * cp * cy - cr * sp * sy,
        cr * sp * cy + sr * cp * sy,
        cr * cp * sy - sr * sp * cy,
        cr * cp * cy + sr * sp * sy,
    ]


def _rpy_steps_quaternion(
    t: float,
    amplitude_rad: float,
    hold_seconds: float,
    center_seconds: float,
    initial_center_seconds: float,
) -> tuple[str, list[float]]:
    if t < initial_center_seconds:
        return "ROT center / lock origin", [0.0, 0.0, 0.0, 1.0]

    segment_seconds = hold_seconds + center_seconds
    period = segment_seconds * len(RPY_DIRECTIONS)
    cycle_t = (t - initial_center_seconds) % period
    direction_index = int(cycle_t // segment_seconds)
    segment_t = cycle_t - direction_index * segment_seconds
    label, direction = RPY_DIRECTIONS[direction_index]
    if segment_t >= hold_seconds:
        return f"ROT center after {label}", [0.0, 0.0, 0.0, 1.0]

    return label, _euler_to_quaternion(
        amplitude_rad * direction[0],
        amplitude_rad * direction[1],
        amplitude_rad * direction[2],
    )


def _rotation_quaternion(
    pattern: str,
    t: float,
    amplitude_rad: float,
    hold_seconds: float,
    center_seconds: float,
    initial_center_seconds: float,
) -> tuple[str, list[float]]:
    if pattern == "rpy_steps":
        return _rpy_steps_quaternion(
            t,
            amplitude_rad,
            hold_seconds,
            center_seconds,
            initial_center_seconds,
        )
    return "ROT none", [0.0, 0.0, 0.0, 1.0]


def _pattern_position(
    pattern: str,
    t: float,
    amplitude: float,
    hold_seconds: float,
    center_seconds: float,
    initial_center_seconds: float,
    phase: float,
) -> tuple[str, list[float]]:
    if pattern == "sine":
        return _sine_position(t, amplitude, phase)
    if pattern == "axis_steps":
        return _axis_step_position(
            t,
            amplitude,
            hold_seconds,
            center_seconds,
            initial_center_seconds,
        )
    return _axis_sweep_position(
        t,
        amplitude,
        hold_seconds,
        center_seconds,
        initial_center_seconds,
    )


def main() -> None:
    parser = argparse.ArgumentParser(description="通过 UDP 发送一段模拟 XR 左/右手柄数据。")
    parser.add_argument("--host", default="127.0.0.1")
    parser.add_argument("--port", type=int, default=15000)
    parser.add_argument("--seconds", type=_positive_float("seconds"), default=30.0)
    parser.add_argument("--hz", type=_positive_float("hz"), default=60.0)
    parser.add_argument("--amplitude", type=_positive_float("amplitude"), default=0.50)
    parser.add_argument(
        "--pattern",
        choices=("axis_sweep", "axis_steps", "sine"),
        default="axis_sweep",
        help="axis_sweep 成对扫轴最适合肉眼看方向；axis_steps 每个方向后回中心；sine 保留旧轨迹。",
    )
    parser.add_argument(
        "--hold-seconds",
        type=_positive_float("hold-seconds"),
        default=4.0,
        help="axis_sweep/axis_steps 模式下，每个方向保持多久。增大可让机械臂走得更明显。",
    )
    parser.add_argument(
        "--center-seconds",
        type=_positive_float("center-seconds"),
        default=1.2,
        help="axis_sweep/axis_steps 模式下，回中心段保持多久。",
    )
    parser.add_argument(
        "--initial-center-seconds",
        type=_positive_float("initial-center-seconds"),
        default=1.0,
        help="axis_sweep/axis_steps 模式下，开头保持中心点多久，用于让遥操节点锁定原点。",
    )
    parser.add_argument(
        "--trigger",
        type=float,
        default=0.0,
        help="发送的 trigger 值。方向调试默认不驱动夹爪。",
    )
    parser.add_argument(
        "--rotation-pattern",
        choices=("none", "rpy_steps"),
        default="none",
        help="可选姿态测试轨迹；none 保持 identity 四元数，rpy_steps 逐轴发送小角度姿态。",
    )
    parser.add_argument(
        "--rotation-amplitude-deg",
        type=_positive_float("rotation-amplitude-deg"),
        default=25.0,
        help="rpy_steps 模式下的姿态幅度，单位 degree。",
    )
    parser.add_argument("--hand", choices=("left", "right", "both"), default="right")
    parser.add_argument(
        "--both-mode",
        choices=("synchronized", "staggered"),
        default="synchronized",
        help="hand=both 时 synchronized 左右臂同时动；staggered 先左后右，便于逐只观察。",
    )
    args = parser.parse_args()

    sock = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
    dt = 1.0 / args.hz
    t0 = time.monotonic()

    hands = ("left", "right") if args.hand == "both" else (args.hand,)
    # 每只手柄独立发送同格式数据包，接收端根据 hand 字段分发左右话题。
    packets = {
        hand: {
            "t": 0.0,
            "hand": hand,
            "grip": False,
            "trigger": 0.0,
            "pos": BASE_POS.copy(),
            "quat": [0.0, 0.0, 0.0, 1.0],
        }
        for hand in hands
    }
    pattern_period = max(
        _pattern_period(args.pattern, args.hold_seconds, args.center_seconds),
        _rotation_pattern_period(args.rotation_pattern, args.hold_seconds, args.center_seconds),
    )
    one_hand_seconds = args.initial_center_seconds + pattern_period
    full_sweep_seconds = one_hand_seconds * len(hands) if args.both_mode == "staggered" else one_hand_seconds
    rotation_amplitude_rad = math.radians(args.rotation_amplitude_deg)
    print(
        "Sample UDP sender: "
        f"hand={args.hand}, both_mode={args.both_mode}, pattern={args.pattern}, "
        f"rotation_pattern={args.rotation_pattern}, "
        f"amplitude={args.amplitude:.3f}m, hold={args.hold_seconds:.2f}s, "
        f"center={args.center_seconds:.2f}s, seconds={args.seconds:.1f}s",
        flush=True,
    )
    if (args.pattern != "sine" or args.rotation_pattern != "none") and args.seconds < full_sweep_seconds:
        print(
            "提示：当前 seconds 不足以跑完一轮完整方向序列；"
            f"建议至少 {full_sweep_seconds:.1f}s。",
            flush=True,
        )
    last_labels = {hand: None for hand in hands}

    while time.monotonic() - t0 < args.seconds:
        t = time.monotonic() - t0
        # 默认用清晰的成对扫轴测试相对位移链路，便于肉眼确认坐标映射。
        for index, packet in enumerate(packets.values()):
            hand = packet["hand"]
            pattern_t = t
            if args.hand == "both" and args.both_mode == "staggered":
                hand_start_t = index * one_hand_seconds
                hand_end_t = hand_start_t + one_hand_seconds
                if t < hand_start_t:
                    label, pos = "WAIT / center", BASE_POS.copy()
                    rot_label, quat = "ROT wait / center", [0.0, 0.0, 0.0, 1.0]
                elif t >= hand_end_t:
                    label, pos = "DONE / center", BASE_POS.copy()
                    rot_label, quat = "ROT done / center", [0.0, 0.0, 0.0, 1.0]
                else:
                    pattern_t = t - hand_start_t
                    label, pos = _pattern_position(
                        args.pattern,
                        pattern_t,
                        args.amplitude,
                        args.hold_seconds,
                        args.center_seconds,
                        args.initial_center_seconds,
                        index * math.pi,
                    )
                    rot_label, quat = _rotation_quaternion(
                        args.rotation_pattern,
                        pattern_t,
                        rotation_amplitude_rad,
                        args.hold_seconds,
                        args.center_seconds,
                        args.initial_center_seconds,
                    )
            else:
                label, pos = _pattern_position(
                    args.pattern,
                    pattern_t,
                    args.amplitude,
                    args.hold_seconds,
                    args.center_seconds,
                    args.initial_center_seconds,
                    index * math.pi,
                )
                rot_label, quat = _rotation_quaternion(
                    args.rotation_pattern,
                    pattern_t,
                    rotation_amplitude_rad,
                    args.hold_seconds,
                    args.center_seconds,
                    args.initial_center_seconds,
                )

            combined_label = label if args.rotation_pattern == "none" else f"{label} | {rot_label}"
            if combined_label != last_labels[hand]:
                quat_text = [round(value, 4) for value in quat]
                print(
                    f"[{t:5.2f}s] {hand:<5} {combined_label:<46} "
                    f"pos={pos} quat={quat_text}",
                    flush=True,
                )
                last_labels[hand] = combined_label

            packet["t"] = t
            packet["grip"] = True
            packet["trigger"] = min(max(args.trigger, 0.0), 1.0)
            packet["pos"] = pos
            packet["quat"] = quat
            sock.sendto(json.dumps(packet).encode("utf-8"), (args.host, args.port))
        time.sleep(dt)

    for packet in packets.values():
        packet["grip"] = False
        packet["trigger"] = 0.0
        packet["pos"] = BASE_POS.copy()
        packet["quat"] = [0.0, 0.0, 0.0, 1.0]
        sock.sendto(json.dumps(packet).encode("utf-8"), (args.host, args.port))


if __name__ == "__main__":
    main()
