"""Lab tín hiệu tổng hợp: codec RVQ, LM codec token và nhận dạng theo khung.
Chỉ cần NumPy; không tải dữ liệu hay mô hình. Đây không phải bộ tổng hợp tiếng nói.
"""
import numpy as np

SEED = 17
CODEC_STEPS = 400
LM_STEPS = 350
ASR_STEPS = 200
LR = 0.01
CODEBOOKS = 2
CODES = 8
BETA = 0.1
SAMPLE_RATE = 8000
FRAME = 64
rng = np.random.default_rng(SEED)


def mm(a, b):
    return np.einsum("ij,jk->ik", a, b, optimize=False)


def block(n, title, observation, conclusion, limit):
    print(f"\n[{n} · {title}]")
    print("Quan sát:", observation)
    print("Kết luận:", conclusion)
    print("Không chứng minh:", limit)


class AdamW:
    """m nhớ hướng trung bình; v nhớ bình phương độ lớn gradient."""
    def __init__(self, params):
        self.m = {k: np.zeros_like(v) for k, v in params.items()}
        self.v = {k: np.zeros_like(v) for k, v in params.items()}
        self.t = 0

    def step(self, params, grads):
        self.t += 1
        for k in params:
            g = grads[k]
            self.m[k] = 0.9 * self.m[k] + 0.1 * g
            self.v[k] = 0.999 * self.v[k] + 0.001 * g*g
            params[k] -= LR * ((self.m[k]/(1-0.9**self.t)) /
                              (np.sqrt(self.v[k]/(1-0.999**self.t))+1e-8)
                              + 0.001 * params[k])


def ce(logits, target):
    logits = logits - logits.max(axis=1, keepdims=True)
    prob = np.exp(logits)
    prob /= prob.sum(axis=1, keepdims=True)
    loss = -np.log(prob[np.arange(len(target)), target] + 1e-12).mean()
    grad = prob.copy()
    grad[np.arange(len(target)), target] -= 1
    return float(loss), grad / len(target), prob


# Mỗi hàng là một khung của một "âm vị" giả A/B/C, không phải âm vị ngôn ngữ.
labels = np.tile(np.arange(3), 64)
speakers = np.repeat(np.arange(2), 96)
t = np.arange(FRAME) / SAMPLE_RATE
freq = np.array([500., 1000., 1500.])
phase = rng.normal(0, 0.08, len(labels))
wave = ((0.6 + 0.15*speakers[:, None]) *
        np.sin(2*np.pi*freq[labels, None]*t + phase[:, None]) +
        rng.normal(0, 0.015, (len(labels), FRAME)))
# Chia trước khi học. Tập cuối chứa cả ba nhãn và cả hai biên độ.
# Tránh chọn chỉ nhãn A khi nhãn lặp chu kỳ 3.
valid = np.arange(len(labels)) % 7 == 0
train = ~valid
x = wave[train]
print("LAB HUẤN LUYỆN AUDIO · seed", SEED)
print("Cần biết: cross-entropy chấm xác suất nhãn đúng; gradient là đạo hàm; AdamW dùng gradient để sửa tham số.")
block(1, "Dữ liệu sóng và nhãn", f"{len(wave)} khung, {train.sum()} học / {valid.sum()} kiểm định; "
      f"{SAMPLE_RATE} mẫu/giây, {FRAME} mẫu/khung, {SAMPLE_RATE/FRAME:.1f} khung/giây.",
      "Mỗi khung có nhãn A/B/C và mã người nói giả 0/1; chia tập trước khi khớp tham số.",
      "Giọng người thật, resampling hay phân đoạn tự động; dữ liệu đã cùng tần số lấy mẫu và có ranh giới sẵn.")

# Bộ lọc tam giác có tâm cách đều trên thang Mel, áp lên năng lượng FFT cửa sổ Hann.
nfft = FRAME
hz = np.fft.rfftfreq(nfft, 1/SAMPLE_RATE)
mel = lambda f: 2595*np.log10(1+f/700)
inv_mel = lambda m: 700*(10**(m/2595)-1)
edges = inv_mel(np.linspace(mel(0), mel(SAMPLE_RATE/2), 14))
filters = np.stack([np.maximum(0, np.minimum((hz-edges[i])/(edges[i+1]-edges[i]),
                  (edges[i+2]-hz)/(edges[i+2]-edges[i+1]))) for i in range(12)])
power = abs(np.fft.rfft(wave*np.hanning(FRAME), axis=1))**2
features = np.log1p(mm(power, filters.T))
mean, std = features[train].mean(0), features[train].std(0)+1e-6
features = (features-mean)/std
block(2, "Đặc trưng Mel thu nhỏ", f"Mỗi khung -> {features.shape[1]} dải; trung bình chuẩn hóa tập học "
      f"lệch 0 tối đa {abs(features[train].mean(0)).max():.2e}.",
      "Năng lượng theo dải giúp nhận dạng phân biệt cao độ; thống kê chuẩn hóa chỉ học từ tập train.",
      "Đây là log1p năng lượng Mel của khung 8 ms, không phải cấu hình log-Mel tiêu chuẩn cho tiếng nói.")

p = {'enc': rng.normal(0, .06, (FRAME, 12)),
     'dec': rng.normal(0, .06, (12, FRAME)),
     'books': rng.normal(0, .15, (CODEBOOKS, CODES, 12))}


def codec(inp, params):
    z = mm(inp, params['enc'])
    residual = z.copy()
    q = np.zeros_like(z)
    ids, residuals = [], []
    for book in params['books']:
        residuals.append(residual.copy())
        idx = ((residual[:, None, :]-book[None, :, :])**2).sum(2).argmin(1)
        chosen = book[idx]
        q += chosen
        residual -= chosen
        ids.append(idx)
    return mm(q, params['dec']), z, q, np.stack(ids, axis=1), residuals


def codec_grads(inp, params):
    recon, z, q, ids, residuals = codec(inp, params)
    err = recon-inp
    mse = float(np.mean(err**2))
    dr = 2*err/err.size
    dq = mm(dr, params['dec'].T)
    # Argmin không khả vi: straight-through truyền đạo hàm của q sang z như đồng nhất.
    dz = dq + 2*BETA*(z-q)/z.size
    db = np.zeros_like(params['books'])
    for j in range(CODEBOOKS):
        # Nhánh codebook regression dùng residual đã dừng gradient.
        np.add.at(db[j], ids[:, j], dq + 2*(params['books'][j, ids[:, j]]-residuals[j])/z.size)
    return mse, {'enc': mm(inp.T, dz), 'dec': mm(q.T, dr), 'books': db}

before = codec_grads(x, p)[0]
opt = AdamW(p)
for _ in range(CODEC_STEPS):
    _, grads = codec_grads(x, p)
    opt.step(p, grads)
recon, z, q, ids, _ = codec(wave, p)
after = np.mean((recon[train]-x)**2)
val_mse = np.mean((recon[valid]-wave[valid])**2)
usage = [int(len(np.unique(ids[train, j]))) for j in range(CODEBOOKS)]
commitment = np.mean((z[train]-q[train])**2)
block(3, "Codec RVQ học tái tạo", f"MSE học {before:.6f} -> {after:.6f}; kiểm định {val_mse:.6f}; "
      f"mã hoạt động {usage}/{CODES}; commitment {commitment:.6f}.",
      "Encoder tuyến tính học latent, RVQ cộng mã của các bảng để sửa phần dư, decoder tuyến tính dựng lại sóng. MSE là sai số bình phương trung bình.",
      "MSE thấp không bảo đảm nghe tự nhiên. Không có discriminator, loss cảm nhận hay semantic distillation; c0 ở đây không được dạy ngữ nghĩa.")

# Sai phân hữu hạn chỉ kiểm decoder với mã đã chọn cố định: không giả vờ kiểm đạo hàm argmin.
_, grads = codec_grads(x, p)
eps = 1e-5
p['dec'][0, 0] += eps
plus = codec_grads(x, p)[0]
p['dec'][0, 0] -= 2*eps
minus = codec_grads(x, p)[0]
p['dec'][0, 0] += eps
grad_error = abs((plus-minus)/(2*eps)-grads['dec'][0, 0])
block(4, "Backward và AdamW", f"Sai lệch gradient decoder với sai phân hữu hạn {grad_error:.2e}; "
      f"AdamW đã cập nhật {opt.t} bước; {sum(v.size for v in p.values())} tham số codec.",
      "Backward truyền độ nhạy của loss; AdamW giữ m và v để đổi độ nhạy thành bước sửa. Đạo hàm decoder khớp phép đo nhiễu nhỏ.",
      "Không kiểm chứng argmin khả vi: straight-through là đạo hàm thay thế, không phải đạo hàm thật của phép chọn mã.")

# LM nhỏ: text + speaker + vị trí codebook + token trước -> tanh -> phân phối mã kế.
# Token trước chỉ là quá khứ của chuỗi đã làm phẳng (frame, codebook), không thấy mã đích.
def lm_rows(token_ids):
    rows, targets, row_train = [], [], []
    for i in range(len(wave)):
        previous = CODES  # BOS riêng cho mỗi ví dụ một khung.
        for j in range(CODEBOOKS):
            rows.append(np.r_[np.eye(3)[labels[i]], np.eye(2)[speakers[i]],
                              np.eye(CODEBOOKS)[j], np.eye(CODES+1)[previous]])
            targets.append(token_ids[i, j])
            row_train.append(train[i])
            previous = token_ids[i, j]
    return np.array(rows), np.array(targets), np.array(row_train)

lm_x, lm_y, lm_train = lm_rows(ids)
lm = {'w': rng.normal(0, .1, (lm_x.shape[1], 24)), 'b': np.zeros(24),
      'out': rng.normal(0, .1, (24, CODES)), 'bias': np.zeros(CODES)}


def lm_forward(inp):
    h = np.tanh(mm(inp, lm['w'])+lm['b'])
    return h, mm(h, lm['out'])+lm['bias']

initial_ce = ce(lm_forward(lm_x[lm_train])[1], lm_y[lm_train])[0]
lm_opt = AdamW(lm)
for _ in range(LM_STEPS):
    h, logits = lm_forward(lm_x[lm_train])
    _, dl, _ = ce(logits, lm_y[lm_train])
    dh = (mm(dl, lm['out'].T))*(1-h*h)
    lm_opt.step(lm, {'w': mm(lm_x[lm_train].T, dh), 'b': dh.sum(0),
                     'out': mm(h.T, dl), 'bias': dl.sum(0)})
final_ce = ce(lm_forward(lm_x[lm_train])[1], lm_y[lm_train])[0]
val_ce = ce(lm_forward(lm_x[~lm_train])[1], lm_y[~lm_train])[0]
# Suy luận tự hồi quy: c1 nhận c0 vừa sinh, không nhận c0 đáp án.
generated = []
for label, speaker in zip(labels[valid], speakers[valid]):
    prev, codes = CODES, []
    for j in range(CODEBOOKS):
        inp = np.r_[np.eye(3)[label], np.eye(2)[speaker], np.eye(CODEBOOKS)[j], np.eye(CODES+1)[prev]][None]
        prev = int(lm_forward(inp)[1].argmax(1)[0])
        codes.append(prev)
    generated.append(codes)
generated = np.array(generated)
gen_q = sum(p['books'][j, generated[:, j]] for j in range(CODEBOOKS))
gen_wave = mm(gen_q, p['dec'])
gen_mse = np.mean((gen_wave-wave[valid])**2)
block(5, "LM dự đoán codec token", f"Cross-entropy học {initial_ce:.4f} -> {final_ce:.4f}; "
      f"kiểm định teacher-forced {val_ce:.4f}; MSE sóng sinh tự hồi quy {gen_mse:.6f}; mã đầu {generated[0].tolist()}.",
      "Text giả A/B/C và speaker làm điều kiện; codebook j và token trước giúp đoán mã kế. Loss là trung bình -ln xác suất mã đúng. Codec được đóng băng khi học LM.",
      "LM này chỉ nhớ một token trước qua MLP, không có attention hay trí nhớ qua nhiều khung. Teacher forcing không đo lỗi tích lũy khi sinh.")

asr = {'w': rng.normal(0, .1, (12, 3)), 'b': np.zeros(3)}
asr_opt = AdamW(asr)
asr_before = ce(mm(features[train], asr['w'])+asr['b'], labels[train])[0]
for _ in range(ASR_STEPS):
    _, dl, _ = ce(mm(features[train], asr['w'])+asr['b'], labels[train])
    asr_opt.step(asr, {'w': mm(features[train].T, dl), 'b': dl.sum(0)})
asr_after = ce(mm(features[train], asr['w'])+asr['b'], labels[train])[0]
acc = np.mean((mm(features[valid], asr['w'])+asr['b']).argmax(1)==labels[valid])
# Chỉ minh họa collapse CTC, không học CTC: blank=3; gộp lặp TRƯỚC khi bỏ blank.
path = [0, 0, 3, 0, 1, 1, 3, 2]
collapsed = [v for i, v in enumerate(path) if (i == 0 or v != path[i-1]) and v != 3]
block(6, "ASR theo khung và căn chỉnh", f"CE học {asr_before:.4f} -> {asr_after:.4f}; "
      f"độ đúng khung kiểm định {acc:.1%}; collapse {path} -> {collapsed}.",
      "Đầu tuyến tính nhận Mel và đoán nhãn theo khung đã căn sẵn. Ví dụ collapse giữ hai A ngăn bởi blank; CTC thật cộng xác suất mọi đường hợp lệ.",
      "Không huấn luyện CTC, encoder ngữ cảnh hay decoder văn bản; accuracy khung không phải WER của tiếng nói thật.")

# Codec độc lập từng khung: đổi tương lai không ảnh hưởng các khung trước.
changed = wave.copy()
changed[10:] += 0.4
prefix_error = np.max(abs(codec(changed, p)[0][:10]-recon[:10]))
flat_rate = SAMPLE_RATE/FRAME*CODEBOOKS
raw_bits = SAMPLE_RATE*16
code_bits = flat_rate*np.log2(CODES)
block(7, "Hợp đồng streaming và bộ nhớ", f"Đổi tương lai: lệch tiền tố {prefix_error:.1e}; "
      f"{flat_rate:.0f} mã/giây khi làm phẳng; {code_bits:.0f} bit/giây mã so với PCM16 {raw_bits} bit/giây (tính lý thuyết).",
      "Codec không nhìn khung sau nên có thể phát sau khi đủ 64 mẫu. Nhiều codebook làm tăng số quyết định nếu nối phẳng; depth predictor là cách tổ chức khác.",
      "Không đo độ trễ phần cứng, băng thông tệp hay chất lượng; chưa tính header. Không thực thi sharding/recompute hay delay/depth predictor.")

checkpoint = {k: v.copy() for k, v in p.items()}
reload_error = np.max(abs(codec(wave, checkpoint)[0]-recon))
block(8, "Checkpoint và giới hạn phương pháp", f"Nạp lại bản sao tham số codec: lệch {reload_error:.1e}; "
      f"LM có {sum(v.size for v in lm.values())} tham số; optimizer codec giữ {sum(v.nbytes for v in opt.m.values())*2} byte cho m và v.",
      "Checkpoint suy luận cần codec, LM, bảng nhãn và cấu hình. Muốn tiếp tục học phải giữ thêm trạng thái AdamW, bước và RNG. Bản sao ở đây chỉ kiểm codec, không ghi tệp.",
      "Không thực thi pretraining quy mô lớn, SFT hội thoại, LoRA, distillation hay preference; chúng được dạy ở bài đọc. Không chứng minh chất lượng Qwen3-Omni hoặc Qwen3-TTS.")
assert np.isfinite(after) and after < before
assert final_ce < initial_ce and asr_after < asr_before
assert grad_error < 1e-6 and prefix_error < 1e-12 and reload_error < 1e-12
print("\nPASS · MSE codec và CE của LM/ASR giảm; gradient decoder, tiền tố nhân quả và nạp lại codec đạt ngưỡng đã nêu.")
