import time
import random
import board
import digitalio
import audiomp3
import audiopwmio
import neopixel

CHASE_LENGTH = 10
NUM_PIXELS = 28

is_lightning = False
audio_enabled = True
lightning_chance = 0.0005
current_palette = 0
last_lightning_time = 0  # Last time when lightning occurred
i = 0

button_a = digitalio.DigitalInOut(board.BUTTON_A)
button_a.direction = digitalio.Direction.INPUT
button_a.pull = digitalio.Pull.DOWN

button_b = digitalio.DigitalInOut(board.BUTTON_B)
button_b.direction = digitalio.Direction.INPUT
button_b.pull = digitalio.Pull.DOWN

speaker = audiopwmio.PWMAudioOut(board.SPEAKER)
strip = neopixel.NeoPixel(board.A1, NUM_PIXELS, brightness=1)

whirl_mp3 = audiomp3.MP3Decoder(open("whirl.mp3", "rb"))
bolt_mp3 = audiomp3.MP3Decoder(open("bolt.mp3", "rb"))

weather_palettes = [
    [(0, 255, 0), (0, 128, 0), (0, 192, 0), (0, 100, 0), (0, 139, 139), (0, 0, 255), (0, 0, 128), (0, 0, 192), (0, 0, 205), (65, 105, 225)],  
    [(128, 0, 128), (0, 0, 255), (0, 255, 0), (255, 255, 0), (255, 165, 0), (255, 0, 0)],  
    [(255, 255, 0), (255, 165, 0), (255, 223, 0), (255, 236, 139)],  
    [(0, 0, 255), (0, 0, 128), (0, 0, 192), (65, 105, 225)],  
    [(0, 0, 0), (128, 128, 128), (47, 79, 79), (112, 128, 144)], 
    [(0, 128, 0), (255, 255, 255), (255, 136, 0)], 
    [(255, 0, 0), (0, 0, 255), (255, 255, 255)], 
    [(255, 165, 0), (0, 0, 0), (139, 0, 0)], 
    [(0, 128, 128), (0, 139, 139), (0, 205, 205), (72, 209, 204)],  
    [(25, 25, 112), (148, 0, 211), (72, 61, 139), (106, 90, 205)], 
    [(255, 0, 0), (255, 69, 0), (255, 99, 71), (220, 20, 60)], 
    [(255, 69, 0), (220, 20, 60), (255, 160, 122), (184, 134, 11)] 
]

def update_leds(palette_index, i):
    if is_lightning:
        strip.fill((255, 255, 255))
        time.sleep(0.1)
        strip.fill((255, 255, 0))
        time.sleep(0.1)
    else:
        colors = weather_palettes[palette_index]
        for j in range(NUM_PIXELS):
            strip[(j + i) % NUM_PIXELS] = colors[j % len(colors)]

while True:
    now = time.monotonic() * 1000  

    if audio_enabled and not speaker.playing:
        speaker.play(whirl_mp3)

    if button_a.value:
        if audio_enabled:
            speaker.stop()
            audio_enabled = False
        else:
            audio_enabled = True
        time.sleep(0.2)  

    if button_b.value:
        current_palette = (current_palette + 1) % len(weather_palettes)
        time.sleep(0.2) 

    if now - last_lightning_time > random.randint(30000, 60000):
        is_lightning = True
        last_lightning_time = now

 
    elif now - last_lightning_time < 10000:
        if random.random() < lightning_chance:
            is_lightning = True
            last_lightning_time = now

    update_leds(current_palette, i)
    is_lightning = False 
    i += 1
    if i > 255:
        i = 0