import numpy as np
import matplotlib.pyplot as plt
import sounddevice as sd
import pywt
import threading

# PARAMETROS ---------------------------
N_CELLS = 120
FREQ_MIN = 100
FREQ_MAX = 8000
fs = 44100
frame = 2048  # ~46ms

center_freqs = np.logspace(np.log10(FREQ_MIN), np.log10(FREQ_MAX), N_CELLS)

latest_audio = None
lock = threading.Lock()

def cochlea_waveletbank(frame, fs):
    wavelet = 'morl'
    fc = pywt.central_frequency(wavelet)
    scales = fc * fs / center_freqs

    coeffs, freqs = pywt.cwt(frame, scales, wavelet, sampling_period=1/fs)
    power = np.mean(np.abs(coeffs) ** 2, axis=1)

    # Converte energia para dB (escala logarítmica)
    power_db = 10 * np.log10(power + 1e-12)

    return power_db

def audio_callback(indata, frames, time, status):
    global latest_audio
    mono = indata[:, 0]

    with lock:
        latest_audio = mono.copy()   # só copia o frame, nada pesado

plt.ion()
fig, ax = plt.subplots(figsize=(14,5))

larguras = np.diff(np.logspace(np.log10(FREQ_MIN), np.log10(FREQ_MAX), N_CELLS+1))
bars = ax.bar(center_freqs,
              np.zeros(N_CELLS),
              width=larguras,
              bottom=-80,          # <<--- faz as barras nascerem em -80 dB
              color="dodgerblue",
              alpha=0.7,
              align="edge")

ax.set_xscale("log")
ax.set_ylim(-80, 20)   # de -80 dB (silêncio) até 0 dB (máximo)
ax.set_xlabel("Frequência (Hz)")
ax.set_ylabel("Intensidade (dB)")
ax.set_title("Demonstração de Cóclea em Tempo Real")

# LOOP PRINCIPAL ---------------------------------------------------
with sd.InputStream(channels=1, samplerate=fs, blocksize=frame, callback=audio_callback):

    print("!! rodando em tempo real !!")

    while True:
        plt.pause(0.001)

        with lock:
            audio = latest_audio

        if audio is None:
            continue

        # wavelet
        ativ = cochlea_waveletbank(audio, fs)

        # atualiza grafico
        for b, h in zip(bars, ativ):
            b.set_height(h + 80)

        regions = [
            (3000, 6000,  "#0f62fe"),   # azul
            (0, 150,  "#42be65"),   # verde
            (1500, 3000, "#be95ff"),   # roxo
            (150, 375,"#ff832b"),   # laranja
            (375, 750,"#fa4d56"),   # vermelho
            (750, 1500,"#ff7eb6"),   # rosa
            (6000, 10000,"#08bdba")   # ciano
        ]

        for b, f in zip(bars, center_freqs):
            for fmin, fmax, color in regions:
                if fmin <= f < fmax:
                    b.set_color(color)
                    break

        fig.canvas.draw()
        fig.canvas.flush_events()