#-----------------------------------------------------------------------------#
# Add libraries 
from machine import Pin, PWM
from time import sleep
from math import sqrt, acos, atan2, sin, cos, pi # For inverse kinematic
#-----------------------------------------------------------------------------#
# Function to set angle to servomotors
def angle2duty(angle):
    # Convert angle to duty cycle
    min_us = 500   # for 0°
    max_us = 2500  # for 180°
    us = min_us + (angle / 180) * (max_us - min_us)
    duty = int((us / 20000) * 1023)  # 20ms period → duty cycle (0–1023)
    return duty
#-----------------------------------------------------------------------------#
def inverse_kinematic(x,z,L1,L2):
    r = sqrt(x**2 + z**2)  # Distance to foot
    # Check reachability
    if r > (L1 + L2):
        return None, None, False
    # Law of Cosines for q2
    cos_q2 = (x**2 + z**2 - L1**2 - L2**2)/(2*L1*L2)
    q2 = acos(cos_q2)
    # Law of Sines/geometry for q1
    alpha = atan2(z, x)
    beta = atan2(L2*sin(q2), L1+L2*cos(q2))
    q1 = alpha - beta
    q1 = q1 - pi/4
    # Convert to positive if needed
    if q1 < 0:
        q1 = abs(pi+q1)
    if q2 < 0:
        q2 = abs(pi+q2)
    return q1, q2, True
#-----------------------------------------------------------------------------#
def rad2deg(angle_rad):
    angle_deg = angle_rad*180/pi
    return angle_deg
#-----------------------------------------------------------------------------#
def rotXYZ(rpy):
    roll, pitch, yaw = rpy
    cr = cos(roll)
    sr = sin(roll)
    cp = cos(pitch)
    sp = sin(pitch)
    cy = cos(yaw)
    sy = sin(yaw)

    R = [
        [cy*cp, cy*sp*sr - sy*cr, cy*sp*cr + sy*sr],
        [sy*cp, sy*sp*sr + cy*cr, sy*sp*cr - cy*sr],
        [-sp,   cp*sr,             cp*cr]
    ]
    return R
#-----------------------------------------------------------------------------#
def matvec3x3(R, v):
    return [
        R[0][0]*v[0] + R[0][1]*v[1] + R[0][2]*v[2],
        R[1][0]*v[0] + R[1][1]*v[1] + R[1][2]*v[2],
        R[2][0]*v[0] + R[2][1]*v[1] + R[2][2]*v[2]
    ]
#-----------------------------------------------------------------------------#
def ik_leg_floating(xyz_foot_world, rpy_base, legID, leg_origins, xyz_base, link_lengths):
    # 1. Obtener matriz de rotación del cuerpo
    R = rotXYZ(rpy_base)

    # 2. Obtener origen de la pierna en el mundo
    leg_origin_body = leg_origins[legID]
    origin_leg_world = matvec3x3(R, leg_origin_body)
    origin_leg_world = [
        origin_leg_world[0] + xyz_base[0],
        origin_leg_world[1] + xyz_base[1],
        origin_leg_world[2] + xyz_base[2]
    ]

    # 3. Posición relativa del pie respecto al origen de la pierna
    p_leg = [
        xyz_foot_world[0] - origin_leg_world[0],
        xyz_foot_world[1] - origin_leg_world[1],
        xyz_foot_world[2] - origin_leg_world[2]
    ]

    # 4. Solo usamos X-Z
    x = p_leg[0]
    z = p_leg[2]

    # 5. Longitudes
    L1 = link_lengths[0]
    L2 = link_lengths[1]

    # 6. Resolver cinemática inversa
    q1, q2, reachable = inverse_kinematic(x, z, L1, L2)

    if not reachable:
        return None, None, False
    else:
        return q1, q2, True
#-----------------------------------------------------------------------------#
# Function to lift a leg

def lift_a_leg(leg, l1, l2):
    x = 1.5
    z = -0.3
    reachable = False
    q1, q2 = 0.0, 0.0

    if leg == "ULL":
        q1, q2, reachable = legIK(x, z, l1, l2)

    elif leg == "URL":
        q1, q2, reachable = legIK(x, z, l1, l2)

    return q1, q2, reachable
#-----------------------------------------------------------------------------#



# Create PWM object at 50Hz
servo1_ULL = PWM(Pin(18), freq=50)
servo2_ULL = PWM(Pin(19), freq=50)
servo3_URL = PWM(Pin(21), freq=50)
servo4_URL = PWM(Pin(22), freq=50)
servo5_LLL = PWM(Pin(23), freq=50)
servo6_LLL = PWM(Pin(25), freq=50)
servo7_LRL = PWM(Pin(26), freq=50)
servo8_LRL = PWM(Pin(27), freq=50)

servos = [servo1_ULL,servo2_ULL,servo3_URL,servo4_URL,servo5_LLL,servo6_LLL,servo7_LRL,servo8_LRL]

# Longitudes de los eslabones
link_lengths = [1, 1]  # L1, L2

z_desired = 1
x_desired = 0

# Posición y orientación del cuerpo flotante
xyz_base = [x_desired, 0, z_desired]
rpy_base = [0, 0, 0]  # roll, pitch, yaw en radianes 


# Posiciones deseadas de los pies (una para cada pierna, en el mundo)
foot_targets = [
    [2,  1, 0],  # ULL
    [2, -1, 0],  # URL
    [-2, 1, 0],  # LLL
    [-2, -1, 0]  # LRL
]

# Orígenes de las 4 piernas (en coordenadas del cuerpo)
leg_origins = [
    [ 2, 1, 0],  # ULL
    [ 2,-1, 0],  # URL
    [-2, 1, 0],  # LLL
    [-2,-1, 0]   # LRL
]


while True:

    # Posición y orientación del cuerpo flotante
    xyz_base = [x_desired, 0, z_desired]
    rpy_base = [0, 0, 0]  # roll, pitch, yaw en radianes 

    # Para cada pierna
    for i in range(4):
        xyz_foot = foot_targets[i]

        q1, q2, reachable = ik_leg_floating(xyz_foot, rpy_base, i, leg_origins, xyz_base, link_lengths)

        if reachable:
            angle1 = rad2deg(q1)
            angle2 = rad2deg(q2)

            duty1 = angle2duty(angle1)
            duty2 = angle2duty(angle2)

            servos[2*i].duty(duty1)
            servos[2*i + 1].duty(duty2)

            print(f"Leg {i}: q1={q1:.2f}, q2={q2:.2f} → angles=({angle1:.1f}, {angle2:.1f}) → duty=({duty1}, {duty2})")
            print(x_desired)
            print(z_desired)
        else:
            print(f"Leg {i}: Target not reachable")