"""
audio_lab.py —— 第 24 讲的共享模块：声音怎样变成模型能处理的东西
=====================================================================

【这个文件要回答的问题】
    声音在计算机里是一串数字（每秒几千到几万个"气压采样"）。模型要理解或生成它，有三条路：
        1. 直接处理波形（WaveNet）——序列太长：1 秒的 16 kHz 语音就是 16000 个数；
        2. 变成**频谱图**（每 10 毫秒一列"各个频率有多响"）——Whisper 的输入是对数梅尔频谱，像一张图；
        3. 变成**离散的 token**（神经音频编解码器 + 残差向量量化）——AudioLM、MusicGen 在 token 上训练语言模型。
    本模块从零实现每一步要用的工具，并且都和 numpy / PyTorch 的标准实现逐位核对。

【本文件提供】
    信号与 PCM：tone、quantize_pcm、snr_db
    傅里叶变换：dft_naive（O(N²) 的矩阵乘法）、fft_radix2（O(N log N) 的递归）、
        stft（分帧 + 加窗 + FFT，和 torch.stft 的约定相同）、hz_to_mel / mel_to_hz、mel_filterbank、log_mel
    电话拨号音（DTMF）：DTMF_ROWS / DTMF_COLS / DTMF_KEYS、dtmf_utterance（随机的按键序列 → 波形 + 标签）、
        goertzel_power（经典的单频能量检测）、classic_dtmf_decode（不用学习的传统解码器）
    CTC：ctc_forward（前向算法，对数空间）、ctc_bruteforce（枚举全部对齐，只用于核对）、ctc_greedy、edit_distance、
        CTCModel（卷积 + 双向 GRU 的小声学模型）
    音频 token：synth_music（合成的"旋律"）、frame_notes、kmeans、rvq_fit / rvq_encode / rvq_decode（残差向量量化）

【边界】
    - 所有信号都是合成的（正弦波、拨号音、简单的"乐器"），不下载任何音频；结论是机制上的。
    - 拨号音（DTMF）是真实的电话标准（ITU-T Q.23）：每个按键 = 一个低频 + 一个高频同时响。我们把它当成一门
      只有 12 个"音素"的"语言"，用来演示语音识别的完整流程（特征 → 声学模型 → CTC 解码）。
"""

import math
import random

import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F


# =====================================================================
# 1. 信号与 PCM
# =====================================================================

def tone(freqs, dur, sr, amps=None, phase=0.0):
    """
    若干个正弦波相加：Σ a_i sin(2π f_i t + φ)，t = 0, 1/sr, 2/sr, …（共 round(dur·sr) 个采样）。
    【返回】float64 的 numpy 数组
    """
    t = np.arange(int(round(dur * sr))) / sr
    amps = amps or [1.0] * len(freqs)
    return sum(a * np.sin(2 * np.pi * f * t + phase) for f, a in zip(freqs, amps))


def quantize_pcm(x, bits):
    """
    脉冲编码调制（PCM）：把 [−1, 1] 的实数量化成 bits 位的有符号整数，再换回实数（看量化误差用）。
    【思路】bits 位一共 2^bits 个台阶，台阶的间距是 2 / 2^bits；每个采样四舍五入到最近的台阶。
    【返回】(整数数组, 换回实数后的数组)
    【边界】超出 [−1, 1] 的部分被截断（"削波"）。
    """
    q = 2 ** (bits - 1)
    ints = np.clip(np.round(x * q), -q, q - 1).astype(np.int64)
    return ints, ints / q


def snr_db(clean, noisy):
    """信噪比（分贝）= 10·log₁₀(信号能量 / 误差能量)。"""
    err = np.sum((clean - noisy) ** 2)
    return float("inf") if err == 0 else 10 * math.log10(np.sum(clean ** 2) / err)


# =====================================================================
# 2. 傅里叶变换、STFT、梅尔频谱
# =====================================================================

def dft_naive(x):
    """
    离散傅里叶变换的定义：X_k = Σ_n x_n · e^{−2πi·kn/N}，k = 0 … N−1。写成一个 N × N 的矩阵乘向量，O(N²)。
    【返回】复数 numpy 数组（和 np.fft.fft 的约定相同）
    """
    N = len(x)
    n = np.arange(N)
    W = np.exp(-2j * np.pi * np.outer(n, n) / N)
    return W @ np.asarray(x, dtype=np.complex128)


def fft_radix2(x):
    """
    Cooley–Tukey 快速傅里叶变换（基 2，递归）：把长度 N 的 DFT 拆成偶数项和奇数项两个长度 N/2 的 DFT，
        X_k = E_k + e^{−2πik/N} O_k,  X_{k+N/2} = E_k − e^{−2πik/N} O_k
    每一层 O(N)，一共 log₂N 层，所以 O(N log N)。
    【边界】N 必须是 2 的幂（不是的话请先补零）。
    """
    x = np.asarray(x, dtype=np.complex128)
    N = len(x)
    if N == 1:
        return x
    if N % 2:
        raise ValueError("fft_radix2 要求长度是 2 的幂")
    even, odd = fft_radix2(x[0::2]), fft_radix2(x[1::2])
    tw = np.exp(-2j * np.pi * np.arange(N // 2) / N) * odd
    return np.concatenate([even + tw, even - tw])


def hann(n):
    """周期 Hann 窗（torch.hann_window 的默认）：w_k = 0.5 − 0.5·cos(2πk/n)。让每一帧的两端平滑地降到 0，减少频谱泄漏。"""
    return 0.5 - 0.5 * np.cos(2 * np.pi * np.arange(n) / n)


def stft(x, n_fft, hop, center=True):
    """
    短时傅里叶变换：把信号切成重叠的帧（每帧 n_fft 个采样、相邻两帧相隔 hop 个采样），每帧乘上 Hann 窗再做 FFT。

    【参数】center=True：先在两端各反射填充 n_fft/2 个采样，让第 t 帧的中心落在第 t·hop 个采样上（和 torch.stft 相同）
    【返回】复数数组 (n_fft/2 + 1, 帧数)：只保留非负频率（实信号的频谱是共轭对称的）
    【边界】帧数 = 1 + (填充后的长度 − n_fft) // hop。
    """
    x = np.asarray(x, dtype=np.float64)
    if center:
        x = np.pad(x, n_fft // 2, mode="reflect")
    n_frames = 1 + (len(x) - n_fft) // hop
    idx = np.arange(n_fft)[None, :] + hop * np.arange(n_frames)[:, None]
    frames = x[idx] * hann(n_fft)[None, :]
    return np.fft.rfft(frames, axis=1).T


def hz_to_mel(f):
    """HTK 的梅尔刻度：m = 2595·log₁₀(1 + f/700)。低频近似线性、高频近似对数——人耳对低频的音高差更敏感。"""
    return 2595.0 * np.log10(1.0 + np.asarray(f, dtype=np.float64) / 700.0)


def mel_to_hz(m):
    return 700.0 * (10 ** (np.asarray(m, dtype=np.float64) / 2595.0) - 1.0)


def mel_filterbank(n_mels, n_fft, sr, fmin=0.0, fmax=None):
    """
    三角形梅尔滤波器组：在梅尔刻度上等间距地放 n_mels + 2 个点，相邻三个点构成一个三角形滤波器。
    【返回】(n_mels, n_fft/2 + 1) 的矩阵；梅尔能量 = 矩阵 @ 功率谱
    【边界】这里用 HTK 刻度、不做面积归一化（Whisper 用 librosa 的 Slaney 刻度和归一化，形状相同、数值略不同）。
        n_mels 太大、n_fft 太小时，低频的三角形可能窄到一个频点都不包含（全 0 的行）。
    """
    fmax = fmax or sr / 2
    freqs = np.arange(n_fft // 2 + 1) * sr / n_fft
    pts = mel_to_hz(np.linspace(hz_to_mel(fmin), hz_to_mel(fmax), n_mels + 2))
    fb = np.zeros((n_mels, len(freqs)))
    for m in range(n_mels):
        lo, mid, hi = pts[m], pts[m + 1], pts[m + 2]
        up = (freqs - lo) / (mid - lo)
        down = (hi - freqs) / (hi - mid)
        fb[m] = np.clip(np.minimum(up, down), 0.0, None)
    return fb


def log_mel(x, sr, n_fft, hop, n_mels):
    """
    对数梅尔频谱：|STFT|² → 梅尔滤波器组 → log₁₀，再按 Whisper 的做法限制动态范围并缩放：
        L = log₁₀(max(mel, 10⁻¹⁰));  L = max(L, max(L) − 8);  L = (L + 4) / 4
    【返回】(n_mels, 帧数) 的 float32 数组
    """
    power = np.abs(stft(x, n_fft, hop)) ** 2
    mel = mel_filterbank(n_mels, n_fft, sr) @ power
    L = np.log10(np.maximum(mel, 1e-10))
    L = np.maximum(L, L.max() - 8.0)
    return ((L + 4.0) / 4.0).astype(np.float32)


# =====================================================================
# 3. 电话拨号音（DTMF）
# =====================================================================

DTMF_ROWS = (697, 770, 852, 941)          # 低频组（Hz）
DTMF_COLS = (1209, 1336, 1477, 1633)      # 高频组（Hz）
DTMF_KEYS = ("123A", "456B", "789C", "*0#D")
SYMBOLS = list("0123456789*#")            # 我们只用电话上常见的 12 个键
SYM_FREQ = {DTMF_KEYS[r][c]: (DTMF_ROWS[r], DTMF_COLS[c]) for r in range(4) for c in range(4)}
SR_PHONE = 8000                            # 电话的采样率
N_FFT, HOP, N_MELS = 200, 80, 40           # 25 ms 的窗、10 ms 的步长、40 个梅尔通道（Whisper 是 16 kHz 下的 400 / 160 / 80）


def dtmf_utterance(rng, n_keys=None, snr=None, sr=SR_PHONE):
    """
    一段随机的"拨号录音"：n_keys 个按键，每个按键响 40–120 毫秒，按键之间停 30–120 毫秒；
    每个按键的两个频率各有 ±1.5% 的偏差（标准允许的范围）、音量随机；最后加白噪声到指定的信噪比。

    【参数】n_keys：默认 3–8 随机；snr：信噪比（分贝），默认在 0–25 dB 之间均匀随机
    【返回】(波形 float64 数组, 按键字符串, 信噪比)
    【边界】同一个键可以连续出现（例如 "55"）：两次之间一定有停顿——CTC 正是靠这个停顿（空白）区分 "5" 和 "55"。
    """
    n_keys = n_keys or rng.randint(3, 8)
    snr = rng.uniform(0, 25) if snr is None else snr
    keys = "".join(rng.choice(SYMBOLS) for _ in range(n_keys))
    parts = [np.zeros(int(sr * rng.uniform(0.03, 0.12)))]
    for k in keys:
        lo, hi = SYM_FREQ[k]
        jit = [1 + rng.uniform(-0.015, 0.015) for _ in range(2)]
        amp = rng.uniform(0.2, 0.5)
        parts.append(tone([lo * jit[0], hi * jit[1]], rng.uniform(0.04, 0.12), sr, [amp, amp * rng.uniform(0.7, 1.0)],
                          phase=rng.uniform(0, 2 * np.pi)))
        parts.append(np.zeros(int(sr * rng.uniform(0.03, 0.12))))
    x = np.concatenate(parts)
    p_sig = np.mean(x ** 2)
    noise = np.array([rng.gauss(0.0, 1.0) for _ in range(len(x))]) * math.sqrt(p_sig / 10 ** (snr / 10))
    return x + noise, keys, snr


def goertzel_power(frame, f, sr):
    """
    Goertzel 算法：只算一个频率 f 上的能量（相当于 DFT 的一个分量），O(N)。电话交换机几十年来就用它识别按键。
        s_n = x_n + 2cos(ω)·s_{n−1} − s_{n−2},   功率 = s_{N−1}² + s_{N−2}² − 2cos(ω)·s_{N−1}·s_{N−2}
    """
    w = 2 * math.pi * f / sr
    c = 2 * math.cos(w)
    s1 = s2 = 0.0
    for v in frame:
        s1, s2 = v + c * s1 - s2, s1
    return s1 * s1 + s2 * s2 - c * s1 * s2


def classic_dtmf_decode(x, sr=SR_PHONE, frame_len=N_FFT, hop=HOP, rel=0.25, min_run=2, max_gap=2):
    """
    不用学习的传统解码器：每一帧用 Goertzel 算 8 个标准频率的能量；低频组、高频组各取最强的那个；
    如果两者都足够突出（各占本帧 8 个频率总能量的比例超过 rel），就认为这一帧是那个键，否则是"静音"。
    然后按真实解码器的做法"去抖动"：同一个键中间断了不超过 max_gap 帧，算作没断；持续不到 min_run 帧的，当作噪声丢掉。
    【返回】识别出的按键字符串
    【边界】rel、min_run、max_gap 是在 80 段另外生成的录音上网格搜索选的（03 号脚本的测试录音没有参与）：噪声大时，静音帧可能被误判成按键，短促的按键可能被漏掉。
    """
    labels = []
    for start in range(0, len(x) - frame_len + 1, hop):
        fr = x[start:start + frame_len]
        lo = [goertzel_power(fr, f, sr) for f in DTMF_ROWS]
        hi = [goertzel_power(fr, f, sr) for f in DTMF_COLS]
        tot = sum(lo) + sum(hi) + 1e-12
        r, c = int(np.argmax(lo)), int(np.argmax(hi))
        k = DTMF_KEYS[r][c]
        labels.append(k if (lo[r] / tot > rel and hi[c] / tot > rel and k in SYMBOLS) else "-")
    # 把帧标签变成 [键, 起始帧, 结束帧] 的段；同一个键之间只隔了很短的静音就连起来
    runs = []
    for t, lab in enumerate(labels):
        if lab == "-":
            continue
        if runs and runs[-1][0] == lab and t - runs[-1][2] <= max_gap + 1:
            runs[-1][2] = t
        else:
            runs.append([lab, t, t])
    return "".join(r[0] for r in runs if r[2] - r[1] + 1 >= min_run)


# =====================================================================
# 4. CTC
# =====================================================================

BLANK = 0
CTC_CLASSES = ["_"] + SYMBOLS               # 第 0 类是空白


def ctc_forward(logp, labels, blank=BLANK):
    """
    CTC 的前向算法：log p(labels | x) = log Σ_{所有能塌缩成 labels 的对齐} Π_t p(对齐_t)。

    【思路】在标签之间和两端插入空白：l' = (_, l₁, _, l₂, …, _)，长度 2U + 1。
        α_t(s) = 走到第 t 帧、停在 l' 的第 s 个位置的所有路径的概率之和。转移只能是：留在原地、前进 1、
        或者跳过一个空白前进 2（只有当 l'_s 不是空白、并且和 l'_{s−2} 不同时才允许——相同的标签之间必须有空白）。
        最后的答案 = α_{T−1}(最后一个标签) + α_{T−1}(最后一个空白)。在对数空间里做，避免下溢。
    【参数】logp：(T, C) 的对数概率张量（每帧一个分布）；labels：整数列表（不含空白）
    【返回】标量张量 log p(labels | x)（可以反向传播）
    【边界】T < U + (相邻重复标签的个数) 时不可能对齐，返回 −inf。
    """
    ext = [blank]
    for l in labels:
        ext += [l, blank]
    S, T = len(ext), logp.shape[0]
    neg = torch.tensor(float("-inf"), dtype=logp.dtype)
    alpha = [neg] * S
    alpha[0] = logp[0, ext[0]]
    if S > 1:
        alpha[1] = logp[0, ext[1]]
    for t in range(1, T):
        new = []
        for s in range(S):
            cands = [alpha[s]]
            if s >= 1:
                cands.append(alpha[s - 1])
            if s >= 2 and ext[s] != blank and ext[s] != ext[s - 2]:
                cands.append(alpha[s - 2])
            new.append(torch.logsumexp(torch.stack(cands), 0) + logp[t, ext[s]])
        alpha = new
    return torch.logsumexp(torch.stack(alpha[-2:] if S > 1 else alpha[-1:]), 0)


def collapse(path, blank=BLANK):
    """CTC 的塌缩规则：先合并连续相同的符号，再删掉空白。例如 _ 5 5 _ _ 5 3 → 5 5 3。"""
    out, prev = [], None
    for p in path:
        if p != prev and p != blank:
            out.append(p)
        prev = p
    return out


def ctc_bruteforce(logp, labels, blank=BLANK):
    """枚举全部 C^T 条路径，加起来塌缩成 labels 的那些的概率（只用于在很小的例子上核对 ctc_forward）。"""
    import itertools
    T, C = logp.shape
    total = []
    for path in itertools.product(range(C), repeat=T):
        if collapse(path, blank) == list(labels):
            total.append(sum(logp[t, path[t]] for t in range(T)))
    return torch.logsumexp(torch.stack(total), 0) if total else torch.tensor(float("-inf"))


def ctc_greedy(logp):
    """贪心（最佳路径）解码：每帧取概率最大的类，再按 CTC 规则塌缩。"""
    return collapse(logp.argmax(-1).tolist())


def edit_distance(a, b):
    """编辑距离（插入、删除、替换各算 1）：识别错误率 = 编辑距离 / 参考的长度。"""
    prev = list(range(len(b) + 1))
    for i, x in enumerate(a, 1):
        cur = [i]
        for j, y in enumerate(b, 1):
            cur.append(min(prev[j] + 1, cur[j - 1] + 1, prev[j - 1] + (x != y)))
        prev = cur
    return prev[-1]


class CTCModel(nn.Module):
    """
    一个很小的声学模型：对数梅尔帧 → 两层一维卷积（看相邻几帧）→ 双向 GRU（看整段上下文）→ 每帧一个 13 类的分布（含空白）。
    【参数】n_mels、hidden、n_classes；stride：第二层卷积的步长（2 = 输出帧数减半，每帧 20 毫秒）
    【返回】forward(feats (B, n_mels, T)) -> 对数概率 (T', B, C)（torch 的 ctc_loss 要求的形状），T' = out_len(T)
    【GenAI 联系】Whisper 的编码器前面也是两层一维卷积（第二层步长 2），后面接 Transformer；
        NVIDIA 的 Parakeet 系列里有一个就用 CTC 损失训练。
    """

    def __init__(self, n_mels=N_MELS, hidden=64, n_classes=len(CTC_CLASSES), stride=1):
        super().__init__()
        self.stride = stride
        self.conv = nn.Sequential(nn.Conv1d(n_mels, hidden, 5, padding=2), nn.ReLU(),
                                  nn.Conv1d(hidden, hidden, 5, stride=stride, padding=2), nn.ReLU())
        self.rnn = nn.GRU(hidden, hidden, batch_first=True, bidirectional=True)
        self.out = nn.Linear(2 * hidden, n_classes)

    def out_len(self, T):
        """输入 T 帧时输出多少帧（卷积的 padding = 2、核 = 5，所以只有步长改变长度）。"""
        return (T - 1) // self.stride + 1

    def forward(self, feats):
        h = self.conv(feats).transpose(1, 2)          # (B, T', hidden)
        h, _ = self.rnn(h)
        return F.log_softmax(self.out(h), dim=-1).transpose(0, 1)


def dtmf_batch(seed, batch_size=32, snr=None):
    """
    一批训练数据：随机的拨号录音 → 对数梅尔特征（补齐到同样的帧数）+ 标签。
    【返回】(feats (B, n_mels, T), 帧数列表, 标签张量（拼在一起）, 每条的标签长度, 按键字符串列表)
    """
    rng = random.Random(seed)
    feats, keys = [], []
    for _ in range(batch_size):
        x, k, _ = dtmf_utterance(rng, snr=snr)
        feats.append(log_mel(x, SR_PHONE, N_FFT, HOP, N_MELS))
        keys.append(k)
    T = max(f.shape[1] for f in feats)
    out = np.zeros((batch_size, N_MELS, T), dtype=np.float32)
    for i, f in enumerate(feats):
        out[i, :, :f.shape[1]] = f
        out[i, :, f.shape[1]:] = f.min()
    targets = [CTC_CLASSES.index(c) for k in keys for c in k]
    return (torch.from_numpy(out), [f.shape[1] for f in feats], torch.tensor(targets),
            [len(k) for k in keys], keys)


# =====================================================================
# 5. 音频 token：合成的旋律 + 残差向量量化
# =====================================================================

SR_MUSIC = 16000
SCALE = (261.63, 293.66, 329.63, 392.00, 440.00, 523.25, 587.33, 659.25)   # C 大调五声音阶（两个八度的一部分）


def synth_music(rng, n_notes=12, sr=SR_MUSIC):
    """
    一段合成的"旋律"：n_notes 个音，每个音从 SCALE 里随机挑（不和上一个音重复）、长 80–250 毫秒，
    用基频 + 3 个泛音和"快起慢落"的包络合成；乐器的"音色"（泛音的比例）每段随机微调。
    【返回】(float64 波形, 音符列表 [(音阶里的编号, 起始采样, 结束采样), ...])
    """
    timbre = [1.0, rng.uniform(0.3, 0.7), rng.uniform(0.1, 0.4), rng.uniform(0.05, 0.2)]
    parts, notes, pos, prev = [], [], 0, None
    for _ in range(n_notes):
        k = rng.choice([i for i in range(len(SCALE)) if i != prev])
        prev = k
        n = int(rng.uniform(0.08, 0.25) * sr)
        env = np.minimum(1.0, np.arange(n) / (0.01 * sr)) * np.exp(-3.0 * np.arange(n) / n)
        parts.append(env * tone([SCALE[k] * (h + 1) for h in range(4)], n / sr, sr, timbre)[:n] * 0.3)
        notes.append((k, pos, pos + n))
        pos += n
    return np.concatenate(parts), notes


def frame_notes(notes, n_frames, hop):
    """每一帧（中心在 t·hop 个采样处）属于哪个音：返回长度 n_frames 的列表（音阶编号）。"""
    out = []
    for t in range(n_frames):
        c = t * hop
        out.append(next((k for k, a, b in notes if a <= c < b), notes[-1][0]))
    return out


def kmeans(X, K, iters=20, seed=0):
    """
    k 均值聚类（Lloyd 算法）：随机挑 K 个点当初始中心，反复"分配到最近的中心 → 中心移到所分配点的平均"。
    【参数】X：(N, D) 的 torch 张量；K：中心个数
    【返回】(K, D) 的中心
    【边界】某个中心一个点都没分到时，把它重新放到一个随机的数据点上（防止"死码字"）。
    """
    g = torch.Generator().manual_seed(seed)
    C = X[torch.randperm(len(X), generator=g)[:K]].clone()
    for _ in range(iters):
        assign = torch.cdist(X, C).argmin(1)
        for k in range(K):
            sel = X[assign == k]
            C[k] = sel.mean(0) if len(sel) else X[torch.randint(len(X), (1,), generator=g)].squeeze(0)
    return C


def rvq_fit(X, K, levels, iters=20, seed=0):
    """
    残差向量量化（RVQ，SoundStream / EnCodec 用的方法）：第 1 个码本量化原始向量，第 2 个码本量化"第 1 层没量化掉的残差"，……
        r₀ = x；  c_q = 码本 q 里离 r_{q−1} 最近的码字；  r_q = r_{q−1} − c_q
    【返回】码本列表，每个 (K, D)
    【GenAI 联系】每多一层码本，每帧多 log₂K 比特；EnCodec 用 1–32 个码本对应 1.5–24 kbps 的不同码率，
        训练时随机丢掉后面几层，于是同一个模型能在不同码率下工作。
    """
    books, R = [], X.clone()
    for q in range(levels):
        C = kmeans(R, K, iters, seed + q)
        books.append(C)
        R = R - C[torch.cdist(R, C).argmin(1)]
    return books


def rvq_encode(X, books):
    """把每个向量编码成 len(books) 个码字编号。返回 (N, levels) 的整数张量。"""
    codes, R = [], X.clone()
    for C in books:
        idx = torch.cdist(R, C).argmin(1)
        codes.append(idx)
        R = R - C[idx]
    return torch.stack(codes, 1)


def rvq_decode(codes, books, levels=None):
    """用前 levels 层码本重建：x̂ = Σ_q C_q[code_q]。"""
    levels = levels or len(books)
    return sum(books[q][codes[:, q]] for q in range(levels))