"""SonicForge audio_features - librosa-free DSP shim (numpy/scipy/soundfile only).

Thay the toan bo phan librosa duoc dung trong app bang cac ham nhe, cung
ngu nghia, khong keo theo numba/llvmlite (~171MB) + scikit-learn (~17MB).

Cac ham duoc clone theo ngu nghia cua librosa 0.11 tai cac call-site:
  - load()                ~  librosa.load (sr=None, mono=True)
  - frames_to_time()      ~  librosa.frames_to_time
  - beat_track()          ~  librosa.beat.beat_track (onset spectral flux
                            + autocorrelation tempo + adaptive peak picking)
  - spectral_centroid()   ~  librosa.feature.spectral_centroid
  - rms()                 ~  librosa.feature.rms
  - zero_crossing_rate()  ~  librosa.feature.zero_crossing_rate
  - time_stretch()        ~  librosa.effects.time_stretch (phase vocoder)
  - pitch_shift()         ~  librosa.effects.pitch_shift
  - chroma_stft()         ~  librosa.feature.chroma_cqt (xap xi STFT-based,
                            dung cho fingerprint/similarity, KHONG dung cho
                            hien thi pitch chinh xac)

Chi phu thuoc: numpy, scipy.signal, soundfile - tat ca da co trong bundle.
"""
import numpy as np
import soundfile as sf
from scipy import signal as _signal

__all__ = [
    "load", "frames_to_time", "beat_track",
    "spectral_centroid", "rms", "zero_crossing_rate",
    "time_stretch", "pitch_shift", "chroma_stft",
]

# Mat dinh giong librosa (hop_length=512, n_fft=2048, win_length=2048)
HOP_LENGTH = 512
N_FFT = 2048
WIN_LENGTH = 2048


# ── Load / time ──────────────────────────────────────────────────────────────
def load(path, sr=None, mono=True, offset=0.0, duration=None):
    """Doc audio giong librosa.load: float32 [-1,1], mono = mean cac channel.

    sr=None -> giu nguyen sample rate goc (tat ca call-site deu dung sr=None).
    Neu truyen sr -> resample bang scipy.signal.resample_poly.
    """
    if offset or duration:
        info = sf.info(path)
        start = int(offset * info.samplerate) if offset else 0
        n_frames = int(duration * info.samplerate) if duration else -1
        data, file_sr = sf.read(path, dtype="float32", start=start, frames=n_frames)
    else:
        data, file_sr = sf.read(path, dtype="float32")

    if data.ndim > 1:
        if mono:
            data = data.mean(axis=1)
        else:
            data = data.T  # (channels, samples) giong librosa

    if sr is not None and sr != file_sr:
        from fractions import Fraction
        ratio = Fraction(int(sr), int(file_sr))
        up, down = ratio.numerator, ratio.denominator
        data = _signal.resample_poly(data, up, down).astype(np.float32)
        file_sr = sr

    return data, file_sr


def frames_to_time(frames, sr=22050, hop_length=HOP_LENGTH, n_fft=None):
    """Chuyen frame index sang giay: frames * hop_length / sr (giong librosa)."""
    return np.asanyarray(frames) * float(hop_length) / float(sr)


# ── Framing / STFT (center=True, reflect pad, giong librosa) ────────────────
def _frame(y, frame_length=WIN_LENGTH, hop_length=HOP_LENGTH):
    """Cua so hoa tin hieu voi center padding reflect (nhu librosa center=True)."""
    pad = frame_length // 2
    yp = np.pad(np.asarray(y, dtype=np.float64), pad, mode="reflect")
    n_frames = 1 + (len(yp) - frame_length) // hop_length
    if n_frames < 1:
        n_frames = 1
    idx = np.arange(frame_length)[:, None] + hop_length * np.arange(n_frames)[None, :]
    return yp[idx]


def _stft(y, n_fft=N_FFT, hop_length=HOP_LENGTH, win_length=WIN_LENGTH):
    """STFT mot phia (rfft) voi cua so hann periodic, reflect pad."""
    y = np.asarray(y, dtype=np.float64)
    window = _signal.get_window("hann", win_length, fftbins=False)
    f, _t, Zxx = _signal.stft(
        y, fs=1.0, window=window, nperseg=win_length,
        noverlap=win_length - hop_length, nfft=n_fft,
        boundary="even", padded=True,
    )
    return Zxx


def _istft(Zxx, n_fft=N_FFT, hop_length=HOP_LENGTH, win_length=WIN_LENGTH,
           length=None):
    """ISTFT nguoc voi _stft (boi so chinh xac, rate=1 -> ~identity).

    boundary=True: cat padding (nperseg//2 moi ben) nhu librosa center=True.
    """
    window = _signal.get_window("hann", win_length, fftbins=False)
    _t, y = _signal.istft(
        Zxx, fs=1.0, window=window, nperseg=win_length,
        noverlap=win_length - hop_length, nfft=n_fft,
        input_onesided=True, boundary=True,
    )
    if length is not None and len(y) > length:
        y = y[:length]
    return y


# ── Features ─────────────────────────────────────────────────────────────────
def spectral_centroid(y=None, sr=22050, n_fft=N_FFT, hop_length=HOP_LENGTH,
                      S=None):
    """Trong tam pho (brightness) - (1, n_frames) Hz, dung power spectrogram."""
    if S is None:
        S = np.abs(_stft(y, n_fft, hop_length)) ** 2
    freqs = np.fft.rfftfreq(n_fft, d=1.0 / sr)
    mag = np.abs(S)
    denom = mag.sum(axis=0)
    cent = np.divide(
        np.sum(freqs[:, None] * mag, axis=0), denom,
        out=np.zeros_like(denom), where=denom > 1e-10,
    )
    return cent[None, :]


def rms(y=None, frame_length=WIN_LENGTH, hop_length=HOP_LENGTH, S=None):
    """RMS nang luong moi frame - (1, n_frames)."""
    if S is not None:
        frames = S  # caller truyen power spectrogram
    else:
        frames = _frame(y, frame_length, hop_length)
    return np.sqrt(np.mean(frames ** 2, axis=0))[None, :]


def zero_crossing_rate(y, frame_length=WIN_LENGTH, hop_length=HOP_LENGTH):
    """Ti le zero-crossing moi frame - (1, n_frames)."""
    frames = _frame(y, frame_length, hop_length)
    signs = np.signbit(frames).astype(np.int8)
    zcr = np.mean(np.abs(np.diff(signs, axis=0)), axis=0)
    return zcr[None, :]


def chroma_stft(y=None, sr=22050, n_fft=4096, hop_length=HOP_LENGTH):
    """Chroma 12 pitch class (xap xi chroma_cqt bang STFT bin folding).

    Tra ve (12, n_frames), chuan hoa L2 tung frame - tuong thich voi
    cosine_similarity trong ai_dsp_engine.
    """
    mag = np.abs(_stft(y, n_fft, hop_length))
    freqs = np.fft.rfftfreq(n_fft, d=1.0 / sr)
    # Chi giu bin <= 5kHz (tranh nhieu alias o high freq)
    keep = freqs <= 5000.0
    freqs = freqs[keep]
    mag = mag[keep]
    # note number -> pitch class
    note = 12.0 * np.log2(np.maximum(freqs, 1e-6) / 440.0) + 69.0
    pc = np.mod(np.round(note).astype(int), 12)
    chroma = np.zeros((12, mag.shape[1]), dtype=np.float64)
    np.add.at(chroma, pc, mag)
    # L2 normalize tung frame (giong librosa)
    norms = np.linalg.norm(chroma, axis=0)
    chroma = np.divide(chroma, norms, out=np.zeros_like(chroma), where=norms > 1e-10)
    return chroma


# ── Onset / tempo / beat (thay librosa.beat) ─────────────────────────────────
def _onset_strength(y, sr, hop_length=HOP_LENGTH, n_fft=N_FFT):
    """Onset envelope: spectral flux (log-magnitude diff, chi chieu duong)."""
    mag = np.abs(_stft(y, n_fft, hop_length))
    logmag = np.log1p(1000.0 * mag)
    flux = np.diff(logmag, axis=1)
    onset = np.maximum(flux, 0.0).sum(axis=0)
    if onset.size == 0:
        return onset
    # Tru moving-average ~1s de loai trend (giong librosa detrend)
    win = max(1, int(round(1.0 * sr / hop_length)))
    if len(onset) >= win:
        kernel = np.ones(win) / win
        ma = np.convolve(onset, kernel, mode="same")
        onset = np.maximum(onset - ma, 0.0)
    return onset


def _autocorr(x):
    """Autocorrelation chuan hoa (FFT, O(n log n)), r[0]=1."""
    n = len(x)
    if n < 2:
        return np.ones(n)
    x = x - x.mean()
    nfft = 2 ** int(np.ceil(np.log2(2 * n)))
    X = np.fft.rfft(x, nfft)
    r = np.fft.irfft(X * np.conj(X), nfft)[:n]
    denom = np.maximum(n - np.arange(n), 1)
    r = r / denom
    r0 = r[0] if r[0] != 0 else 1.0
    return r / r0


def _estimate_tempo(onset, sr, hop_length=HOP_LENGTH, bpm_range=(30.0, 300.0),
                    start_bpm=120.0):
    """Uoc luong BPM bang autocorrelation cua onset envelope.

    Co them prior Gaussian quanh start_bpm (mac dinh 120, nhu librosa) de
    chon dung octave (tranh roi vao nua/double tempo khi autocorrelation
    bi mo ho giua cac harmonic).
    """
    if len(onset) < 4:
        return float(start_bpm)
    min_lag = int(np.ceil(60.0 * sr / (bpm_range[1] * hop_length)))
    max_lag = int(np.floor(60.0 * sr / (bpm_range[0] * hop_length)))
    if max_lag <= min_lag or max_lag >= len(onset):
        return float(start_bpm)
    ac = _autocorr(onset)
    lags = np.arange(min_lag, max_lag + 1)
    tempi = 60.0 * sr / (hop_length * lags)
    # prior rong ~0.7 octave quanh start_bpm (log2 scale)
    prior = np.exp(-0.5 * ((np.log2(np.maximum(tempi, 1.0)) - np.log2(start_bpm)) / 0.7) ** 2)
    seg = ac[lags] * prior
    best = lags[int(np.argmax(seg))]
    tempo = 60.0 * sr / (hop_length * best)
    # Neu tempo > 200 -> kha nang la harmonic (half-time) -> chia doi
    if tempo > 200.0 and best * 2 <= max_lag:
        tempo = 60.0 * sr / (hop_length * best * 2)
    return float(tempo)


def _localmax(x):
    """Boolean mask cac diem cuc dai dia phuong (lon hon 2 lan can)."""
    n = len(x)
    if n < 3:
        return np.zeros(n, dtype=bool)
    out = np.zeros(n, dtype=bool)
    out[1:-1] = (x[1:-1] > x[:-2]) & (x[1:-1] >= x[2:])
    return out


def _beat_frames(onset, sr, hop_length=HOP_LENGTH, tempo=120.0):
    """Chon beat frames bang peak-picking thich nghi + rang buoc tempo grid."""
    n = len(onset)
    if n == 0:
        return np.array([], dtype=int)
    period = 60.0 * sr / (hop_length * max(tempo, 1.0))  # frames/beat
    win = max(1, int(round(period)))
    kernel = np.ones(win) / win
    ma = np.convolve(onset, kernel, mode="same")
    thresh = 1.25 * ma + 1e-9

    cand = np.where(_localmax(onset) & (onset >= thresh))[0]
    if cand.size == 0:
        cand = np.where(_localmax(onset))[0]
    if cand.size == 0:
        cand = np.arange(0, n, max(1, int(round(period))))

    beats = [int(cand[0])]
    while True:
        expected = beats[-1] + period
        if expected >= n:
            break
        lo, hi = expected - 0.45 * period, expected + 0.45 * period
        in_win = cand[(cand >= lo) & (cand <= hi)]
        if in_win.size == 0:
            nxt = int(round(expected))
            if nxt >= n:
                break
            beats.append(nxt)
        else:
            beats.append(int(in_win[np.argmin(np.abs(in_win - expected))]))
        # Chong beat kep (khoang cach < 0.5 period)
        if len(beats) >= 2 and beats[-1] - beats[-2] < 0.5 * period:
            beats.pop()
            continue
        if len(beats) > 2000:
            break
    return np.array(beats, dtype=int)


def beat_track(y=None, sr=22050, hop_length=HOP_LENGTH, start_bpm=120.0,
               tightness=100):
    """Beat tracking don gian: (tempo: float, beat_frames: np.ndarray int).

    Tempo bang autocorrelation onset; beats bang peak-picking thich nghi.
    Tuong thich kieu tra ve cua librosa.beat.beat_track tai call-site
    (analyzer xu ly ca scalar lan ndarray).
    """
    onset = _onset_strength(y, sr, hop_length)
    tempo = _estimate_tempo(onset, sr, hop_length, start_bpm=start_bpm)
    beats = _beat_frames(onset, sr, hop_length, tempo)
    return tempo, beats


# ── Effects (thay librosa.effects) ───────────────────────────────────────────
def _phase_vocoder(D, rate, hop_length=HOP_LENGTH):
    """Phase vocoder time-stretch kinh dien (DAFX/Puckette).

    D: STFT (freq_bins, n_frames). rate > 1 -> nhanh hon (ngan hon).
    Tra ve STFT da stretch voi so frame ~ n_frames / rate.
    """
    n_freq, n_frames = D.shape
    if rate <= 0:
        raise ValueError("rate phai > 0")
    if rate == 1.0:
        return D
    time_steps = np.arange(0, n_frames, rate, dtype=float)
    n_out = len(time_steps)
    if n_out == 0:
        return D[:, :0]
    out = np.zeros((n_freq, n_out), dtype=np.complex128)
    # Phase advance moi hop cua tung bin tan so
    phase_adv = np.linspace(0.0, np.pi * hop_length, n_freq)
    mag = np.abs(D)
    phase_acc = np.angle(D[:, 0])
    for t, step in enumerate(time_steps):
        idx = int(step)
        if idx >= n_frames:
            break
        if idx + 1 >= n_frames:
            out[:, t] = mag[:, idx] * np.exp(1j * phase_acc)
            break
        # Phase difference that giua 2 frame lien tiep (true frequency)
        dphase = np.angle(D[:, idx + 1]) - np.angle(D[:, idx]) - phase_adv
        dphase -= 2.0 * np.pi * np.round(dphase / (2.0 * np.pi))
        phase_acc = phase_acc + phase_adv + dphase
        out[:, t] = 0.5 * (mag[:, idx] + mag[:, idx + 1]) * np.exp(1j * phase_acc)
    return out


def time_stretch(y, rate, **kwargs):
    """Time stretch giu nguyen pitch. rate > 1 -> nhanh/ngan hon."""
    if rate <= 0:
        raise ValueError("rate phai > 0")
    if rate == 1.0:
        return np.asarray(y, dtype=np.float32)
    y = np.asarray(y, dtype=np.float64)
    D = _stft(y)
    D_stretch = _phase_vocoder(D, rate)
    y_out = _istft(D_stretch)
    # Cat ve dung do dai ky vong: len(y) / rate
    target = int(round(len(y) / rate))
    if len(y_out) > target:
        y_out = y_out[:target]
    return y_out.astype(np.float32)


def pitch_shift(y, sr=22050, n_steps=1, **kwargs):
    """Dich pitch n semitone (positive = cao hon), giu nguyen duration.

    Co che (giong librosa): time_stretch voi rate=2^(-n/12) roi resample
    nguoc lai ve dung do dai goc -> pitch doi, duration giu nguyen.
    """
    if n_steps == 0:
        return np.asarray(y, dtype=np.float32)
    rate = 2.0 ** (-float(n_steps) / 12.0)
    y_shift = time_stretch(y, rate)
    # Resample (FFT) ve dung do dai goc: factor = rate
    target = int(round(len(y_shift) * rate))
    if target != len(y_shift) and target > 0:
        y_shift = _signal.resample(y_shift, target)
    return np.asarray(y_shift, dtype=np.float32)