"""
Track 1 Lab C — Test: cached vs uncached logit parity over several decode steps.

Runs a tiny deterministic CPU model through:
  1. Prefill on a prompt (cached path produces initial cache)
  2. Several decode steps, each comparing cached-path logits against an
     uncached full-sequence reference at the last position
Asserts parity within FP32 tolerance (1e-5) for every decode step.

Run: python test_logit_parity.py
"""

from __future__ import annotations

import torch

from cache_model import set_seed, TinyDecoder

set_seed(42)

# ---- Lab dimensions ----
D, H, D_HEAD, L, V = 32, 2, 16, 2, 50
B, P = 1, 3            # batch size, prompt length
NUM_DECODE_STEPS = 5   # how many tokens to generate
TOL = 1e-5             # FP32 parity tolerance

model = TinyDecoder(D=D, H=H, D_HEAD=D_HEAD, L=L, V=V)
model.eval()
torch.set_grad_enabled(False)

# Prompt: P distinct token ids.
prompt = torch.tensor([[7, 3, 9]])   # [B=1, P=3]

# ---- Cached path: prefill ----
logits_prefill, past = model(prompt, use_cache=True)
cache_len_after_prefill = past[0][0].shape[2]
print(f"Cache length after prefill (P={P}): {cache_len_after_prefill}")
assert cache_len_after_prefill == P, f"Expected cache length {P}, got {cache_len_after_prefill}"

# The final prompt-position logits predict the first generated token.
next_id = logits_prefill[:, -1, :].argmax(dim=-1, keepdim=True)   # [B, 1]
generated = [next_id]   # accumulate all generated tokens

all_parity_ok = True

print(f"\nRunning {NUM_DECODE_STEPS} decode steps with cached-vs-uncached parity checks:\n")
print(f"{'step':>5} {'cache_len':>10} {'max|diff|':>12} {'parity_ok':>10}")
print("-" * 45)

for step in range(1, NUM_DECODE_STEPS + 1):
    # Cached decode step: feed next_id against the growing cache.
    logits_decode, past = model(next_id, past_key_values=past, use_cache=True)
    cache_len = past[0][0].shape[2]

    # Uncached reference: full sequence (prompt + all generated tokens so far)
    # in a single forward pass without a cache.
    full_ids = torch.cat([prompt] + generated, dim=1)   # [B, P + step]
    logits_full, _ = model(full_ids, use_cache=False)

    # Parity: last cached-decode logits == last uncached logits.
    diff = (logits_decode[:, -1, :] - logits_full[:, -1, :]).abs().max().item()
    ok = diff < TOL
    all_parity_ok = all_parity_ok and ok

    print(f"{step:>5} {cache_len:>10} {diff:>12.3e} {str(ok):>10}")

    # Prepare next token for the following step.
    next_id = logits_decode[:, -1, :].argmax(dim=-1, keepdim=True)
    generated.append(next_id)

print("-" * 45)
print(f"All {NUM_DECODE_STEPS} steps within tolerance {TOL}: {all_parity_ok}")
assert all_parity_ok, "Logit parity failed on at least one decode step"
print("\nPASS: cached and uncached next-token logits agree for every decode step.")