import pygame, sys, random, time
from pygame.locals import *

pygame.init()

snakeSize = 30 #determines the size of snake squares
displayWidth = 500  #determines width of window in pixels
displayHeight = 500 #determines height of window in pixels
topPadding = snakeSize * 2 #how much space to leave at top for info, default 2*snakeSize

#set up window
DISPLAYSURF = pygame.display.set_mode((displayWidth,displayHeight),0,32)
pygame.display.set_caption('Snake by Isabella')

#set up colours to be used later
black = (0,0,0)
white = (230,230,230)
red = (150,0,0)
green = (0,100,0)

#set up fonts to be used later
bigFont = pygame.font.Font('freesansbold.ttf', round(snakeSize))
smallFont = pygame.font.Font('freesansbold.ttf', round(snakeSize/2))

class menuButton:
    registry = []
    def __init__(self,text,pos,size):
        self.registry.append(self)
        self.buttonText = smallFont.render(text, True, white)
        self.buttonTextRect = self.buttonText.get_rect()
        self.buttonTextRect.center = (pos)
        DISPLAYSURF.blit(self.buttonText,self.buttonTextRect)
        self.buttonRect = Rect(0,0,snakeSize*3,snakeSize)
        self.buttonRect.center = pos
        pygame.draw.rect(DISPLAYSURF,white,self.buttonRect, 2)
        self.size = size

def startMenu():
    #function to display screen that greets player at start of game
    #buttons determine size of playing field
    mouseClicked = False
    mouseX = 0
    mouseY = 0
    size = 0

    #define buttons
    tenByten = menuButton('10 x 10',(displayWidth/2,displayHeight/2),(10,10)) #(text,position, size of playing field)
    fifteenByFifteen = menuButton('15 x 15',(displayWidth/2,displayHeight/(5/3)),(15,15))
    twentyByTwenty = menuButton('20 x 20',(displayWidth/2,displayHeight/(10/7)),(20,20))
    thirtyByThirty = menuButton('30 x 30',(displayWidth/2,displayHeight/(5/4)),(30,30))

    #show text
    welcomeText = bigFont.render('WELCOME', True, green) #choose text and colour
    welcomeRect = welcomeText.get_rect() #get Rect object containing text
    welcomeRect.center = (displayWidth/2, displayHeight/3) #center rect on given coord
    DISPLAYSURF.blit(welcomeText, welcomeRect) #blit to screen

    pleaseText = smallFont.render('Please pick the size of the playing field', True, white)
    pleaseRect = pleaseText.get_rect()
    pleaseRect.center = (displayWidth/2, displayHeight/2.5)
    DISPLAYSURF.blit(pleaseText, pleaseRect)
    #update screen
    pygame.display.update()
    while size == 0: #loop to check for user input, exits as soon as a size if found
        for event in pygame.event.get(): #get events
            if event.type == QUIT: #checks if a quit action has been undertaken
                pygame.quit() #quits pygame
                sys.exit() #quits the python program
            elif event.type == MOUSEBUTTONUP: #check if mouse has been clicked
                mouseX, mouseY == event.pos #get coordinates of mouseclick
                for i in menuButton.registry: #loop through buttons
                    if i.buttonRect.collidepoint(mouseX,mouseY): #if mourse coord contained in button
                        size = i.size
                        gameLoop(highScore,size)    #start main game loop



def thankYou():
    #function to display thank you screen at end of play, disabled by default
    DISPLAYSURF.fill(black)
    thankYouText = bigFont.render('THANK YOU FOR PLAYING!', True, white)
    thankYouRect = thankYouText.get_rect()
    thankYouRect.center = (displayWidth/2, displayHeight/3)
    DISPLAYSURF.blit(thankYouText, thankYouRect)
    pygame.display.update()
    time.sleep(5) #sleep for x seconds

def fruitPicker():
        #function to find location of next fruit by picking random numbers in range width and height in boxes
        widthInBoxes = displayWidth/snakeSize
        heightInBoxes = (displayHeight-topPadding)/snakeSize #determines width and height in boxes for easy calculation
        fruit = [0,0]
        fruit[0]=random.randrange(0,widthInBoxes,1)*snakeSize+1 #picks random number and range and finds the coordinate of the corrosponding box
        fruit[1]=random.randrange(2,heightInBoxes,1)*snakeSize+1
        return fruit

def drawScreen(score,size):
    global displayWidth, displayHeight #global because fuck you that's why
    displayWidth = snakeSize*size[0]  #determines width of window in pixels
    displayHeight = snakeSize*size[1]
    DISPLAYSURF = pygame.display.set_mode((displayWidth,displayHeight),0,32)

    #function to draw the background the game takes place on
    DISPLAYSURF.fill(black)
    #draw horizontal lines
    lineX = 0
    while lineX < displayWidth:
        lineX += snakeSize
        pygame.draw.line(DISPLAYSURF, white, (lineX,topPadding),(lineX,displayHeight),1) #draw at x value lineX from top to bottom
    #draw vertical lines
    lineY = topPadding
    while lineY < displayHeight:
        pygame.draw.line(DISPLAYSURF, white, (0,lineY),(displayWidth,lineY),1) #draw at y value lineY from left edge to right edge
        lineY += snakeSize
    #The name of the game displayed in top
    byline = smallFont.render('Snake by Isabella', True, green)
    bylineRect = byline.get_rect()
    bylineRect.topleft = (20,snakeSize/2)
    DISPLAYSURF.blit(byline, bylineRect)

    #the current score displayed in top
    scoreText = smallFont.render('Score: ' + str(score), True, white)
    scoreRect = scoreText.get_rect()
    scoreRect.topright = (displayWidth - 20, snakeSize/2)
    DISPLAYSURF.blit(scoreText, scoreRect)

def gameOver(highScore,size):
    #a function to handle the game over scenario
    #starts by drawing the game over screen
    pygame.draw.rect(DISPLAYSURF,black,(0,topPadding,displayWidth,displayHeight))
    gameOverText = bigFont.render('GAME OVER!', True, green)
    tryAgainText = smallFont.render('Try Again? y/n', True, white)
    highScoreText = smallFont.render('High score: '+  str(highScore), True, white)

    gameOverRect = gameOverText.get_rect()
    gameOverRect.center = (round(displayWidth/2), round(displayHeight/4))

    tryAgainRect = tryAgainText.get_rect()
    tryAgainRect.center = (round(displayWidth/2), round(displayHeight/3))

    highScoreRect = highScoreText.get_rect()
    highScoreRect.center = (round(displayWidth/2),round(displayHeight/2))


    DISPLAYSURF.blit(gameOverText, gameOverRect)
    DISPLAYSURF.blit(tryAgainText, tryAgainRect)
    DISPLAYSURF.blit(highScoreText, highScoreRect)
    pygame.display.update()
    while True: #loop to check for user input
        for event in pygame.event.get(): #get events
            if event.type == QUIT or event.type == KEYUP and event.key == K_n: #checks if n key or a quit action has been undertaken
                pygame.quit() #quits pygame
                sys.exit() #quits the python program
            elif event.type ==  KEYDOWN and event.key == K_y: # checks if y has been pressed
                gameLoop(highScore,size) #start game over

def gameLoop(highScore,size):
    #main game loop
    snakeList = [] #list of coordinates for the snake segments
    score = 0 #the current score
    fps = 10 #framerate of game, increase to increase difficulty and smoothness.
    fpsClock = pygame.time.Clock()
    snakePosx = snakeSize + 1 #x-coordinate of the snake head
    snakePosy = snakeSize*2 + 1 #y-coordinate of the snake head
    snakeList = [(snakePosx,snakePosy)] #add starting position to segment list
    direction = '' #keeps track of the direction of the snake head
    fruit = (snakeSize*5+1,snakeSize*5+1) #stores position of fruit
    fruitCheck = False #did the snake eat a fruit last loop?
    pygame.draw.rect(DISPLAYSURF, green, (snakePosx,snakePosy, snakeSize-1, snakeSize-1))#draw starting position
    while True:
        #get events
        for event in pygame.event.get():
            if event.type == QUIT:
                #thankYou() #never enable this, I hate it.
                pygame.quit()
                sys.exit()
            #check if user pressed an arrow key
            elif event.type == KEYDOWN:
                previousDirection = direction #store direction
                if event.key == K_RIGHT:
                    if len(snakeList) == 1 or previousDirection is not 'left': #check if direction is applicable
                        direction = 'right' #match direction to arrow key
                elif event.key == K_LEFT:
                    if len(snakeList) == 1 or previousDirection is not 'right':
                        direction = 'left'
                elif event.key == K_DOWN:
                    if len(snakeList) == 1 or previousDirection  is not 'up':
                        direction = 'down'
                elif event.key == K_UP:
                    if len(snakeList) == 1 or previousDirection is not 'down':
                        direction = 'up'
        #move snake in accordance with current direction
        if direction == 'right':
            snakePosx += snakeSize
        elif direction == 'left':
            snakePosx -= snakeSize
        elif direction == 'down':
            snakePosy += snakeSize
        elif direction == 'up':
            snakePosy -= snakeSize
        snakeList.append((snakePosx,snakePosy))
        if fruitCheck == False:
            snakeList.pop(0)
        else:
            fruitCheck = False
        #check oob or collision
        if snakePosx < 0 or snakePosx > displayWidth or snakePosy < topPadding or snakePosy > displayHeight \
        or (snakePosx,snakePosy) in snakeList[0:len(snakeList)-2]:
            if score > highScore: #checks if the highscore has been beaten and corrects
                highScore = score
            gameOver(highScore,size)
        elif snakePosx == fruit[0] and snakePosy == fruit[1]:
            fruitCheck = True
            score += 1
        #draw the background
        drawScreen(score,size)
        #check if we need to find new location for fruit and draw fruit
        if fruitCheck == True:
            fruit = fruitPicker()
        pygame.draw.rect(DISPLAYSURF, red, (fruit[0],fruit[1], snakeSize-1, snakeSize-1))
        #draw new snake
        for snakePiece in snakeList:
            pygame.draw.rect(DISPLAYSURF, green, (snakePiece[0],snakePiece[1], snakeSize-1, snakeSize-1))
        pygame.display.update()
        fpsClock.tick(fps)
highScore = 0 #keeps track of the highscore
startMenu()