# main.py — MicroPython ESP32 robot (ported from Arduino)
# Preserves states (SLEEPING/ACTIVE/UPSIDE_DOWN), double-tap wake,
# idle→sleep, walking/dancing/sitting, OLED boot + rich expressions.
#
# 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

import time
import math
import utime  # Expressions uses utime.ticks_ms / sleep_ms
from machine import Pin, PWM, I2C
from urandom import getrandbits

# =========================
# --- USER CONFIGURATION ---
# =========================
I2C_SCL = 22
I2C_SDA = 21
OLED_ADDR = 0x3C
SCREEN_WIDTH = 128
SCREEN_HEIGHT = 64

# Display color constants for SSD1306 (1-bit)
WHITE = 1
BLACK = 0
OLED_WIDTH = SCREEN_WIDTH
OLED_HEIGHT = SCREEN_HEIGHT

# Pick PWM-capable pins on your ESP32:
GPIO_FRONT_LEFT  = 15
GPIO_FRONT_RIGHT = 2
GPIO_BACK_LEFT   = 4
GPIO_BACK_RIGHT  = 16

# Tap detection tuning (units are "g" because we normalize)
TAP_THRESHOLD_G = 1.0        # tweak based on your build
TAP_WINDOW_MS   = 500

IDLE_TIMEOUT_MS = 30000

# Servo geometry
STAND_ANGLE   = 90
SLEEP_ANGLE   = 0
FORWARD_STEP  = 60
BACKWARD_STEP = 150
SIT_ANGLE     = 20
DANCE_1       = 60
DANCE_2       = 120

DELAY_TIME = 200  # ms for gait steps

# ===== Personality / scheduler tuning =====
ACTIVE_BURST_ACTIONS = (2, 4)   # min/max actions per active burst before a short chill
SHORT_CHILL_MS = (300, 900)     # rest a beat between bursts
MICRO_GESTURE_CHANCE = 18       # % chance to wink/cute between actions
REPEAT_COOLDOWN_MS = 4000       # avoid repeating same action inside this window
NOVELTY_BONUS = 1.25            # mild boost to less-used actions

# Base weights (scheduler will modulate by mood & cooldown)
ACTION_WEIGHTS = {
    "walk": 3,
    "turn": 2,
    "dance": 1,
    "sit": 1,
    "react": 1,
    "shimmy": 2,
    "wag": 2,
    "hop": 1,
    "pose": 2,
}

# =========================
# --- OLED / SSD1306 ---
# =========================
try:
    from ssd1306 import SSD1306_I2C
except ImportError:
    raise RuntimeError("Install ssd1306 MicroPython driver on the board.")

# ESP32/Wokwi-friendly constructor (id-less)
i2c = I2C(scl=Pin(I2C_SCL), sda=Pin(I2C_SDA), freq=400000)
display = SSD1306_I2C(SCREEN_WIDTH, SCREEN_HEIGHT, i2c, addr=OLED_ADDR)

def oled_clear():
    display.fill(0)
    display.show()

def oled_center_text(txt, y, scale=1):
    # crude centering for 6x8 font
    w = len(txt) * 8 * scale * 0.75
    x = max(0, (SCREEN_WIDTH - int(w)) // 2)
    if scale == 1:
        display.text(txt, x, y, 1)
    else:
        # simple "scale 2" spacing using default font
        for i, ch in enumerate(txt):
            display.text(ch, x + i*12, y, 1)

def show_boot_animation():
    oled_clear()
    oled_center_text("TN24 Robodog", 24, scale=2)
    display.show()
    time.sleep_ms(2000)
    oled_clear()
    # Start with neutral eyes after boot
    try:
        expressions.normal()
    except NameError:
        pass

# =========================
# ====== 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)

# Create expressions instance
expressions = Expressions(display)

# Keep the same call sites by exposing thin wrappers:
def happy():        expressions.happy()
def sleepy():       expressions.sleepy()
def mischievous():  expressions.mischievous()
def cute():         expressions.cute()
def wink():         expressions.wink()
def thinking():     expressions.thinking()
def love():         expressions.love()
def dizzy():        expressions.dizzy()
def crying():       expressions.crying()
def play():         expressions.surprised()   # mapping
def react():        expressions.confused()    # mapping

# =========================
# --- MPU6050 (tiny driver) ---
# Returns accel in g, gyro in deg/s approximately.
# =========================
class MPU6050:
    ADDR = 0x68
    REG_PWR1 = 0x6B
    REG_ACCEL = 0x3B
    REG_GYRO = 0x43

    def __init__(self, i2c, addr=ADDR):
        self.i2c = i2c
        self.addr = addr
        # wake up
        self.i2c.writeto_mem(self.addr, self.REG_PWR1, b'\x00')
        time.sleep_ms(50)

    def _read_vector(self, reg):
        raw = self.i2c.readfrom_mem(self.addr, reg, 6)
        # MicroPython v1.21: use positional 'signed' arg
        x = int.from_bytes(raw[0:2], 'big', True)
        y = int.from_bytes(raw[2:4], 'big', True)
        z = int.from_bytes(raw[4:6], 'big', True)
        return x, y, z

    def get_accel(self):
        ax, ay, az = self._read_vector(self.REG_ACCEL)
        # default ±2g -> 16384 LSB/g
        return ax/16384.0, ay/16384.0, az/16384.0

    def get_gyro(self):
        gx, gy, gz = self._read_vector(self.REG_GYRO)
        # default ±250 dps -> 131 LSB/(deg/s)
        return gx/131.0, gy/131.0, gz/131.0

mpu = MPU6050(i2c)

# =========================
# --- Servo helper (ESP32/Wokwi robust) ---
# Uses duty_ns if available, else duty(0..1023), else duty_u16 fallback.
# =========================
class Servo:
    def __init__(self, pin_num, freq=50, min_us=500, max_us=2500):
        self.pwm = PWM(Pin(pin_num), freq=freq)
        self.min_us = min_us
        self.max_us = max_us
        self.freq = freq
        self.period_ns = int(1_000_000_000 // freq)

    def write(self, angle):
        angle = max(0, min(180, int(angle)))
        pulse_us = self.min_us + (self.max_us - self.min_us) * angle // 180
        pulse_ns = pulse_us * 1000

        if hasattr(self.pwm, "duty_ns"):
            self.pwm.duty_ns(pulse_ns)
        elif hasattr(self.pwm, "duty"):  # 0..1023 on ESP32
            duty = int((pulse_ns / self.period_ns) * 1023)
            duty = max(0, min(1023, duty))
            self.pwm.duty(duty)
        else:
            # Fallback
            duty_u16 = int((pulse_ns / self.period_ns) * 65535)
            duty_u16 = max(0, min(65535, duty_u16))
            self.pwm.duty_u16(duty_u16)

# Instantiate servos
frontLeftServo  = Servo(GPIO_FRONT_LEFT)
frontRightServo = Servo(GPIO_FRONT_RIGHT)
backLeftServo   = Servo(GPIO_BACK_LEFT)
backRightServo  = Servo(GPIO_BACK_RIGHT)

# =========================
# --- State machine ---
# =========================
SLEEPING = 0
ACTIVE = 1
UPSIDE_DOWN = 2
currentState = SLEEPING
lastTapTime = 0
tapCount = 0
isUpright = True
lastActionTime = time.ticks_ms()

# Gait/face gating
forwardCueArmed = True  # flash mischievous only on first forward step until re-armed

def _arm_forward_cue():
    global forwardCueArmed
    forwardCueArmed = True

def _disarm_forward_cue():
    global forwardCueArmed
    forwardCueArmed = False

def _show_turn_face(step_deg, speed_ms):
    """
    Angry 'intensity' scales with turn 'step' (angle) and speed (shorter ms = faster).
    We map both to an integer 1..6 and repeat expressions.angry() that many times.
    """
    a = max(0, step_deg - 40) / 80.0        # 0..~1
    v = max(0, 800 - speed_ms) / 600.0      # 0..~1 (faster -> bigger)
    intensity = 1 + int(5 * min(1.0, 0.6*a + 0.6*v))  # 1..6
    for _ in range(intensity):
        expressions.angry()

# ===== Mood/energy =====
EXCITED, CURIOUS, CHILL = 0, 1, 2
mood = CURIOUS
energy = 65  # 0..100

def _set_mood(new_mood):
    global mood
    mood = new_mood

def _mood_speed(base_ms):
    # faster when excited, slower when chill
    if mood == EXCITED: return max(120, int(base_ms * 0.75))
    if mood == CHILL:   return int(base_ms * 1.20)
    return base_ms

def _tick_energy(delta):
    # energy drifts; sleep if too low
    global energy
    energy = max(0, min(100, energy + delta))
    if energy < 10:
        _set_mood(CHILL)

# =========================
# --- Motion primitives ---
# =========================
def stand():
    frontLeftServo.write(STAND_ANGLE)
    frontRightServo.write(STAND_ANGLE)
    backLeftServo.write(STAND_ANGLE)
    backRightServo.write(STAND_ANGLE)

def sleep_pose():
    sleepy()  # cue sleepy during pose
    frontLeftServo.write(SLEEP_ANGLE)
    frontRightServo.write(180 - SLEEP_ANGLE)
    backLeftServo.write(180)
    backRightServo.write(SLEEP_ANGLE)
    time.sleep_ms(1000)

def walkForward():
    # Show mischievous only on the first forward step after any non-forward action
    if forwardCueArmed:
        mischievous()
        _disarm_forward_cue()

    # FL + BR
    frontLeftServo.write(FORWARD_STEP)
    backRightServo.write(180 - FORWARD_STEP)
    time.sleep_ms(DELAY_TIME)
    frontLeftServo.write(STAND_ANGLE)
    backRightServo.write(STAND_ANGLE)
    time.sleep_ms(DELAY_TIME)
    # FR + BL
    frontRightServo.write(180 - FORWARD_STEP)
    backLeftServo.write(FORWARD_STEP)
    time.sleep_ms(DELAY_TIME)
    frontRightServo.write(STAND_ANGLE)
    backLeftServo.write(STAND_ANGLE)
    time.sleep_ms(DELAY_TIME)
    expressions.normal()

def walkBackward():
    expressions.upset()
    frontLeftServo.write(BACKWARD_STEP)
    backRightServo.write(180 - BACKWARD_STEP)
    time.sleep_ms(500)
    frontLeftServo.write(STAND_ANGLE)
    backRightServo.write(STAND_ANGLE)
    time.sleep_ms(500)
    frontRightServo.write(180 - BACKWARD_STEP)
    backLeftServo.write(BACKWARD_STEP)
    time.sleep_ms(500)
    frontRightServo.write(STAND_ANGLE)
    backLeftServo.write(STAND_ANGLE)
    time.sleep_ms(500)
    expressions.normal()
    _arm_forward_cue()

def turnLeft(step_deg=FORWARD_STEP, speed_ms=500):
    _arm_forward_cue()  # next forward step gets the mischievous flash
    _show_turn_face(step_deg, speed_ms)
    frontLeftServo.write(step_deg)
    backLeftServo.write(step_deg)
    time.sleep_ms(speed_ms)
    frontLeftServo.write(STAND_ANGLE)
    backLeftServo.write(STAND_ANGLE)
    time.sleep_ms(speed_ms)
    expressions.normal()

def turnRight(step_deg=FORWARD_STEP, speed_ms=500):
    _arm_forward_cue()
    _show_turn_face(step_deg, speed_ms)
    frontRightServo.write(180 - step_deg)
    backRightServo.write(180 - step_deg)
    time.sleep_ms(speed_ms)
    frontRightServo.write(STAND_ANGLE)
    backRightServo.write(STAND_ANGLE)
    time.sleep_ms(speed_ms)
    expressions.normal()

def sit():
    thinking()
    frontLeftServo.write(90)
    frontRightServo.write(90)
    backLeftServo.write(180 - SIT_ANGLE)
    backRightServo.write(SIT_ANGLE)
    time.sleep_ms(1000)
    expressions.normal()
    _arm_forward_cue()

def dance():
    cute()
    for _ in range(3):
        frontLeftServo.write(DANCE_1)
        backRightServo.write(180 - DANCE_1)
        time.sleep_ms(200)
        frontRightServo.write(180 - DANCE_2)
        backLeftServo.write(DANCE_2)
        time.sleep_ms(200)
        frontLeftServo.write(DANCE_2)
        backRightServo.write(180 - DANCE_2)
        time.sleep_ms(200)
        frontRightServo.write(180 - DANCE_1)
        backLeftServo.write(DANCE_1)
        time.sleep_ms(200)
    stand()
    love()
    time.sleep_ms(300)
    expressions.normal()
    _arm_forward_cue()

def react_motion():
    expressions.surprised()
    frontLeftServo.write(90)
    frontRightServo.write(90)
    backLeftServo.write(180 - SIT_ANGLE)
    backRightServo.write(SIT_ANGLE)
    time.sleep_ms(1000)
    stand()
    time.sleep_ms(500)

    expressions.nervous()
    frontLeftServo.write(SIT_ANGLE)
    frontRightServo.write(180 - SIT_ANGLE)
    backLeftServo.write(90)
    backRightServo.write(90)
    time.sleep_ms(1000)

    d = 250
    seq = [
        (0, 160),
        (20, 180),
        (0, 180 - SIT_ANGLE),
        (SIT_ANGLE, 180),
        (0, 160),
        (20, 180),
        (0, 180 - SIT_ANGLE),
        (SIT_ANGLE, 170),
    ]
    for fl, fr in seq:
        frontLeftServo.write(fl)
        frontRightServo.write(fr)
        time.sleep_ms(d)

    stand()
    time.sleep_ms(500)
    expressions.normal()
    _arm_forward_cue()

def panicMovement():
    expressions.nervous()
    for _ in range(4):
        frontLeftServo.write(getrandbits(8) % 181)
        frontRightServo.write(getrandbits(8) % 181)
        backLeftServo.write(getrandbits(8) % 181)
        backRightServo.write(getrandbits(8) % 181)
        time.sleep_ms(200)
    stand()
    expressions.normal()
    _arm_forward_cue()

def selfRight():
    print("Attempting to self-right...")
    expressions.upset()
    frontLeftServo.write(SLEEP_ANGLE)
    frontRightServo.write(180 - SLEEP_ANGLE)
    backLeftServo.write(180)
    backRightServo.write(SLEEP_ANGLE)
    time.sleep_ms(DELAY_TIME)
    stand()
    time.sleep_ms(DELAY_TIME)
    happy()
    _arm_forward_cue()

# ===== New micro-motions =====
def shimmy(cycles=4, amp=18, hold_ms=120):
    cute()
    for _ in range(cycles):
        frontLeftServo.write(STAND_ANGLE - amp)
        backRightServo.write(STAND_ANGLE + amp)
        time.sleep_ms(_mood_speed(hold_ms))
        frontLeftServo.write(STAND_ANGLE + amp)
        backRightServo.write(STAND_ANGLE - amp)
        time.sleep_ms(_mood_speed(hold_ms))
    stand()
    expressions.normal()
    _arm_forward_cue()
    _tick_energy(-3)

def wag(cycles=6, amp=22, hold_ms=90):
    # "tail" vibe: mostly back legs, eyes happy
    happy()
    for _ in range(cycles):
        backLeftServo.write(STAND_ANGLE + amp)
        backRightServo.write(STAND_ANGLE - amp)
        time.sleep_ms(_mood_speed(hold_ms))
        backLeftServo.write(STAND_ANGLE - amp)
        backRightServo.write(STAND_ANGLE + amp)
        time.sleep_ms(_mood_speed(hold_ms))
    stand()
    expressions.normal()
    _arm_forward_cue()
    _tick_energy(-3)

def hop(times=3, amp=30, hold_ms=120):
    expressions.surprised()
    for _ in range(times):
        # crouch
        frontLeftServo.write(STAND_ANGLE - amp)
        frontRightServo.write(STAND_ANGLE - amp)
        backLeftServo.write(STAND_ANGLE + amp)
        backRightServo.write(STAND_ANGLE + amp)
        time.sleep_ms(_mood_speed(hold_ms))
        # spring back
        stand()
        time.sleep_ms(_mood_speed(hold_ms))
    expressions.normal()
    _arm_forward_cue()
    _tick_energy(-4)

def pose(kind=None):
    # quick eye pose with tiny body bias
    if kind is None:
        kind = getrandbits(8) % 4
    if kind == 0: love()
    elif kind == 1: wink()
    elif kind == 2: mischievous()
    else: cute()
    # micro-tilt
    frontLeftServo.write(STAND_ANGLE - 8)
    frontRightServo.write(STAND_ANGLE + 8)
    time.sleep_ms(_mood_speed(250))
    stand()
    expressions.normal()
    _tick_energy(-1)

# ===== Weighted action scheduler =====
_last_action_at = {}  # name -> ticks_ms

def _age_ms(name):
    t = _last_action_at.get(name, None)
    if t is None: return 1_000_000
    return time.ticks_diff(time.ticks_ms(), t)

def _mark(name):
    _last_action_at[name] = time.ticks_ms()

def _weighted_choice():
    # base weights
    weights = ACTION_WEIGHTS.copy()

    # mood modulation
    if mood == EXCITED:
        for k in ("walk", "turn", "dance", "hop", "shimmy", "wag"):
            weights[k] = int(weights.get(k,1)*1.3) or 1
    elif mood == CHILL:
        for k in ("sit","pose","wag"):
            weights[k] = int(weights.get(k,1)*1.4) or 1
        weights["dance"] = max(1, int(weights["dance"]*0.7))

    # novelty: boost actions not seen recently
    for k in list(weights):
        age = _age_ms(k)
        if age > REPEAT_COOLDOWN_MS*2:
            weights[k] = int(weights[k]*NOVELTY_BONUS)

    # cooldown: zero out the one we just did if still “hot”
    for k in list(weights):
        if _age_ms(k) < REPEAT_COOLDOWN_MS:
            weights[k] = 0

    # build roulette
    bag = []
    for k, w in weights.items():
        bag += [k]*max(0, int(w))

    # fallback if all cooled down
    if not bag:
        bag = [k for k in ACTION_WEIGHTS]

    # pick
    idx = getrandbits(16) % len(bag)
    return bag[idx]

def _do(name):
    # map names to actions; vary parameters by mood
    if name == "walk":
        if mood == EXCITED:
            for _ in range(2): walkForward()
        else:
            walkForward()
        _tick_energy(-2)
    elif name == "turn":
        # random left/right, speed tied to mood
        step = FORWARD_STEP if mood != EXCITED else min(120, FORWARD_STEP+30)
        spd  = _mood_speed(500)
        if getrandbits(1): turnLeft(step_deg=step, speed_ms=spd)
        else:               turnRight(step_deg=step, speed_ms=spd)
        _tick_energy(-2)
    elif name == "dance":
        dance(); _tick_energy(-5)
    elif name == "sit":
        sit();  _tick_energy(+2)  # rest gives a bit back
    elif name == "react":
        react_motion(); _tick_energy(-3)
    elif name == "shimmy":
        shimmy()
    elif name == "wag":
        wag()
    elif name == "hop":
        hop()
    elif name == "pose":
        pose()
    _mark(name)

def _maybe_micro_gesture():
    if (getrandbits(7) % 100) < MICRO_GESTURE_CHANCE:
        if getrandbits(1): wink()
        else: cute()
        time.sleep_ms(_mood_speed(200))
        expressions.normal()

# =========================
# --- State helpers ---
# =========================
def wakeUp():
    global currentState, lastActionTime
    currentState = ACTIVE
    happy()
    _arm_forward_cue()
    for angle in range(SLEEP_ANGLE, STAND_ANGLE + 1, 5):
        frontLeftServo.write(angle)
        frontRightServo.write(angle)
        backLeftServo.write(angle)
        backRightServo.write(angle)
        time.sleep_ms(20)
    lastActionTime = time.ticks_ms()

def goToSleep():
    global currentState
    currentState = SLEEPING
    expressions.sad()
    time.sleep_ms(250)
    sleepy()
    for angle in range(STAND_ANGLE, SLEEP_ANGLE - 1, -5):
        frontLeftServo.write(angle)
        frontRightServo.write(180 - angle)
        backLeftServo.write(180)
        backRightServo.write(angle)
        time.sleep_ms(20)
    _arm_forward_cue()

# =========================
# --- Sensors & events ---
# =========================
def checkOrientation():
    global isUpright, currentState
    ax, ay, az = mpu.get_accel()
    wasUpright = isUpright
    isUpright = (az > 0.1)  # adjust if needed for your mounting

    if not isUpright and wasUpright:
        currentState = UPSIDE_DOWN
        crying()
    elif isUpright and not wasUpright:
        currentState = ACTIVE
        happy()

def checkTaps():
    """
    Emulates motion-interrupt double-tap using accel magnitude spikes.
    """
    global tapCount, lastTapTime
    ax, ay, az = mpu.get_accel()
    mag = math.sqrt(ax*ax + ay*ay + az*az)  # in g
    now = time.ticks_ms()

    if mag > (1.0 + TAP_THRESHOLD_G):  # baseline ~1g at rest
        if time.ticks_diff(now, lastTapTime) > 100:  # debounce
            tapCount += 1
            lastTapTime = now

    # Double tap window handling
    if tapCount >= 2 and time.ticks_diff(now, lastTapTime) <= TAP_WINDOW_MS:
        if currentState == SLEEPING:
            wakeUp()
        else:
            # greet burst when already active
            expressions.surprised()
            wag(cycles=4)
            love()
            time.sleep_ms(250)
            expressions.normal()
            _tick_energy(+3)  # social boost
        tapCount = 0

    # If window expired, reset
    if time.ticks_diff(now, lastTapTime) > TAP_WINDOW_MS:
        tapCount = 0

# =========================
# --- State handlers ---
# =========================
def handleSleepingState():
    global lastActionTime
    sleepy()
    if time.ticks_diff(time.ticks_ms(), lastActionTime) > 5000:
        # Occasionally show sleeping animation (10% chance)
        if (getrandbits(7) % 100) < 10:
            for _ in range(3):
                expressions.dizzy()
                time.sleep_ms(500)

def handleActiveState():
    global lastActionTime

    # auto-sleep if truly idle
    if time.ticks_diff(time.ticks_ms(), lastActionTime) > IDLE_TIMEOUT_MS:
        expressions.sad()
        time.sleep_ms(350)
        goToSleep()
        return

    # mood drifts with energy
    if energy > 70: _set_mood(EXCITED)
    elif energy < 30: _set_mood(CHILL)
    else: _set_mood(CURIOUS)

    # run a short burst of varied actions
    burst_len = (ACTIVE_BURST_ACTIONS[0] +
                 getrandbits(3) % (ACTIVE_BURST_ACTIONS[1]-ACTIVE_BURST_ACTIONS[0]+1))
    for _ in range(burst_len):
        name = _weighted_choice()
        _do(name)
        _maybe_micro_gesture()
        lastActionTime = time.ticks_ms()

        # quick breather between actions
        time.sleep_ms(_mood_speed(120))

    # short chill between bursts (adds a bit of variability)
    rest = SHORT_CHILL_MS[0] + (getrandbits(10) % (SHORT_CHILL_MS[1]-SHORT_CHILL_MS[0]+1))
    time.sleep_ms(_mood_speed(rest))

def handleUpsideDownState():
    crying()
    panicMovement()
    time.sleep_ms(1000)

# =========================
# --- Main ---
# =========================
def setup():
    show_boot_animation()
    stand()
    time.sleep_ms(1000)
    sleep_pose()
    sleepy()

def loop():
    while True:
        checkOrientation()
        checkTaps()

        # Keep the face lively
        expressions.blink()

        if currentState == SLEEPING:
            handleSleepingState()
        elif currentState == ACTIVE:
            handleActiveState()
        elif currentState == UPSIDE_DOWN:
            handleUpsideDownState()

        # Small idle delay to keep CPU cool
        time.sleep_ms(10)

# Run
setup()
loop()