import pygame
import random
import time

# Initialize pygame
pygame.init()

# Constants
SCREEN_WIDTH = 800
SCREEN_HEIGHT = 600
INITIAL_NOTE_SPEED = 16  # Adjusted to match Osu!Mania scroll speed 16 
NOTE_WIDTH = 50
NOTE_HEIGHT = 50
KEYBIND = {'e': pygame.K_e, 'r': pygame.K_r, 'u': pygame.K_u, 'i': pygame.K_i}  # Proper pygame key bindings
KEY_POSITIONS = {'e': 200, 'r': 300, 'u': 400, 'i': 500}  # Adjusted order (left to right)

# Set up the screen
screen = pygame.display.set_mode((SCREEN_WIDTH, SCREEN_HEIGHT))
pygame.display.set_caption('Osu!Mania Type Game')

# Set up fonts
font = pygame.font.SysFont('Arial', 30)

# Colors
WHITE = (255, 255, 255)
RED = (255, 0, 0)
GREEN = (0, 255, 0)
BLUE = (0, 0, 255)
BLACK = (0, 0, 0)
GRAY = (100, 100, 100)

# Set up clock for FPS
clock = pygame.time.Clock()

# Note Class
class Note:
    def __init__(self, x, y, key, speed):
        self.x = x
        self.y = y
        self.key = key
        self.speed = speed

    def move(self):
        self.y += self.speed

    def draw(self):
        pygame.draw.rect(screen, WHITE, (self.x, self.y, NOTE_WIDTH, NOTE_HEIGHT))

# Function to check if a key press matches
def check_hit(note, key):
    return note.key == key and SCREEN_HEIGHT - 120 <= note.y <= SCREEN_HEIGHT - 80

# Main game function
def game():
    notes = []
    score = 0
    missed_notes = 0
    note_speed = INITIAL_NOTE_SPEED
    running = True
    hit_notifier = ""
    hit_timer = 0

    while running:
        screen.fill(BLACK)

        # Draw hit area
        for key, x in KEY_POSITIONS.items():
            pygame.draw.rect(screen, GRAY, (x, SCREEN_HEIGHT - 100, NOTE_WIDTH, NOTE_HEIGHT))
        
        # Increase note speed over time
        note_speed = min(40, note_speed + 0.002)  # Adjusted to maintain challenge at high speeds
        
        # Generate new notes at random intervals
        if random.randint(1, 100) <= 10:  # Increased spawn rate for difficulty
            key = random.choice(list(KEY_POSITIONS.keys()))
            x = KEY_POSITIONS[key]
            note = Note(x, -NOTE_HEIGHT, key, note_speed)
            notes.append(note)

        # Handle events
        for event in pygame.event.get():
            if event.type == pygame.QUIT:
                running = False
            elif event.type == pygame.KEYDOWN:
                for key, pygame_key in KEYBIND.items():
                    if event.key == pygame_key:
                        hit = False
                        for note in notes[:]:  # Copy list to avoid modification issues
                            if check_hit(note, key):
                                notes.remove(note)
                                score += 15  # Increased reward for hits
                                hit = True
                                hit_notifier = "Hit!"
                                hit_timer = pygame.time.get_ticks()
                                break
                        
        # Move and draw notes
        for note in notes[:]:  # Copy list to avoid modification issues
            note.move()
            note.draw()
            if note.y > SCREEN_HEIGHT:
                notes.remove(note)
                missed_notes += 1
                score -= 5  # Increased punishment for missing

        # Display score and missed notes
        score_text = font.render(f"Score: {score}", True, WHITE)
        miss_text = font.render(f"Missed: {missed_notes}", True, RED)
        screen.blit(score_text, (10, 10))
        screen.blit(miss_text, (10, 50))

        # Display hit notifier
        if hit_notifier and pygame.time.get_ticks() - hit_timer < 500:
            hit_text = font.render(hit_notifier, True, GREEN)
            screen.blit(hit_text, (SCREEN_WIDTH // 2 - 40, SCREEN_HEIGHT // 2))
        else:
            hit_notifier = ""

        # Refresh the screen
        pygame.display.flip()

        # Cap the frame rate
        clock.tick(60)

    pygame.quit()

if __name__ == "__main__":
    game()