from machine import Pin, PWM
import time
import network
import espnow

mainESC = PWM(Pin(18), freq=50)
tailESC = PWM(Pin(9), freq=50)

STOP_DUTY = 51     
FULL_DUTY = 115  

def set_throttle(percent):
    if percent < 0:
        percent = 0
    if percent > 100:
        percent = 100

    duty = STOP_DUTY + (FULL_DUTY - STOP_DUTY) * percent // 100
    mainESC.duty(int(duty))
    
def rotate(percent):
    if percent < 0:
        percent = 0
    if percent > 100:
        percent = 100

    duty = STOP_DUTY + (FULL_DUTY - STOP_DUTY) * percent // 100
    tailESC.duty(int(duty))

# ---- ARM ESC ----
rotate(100)
time.sleep(3)
rotate(0)
set_throttle(0)
time.sleep(3)

wlan = network.WLAN(network.STA_IF)
wlan.active(True)

e = espnow.ESPNow()
e.active(True)

print("Waiting for messages...")

# Setup servos
servoA = PWM(Pin(6), freq=50)
servoB = PWM(Pin(5), freq=50)
servoC = PWM(Pin(4), freq=50)

def angle_to_duty(angle):
    angle = max(0, min(180, angle))
    return int(26 + (angle / 180) * 102)

def set_servo(servo, angle):
    servo.duty(angle_to_duty(angle))
    print(angle)

import math

rotation_angle = 2  # 120 degrees in radians

def swash_rotated(collective, pitch, roll):
    # rotate pitch/roll vector
    new_pitch = pitch * math.cos(rotation_angle) + roll * math.sin(rotation_angle)
    new_roll  = -pitch * math.sin(rotation_angle) + roll * math.cos(rotation_angle)
    
    a = collective + new_pitch
    b = collective - new_pitch/2 + new_roll
    c = collective - new_pitch/2 - new_roll

    set_servo(servoA, a)
    set_servo(servoB, b)
    set_servo(servoC, c)

swash_rotated(90, 40, 0)

while True:
    host, msg = e.recv()
    if msg:
        data = eval(msg.decode(msg))
        collective = ((data[3]/4095)*90)+45
        pitch = ((data[1]/4095)*90)-45
        roll = ((data[0]/4095)*90)-45
        swash_rotated(collective, pitch, roll)
        print(collective, pitch, roll)
        throttle = (data[4]/4095)*100
        rot = (data[2]/4095)*100
        rotate(rot)
        set_throttle(throttle)
        print(throttle, rot)