"""
spectro.py — Espectrogramas de antes, después y su diferencia (para curiosos).

Frecuencia lineal hasta la máxima del fichero, para que se vea bien lo que hay por encima de 20 kHz.
Para que audios de media hora no cuesten nada, cada columna de la imagen es la media de unas pocas
FFT repartidas por su tramo (como mucho ~1600 columnas por imagen).
Diferencia: si los dos audios están alineados (misma frecuencia y misma duración) es «después − antes»
en cada punto (azul = se ha quitado, rojo = se ha añadido); si no (se han recortado silencios o pausas,
o se ha cambiado la frecuencia), se compara el espectro medio de los dos.
"""

from __future__ import annotations

from pathlib import Path

import numpy as np

import dsp
from audio_io import read_audio

NFFT = 2048
MAX_COLS = 1600
SUB = 4


def _spec(mono: np.ndarray, sr: int) -> tuple[np.ndarray, np.ndarray, float]:
    """Devuelve (dB [frecuencias × columnas], frecuencias, segundos por columna)."""
    n = len(mono)
    cols = int(min(MAX_COLS, max(1, (n - NFFT) // (NFFT // 2))))
    hop = max(1, (n - NFFT) / cols) if n > NFFT else 1
    win = np.hanning(NFFT).astype(np.float32)
    out = np.empty((NFFT // 2 + 1, cols), dtype=np.float32)
    x = np.pad(mono.astype(np.float32), (0, NFFT))
    for c in range(cols):
        starts = (c * hop + np.linspace(0, max(0.0, hop - NFFT), SUB)).astype(int) if hop > NFFT else [int(c * hop)]
        p = np.mean([np.abs(np.fft.rfft(x[s : s + NFFT] * win)) ** 2 for s in starts], axis=0)
        out[:, c] = 10 * np.log10(np.maximum(p / (NFFT * win.sum() / 2), 1e-20))
    return out, np.fft.rfftfreq(NFFT, 1 / sr), hop / sr


def make_png(before: str, after: str, png: Path) -> str | None:
    try:
        import matplotlib

        matplotlib.use("Agg")
        import matplotlib.pyplot as plt

        da, sra, _ = read_audio(before)
        db, srb, _ = read_audio(after)
        ma, mb = dsp.to_mono(da)[:, 0], dsp.to_mono(db)[:, 0]
        sa, fa, ta = _spec(ma, sra)
        sb, fb, tb = _spec(mb, srb)
        aligned = sra == srb and len(ma) == len(mb)
        fig, axes = plt.subplots(3, 1, figsize=(11, 8.4), gridspec_kw={"height_ratios": [1, 1, 1]})
        fig.patch.set_facecolor("#14141f")
        vmax = float(max(np.percentile(sa, 99.9), np.percentile(sb, 99.9)))
        for ax, S, f, t, title in ((axes[0], sa, fa, ta, "Antes"), (axes[1], sb, fb, tb, "Después")):
            im = ax.imshow(S, origin="lower", aspect="auto", cmap="magma", vmin=vmax - 110, vmax=vmax,
                           extent=[0, S.shape[1] * t, 0, f[-1] / 1000])
            ax.set_title(title, color="white", fontsize=10, loc="left")
            ax.set_ylabel("kHz", color="#ddd")
            cb = fig.colorbar(im, ax=ax, pad=0.01)
            cb.set_label("dB", color="#ddd")
            cb.ax.tick_params(colors="#ddd", labelsize=8)
        ax = axes[2]
        if aligned:
            d = np.clip(sb - sa, -30, 30)
            im = ax.imshow(d, origin="lower", aspect="auto", cmap="RdBu_r", vmin=-30, vmax=30,
                           extent=[0, d.shape[1] * ta, 0, fa[-1] / 1000])
            ax.set_title("Diferencia (después − antes): azul = se ha quitado · rojo = se ha añadido",
                         color="white", fontsize=10, loc="left")
            ax.set_ylabel("kHz", color="#ddd")
            cb = fig.colorbar(im, ax=ax, pad=0.01)
            cb.set_label("dB", color="#ddd")
            cb.ax.tick_params(colors="#ddd", labelsize=8)
            ax.set_xlabel("Tiempo (s)", color="#ddd")
        else:
            la = 10 * np.log10(np.mean(10 ** (sa / 10), axis=1))
            lb = 10 * np.log10(np.mean(10 ** (sb / 10), axis=1))
            fmax = min(fa[-1], fb[-1])
            grid = np.linspace(0, fmax, 1000)
            ia, ib = np.interp(grid, fa, la), np.interp(grid, fb, lb)
            ax.plot(grid / 1000, ia, color="#f0a050", lw=1, label="Antes")
            ax.plot(grid / 1000, ib, color="#50c0f0", lw=1, label="Después")
            ax2 = ax.twinx()
            diff = np.convolve(ib - ia, np.ones(15) / 15, mode="same")  # suavizado para leer la tendencia
            ax2.plot(grid / 1000, diff, color="#e05080", lw=1.2, alpha=.9, label="Diferencia")
            ax2.axhline(0, color="#888", lw=.6)
            ax2.set_ylabel("Diferencia (dB)", color="#e05080")
            ax2.tick_params(colors="#ddd", labelsize=8)
            ax.set_title("Espectro medio (la duración ha cambiado, así que se compara el promedio)",
                         color="white", fontsize=10, loc="left")
            ax.set_xlabel("kHz", color="#ddd")
            ax.set_ylabel("dB", color="#ddd")
            ax.legend(loc="upper right", fontsize=8)
            ax.set_facecolor("#14141f")
        for a in axes:
            a.tick_params(colors="#ddd", labelsize=8)
        axes[0].set_xlabel("")
        axes[1].set_xlabel("Tiempo (s)", color="#ddd")
        fig.tight_layout()
        fig.savefig(png, dpi=100, facecolor=fig.get_facecolor())
        plt.close(fig)
        return str(png)
    except Exception as e:  # una imagen nunca debe romper la corrección
        print(f"[ESPECTRO] {e}")
        return None