from time import sleep, ticks_us, ticks_diff
sleep(0.1) # Wait for USB to become ready
from SSD1306 import SSD1306, SSD1306_I2C
print("Hello, Pi Pico!")

from machine import Pin, I2C, ADC
import framebuf
import math, array, gc
from random import randint
from sys import exit

SCREEN_WIDTH  = const(128)    # oled display width
SCREEN_HEIGHT = const(64)     # oled display height
MAX_SEGS = const(300)
FOOD_SIZE = const(5)
TAIL_GROW = const(20)
START_LENGTH = const(30)

SNAKE_PARAMS = const(10)
POS = const(0)
LEN = const(1)
FOOD_X = const(2)
FOOD_Y = const(3)
FOOD_READY = const(4)
NEW_LEN = const(5)

POT_PARAMS = const(4)
X = const(0)
Y = const(1)
XV = const(2)
YV = const(3)

def init_IO():
    global oled, POT, ADC_VERT, ADC_HORZ
    sda=machine.Pin(4)
    scl=machine.Pin(5)
    i2c=machine.I2C(0,sda=sda, scl=scl, freq=400_000) # 400_000
    oled = SSD1306_I2C(128, 64, i2c)
    ADC_VERT = ADC(Pin(26))
    ADC_HORZ = ADC(Pin(27))
    POT = array.array('i',0 for _ in range(POT_PARAMS))
    POT[X] = 64
    POT[Y] = 32
    POT[XV] = 1

def init_main():
    global SNAKE, BODY
    SNAKE = array.array('i',0 for _ in range(SNAKE_PARAMS))
    BODY = array.array('i',0 for _ in range(MAX_SEGS))
    SNAKE[LEN] = START_LENGTH
    SNAKE[NEW_LEN] = START_LENGTH
    for i in range(5):
        BODY[i] = SCREEN_HEIGHT//2 * SCREEN_WIDTH + SCREEN_WIDTH//2 - i 

def blk():
    oled.fill(0)
    oled.show()

def read_joystick():
    x = POT[X]
    y = POT[Y]
    x_inc = POT[XV]
    y_inc = POT[YV]
    vert = -ADC_VERT.read_u16() + 32759
    horz = -ADC_HORZ.read_u16() + 32759
    if not SNAKE[FOOD_READY]:
        if (vert + horz) != 0:
            SNAKE[FOOD_READY] = 1
            new_food()
        else:
            dummy = randint(0,100)    
    if vert > 0:
        y_inc = 1
        x_inc = 0
    if vert < 0:
        y_inc = -1
        x_inc = 0
    if horz > 0:
        x_inc = 1
        y_inc = 0
    if horz < 0:
        x_inc = -1
        y_inc = 0
    POT[XV] = x_inc
    POT[YV] = y_inc
    x += x_inc
    y += y_inc
    if 0 < x < 128 and 0 < y < 64:
        POT[X] = x
        POT[Y] = y
        BODY[0] = y * SCREEN_WIDTH + x

def new_food():
    SNAKE[FOOD_X] = randint(FOOD_SIZE,SCREEN_WIDTH-FOOD_SIZE)
    SNAKE[FOOD_Y] = randint(FOOD_SIZE,SCREEN_HEIGHT-FOOD_SIZE)

@micropython.viper
def move_snake():
    snake = ptr32(SNAKE)
    body = ptr32(BODY)
    segments = snake[LEN]
    new_length = snake[NEW_LEN]
    if segments < new_length:
        snake[LEN] += 1    
    head = body[0]
    for i in range(segments,0,-1):
        body[i] = body[i-1]
        if i > 1 and head == body[i]:
            game_over()
    

@micropython.viper
def draw():
    pot = ptr32(POT)
    snake = ptr32(SNAKE)
    body = ptr32(BODY)
    x = pot[X]
    y = pot[Y]
    segments = snake[LEN]
    for i in range(segments):
        position = body[i]
        x = position % SCREEN_WIDTH
        y = position // SCREEN_WIDTH
        oled.pixel(x,y,1)
    if SNAKE[FOOD_READY]:
        x_f = snake[FOOD_X]
        y_f = snake[FOOD_Y]
        oled.ellipse(x_f,y_f,FOOD_SIZE,FOOD_SIZE,1,1)
    oled.show()
    oled.fill(0) 
      

@micropython.viper
def check_food_collision() -> int:
    body = ptr32(BODY)
    snake = ptr32(SNAKE)  # Use Viper to access the snake array directly
    x_head = int(body[POS] % SCREEN_WIDTH)  # Calculate X position of snake's head
    y_head = int(body[POS] // SCREEN_WIDTH)  # Calculate Y position of snake's head
    x_food = int(snake[FOOD_X])  # X position of food
    y_food = int(snake[FOOD_Y])  # Y position of food

    # Calculate the distance between the snake head and the food
    # Use integer arithmetic for distance squared to avoid the need for floating-point operations
    distance_squared = ((x_head - x_food) * (x_head - x_food)) + ((y_head - y_food) * (y_head - y_food))
    radius_squared = FOOD_SIZE * FOOD_SIZE # Square of the food's radius for comparison
    # Check if the distance between the snake's head and the food is less than the food's radius
    if distance_squared <= radius_squared:
        return 1  # Collision detected
    else:
        return 0  # No collision


def game_over():
    oled.text('GAME OVER',30,20,1)
    oled.show()
    exit()


def main():
    while(1):
        gticks = ticks_us()
        read_joystick()
        move_snake()
        if SNAKE[FOOD_READY] and check_food_collision():
            SNAKE[NEW_LEN] += TAIL_GROW
            new_food()
        draw()
        diff = ticks_diff(ticks_us(),gticks)
        #print(1_000_000//diff)
        #print(gc.mem_free())

if __name__=='__main__':
    init_IO()
    init_main()
    oled.fill(0) # Black
    #oled.rect(0,0,WIDTH,HEIGHT,1,0)
    #oled.ellipse(50,43,20,20,1,1)  # Empty circle             
    oled.show()
    main()