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

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

audio = pyaudio.PyAudio()

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

pastBassSignal = [0]*100
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)
        board.neopixel_fill(r=int(min(255,power_bass*AMPLITUDE_MULTIPLIER)),g=0,b=0,auto_show=True)
    else:
        board.neopixel_clear(auto_show=True)
audioStream.stop_stream()
audioStream.close()
audio.terminate()