"""
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 = (0, 200, 255)


# === Animation state ===
mood_mode = "idle"
jaw_mode = "off"
jaw_amplitude = 0  # 0-100, set by bare integer commands

# Per-animation persistent state
larson_pos = 0
larson_dir = 1
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.5


def tick_larson(color):
    """Larson scanner (bouncing dot with trail)."""
    global larson_pos, larson_dir
    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()
    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 ~15%: multiply by 217/256
        v = (buf[i] * 217) >> 8
        # Random spawn: ~15% chance = getrandbits check
        if getrandbits(3) == 0:  # 1/8 = 12.5%, close enough to 15%
            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, 500, 51, 204)      # 0.5 Hz, 20-80%
    elif mode == "listening":
        tick_pulse(now_ms, color, 2000, 153, 255)     # 2.0 Hz, 60-100%
    elif mode == "responding":
        tick_wave(color)
    elif mode == "pleasure":
        tick_pulse(now_ms, color, 800, 77, 230)       # 0.8 Hz, 30-90%
    elif mode == "thinking":
        tick_larson(color)
    elif mode == "playful":
        tick_sparkle(color)
    elif mode == "commanding":
        tick_pulse(now_ms, color, 3000, 128, 255)     # 3.0 Hz, 50-100%
    elif mode == "love":
        tick_pulse(now_ms, color, 600, 102, 230)      # 0.6 Hz, 40-90%
    elif mode == "sleep":
        tick_pulse(now_ms, color, 200, 26, 77)        # 0.2 Hz, 10-30%


# === Jaw animation ticks ===

def jaw_set_amplitude(amp):
    """Set jaw LEDs based on amplitude 0-100."""
    num_lit = (amp * JAW_USABLE) // 100
    r, g, b = JAW_COLOR
    on = pack_grb(r, g, b)
    ar = jaw.ar
    ar[0] = 0  # Damaged LED always off
    for i in range(JAW_OFFSET, JAW_LEDS):
        if i - JAW_OFFSET < num_lit:
            ar[i] = on
        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
        ar = jaw.ar
        ar[0] = 0
        dim = scale_color_grb(JAW_COLOR[0], JAW_COLOR[1], JAW_COLOR[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, wave_offset
    larson_pos = 0
    larson_dir = 1
    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)