import pygame
import random

# Initialize pygame
pygame.init()

# Constants
WIDTH, HEIGHT = 800, 600
LANE_WIDTH = WIDTH // 4
NOTE_SIZE = 50  # Making notes square
NOTE_SPEED = 7  # Increased speed for difficulty
WHITE = (255, 255, 255)
BLACK = (0, 0, 0)
RED = (255, 0, 0)
GREEN = (0, 255, 0)

# Key bindings
KEYS = [pygame.K_e, pygame.K_r, pygame.K_u, pygame.K_i]

# Setup screen
screen = pygame.display.set_mode((WIDTH, HEIGHT))
pygame.display.set_caption("osu!mania Clone")
font = pygame.font.Font(None, 36)

# Note class
class Note:
    def __init__(self, lane, y):
        self.lane = lane
        self.x = lane * LANE_WIDTH + (LANE_WIDTH - NOTE_SIZE) // 2
        self.y = y
        self.size = NOTE_SIZE
        self.hit = False
    
    def move(self):
        self.y += NOTE_SPEED
    
    def draw(self, screen):
        pygame.draw.rect(screen, RED, (self.x, self.y, self.size, self.size))

# Game variables
notes = []
score = 0
running = True
clock = pygame.time.Clock()

# Generate notes randomly
for i in range(100):
    lane = random.randint(0, 3)
    y = -i * 100  # Spaced out falling notes (harder timing)
    notes.append(Note(lane, y))

# Game loop
while running:
    screen.fill(BLACK)
    
    # Draw hit zone
    pygame.draw.rect(screen, GREEN, (0, HEIGHT - 100, WIDTH, 10))
    
    # Event handling
    for event in pygame.event.get():
        if event.type == pygame.QUIT:
            running = False
        elif event.type == pygame.KEYDOWN:
            if event.key in KEYS:
                index = KEYS.index(event.key)
                for note in notes:
                    if note.lane == index and HEIGHT - 100 <= note.y <= HEIGHT - 50:
                        note.hit = True
                        score += 10
                        break
    
    # Move and draw notes
    for note in notes:
        if not note.hit:
            note.move()
            note.draw(screen)
    
    # Draw score
    score_text = font.render(f"Score: {score}", True, WHITE)
    screen.blit(score_text, (10, 10))
    
    pygame.display.flip()
    clock.tick(60)
    
pygame.quit()