"""SBR-style high-frequency bandwidth extension (the "sounds like AI" fix).

Suno/codec output rolls off ~14-16 kHz. Copy the octave below the detected
roll-off above it (STFT bin shift), shape with a decaying envelope matched to
the spectral slope below roll-off, blend shaped noise, boost on transient
frames. No-op when input already extends past 0.85*Nyquist.
"""
import numpy as np
from scipy.ndimage import median_filter, uniform_filter1d
from scipy.signal import stft, istft

NFFT = 4096
HOP = NFFT // 4
NOISE_BLEND = 0.45
PEAK_ABOVE_MEDIAN_DB = 6.0
SLOPE_DB_PER_OCT = -3.0
TRANSIENT_BOOST_DB = 3.0


def detect_rolloff(audio, sr):
    mono = audio.mean(axis=1) if audio.ndim == 2 else audio
    f, t, Z = stft(mono, fs=sr, nperseg=NFFT, noverlap=NFFT - HOP)
    mag = np.mean(np.abs(Z), axis=1)
    # smooth 1/6 octave
    smooth = np.convolve(mag, np.hanning(9) / np.hanning(9).sum(), mode="same")
    ref_band = smooth[(f >= 1000) & (f <= 4000)].mean()
    floor = ref_band * 10 ** (-30 / 20.0)
    above = np.where((f > 4000) & (smooth < floor))[0]
    return float(f[above[0]]) if len(above) else float(f[-1])


def extend_hf(audio, sr, rolloff_hz=None):
    nyq = sr / 2.0
    roll = detect_rolloff(audio, sr) if rolloff_hz is None else rolloff_hz
    if roll >= 0.85 * nyq:
        return audio

    out = np.empty_like(audio)
    rng = np.random.default_rng(0x5B12)
    for ch in range(audio.shape[1]):
        f, t, Z = stft(audio[:, ch], fs=sr, nperseg=NFFT, noverlap=NFFT - HOP)
        df = f[1] - f[0]
        i_roll = int(roll / df)
        i_src_lo = max(1, i_roll // 2)                  # octave below roll-off
        n_copy = min(i_roll - i_src_lo, len(f) - i_roll)
        if n_copy <= 0:
            out[:, ch] = audio[:, ch]
            continue

        src = Z[i_src_lo:i_src_lo + n_copy, :].copy()
        # suppress tonal peaks before bin-shift (avoids replicated whistle images)
        mag = np.abs(src)
        med_mag = median_filter(mag, size=(5, 1))
        capped = np.minimum(mag, med_mag * 10 ** (PEAK_ABOVE_MEDIAN_DB / 20.0))
        src = capped * np.exp(1j * np.angle(src))

        ext_freqs = f[i_roll:i_roll + n_copy]
        octaves_up = np.log2((ext_freqs + 1e-9) / roll)
        decay = 10 ** ((SLOPE_DB_PER_OCT * octaves_up) / 20.0)
        # per-frame splice matching: the extension level follows the program
        # instead of sitting at a constant track-mean level — a static gain
        # leaves an audible constant hiss bed in quiet sections (the
        # tell-tale "sounds like AI" artifact this module exists to fix)
        splice_ref_t = np.abs(Z[max(1, i_roll - 4):i_roll, :]).mean(axis=0)
        src_ref_t = np.abs(src[:4, :]).mean(axis=0)
        ratio_t = splice_ref_t / (src_ref_t + 1e-12)
        ratio_t = uniform_filter1d(ratio_t, size=9)        # ~smooth over 200 ms
        shifted = src * ratio_t[None, :] * decay[:, None]

        # noise obeys the same decay envelope (prevents hiss overshoot at Nyquist)
        noise_mag = np.abs(shifted) / np.sqrt(2)
        noise = (rng.standard_normal(shifted.shape) +
                 1j * rng.standard_normal(shifted.shape)) * noise_mag
        blended = shifted * (1 - NOISE_BLEND) + noise * NOISE_BLEND

        # hard-cap every extended bin to the per-frame slope-predicted level
        # (cap BEFORE the transient boost so the boost is not erased)
        cap = decay[:, None] * uniform_filter1d(splice_ref_t, size=9)[None, :]
        bmag = np.abs(blended)
        blended = np.minimum(bmag, cap) * np.exp(1j * np.angle(blended))

        # transient boost only in the first extended octave (not the top octave);
        # flux measured on the HF source band, not the full spectrum — otherwise
        # kick/bass hits trigger HF extension boosts
        frame_mag = np.abs(Z[i_src_lo:i_roll, :]).sum(axis=0)
        flux = np.maximum(0, np.diff(frame_mag, prepend=frame_mag[0]))
        thresh = np.quantile(flux, 0.9)
        boost = np.where(flux >= thresh, 10 ** (TRANSIENT_BOOST_DB / 20.0), 1.0)
        boost_mask = (octaves_up <= 1.0)[:, None]
        blended *= 1.0 + (boost[None, :] - 1.0) * boost_mask

        # crossfade +-1 kHz around the splice
        xfade_bins = max(1, int(1000 / df))
        ramp = np.linspace(0, 1, min(xfade_bins, n_copy))[:, None]
        Z[i_roll:i_roll + len(ramp), :] = (Z[i_roll:i_roll + len(ramp), :] * (1 - ramp) +
                                           blended[:len(ramp), :] * ramp)
        if n_copy > len(ramp):
            Z[i_roll + len(ramp):i_roll + n_copy, :] = blended[len(ramp):, :]

        _, y = istft(Z, fs=sr, nperseg=NFFT, noverlap=NFFT - HOP)
        out[:, ch] = y[: audio.shape[0]]
    return out