"""Lab Transformer: một lớp decoder-only chạy NumPy thuần, đồng bộ thuật ngữ
với khung kiến trúc tương tác AttentionArchitectureOverlay (mã 1A..4C).

Con đường kiến trúc lab đã chọn (phải nói rõ, không nói chung chung):
  - MHA (Multi-Head Attention): n_q = n_kv = 2, mỗi Q head có K/V riêng.
  - RoPE chuẩn cho từng head, chỉ xoay Q và K, theo VỊ TRÍ TUYỆT ĐỐI
    (= độ dài cache trước đó + chỉ số token trong pass hiện tại).
  - Pre-Norm RMSNorm, causal scaled dot-product attention (chia √d_head),
    output projection W_O, và đủ hai phép cộng residual.
  - FFN Dense SwiGLU (không MoE), cuối cùng là RMSNorm + LM head.

KV cache bắt chước cache thật: lưu K đã xoay RoPE và V chưa xoay, theo từng
head; decode nối 1 token mới vào cache thay vì tính lại token cũ. Đây là mảng
NumPy giáo dục: KHÔNG có HBM thật, block phân trang (PagedAttention) hay GPU.
GQA/MQA/MLA/MoE chỉ được nêu tên ở 3B/4B như các lựa chọn kiến trúc, KHÔNG
được chạy trong lab.
"""
import numpy as np

# ── Cấu hình demo: số nhỏ để giá trị có thể in thẳng ra màn hình ─────────────
VOCAB = 6           # bảng từ vựng demo
D_MODEL = 4         # chiều vector ẩn
N_HEADS = 2         # MHA: n_q = n_kv = 2
D_HEAD = 2          # chiều mỗi head (n_heads * d_head == d_model)
MAX_CONTEXT = 12
FFN_HIDDEN = 8      # chiều ẩn mạng Dense SwiGLU
ROPE_BASE = 10000.0

# weights chưa học: ngẫu nhiên có seed cố định, chạy lại cho ra cùng số.
rng = np.random.default_rng(4)
E = rng.normal(0, 0.25, (VOCAB, D_MODEL))                # bảng nhúng token
# MHA: MỖI Q head có W_Q / W_K / W_V riêng — không chia sẻ như GQA/MQA/MLA.
Wq = rng.normal(0, 0.30, (N_HEADS, D_MODEL, D_HEAD))
Wk = rng.normal(0, 0.30, (N_HEADS, D_MODEL, D_HEAD))
Wv = rng.normal(0, 0.30, (N_HEADS, D_MODEL, D_HEAD))
Wo = rng.normal(0, 0.30, (D_MODEL, D_MODEL))             # output projection
Wg = rng.normal(0, 0.30, (D_MODEL, FFN_HIDDEN))          # nhánh gate SwiGLU
Wz = rng.normal(0, 0.30, (D_MODEL, FFN_HIDDEN))          # nhánh up SwiGLU
Wf = rng.normal(0, 0.20, (FFN_HIDDEN, D_MODEL))          # chiếu về d_model
U = rng.normal(0, 0.20, (D_MODEL, VOCAB))                # LM head

DEMO_PROMPT = [1, 2, 3]        # Giữ làm mốc so sánh của bài học.
TRY_PROMPT = [1, 2, 4]         # Bài tập: chỉ đổi dòng này, dự đoán trước khi chạy.
DEMO_CONTINUATION = [4, 0, 2]  # token demo decode thêm (chỉ minh họa số học)


def rmsnorm(x):
    """RMSNorm không tham số (như cũ): giữ hướng, chỉ chuẩn hóa biên độ."""
    return x / np.sqrt(np.mean(x * x, axis=-1, keepdims=True) + 1e-6)


def rope_angles(positions):
    """Góc RoPE cho các vị trí tuyệt đối: [half, T]."""
    half = D_HEAD // 2
    freqs = ROPE_BASE ** (-2.0 * np.arange(half) / D_HEAD)
    return np.asarray(positions, dtype=float)[None, :] * freqs[:, None]


def rope_rotate(x, positions):
    """Xoay RoPE từng cặp nửa chiều của x: áp dụng cho Q và K, KHÔNG cho V.

    x: [n_heads, T, d_head]; positions là vị trí tuyệt đối (đã cộng offset
    của cache), cho nên token decode thứ 4 xoay đúng góc của vị trí 3.
    """
    half = D_HEAD // 2
    ang = rope_angles(positions)                 # [half, T]
    cos = np.cos(ang).T[None, :, :]              # [1, T, half]
    sin = np.sin(ang).T[None, :, :]              # [1, T, half]
    x1, x2 = x[..., :half], x[..., half:]
    r1 = x1 * cos - x2 * sin
    r2 = x1 * sin + x2 * cos
    return np.concatenate([r1, r2], axis=-1)


def empty_cache():
    """KV cache trống theo MHA: K và V riêng cho từng head, chưa có token nào."""
    return {
        "k": np.zeros((N_HEADS, 0, D_HEAD)),
        "v": np.zeros((N_HEADS, 0, D_HEAD)),
    }


def attention_traced(ids, cache=None):
    """Tính một pass forward và trả về MỌI tensor trung gian (B=1 bị ẩn).

    Thứ tự tính theo phụ thuộc số học (không theo cột vẽ trong sơ đồ):
    1A → 1B → 3A/3B → 2A → 2B → 2C → 2D → 2E → 4A → 4B → 4C.
    Cột 3 của sơ đồ là nhánh hỗ trợ attention: cache phải sẵn sàng TRƯỚC
    khi 2B tạo scores và 2E nhân với V, không phải bước chạy sau attention.
    """
    ids = np.asarray(ids, dtype=int)
    assert ids.ndim == 1 and len(ids) > 0
    assert np.all((ids >= 0) & (ids < VOCAB))
    if cache is None:
        cache = empty_cache()
    Kc, Vc = cache["k"], cache["v"]
    Tc = int(Kc.shape[1])
    assert Vc.shape[1] == Tc, "K cache và V cache lệch độ dài token"
    assert Kc.shape == (N_HEADS, Tc, D_HEAD) and Vc.shape == (N_HEADS, Tc, D_HEAD), (
        "cache phải có shape (n_heads, token, d_head) theo MHA"
    )
    T = len(ids)
    assert Tc + T <= MAX_CONTEXT, "Vượt demo ngữ cảnh tối đa"

    # Vị trí tuyệt đối: offset theo độ dài cache, tức là 3A cung cấp "đã có bao nhiêu token".
    q_pos = np.arange(Tc, Tc + T)          # ví dụ decode: [3] thay vì [0]
    key_index = np.arange(Tc + T)

    # 1A — Hidden State & RMSNorm (Pre-Norm)
    X = E[ids].astype(float)               # [T, d_model] (chiều B=1 đã ẩn)
    N = rmsnorm(X)

    # 1B — Chiếu Q/K/V theo head (MHA: W riêng từng head)
    Q = np.einsum("td,hde->hte", N, Wq)    # [n_heads, T, d_head]
    Knew = np.einsum("td,hde->hte", N, Wk)
    Vnew = np.einsum("td,hde->hte", N, Wv)
    Qr = rope_rotate(Q, q_pos)             # RoPE trên Q
    Kr = rope_rotate(Knew, q_pos)          # RoPE trên K; V KHÔNG xoay

    # 3A — KV Cache theo MHA: nối K/V mới vào cache (K đã xoay, V chưa xoay)
    K = np.concatenate([Kc, Kr], axis=1)   # [n_heads, Tc+T, d_head]
    V = np.concatenate([Vc, Vnew], axis=1)

    # 2A/2B — mặt nạ nhân quả & scores chia √d_head (KHÔNG phải √d_model)
    scores = np.einsum("htd,hkd->htk", Qr, K) / np.sqrt(D_HEAD)   # [n_heads, Tq, Tkey]
    allow = key_index[None, :] <= q_pos[:, None]                  # j ≤ i (kể cả cache)
    masked = np.where(allow[None, :, :], scores, -np.inf)         # 2C — cộng mask

    # 2D — softmax từng hàng (trọng số chú ý trên các token)
    expd = np.exp(masked - masked.max(axis=-1, keepdims=True))
    prob = expd / expd.sum(axis=-1, keepdims=True)                # [n_heads, Tq, Tkey]

    # 2E — tổng hợp ngữ cảnh P·V (V giữ nguyên, không xoay)
    ctx = np.einsum("htk,hkd->htd", prob, V)                      # [n_heads, Tq, d_head]
    o = ctx.transpose(1, 0, 2).reshape(T, D_MODEL)                # ghép 2 head về d_model

    # 4A — Output projection W_O + residual thứ nhất
    attn_out = o @ Wo
    H = X + attn_out

    # 4B — RMSNorm + FFN Dense SwiGLU (không MoE) + residual thứ hai
    M2 = rmsnorm(H)
    gate = M2 @ Wg
    up = M2 @ Wz
    ffn = (gate / (1.0 + np.exp(-gate))) * up                     # silu(gate) · up
    F = H + ffn @ Wf

    # 4C — RMSNorm cuối + LM head → logits (điểm số, KHÔNG phải xác suất)
    logits = rmsnorm(F) @ U                                       # [T, VOCAB]

    return {
        "ids": ids, "positions": q_pos, "cache_past": Tc,
        "x": X, "n": N, "q": Q, "k_new": Knew, "v_new": Vnew,
        "q_rot": Qr, "k_rot": Kr,
        "k": K, "v": V, "scores": scores, "allow": allow,
        "masked": masked, "prob": prob, "ctx": ctx, "attn_out": attn_out,
        "h1": H, "gate": gate, "up": up, "ffn": ffn, "f2": F,
        "logits": logits, "cache": {"k": K, "v": V},
    }


def forward(ids, cache=None, trace=False):
    """Forward một lớp; im lặng mặc định, in chi tiết 1A..4C khi trace=True.

    Trả về (logits, cache, shapes) — cache là dict {"k", "v"} theo head.
    """
    st = attention_traced(ids, cache)
    if trace:
        _print_pass(st, "Prefill" if st["cache_past"] == 0 else "Decode")
    return (
        st["logits"],
        st["cache"],
        {
            "embedding": st["x"].shape,
            "attention weights": st["prob"].shape,
            "MLP hidden": st["ffn"].shape,
            "logits": st["logits"].shape,
        },
    )


# ── Bộ in câu chuyện 1A..4C (formatter dùng chung, giữ output dễ đọc) ───────

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


def fmt_rows(a, nd=3, max_rows=None):
    a = np.asarray(a, dtype=float)
    rows = a if a.ndim > 1 else a[None, :]
    if max_rows:
        rows = rows[:max_rows]
    return ["  " + fmt_vec(r, nd) for r in rows]


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


def _print_pass(st, phase):
    """In lần lượt từng mã overlay cho một pass; giá trị đều là giá trị thật."""
    T = len(st["ids"])
    Tkey = st["k"].shape[1]
    print(f"--- {phase}: xử lý mã token {st['ids'].tolist()} tại vị trí {st['positions'].tolist()} ---")
    say("1A", "Hidden State x_t & RMSNorm (Pre-Norm) — Chuẩn bị vector", [
        "Quan sát: mã token chỉ dùng để tra bảng E. Vector mới là dữ liệu đi qua lớp.",
        f"shape x=({T}, {D_MODEL}) (nghĩa là {T} vị trí × d_model={D_MODEL}; batch B=1 bị ẩn).",
        "hàng 0 của X: " + fmt_vec(st["x"][0]),
        "hàng 0 sau RMSNorm: " + fmt_vec(st["n"][0]),
        "Kết luận: số chiều giữ nguyên; RMSNorm điều chỉnh biên độ, chưa trộn thông tin giữa token.",
        "Giữ X gốc để cộng residual ở 4A. Không dùng vector chuẩn hóa thay cho nhánh gốc này.",
    ])
    rope_idx = 1 if T > 1 else 0
    rope_pos = int(st["positions"][rope_idx])
    say("1B", "Chiếu Query/Key/Value (W_Q, W_K, W_V) & RoPE", [
        "Quan sát: một vector được chiếu thành Q để hỏi, K để đối chiếu, V làm nội dung cần lấy.",
        f"shape Q, K, V=({N_HEADS}, {T}, {D_HEAD}) (n_heads · n_tokens · d_head).",
        f"Q head 0, hàng 0 (vị trí {int(st['positions'][0])}): " + fmt_vec(st["q_rot"][0][0]),
        f"Q head 0, hàng {rope_idx} (vị trí {rope_pos} — đã xoay RoPE): " + fmt_vec(st["q_rot"][0][rope_idx])
        + " | chưa xoay: " + fmt_vec(st["q"][0][rope_idx]),
        "K head 0 sau RoPE, hàng đang xem: " + fmt_vec(st["k_rot"][0][rope_idx]),
        "V head 0, hàng 0: " + fmt_vec(st["v_new"][0][0]),
        f"Kết luận: RoPE xoay Q/K theo vị trí {st['positions'].tolist()}, V không xoay.",
        "Ở vị trí 0 góc xoay bằng 0; ở vị trí khác, hãy so số trước/sau để nhận ra thông tin vị trí.",
    ])
    say("3A", "KV Cache theo MHA", [
        f"shape K cache, V cache=({N_HEADS}, {Tkey}, {D_HEAD}) — head × token × d_head.",
        f"MHA: mỗi Q head có K/V riêng — Q0 → (K0, V0), Q1 → (K1, V1).",
        f"Độ dài token cache của pass này: {st['cache_past']} → {Tkey}.",
        "K head 0, vị trí cuối cache: " + fmt_vec(st["k"][0][-1]),
        "V head 0, vị trí cuối cache: " + fmt_vec(st["v"][0][-1]),
        "Quan sát: cache lưu K đã xoay RoPE và V chưa xoay của từng head.",
        ("Kết luận: prefill vừa lưu K/V của cả prompt; lượt decode sau có thể dùng lại chúng."
         if st["cache_past"] == 0 else
         f"Kết luận: giữ {st['cache_past']} vị trí cũ, chỉ tính rồi nối K/V của {T} token mới."),
        "Đây là mảng NumPy, không phải HBM hay bộ nhớ phân trang thật trên hình.",
    ])
    say("3B", "Tổ chức head: Lab Transformer chọn MHA", [
        f"n_q = n_kv = {N_HEADS} (tỉ lệ 1:1:1). GQA (g:1:1), MQA (n:1:1), MLA (latent) KHÔNG chạy trong lab.",
        "Quan sát hai head tại cùng vị trí 0:",
        "K0: " + fmt_vec(st["k"][0][0]) + " | K1: " + fmt_vec(st["k"][1][0]),
        "Kết luận: MHA là cách TỔ CHỨC HEAD; Q0 chỉ đọc K0/V0, Q1 chỉ đọc K1/V1.",
        "Cache ở 3A tái sử dụng THEO TOKEN qua các lượt; giữ token cũ không có nghĩa là chia sẻ giữa head.",
    ])
    say("2A", "Causal Mask (0 tại j≤i, -inf tại j>i)", [
        f"shape mask=({T}, {Tkey}); cột j là vị trí nguồn, hàng là token đang hỏi.",
        f"Hàng đầu đang ở vị trí tuyệt đối {int(st['positions'][0])}; chỉ được đọc j không lớn hơn vị trí này.",
        "Mask hàng đầu: " + ", ".join("0" if a else "-inf" for a in st["allow"][0]),
        "Kết luận: 0 nghĩa là cho phép, -inf nghĩa là chặn. Khi decode ở vị trí 3, các cột 0..3 đều hợp lệ.",
    ])
    say("2B", "Scaled Dot-Product Attention: Q · K^T / √d_head", [
        f"shape scores=({N_HEADS}, {T}, {Tkey}).",
        "scores head 0:"
        ] + fmt_rows(st["scores"][0]) + [
        "Cách đọc: ô hàng i, cột j là điểm đối chiếu Q của token i với K ở vị trí j, cùng head.",
        f"Chia √d_head = √{D_HEAD} ≈ {np.sqrt(D_HEAD):.3f} để kiểm soát độ lớn điểm.",
        "Kết luận: đây mới là điểm liên quan thô, có thể âm; chưa phải trọng số để gom nội dung.",
    ])
    say("2C", "Add Causal Mask (cộng mask vào scores)", [
        f"shape masked=({N_HEADS}, {T}, {Tkey}); head 0 sau khi cộng mask:"
        ] + fmt_rows(st["masked"][0]) + [
        "Ô hợp lệ (j≤i) cộng 0 nên giữ nguyên điểm; ô tương lai (j>i) cộng -∞ và sau softmax sẽ có xác suất 0.",
    ])
    say("2D", "Softmax Attention Weights", [
        f"shape P=({N_HEADS}, {T}, {Tkey}) — xác suất TRỌNG SỐ trên CÁC TOKEN, mỗi hàng tổng 1, không âm.",
        "P head 0:"
        ] + fmt_rows(st["prob"][0]) + [
        "Quan sát hàng đầu: khi đang ở vị trí 0, [1, 0, 0] nghĩa là chỉ lấy thông tin từ chính nó.",
        "Tổng hàng cuối head 0: " + f"{st['prob'][0][-1].sum():.6f}",
        "Kết luận: mỗi hàng phân bổ trọng số trên vị trí đầu vào, không phải trên sáu token ứng viên.",
        "Đây là trọng số chọn nguồn thông tin; điểm chọn token kế tiếp chỉ có ở 4C.",
    ])
    say("2E", "Context Aggregation: O = P · V", [
        f"shape mỗi head=({T}, {D_HEAD}); ghép 2 head thành ({T}, {D_MODEL}).",
        "O head 0, hàng 0: " + fmt_vec(st["ctx"][0][0]),
        "V head 0, vị trí nguồn đầu: " + fmt_vec(st["v"][0][0]),
        "Cách tính: nhân mỗi trọng số trong hàng P với V cùng cột, rồi cộng các vector lại.",
        "Kết luận: ở vị trí 0, trọng số [1, 0, 0] trả về đúng V đầu; các vị trí sau gom nhiều nguồn hợp lệ.",
    ])
    say("4A", "Output Projection W_O + Residual", [
        f"shape AttnOut=({T}, {D_MODEL}); H = X + AttnOut=({T}, {D_MODEL}).",
        "AttnOut hàng 0: " + fmt_vec(st["attn_out"][0]),
        "X gốc hàng 0: " + fmt_vec(st["x"][0]),
        "H sau cộng hàng 0: " + fmt_vec(st["h1"][0]),
        "Kết luận: W_O trộn kết quả hai head; residual cộng phần cập nhật vào X, không bỏ thông tin gốc.",
    ])
    say("4B", "RMSNorm + FFN (Dense SwiGLU, không MoE)", [
        f"shape gate/up=({T}, {FFN_HIDDEN}), đầu SwiGLU=({T}, {FFN_HIDDEN}); F = H + FFN=({T}, {D_MODEL}).",
        "gate hàng 0: " + fmt_vec(st["gate"][0]),
        "F sau cộng residual lần hai, hàng 0: " + fmt_vec(st["f2"][0]),
        "Quan sát: mỗi vị trí được chuẩn hóa rồi xử lý riêng bằng Dense SwiGLU, sau đó cộng vào H.",
        "Kết luận: attention đọc giữa các vị trí; FFN xử lý từng vector với cùng weights. Nhánh MoE không chạy.",
    ])
    say("4C", "Next Layer Hidden State / LM Head Logits", [
        f"shape logits=({T}, {VOCAB}) — MỖI VỊ TRÍ có một hàng điểm số cho {VOCAB} token.",
        "logits (3 chữ số, mỗi hàng một vị trí):"
        ] + fmt_rows(st["logits"]) + [
        "Cột theo thứ tự mã token: " + ", ".join(str(i) for i in range(VOCAB)),
        "Cách đọc: hàng cuối cho điểm của token đi sau toàn bộ chuỗi đang có; mỗi cột ứng với một mã ứng viên.",
        "Kết luận: logits là điểm số chưa chuẩn hóa (không phải xác suất). Không cần cộng chúng thành 1.",
        "Lab chỉ có một lớp nên chiếu LM head ngay ở đây; mô hình nhiều lớp phải đi hết stack trước.",
    ])


# ── Demo chạy khi thực thi file (trình duyệt/CLI/test đều dùng chung) ────────

print("=== THEO DÕI BA TOKEN ĐI QUA MỘT LỚP TRANSFORMER ===")
print("MỤC TIÊU: giải thích token lấy thông tin ở đâu, tạo điểm dự đoán ra sao và dùng lại K/V thế nào.")
print(f"ĐẦU VÀO: {DEMO_PROMPT}, ba mã giả lập trong bảng từ vựng 0..5, không phải ba từ tiếng Việt.")
print("CÁCH HỌC: đọc câu Quan sát, xem số liệu rồi thử tự nói lại Kết luận của từng bước 1A..4C.")
print("Lab chọn MHA: mỗi Q head có K/V riêng, cùng FFN Dense SwiGLU và RoPE trên Q/K.")
print("Chỉ một chuỗi (B=1, ẩn trong shape), một lớp NumPy; weights chưa huấn luyện.")
print("Các số kiểm chứng cơ chế, không chứng minh khả năng ngôn ngữ hay tốc độ GPU.")
print("THỨ TỰ THỰC THI theo phụ thuộc số học (không theo cột vẽ):")
print("  1A → 1B → 3A/3B → 2A → 2B → 2C → 2D → 2E → 4A → 4B → 4C")
print("  Cột 3 là nhánh HỖ TRỢ attention: cache phải sẵn sàng TRƯỚC khi 2B/2E dùng K,V.")
print("")

# Prefill: in toàn bộ câu chuyện 1A..4C với giá trị thật.
prefill_logits, cache, prefill_shapes = forward(DEMO_PROMPT, trace=True)
assert prefill_shapes["logits"] == (len(DEMO_PROMPT), VOCAB)
print("")
print(f"argmax hàng cuối prefill → token demo {int(np.argmax(prefill_logits[-1]))} "
      "(chọn bằng số, không phải chất lượng ngôn ngữ).")
demo_token = int(np.argmax(prefill_logits[-1]))
print("")

# Decode đúng 1 token, mượn cache vừa ghi ở prefill.
past_before = int(cache["k"].shape[1])
assert cache["k"].shape == (N_HEADS, past_before, D_HEAD)
history_prompt_prefix = {key: value.copy() for key, value in cache.items()}
st_dec = attention_traced([demo_token], cache)
print(f"--- DECODE: đưa token vừa chọn ({demo_token}) trở lại mạng với cache {past_before} token ---")
print("Token thứ tư đã được chọn từ hàng cuối prefill. Lượt này xử lý nó để dự đoán token thứ năm.")
_print_pass(st_dec, "Decode")

# Logits theo từng vị trí của prompt + điểm next-token từ hàng cuối.
print(f"shape logits prefill: {prefill_logits.shape} — hàng cuối là next-token scores:")
print("  " + fmt_vec(prefill_logits[-1]))

# ── Kiểm tra số học (các forward nội bộ im lặng) ────────────────────────────
full_next_logits, _, _ = forward(list(DEMO_PROMPT) + [demo_token])
error = float(np.max(np.abs(st_dec["logits"][0] - full_next_logits[-1])))
assert error <= 1e-9
assert int(st_dec["cache"]["k"].shape[1]) == past_before + 1
# Prefix cache bất biến trước khi decode.
assert np.array_equal(st_dec["cache"]["k"][:, :past_before], history_prompt_prefix["k"])
assert np.array_equal(st_dec["cache"]["v"][:, :past_before], history_prompt_prefix["v"])
# Deterministic: chạy lại cùng prompt cho cùng logits/cache.
again_logits, again_cache, _ = forward(DEMO_PROMPT)
assert np.array_equal(again_logits, prefill_logits)
assert np.array_equal(again_cache["k"], cache["k"]) and np.array_equal(again_cache["v"], cache["v"])
# Causal mask: xác suất tương lai đúng bằng 0, mỗi hàng tổng 1 và không âm.
st3 = attention_traced(DEMO_PROMPT)
P = st3["prob"]
future = ~st3["allow"]
assert np.all(P[:, future] == 0.0)
assert np.all(P >= 0.0)
assert np.allclose(P.sum(axis=-1), 1.0, atol=1e-12)
# Hàng đầu mỗi head chỉ được nhìn chính nó.
assert np.allclose(P[:, 0, 0], 1.0) and np.all(P[:, 0, 1:] == 0.0)
# Thay token tương lai không đổi logits trước đó.
base, _, _ = forward(DEMO_PROMPT)
changed, _, _ = forward([1, 2, 4])
assert np.allclose(base[:2], changed[:2], atol=1e-10)
# Decode nhiều token (vị trí khác 0) qua cache khớp tính lại toàn bộ — im lặng.
full_ids = list(DEMO_PROMPT) + [demo_token]
cache = st_dec["cache"]   # bắt đầu từ cache sau pass decode demo (4 token)
for token in DEMO_CONTINUATION:
    full_ids.append(token)
    full_logits, _, _ = forward(full_ids)
    cached_logits, cache, _ = forward([token], cache)
    step_error = float(np.max(np.abs(cached_logits[-1] - full_logits[-1])))
    assert step_error <= 1e-9, f"decode lệch full recompute tại token {token}: {step_error}"
    assert int(cache["k"].shape[1]) == len(full_ids)

print("\nTHỬ NGHIỆM: Token tương lai có làm đổi quá khứ không?")
print(f"Mốc so sánh: {DEMO_PROMPT}; prompt bạn thử: {TRY_PROMPT}")
trial_logits, _, _ = forward(TRY_PROMPT)
if len(TRY_PROMPT) == len(DEMO_PROMPT):
    changed_positions = np.flatnonzero(np.asarray(TRY_PROMPT) != np.asarray(DEMO_PROMPT))
    first_changed = int(changed_positions[0]) if len(changed_positions) else len(DEMO_PROMPT)
    row_errors = np.max(np.abs(trial_logits - prefill_logits), axis=1)
    for i, delta in enumerate(row_errors):
        expectation = "phải giữ nguyên" if i < first_changed else "được phép đổi"
        print(f"  Hàng {i}: lệch lớn nhất = {delta:.3e}; dự đoán: {expectation}.")
    assert np.all(row_errors[:first_changed] <= 1e-9)
    print("Kết luận: các vị trí trước token đầu tiên bị sửa không đọc được thay đổi đó.")
    print("Được phép đổi không có nghĩa là bắt buộc đổi: hãy phân biệt quy tắc mask với giá trị cụ thể.")
else:
    print("Để so từng hàng, hãy dùng TRY_PROMPT có cùng số token với DEMO_PROMPT.")

print("\nTỔNG KẾT: Bạn vừa quan sát được gì?")
print("1. Causal mask cho phép tính cả prompt mà token trước không nhìn token sau.")
print("2. Attention weights chọn nguồn thông tin; logits ở LM head chấm điểm token ứng viên.")
print(f"3. Dùng cache và tính lại cả chuỗi có độ lệch logits lớn nhất {error:.3e} (ngưỡng 1e-9).")
print("   Các vị trí cũ trong cache giữ nguyên; cùng weights và input cho cùng kết quả.")
print("Các kiểm tra này xác nhận trường hợp demo, không chứng minh model hiểu ngôn ngữ hoặc tăng tốc bao nhiêu.")
print("PASS: decoder một lớp MHA + RoPE + SwiGLU chạy đúng và khớp 12 mã overlay kiến trúc.")
