import numpy as np
from math import gcd

SR = 48000
N_FFT = 2048
HOP = 480
N_MELS = 256
FMIN = 20.0
FMAX = 24000.0
LOG_FLOOR = 1e-5
MAG_EPS = 1e-9

_MEL = None
_WIN = None


def _hz_to_mel(f):
    f = np.asarray(f, dtype=np.float64)
    f_sp, min_log_hz = 200.0 / 3.0, 1000.0
    min_log_mel = min_log_hz / f_sp
    logstep = np.log(6.4) / 27.0
    mel = np.empty_like(f)
    lin = f < min_log_hz
    mel[lin] = f[lin] / f_sp
    mel[~lin] = min_log_mel + np.log(f[~lin] / min_log_hz) / logstep
    return mel


def _mel_to_hz(m):
    m = np.asarray(m, dtype=np.float64)
    f_sp, min_log_hz = 200.0 / 3.0, 1000.0
    min_log_mel = min_log_hz / f_sp
    logstep = np.log(6.4) / 27.0
    hz = np.empty_like(m)
    lin = m < min_log_mel
    hz[lin] = m[lin] * f_sp
    hz[~lin] = min_log_hz * np.exp(logstep * (m[~lin] - min_log_mel))
    return hz


def mel_filterbank():
    global _MEL
    if _MEL is None:
        freqs = np.linspace(0.0, SR / 2.0, N_FFT // 2 + 1)
        edges = np.linspace(_hz_to_mel([FMIN])[0], _hz_to_mel([FMAX])[0], N_MELS + 2)
        hz = _mel_to_hz(edges)
        fb = np.zeros((N_MELS, freqs.shape[0]), dtype=np.float32)
        for m in range(N_MELS):
            left, center, right = hz[m], hz[m + 1], hz[m + 2]
            up = (freqs - left) / (center - left)
            down = (right - freqs) / (right - center)
            fb[m] = np.maximum(0.0, np.minimum(up, down)) * (2.0 / max(1e-8, right - left))
        _MEL = fb
    return _MEL


def hann():
    global _WIN
    if _WIN is None:
        _WIN = np.hanning(N_FFT + 1)[:-1].astype(np.float64)
    return _WIN


def log_mel(audio):
    x = np.asarray(audio, dtype=np.float64).reshape(-1)
    pad = (N_FFT - HOP) // 2
    x = np.pad(x, (pad, pad), mode="reflect")
    frames = 1 + (x.size - N_FFT) // HOP
    idx = np.arange(N_FFT)[None, :] + HOP * np.arange(frames)[:, None]
    spec = np.fft.rfft(x[idx] * hann(), n=N_FFT, axis=1).T
    mag = np.sqrt(spec.real ** 2 + spec.imag ** 2 + MAG_EPS)
    mel = np.log(np.clip(mel_filterbank() @ mag, LOG_FLOOR, None))
    return mel.T[None].astype(np.float32)


def stft(x, center=True, pad_mode="constant"):
    x = np.asarray(x, dtype=np.float64).reshape(-1)
    if center:
        p = N_FFT // 2
        x = np.pad(x, (p, p), mode=pad_mode)
    frames = 1 + (x.size - N_FFT) // HOP
    idx = np.arange(N_FFT)[None, :] + HOP * np.arange(frames)[:, None]
    return np.fft.rfft(x[idx] * hann(), n=N_FFT, axis=1).T


def istft(spec, length=None, center=True):
    spec = np.asarray(spec)
    frames = np.fft.irfft(spec.T, n=N_FFT, axis=1) * hann()
    out_len = N_FFT + HOP * (frames.shape[0] - 1)
    ola = np.zeros(out_len, dtype=np.float64)
    wsum = np.zeros(out_len, dtype=np.float64)
    w2 = hann() ** 2
    for i in range(frames.shape[0]):
        s = i * HOP
        ola[s:s + N_FFT] += frames[i]
        wsum[s:s + N_FFT] += w2
    ola /= np.where(wsum > 1e-11, wsum, 1.0)
    if center:
        ola = ola[N_FFT // 2:len(ola) - N_FFT // 2]
    if length is not None:
        if ola.size > length:
            ola = ola[:length]
        elif ola.size < length:
            ola = np.pad(ola, (0, length - ola.size))
    return ola.astype(np.float32)


def cutoff_bin(spec):
    energy = np.cumsum(np.abs(spec).sum(axis=-1))
    if energy[-1] <= 0:
        return 0
    limit = energy[-1] * 0.99
    below = np.nonzero(energy < limit)[0]
    return int(below[-1]) + 1 if below.size else 0


def postprocess(pred, src, length):
    ps = stft(pred, center=True, pad_mode="constant")
    ss = stft(src, center=True, pad_mode="constant")
    n = min(ps.shape[-1], ss.shape[-1])
    ps, ss = ps[:, :n], ss[:, :n]
    cr = cutoff_bin(ss)
    merged = ps.copy()
    merged[:cr] = ss[:cr]
    out = istft(merged, length=length, center=True)
    peak = np.max(np.abs(out))
    if peak > 0:
        out = out / peak * 0.99
    return out.astype(np.float32)


def resample(x, sr_in, sr_out=SR):
    x = np.asarray(x, dtype=np.float32).reshape(-1)
    if sr_in == sr_out:
        return x
    g = gcd(int(sr_out), int(sr_in))
    from scipy.signal import resample_poly
    return resample_poly(x, sr_out // g, sr_in // g).astype(np.float32)