"""
Track 1 Lab C — Test: one KV append per layer per step (cache growth).

Verifies that each decode step appends exactly one K/V column per layer,
and that the usable cache shape follows [B, H_kv, P + S, D_HEAD].

Run: python test_cache_growth.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
NUM_DECODE_STEPS = 4

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

prompt = torch.tensor([[7, 3, 9]])   # [B=1, P=3]

# Prefill
logits_prefill, past = model(prompt, use_cache=True)
expected_cache_len = P
all_growth_ok = True

print(f"Cache length after prefill (P={P}): {past[0][0].shape[2]}")
assert past[0][0].shape[2] == expected_cache_len, \
    f"Expected {expected_cache_len}, got {past[0][0].shape[2]}"

print(f"\nPer-step cache growth check ({NUM_DECODE_STEPS} decode steps):\n")
print(f"{'step':>5} {'expected':>10} {'actual':>10} {'per_layer':>12} {'ok':>6}")
print("-" * 49)

next_id = logits_prefill[:, -1, :].argmax(dim=-1, keepdim=True)

for step in range(1, NUM_DECODE_STEPS + 1):
    logits_decode, past = model(next_id, past_key_values=past, use_cache=True)

    expected_len = P + step
    actual_len = past[0][0].shape[2]

    # Verify every layer has the same cache length.
    layer_lens = [p[0].shape[2] for p in past]
    per_layer_ok = all(l == expected_len for l in layer_lens)
    all_growth_ok = all_growth_ok and per_layer_ok and (actual_len == expected_len)

    print(f"{step:>5} {expected_len:>10} {actual_len:>10} {str(per_layer_ok):>12} {str(per_layer_ok):>6}")

    # Also verify cache shape follows [B, H_kv, P + step, D_HEAD]
    expected_shape = (B, H, expected_len, D_HEAD)
    actual_shape = past[0][0].shape
    assert actual_shape == expected_shape, \
        f"Step {step}: expected cache shape {expected_shape}, got {actual_shape}"

    next_id = logits_decode[:, -1, :].argmax(dim=-1, keepdim=True)

print("-" * 49)
print(f"Final cache length: {past[0][0].shape[2]} (expected {P + NUM_DECODE_STEPS})")
print(f"All layers grow by exactly one column per step: {all_growth_ok}")
assert all_growth_ok, "Cache growth mismatch on at least one step"
print("\nPASS: one KV append per layer per step; cache shape [B, H_kv, P+S, D_HEAD] confirmed.")