import serial
import pygame
import sys
import time

# Serial setup
ser = serial.Serial('/dev/ttyACM1', 115200, timeout=0.05)

# Pygame setup
pygame.init()
SCREEN_WIDTH, SCREEN_HEIGHT = 640, 480
screen = pygame.display.set_mode((SCREEN_WIDTH, SCREEN_HEIGHT))
pygame.display.set_caption("Controller GUI")

# Joystick viewer (top right)
JOYSTICK_RADIUS = 35
CIRCLE_RADIUS = 7
circle_color = (0, 200, 255)
path_color = (255, 100, 100)
bg_color = (30, 30, 30)
boundary_color = (120, 120, 120)
MARGIN = 10
CENTER_X = SCREEN_WIDTH - JOYSTICK_RADIUS - MARGIN
CENTER_Y = JOYSTICK_RADIUS + MARGIN

# Fading path
FADE_TIME = 1.5  # seconds
FPS = 60
clock = pygame.time.Clock()
path_points = []

# Calibration variables (default to full joystick range)
x_min, x_max = -127, 127
y_min, y_max = -127, 127
calibrated = False

def calibrate_joystick():
    global x_min, x_max, y_min, y_max, calibrated
    cal_duration = 3
    cal_start = time.time()
    x_min = y_min = float('inf')
    x_max = y_max = float('-inf')
    while time.time() - cal_start < cal_duration:
        try:
            line = ser.readline().decode().strip()
            if not line: continue
            parts = line.split(',')
            if len(parts) < 2: continue
            x, y = int(parts[0]), int(parts[1])
            x_min = min(x_min, x)
            x_max = max(x_max, x)
            y_min = min(y_min, y)
            y_max = max(y_max, y)
        except Exception:
            continue
    calibrated = True

def map_to_circle(val, in_min, in_max, radius):
    # Map input to -radius ... +radius
    val = max(min(val, in_max), in_min)
    return int((val - in_min) * 2 * radius / (in_max - in_min) - radius)

# Button properties
button_font = pygame.font.SysFont(None, 36)
button_rect = pygame.Rect(SCREEN_WIDTH - 220, 30, 180, 50)
button_color = (80, 180, 80)
button_color_active = (120, 220, 120)
button_text = button_font.render("Calibrate", True, (255, 255, 255))

# Instruction overlay
instruction_font = pygame.font.SysFont(None, 28)
show_instruction = False
instruction_start = 0
INSTRUCTION_TIME = 3  # seconds

# Button indicator properties
BTN_RADIUS = 18
BTN_MARGIN = 12
BTN_ROW_Y = SCREEN_HEIGHT - 60

running = True
button_pressed = False

# For displaying last values
last_joy = (0, 0)
last_angles = (0, 0, 0)
last_buttons = [0] * 9

while running:
    for event in pygame.event.get():
        if event.type == pygame.QUIT:
            running = False
        elif event.type == pygame.MOUSEBUTTONDOWN:
            if button_rect.collidepoint(event.pos):
                button_pressed = True
                show_instruction = True
                instruction_start = time.time()
                calibrate_joystick()
        elif event.type == pygame.MOUSEBUTTONUP:
            button_pressed = False

    # Read serial data
    try:
        line = ser.readline().decode().strip()
        if line:
            parts = line.split(',')
            if len(parts) >= 14:
                jx, jy = int(parts[0]), int(parts[1])
                pitch, roll, yaw = int(parts[2]), int(parts[3]), int(parts[4])
                buttons = [int(b) for b in parts[5:14]]
                last_joy = (jx, jy)
                last_angles = (pitch, roll, yaw)
                last_buttons = buttons
    except Exception:
        pass

    # Map joystick to circle
    x, y = last_joy
    dx = map_to_circle(x, x_min, x_max, JOYSTICK_RADIUS)
    dy = map_to_circle(y, y_min, y_max, JOYSTICK_RADIUS)
    dist = (dx**2 + dy**2) ** 0.5
    if dist > JOYSTICK_RADIUS:
        scale = JOYSTICK_RADIUS / dist
        dx = int(dx * scale)
        dy = int(dy * scale)
    pos = (CENTER_X + dx, CENTER_Y + dy)
    path_points.append((pos, time.time()))
    if len(path_points) > 1000:
        path_points.pop(0)

    # Draw background
    screen.fill(bg_color)

    # Draw joystick boundary
    pygame.draw.circle(screen, boundary_color, (CENTER_X, CENTER_Y), JOYSTICK_RADIUS, 3)

    # Draw fading path
    now = time.time()
    faded_points = []
    for i in range(1, len(path_points)):
        (p1, t1) = path_points[i - 1]
        (p2, t2) = path_points[i]
        age = now - t2
        if age > FADE_TIME:
            continue
        alpha = max(0, 255 - int(255 * (age / FADE_TIME)))
        path_surf = pygame.Surface((SCREEN_WIDTH, SCREEN_HEIGHT), pygame.SRCALPHA)
        pygame.draw.line(path_surf, path_color + (alpha,), p1, p2, 3)
        screen.blit(path_surf, (0, 0))
        faded_points.append((p2, t2))
    path_points = faded_points

    # Draw joystick position
    pygame.draw.circle(screen, circle_color, pos, CIRCLE_RADIUS)

    # Draw calibrate button with feedback
    current_button_color = button_color_active if button_pressed else button_color
    pygame.draw.rect(screen, current_button_color, button_rect, border_radius=8)
    screen.blit(button_text, (button_rect.x + 20, button_rect.y + 10))

    # Show instructions overlay if needed
    if show_instruction:
        if time.time() - instruction_start < INSTRUCTION_TIME:
            instruction_text = instruction_font.render(
                "Move joystick to all extremes for 3 seconds...", True, (255, 255, 255)
            )
            instruction_bg = pygame.Surface((instruction_text.get_width() + 40, instruction_text.get_height() + 20))
            instruction_bg.set_alpha(200)
            instruction_bg.fill((40, 40, 40))
            screen.blit(
                instruction_bg,
                ((SCREEN_WIDTH - instruction_bg.get_width()) // 2, SCREEN_HEIGHT - 100)
            )
            screen.blit(
                instruction_text,
                ((SCREEN_WIDTH - instruction_text.get_width()) // 2, SCREEN_HEIGHT - 90)
            )
        else:
            show_instruction = False

    # Draw pitch, roll, yaw
    angle_font = pygame.font.SysFont(None, 32)
    pitch, roll, yaw = last_angles
    angle_text = angle_font.render(f"Pitch: {pitch}°  Roll: {roll}°  Yaw: {yaw}°", True, (220, 220, 220))
    screen.blit(angle_text, (40, 40))

    # Draw button indicators
    for i, state in enumerate(last_buttons):
        cx = 60 + i * (BTN_RADIUS * 2 + BTN_MARGIN)
        cy = BTN_ROW_Y
        color = (80, 220, 80) if state else (80, 80, 80)
        pygame.draw.circle(screen, color, (cx, cy), BTN_RADIUS)
        label = angle_font.render(str(i + 1), True, (0, 0, 0) if state else (180, 180, 180))
        label_rect = label.get_rect(center=(cx, cy))
        screen.blit(label, label_rect)

    pygame.display.flip()
    clock.tick(FPS)

pygame.quit()
sys.exit()