Video game not responding correctly to input - pygame

Question

I'm developing a pygame video game, and everything was working out perfectly until yesterday. The issues began after I formatted my pc. So when i run the game, the first screen to show up is the 'Menu'. So in this state class I have an event method where when you press the 'p' key it gets you to the 'Play' state. So now it is not working, I don't know why.
I've changed nothing. I just formatted my pc and reinstalled python, pygame and pgu module. But the strange thing comes when I reprogram the videogame so that the first state to show up when you run the game is the 'Play' state, everything works perfectly. It also has an event method where when you press the arrows, the character moves, and when the player presses ESC it takes you to the 'Menu' state.
So again when I'm at the 'Menu' state the game doesn't respond to the input I'm giving to it. I don't really know what's happening.

Answer 1

Here is an example of what I was saying in comments :
Sorry if all comments are not appropriate to your level but I wanted to be sure you understand it all.
import pygame
screen = pygame.display.set_mode((1000, 1000))
class Menu:
pass
class Play:
pass
def main():
running = True # here we define the main variable of the main loop
main_menu = True # when the main loop will begin, main_menu will begin too
game = False # game is false because player didn't press p
options = False # options is false because player didn't go to options
while running: # begin the main loop
while main_menu:
for event in pygame.event.get(): # listen for events
if event.type == pygame.QUIT: # if the player quit, ends up the 'main_menu' loop and 'running' loop too
running = False
main_menu = False
if event.type == pygame.KEYDOWN:
if event.key == pygame.K_p: # if p key is pressed, exit from main menu and begin 'game' loop
main_menu = False
game = True
if event.key == pygame.K_o: # if o key is pressed, exit from main menu and begin 'options' loop
main_menu = False
options = True
screen.fill((255, 0, 0)) # I fill the screen in red to make the example more explicit
pygame.display.flip() # I update the screen every frame
while game:
for event in pygame.event.get(): # listen for events
if event.type == pygame.QUIT: # if the player quit, ends up 'running' and 'options' loop
running = False
game = False
if event.type == pygame.KEYDOWN: # listen for keys
if event.key == pygame.K_BACKSPACE: # if the player press backspace (delete), ends up 'game' loop
game = False # and begin (again) the 'main_menu' loop
main_menu = True
screen.fill((255, 255, 255)) # I fill the screen in white to make the example more explicit
pygame.display.flip() # I update the screen every frame
while options: # same logic for this one
for event in pygame.event.get():
if event.type == pygame.QUIT:
running = False
options = False
if event.type == pygame.KEYDOWN:
if event.key == pygame.K_BACKSPACE:
options = False
main_menu = True
screen.fill((0, 255, 0))
pygame.display.flip()
pygame.init() # initialize pygame
main() # begin main loop
pygame.quit() # quit pygame