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

# OLED (SSD1306)
import ssd1306

# =============================
# I2C / MPU6050 / OLED
# =============================
i2c = I2C(0, scl=Pin(22), sda=Pin(21), freq=400000)

mpu = MPU6050(i2c, addr=0x68)
print("I2C scan:", [hex(x) for x in i2c.scan()])

# OLED is usually 0x3C in Wokwi
OLED_W = 128
OLED_H = 64
oled = ssd1306.SSD1306_I2C(OLED_W, OLED_H, i2c, addr=0x3C)

oled.fill(0)
oled.text("MPU + Servo", 0, 0)
oled.text("OLED ready", 0, 12)
oled.show()

# =============================
# Servo (GPIO 18)
# =============================
servo = PWM(Pin(18), freq=50)

def clamp(x, lo, hi):
    return lo if x < lo else hi if x > hi else x

SERVO_CENTER_US = 1500
SERVO_RANGE_US  = 500      # ~1000–2000 us
MAX_CMD_DEG     = 45.0     # authority

def servo_write_us(us):
    us = int(clamp(us, 900, 2100))
    duty = int(us * 65535 / 20000)  # 50Hz => 20ms
    servo.duty_u16(duty)

def servo_write_cmd(cmd_deg):
    cmd_deg = clamp(cmd_deg, -MAX_CMD_DEG, MAX_CMD_DEG)
    us = SERVO_CENTER_US + (cmd_deg / MAX_CMD_DEG) * SERVO_RANGE_US
    servo_write_us(us)
    return us  # return actual microseconds for display/debug

# Center servo on boot
servo_us = servo_write_cmd(0.0)
time.sleep_ms(300)

# =============================
# Attitude estimation (pitch)
# =============================
ALPHA = 0.995   # gyro-heavy

pitch = 0.0

def accel_pitch_deg(ax, ay, az):
    # pitch from accelerometer
    return math.degrees(math.atan2(-ax, math.sqrt(ay*ay + az*az)))

# =============================
# PID Controller
# =============================
TARGET_PITCH = 0.0

Kp = 4.0
Ki = 0.0
Kd = 0.15

integral = 0.0
prev_err = 0.0
I_LIMIT  = 20.0

SERVO_DIR = 1   # flip to -1 if correction is backwards

# =============================
# Timing
# =============================
t_prev = time.ticks_ms()

# OLED refresh timing (don’t refresh every loop)
OLED_PERIOD_MS = 100  # 10 Hz
t_oled = time.ticks_ms()

print("Pitch stabilizer running")

# =============================
# Main loop
# =============================
while True:
    t_now = time.ticks_ms()
    dt = time.ticks_diff(t_now, t_prev) / 1000.0
    t_prev = t_now

    if dt <= 0:
        dt = 0.002
    if dt > 0.05:
        dt = 0.05

    v = mpu.get_values()

    # Accel (g)
    ax = v["AcX"] / 16384.0
    ay = v["AcY"] / 16384.0
    az = v["AcZ"] / 16384.0

    # Gyro (deg/s) - using GyY as your pitch rate axis
    gy = v["GyY"] / 131.0

    # Complementary filter
    pitch_acc = accel_pitch_deg(ax, ay, az)
    pitch = ALPHA * (pitch + gy * dt) + (1.0 - ALPHA) * pitch_acc

    # PID
    err = TARGET_PITCH - pitch

    integral += err * dt
    integral = clamp(integral, -I_LIMIT, I_LIMIT)

    deriv = (err - prev_err) / dt
    prev_err = err

    out = Kp * err + Ki * integral + Kd * deriv
    out = clamp(out, -MAX_CMD_DEG, MAX_CMD_DEG)

    # Servo output
    servo_us = servo_write_cmd(SERVO_DIR * out)

    # OLED update (10 Hz)
    if time.ticks_diff(t_now, t_oled) >= OLED_PERIOD_MS:
        t_oled = t_now

        oled.fill(0)
        oled.text("ALLAH stabilizer", 0, 0)

        # pitch
        oled.text("pitch:", 0, 14)
        oled.text("{:+6.2f} deg".format(pitch), 54, 14)

        # accel x (what you mentioned)
        oled.text("ax:", 0, 28)
        oled.text("{:+5.2f} g".format(ax), 54, 28)

        # servo output
        oled.text("servo:", 0, 42)
        oled.text("{} us".format(int(servo_us)), 54, 42)

        oled.show()

    # FAST loop (~300–500 Hz)
    time.sleep_ms(2)