from machine import Pin, I2C
import math, time

# ---------------------------
# MPU6050 Class
# ---------------------------
class MPU6050:
    def __init__(self, i2c, addr=0x68):
        self.i2c = i2c
        self.addr = addr
        self.i2c.writeto_mem(self.addr, 0x6B, b'\x00')  # Wake up sensor

    def read_raw_data(self, reg):
        high = self.i2c.readfrom_mem(self.addr, reg, 1)[0]
        low = self.i2c.readfrom_mem(self.addr, reg + 1, 1)[0]
        value = (high << 8) | low
        if value > 32768:
            value -= 65536
        return value

    def get_accel(self):
        ax = self.read_raw_data(0x3B) / 16384.0
        ay = self.read_raw_data(0x3D) / 16384.0
        az = self.read_raw_data(0x3F) / 16384.0
        return ax, ay, az

    def get_tilt_angle(self):
        ax, ay, az = self.get_accel()
        angle_x = math.degrees(math.atan2(ax, math.sqrt(ay**2 + az**2)))
        angle_y = math.degrees(math.atan2(ay, math.sqrt(ax**2 + az**2)))
        return angle_x, angle_y


# ---------------------------
# Setup I2C & Sensor
# ---------------------------
i2c = I2C(0, scl=Pin(1), sda=Pin(0), freq=400000)
mpu = MPU6050(i2c)

# ---------------------------
# Setup LEDs
# ---------------------------
led_pins = [Pin(p, Pin.OUT) for p in range(2, 10)]
n_leds = len(led_pins)
center = n_leds // 2

def clear_leds():
    for led in led_pins:
        led.value(0)

def show_led(pos):
    clear_leds()
    pos = max(0, min(n_leds - 1, pos))
    led_pins[pos].value(1)

# ---------------------------
# Game Logic
# ---------------------------
print("🎮 Balance Game Started!")
print("Tilt left/right to keep the LED centered!")

score = 0
start_time = time.ticks_ms()

while True:
    angle_x, angle_y = mpu.get_tilt_angle()

    # Adjust sensitivity
    pos = int(center + (angle_x / 10))

    # Update LED display
    show_led(pos)

    # Check for losing condition
    if pos <= 0 or pos >= n_leds - 1:
        print("\n💥 You lost balance!")
        for _ in range(3):
            for led in led_pins:
                led.value(1)
            time.sleep(0.15)
            clear_leds()
            time.sleep(0.15)

        # Show score
        duration = (time.ticks_ms() - start_time) / 1000
        print(f"🕓 Survived: {duration:.2f} seconds")
        score = 0
        start_time = time.ticks_ms()
        time.sleep(1)

    # Update score based on time
    score = (time.ticks_ms() - start_time) // 1000
    time.sleep(0.05)