# esp32_animated_face.py
# Animated eyes for ESP32 + SSD1306 128x64 OLED display
# Simplified MicroPython port of Cozmo-style eyes
# =============================================================================

import math
import time
import ujson as json
from machine import Pin, I2C
import ssd1306

# ---------------------------------------------------------------------------
# CONSTANTS
# ---------------------------------------------------------------------------
FB_W, FB_H = 128, 64
EYE_COLOR = 1  # White for monochrome display
BG_COLOR = 0   # Black background

# Blink timing (seconds)
BLINK_INTERVAL_MIN = 2.0
BLINK_INTERVAL_MAX = 7.5
BLINK_CLOSE_DUR = 0.055
BLINK_HOLD_MIN = 0.018
BLINK_HOLD_MAX = 0.062
BLINK_OPEN_DUR = 0.180

# Gaze limits (pixels)
GAZE_RANGE_X = 10
GAZE_RANGE_Y = 6

# Spring damping for gaze
SPRING_OMEGA = 18.0

# ---------------------------------------------------------------------------
# HELPER FUNCTIONS
# ---------------------------------------------------------------------------
def clamp(x, a, b):
    return a if x < a else (b if x > b else x)

def lerp(a, b, t):
    return a + (b - a) * t

def ease_in_quart(t):
    t = clamp(t, 0, 1)
    return t * t * t * t

def ease_out_quint(t):
    t = clamp(t, 0, 1)
    return 1 - (1 - t) ** 5

def ease_in_out_cubic(t):
    t = clamp(t, 0, 1)
    if t < 0.5:
        return 4 * t * t * t
    return 1 - (-2 * t + 2) ** 3 / 2

# ---------------------------------------------------------------------------
# SPRING CLASS (for smooth gaze movement)
# ---------------------------------------------------------------------------
class Spring1D:
    def __init__(self, x=0.0):
        self.x = x
        self.v = 0.0
        self.target = x
    
    def step(self, dt, omega=SPRING_OMEGA):
        a = -2 * omega * self.v - omega * omega * (self.x - self.target)
        self.v += a * dt
        self.x += self.v * dt
        return self.x

# ---------------------------------------------------------------------------
# BLINK CONTROLLER
# ---------------------------------------------------------------------------
class BlinkController:
    def __init__(self):
        self.phase = -1.0  # -1 = open, 0-1 = closing/opening
        self.next_blink_time = time.time() + 3.0
        self.close_start = 0
        self.hold_dur = 0
        
    def trigger_blink(self):
        self.phase = 0.0
        self.close_start = time.time()
        self.hold_dur = BLINK_HOLD_MIN + (BLINK_HOLD_MAX - BLINK_HOLD_MIN) * 0.5
    
    def update(self, dt):
        now = time.time()
        
        # Check if it's time for automatic blink
        if self.phase < 0 and now >= self.next_blink_time:
            self.trigger_blink()
            self.next_blink_time = now + BLINK_INTERVAL_MIN + \
                                  (BLINK_INTERVAL_MAX - BLINK_INTERVAL_MIN) * 0.5
        
        # Update blink phase
        if self.phase >= 0:
            elapsed = now - self.close_start
            
            # Closing phase
            if elapsed < BLINK_CLOSE_DUR:
                self.phase = elapsed / BLINK_CLOSE_DUR
            # Hold closed
            elif elapsed < BLINK_CLOSE_DUR + self.hold_dur:
                self.phase = 1.0
            # Opening phase
            elif elapsed < BLINK_CLOSE_DUR + self.hold_dur + BLINK_OPEN_DUR:
                t = (elapsed - BLINK_CLOSE_DUR - self.hold_dur) / BLINK_OPEN_DUR
                self.phase = 1.0 - ease_out_quint(t)
            else:
                self.phase = -1.0  # Back to open
        
        return self.phase

# ---------------------------------------------------------------------------
# GAZE CONTROLLER
# ---------------------------------------------------------------------------
class GazeController:
    def __init__(self):
        self.gx = Spring1D(0)
        self.gy = Spring1D(0)
        self.next_saccade_time = time.time() + 1.5
        
    def update(self, dt):
        now = time.time()
        
        # Automatic saccades (random eye movements)
        if now >= self.next_saccade_time:
            self.gx.target = (2 * (int(time.ticks_us()) % 1000) / 1000.0 - 1.0) * 0.8
            self.gy.target = (2 * (int(time.ticks_us() / 1000) % 1000) / 1000.0 - 1.0) * 0.6
            self.next_saccade_time = now + 0.8 + (int(time.ticks_us()) % 1000) / 1000.0
        
        # Update spring physics
        self.gx.step(dt)
        self.gy.step(dt)
    
    def set_target(self, x, y):
        """Manually set gaze target (-1 to 1 range)"""
        self.gx.target = clamp(x, -1, 1)
        self.gy.target = clamp(y, -1, 1)

# ---------------------------------------------------------------------------
# EYE SHAPE PARAMETERS (from JSON)
# ---------------------------------------------------------------------------
class EyeShape:
    def __init__(self, json_data=None):
        if json_data is None:
            # Default neutral expression
            self.width_L = 34
            self.height_L = 36
            self.width_R = 34
            self.height_R = 36
            self.gap = 5
            self.center_x = 64
            self.center_y = 32
            self.corner_radius = 3
            self.upper_cover_L = 0
            self.upper_cover_R = 0
            self.lower_cover_L = 0
            self.lower_cover_R = 0
        else:
            geo = json_data.get('Geometry', {})
            pos = json_data.get('Position', {})
            lids = json_data.get('Lids', {})
            
            self.width_L = geo.get('width_L', 34)
            self.height_L = geo.get('height_L', 36)
            self.width_R = geo.get('width_R', 34)
            self.height_R = geo.get('height_R', 36)
            self.gap = geo.get('gap', 5)
            self.center_x = pos.get('center_x', 64)
            self.center_y = pos.get('center_y', 32)
            self.corner_radius = geo.get('corner_TL_L', 3)
            self.upper_cover_L = lids.get('upper_cover_L', 0)
            self.upper_cover_R = lids.get('upper_cover_R', 0)
            self.lower_cover_L = lids.get('lower_cover_L', 0)
            self.lower_cover_R = lids.get('lower_cover_R', 0)

# ---------------------------------------------------------------------------
# MAIN FACE CLASS
# ---------------------------------------------------------------------------
class AnimatedFace:
    def __init__(self, display):
        self.oled = display
        self.blink = BlinkController()
        self.gaze = GazeController()
        
        # Load eye shape
        try:
            with open('neutral.json', 'r') as f:
                json_data = json.load(f)
            self.shape = EyeShape(json_data)
        except:
            print("Could not load neutral.json, using defaults")
            self.shape = EyeShape()
        
    def draw_rounded_rect(self, x, y, w, h, r, color):
        """Draw a rounded rectangle directly on OLED"""
        # Convert float to int
        x, y, w, h, r = int(x), int(y), int(w), int(h), int(r)
        
        if w <= 0 or h <= 0:
            return
            
        # Clamp radius
        r = min(r, w // 2, h // 2)
        
        # Draw main rectangles
        self.oled.fill_rect(x + r, y, w - 2 * r, h, color)
        self.oled.fill_rect(x, y + r, w, h - 2 * r, color)
        
        # Draw corners (approximated with pixels)
        for i in range(r):
            for j in range(r):
                # Check if pixel is inside rounded corner
                dx = r - i
                dy = r - j
                if dx * dx + dy * dy <= r * r:
                    # Top-left
                    if x + i >= 0 and y + j >= 0:
                        self.oled.pixel(x + i, y + j, color)
                    # Top-right
                    if x + w - 1 - i < FB_W and y + j >= 0:
                        self.oled.pixel(x + w - 1 - i, y + j, color)
                    # Bottom-left
                    if x + i >= 0 and y + h - 1 - j < FB_H:
                        self.oled.pixel(x + i, y + h - 1 - j, color)
                    # Bottom-right
                    if x + w - 1 - i < FB_W and y + h - 1 - j < FB_H:
                        self.oled.pixel(x + w - 1 - i, y + h - 1 - j, color)
    
    def draw_eye(self, cx, cy, w, h, blink_scale, color):
        """Draw a single eye with blink animation"""
        # Apply blink scaling
        w_scaled = w * (1.0 + 0.1 * (1.0 - blink_scale))  # Slight width increase when blinking
        h_scaled = h * blink_scale
        
        if h_scaled < 2:
            # Draw blink line
            line_y = int(cy)
            if 0 <= line_y < FB_H:
                self.oled.hline(int(cx - w/2), line_y, int(w), color)
            return
        
        # Calculate eye position
        x = cx - w_scaled / 2
        y = cy - h_scaled / 2
        
        # Draw rounded rectangle eye
        self.draw_rounded_rect(x, y, w_scaled, h_scaled, 
                              self.shape.corner_radius, color)
    
    def update(self, dt):
        """Update all controllers"""
        self.blink.update(dt)
        self.gaze.update(dt)
    
    def render(self):
        """Render the face to the OLED display"""
        # Clear display
        self.oled.fill(BG_COLOR)
        
        # Calculate blink scale (0 = closed, 1 = open)
        phase = self.blink.phase
        if phase < 0:
            blink_scale = 1.0
        elif phase <= 1.0:
            blink_scale = 1.0 - ease_in_quart(phase)
        else:
            blink_scale = 1.0
        
        # Get gaze offset
        gx = self.gaze.gx.x * GAZE_RANGE_X
        gy = self.gaze.gy.x * GAZE_RANGE_Y
        
        # Calculate eye positions
        left_cx = self.shape.center_x - self.shape.gap/2 - self.shape.width_L/2 + gx
        left_cy = self.shape.center_y + gy
        right_cx = self.shape.center_x + self.shape.gap/2 + self.shape.width_R/2 + gx
        right_cy = self.shape.center_y + gy
        
        # Draw both eyes
        self.draw_eye(left_cx, left_cy, 
                     self.shape.width_L, self.shape.height_L, 
                     blink_scale, EYE_COLOR)
        self.draw_eye(right_cx, right_cy,
                     self.shape.width_R, self.shape.height_R,
                     blink_scale, EYE_COLOR)
        
        # Update display
        self.oled.show()
    
    def trigger_blink(self):
        """Manually trigger a blink"""
        self.blink.trigger_blink()
    
    def set_gaze(self, x, y):
        """Set gaze direction (-1 to 1 range)"""
        self.gaze.set_target(x, y)

# ---------------------------------------------------------------------------
# MAIN FUNCTION
# ---------------------------------------------------------------------------
def main():
    # Initialize I2C and display
    i2c = I2C(0, scl=Pin(22), sda=Pin(21), freq=400000)
    oled = ssd1306.SSD1306_I2C(FB_W, FB_H, i2c)
    
    # Create animated face
    face = AnimatedFace(oled)
    
    print("Animated Face starting...")
    print("The face will blink and look around automatically")
    
    # Main animation loop
    last_time = time.ticks_ms()
    
    try:
        while True:
            # Calculate delta time
            current_time = time.ticks_ms()
            dt = time.ticks_diff(current_time, last_time) / 1000.0
            last_time = current_time
            
            # Update and render
            face.update(dt)
            face.render()
            
            # Small delay to control frame rate (~30 FPS)
            time.sleep_ms(33)
            
    except KeyboardInterrupt:
        print("\nStopping...")
        oled.fill(0)
        oled.show()

if __name__ == "__main__":
    main()