#!/usr/bin/env python3
"""Move one calibrated SO-101 follower joint a few degrees and return."""

import argparse
import os
import time
from pathlib import Path

from lerobot.robots.so_follower import SO101Follower, SO101FollowerConfig


JOINTS = (
    "shoulder_pan",
    "shoulder_lift",
    "elbow_flex",
    "wrist_flex",
    "wrist_roll",
    "gripper",
)


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument(
        "--port",
        default=os.environ.get("FOLLOWER_PORT"),
        help="Follower serial port; defaults to FOLLOWER_PORT from env/activate_lerobot.sh.",
    )
    parser.add_argument("--robot-id", default="my_follower", help="ID used when calibrating the follower.")
    parser.add_argument("--joint", choices=JOINTS, default="wrist_roll")
    parser.add_argument("--delta", type=float, default=3.0, help="Signed movement in degrees (gripper: percent).")
    parser.add_argument("--duration", type=float, default=1.0, help="Seconds for each outbound/return motion.")
    return parser.parse_args()


def interpolate(robot: SO101Follower, joint: str, start: float, delta: float, duration: float) -> None:
    steps = max(2, round(duration * 20))
    for step in range(1, steps + 1):
        target = start + delta * step / steps
        robot.bus.write("Goal_Position", joint, target)
        time.sleep(duration / steps)


def main() -> None:
    args = parse_args()
    if not args.port:
        raise SystemExit("Missing port. Run: source env/activate_lerobot.sh")
    port = Path(args.port)
    if not port.exists():
        raise SystemExit(f"Follower port is not connected: {port}")
    if not os.access(port, os.R_OK | os.W_OK):
        raise SystemExit(
            f"No read/write permission for {port}. Add the user to dialout, then log in again or run newgrp dialout."
        )
    if not 0 < abs(args.delta) <= 5:
        raise SystemExit("For this first-motion test, --delta must be non-zero and at most 5.")
    if not 0.5 <= args.duration <= 5:
        raise SystemExit("--duration must be between 0.5 and 5 seconds.")

    config = SO101FollowerConfig(port=str(port), id=args.robot_id)
    robot = SO101Follower(config)
    if not robot.calibration_fpath.is_file():
        raise SystemExit(
            f"Calibration file not found: {robot.calibration_fpath}\n"
            "Calibrate first with lerobot-calibrate, using the same --robot.id."
        )

    print(f"Calibration: {robot.calibration_fpath}")
    print(f"Test: {args.joint} {args.delta:+g}, then return to the measured start position.")
    print("Place the arm in a supported pose, clear the workspace, and keep power cutoff within reach.")
    if input("Type MOVE and press Enter to continue: ").strip() != "MOVE":
        raise SystemExit("Cancelled; no movement command was sent.")

    selected_torque_enabled = False
    try:
        # Use the motor bus directly: SOFollower.connect() configures and enables
        # every motor, which is inappropriate for a first single-joint test.
        robot.bus.connect()
        robot.bus.disable_torque()
        if not robot.bus.is_calibrated:
            raise RuntimeError("Motor calibration does not match the local calibration file.")

        start = float(robot.bus.read("Present_Position", args.joint))
        print(f"Start {args.joint}.pos: {start:.2f}")
        # Set a non-moving goal before enabling torque, avoiding a stale goal
        # register from a previous session.
        robot.bus.write("Goal_Position", args.joint, start)
        robot.bus.enable_torque(args.joint)
        selected_torque_enabled = True

        interpolate(robot, args.joint, start, args.delta, args.duration)
        time.sleep(0.3)
        outbound = float(robot.bus.read("Present_Position", args.joint))
        actual_delta = outbound - start
        print(f"Outbound {args.joint}.pos: {outbound:.2f} (actual {actual_delta:+.2f})")
        if actual_delta * args.delta <= 0 or abs(actual_delta) < abs(args.delta) * 0.5:
            raise RuntimeError(
                f"Joint feedback did not follow the command: requested {args.delta:+.2f}, observed {actual_delta:+.2f}."
            )

        interpolate(robot, args.joint, start + args.delta, -args.delta, args.duration)
        time.sleep(0.3)
        returned = float(robot.bus.read("Present_Position", args.joint))
        return_error = returned - start
        print(f"Returned {args.joint}.pos: {returned:.2f} (error {return_error:+.2f})")
        if abs(return_error) > 1.0:
            raise RuntimeError(f"Joint did not return within 1 degree of start; error is {return_error:+.2f}.")
        print("Motion test completed and verified by position feedback.")
    finally:
        if robot.bus.is_connected:
            try:
                if selected_torque_enabled:
                    print(f"Disabling torque on {args.joint}.")
                    robot.bus.disable_torque(args.joint, num_retry=5)
            finally:
                robot.bus.disconnect(disable_torque=False)


if __name__ == "__main__":
    main()