# === ADC -> FFT -> OLED (SSD1306) ===
# Plataforma: ESP32 / RPi Pico W / Pico 2W
# - FFT con 'ulab' (N=1024) o DFT de respaldo (N=256)
# - Grafica señal temporal (arriba) y espectro en dB (abajo)
# - Muestra picos principales en consola

import sys, math, time, micropython
from machine import Pin, ADC, Timer, I2C

# ===== Detección de plataforma =====
PLAT = sys.platform        # 'esp32', 'rp2', 'rp2350', etc.
IS_ESP32 = (PLAT == 'esp32')
IS_RP    = ('rp2' in PLAT) or ('rp2350' in PLAT)

# ===== Parámetros =====
FS_HZ          = 4000       # Frecuencia de muestreo (Hz)
N_ULAB         = 1024       # Tamaño FFT con ulab
N_FALLBACK     = 256        # Tamaño DFT sin ulab
PEAKS_TO_SHOW  = 5          # Número de picos a listar en consola
DB_DYN         = 60.0       # Rango dinámico para graficar espectro (dB)

# Pines por defecto
ADC_PIN_ESP32  = 34         # GPIO34 (ADC1_CH6)
ADC_PIN_RP     = 26         # GP26 (ADC0)
I2C_ID         = 0
I2C_SDA_ESP32  = 21
I2C_SCL_ESP32  = 22
I2C_SDA_RP     = 4
I2C_SCL_RP     = 5
OLED_W, OLED_H = 128, 64
OLED_ADDR      = 0x3C

# ===== ulab opcional =====
HAVE_ULAB = False
try:
    from ulab import numpy as np
    HAVE_ULAB = True
except Exception:
    HAVE_ULAB = False

N = N_ULAB if HAVE_ULAB else N_FALLBACK
print("Plataforma:", PLAT, "| FS:", FS_HZ, "Hz | N:", N, "| ulab:", HAVE_ULAB)

# ===== OLED =====
from ssd1306 import SSD1306_I2C
if IS_ESP32:
    i2c = I2C(I2C_ID, scl=Pin(I2C_SCL_ESP32), sda=Pin(I2C_SDA_ESP32), freq=400_000)
else:
    i2c = I2C(I2C_ID, scl=Pin(I2C_SCL_RP),     sda=Pin(I2C_SDA_RP),     freq=400_000)
oled = SSD1306_I2C(OLED_W, OLED_H, i2c, addr=OLED_ADDR)

# ===== ADC =====
if IS_ESP32:
    adc = ADC(Pin(ADC_PIN_ESP32))
    if hasattr(adc, "atten"):
        adc.atten(ADC.ATTN_11DB)     # ~0..3.6V
    if hasattr(adc, "width"):
        adc.width(ADC.WIDTH_12BIT)   # 0..4095
    def adc_read():
        return adc.read()            # 0..4095
else:
    adc = ADC(Pin(ADC_PIN_RP))
    def adc_read():
        return (adc.read_u16() >> 4) # normaliza a ~12 bits 0..4095

# ===== Muestreo con Timer =====
samples = [0] * N
idx = 0
full = False
tim = Timer(0)
micropython.alloc_emergency_exception_buf(200)

def isr_sample(_t):
    global idx, full
    if full:
        return
    samples[idx] = adc_read()
    idx += 1
    if idx >= N:
        full = True

def start_timer(fs_hz):
    try:
        tim.init(freq=fs_hz, mode=Timer.PERIODIC, callback=isr_sample)
    except Exception:
        # Fallback aproximado por periodo en ms (menos preciso)
        period_ms = max(1, int(1000 // fs_hz))
        tim.init(period=period_ms, mode=Timer.PERIODIC, callback=isr_sample)

def stop_timer():
    tim.deinit()

# ===== DSP =====
def hann_window(n, N_):
    return 0.5 - 0.5 * math.cos(2*math.pi*n/(N_-1))

def compute_fft_ulab(samps, fs):
    x = np.array(samps, dtype=np.float)
    x = x - np.mean(x)
    n = np.arange(N)
    w = 0.5 - 0.5*np.cos(2*np.pi*n/(N-1))
    xw = x * w
    X = np.fft.fft(xw)
    mag = np.abs(X[:N//2]) / (np.sum(w)/2)  # normalización aprox
    freqs = (np.arange(N//2) * fs) / N
    return freqs, mag

def compute_dft_fallback(samps, fs):
    Nloc = len(samps)
    mean = sum(samps)/Nloc
    w = [hann_window(n, Nloc) for n in range(Nloc)]
    xw = [(samps[n]-mean)*w[n] for n in range(Nloc)]
    half = Nloc//2
    mag = [0.0]*half
    for k in range(half):
        re = 0.0; im = 0.0
        angk = -2*math.pi*k/Nloc
        for n in range(Nloc):
            ang = angk*n
            c = math.cos(ang); s = math.sin(ang)
            re += xw[n]*c; im += xw[n]*s
        mag[k] = math.sqrt(re*re + im*im)
    freqs = [k*fs/Nloc for k in range(half)]
    return freqs, mag

def top_peaks(freqs, mag, n_peaks=5, fmin=1.0, fmax=None):
    if not mag: return []
    if fmax is None: fmax = freqs[-1]
    idxs = [i for i,f in enumerate(freqs) if i>0 and f>=fmin and f<=fmax]
    idxs.sort(key=lambda i: mag[i], reverse=True)
    return [(freqs[i], mag[i], i) for i in idxs[:n_peaks]]

# ===== Gráficos en OLED =====
def draw_waveform(oled, x0, y0, w, h, samps):
    """Dibuja señal temporal en rectángulo (x0,y0,w,h)."""
    oled.fill_rect(x0, y0, w, h, 0)
    Nsrc = len(samps)
    if Nsrc < 2: return
    step = (Nsrc-1) / (w-1)
    # auto-escala simple: usa min/max
    vmin = min(samps); vmax = max(samps)
    if vmax == vmin: vmax = vmin + 1
    def mapy(v):
        # 0 en la parte baja (y0+h-1), vmax en y0
        return y0 + (h-1) - int((v - vmin) * (h-1) / (vmax - vmin))
    x_prev = x0
    y_prev = mapy(samps[0])
    for xi in range(1, w):
        idxf = int(xi * step)
        if idxf >= Nsrc: idxf = Nsrc-1
        y = mapy(samps[idxf])
        oled.line(x_prev, y_prev, x0 + xi, y, 1)
        x_prev, y_prev = x0 + xi, y

def draw_spectrum(oled, x0, y0, w, h, freqs, mag, dyn_db=60.0):
    """Dibuja espectro (dB) en rectángulo (x0,y0,w,h)."""
    oled.fill_rect(x0, y0, w, h, 0)
    if not mag: return
    mmax = max(mag)
    if mmax <= 0: mmax = 1e-12
    # Mapeo log: 20*log10(m/mmax) -> [-dyn, 0] dB
    def mapy(m):
        db = 20*math.log10(m/mmax + 1e-12)
        if db < -dyn_db: db = -dyn_db
        # db=-dyn -> y= y0+h-1 ; db=0 -> y=y0
        return y0 + (h-1) - int((db + dyn_db) * (h-1) / dyn_db)

    # Re-sample a 128 columnas
    half = len(mag)
    for xi in range(w):
        k = int(xi * (half-1) / (w-1))  # bin
        y = mapy(mag[k])
        # columna tipo "barra"
        oled.vline(x0 + xi, y, y0 + h - y, 1)

def banner(oled, txt, y=0):
    oled.fill_rect(0, y, OLED_W, 8, 0)
    oled.text(txt, 0, y, 1)

# ===== Main =====
print("Muestreando... Ctrl+C para salir.")
banner(oled, "ADC->FFT->OLED", 0); oled.show()
start_timer(FS_HZ)

try:
    while True:
        if full:
            # Detener para procesar y dibujar este frame
            stop_timer()
            local = samples[:]  # copia

            # FFT / DFT
            if HAVE_ULAB:
                freqs, mag = compute_fft_ulab(local, FS_HZ)
            else:
                print("(Sin ulab) DFT de respaldo; puede tardar…")
                freqs, mag = compute_dft_fallback(local, FS_HZ)

            # Picos (consola)
            peaks = top_peaks(freqs, mag, n_peaks=PEAKS_TO_SHOW, fmin=1.0)
            print("\n=== FFT ===  FS:", FS_HZ, "Hz | N:", N, "| Resol:", FS_HZ/N, "Hz/bin")
            if peaks:
                for f, m, k in peaks:
                    mmax = max(mag) if mag else 1.0
                    rel_db = 20*math.log10((m/(mmax+1e-12))+1e-12)
                    print("Pico: {:8.2f} Hz | bin {:4d} | {:6.2f} dB".format(f, k, rel_db))
                f0 = peaks[0][0]
            else:
                print("Sin picos notables.")
                f0 = 0.0

            # ====== Dibujo en OLED ======
            oled.fill(0)
            banner(oled, "FS:{}Hz N:{} f0:{:.0f}Hz".format(FS_HZ, N, f0), 0)

            # Área superior: time domain (alto 22 px)
            draw_waveform(oled, x0=0, y0=10, w=OLED_W, h=22, samps=local)

            # Área inferior: spectrum (alto 32 px)
            draw_spectrum(oled, x0=0, y0=36, w=OLED_W, h=28, freqs=freqs, mag=mag, dyn_db=DB_DYN)

            oled.show()

            # Preparar nueva captura
            idx = 0
            full = False
            start_timer(FS_HZ)

        time.sleep_ms(10)

except KeyboardInterrupt:
    stop_timer()
    banner(oled, "Detenido", 0); oled.show()
    print("\nDetenido por usuario.")