from pypot.dynamixel.io import DxlIO
import pypot, time

find_angles = False


def calculate_speeds_and_max_time(goal_positions, current_positions, base_speed=80):
    """
    Calculate speeds for synchronized motion and the required maximum time.

    :param goal_positions: List of target positions for each motor.
    :param current_positions: List of current positions for each motor.
    :param base_time_per_degree: Time it takes to move 1 degree at the base speed (e.g., 0.1 seconds at speed 10).
    :param base_speed: Base speed corresponding to the given time-per-degree.
    :return: (List of speeds, maximum time required for the movement)
    """
    distances = [abs(goal - current) for goal, current in zip(goal_positions, current_positions)]
    max_distance = max(distances) if max(distances) != 0 else 1  # Avoid division by zero

    # Calculate speeds for each motor, proportional to their distances
    speeds = [int((distance / max_distance) * base_speed) for distance in distances]

    # Calculate time-per-degree for each speed
    times = [
        distance * (0.9 / base_speed) * (base_speed / speed) if speed > 0 else 0
        for distance, speed in zip(distances, speeds)
    ]

    # Use the maximum time among all motors
    max_time = max(times) if times else 0
    print(speeds, max_time)
    return speeds, max_time



print(pypot.dynamixel.get_available_ports())

# with DxlIO('/dev/tty.usbserial-AL02L1UN', baudrate=1000000) as dxl_io:
with DxlIO('/dev/tty.usbserial-AH01FPVO', baudrate=1000000) as dxl_io:
    # motor_IDs = dxl_io.scan()
    # print(motor_IDs)

    # configure motors
    motor_IDs = [1,2,3,4]
    current_positions = dxl_io.get_present_position(motor_IDs)  # Initial positions of motors

    # Define goal positions
    goal_positions_list = [
        ([150, -32, -125, -45], 3),
        ([0, -71, -85, 50], 0),
        ([0, 46, 113, 116], 0),
        ([0, 31, 125, -70], 0),
        ([0, -68, 75, -25], 2), # touchdown
        ([0, -50, 115, 0], 2), # show off
        ([0, -72, 75, -25], 2), # pick up again
        ([0, -55, 90, -15], 0), # lift up
        ([150, -55, 90, -15], 10), #move sideways
        ([150, -65, 100, -10], 1), # touchdown
    ]

    num_motors = len(motor_IDs)
    print("Found", num_motors, "motors with current angles",  dxl_io.get_present_position(motor_IDs))


    # dxl_io.set_torque_limit(dict(zip(motor_IDs, ([0,0,0,100]))))
    # dxl_io.set_moving_speed(dict(zip(motor_IDs, [100,100,100,50])))
    # time.sleep(2)
    # dxl_io.set_goal_position(dict(zip(motor_IDs, [0,0,0,0])))
    # time.sleep(5)
    # dxl_io.set_goal_position(dict(zip(motor_IDs, [0,0,0,90])))
    # dxl_io.set_torque_limit(dict(zip(motor_IDs, ([0,0,0,0]))))
    # time.sleep(20)



    if find_angles:
        dxl_io.set_torque_limit(dict(zip(motor_IDs, ([0,0,0,0]))))
        while True:
            print(dxl_io.get_present_position(motor_IDs))
    else:
        # configure motors, set goal position as present goal position to not use last goal position.
        dxl_io.set_goal_position(dict(zip(motor_IDs, dxl_io.get_present_position(motor_IDs))))
        time.sleep(1)
        dxl_io.set_angle_limit(dict(zip(motor_IDs, ([(-150, 150), (-110, 110), (-128, 128), (-128, 128)]))))
        dxl_io.set_moving_speed(dict(zip(motor_IDs, ([10,10,10,10]))))
        dxl_io.set_torque_limit(dict(zip(motor_IDs, ([15,100,100,100]))))
        time.sleep(2)


        for goal_positions, extra_time in goal_positions_list:
            # Calculate speeds and dynamic sleep time
            speeds, max_time = calculate_speeds_and_max_time(goal_positions, current_positions)

            speeds[0] *= 4
            
            # Set the moving speeds for each motor
            dxl_io.set_moving_speed(dict(zip(motor_IDs, speeds)))
            
            # Command the motors to move to their goal positions
            dxl_io.set_goal_position(dict(zip(motor_IDs, goal_positions)))
            
            # Update current positions to the new goal
            current_positions = goal_positions.copy()
            
            # Sleep dynamically based on the maximum time required, plus extra time
            time.sleep(max_time + extra_time)


        dxl_io.set_torque_limit(dict(zip(motor_IDs, ([0,0,0,0]))))
        time.sleep(2)