import numpy as np
from scipy.optimize import minimize

find_angles = True


# Define variables that can be changed later
d1 = 26.12  # Adjust this value as needed (height of base)
a2 = 18.6  # Adjust this value as needed (length of first arm segment)
a3 = 18.6  # Adjust this value as needed (length of second arm segment)
a4 = 12.0  # Adjust this value as needed (length of end-effector link)

# Define the DH transformation matrix
def dh_matrix(theta, d, a, alpha):
    return np.array([
        [np.cos(theta), -np.sin(theta) * np.cos(alpha), np.sin(theta) * np.sin(alpha), a * np.cos(theta)],
        [np.sin(theta), np.cos(theta) * np.cos(alpha), -np.cos(theta) * np.sin(alpha), a * np.sin(theta)],
        [0, np.sin(alpha), np.cos(alpha), d],
        [0, 0, 0, 1]
    ])

# Desired position and orientation for the end-effector
desired_position = np.array([1, -6, -14.8])  # Target position
desired_orientation = np.array([
    [1, 0, 0],
    [0, 1, 0],
    [0, 0, -1]  # Z-axis of end-effector pointing downward
])

# Inverse Kinematics Cost Function
def ik_cost(angles):
    theta1, theta2, theta3, theta4 = angles
    
    # Compute forward kinematics based on the angles
    T1_0 = dh_matrix(theta1, d1, 0, np.pi/2)
    T2_1 = dh_matrix(theta2, 0, a2, 0)
    T3_2 = dh_matrix(theta3, 0, a3, 0)
    T4_3 = dh_matrix(theta4, 0, a4, 0)
    T_final = T1_0 @ T2_1 @ T3_2 @ T4_3

    
    # Position and orientation of the end-effector
    end_effector_pos = T_final[:3, 3]
    end_effector_orientation = T_final[:3, :3]

    print(T_final)
    
    # Calculate position error
    position_error = np.linalg.norm(end_effector_pos - desired_position)
    
    # Calculate orientation error (Frobenius norm of the rotation matrix difference)
    orientation_error = np.linalg.norm(end_effector_orientation - desired_orientation, ord='fro')
    
    # Total cost as a weighted sum of position and orientation errors
    return position_error + orientation_error

# Initial guess for joint angles
initial_guess = [0, 0, 0, 0]

# Joint angle limits in radians (convert degrees to radians)
bounds = [
    (np.radians(-180), np.radians(180)),   # ¦È2 limits
    # (np.radians(-33), np.radians(35)),    # ¦È1 limits
    (np.radians(-120), np.radians(120)),   # ¦È2 limits
    (np.radians(-120), np.radians(120)),     # ¦È3 limits
    (np.radians(-120), np.radians(120))     # ¦È4 limits
]

# Minimize the IK cost function with joint angle constraints
result = minimize(ik_cost, initial_guess, bounds=bounds)

if result.success:
    theta1_sol, theta2_sol, theta3_sol, theta4_sol = result.x
    print("\nInverse Kinematics Solution (Joint Angles):")
    print(f"Theta1: {np.degrees(theta1_sol):.2f}¡ã")
    print(f"Theta2: {np.degrees(theta2_sol):.2f}¡ã")
    print(f"Theta3: {np.degrees(theta3_sol):.2f}¡ã")
    print(f"Theta4: {np.degrees(theta4_sol):.2f}¡ã")
else:
    print("Inverse kinematics solver failed to find a solution.")



def forward_kinematics(angles):
    theta1, theta2, theta3, theta4 = angles

    T1_0 = dh_matrix(theta1, d1, 0, np.pi/2)
    T2_1 = dh_matrix(theta2, 0, a2, 0)
    T3_2 = dh_matrix(theta3, 0, a3, 0)
    T4_3 = dh_matrix(theta4, 0, a4, 0)

    T_final = T1_0 @ T2_1 @ T3_2 @ T4_3

    return T_final[:3, 3]


print(forward_kinematics((-7.77,-95.75, 32.11, -51.76)))



if find_angles:
    from pypot.dynamixel.io import DxlIO
    import pypot, time

    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]
        num_motors = len(motor_IDs)
        print("Found", num_motors, "motors with current angles",  dxl_io.get_present_position(motor_IDs))

        dxl_io.set_moving_speed(dict(zip(motor_IDs, ([0,26, 33, 133]))))
        time.sleep(2)
        dxl_io.set_angle_limit(dict(zip(motor_IDs, ([(-150, 150), (-75, 75), (-150, 150), (-150, 150)]))))
        time.sleep(2)
        dxl_io.set_torque_limit(dict(zip(motor_IDs, ([1,1,1,1]))))
        time.sleep(2)


        while True:
            print(forward_kinematics(dxl_io.get_present_position(motor_IDs)))