dsp.py

"""
dsp.py — Procesado de señal: remuestreo, filtros, detección de voz, silencios y nivel.

Convención: el audio va como float32 con forma (muestras, canales).
"""

from __future__ import annotations

from dataclasses import dataclass

import numpy as np
import soxr
from numpy.lib.stride_tricks import sliding_window_view
from scipy import signal
from scipy.ndimage import minimum_filter1d

FRAME_MS = 20
DIGITAL_SILENCE_DB = -140.0


def db(x: float) -> float:
    return 20 * np.log10(x) if x > 0 else -np.inf


def to_mono(x: np.ndarray) -> np.ndarray:
    return x.mean(axis=1, keepdims=True) if x.shape[1] > 1 else x


def resample(x: np.ndarray, sr_in: int, sr_out: int) -> np.ndarray:
    """Remuestreo de muy alta calidad (soxr VHQ, fase lineal, sin desfase)."""
    if sr_in == sr_out:
        return x
    y = soxr.resample(x, sr_in, sr_out, quality="VHQ")
    return np.asarray(y, dtype=np.float32).reshape(-1, x.shape[1])


def highpass(x: np.ndarray, sr: int, fc: float | None) -> np.ndarray:
    """Paso alto Butterworth de fase cero (−24 dB/oct efectivos). Quita retumbes, golpes graves y DC."""
    if not fc:
        return x
    sos = signal.butter(2, fc, btype="highpass", fs=sr, output="sos")
    return signal.sosfiltfilt(sos, x, axis=0).astype(np.float32)


# ── Reducción de ruido a banda completa ──────────────────────────────────────

HF_SPLIT_HZ = 20000.0


def _crossover_lowpass(x: np.ndarray, sr: int) -> np.ndarray:
    """Paso bajo FIR de fase lineal en HF_SPLIT_HZ (transición ±1 kHz); complementario exacto con x − LP(x)."""
    taps = signal.firwin(1023, HF_SPLIT_HZ, width=2000, fs=sr, window=("kaiser", 10.0))
    return signal.oaconvolve(x, taps[:, None], mode="same", axes=0).astype(np.float32)


def _hf_gain_track(x48: np.ndarray, y48: np.ndarray, n_out: int, sr_out: int, floor: float) -> np.ndarray:
    """
    Ganancia temporal para la banda > 20 kHz: cuánto ha bajado la energía de 12–20 kHz con
    DeepFilterNet en cada instante. Así los ultrasonidos se atenúan igual que el ruido audible
    y se conservan donde hay voz (sibilantes, aire).
    """
    hop = 480  # 10 ms a 48 kHz
    sos = signal.butter(4, [12000, 20000], btype="bandpass", fs=48000, output="sos")
    bx = signal.sosfiltfilt(sos, x48.mean(axis=1))
    by = signal.sosfiltfilt(sos, y48.mean(axis=1))
    nfr = len(bx) // hop
    if nfr == 0:
        return np.ones(n_out, dtype=np.float32)
    ex = (bx[: nfr * hop].reshape(nfr, hop) ** 2).mean(axis=1)
    ey = (by[: nfr * hop].reshape(nfr, hop) ** 2).mean(axis=1)
    g = np.sqrt(np.clip(ey / np.maximum(ex, 1e-20), 0, 1))
    g = signal.filtfilt(np.ones(5) / 5, [1.0], g) if nfr > 15 else g  # suaviza ±25 ms
    g = np.clip(g, floor, 1.0)
    t_fr = (np.arange(nfr) + 0.5) * hop / 48000
    t_out = np.arange(n_out) / sr_out
    return np.interp(t_out, t_fr, g).astype(np.float32)


def denoise(x: np.ndarray, sr: int, atten_db: float | None, keep_ultrasonic: bool = True,
            use_gpu: bool = True, progress=None) -> tuple[np.ndarray, dict]:
    """
    Reduce el ruido conservando la frecuencia de muestreo original.
    Si sr > 48 kHz y keep_ultrasonic, la banda > 20 kHz no pasa por el modelo (que llega a 24 kHz):
    se conserva del original con una ganancia que sigue la reducción aplicada a 12–20 kHz.
    """
    from denoise_core import DF_SR, enhance_48k

    x48 = resample(x, sr, DF_SR)
    y48, info = enhance_48k(x48.T, atten_db, use_gpu=use_gpu, progress=progress)
    y48 = y48.T
    y = resample(y48, DF_SR, sr)
    n = x.shape[0]
    y = y[:n] if y.shape[0] >= n else np.pad(y, ((0, n - y.shape[0]), (0, 0)))

    if keep_ultrasonic and sr > DF_SR:
        floor = 10 ** (-atten_db / 20) if atten_db else 0.0
        high = x - _crossover_lowpass(x, sr)
        g = _hf_gain_track(x48, y48, n, sr, floor)
        y = _crossover_lowpass(y, sr) + high * g[:, None]
        info["avisos"].append("Agudos extremos (más de 20 kHz) conservados del original")
    return y.astype(np.float32), info


# ── Detección de voz y silencios ─────────────────────────────────────────────

@dataclass
class VoiceRegion:
    start: int              # primera muestra con voz
    end: int                # última muestra con voz
    noise_db: float         # nivel estimado de ruido de fondo por trama
    start_th_db: float
    frame_db: np.ndarray    # nivel por trama (20 ms)
    frame_len: int
    found: bool


def frame_levels(mono: np.ndarray, sr: int) -> tuple[np.ndarray, int]:
    flen = max(64, int(sr * FRAME_MS / 1000))
    nfr = max(1, int(np.ceil(len(mono) / flen)))
    padded = np.pad(mono.astype(np.float64), (0, nfr * flen - len(mono)))
    p = (padded.reshape(nfr, flen) ** 2).mean(axis=1)
    return 10 * np.log10(np.maximum(p, 1e-30)), flen


def detect_voice(mono: np.ndarray, sr: int, start_above: float = 14.0, end_above: float = 10.0,
                 min_ms: int = 120) -> VoiceRegion:
    """
    Busca el primer y el último tramo con voz sostenida (≥ min_ms por encima del umbral).
    Un clic o un golpe aislado, por cortos, no cuentan como voz.
    El ruido de fondo se estima con el percentil 10 de las tramas no digitales.
    """
    fdb, flen = frame_levels(mono, sr)
    live = fdb[fdb > DIGITAL_SILENCE_DB]
    if live.size == 0:
        return VoiceRegion(0, len(mono) - 1, -np.inf, -np.inf, fdb, flen, False)
    noise = float(np.percentile(live, 10))
    loud = float(np.percentile(live, 95))
    start_th = max(noise + start_above, loud - 45)
    end_th = max(noise + end_above, loud - 50)
    k = max(2, int(round(min_ms / FRAME_MS)))
    if len(fdb) < k:
        return VoiceRegion(0, len(mono) - 1, noise, start_th, fdb, flen, False)
    win_min = sliding_window_view(fdb, k).min(axis=1)
    starts = np.flatnonzero(win_min >= start_th)
    if starts.size == 0:
        return VoiceRegion(0, len(mono) - 1, noise, start_th, fdb, flen, False)
    first = int(starts[0])
    ends = np.flatnonzero(win_min[first:] >= end_th)
    last = first + int(ends[-1]) + k - 1
    start = first * flen
    end = min(len(mono) - 1, (last + 1) * flen - 1)
    return VoiceRegion(start, end, noise, start_th, fdb, flen, True)


def detect_voice_v4(mono: np.ndarray, sr: int) -> VoiceRegion:
    """
    Réplica exacta de detectVoiceRegion() del «Analizador de audio v4» (HTML):
    ruido = media en dB de las tramas de los primeros y últimos 0,4 s; inicio = primeras 6 tramas
    seguidas ≥ ruido + 14 dB; final = últimas 6 tramas seguidas ≥ ruido + 10 dB.
    Así lo que mide y corrige la app coincide con lo que mide esa herramienta.
    """
    flen = max(64, round(sr * FRAME_MS / 1000))
    nfr = max(1, int(np.ceil(len(mono) / flen)))
    x = np.asarray(mono, dtype=np.float64)
    padded = np.pad(x, (0, nfr * flen - len(x)))
    sums = (padded.reshape(nfr, flen) ** 2).sum(axis=1)
    counts = np.full(nfr, flen, dtype=np.float64)
    counts[-1] = len(x) - (nfr - 1) * flen  # la última trama puede ser más corta (como en el HTML)
    with np.errstate(divide="ignore"):
        fdb = 10 * np.log10(sums / counts)
    edge = max(3, int(0.4 * sr / flen))
    vals = np.concatenate([fdb[: min(edge, nfr)], fdb[max(0, nfr - edge) :]])
    vals = vals[np.isfinite(vals)]
    noise = float(vals.mean()) if vals.size else -90.0
    start_th, end_th = noise + 14, noise + 10
    k = max(2, round(120 / FRAME_MS))
    fdb_safe = np.where(np.isfinite(fdb), fdb, -1e9)
    if nfr < k:
        return VoiceRegion(0, len(mono) - 1, noise, start_th, fdb_safe, flen, False)
    win_min = sliding_window_view(fdb_safe, k).min(axis=1)
    starts = np.flatnonzero(win_min >= start_th)
    if starts.size == 0:
        return VoiceRegion(0, len(mono) - 1, noise, start_th, fdb_safe, flen, False)
    first = int(starts[0])
    ends = np.flatnonzero(win_min[first:] >= end_th)
    last = first + int(ends[-1]) + k - 1 if ends.size else nfr - 1
    return VoiceRegion(first * flen, min(len(mono) - 1, (last + 1) * flen - 1), noise, start_th, fdb_safe, flen, True)


def voice_region(mono: np.ndarray, sr: int, method: str | None = None) -> VoiceRegion:
    """Detector de voz según el perfil: «analizador_v4» (compatible con el HTML) o el propio por defecto."""
    return detect_voice_v4(mono, sr) if method == "analizador_v4" else detect_voice(mono, sr)


def find_pauses(mono: np.ndarray, sr: int, region: VoiceRegion, min_s: float = 0.3) -> list[tuple[int, int]]:
    """
    Pausas internas (entre la primera y la última voz) de al menos `min_s`, en muestras [ini, fin).
    Una trama es pausa si queda por debajo del umbral de «fin de voz» (ruido + 10 dB).
    """
    if not region.found:
        return []
    flen, fdb = region.frame_len, region.frame_db
    quiet = fdb < max(region.start_th_db - 4, DIGITAL_SILENCE_DB)
    f0, f1 = region.start // flen, region.end // flen
    pauses, run = [], None
    for f in range(f0, f1 + 1):
        if quiet[f] and run is None:
            run = f
        elif not quiet[f] and run is not None:
            if (f - run) * flen >= min_s * sr:
                pauses.append((run * flen, f * flen))
            run = None
    return pauses


def shorten_pauses(x: np.ndarray, sr: int, region: VoiceRegion, max_s: float) -> tuple[np.ndarray, list[float]]:
    """
    Acorta a `max_s` las pausas internas más largas. Conserva la mitad de `max_s` pegada a cada lado
    (respiraciones y finales de palabra intactos), quita el centro y une con un fundido de 30 ms.
    Devuelve (audio, duraciones originales de las pausas acortadas).
    """
    mono = to_mono(x)[:, 0]
    long_ones = [(a, b) for a, b in find_pauses(mono, sr, region) if (b - a) > max_s * sr + 0.05 * sr]
    if not long_ones:
        return x, []
    half = int(max_s * sr / 2)
    xf = min(int(0.03 * sr), half // 2)
    fade = np.sin(np.linspace(0, np.pi / 2, xf, dtype=np.float32))[:, None] if xf else None
    pieces, pos = [], 0
    for a, b in long_ones:
        cut_a, cut_b = a + half, b - half - xf  # se elimina [cut_a, cut_b); el fundido solapa xf muestras
        left = x[pos:cut_a].copy()
        right_start = cut_b
        if xf:
            # fundido de potencia constante entre el final del tramo izquierdo y el inicio del derecho
            tail = left[-xf:] * fade[::-1]
            head = x[right_start : right_start + xf] * fade
            left[-xf:] = tail + head
            right_start += xf
        pieces.append(left)
        pos = right_start
    pieces.append(x[pos:])
    return np.concatenate(pieces).astype(np.float32), [(b - a) / sr for a, b in long_ones]


def _room_tone(x: np.ndarray, region: VoiceRegion, length: int, sr: int) -> np.ndarray:
    """
    Construye `length` muestras de "tono de sala" a partir del tramo sin voz más silencioso
    de la propia grabación, encadenando trozos con fundidos de potencia constante.
    """
    flen = region.frame_len
    seg_frames = max(5, int(0.25 * sr / flen))  # 250 ms
    fdb = region.frame_db
    if len(fdb) <= seg_frames:
        return np.zeros((length, x.shape[1]), dtype=np.float32)
    # trama inicial del tramo de 250 ms con menos energía (y sin silencio digital)
    win_mean = sliding_window_view(fdb, seg_frames).mean(axis=1)
    win_mean[sliding_window_view(fdb, seg_frames).min(axis=1) <= DIGITAL_SILENCE_DB] = np.inf
    if not np.isfinite(win_mean).any():
        return np.zeros((length, x.shape[1]), dtype=np.float32)
    i = int(np.argmin(win_mean))
    piece = x[i * flen : (i + seg_frames) * flen]
    xf = min(len(piece) // 4, int(0.02 * sr))
    fade_in = np.sin(np.linspace(0, np.pi / 2, xf))[:, None] if xf else None
    out = np.zeros((length + len(piece), x.shape[1]), dtype=np.float32)
    pos, flip = 0, False
    while pos < length:
        p = piece[::-1] if flip else piece  # alternar sentido evita repeticiones evidentes
        if pos and xf:
            out[pos - xf : pos] *= fade_in[::-1]
            p = p.copy()
            p[:xf] *= fade_in
            out[pos - xf : pos - xf + len(p)] += p
            pos += len(p) - xf
        else:
            out[pos : pos + len(p)] = p
            pos += len(p)
        flip = not flip
    return out[:length]


def set_edge_silences(x: np.ndarray, sr: int, region: VoiceRegion, target_s: float,
                      fill: str = "tono_sala", sides=("ini", "fin")) -> tuple[np.ndarray, list[str]]:
    """
    Deja exactamente `target_s` antes de la primera voz («ini») y/o después de la última («fin»).
    Solo toca los lados indicados en `sides`; el resto queda idéntico.
    """
    notes = []
    tgt = int(round(target_s * sr))
    head, tail = region.start, x.shape[0] - 1 - region.end
    core = x[region.start : region.end + 1]
    filler = (lambda n: _room_tone(x, region, n, sr)) if fill == "tono_sala" else \
             (lambda n: np.zeros((n, x.shape[1]), dtype=np.float32))

    def edge(seg_len, side):
        if side not in sides:  # lado que ya cumple: intacto
            return x[: region.start] if side == "ini" else x[region.end + 1 :]
        if seg_len >= tgt:  # recortar (se conserva el tramo pegado a la voz)
            return x[region.start - tgt : region.start] if side == "ini" else x[region.end + 1 : region.end + 1 + tgt]
        extra = filler(tgt - seg_len)
        if side == "ini":
            return np.concatenate([extra, x[: region.start]])
        return np.concatenate([x[region.end + 1 :], extra])

    pre, post = edge(head, "ini"), edge(tail, "fin")
    out = np.concatenate([pre, core, post]).astype(np.float32)
    # fundidos muy cortos en los extremos del fichero (evitan clics al cortar)
    f = min(int(0.005 * sr), len(out) // 2)
    if f:
        ramp = np.linspace(0, 1, f, dtype=np.float32)[:, None]
        if "ini" in sides:
            out[:f] *= ramp
        if "fin" in sides:
            out[-f:] *= ramp[::-1]
    for nombre, antes, side in (("inicial", head, "ini"), ("final", tail, "fin")):
        if side in sides and abs(antes - tgt) > sr * 0.001:
            accion = "recortado" if antes > tgt else ("alargado con el silencio natural de la sala" if fill == "tono_sala" else "alargado con silencio")
            notes.append(f"Silencio {nombre}: {antes / sr:.2f} s → {target_s:.2f} s ({accion})")
    return out, notes


def pause_expander(x: np.ndarray, sr: int, range_db: float = 15.0) -> tuple[np.ndarray, dict]:
    """
    Expansor suave: baja `range_db` el ruido residual en las pausas sin tocar la voz.
    La zona de voz se ensancha 100 ms antes y 200 ms después (no se comen finales de palabra
    ni respiraciones) y la ganancia se suaviza ~100 ms para que no haya «bombeo».
    """
    mono = to_mono(x)[:, 0]
    fdb, flen = frame_levels(mono, sr)
    live = fdb[fdb > DIGITAL_SILENCE_DB]
    if live.size == 0:
        return x, {"pausas_pct": 0.0}
    noise = float(np.percentile(live, 10))
    loud = float(np.percentile(live, 95))
    core = fdb >= max(noise + 10, loud - 50)
    voice = core.copy()
    for s in range(1, int(100 / FRAME_MS) + 1):   # 100 ms antes de cada trama con voz
        voice[:-s] |= core[s:]
    for s in range(1, int(200 / FRAME_MS) + 1):   # 200 ms después
        voice[s:] |= core[:-s]
    g = np.where(voice, 1.0, 10 ** (-range_db / 20))
    k = np.hanning(7)[1:-1]
    g = np.convolve(np.pad(g, 2, mode="edge"), k / k.sum(), mode="valid")
    g = np.maximum(g, np.where(voice, 1.0, 0.0))  # la voz nunca se atenúa
    t_fr = (np.arange(len(g)) + 0.5) * flen
    gs = np.interp(np.arange(x.shape[0]), t_fr, g).astype(np.float32)
    return (x * gs[:, None]).astype(np.float32), {"pausas_pct": float(100 * (~voice).mean())}


# ── Nivel ────────────────────────────────────────────────────────────────────

def peak_limit(x: np.ndarray, sr: int, ceiling_db: float, window_ms: float = 5.0) -> tuple[np.ndarray, dict]:
    """
    Limitador de picos sin sobreoscilaciones (offline, con anticipación):
    ganancia = mínimo deslizante de (techo / |x|) suavizado con una ventana del mismo ancho.
    Garantiza |salida| ≤ techo. Solo actúa alrededor de los picos que superan el techo.
    """
    ceiling = 10 ** (ceiling_db / 20)
    peak = np.abs(x).max(axis=1)
    if peak.max() <= ceiling:
        return x, {"muestras_afectadas": 0, "reduccion_max_db": 0.0}
    w = max(1, int(sr * window_ms / 1000))
    g = np.minimum(1.0, ceiling / np.maximum(peak, 1e-12))
    g = minimum_filter1d(g, size=2 * w + 1, mode="nearest")
    win = np.hanning(2 * w + 3)[1:-1]
    g = signal.oaconvolve(np.pad(g, w, mode="edge"), win / win.sum(), mode="valid")
    y = (x * g[:, None]).astype(np.float32)
    # seguridad numérica: nunca por encima del techo
    over = np.abs(y) > ceiling
    if over.any():
        y[over] = np.sign(y[over]) * ceiling
    return y, {"muestras_afectadas": int((g < 0.999).sum()), "reduccion_max_db": float(-db(g.min()))}


def rms_db(x: np.ndarray) -> float:
    return db(float(np.sqrt(np.mean(np.asarray(x, dtype=np.float64) ** 2)))) if x.size else -np.inf


# ── Ruidos sueltos en las pausas (cláxones, golpes, puertas…) ────────────────

EVENT_ABOVE_DB = 12.0    # un ruido suelto destaca al menos 12 dB sobre el ruido de fondo
EVENT_GUARD_S = 0.25     # no se toca nada a menos de 250 ms de la voz (inicios y finales de palabra)
EVENT_LOUD_DB = 15.0     # y llega a menos de 15 dB del volumen de la voz (respiraciones y ruiditos se quedan)
EVENT_MAX_S = 8.0        # más largo que esto no es un «ruido suelto» (música, obra…): se deja


def remove_events(x: np.ndarray, sr: int, speech_prob: np.ndarray, hop_s: float) -> tuple[np.ndarray, list]:
    """
    Sustituye por tono de sala los ruidos fuertes que caen en las pausas (donde el detector neuronal
    dice que NO hay voz, el nivel sube ≥ 12 dB sobre el fondo y llega cerca del volumen de la voz).
    La voz y sus alrededores no se tocan; las respiraciones y ruiditos suaves, tampoco. Si el detector no
    encuentra voz en todo el audio (risas, llantos…), no se toca nada.
    speech_prob: probabilidad de voz por ventana de `hop_s` segundos. Devuelve (audio, [(inicio_s, dur_s, dB)]).
    """
    mono = to_mono(x)[:, 0]
    hop = hop_s * sr
    n = min(len(speech_prob), int(len(mono) / hop))
    if n < 10:
        return x, []
    idx = (np.arange(n + 1) * hop).astype(int)
    lvl = np.array([rms_db(mono[idx[i] : idx[i + 1]]) for i in range(n)])
    live = lvl[lvl > DIGITAL_SILENCE_DB]
    if live.size < 10:
        return x, []
    floor = float(np.percentile(live, 10))
    p = speech_prob[:n]
    speech = p > 0.4
    if speech.sum() < 10:
        return x, []
    loud = float(np.median(lvl[p > 0.5])) - EVENT_LOUD_DB if (p > 0.5).any() else float(np.median(lvl[speech])) - EVENT_LOUD_DB
    g = max(1, int(round(EVENT_GUARD_S / hop_s)))
    near = np.convolve(speech.astype(int), np.ones(2 * g + 1, dtype=int), mode="same") > 0
    cand = (p < 0.1) & (lvl > floor + EVENT_ABOVE_DB) & ~near
    # ensanchar 2 ventanas (colas del ruido) sin entrar en la zona protegida
    cand = (np.convolve(cand.astype(int), np.ones(5, dtype=int), mode="same") > 0) & ~near
    events, i = [], 0
    while i < n:
        if not cand[i]:
            i += 1
            continue
        j = i
        while j + 1 < n and (cand[j + 1] or (j + 3 < n and cand[j + 2:j + 4].any())):
            j += 1
        if (j + 1 - i) * hop_s <= EVENT_MAX_S and lvl[i : j + 1].max() >= loud:
            events.append((i, j + 1))
        i = j + 1
    if not events:
        return x, []
    region = voice_region(mono, sr)
    y = x.copy()
    xf = int(0.015 * sr)
    out = []
    for a, b in events:
        s0, s1 = idx[a], idx[min(b, n)]
        tone = _room_tone(x, region, s1 - s0, sr)
        w = np.ones(s1 - s0, dtype=np.float32)
        k = min(xf, (s1 - s0) // 2)
        if k:
            ramp = np.sin(np.linspace(0, np.pi / 2, k)) ** 2
            w[:k], w[-k:] = ramp, ramp[::-1]
        y[s0:s1] = x[s0:s1] * (1 - w[:, None]) + tone[: s1 - s0] * w[:, None]
        out.append((s0 / sr, (s1 - s0) / sr, float(lvl[a:b].max() - floor)))
    return y.astype(np.float32), out


def remove_hum(x: np.ndarray, sr: int, harmonics: list[float], q: float = 35.0) -> np.ndarray:
    """Filtros de muesca muy estrechos (Q 35: ~1,5 Hz a 50 Hz) en cada armónico del zumbido, sin desfase."""
    y = x.astype(np.float64)
    for f0 in harmonics:
        if f0 < sr / 2 - 50:
            b, a = signal.iirnotch(f0, q, fs=sr)
            y = signal.filtfilt(b, a, y, axis=0)
    return y.astype(np.float32)