"""
Pico 2 LED Controller for Vixy's Head - Animation Engine

Runs all animations locally. Pi sends mode commands, Pico does the rest.

Hardware:
  - Strip 0 (Jaw):  GP0, 14 LEDs (index 0 damaged, use 1-13)
  - Strip 1 (Mood): GP1, 56 LEDs

Protocol:
  <integer>        - Jaw amplitude 0-100 (fire-and-forget, no response)
  mood <mode>      - Set mood animation (responds OK)
  jaw <mode>       - Set jaw animation (responds OK)

Mood modes: idle, listening, responding, pleasure, thinking, playful, commanding, love, sleep, off
Jaw modes: idle, talking, off
"""

import sys
import array
import math
import time
from machine import Pin
import rp2

try:
    from urandom import getrandbits
except ImportError:
    from random import getrandbits

# Configuration
JAW_PIN = 0
MOOD_PIN = 1
JAW_LEDS = 14
MOOD_LEDS = 56
JAW_OFFSET = 1  # Skip index 0 (damaged)
JAW_USABLE = JAW_LEDS - JAW_OFFSET  # 13

FRAME_MS = 16  # ~60fps target


# WS2812 PIO program - 800kHz timing
@rp2.asm_pio(sideset_init=rp2.PIO.OUT_LOW, out_shiftdir=rp2.PIO.SHIFT_LEFT, autopull=True, pull_thresh=24)
def ws2812():
    T1 = 2
    T2 = 5
    T3 = 3
    wrap_target()
    label("bitloop")
    out(x, 1)               .side(0)    [T3 - 1]
    jmp(not_x, "do_zero")   .side(1)    [T1 - 1]
    jmp("bitloop")          .side(1)    [T2 - 1]
    label("do_zero")
    nop()                   .side(0)    [T2 - 1]
    wrap()


class WS2812Strip:
    def __init__(self, pin, num_leds, sm_id):
        self.num_leds = num_leds
        self.sm = rp2.StateMachine(sm_id, ws2812, freq=8_000_000, sideset_base=Pin(pin))
        self.sm.active(1)
        self.ar = array.array("I", [0] * num_leds)

    def set_led(self, index, r, g, b):
        if 0 <= index < self.num_leds:
            self.ar[index] = (g << 16) | (r << 8) | b

    def fill(self, r, g, b):
        color = (g << 16) | (r << 8) | b
        for i in range(self.num_leds):
            self.ar[i] = color

    def clear(self):
        for i in range(self.num_leds):
            self.ar[i] = 0

    def show(self):
        for c in self.ar:
            self.sm.put(c, 8)
        time.sleep_us(60)


# Initialize strips
jaw = WS2812Strip(JAW_PIN, JAW_LEDS, 0)
mood = WS2812Strip(MOOD_PIN, MOOD_LEDS, 1)

jaw.clear()
jaw.show()
mood.clear()
mood.show()


# === Color definitions ===
# (r, g, b) tuples
MOOD_COLORS = {
    "idle":       (0, 255, 255),
    "listening":  (160, 255, 255),
    "responding": (100, 220, 255),
    "pleasure":   (180, 150, 255),
    "thinking":   (255, 180, 50),
    "playful":    (255, 130, 90),
    "commanding": (220, 50, 120),
    "love":       (255, 160, 180),
    "sleep":      (20, 22, 35),
}

# Jaw color breakpoints for amplitude-based color shift
# (threshold 0-100, r, g, b)
JAW_COLOR_LO = (0, 200, 255)     # Cyan at low amplitude
JAW_COLOR_MID = (200, 240, 255)  # White-ish at mid amplitude
JAW_COLOR_HI = (255, 220, 180)   # Warm at high amplitude
JAW_CENTER = JAW_OFFSET + 6      # Center LED index in strip (index 7)


# === Animation state ===
mood_mode = "idle"
jaw_mode = "off"
jaw_amplitude = 0  # 0-100, set by bare integer commands
jaw_peak = 0       # Current display level with peak+decay (0-100)

# Per-animation persistent state
larson_pos = 0
larson_dir = 1
larson_tick = 0
wave_offset = 0.0
# Sparkle buffer: store per-LED brightness as fixed-point (0-255)
sparkle_buf = array.array("B", [0] * MOOD_LEDS)


# === Helpers ===
def scale_color_grb(r, g, b, factor_256):
    """Scale RGB and return GRB packed int. factor_256 is 0-256 (fixed-point for 0.0-1.0)."""
    rr = (r * factor_256) >> 8
    gg = (g * factor_256) >> 8
    bb = (b * factor_256) >> 8
    return (gg << 16) | (rr << 8) | bb


def pack_grb(r, g, b):
    return (g << 16) | (r << 8) | b


# Precompute a sine table (64 entries, values 0-255 representing 0.0-1.0)
SINE_TABLE_SIZE = 64
_sine_table = array.array("B", [0] * SINE_TABLE_SIZE)
for _i in range(SINE_TABLE_SIZE):
    # sin maps to 0..1 range: (sin(x) + 1) / 2
    _v = (math.sin(2.0 * math.pi * _i / SINE_TABLE_SIZE) + 1.0) / 2.0
    _sine_table[_i] = int(_v * 255)


def sine_lookup(phase_256):
    """Look up sine value from table. phase_256 is 0-255 mapping to 0-2pi. Returns 0-255."""
    idx = (phase_256 * SINE_TABLE_SIZE) >> 8
    return _sine_table[idx % SINE_TABLE_SIZE]


# === Mood animation ticks ===

def tick_pulse(now_ms, color, speed_mhz, min_b, max_b):
    """Solid pulse animation. speed_mhz is milli-Hz (500 = 0.5 Hz)."""
    # phase = (now_ms * speed_mhz / 1000) mod 256
    phase = ((now_ms * speed_mhz) // 1000) & 0xFF
    raw = sine_lookup(phase)  # 0-255
    # Map to min_b..max_b range (both 0-255 scale)
    brightness = min_b + ((max_b - min_b) * raw >> 8)
    r, g, b = color
    packed = scale_color_grb(r, g, b, brightness)
    ar = mood.ar
    for i in range(MOOD_LEDS):
        ar[i] = packed
    mood.show()


def tick_wave(color):
    """Traveling wave animation."""
    global wave_offset
    r, g, b = color
    ar = mood.ar
    for i in range(MOOD_LEDS):
        # Wave phase for this LED
        phase = int(((i - wave_offset) / 15.0) * 256) & 0xFF
        raw = sine_lookup(phase)  # 0-255
        # Map to 0.2-1.0 range: 51 + (205 * raw / 255)
        brightness = 51 + ((204 * raw) >> 8)
        ar[i] = scale_color_grb(r, g, b, brightness)
    mood.show()
    wave_offset += 0.25


def tick_larson(color):
    """Larson scanner (bouncing dot with trail). Advances every 2nd frame."""
    global larson_pos, larson_dir, larson_tick
    r, g, b = color
    ar = mood.ar
    # Clear
    for i in range(MOOD_LEDS):
        ar[i] = 0
    # Bright point
    if 0 <= larson_pos < MOOD_LEDS:
        ar[larson_pos] = pack_grb(r, g, b)
    # Trail (8 LEDs behind)
    for t in range(1, 9):
        idx = larson_pos - larson_dir * t
        if 0 <= idx < MOOD_LEDS:
            # cosine fade: cos((t/8) * pi/2) mapped to 0-256
            fade = _sine_table[(SINE_TABLE_SIZE // 4 - (t * SINE_TABLE_SIZE // 32)) % SINE_TABLE_SIZE]
            ar[idx] = scale_color_grb(r, g, b, fade)
    mood.show()
    # Advance position every 2nd frame
    larson_tick += 1
    if larson_tick >= 2:
        larson_tick = 0
        larson_pos += larson_dir
        if larson_pos >= MOOD_LEDS - 1:
            larson_pos = MOOD_LEDS - 1
            larson_dir = -1
        elif larson_pos <= 0:
            larson_pos = 0
            larson_dir = 1


def tick_sparkle(color):
    """Random sparkle with decay."""
    r, g, b = color
    ar = mood.ar
    buf = sparkle_buf
    for i in range(MOOD_LEDS):
        # Decay existing brightness by ~10%: multiply by 230/256
        v = (buf[i] * 230) >> 8
        # Random spawn: ~6% chance
        if getrandbits(4) == 0:  # 1/16 = 6.25%
            v = 255
        buf[i] = v
        ar[i] = scale_color_grb(r, g, b, v)
    mood.show()


def tick_mood(now_ms):
    """Run one frame of the current mood animation."""
    mode = mood_mode
    if mode == "off":
        return

    color = MOOD_COLORS.get(mode)
    if color is None:
        return

    if mode == "idle":
        tick_pulse(now_ms, color, 250, 51, 204)      # 0.25 Hz, 20-80%
    elif mode == "listening":
        tick_pulse(now_ms, color, 1000, 153, 255)     # 1.0 Hz, 60-100%
    elif mode == "responding":
        tick_wave(color)
    elif mode == "pleasure":
        tick_pulse(now_ms, color, 400, 77, 230)       # 0.4 Hz, 30-90%
    elif mode == "thinking":
        tick_larson(color)
    elif mode == "playful":
        tick_sparkle(color)
    elif mode == "commanding":
        tick_pulse(now_ms, color, 1500, 128, 255)     # 1.5 Hz, 50-100%
    elif mode == "love":
        tick_pulse(now_ms, color, 300, 102, 230)      # 0.3 Hz, 40-90%
    elif mode == "sleep":
        tick_pulse(now_ms, color, 100, 26, 77)        # 0.1 Hz, 10-30%


# === Jaw animation ticks ===

def _jaw_color_for_amp(amp):
    """Interpolate jaw color based on amplitude 0-100. Returns (r, g, b)."""
    if amp <= 33:
        # Blend lo -> mid over 0-33
        t = (amp * 256) // 33  # 0-256
        r = JAW_COLOR_LO[0] + ((JAW_COLOR_MID[0] - JAW_COLOR_LO[0]) * t >> 8)
        g = JAW_COLOR_LO[1] + ((JAW_COLOR_MID[1] - JAW_COLOR_LO[1]) * t >> 8)
        b = JAW_COLOR_LO[2] + ((JAW_COLOR_MID[2] - JAW_COLOR_LO[2]) * t >> 8)
    else:
        # Blend mid -> hi over 34-100
        t = ((amp - 34) * 256) // 66  # 0-256
        r = JAW_COLOR_MID[0] + ((JAW_COLOR_HI[0] - JAW_COLOR_MID[0]) * t >> 8)
        g = JAW_COLOR_MID[1] + ((JAW_COLOR_HI[1] - JAW_COLOR_MID[1]) * t >> 8)
        b = JAW_COLOR_MID[2] + ((JAW_COLOR_HI[2] - JAW_COLOR_MID[2]) * t >> 8)
    return (r, g, b)


def jaw_set_amplitude(amp):
    """Symmetric VU bar from center with peak+decay and color shift."""
    global jaw_peak
    # Peak+decay: snap up, smooth decay down
    if amp >= jaw_peak:
        jaw_peak = amp
    else:
        jaw_peak = jaw_peak - (jaw_peak - amp) * 0.3
        if jaw_peak < 0.5:
            jaw_peak = 0

    display = int(jaw_peak)
    ar = jaw.ar
    ar[0] = 0  # Damaged LED always off

    if display == 0:
        for i in range(JAW_OFFSET, JAW_LEDS):
            ar[i] = 0
        jaw.show()
        return

    r, g, b = _jaw_color_for_amp(display)
    on = pack_grb(r, g, b)

    # How many LEDs per side (0-6), plus center = 1 + 2*num_side
    # At amp 100: num_side = 6 (all 13 lit). At amp ~8: num_side = 0 (center only)
    num_side_256 = (display * 6 * 256) // 100  # fixed-point 0..6*256
    num_side_full = num_side_256 >> 8           # whole LEDs per side
    edge_brightness = num_side_256 & 0xFF       # fractional part for edge LED

    for i in range(JAW_OFFSET, JAW_LEDS):
        dist = abs(i - JAW_CENTER)  # distance from center (0-6)
        if dist == 0:
            # Center always on when amp > 0
            ar[i] = on
        elif dist <= num_side_full:
            ar[i] = on
        elif dist == num_side_full + 1 and edge_brightness > 0:
            # Edge LED with fractional brightness
            ar[i] = scale_color_grb(r, g, b, edge_brightness)
        else:
            ar[i] = 0
    jaw.show()


def tick_jaw(now_ms):
    """Run one frame of jaw animation."""
    mode = jaw_mode

    if mode == "level":
        jaw_set_amplitude(jaw_amplitude)
    elif mode == "talking":
        # Autonomous talking: oscillate with some randomness
        phase = ((now_ms * 6000) // 1000) & 0xFF  # 6 Hz base
        base = sine_lookup(phase)  # 0-255
        # Add randomness: +/- 30
        jitter = (getrandbits(6) - 32)  # -32 to 31
        val = base + jitter
        if val < 0:
            val = 0
        elif val > 255:
            val = 255
        amp = (val * 100) >> 8  # Scale to 0-100
        jaw_set_amplitude(amp)
    elif mode == "idle":
        # Subtle dim glow from center
        ar = jaw.ar
        ar[0] = 0
        dim = scale_color_grb(JAW_COLOR_LO[0], JAW_COLOR_LO[1], JAW_COLOR_LO[2], 15)
        for i in range(JAW_OFFSET, JAW_LEDS):
            ar[i] = dim
        jaw.show()
    elif mode == "off":
        pass  # Already cleared on mode switch


# === Command processing ===

def process_command(line):
    """Parse incoming command. Returns True if response needed."""
    global mood_mode, jaw_mode, jaw_amplitude

    line = line.strip()
    if not line:
        return

    # Fast path: bare integer = jaw amplitude
    try:
        val = int(line)
        jaw_amplitude = max(0, min(100, val))
        jaw_mode = "level"
        return  # Fire-and-forget, no response
    except ValueError:
        pass

    parts = line.split()
    cmd = parts[0]

    if cmd == "mood" and len(parts) >= 2:
        new_mode = parts[1]
        if new_mode == "off":
            mood_mode = "off"
            mood.clear()
            mood.show()
        elif new_mode in MOOD_COLORS:
            mood_mode = new_mode
            # Reset animation state on mode switch
            _reset_mood_state()
        else:
            print("ERR: unknown mood")
            return
        print("OK")

    elif cmd == "jaw" and len(parts) >= 2:
        new_mode = parts[1]
        if new_mode in ("idle", "talking", "off"):
            jaw_mode = new_mode
            if new_mode == "off":
                jaw.clear()
                jaw.show()
            print("OK")
        else:
            print("ERR: unknown jaw mode")

    else:
        print("ERR: unknown command")


def _reset_mood_state():
    """Reset per-animation state on mode switch."""
    global larson_pos, larson_dir, larson_tick, wave_offset
    larson_pos = 0
    larson_dir = 1
    larson_tick = 0
    wave_offset = 0.0
    for i in range(MOOD_LEDS):
        sparkle_buf[i] = 0


# === Main loop ===

print("Pico LED Controller ready")
print(f"Jaw: GP{JAW_PIN}, {JAW_LEDS} LEDs")
print(f"Mood: GP{MOOD_PIN}, {MOOD_LEDS} LEDs")

# Non-blocking serial setup
import select as _sel
_poller = _sel.poll()
_poller.register(sys.stdin, _sel.POLLIN)
_read_buf = ""

while True:
    now = time.ticks_ms()

    # Drain all available serial data (non-blocking)
    while _poller.poll(0):
        ch = sys.stdin.read(1)
        if ch == '\n':
            process_command(_read_buf)
            _read_buf = ""
        else:
            _read_buf += ch

    # Tick animations
    tick_mood(now)
    tick_jaw(now)

    # Frame pacing
    elapsed = time.ticks_diff(time.ticks_ms(), now)
    if elapsed < FRAME_MS:
        time.sleep_ms(FRAME_MS - elapsed)