import machine
import utime
import urandom

# Pin setup
buttons = [
    machine.Pin(0, machine.Pin.IN, machine.Pin.PULL_DOWN),
    machine.Pin(1, machine.Pin.IN, machine.Pin.PULL_DOWN),
    machine.Pin(2, machine.Pin.IN, machine.Pin.PULL_DOWN),
    machine.Pin(3, machine.Pin.IN, machine.Pin.PULL_DOWN)
]

leds = [
    machine.Pin(4, machine.Pin.OUT),  # Red
    machine.Pin(5, machine.Pin.OUT),  # Green
    machine.Pin(6, machine.Pin.OUT),  # Blue
    machine.Pin(7, machine.Pin.OUT)   # Yellow
]

# RGB LED setup
rgb_red = machine.PWM(machine.Pin(12))
rgb_green = machine.PWM(machine.Pin(13))
rgb_blue = machine.PWM(machine.Pin(14))
rgb_red.freq(1000)
rgb_green.freq(1000)
rgb_blue.freq(1000)

speaker = machine.PWM(machine.Pin(11))
speaker.freq(1000)
speaker.duty_u16(0)

# Game settings
GAME_MODES = ["SOLO", "TIMED", "ADVANCED", "PVP"]
MAX_LEVEL = 40
LIVES = 3

# Game state
sequence = []
player_input = []
game_mode = 0
level = 1
lives = LIVES
game_active = False
player_lives = [LIVES, LIVES]
current_player = 0
streak_count = 0

def set_rgb(r, g, b):
    rgb_red.duty_u16(r * 257)
    rgb_green.duty_u16(g * 257)
    rgb_blue.duty_u16(b * 257)

def show_health(lives):
    if lives == 3: set_rgb(0, 255, 0)
    elif lives == 2: set_rgb(255, 255, 0)
    elif lives == 1: set_rgb(255, 0, 0)
    else: set_rgb(0, 0, 0)

def play_sound(freq, duration=300):
    speaker.freq(freq)
    speaker.duty_u16(32768)
    utime.sleep_ms(duration)
    speaker.duty_u16(0)

def success_sound():
    play_sound(523, 100)
    play_sound(659, 100)
    play_sound(784, 200)

def fail_sound():
    play_sound(349, 200)
    play_sound(293, 200)

def win_sound():
    play_sound(659, 200)
    play_sound(784, 200)
    play_sound(1046, 400)
    play_sound(1046, 800)

def gameover_sound():
    play_sound(220, 300)
    play_sound(196, 300)
    play_sound(165, 500)

def led_on(index):
    leds[index].value(1)

def led_off(index):
    leds[index].value(0)

def all_leds_off():
    for led in leds: led.value(0)

def blink_led(index, times=1, delay=300):
    for _ in range(times):
        led_on(index)
        utime.sleep_ms(delay)
        led_off(index)
        utime.sleep_ms(delay)

def intro_animation():
    for i in range(4):
        led_on(i)
        play_sound(300 + i*100, 100)
        utime.sleep_ms(100)
        led_off(i)
    
    for i in range(3, -1, -1):
        led_on(i)
        play_sound(700 - i*100, 100)
        utime.sleep_ms(100)
        led_off(i)
    
    for _ in range(3):
        for i in range(4): led_on(i)
        play_sound(800, 100)
        utime.sleep_ms(100)
        all_leds_off()
        utime.sleep_ms(100)

def win_animation():
    for _ in range(5):
        for i in range(4): led_on(i)
        set_rgb(255, 0, 0)
        play_sound(784, 100)
        utime.sleep_ms(100)
        
        all_leds_off()
        set_rgb(0, 255, 0)
        play_sound(988, 100)
        utime.sleep_ms(100)
        
        set_rgb(0, 0, 255)
        play_sound(1046, 100)
        utime.sleep_ms(100)
        
        set_rgb(0, 0, 0)
        utime.sleep_ms(100)

def lose_life_animation(player):
    for _ in range(3):
        if player_lives[player] == 3: set_rgb(0, 255, 0)
        elif player_lives[player] == 2: set_rgb(255, 255, 0)
        elif player_lives[player] == 1: set_rgb(255, 0, 0)
        play_sound(200, 200)
        utime.sleep_ms(200)
        set_rgb(0, 0, 0)
        utime.sleep_ms(200)

def rainbow_animation():
    colors = [(255,0,0), (255,127,0), (255,255,0), (0,255,0), (0,0,255), (75,0,130), (148,0,211)]
    for _ in range(2):
        for r,g,b in colors:
            set_rgb(r, g, b)
            utime.sleep_ms(50)
    set_rgb(0, 0, 0)

def generate_sequence(length):
    return [urandom.randint(0, 3) for _ in range(length)]

def show_sequence():
    for color in sequence:
        led_on(color)
        
        if game_mode == 2:  # Advanced mode
            show_time = urandom.randint(300, 500)
        else:
            show_time = 500
            
        play_sound(300 + color*100, show_time)
        utime.sleep_ms(100)
        led_off(color)
        utime.sleep_ms(200)

def handle_button(index):
    led_on(index)
    play_sound(300 + index*100, 100)
    player_input.append(index)
    
    if player_input[-1] != sequence[len(player_input)-1]:
        return False
    
    if len(player_input) == len(sequence):
        return True
    
    led_off(index)
    return None

def pvp_game():
    global sequence, player_input, level, player_lives, current_player, streak_count
    
    # Both players attempt same sequence + new color
    for player in range(2):
        if player_lives[player] <= 0: 
            continue
            
        current_player = player
        player_input = []
        show_health(player_lives[player])
        utime.sleep_ms(500)
        
        print(f"Player {player+1} - Level {level}")
        show_sequence()
        
        last_press = utime.ticks_ms()
        input_done = False
        correct = True
        
        while not input_done:
            now = utime.ticks_ms()
            
            # Timed mode check
            if game_mode == 1 and utime.ticks_diff(now, last_press) > 2000:
                print("Time out!")
                correct = False
                break
                
            for i in range(4):
                if buttons[i].value():
                    result = handle_button(i)
                    if result is False:
                        correct = False
                        input_done = True
                    elif result is True:
                        input_done = True
                    utime.sleep_ms(200)  # Debounce
                    last_press = now
            
            if len(player_input) == len(sequence):
                input_done = True
                
        all_leds_off()
        
        if correct:
            print("Correct!")
            streak_count += 1
            
            # Special effects every 3 levels
            if streak_count % 3 == 0:
                rainbow_animation()
                success_sound()
            
            # Add new color for next player
            if player == 1:
                sequence.append(urandom.randint(0, 3))
                level += 1
        else:
            print("Wrong!")
            fail_sound()
            streak_count = 0
            player_lives[player] -= 1
            lose_life_animation(player)
            
    # Check win/lose conditions
    if level > MAX_LEVEL:
        return "win"
    elif player_lives[0] <= 0 and player_lives[1] <= 0:
        return "lose"
    return "continue"

def solo_game():
    global sequence, player_input, level, lives, streak_count
    
    sequence = generate_sequence(level)
    print(f"Level {level}")
    show_sequence()
    
    player_input = []
    last_press = utime.ticks_ms()
    input_done = False
    correct = True
    
    while not input_done:
        now = utime.ticks_ms()
        
        # Timed mode check
        if game_mode == 1 and utime.ticks_diff(now, last_press) > urandom.randint(1000, 2000):
            print("Time out!")
            correct = False
            break
            
        for i in range(4):
            if buttons[i].value():
                result = handle_button(i)
                if result is False:
                    correct = False
                    input_done = True
                elif result is True:
                    input_done = True
                utime.sleep_ms(200)
                last_press = now
    
    all_leds_off()
    
    if correct:
        print("Correct!")
        streak_count += 1
        
        if streak_count % 3 == 0:
            rainbow_animation()
            success_sound()
        
        level += 1
        if level > MAX_LEVEL:
            return "win"
    else:
        print("Wrong!")
        fail_sound()
        streak_count = 0
        lives -= 1
        show_health(lives)
        lose_life_animation(0)
        
        if lives <= 0:
            return "lose"
    
    return "continue"

def game_loop():
    global game_mode, level, lives, player_lives, game_active, streak_count
    
    while True:
        if not game_active:
            all_leds_off()
            set_rgb(0, 0, 0)
            
            # Mode selection
            print("Select game mode:")
            for i, mode in enumerate(GAME_MODES):
                print(f"{i+1}. {mode}")
                blink_led(i, 1, 500)
                utime.sleep_ms(300)
            
            mode_chosen = False
            while not mode_chosen:
                for i in range(4):
                    if buttons[i].value():
                        game_mode = i
                        print(f"Chose: {GAME_MODES[i]}")
                        blink_led(i, 3, 200)
                        mode_chosen = True
                        utime.sleep_ms(500)
                        break
            
            # Reset game state
            game_active = True
            sequence.clear()
            player_input.clear()
            level = 1
            lives = LIVES
            player_lives = [LIVES, LIVES]
            streak_count = 0
            show_health(lives)
            
            intro_animation()
            print("Press any button to start")
            
            while not any(button.value() for button in buttons):
                utime.sleep_ms(100)
        
        # Gameplay
        if game_mode == 3:  # PVP Mode
            result = pvp_game()
        else:
            result = solo_game()
        
        # Handle game result
        if result == "win":
            win_sound()
            win_animation()
            game_active = False
            utime.sleep_ms(2000)
        elif result == "lose":
            gameover_sound()
            for _ in range(5):
                for i in range(4): led_on(i)
                utime.sleep_ms(200)
                all_leds_off()
                utime.sleep_ms(200)
            game_active = False
            utime.sleep_ms(2000)

# Start game
print("Rainbow Memory Game Ready!")
set_rgb(0, 0, 0)
all_leds_off()
game_loop()