import time
import board
import array
import math
import neopixel
import ulab.numpy as np

LED_PIN = board.GP26 
NUM_LEDS = 120
BRIGHTNESS = 0.3
ORDER = neopixel.GRB
pixels = neopixel.NeoPixel(LED_PIN, NUM_LEDS, brightness=BRIGHTNESS, auto_write=False, pixel_order=ORDER)

NUM_SAMPLES = 256
SAMPLE_RATE = 16000

def constrain(value, min_val, max_val):
    return max(min_val, min(value, max_val))

def hsv_to_rgb(h, s, v):
    if s == 0.0: return (v, v, v)
    i = int(h * 6.)
    f = (h * 6.) - i
    p = v * (1. - s)
    q = v * (1. - s * f)
    t = v * (1. - s * (1. - f))
    i %= 6
    if i == 0: return (v, t, p)
    if i == 1: return (q, v, p)
    if i == 2: return (p, v, t)
    if i == 3: return (p, q, v)
    if i == 4: return (t, p, v)
    if i == 5: return (v, p, q)

NUM_BANDS = min(NUM_LEDS, NUM_SAMPLES // 2)

mock_time = 0.0
mock_peak_positions = [NUM_SAMPLES // 16, NUM_SAMPLES // 8, NUM_SAMPLES // 4]
mock_peak_widths = [2, 3, 4]
mock_peak_amps_base = [0.7, 0.9, 0.6]

def generate_mock_audio_data(num_fft_bins):
    global mock_time
    magnitudes = [0.0] * num_fft_bins

    amp1 = mock_peak_amps_base[0] * (0.5 + 0.5 * math.sin(mock_time * 1.5))
    pos1 = int(mock_peak_positions[0] + 3 * math.sin(mock_time * 0.5))
    for i in range(-mock_peak_widths[0], mock_peak_widths[0] + 1):
        bin_idx = constrain(pos1 + i, 0, num_fft_bins - 1)
        dist_from_center = abs(i) / mock_peak_widths[0]
        mag_val = amp1 * math.exp(-dist_from_center**2)
        magnitudes[bin_idx] = max(magnitudes[bin_idx], mag_val * 700)

    amp2 = mock_peak_amps_base[1] * (0.6 + 0.4 * math.sin(mock_time * 1.0 + 1.0))
    pos2 = int(mock_peak_positions[1] + 5 * math.sin(mock_time * 0.8 + 0.5))
    for i in range(-mock_peak_widths[1], mock_peak_widths[1] + 1):
        bin_idx = constrain(pos2 + i, 0, num_fft_bins - 1)
        dist_from_center = abs(i) / mock_peak_widths[1]
        mag_val = amp2 * math.exp(-dist_from_center**2)
        magnitudes[bin_idx] = max(magnitudes[bin_idx], mag_val * 900)

    amp3 = mock_peak_amps_base[2] * (0.4 + 0.6 * abs(math.sin(mock_time * 2.5)))
    pos3 = int(mock_peak_positions[2])
    for i in range(-mock_peak_widths[2], mock_peak_widths[2] + 1):
        bin_idx = constrain(pos3 + i, 0, num_fft_bins - 1)
        dist_from_center = abs(i) / mock_peak_widths[2]
        mag_val = amp3 * math.exp(-dist_from_center**2)
        magnitudes[bin_idx] = max(magnitudes[bin_idx], mag_val * 600)

    for i in range(num_fft_bins):
        magnitudes[i] += (0.02 + 0.02 * math.sin(mock_time * 0.2 + i * 0.1)) * 50

    mock_time += 0.1
    return magnitudes


print("Starting MOCK audio visualizer for Raspberry Pi Pico...")
print(f"LED Pin: {LED_PIN}, Num LEDs: {NUM_LEDS}")
print(f"FFT Size (mock): {NUM_SAMPLES}, Bands: {NUM_BANDS}")

prev_mags = [0.0] * NUM_BANDS
smoothing_factor = 0.4
color_offset_speed = 0.005
current_color_offset = 0.0

while True:
    try:
        magnitudes = generate_mock_audio_data(NUM_SAMPLES // 2)
        pixels.fill(0)

        for i in range(NUM_BANDS):
            if i < len(magnitudes):
                raw_mag = magnitudes[i]
            else:
                raw_mag = 0

            current_mag = (raw_mag * (1 - smoothing_factor)) + (prev_mags[i] * smoothing_factor)
            prev_mags[i] = current_mag

            visualization_scaling_factor = 700
            normalized_mag = constrain(current_mag / visualization_scaling_factor, 0, 1.0)

            if normalized_mag > 0.01:
                hue = (i / NUM_BANDS + current_color_offset) % 1.0
                r, g, b = hsv_to_rgb(hue, 1.0, normalized_mag)

                r_int = int(r * 255)
                g_int = int(g * 255)
                b_int = int(b * 255)

                if i < NUM_LEDS:
                    pixels[i] = (r_int, g_int, b_int)

        pixels.show()
        current_color_offset = (current_color_offset + color_offset_speed) % 1.0
        time.sleep(0.05)

    except Exception as e:
        print(f"Error in main loop: {e}")
        time.sleep(0.1)