# -*- coding: utf-8 -*-
"""
audio_utils.py —— 本教程公共工具模块
统一约定：16000 Hz、单声道、float32、幅度 [-1, 1]
"""
from __future__ import annotations
from math import gcd
import numpy as np

SR = 16000                      # 教程统一采样率
FRAME_LEN = 400                 # 25 ms @16 kHz，默认帧长（点数）
HOP_LEN = 160                   # 10 ms @16 kHz，默认帧移（点数）


# ============================================================
# 0) 测试信号生成（没有音频文件时用它）
# ============================================================
def _resonator(x, f, bw, sr):
    """二阶谐振器：模拟一个共振峰（formant），f 为共振峰频率，bw 为带宽"""
    from scipy.signal import lfilter
    r = np.exp(-np.pi * bw / sr)
    th = 2 * np.pi * f / sr
    b = [1.0 - 2 * r * np.cos(th) + r * r]      # 直流增益归一化的分子
    a = [1.0, -2 * r * np.cos(th), r * r]
    return lfilter(b, a, x)


def glottal_pulse_train(n, sr=SR, f0=200.0, jitter=0.0, seed=0):
    """声门脉冲串（基频 f0），jitter>0 时加入基频微扰，更接近真实嗓音"""
    rng = np.random.default_rng(seed)
    x = np.zeros(n, dtype=np.float64)
    pos = 0.0
    while pos < n - 1:
        i = int(round(pos))
        if i < n:
            x[i] = 1.0
        T = sr / max(1e-6, f0 * (1.0 + jitter * rng.standard_normal()))
        pos += T
    return x


def voiced_source(dur=1.0, sr=SR, f0=200.0, amp=0.5,
                  formants=((700, 80), (1220, 90), (2600, 120)), seed=0):
    """
    合成“元音样”浊音信号：声门脉冲串 → 级联共振峰谐振器
    用于基音周期、共振峰、LPC 等章节的演示（近似 /a/ 的三个共振峰）
    """
    n = int(round(dur * sr))
    x = glottal_pulse_train(n, sr, f0=f0, seed=seed)
    for f, bw in formants:
        x = _resonator(x, f, bw, sr)
    m = np.max(np.abs(x))
    if m > 1e-12:
        x = x / m * amp
    return x.astype(np.float32)


def gen_test_signal(kind="sine", dur=1.0, sr=SR, f0=220.0, f1=4000.0,
                    amp=0.5, seed=0):
    """
    生成各种测试信号
    kind: sine(正弦) | sweep(线性扫频) | square(方波) | white(白噪声)
          | am(幅度调制) | voiced(元音样浊音) | silence(静音)
    """
    n = int(round(dur * sr))
    t = np.arange(n) / sr
    rng = np.random.default_rng(seed)

    if kind == "sine":
        x = amp * np.sin(2 * np.pi * f0 * t)
    elif kind == "sweep":                     # 线性调频，测语谱图最直观
        k = (f1 - f0) / max(dur, 1e-9)
        x = amp * np.sin(2 * np.pi * (f0 * t + 0.5 * k * t * t))
    elif kind == "square":
        x = amp * np.sign(np.sin(2 * np.pi * f0 * t))
    elif kind == "white":
        x = amp * rng.standard_normal(n)
    elif kind == "am":                        # 4 Hz 包络调制，相当于“音节率”
        x = amp * (0.5 + 0.5 * np.sin(2 * np.pi * 4 * t)) * np.sin(2 * np.pi * f0 * t)
    elif kind == "voiced":
        x = voiced_source(dur=dur, sr=sr, f0=f0, amp=amp, seed=seed)
    elif kind == "silence":
        x = np.zeros(n)
    else:
        raise ValueError("未知信号类型: %s" % kind)
    return np.asarray(x, dtype=np.float32)


def gen_vad_demo(dur=3.0, sr=SR, snr=15.0, seed=0):
    """
    生成“静音-语音-静音-语音-静音”的准语音信号，用于端点检测（4.1 节）
    返回 (x, sr, segments)，segments 为真实的语音区间 [(起点样点, 终点样点), ...]
    """
    segs = [(0.50, 1.30, 120.0), (1.70, 2.70, 220.0)]     # (起点s, 终点s, 基频)
    n = int(round(dur * sr))
    x = np.zeros(n, dtype=np.float32)
    for s, e, f0 in segs:
        v = voiced_source(dur=e - s, sr=sr, f0=f0, amp=0.6, seed=seed)
        x[int(s * sr):min(n, int(e * sr))] = v[: max(0,min(n, int(e * sr)) - int(s * sr))]
    x = add_noise(x, snr=snr, seed=seed)
    segments = [(int(s * sr), int(e * sr)) for s, e, _ in segs]
    return x, sr, segments


def demo_signal(dur=1.0, sr=SR, path=None, seed=0, kind="voiced"):
    """
    统一的数据入口：优先读取 path（默认 data/demo.wav），
    不存在则返回合成信号。所有示例都用它，保证“复制即可运行”。
    """
    import os
    if path is None:
        path = os.path.join("data", "demo.wav")
    if os.path.exists(path):
        return load_wav(path, sr=sr)
    return gen_test_signal(kind=kind, dur=dur, sr=sr, seed=seed), sr


# ============================================================
# 1) 读写（soundfile / librosa / 标准库 wave）
# ============================================================
def load_wav(path, sr=SR, mono=True):
    """读取音频 → (float32 单声道信号, 采样率)。自动重采样到 sr"""
    import soundfile as sf
    x, fs = sf.read(path, dtype="float32", always_2d=False)
    if x.ndim > 1 and mono:
        x = x.mean(axis=1)
    if fs != sr:
        x = resample(x, fs, sr)
        fs = sr
    return np.ascontiguousarray(x, dtype=np.float32), int(fs)


def save_wav(path, x, sr=SR, subtype="PCM_16"):
    """保存为 WAV（默认 16 bit PCM），自动削波保护"""
    import soundfile as sf
    y = np.clip(np.asarray(x, dtype=np.float32), -1.0, 1.0)
    sf.write(path, y, int(sr), subtype=subtype)


def read_wave_int16(path):
    """只用标准库读取 PCM WAV → (int16 数组, 采样率, 声道数)。第 6.1 节会用到"""
    import wave
    with wave.open(path, "rb") as w:
        n_ch = w.getnchannels()
        width = w.getsampwidth()
        fs = w.getframerate()
        n = w.getnframes()
        raw = w.readframes(n)
    if width != 2:
        raise ValueError("仅支持 16 bit PCM，当前位宽 = %d 字节" % width)
    data = np.frombuffer(raw, dtype="<i2")
    if n_ch > 1:
        data = data.reshape(-1, n_ch).mean(axis=1).astype(np.int16)
    return data, fs, n_ch


def write_wave_int16(path, data, sr=SR):
    """把 int16 数组（或 float32 数组）写成 16 bit 单声道 PCM WAV"""
    import wave
    data = np.asarray(data)
    if data.dtype != np.int16:
        data = (np.clip(data, -1.0, 1.0) * 32767).astype("<i2")
    with wave.open(path, "wb") as w:
        w.setnchannels(1)
        w.setsampwidth(2)
        w.setframerate(int(sr))
        w.writeframes(data.astype("<i2").tobytes())


def resample(x, sr_in, sr_out):
    """多相滤波重采样（比 np.interp 线性插值质量高得多）"""
    from scipy.signal import resample_poly
    x = np.asarray(x, dtype=np.float32)
    if sr_in == sr_out:
        return x
    g = gcd(int(sr_in), int(sr_out))
    return resample_poly(x, int(sr_out // g), int(sr_in // g)).astype(np.float32)


def play(x, sr=SR, blocking=True):
    """播放（需要 sounddevice）。没有声卡/库时静默跳过，不打断示例"""
    try:
        import sounddevice as sd
        sd.play(np.asarray(x, dtype=np.float32), sr, blocking=blocking)
    except Exception as e:
        print("[play 跳过] %s" % e)


def record(dur=3.0, sr=SR, channels=1):
    """录音（网页演示版）→ (float32 信号, sr)
    替代方案：使用合成语音信号模拟录音结果，便于在线演示。
    本地有麦克风时可替换回 sounddevice 版本。"""
    print("开始录音 %.1f 秒 ...（网页演示模式：使用合成语音信号）" % dur)
    y = gen_test_signal(kind="voiced", dur=dur, f0=220.0)
    print("录音结束")
    return np.asarray(y).reshape(-1), sr


# ============================================================
# 2) 幅度/dB/信噪比/归一化/预加重
# ============================================================
def to_db(x, ref=None, floor=-120.0):
    """幅度谱 → dB。ref=None 时以自身最大值为 0 dB（相对峰值 dB；计算 dBFS 请显式传 ref=1）"""
    a = np.abs(np.asarray(x, dtype=np.float64))
    if ref is None:
        ref = np.max(a) if a.size and np.max(a) > 0 else 1.0
    ref = max(float(ref), 1e-12)
    with np.errstate(divide="ignore"):
        y = 20.0 * np.log10(np.maximum(a, 1e-12) / ref)
    return np.maximum(y, floor)


def snr_db(clean, test):
    """计算 test 相对 clean 的信噪比(dB)：SNR = 10log10(Σs² / Σ(s-t)²)"""
    s = np.asarray(clean, dtype=np.float64).ravel()
    t = np.asarray(test, dtype=np.float64).ravel()
    n = min(len(s), len(t))
    s, t = s[:n], t[:n]
    noise = t - s
    return float(10 * np.log10(np.sum(s ** 2) / max(np.sum(noise ** 2), 1e-12)))


def add_noise(x, snr=10.0, kind="white", seed=0):
    """按指定 SNR(dB) 加噪声。snr 越小噪声越大"""
    rng = np.random.default_rng(seed)
    x = np.asarray(x, dtype=np.float64)
    if kind == "white":
        n = rng.standard_normal(len(x))
    elif kind == "pink":                       # 1/f 噪声：低通白噪声的近似
        from scipy.signal import lfilter
        n = lfilter([1.0], [1.0, -0.95], rng.standard_normal(len(x)))
    else:
        raise ValueError("未知噪声类型: %s" % kind)
    n = n / max(np.sqrt(np.mean(n ** 2)), 1e-12) * np.sqrt(np.mean(x ** 2) + 1e-12)
    n = n / (10 ** (snr / 20.0))
    return (x + n).astype(np.float32)


def normalize(x, peak=0.99):
    """峰值归一化到 peak"""
    x = np.asarray(x, dtype=np.float32)
    m = float(np.max(np.abs(x))) if x.size else 0.0
    if m < 1e-12:
        return x.copy()
    return (x * (peak / m)).astype(np.float32)


def preemphasis(x, a=0.97):
    """预加重 y[n] = x[n] - a·x[n-1]：提升高频，补偿声门谱 -6 dB/oct 衰减"""
    x = np.asarray(x, dtype=np.float32)
    if x.size == 0:
        return x
    y = np.empty_like(x)
    y[0] = x[0]
    y[1:] = x[1:] - a * x[:-1]
    return y


def deemphasis(x, a=0.97):
    """去加重（预加重的逆运算）：y[n] = x[n] + a·y[n-1]"""
    from scipy.signal import lfilter
    return lfilter([1.0], [1.0, -a], np.asarray(x, dtype=np.float32)).astype(np.float32)


# ============================================================
# 3) 分帧 / 加窗 / 还原
# ============================================================
def enframe(x, frame_len=FRAME_LEN, hop=HOP_LEN, pad=True, window=None,
            pad_mode="constant"):
    """分帧 → (n_frames, frame_len)。window 传入窗数组则顺带加窗"""
    x = np.asarray(x, dtype=np.float32)
    if pad:
        n_need = frame_len + hop * int(np.ceil(max(0, len(x) - frame_len) / hop))
        x = np.pad(x, (0, max(0, n_need - len(x))), mode=pad_mode)
    if len(x) < frame_len:
        x = np.pad(x, (0, frame_len - len(x)))
    n_frames = 1 + (len(x) - frame_len) // hop
    idx = np.arange(frame_len)[None, :] + hop * np.arange(n_frames)[:, None]
    frames = x[idx].astype(np.float32)
    if window is not None:
        frames = frames * np.asarray(window, dtype=np.float32)[None, :]
    return frames


def deframe(frames, hop=HOP_LEN, length=None, window=None):
    """
    重叠相加（OLA）还原时域信号。
    注意：本函数接收未加窗的帧，再施加合成窗并除以窗叠加和。
    已加窗的分析帧不要直接传入；需明确分析/合成窗和边界处理。
    """
    frames = np.asarray(frames, dtype=np.float32)
    n_frames, flen = frames.shape
    total = hop * (n_frames - 1) + flen
    out = np.zeros(total, dtype=np.float64)      # 加窗后的加权叠加
    wsum = np.zeros(total, dtype=np.float64)     # 窗的叠加和（用于补偿）
    raw = np.zeros(total, dtype=np.float64)      # 不加权的叠加（首尾边界兜底）
    cnt = np.zeros(total, dtype=np.float64)
    win = hann(flen).astype(np.float64) if window is None else np.asarray(window, np.float64)
    for i in range(n_frames):
        s, e = i * hop, i * hop + flen
        seg = frames[i].astype(np.float64)
        out[s:e] += seg * win
        wsum[s:e] += win
        raw[s:e] += seg
        cnt[s:e] += 1
    res = np.zeros(total, dtype=np.float64)
    good = wsum > 1e-6                            # 正常重叠区：除以窗叠加和（COLA 补偿）
    res[good] = out[good] / wsum[good]
    res[~good] = raw[~good] / np.maximum(cnt[~good], 1.0)   # 首尾边界（窗值为 0）：直接取平均
    return (res[:length] if length else res).astype(np.float32)


def hamming(n):
    """周期汉明窗（分析用）"""
    from scipy.signal import windows
    return windows.hamming(n, sym=False).astype(np.float32)


def hann(n):
    """周期 Hann 窗（合成/重叠相加用）"""
    from scipy.signal import windows
    return windows.hann(n, sym=False).astype(np.float32)


# ============================================================
# 4) MFCC 等分类特征（第 10/11/12 章共用，保证量纲一致）
# ============================================================
def utt_feature(feat, stat="mean_std", cmvn=False):
    """
    帧级特征 (帧数, 维数) → 整句级定长向量（供 KNN/SVM/LDA 等分类器使用）
    stat: mean(维数) | mean_std(2×维数) | mean_std_minmax(4×维数)

    ⚠️ 关键设计说明（这里曾经踩过坑）：
    本函数**默认不做** CMVN，直接统计各维的均值/标准差/极值。
    原因：CMVN 会逐维减掉整句均值——对“稳态元音/孤立字”这类信号，被减掉的
    恰恰就是共振峰位置等身份信息，mean 统计量会退化成恒 0、std 退化成恒 1，
    分类器立刻失效。
    什么时候才开 cmvn=True：说话人/信道差异远大于内容差异、且每句内容足够丰富
    （连续语音、大词汇量）时，CMVN 才利大于弊。跨样本的尺度差异请交给
    sklearn 的 StandardScaler 处理，不要在这里做。
    """
    f = np.asarray(feat, dtype=np.float32)
    if cmvn:
        f = (f - f.mean(axis=0, keepdims=True)) / (f.std(axis=0, keepdims=True) + 1e-8)
    mu = f.mean(axis=0)
    sd = f.std(axis=0)
    if stat == "mean":
        return mu.astype(np.float32)
    if stat == "mean_std":
        return np.concatenate([mu, sd]).astype(np.float32)
    return np.concatenate([mu, sd, f.min(axis=0), f.max(axis=0)]).astype(np.float32)


# ============================================================
# 5) 画图（统一风格，避免每个例子重复写 matplotlib 样板）
# ============================================================
_CJK_READY = [False]


def setup_cjk_font():
    """
    自动挑一个系统中可用的中文字体，避免图上的中文变成“豆腐块”。
    Windows 通常命中 Microsoft YaHei / SimHei，macOS 命中 PingFang SC，Linux 命中 Noto Sans CJK。
    """
    if _CJK_READY[0]:
        return
    _CJK_READY[0] = True
    try:
        import matplotlib
        from matplotlib import font_manager
        prefer = ["Microsoft YaHei", "SimHei", "PingFang SC", "Hiragino Sans GB",
                  "Noto Sans CJK SC", "Source Han Sans SC", "WenQuanYi Zen Hei",
                  "Arial Unicode MS"]
        have = {f.name for f in font_manager.fontManager.ttflist}
        for name in prefer:
            if name in have:
                matplotlib.rcParams["font.sans-serif"] = [name] + list(
                    matplotlib.rcParams.get("font.sans-serif", []))
                matplotlib.rcParams["axes.unicode_minus"] = False    # 负号别变方框
                return
    except Exception:
        pass


def plot_waveform(x, sr=SR, title="时域波形", ax=None, t0=0.0, color="#1f77b4"):
    import matplotlib.pyplot as plt
    setup_cjk_font()
    if ax is None:
        _, ax = plt.subplots(figsize=(10, 2.6))
    t = t0 + np.arange(len(x)) / float(sr)
    ax.plot(t, x, lw=0.8, color=color)
    ax.set_xlabel("时间 / s"); ax.set_ylabel("幅度"); ax.set_title(title)
    ax.grid(alpha=0.3)
    return ax


def plot_spectrogram(x, sr=SR, n_fft=512, hop=160, title="语谱图",
                     ax=None, vmin=-80, vmax=0, cmap="magma"):
    import matplotlib.pyplot as plt
    setup_cjk_font()
    if ax is None:
        _, ax = plt.subplots(figsize=(10, 3.4))
    ax.specgram(np.asarray(x, dtype=np.float32), NFFT=n_fft, Fs=sr,
                noverlap=max(0, n_fft - hop), cmap=cmap, vmin=vmin, vmax=vmax)
    ax.set_xlabel("时间 / s"); ax.set_ylabel("频率 / Hz"); ax.set_title(title)
    return ax


def plot_frames(frames, sr=SR, hop=HOP_LEN, n=5, title="分帧波形（前 %d 帧）"):
    import matplotlib.pyplot as plt
    setup_cjk_font()
    frames = np.asarray(frames, dtype=np.float32)
    n = min(n, len(frames))
    _, axes = plt.subplots(n, 1, figsize=(10, 1.2 * n), sharex=True)
    axes = np.atleast_1d(axes)
    for i, ax in enumerate(axes):
        t = (i * hop + np.arange(frames.shape[1])) / float(sr)
        ax.plot(t, frames[i], lw=0.9)
        ax.set_ylabel("帧%d" % i)
        ax.grid(alpha=0.3)
    axes[-1].set_xlabel("时间 / s")
    axes[0].set_title(title % n)
    return axes


def plot_spectrum(x, sr=SR, n_fft=512, title="幅度谱", ax=None, db=True):
    import matplotlib.pyplot as plt
    setup_cjk_font()
    if ax is None:
        _, ax = plt.subplots(figsize=(10, 3))
    X = np.fft.rfft(np.asarray(x, dtype=np.float64) * hamming(len(x)), n=n_fft)
    f = np.fft.rfftfreq(n_fft, 1.0 / sr)
    mag = np.abs(X) / max(1, len(x))
    ax.plot(f, to_db(mag) if db else mag, lw=1.0)
    ax.set_xlabel("频率 / Hz"); ax.set_ylabel("幅度 / dB" if db else "幅度")
    ax.set_title(title); ax.grid(alpha=0.3)
    return ax

def mel_filterbank(sr=SR,n_fft=512,n_mels=26,fmin=20.,fmax=None):
    """HTK-style Mel scale, unit-height triangular filters."""
    fmax=sr/2 if fmax is None else min(fmax,sr/2)
    if not 0 <= fmin < fmax: raise ValueError("Invalid Mel frequency range")
    mel=lambda f:2595*np.log10(1+np.asarray(f)/700)
    hz=lambda m:700*(10**(np.asarray(m)/2595)-1)
    edges=hz(np.linspace(mel(fmin),mel(fmax),n_mels+2))
    freq=np.fft.rfftfreq(n_fft,1/sr)
    bank=np.zeros((n_mels,len(freq)))
    for i in range(n_mels):
        left=(freq-edges[i])/(edges[i+1]-edges[i])
        right=(edges[i+2]-freq)/(edges[i+2]-edges[i+1])
        bank[i]=np.maximum(0,np.minimum(left,right))
    return bank


def mfcc(x,sr=SR,n_mfcc=13,n_fft=512,win=400,hop=160,n_mels=26,
         fmin=20.,fmax=None,cmvn=False,use_librosa=None):
    """Teaching MFCC: pre-emphasis .97, periodic Hamming, power, ln, DCT-II.

    Returns frames x dimensions. Not a bit-identical replacement for librosa/
    python_speech_features defaults. CMVN is off to preserve utterance means.
    """
    from scipy.fft import dct
    frames=enframe(preemphasis(x),win,hop,window=hamming(win))
    power=np.abs(np.fft.rfft(frames,n_fft,axis=1))**2/n_fft
    mel=power @ mel_filterbank(sr,n_fft,n_mels,fmin,fmax).T
    out=dct(np.log(np.maximum(mel,1e-12)),type=2,norm="ortho",axis=1)[:,:n_mfcc]
    if cmvn: out=(out-out.mean(0))/(out.std(0)+1e-8)
    return out.astype(np.float32)


def mfcc_delta(m,order=1,width=9):
    """Regression deltas with replicated edge frames (frames x dimensions)."""
    m=np.asarray(m,dtype=np.float32)
    if width<3 or width%2!=1: raise ValueError("width must be odd and >=3")
    n=width//2; denom=2*sum(i*i for i in range(1,n+1))
    out=[m]; cur=m
    for _ in range(order):
        padded=np.pad(cur,((n,n),(0,0)),mode="edge")
        delta=sum(i*(padded[n+i:n+i+len(cur)]-padded[n-i:n-i+len(cur)])
                  for i in range(1,n+1))/denom
        out.append(delta); cur=delta
    return np.concatenate(out,axis=1).astype(np.float32)


def lpc(frame,order=12):
    """Autocorrelation LPC, denominator a[0]=1, diagonal regularization."""
    from scipy.linalg import solve_toeplitz
    x=np.asarray(frame,dtype=float); x=x-x.mean()
    if len(x)<=order: raise ValueError("Frame must be longer than LPC order")
    r=np.correlate(x,x,"full")[len(x)-1:len(x)+order]
    r[0]+=max(r[0]*1e-8,1e-12)
    coeff=solve_toeplitz(r[:order],-r[1:order+1])
    error=max(float(r[0]+coeff@r[1:order+1]),1e-12)
    return np.r_[1.,coeff],error


def emotion_demo(seed=0):
    """Synthetic features; not real audio or evidence of emotion accuracy.

    Per-speaker shifts and independent within-class noise. Three abstract
    prosodic features plus three nuisance features. Fixed integer RNG seeds.
    """
    rng=np.random.default_rng(seed)
    centers=np.array([[150,4,.5],[215,5.5,.8],[190,5,1.],[118,2.6,.4]])
    X=[];y=[];groups=[]
    for speaker in range(12):
        offset=rng.normal(0,[18,.4,.1])
        for label,center in enumerate(centers):
            for repetition in range(5):
                prosody=center+offset+rng.normal(0,[20,.6,.13])
                X.append(np.r_[prosody,rng.normal(size=3)])
                y.append(label);groups.append(speaker)
    return np.asarray(X),np.asarray(y),np.asarray(groups)