# BÀI LAB TỔNG HỢP QWEN3-OMNI & QWEN3-TTS-12HZ: TOÀN TRÌNH 1A → 6B BẰNG NUMPY
# Chạy một lượt từ nén spectrogram AuT (1A), đan xen TM-RoPE 3D (1B), trích ngang
# Lớp 24 của Thinker (2A), hạ chiều ResizeMLP & Speaker ID (3A), đồng bộ Dual-Track
# 12.5 Hz (3B), Talker MoE sinh mã ngữ nghĩa c_0 (3C), 5-Layer Code Predictor tự
# hồi quy các mã âm học c_1..c_3 (4A), bảng nhúng gộp RVQ code_offset (6A),
# Pre-Transformer cửa sổ trượt (5A) đến giải mã sóng âm tuần hoàn SnakeBeta (6B).

import numpy as np

rng = np.random.default_rng(30)

# ── Cấu hình thu nhỏ (ánh xạ 1:1 với Qwen3-Omni-30B-A3B & Qwen3-TTS-12Hz) ────
MEL_BINS = 8              # prod: 128 mel bins
MEL_FRAMES = 16           # 16 khung Mel → sau 3× stride-2 còn 1 × 2 khung
THINKER_DIM = 8           # prod: thinker_hidden_size = 2048
TALKER_DIM = 6            # prod: talker hidden_size = 1024
HEAD_DIM = 6              # 3 cặp tần số (T, H, W) đan xen (prod: 64 cặp = [24, 20, 20])
NUM_EXPERTS = 6           # prod: 128 routed experts
TOP_K_EXPERTS = 2         # prod: num_experts_per_tok = 6
NUM_QUANTIZERS = 4        # prod: num_code_groups = 16 (1 semantic + 15 acoustic)
CODEBOOK_SIZE = 8         # prod: codebook_size = 2048
SLIDING_WINDOW = 2        # prod: sliding_window = 72 frames
PRE_UPSAMPLE = 2          # prod: upsampling_ratios = [2, 2] (4×)
DEC_UPSAMPLES = (2, 3)    # prod: upsample_rates = [8, 5, 4, 3] (480×)
TOTAL_UPSAMPLE = PRE_UPSAMPLE * DEC_UPSAMPLES[0] * DEC_UPSAMPLES[1]  # 2 × 6 = 12× (prod: 1920×)


def silu(x):
    x = np.asarray(x, dtype=float)
    return x / (1.0 + np.exp(-x))


def gelu_tanh(x):
    x = np.asarray(x, dtype=float)
    return 0.5 * x * (1.0 + np.tanh(np.sqrt(2.0 / np.pi) * (x + 0.044715 * x**3)))


def rmsnorm(x, weight, eps=1e-6):
    x = np.asarray(x, dtype=float)
    ms = np.mean(x * x, axis=-1, keepdims=True)
    return (x / np.sqrt(ms + eps)) * weight


def softmax(x, axis=-1):
    x = np.asarray(x, dtype=float)
    s = x - np.max(x, axis=axis, keepdims=True)
    e = np.exp(s)
    return e / np.sum(e, axis=axis, keepdims=True)


def fmt_vec(v, nd=3):
    arr = np.asarray(v, dtype=float).ravel()
    return "[" + ", ".join(f"{x:.{nd}f}" for x in arr) + "]"


def say(code, title, lines):
    print(f"\n[{code}] {title}")
    for ln in lines:
        print("  ", ln)


# ── [1A] AuT Audio Encoder: 3× Conv2d Stride-2 + Windowed Attention ──────────
def conv2d_stride2_mean(mat):
    # Tích chập 2D bước nhảy 2 (kernel 2×2 trung bình có trọng số cố định)
    h, w = mat.shape
    out = np.zeros((h // 2, w // 2), dtype=float)
    kernel = np.array([[0.35, 0.25], [0.20, 0.20]], dtype=float)
    for i in range(h // 2):
        for j in range(w // 2):
            patch = mat[2 * i : 2 * i + 2, 2 * j : 2 * j + 2]
            out[i, j] = float(np.sum(patch * kernel))
    return out


MEL_INPUT = rng.normal(0.0, 0.8, size=(MEL_BINS, MEL_FRAMES))
W_AUT_PROJ = rng.normal(0.0, 0.35, size=(1, THINKER_DIM))
W_AUT_Q = rng.normal(0.0, 0.35, size=(THINKER_DIM, THINKER_DIM))
W_AUT_K = rng.normal(0.0, 0.35, size=(THINKER_DIM, THINKER_DIM))
W_AUT_V = rng.normal(0.0, 0.35, size=(THINKER_DIM, THINKER_DIM))


def run_aut_encoder(mel_input, cu_seqlens=(0, 1, 2)):
    c1 = conv2d_stride2_mean(mel_input)   # (4, 8)
    c2 = conv2d_stride2_mean(c1)          # (2, 4)
    c3 = conv2d_stride2_mean(c2)          # (1, 2) -> 2 khung âm thanh sau 8× downsample
    frames = c3.T @ W_AUT_PROJ            # (2, 8)
    pos = np.array([[np.sin(t / 4.0 + i * 0.2) for i in range(THINKER_DIM)] for t in range(frames.shape[0])])
    x = frames + 0.15 * pos
    out = np.zeros_like(x)
    for idx in range(len(cu_seqlens) - 1):
        a, b = cu_seqlens[idx], cu_seqlens[idx + 1]
        seg = x[a:b]
        q = seg @ W_AUT_Q
        k = seg @ W_AUT_K
        v = seg @ W_AUT_V
        att = softmax((q @ k.T) / np.sqrt(THINKER_DIM), axis=-1)
        out[a:b] = seg + att @ v
    return {"c1": c1, "c2": c2, "c3": c3, "out": out}


# ── [1B] Interleaved TM-RoPE 3D: đan xen 3 trục (Temporal, Height, Width) ────
def apply_interleaved_tm_rope(vec, pos_thw):
    # vec: (..., 6) gồm 3 cặp đan xen: cặp 0 = Temporal, cặp 1 = Height, cặp 2 = Width
    thetas = np.array([1.0, 0.1, 0.01], dtype=float)
    angles = np.asarray(pos_thw, dtype=float) * thetas
    out = np.array(vec, dtype=float, copy=True)
    for pair_idx in range(3):
        phi = angles[pair_idx]
        c, s = np.cos(phi), np.sin(phi)
        x0 = out[..., 2 * pair_idx].copy()
        x1 = out[..., 2 * pair_idx + 1].copy()
        out[..., 2 * pair_idx] = x0 * c - x1 * s
        out[..., 2 * pair_idx + 1] = x0 * s + x1 * c
    return out


# ── [2A] Thinker MoE & Layer-24 Hidden Handoff ───────────────────────────────
W_THINKER_L12 = rng.normal(0.0, 0.3, size=(THINKER_DIM, THINKER_DIM))
W_THINKER_L24 = rng.normal(0.0, 0.3, size=(THINKER_DIM, THINKER_DIM))
W_THINKER_TOP = rng.normal(0.0, 0.3, size=(THINKER_DIM, THINKER_DIM))
THINKER_TEXT_EMB = rng.normal(0.0, 0.5, size=(2, THINKER_DIM))  # 2 token văn bản sinh ra


def run_thinker_handoff(text_emb):
    h12 = np.tanh(text_emb @ W_THINKER_L12)
    h24 = h12 + silu(h12 @ W_THINKER_L24)      # accept_hidden_layer = 24
    htop = h24 @ W_THINKER_TOP                 # tầng cuối sát LM Head
    return {"h12": h12, "h24": h24, "htop": htop}


# ── [3A] ResizeMLP Bridge (2048 → 1024) & Speaker Conditioning ───────────────
W_RESIZE_TXT1 = rng.normal(0.0, 0.35, size=(THINKER_DIM, THINKER_DIM))
B_RESIZE_TXT1 = rng.normal(0.0, 0.05, size=(THINKER_DIM,))
W_RESIZE_TXT2 = rng.normal(0.0, 0.35, size=(THINKER_DIM, TALKER_DIM))
B_RESIZE_TXT2 = rng.normal(0.0, 0.05, size=(TALKER_DIM,))

W_RESIZE_HID1 = rng.normal(0.0, 0.35, size=(THINKER_DIM, THINKER_DIM))
B_RESIZE_HID1 = rng.normal(0.0, 0.05, size=(THINKER_DIM,))
W_RESIZE_HID2 = rng.normal(0.0, 0.35, size=(THINKER_DIM, TALKER_DIM))
B_RESIZE_HID2 = rng.normal(0.0, 0.05, size=(TALKER_DIM,))

SPEAKER_TABLE = {
    2301: rng.normal(0.0, 0.4, size=(TALKER_DIM,)),  # Chelsie
    2302: rng.normal(0.0, 0.4, size=(TALKER_DIM,)),  # Ethan
    2303: rng.normal(0.0, 0.4, size=(TALKER_DIM,)),  # Aiden
}
TTS_PAD_EMBED = rng.normal(0.0, 0.25, size=(TALKER_DIM,))
CODEC_BOS_EMBED = rng.normal(0.0, 0.25, size=(TALKER_DIM,))


def resize_text_mlp(x):
    return silu(x @ W_RESIZE_TXT1 + B_RESIZE_TXT1) @ W_RESIZE_TXT2 + B_RESIZE_TXT2


def resize_hidden_mlp(x):
    return silu(x @ W_RESIZE_HID1 + B_RESIZE_HID1) @ W_RESIZE_HID2 + B_RESIZE_HID2


# ── [3C] Talker MoE (Top-2 of 6 Routed + 1 Shared) & Semantic Head c_0 ───────
W_TALKER_ROUTER = rng.normal(0.0, 0.4, size=(TALKER_DIM, NUM_EXPERTS))
W_TALKER_EXPERTS = rng.normal(0.0, 0.3, size=(NUM_EXPERTS, TALKER_DIM, TALKER_DIM))
W_TALKER_SHARED = rng.normal(0.0, 0.3, size=(TALKER_DIM, TALKER_DIM))
W_CODEC_HEAD = rng.normal(0.0, 0.4, size=(TALKER_DIM, CODEBOOK_SIZE))
TALKER_RMS_W = np.ones(TALKER_DIM, dtype=float)


def talker_moe_step(dual_track_in, speaker_vec):
    x = rmsnorm(dual_track_in + 0.25 * speaker_vec, TALKER_RMS_W)
    router_probs = softmax(x @ W_TALKER_ROUTER, axis=-1)
    top_idx = np.argsort(-router_probs)[:TOP_K_EXPERTS]
    top_vals = router_probs[top_idx]
    norm_weights = top_vals / np.sum(top_vals)  # norm_topk_prob = True
    routed_out = np.zeros(TALKER_DIM, dtype=float)
    for w, e_id in zip(norm_weights, top_idx):
        routed_out += w * silu(x @ W_TALKER_EXPERTS[int(e_id)])
    shared_out = silu(x @ W_TALKER_SHARED)
    h_talker = dual_track_in + routed_out + shared_out
    logits_c0 = h_talker @ W_CODEC_HEAD
    c0 = int(np.argmax(logits_c0))
    return {
        "h_talker": h_talker,
        "router_probs": router_probs,
        "top_idx": top_idx,
        "norm_weights": norm_weights,
        "logits_c0": logits_c0,
        "c0": c0,
    }


# ── [6A] Unified 16-Group RVQ Table (thu nhỏ: 4 × 8 = 32 dòng) ──────────────
UNIFIED_RVQ_TABLE = rng.normal(0.0, 0.45, size=(NUM_QUANTIZERS * CODEBOOK_SIZE, TALKER_DIM))
CODE_OFFSET = np.arange(NUM_QUANTIZERS, dtype=int).reshape(NUM_QUANTIZERS, 1) * CODEBOOK_SIZE


def lookup_rvq_offset_mean(codes_matrix):
    # codes_matrix: (NUM_QUANTIZERS, T_frames)
    shifted = codes_matrix + CODE_OFFSET
    emb = UNIFIED_RVQ_TABLE[shifted]  # (NUM_QUANTIZERS, T_frames, TALKER_DIM)
    return np.mean(emb, axis=0), shifted  # (T_frames, TALKER_DIM)


def lookup_rvq_separate_tables_mean(codes_matrix):
    # Tra cứu rời rạc từng bảng con để đối chứng với bảng gộp
    separate = [
        UNIFIED_RVQ_TABLE[q * CODEBOOK_SIZE : (q + 1) * CODEBOOK_SIZE][codes_matrix[q]]
        for q in range(NUM_QUANTIZERS)
    ]
    return np.mean(np.stack(separate, axis=0), axis=0)


# ── [4A] 5-Layer Depth Code Predictor (tự hồi quy c_1 .. c_{Q-1}) ────────────
W_DEPTH_ATTN = rng.normal(0.0, 0.3, size=(TALKER_DIM, TALKER_DIM))
W_DEPTH_MLP = rng.normal(0.0, 0.3, size=(TALKER_DIM, TALKER_DIM))
DEPTH_STEP_HEADS = [
    rng.normal(0.0, 0.4, size=(TALKER_DIM, CODEBOOK_SIZE))
    for _ in range(NUM_QUANTIZERS - 1)
]


def run_depth_code_predictor(h_talker, c0):
    # Prefill 2 vị trí: [h_talker, embed_0(c_0)]
    c0_emb = UNIFIED_RVQ_TABLE[0 * CODEBOOK_SIZE + int(c0)]
    seq = [h_talker, c0_emb]
    codes = [int(c0)]
    step_logs = []
    for step in range(NUM_QUANTIZERS - 1):
        mat = np.stack(seq, axis=0)  # (len_seq, TALKER_DIM)
        scores = (mat[-1:] @ mat.T) / np.sqrt(TALKER_DIM)
        ctx = softmax(scores, axis=-1)[0] @ (mat @ W_DEPTH_ATTN)
        h_depth = mat[-1] + ctx + silu((mat[-1] + ctx) @ W_DEPTH_MLP)
        logits_k = h_depth @ DEPTH_STEP_HEADS[step]
        ck = int(np.argmax(logits_k))
        codes.append(ck)
        step_logs.append({"step": step, "seq_len": len(seq), "code_idx": step + 1, "code": ck, "logits": logits_k})
        # Tra bảng nhúng riêng của tầng (step + 1) cho bước vi mô kế tiếp
        next_emb = UNIFIED_RVQ_TABLE[(step + 1) * CODEBOOK_SIZE + ck]
        seq.append(next_emb)
    return np.array(codes, dtype=int), step_logs


# ── [3B] Vòng lặp Dual-Track 12.5 Hz qua 4 khung 80 ms ───────────────────────
def run_dual_track_frames(trailing_text_hidden, speaker_id=2302, n_frames=4):
    speaker_vec = SPEAKER_TABLE[speaker_id]
    prev_audio_emb = CODEC_BOS_EMBED.copy()
    all_codes = []
    frame_records = []
    for t in range(n_frames):
        used_pad = t >= len(trailing_text_hidden)
        text_cond = TTS_PAD_EMBED if used_pad else trailing_text_hidden[t]
        dual_in = prev_audio_emb + text_cond
        t_out = talker_moe_step(dual_in, speaker_vec)
        codes_t, depth_logs = run_depth_code_predictor(t_out["h_talker"], t_out["c0"])
        all_codes.append(codes_t)
        # Nhúng trung bình của 4 sổ mã khung t làm phản hồi cho khung t+1
        frame_emb, _ = lookup_rvq_offset_mean(codes_t.reshape(NUM_QUANTIZERS, 1))
        prev_audio_emb = frame_emb[0]
        frame_records.append({
            "frame": t,
            "used_pad": used_pad,
            "dual_in": dual_in,
            "talker": t_out,
            "codes": codes_t,
            "depth_logs": depth_logs,
        })
    codes_matrix = np.stack(all_codes, axis=1)  # (NUM_QUANTIZERS, n_frames)
    return codes_matrix, frame_records


# ── [5A] Pre-Transformer (Sliding-Window Causal Mask) & ConvNeXt 2× ──────────
W_PRE_Q = rng.normal(0.0, 0.3, size=(TALKER_DIM, TALKER_DIM))
W_PRE_K = rng.normal(0.0, 0.3, size=(TALKER_DIM, TALKER_DIM))
W_PRE_V = rng.normal(0.0, 0.3, size=(TALKER_DIM, TALKER_DIM))
LAYER_SCALE = np.full(TALKER_DIM, 0.01, dtype=float)
W_UP_PRE = rng.normal(0.0, 0.35, size=(PRE_UPSAMPLE, TALKER_DIM, TALKER_DIM))


def build_sliding_causal_mask(T, window=SLIDING_WINDOW):
    mask = np.zeros((T, T), dtype=bool)
    for q in range(T):
        for k in range(T):
            if 0 <= q - k < window:
                mask[q, k] = True
    return mask


def run_pre_transformer_and_upsample(frame_embeds):
    T = frame_embeds.shape[0]
    mask = build_sliding_causal_mask(T, SLIDING_WINDOW)
    q = frame_embeds @ W_PRE_Q
    k = frame_embeds @ W_PRE_K
    v = frame_embeds @ W_PRE_V
    raw_scores = (q @ k.T) / np.sqrt(TALKER_DIM)
    masked_scores = np.where(mask, raw_scores, -1e9)
    attn_w = softmax(masked_scores, axis=-1)
    h_pre = frame_embeds + LAYER_SCALE * (attn_w @ v)
    # Causal Transposed Upsample ×2: mỗi khung t nở thành 2 bước con nhân quả
    up_list = []
    for t in range(T):
        for s in range(PRE_UPSAMPLE):
            up_list.append(h_pre[t] @ W_UP_PRE[s])
    h_up = np.stack(up_list, axis=0)  # (T * PRE_UPSAMPLE, TALKER_DIM)
    return {"mask": mask, "attn_w": attn_w, "h_pre": h_pre, "h_up": h_up}


# ── [6B] SnakeBeta Causal ConvNet & Chunked 24kHz Wave Decode ────────────────
SNAKE_ALPHA = np.linspace(0.2, 0.8, TALKER_DIM)
SNAKE_BETA = np.linspace(0.5, 1.1, TALKER_DIM)
W_DEC_STAGE1 = rng.normal(0.0, 0.3, size=(DEC_UPSAMPLES[0], TALKER_DIM, TALKER_DIM))
W_DEC_STAGE2 = rng.normal(0.0, 0.3, size=(DEC_UPSAMPLES[1], TALKER_DIM, 1))


def snake_beta(x, alpha=SNAKE_ALPHA, beta=SNAKE_BETA):
    a = np.exp(alpha)
    b = np.exp(beta)
    return x + (1.0 / (b + 1e-9)) * (np.sin(x * a) ** 2)


def run_snakebeta_decoder(h_up):
    # Stage 1: SnakeBeta -> upsample ×2
    s1_in = snake_beta(h_up)
    s1_list = []
    for t in range(s1_in.shape[0]):
        for r in range(DEC_UPSAMPLES[0]):
            s1_list.append(s1_in[t] @ W_DEC_STAGE1[r])
    s1_out = np.stack(s1_list, axis=0)
    # Stage 2: SnakeBeta -> upsample ×3 -> 1 kênh sóng âm
    s2_in = snake_beta(s1_out)
    wav_list = []
    for t in range(s2_in.shape[0]):
        for r in range(DEC_UPSAMPLES[1]):
            val = float((s2_in[t] @ W_DEC_STAGE2[r])[0])
            wav_list.append(val)
    wav = np.clip(np.array(wav_list, dtype=float), -1.0, 1.0)
    return wav


def decode_codes_to_wav(codes_matrix):
    frame_embeds, _ = lookup_rvq_offset_mean(codes_matrix)
    pre = run_pre_transformer_and_upsample(frame_embeds)
    return run_snakebeta_decoder(pre["h_up"])


def chunked_decode_codes(codes_matrix, chunk_size=2, left_context_size=1):
    total_frames = codes_matrix.shape[1]
    wav_chunks = []
    start = 0
    while start < total_frames:
        end = min(start + chunk_size, total_frames)
        ctx = left_context_size if start - left_context_size > 0 else start
        sub_codes = codes_matrix[:, start - ctx : end]
        sub_wav = decode_codes_to_wav(sub_codes)
        wav_chunks.append(sub_wav[ctx * TOTAL_UPSAMPLE :])
        start = end
    return np.concatenate(wav_chunks, axis=0)


# ── Chạy toàn trình và in báo cáo 1A → 6B ────────────────────────────────────
print("=== LAB QWEN3-OMNI & QWEN3-TTS-12HZ: TOÀN TRÌNH 1A → 6B BẰNG NUMPY ===")
print("MỤC TIÊU: chạy khép kín đường ống Thinker–Talker MoE + 16-RVQ + SnakeBeta Code2Wav,")
print("kiểm chứng hạ mẫu AuT 8×, TM-RoPE 3D, đồng bộ Dual-Track, Depth Predictor và bảng RVQ gộp.")

aut = run_aut_encoder(MEL_INPUT, cu_seqlens=(0, 1, 2))
say("1A", "AuT AUDIO ENCODER (3× CONV2D STRIDE-2 & WINDOWED ATTENTION)", [
    f"Mel spectrogram đầu vào: shape {MEL_INPUT.shape} (8 dải Mel × 16 khung thời gian).",
    f"Sau conv2d1 (stride=2): {aut['c1'].shape} → sau conv2d2: {aut['c2'].shape} → sau conv2d3: {aut['c3'].shape}.",
    f"Tỷ lệ nén thời gian: {MEL_FRAMES} → {aut['c3'].shape[1]} khung (giảm đúng 2³ = 8×).",
    f"Đặc trưng âm thanh sau Windowed Attn (cu_seqlens=(0,1,2)): shape {aut['out'].shape}",
    f"   khung âm thanh 0: {fmt_vec(aut['out'][0])}",
])

q_sample = aut["out"][0, :HEAD_DIM]
q_rot = apply_interleaved_tm_rope(q_sample, pos_thw=(3, 3, 3))
norm_before = float(np.linalg.norm(q_sample))
norm_after = float(np.linalg.norm(q_rot))
say("1B", "INTERLEAVED TM-ROPE 3D (TEMPORAL, HEIGHT, WIDTH @ 12.5 HZ)", [
    "Prod phân bổ 64 cặp theo mrope_section = [24, 20, 20]; lab đan xen 3 cặp (T, H, W).",
    f"Vector trước xoay : {fmt_vec(q_sample)} (‖q‖ = {norm_before:.6f})",
    f"Vector sau xoay   : {fmt_vec(q_rot)} (‖RoPE(q)‖ = {norm_after:.6f})",
    f"Độ lệch chuẩn Euclid |‖RoPE(q)‖ − ‖q‖| = {abs(norm_after - norm_before):.2e}.",
])

thinker = run_thinker_handoff(THINKER_TEXT_EMB)
say("2A", "THINKER MOE BACKBONE & TRÍCH NGANG TRẠNG THÁI ẨN LỚP 24", [
    f"Thinker sinh 2 token văn bản; trích ngang tại accept_hidden_layer = 24 (shape {thinker['h24'].shape}):",
    f"   trailing_text_hidden[0] (Lớp 24): {fmt_vec(thinker['h24'][0])}",
    f"   trailing_text_hidden[1] (Lớp 24): {fmt_vec(thinker['h24'][1])}",
    f"   So với tầng đỉnh sát LM Head    : {fmt_vec(thinker['htop'][0])}",
])

txt_proj = resize_text_mlp(THINKER_TEXT_EMB)
hid_proj = resize_hidden_mlp(thinker["h24"])
say("3A", "CẦU NỐI RESIZEMLP KÉP (2048 → 1024) & ĐIỀU KIỆN HÓA SPEAKER ID", [
    f"text_projection(embed)   shape {THINKER_TEXT_EMB.shape} → {txt_proj.shape}: {fmt_vec(txt_proj[0])}",
    f"hidden_projection(L24)   shape {thinker['h24'].shape} → {hid_proj.shape}: {fmt_vec(hid_proj[0])}",
    f"Hồ sơ giọng đọc speaker_id = 2302 (Ethan)  : {fmt_vec(SPEAKER_TABLE[2302])}",
    f"Hồ sơ giọng đọc speaker_id = 2301 (Chelsie): {fmt_vec(SPEAKER_TABLE[2301])}",
])

codes_matrix, frame_records = run_dual_track_frames(hid_proj, speaker_id=2302, n_frames=4)
say("3B", "ĐỒNG BỘ DUAL-TRACK STREAMING 12.5 HZ (4 KHUNG 80 MS, 2 TOKEN CHỮ)", [
    "Khung {} (80ms): nguồn chữ = {} → dual_in[:3] = {}".format(
        r["frame"],
        "tts_pad_embed" if r["used_pad"] else f"trailing_text_hidden[{r['frame']}]",
        fmt_vec(r["dual_in"][:3]),
    )
    for r in frame_records
] + [
    "Kết luận: Khung 0 phát ngay khi có chữ đầu tiên (độ trễ 97 ms); Khung 2–3 tự động dùng tts_pad_embed.",
])

f0_talker = frame_records[0]["talker"]
say("3C", "TALKER MOE (TOP-K ROUTING + SHARED EXPERT) & MÃ NGỮ NGHĨA c_0", [
    f"Khung 0: xác suất 6 chuyên gia = {fmt_vec(f0_talker['router_probs'])}",
    f"Chọn Top-{TOP_K_EXPERTS} chuyên gia {list(map(int, f0_talker['top_idx']))} → trọng số chuẩn hóa = {fmt_vec(f0_talker['norm_weights'])} (tổng = {np.sum(f0_talker['norm_weights']):.4f}).",
    f"Logits codec_head (8 mã): {fmt_vec(f0_talker['logits_c0'])} → chọn mã ngữ nghĩa c_0 = {f0_talker['c0']}.",
])

say("4A", "5-LAYER DEPTH CODE PREDICTOR (TỰ HỒI QUY MÃ ÂM HỌC TRONG 80 MS)", [
    f"Bước depth {d['step']} (chiều dài chuỗi nội bộ = {d['seq_len']}): dùng lm_head[{d['step']}] → sinh c_{d['code_idx']} = {d['code']}"
    for d in frame_records[0]["depth_logs"]
] + [
    f"Bộ mã đầy đủ của Khung 0 [c_0 .. c_3]: {codes_matrix[:, 0].tolist()} (prod sinh 16 mã [c_0 .. c_15]).",
])

rvq_mean, shifted_indices = lookup_rvq_offset_mean(codes_matrix)
say("6A", "BẢNG NHÚNG 16-GROUP RVQ HỢP NHẤT (CODE_OFFSET) & MEAN POOLING", [
    f"Ma trận mã 4 tầng × 4 khung:\n      {codes_matrix.tolist()}",
    f"Cộng code_offset {CODE_OFFSET.ravel().tolist()} → chỉ số tra bảng gộp 32 dòng:\n      {shifted_indices.tolist()}",
    f"Vector trung bình .mean(dim=0) của Khung 0: {fmt_vec(rvq_mean[0])}",
])

pre_res = run_pre_transformer_and_upsample(rvq_mean)
mask_rows = [
    f"   q={q} | " + " ".join("1" if pre_res["mask"][q, k] else "." for k in range(4))
    for q in range(4)
]
say("5A", f"CODE2WAV PRE-TRANSFORMER (SLIDING WINDOW = {SLIDING_WINDOW}) & CONVNEXT {PRE_UPSAMPLE}×", [
    f"Mặt nạ nhân quả cửa sổ trượt 4×4 (window = {SLIDING_WINDOW}, 1 = cho phép, . = chặn):",
] + mask_rows + [
    f"Sau Pre-Transformer (LayerScale=0.01): shape {pre_res['h_pre'].shape} → sau TransConv {PRE_UPSAMPLE}×: shape {pre_res['h_up'].shape}.",
])

wav_full = run_snakebeta_decoder(pre_res["h_up"])
wav_chunked = chunked_decode_codes(codes_matrix, chunk_size=2, left_context_size=1)
say("6B", f"SNAKEBETA CAUSAL CONVNET ({DEC_UPSAMPLES[0]}×{DEC_UPSAMPLES[1]} = {DEC_UPSAMPLES[0]*DEC_UPSAMPLES[1]}×, TỔNG {TOTAL_UPSAMPLE}×) → SÓNG ÂM", [
    f"4 khung mã @ 12.5 Hz × {TOTAL_UPSAMPLE} mẫu/khung = {len(wav_full)} mẫu sóng âm (prod: 4 × 1920 = 7680 mẫu @ 24 kHz).",
    f"8 mẫu sóng âm đầu tiên (nằm trong [-1, 1]): {fmt_vec(wav_full[:8])}",
    f"Giải mã theo từng chunk (chunk_size=2, left_context=1): thu được {len(wav_chunked)} mẫu liền mạch.",
])

# ── [KIỂM TRA & THỬ NGHIỆM] ─────────────────────────────────────────────────
print("\n[KIỂM TRA & THỬ NGHIỆM]")

# (1) Kiểm chứng bảng gộp code_offset vs 4 bảng rời rạc
rvq_sep = lookup_rvq_separate_tables_mean(codes_matrix)
diff_rvq = float(np.max(np.abs(rvq_mean - rvq_sep)))
print(f"  (1) Bảng gộp code_offset vs tra cứu từng bảng RVQ rời: lệch tối đa = {diff_rvq:.2e} (ngưỡng 0.00e+00).")
assert diff_rvq == 0.0

# (2) Kiểm chứng bảo toàn chuẩn TM-RoPE 3D và chuẩn hóa Top-K MoE
rope_err = abs(norm_after - norm_before)
moe_sum_err = abs(float(np.sum(f0_talker["norm_weights"])) - 1.0)
print(f"  (2) Sai số chuẩn 3D TM-RoPE = {rope_err:.2e}; sai số tổng trọng số Top-K Talker MoE = {moe_sum_err:.2e}.")
assert rope_err < 1e-12 and moe_sum_err < 1e-12

# (3) Kiểm chứng cô lập nhân quả & khớp Chunked Streaming Decode với Left Context
codes_perturbed = codes_matrix.copy()
codes_perturbed[:, 3] = (codes_perturbed[:, 3] + 3) % CODEBOOK_SIZE  # đổi khung cuối (frame 3)
wav_perturbed = decode_codes_to_wav(codes_perturbed)
causal_leak = float(np.max(np.abs(wav_full[: 3 * TOTAL_UPSAMPLE] - wav_perturbed[: 3 * TOTAL_UPSAMPLE])))
chunk_diff = float(np.max(np.abs(wav_full - wav_chunked)))
print(f"  (3) Đổi mã khung cuối (frame 3): rò rỉ ngược về 3 khung đầu = {causal_leak:.2e} (nhân quả tuyệt đối).")
print(f"      Full decode vs Chunked decode (chunk=2, left_context=1): lệch tối đa = {chunk_diff:.2e} (< 1e-12).")
assert causal_leak == 0.0
assert chunk_diff < 1e-12

print("PASS: toàn trình 1A→6B của Qwen3-Omni & Qwen3-TTS-12Hz chạy đúng trên số thật.")