"""Salience -> f0 for the exported pitch models (numpy only).

CREPE and RMVPE share one 360-bin grid: 20 cents per bin from 1997.38 cents (32.70 Hz) re 10 Hz.
FCPE spaces its 360 bins evenly between 32.7 and 1975.5 Hz in cents (~19.8 cents per bin).
Decoding is the models' own: a salience-weighted mean of the 9 bins around the peak, voiced when
the peak clears a threshold.
"""
import numpy as np

CREPE_CENTS = 1997.3794084376191 + 20.0 * np.arange(360)
FCPE_CENTS = np.linspace(1200 * np.log2(32.7 / 10), 1200 * np.log2(1975.5 / 10), 360)
GENERIC_CENTS = (100.0 / 3) * np.arange(384)          # PESTO activations, 1/3 semitone (parity only)

CENTS = {"crepe": CREPE_CENTS, "rmvpe": CREPE_CENTS, "fcpe": FCPE_CENTS, "pesto-act": GENERIC_CENTS}
# each model's reference threshold on the peak salience (CREPE: torchcrepe's documented 0.21
# periodicity, RMVPE: RVC's 0.03, FCPE: torchfcpe's inference default 0.006)
THRESHOLD = {"crepe": 0.21, "rmvpe": 0.03, "fcpe": 0.006, "pesto-act": 0.0}


def local_average_cents(sal: np.ndarray, cents: np.ndarray, width: int = 4) -> np.ndarray:
    """sal [T, K] -> cents [T]: salience-weighted mean of the bins within `width` of the peak."""
    t, k = sal.shape
    peak = sal.argmax(1)
    idx = np.clip(peak[:, None] + np.arange(-width, width + 1)[None, :], 0, k - 1)
    w = np.take_along_axis(sal, idx, 1)
    w = np.maximum(w, 0)
    return (w * cents[idx]).sum(1) / np.maximum(w.sum(1), 1e-12)


def decode(sal: np.ndarray, kind: str, threshold: float | None = None):
    """sal [1, T, K] or [T, K] -> (f0 Hz [T], voiced bool [T]). Unvoiced frames get f0 0."""
    s = np.asarray(sal, np.float64)
    if s.ndim == 3:
        s = s[0]
    th = THRESHOLD[kind] if threshold is None else threshold
    cents = local_average_cents(s, CENTS[kind])
    voiced = s.max(1) > th
    f0 = np.where(voiced, 10.0 * 2 ** (cents / 1200.0), 0.0)
    return f0, voiced