from machine import Pin, I2C, PWM
import time
import math

# =========================
# MPU6050 SETUP
# =========================

i2c = I2C(0, scl=Pin(1), sda=Pin(0), freq=400000)
MPU_ADDR = 0x68

time.sleep(0.1)
i2c.writeto_mem(MPU_ADDR, 0x6B, b'\x00')
time.sleep(0.1)

def read_raw(reg):
    data = i2c.readfrom_mem(MPU_ADDR, reg, 2)
    value = (data[0] << 8) | data[1]
    if value > 32767:
        value -= 65536
    return value

def get_angles():
    ax = read_raw(0x3B) / 16384
    ay = read_raw(0x3D) / 16384
    az = read_raw(0x3F) / 16384

    pitch = math.degrees(math.atan2(ay, az))
    roll = math.degrees(math.atan2(ax, az))

    return pitch, roll

# =========================
# SERVO SETUP
# =========================

servo_x = PWM(Pin(15))
servo_y = PWM(Pin(14))

servo_x.freq(50)
servo_y.freq(50)

MIN_DUTY = 1000
MAX_DUTY = 9000

def set_servo(servo, angle):
    angle = max(0, min(180, angle))
    duty = int(MIN_DUTY + (angle/180)*(MAX_DUTY-MIN_DUTY))
    servo.duty_u16(duty)

# =========================
# FILTER + PID SETTINGS
# =========================

alpha = 0.85

pitch_filtered = 0
roll_filtered = 0

# PID X (Roll)
kp_x = 0.8
ki_x = 0.05
kd_x = 0.03

# PID Y (Pitch)
kp_y = 0.8
ki_y = 0.05
kd_y = 0.03

ix = 0
iy = 0

prev_ex = 0
prev_ey = 0

# Output smoothing
out_fx = 0
out_fy = 0
out_alpha = 0.85

# =========================
# SOFTWARE ALIGNMENT MENU
# =========================

print("\nTVC 2-AXIS CONTROL")
print("Commands:")
print("x+ x- y+ y-")
print("done = start flight")
print("exit = stop")

x_align = 90
y_align = 90

while True:

    cmd = input("Align > ").lower()

    if cmd == "exit":
        raise SystemExit

    if cmd == "done":
        print("Flight Mode 🚀")
        break

    if "x+" in cmd:
        x_align += 2
    if "x-" in cmd:
        x_align -= 2

    if "y+" in cmd:
        y_align += 2
    if "y-" in cmd:
        y_align -= 2

    x_align = max(0, min(180, x_align))
    y_align = max(0, min(180, y_align))

    set_servo(servo_x, x_align)
    set_servo(servo_y, y_align)

    print("X:", x_align, "Y:", y_align)

# =========================
# FLIGHT STABILIZATION LOOP
# =========================

print("Stabilizing Flight 🚀")

last_time = time.ticks_ms()

while True:

    now = time.ticks_ms()
    dt = max(0.002, time.ticks_diff(now, last_time)/1000)
    last_time = now

    pitch, roll = get_angles()

    pitch_filtered = alpha*pitch_filtered + (1-alpha)*pitch
    roll_filtered = alpha*roll_filtered + (1-alpha)*roll

    # ---------- PID X (Roll) ----------
    ex = 0 - roll_filtered

    ix += ex * dt
    ix = max(-15, min(15, ix))

    dx = (ex - prev_ex) / dt

    out_x = kp_x*ex + ki_x*ix + kd_x*dx
    out_fx = out_alpha*out_fx + (1-out_alpha)*out_x

    # ---------- PID Y (Pitch) ----------
    ey = 0 - pitch_filtered

    iy += ey * dt
    iy = max(-15, min(15, iy))

    dy = (ey - prev_ey) / dt

    out_y = kp_y*ey + ki_y*iy + kd_y*dy
    out_fy = out_alpha*out_fy + (1-out_alpha)*out_y

    # Servo output
    sx = 90 + out_fx
    sy = 90 + out_fy

    set_servo(servo_x, sx)
    set_servo(servo_y, sy)

    prev_ex = ex
    prev_ey = ey

    print("Pitch:", round(pitch_filtered,2),
          "Roll:", round(roll_filtered,2))

    time.sleep(0.01)