"""
Track 1 Lab C — Test: broken cache-position case with diagnostic test and fix.

Demonstrates a broken cache write that overwrites slot 0 instead of appending
the new K/V at `write_position` (the current cache length). The cache never
grows, and each new token replaces the oldest entry, so subsequent reads attend
to the wrong keys/values and produce silently wrong logits.

The test:
  1. Runs cached decode with broken_write=True  -> logits should NOT match uncached
  2. Runs cached decode with broken_write=False -> logits should match uncached (fix)

Run: python test_broken_cache_position.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 = 3
TOL = 1e-5

# ---- Build two models: broken and fixed ----
broken_model = TinyDecoder(D=D, H=H, D_HEAD=D_HEAD, L=L, V=V)
fixed_model = TinyDecoder(D=D, H=H, D_HEAD=D_HEAD, L=L, V=V)
broken_model.eval()
fixed_model.eval()
torch.set_grad_enabled(False)

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

print("=" * 60)
print("BROKEN cache-position case (new K/V overwrites slot 0)")
print("=" * 60)

# ---- Broken path ----
logits_prefill_b, past_b = broken_model(prompt, use_cache=True)
next_id_b = logits_prefill_b[:, -1, :].argmax(dim=-1, keepdim=True)
generated_b = [next_id_b]

broken_diffs = []
for step in range(1, NUM_DECODE_STEPS + 1):
    logits_decode_b, past_b = broken_model(
        next_id_b, past_key_values=past_b, broken_write=True, use_cache=True
    )
    full_ids_b = torch.cat([prompt] + generated_b, dim=1)
    logits_full_b, _ = broken_model(full_ids_b, use_cache=False)
    diff_b = (logits_decode_b[:, -1, :] - logits_full_b[:, -1, :]).abs().max().item()
    broken_diffs.append(diff_b)
    print(f"  step {step}: max|cached - uncached| = {diff_b:.3e}  (tolerance {TOL}) -> "
          f"{'MISMATCH (expected)' if diff_b >= TOL else 'match'}")
    next_id_b = logits_decode_b[:, -1, :].argmax(dim=-1, keepdim=True)
    generated_b.append(next_id_b)

all_broken = all(d >= TOL for d in broken_diffs)
print(f"\n  All {NUM_DECODE_STEPS} steps show mismatch (expected for broken cache): {all_broken}")
assert all_broken, "Broken case unexpectedly matched — bug not triggered"

# ---- Fixed path ----
print("\n" + "=" * 60)
print("FIXED cache-position case (correct write at write_position)")
print("=" * 60)

logits_prefill_f, past_f = fixed_model(prompt, use_cache=True)
next_id_f = logits_prefill_f[:, -1, :].argmax(dim=-1, keepdim=True)
generated_f = [next_id_f]

fixed_diffs = []
all_fixed_ok = True
for step in range(1, NUM_DECODE_STEPS + 1):
    logits_decode_f, past_f = fixed_model(
        next_id_f, past_key_values=past_f, broken_write=False, use_cache=True
    )
    full_ids_f = torch.cat([prompt] + generated_f, dim=1)
    logits_full_f, _ = fixed_model(full_ids_f, use_cache=False)
    diff_f = (logits_decode_f[:, -1, :] - logits_full_f[:, -1, :]).abs().max().item()
    fixed_diffs.append(diff_f)
    all_fixed_ok = all_fixed_ok and (diff_f < TOL)
    print(f"  step {step}: max|cached - uncached| = {diff_f:.3e}  (tolerance {TOL}) -> "
          f"{'OK' if diff_f < TOL else 'FAIL'}")
    next_id_f = logits_decode_f[:, -1, :].argmax(dim=-1, keepdim=True)
    generated_f.append(next_id_f)

print(f"\n  All {NUM_DECODE_STEPS} steps within tolerance: {all_fixed_ok}")
assert all_fixed_ok, "Fixed case failed parity — fix not working"

print("\n" + "=" * 60)
print("DIAGNOSIS AND REPAIR")
print("=" * 60)
print("Bug: Each new K/V pair was written to cache slot 0 instead of being")
print("     appended at write_position (the current cache length). The cache")
print("     never grew, and the newest token overwrote the oldest entry, so")
print("     later reads attended to a corrupted key/value history.")
print()
print("Fix: Write new K/V at slot `k_cache.shape[2]` (exactly the current")
print("     cache length), i.e. the next free position. After the fix, cached")
print("     and uncached logits agree within FP32 tolerance.")
print()
print("PASS: broken cache-position detected and repaired.")