# main.py — ESP32 MicroPython port of TN24
#
# Uses Tim Hanewich's mpu6050 module:
#   import mpu6050; mpu = mpu6050.MPU6050(i2c); mpu.wake()
#
# Features:
# - SSD1306 128x64 OLED expressions (ported from expressions.h)
# - MPU6050 double-tap detection + orientation (sleep/awake/upside-down)
# - 4 servo gait engine: stand, sleep, walk, turn, sit, dance, panic, self-right
# - Simple state machine mirroring TN24.ino structure
#
# Wiring (default pins; change below):
#   I2C: SCL=22, SDA=21  (SSD1306 at 0x3C/0x3D, MPU6050 at 0x68/0x69)
#   Servos (PWM 50 Hz): FRONT_L=15, FRONT_R=2, BACK_L=4, BACK_R=16

from machine import Pin, I2C, PWM
from ssd1306 import SSD1306_I2C
import mpu6050          # <-- use the external module directly
import utime
import math

# =========================
# ====== CONFIG ===========
# =========================
I2C_SCL_PIN = 22
I2C_SDA_PIN = 21
I2C_FREQ = 400_000
OLED_WIDTH = 128
OLED_HEIGHT = 64
OLED_ADDR = 0x3C   # used if present; falls back to 0x3D

# MPU6050 default is 0x68; if AD0 is high it becomes 0x69.
MPU_ADDRS = (0x68, 0x69)

SERVO_PINS = {
    'FL': 15,  # Front Left
    'FR': 2,   # Front Right
    'BL': 4,   # Back Left
    'BR': 16,  # Back Right
}
SERVO_MIN_US = 500
SERVO_MAX_US = 2500
SERVO_FREQ = 50  # Hz

# Gait tuning
STAND_ANGLE = 90
STEP_SMALL = 25
STEP_BIG = 35
TURN_DELTA = 25
SIT_ANGLE = 40
SLEEP_ANGLE = 20
DANCE_A1 = 40
DANCE_A2 = 20

# IMU/logic tuning
TAP_THRESHOLD_G = 1.2   # spike over steady magnitude (|Δ| > this) to count tap
TAP_WINDOW_MS = 500     # double-tap window
UPSIDE_Z_G = -0.6       # consider upside-down when az < this
SLEEP_TIMEOUT_MS = 30_000

# =========================
# ====== UTILITIES ========
# =========================
WHITE = 1
BLACK = 0

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

# =========================
# ====== SERVO DRIVER =====
# =========================
class Servo:
    def __init__(self, pin_no, freq=SERVO_FREQ, min_us=SERVO_MIN_US, max_us=SERVO_MAX_US):
        self.pwm = PWM(Pin(pin_no), freq=freq, duty_u16=0)
        self.min_us = min_us
        self.max_us = max_us
        self.period_ns = int(1_000_000_000 / freq)
        self.angle = None
        self.write(90)

    def write_us(self, us):
        us = clamp(us, self.min_us, self.max_us)
        duty_ns = int(us * 1000)
        val = int((duty_ns / self.period_ns) * 65535)
        self.pwm.duty_u16(clamp(val, 0, 65535))

    def write(self, angle):
        angle = clamp(angle, 0, 180)
        self.angle = angle
        us = self.min_us + (self.max_us - self.min_us) * (angle / 180.0)
        self.write_us(us)

    def deinit(self):
        self.pwm.deinit()

# =========================
# ====== OLED DRAW ========
# =========================
class Expressions:
    SCREEN_WIDTH = OLED_WIDTH
    SCREEN_HEIGHT = OLED_HEIGHT

    def __init__(self, display):
        self.display = display
        self.prevBlinkTime = utime.ticks_ms()
        self.happy_state = False
        self.q = -20

    def _pixel(self, x, y, c=WHITE):
        if 0 <= x < self.SCREEN_WIDTH and 0 <= y < self.SCREEN_HEIGHT:
            self.display.pixel(x, y, c)
    def _hline(self, x, y, w, c=WHITE):
        if w > 0: self.display.hline(x, y, w, c)
    def _vline(self, x, y, h, c=WHITE):
        if h > 0: self.display.vline(x, y, h, c)
    def _rect(self, x, y, w, h, c=WHITE):
        self.display.rect(x, y, w, h, c)
    def _fill_rect(self, x, y, w, h, c=WHITE):
        if w > 0 and h > 0: self.display.fill_rect(x, y, w, h, c)
    def _draw_circle(self, x0, y0, r, c=WHITE):
        f = 1 - r; ddF_x = 1; ddF_y = -2 * r; x = 0; y = r
        self._pixel(x0, y0 + r, c); self._pixel(x0, y0 - r, c)
        self._pixel(x0 + r, y0, c); self._pixel(x0 - r, y0, c)
        while x < y:
            if f >= 0:
                y -= 1; ddF_y += 2; f += ddF_y
            x += 1; ddF_x += 2; f += ddF_x
            for dx, dy in ((x,y),(-x,y),(x,-y),(-x,-y),(y,x),(-y,x),(y,-x),(-y,-x)):
                self._pixel(x0 + dx, y0 + dy, c)
    def _fill_circle(self, x0, y0, r, c=WHITE):
        for y in range(-r, r+1):
            dx = int((r*r - y*y) ** 0.5)
            self._hline(x0 - dx, y0 + y, 2*dx + 1, c)
    def _fill_triangle(self, x0, y0, x1, y1, x2, y2, c=WHITE):
        if y0 > y1: x0, x1, y0, y1 = x1, x0, y1, y0
        if y1 > y2: x1, x2, y1, y2 = x2, x1, y2, y1
        if y0 > y1: x0, x1, y0, y1 = x1, x0, y1, y0
        def interp(xa, ya, xb, yb, y):
            if yb == ya: return xa
            return int(xa + (xb - xa) * (y - ya) / (yb - ya))
        for y in range(y0, y1+1):
            xa = interp(x0,y0,x2,y2,y); xb = interp(x0,y0,x1,y1,y)
            if xa > xb: xa, xb = xb, xa
            self._hline(xa, y, xb - xa + 1, c)
        for y in range(y1, y2+1):
            xa = interp(x0,y0,x2,y2,y); xb = interp(x1,y1,x2,y2,y)
            if xa > xb: xa, xb = xb, xa
            self._hline(xa, y, xb - xa + 1, c)
    def _fill_round_rect(self, x, y, w, h, r, c=WHITE):
        self._fill_rect(x + r, y, w - 2*r, h, c)
        self._fill_rect(x, y + r, r, h - 2*r, c)
        self._fill_rect(x + w - r, y + r, r, h - 2*r, c)
        self._fill_circle(x + r, y + r, r, c)
        self._fill_circle(x + w - r - 1, y + r, r, c)
        self._fill_circle(x + r, y + h - r - 1, r, c)
        self._fill_circle(x + w - r - 1, y + h - r - 1, r, c)
    def _draw_line(self, x0,y0,x1,y1,c=WHITE):
        self.display.line(x0,y0,x1,y1,c)
    def _text(self, s, x, y, c=WHITE):
        self.display.text(s, x, y, c)
    def _clear(self):
        self.display.fill(0)
    def _show(self):
        self.display.show()
    def _delay(self, ms):
        utime.sleep_ms(ms)

    # Expressions (subset)
    def close(self):
        self._clear()
        self._fill_round_rect(5, 19, 55, 18, 6, WHITE)
        self._fill_round_rect(67, 19, 55, 18, 6, WHITE)
        self._fill_rect(5, 1, 55, 18, BLACK)
        self._fill_rect(67, 1, 55, 18, BLACK)
        self._show()
    def normal(self):
        self._clear()
        self._fill_round_rect(8, 12, 50, 35, 9, WHITE)
        self._fill_round_rect(70, 12, 50, 35, 9, WHITE)
        self._show()
        self.happy_state = False
    def blink(self):
        now = utime.ticks_ms()
        if utime.ticks_diff(now, self.prevBlinkTime) > 2000:
            self.close(); self._delay(50); self.normal(); self.prevBlinkTime = now
    def sad(self):
        for i in range(0, 16, 3):
            self._clear()
            self._fill_round_rect(8, 18, 50, 29, 9, WHITE)
            self._fill_round_rect(70, 18, 50, 29, 9, WHITE)
            self._fill_triangle(3, 14, 64, 14, 3, 21+i, BLACK)
            self._fill_triangle(68, 14, 124, 21+i, 124, 14, BLACK)
            self._show()
    def upset(self):
        if not self.happy_state:
            self._clear()
            self._fill_round_rect(8, 12, 50, 35, 9, WHITE)
            self._fill_round_rect(70, 12, 50, 35, 9, WHITE)
            self._fill_rect(8, self.q, 50, 35, BLACK)
            self._fill_rect(70, self.q, 50, 35, BLACK)
            self._show()
            if self.q <= -7: self.q += 3
    def happy(self):
        for i in range(62, 58, -3):
            self._clear()
            self._fill_round_rect(8, 12, 50, 35, 11, WHITE)
            self._fill_round_rect(70, 12, 50, 35, 11, WHITE)
            self._fill_circle(33, i, 38, BLACK)
            self._fill_circle(95, i, 38, BLACK)
            self._show()
        self.happy_state = True
    def cute(self):
        for i in range(0, 3, 2):
            self._clear()
            self._fill_round_rect(8, 12, 50, 35, 12, WHITE)
            self._fill_round_rect(70, 12, 50, 35, 12, WHITE)
            self._fill_circle(30, 66 - i, 40, BLACK)
            self._fill_circle(98, 66 - i, 40, BLACK)
            self._show()
        self.happy_state = True
    def angry(self):
        for i in range(0, 16, 3):
            self._clear()
            self._fill_round_rect(8, 18, 50, 29, 9, WHITE)
            self._fill_round_rect(70, 18, 50, 29, 9, WHITE)
            self._fill_triangle(3, 14, 64, 18 + i, 124, 14, BLACK)
            self._show()
    def sleepy(self):
        for i in range(0, 12, 2):
            self._clear()
            self._fill_round_rect(8, 12 + i, 50, 25, 9, WHITE)
            self._fill_round_rect(70, 12 + i, 50, 25, 9, WHITE)
            self._text("z", 100 - i, 40 - i, WHITE)
            self._text("z", 110 - i, 30 - i, WHITE)
            self._text("z", 120 - i, 20 - i, WHITE)
            self._show(); self._delay(100)
    def wink(self):
        self._clear()
        self._fill_round_rect(70, 12, 50, 35, 9, WHITE)
        self._fill_round_rect(5, 19, 55, 18, 6, WHITE)
        self._fill_rect(5, 1, 55, 18, BLACK)
        self._show(); self._delay(300); self.normal()
    def surprised(self):
        self._clear()
        self._fill_circle(33, 30, 20, WHITE)
        self._fill_circle(95, 30, 20, WHITE)
        self._draw_line(13, 5, 53, 5, WHITE)
        self._draw_line(75, 5, 115, 5, WHITE)
        self._show()
    def confused(self):
        self._clear()
        self._fill_round_rect(8, 12, 50, 35, 9, WHITE)
        self._fill_circle(95, 30, 20, WHITE)
        self._draw_line(75, 5, 115, 15, WHITE)
        self._show()
    def love(self):
        self._clear()
        for x in (33, 95):
            self._fill_circle(x - 7, 25, 10, WHITE)
            self._fill_circle(x + 7, 25, 10, WHITE)
            self._fill_triangle(x - 15, 30, x + 15, 30, x, 45, WHITE)
        self._show()
    def dizzy(self):
        self._clear()
        for centerX in (33, 95):
            centerY = 30
            for r in range(0, 15, 3):
                self._draw_circle(centerX, centerY, r, WHITE)
                self._draw_line(centerX - r, centerY - r, centerX + r, centerY + r, WHITE)
                self._draw_line(centerX - r, centerY + r, centerX + r, centerY - r, WHITE)
        self._show()
    def thinking(self):
        self._clear()
        self._fill_round_rect(8, 15, 50, 25, 9, WHITE)
        self._fill_round_rect(70, 15, 50, 25, 9, WHITE)
        self._fill_circle(110, 15, 3, WHITE)
        self._fill_circle(115, 10, 4, WHITE)
        self._fill_circle(122, 5, 5, WHITE)
        self._show()
    def mischievous(self):
        self._clear()
        self._fill_round_rect(8, 12, 50, 25, 9, WHITE)
        self._fill_round_rect(70, 18, 50, 25, 9, WHITE)
        self._draw_line(8, 5, 58, 15, WHITE)
        self._draw_line(70, 15, 120, 5, WHITE)
        self._show()
    def crying(self):
        for i in range(0, 15, 3):
            self._clear()
            self._fill_round_rect(8, 18, 50, 29, 9, WHITE)
            self._fill_round_rect(70, 18, 50, 29, 9, WHITE)
            self._fill_round_rect(20, 45 + i, 3, 5, 1, WHITE)
            self._fill_round_rect(82, 45 + i, 3, 5, 1, WHITE)
            self._show(); self._delay(100)
    def nervous(self):
        for i in range(0, 4):
            self._clear()
            dx = (i % 2) * 2
            self._fill_round_rect(8 + dx, 12, 50, 35, 9, WHITE)
            self._fill_round_rect(70 + dx, 12, 50, 35, 9, WHITE)
            self._fill_round_rect(115, 20 + i*2, 3, 5, 1, WHITE)
            self._show(); self._delay(100)

# =========================
# ====== ROBOT LOGIC ======
# =========================
class TN24:
    def __init__(self):
        # I2C + devices
        self.i2c = I2C(0, scl=Pin(I2C_SCL_PIN), sda=Pin(I2C_SDA_PIN), freq=I2C_FREQ)
        found = self.i2c.scan()
        print("I2C found:", [hex(a) for a in found])

        # OLED: prefer 0x3C, fallback 0x3D
        oled_addr = OLED_ADDR if OLED_ADDR in found else (0x3D if 0x3D in found else None)
        if oled_addr is None:
            raise OSError("SSD1306 not found at 0x3C/0x3D. Check wiring.")
        self.disp = SSD1306_I2C(OLED_WIDTH, OLED_HEIGHT, self.i2c, addr=oled_addr)

        # MPU6050 via module; pick 0x68 or 0x69 present on the bus
        mpu_addr = None
        for cand in MPU_ADDRS:
            if cand in found:
                mpu_addr = cand
                break
        if mpu_addr is None:
            # still try default (some sims don’t show up in scan until touched)
            mpu_addr = 0x68

        self.mpu = mpu6050.MPU6050(self.i2c, address=mpu_addr)
        self.mpu.wake()

        # Optionally set ranges (0=±250 dps, 0=±2g) for standard scaling:
        try:
            self.mpu.write_gyro_range(0)
            self.mpu.write_accel_range(0)
        except Exception:
            # Some implementations might not expose setters. It's fine.
            pass

        self.exp = Expressions(self.disp)

        # Servos
        self.sv = {
            'FL': Servo(SERVO_PINS['FL']),
            'FR': Servo(SERVO_PINS['FR']),
            'BL': Servo(SERVO_PINS['BL']),
            'BR': Servo(SERVO_PINS['BR']),
        }
        # State
        self.last_tap_ms = 0
        self.tap_count = 0
        self.last_active_ms = utime.ticks_ms()
        self.sleeping = False
        self.upside = False
        self.boot()

    # ---------- setup / boot ----------
    def boot(self):
        self.exp.normal(); self.exp._delay(200)
        self.exp.wink(); self.exp._delay(200)
        self.exp.happy(); self.exp._delay(200)
        self.stand()

    # ---------- IMU read helper (directly uses module) ----------
    def _accel_gyro(self):
        gx, gy, gz = self.mpu.read_gyro_data()
        ax, ay, az = self.mpu.read_accel_data()
        return ax, ay, az, gx, gy, gz

    # ---------- helpers ----------
    def _each(self, a_fl, a_fr, a_bl, a_br, hold_ms=200):
        self.sv['FL'].write(a_fl)
        self.sv['FR'].write(a_fr)
        self.sv['BL'].write(a_bl)
        self.sv['BR'].write(a_br)
        utime.sleep_ms(hold_ms)

    def stand(self):
        self._each(STAND_ANGLE, STAND_ANGLE, STAND_ANGLE, STAND_ANGLE, 300)

    def sleep(self):
        self._each(SLEEP_ANGLE, SLEEP_ANGLE, SLEEP_ANGLE, SLEEP_ANGLE, 300)
        self.exp.sleepy()

    def sit(self):
        self._each(SIT_ANGLE, SIT_ANGLE, SIT_ANGLE, SIT_ANGLE, 300)
        self.exp.sad()

    def walk_forward(self, steps=4, step=STEP_SMALL, delay_ms=200):
        for _ in range(steps):
            self._each(STAND_ANGLE+step, STAND_ANGLE-step, STAND_ANGLE+step, STAND_ANGLE-step, delay_ms)
            self._each(STAND_ANGLE-step, STAND_ANGLE+step, STAND_ANGLE-step, STAND_ANGLE+step, delay_ms)

    def walk_backward(self, steps=4, step=STEP_SMALL, delay_ms=200):
        for _ in range(steps):
            self._each(STAND_ANGLE-step, STAND_ANGLE+step, STAND_ANGLE-step, STAND_ANGLE+step, delay_ms)
            self._each(STAND_ANGLE+step, STAND_ANGLE-step, STAND_ANGLE+step, STAND_ANGLE-step, delay_ms)

    def turn_left(self, steps=4, delta=TURN_DELTA, delay_ms=200):
        for _ in range(steps):
            self._each(STAND_ANGLE+delta, STAND_ANGLE+delta, STAND_ANGLE-delta, STAND_ANGLE-delta, delay_ms)
            self._each(STAND_ANGLE-delta, STAND_ANGLE-delta, STAND_ANGLE+delta, STAND_ANGLE+delta, delay_ms)

    def turn_right(self, steps=4, delta=TURN_DELTA, delay_ms=200):
        for _ in range(steps):
            self._each(STAND_ANGLE-delta, STAND_ANGLE-delta, STAND_ANGLE+delta, STAND_ANGLE+delta, delay_ms)
            self._each(STAND_ANGLE+delta, STAND_ANGLE+delta, STAND_ANGLE-delta, STAND_ANGLE-delta, delay_ms)

    def dance(self):
        for _ in range(3):
            self._each(STAND_ANGLE+DANCE_A1, STAND_ANGLE-DANCE_A1, STAND_ANGLE+DANCE_A1, STAND_ANGLE-DANCE_A1, 250)
            self._each(STAND_ANGLE+DANCE_A2, STAND_ANGLE-DANCE_A2, STAND_ANGLE+DANCE_A2, STAND_ANGLE-DANCE_A2, 250)
            self._each(STAND_ANGLE, STAND_ANGLE, STAND_ANGLE, STAND_ANGLE, 200)
        self.exp.happy()

    def panic_movement(self):
        for _ in range(4):
            self._each(0, 180, 0, 180, 120)
            self._each(180, 0, 180, 0, 120)
        self.exp.angry()
        self.stand()

    def self_right(self):
        for _ in range(6):
            self._each(20, 160, 160, 20, 180)
            self._each(160, 20, 20, 160, 180)
        self.stand()

    # ---------- IMU logic ----------
    def _tap_logic(self, ax, ay, az):
        now = utime.ticks_ms()
        mag = math.sqrt(ax*ax + ay*ay + az*az)
        delta = abs(mag - 1.0)
        tap = delta > (TAP_THRESHOLD_G - 1.0)
        if tap:
            if utime.ticks_diff(now, self.last_tap_ms) <= TAP_WINDOW_MS:
                self.tap_count += 1
            else:
                self.tap_count = 1
            self.last_tap_ms = now
            self.last_active_ms = now
        if self.tap_count >= 2 and utime.ticks_diff(now, self.last_tap_ms) <= TAP_WINDOW_MS:
            self.tap_count = 0
            return True
        if utime.ticks_diff(now, self.last_tap_ms) > TAP_WINDOW_MS:
            self.tap_count = 0
        return False

    def _orientation_logic(self, ax, ay, az):
        return az < UPSIDE_Z_G

    # ---------- state handlers ----------
    def handle_sleeping(self):
        self.exp.sleepy()
        self.sleep()
        for _ in range(10):
            ax, ay, az, *_ = self._accel_gyro()
            if self._tap_logic(ax, ay, az):
                self.wake_up()
                return
            utime.sleep_ms(100)

    def handle_active(self):
        self.exp.normal(); self.exp.blink()
        ax, ay, az, gx, gy, gz = self._accel_gyro()
        if self._tap_logic(ax, ay, az):
            self.react()
        if self._orientation_logic(ax, ay, az):
            self.upside = True
            self.handle_upside_down()
            return
        if utime.ticks_diff(utime.ticks_ms(), self.last_active_ms) > SLEEP_TIMEOUT_MS:
            self.go_to_sleep()

    def handle_upside_down(self):
        self.exp.surprised()
        self.panic_movement()
        self.self_right()
        ax, ay, az, *_ = self._accel_gyro()
        if not self._orientation_logic(ax, ay, az):
            self.upside = False
            self.exp.normal()
        else:
            utime.sleep_ms(300)

    # ---------- action macros ----------
    def react(self):
        self.exp.mischievous()
        self.walk_forward(steps=2)
        self.turn_left(steps=2)
        self.walk_backward(steps=2)
        self.dance()

    def play(self):
        self.exp.happy(); self.dance(); self.walk_forward(steps=3, step=STEP_BIG)
        self.turn_right(steps=3); self.walk_backward(steps=2)
        self.exp.love(); utime.sleep_ms(400)
        self.exp.normal()

    def wake_up(self):
        self.sleeping = False
        self.last_active_ms = utime.ticks_ms()
        self.stand(); self.exp.wink()

    def go_to_sleep(self):
        self.sleeping = True
        self.sit(); self.sleep()

    # ---------- main loop ----------
    def loop(self):
        if self.sleeping:
            self.handle_sleeping()
        elif self.upside:
            self.handle_upside_down()
        else:
            self.handle_active()
        utime.sleep_ms(50)


def run():
    bot = TN24()
    while True:
        bot.loop()

if __name__ == '__main__':
    run()