import pygame
import sys

class main(object):
    def __init__(self):
    #initialize pygame
        pygame.init()
    #create display
        self.window = pygame.display.set_mode([640,480])
    #current menu state
        self.state = 'title'
    #list of menu/game classes
        self.menus = {'title':titlescreen(), 'game':game()}

    #primary update loop
        while True:
        #check for key press events
            for event in pygame.event.get():
                if event.type == pygame.KEYDOWN:
                    key = event
                else:
                    key = None

        #update and pass key presses and menu state to active menu
            self.state = self.menus[self.state].update(key,self.state)

        #update window
            pygame.display.update()


#titlescreen menu class
class titlescreen(object):
    def __init__(self):
    #menu option
        self.option = 0

    def update(self,event,state):
    #if a key has been pressed
        if event != None:
        #move up in the menu
            if event.key == pygame.K_UP:
                if self.option > 0:
                    print(self.option)
                    self.option -= 1

        #move down in the menu
            elif event.key == pygame.K_DOWN:
                if self.option < 1:
                    print(self.option)
                    self.option += 1

        #select a menu option
            elif event.key == pygame.K_RETURN:
            #if option is 0 return game state and play tone
                if self.option == 0:
                    print('game menu')
                    return 'game'
            #quit game
                elif self.option == 1:
                    print('quit')
                    pygame.quit()
                    sys.exit(0)

    #return current state if no changes
        return state


#main game class
class game(object):
    def __init__(self):
        pass

#DEMO: return to titlescreen
    def update(self,key,state):
        return 'title'


a = main()