import scipy.signal as signal
import pyaudio
from matplotlib import pyplot
import numpy as np
from time import sleep, time
import cupy
import atexit
import signal
from telemetrix_rpi_pico import telemetrix_rpi_pico
import colorsys

def sigint_handler(signal=None, frame=None):
    print ('KeyboardInterrupt is caught')
    board.neopixel_clear()
    sleep(0.75)
    board.shutdown()
    exit()
signal.signal(signal.SIGINT, sigint_handler)
atexit.register(sigint_handler)

NUM_LEDS = 60
board = telemetrix_rpi_pico.TelemetrixRpiPico()
board.set_pin_mode_neopixel(pin_number=2,num_pixels=NUM_LEDS)
board.neopixel_clear(auto_show=True)

SAMPLE_SIZE = 4096
SAMPLE_RATE = 48000
LOWPASS_CUTOFF = 50
AMPLITUDE_MULTIPLIER = 0.1
VELOCITY = 6

audio = pyaudio.PyAudio()

audioStream = audio.open(format=pyaudio.paInt16, channels=1, rate=SAMPLE_RATE, input=True, frames_per_buffer=SAMPLE_SIZE)

pastBassSignal = [0]*100
timecounter = 0
while True:
    data = audioStream.read(SAMPLE_SIZE)
    sample = cupy.frombuffer(data, dtype=np.int16)
    power = cupy.sum(cupy.abs(sample))/SAMPLE_SIZE
    if power > 1000:
        freq_dom = cupy.fft.rfft(sample,10000)
        freqs = cupy.fft.rfftfreq(len(freq_dom))
        power_bass = cupy.sum(cupy.abs(freq_dom[0:LOWPASS_CUTOFF]))/cupy.sum(cupy.abs(freq_dom))*power
        power_bass = max(0,power_bass-250)
        if len(pastBassSignal)>100:
            pastBassSignal.pop(0)
        pastBassSignal.append(power_bass)
        print(power_bass)

        idmax = cupy.argmax(cupy.abs(freq_dom))
        freqmax = abs(freqs[idmax]*SAMPLE_RATE)
        brightness = int(min(255,power_bass*AMPLITUDE_MULTIPLIER))
        for j in range(0,NUM_LEDS):
            r, g, b = colorsys.hsv_to_rgb(((-timecounter*VELOCITY+j*4)%360)/360,1,1)
            board.neo_pixel_set_value(j,r=int(r*brightness),g=int(g*brightness),b=int(b*brightness))
        timecounter+=1
        board.neopixel_show()
    else:
        board.neopixel_clear(auto_show=True)
audioStream.stop_stream()
audioStream.close()
audio.terminate()