"""
Track 1 Lab C — Test: raw MHA vs GQA cache bytes from the lab model dimensions.

Uses the lab model's dimensions (L, H_q, H_kv, D_HEAD) and computes raw
KV-cache bytes for MHA and GQA at a stated sequence length and batch size.

Per-layer formula: 2 * H_kv * D_HEAD * S * B * b
Total formula:     L * per-layer
(b = bytes per element)

The GQA figures are calculated from H_kv=1, not measured from an instantiated
GQA cache. The MHA model is instantiated only to read its dimensions.

Run: python test_mha_gqa_memory.py
"""

from __future__ import annotations

from cache_model import (
    TinyDecoder,
    kv_cache_bytes_per_layer,
    mha_gqa_cache_bytes,
    model_summary,
    set_seed,
)

set_seed(42)

# ---- Lab dimensions (read from instantiated model configs) ----
D, H, D_HEAD, L, V = 32, 2, 16, 2, 50
B = 1
b = 2            # FP16 bytes per element
S = 64           # cached tokens (prompt + generated)

mha_model = TinyDecoder(D=D, H=H, D_HEAD=D_HEAD, L=L, V=V)
gqa_model = TinyDecoder(D=D, H=H, D_HEAD=D_HEAD, L=L, V=V, use_gqa=True, H_KV=1)

mha_dims = model_summary(mha_model)
gqa_dims = model_summary(gqa_model)

# MHA: H_kv = H_q = H. GQA: H_kv = 1, H_q unchanged.
mha_per_layer = kv_cache_bytes_per_layer(
    H_kv=mha_dims["H_kv"], D_HEAD=D_HEAD, S=S, B=B, b=b
)
gqa_per_layer = kv_cache_bytes_per_layer(
    H_kv=gqa_dims["H_kv"], D_HEAD=D_HEAD, S=S, B=B, b=b
)

mha_total = mha_gqa_cache_bytes(
    L=L, H_kv=mha_dims["H_kv"], D_HEAD=D_HEAD, S=S, B=B, b=b
)
gqa_total = mha_gqa_cache_bytes(
    L=L, H_kv=gqa_dims["H_kv"], D_HEAD=D_HEAD, S=S, B=B, b=b
)

assert mha_per_layer == 8_192, mha_per_layer
assert gqa_per_layer == 4_096, gqa_per_layer
assert mha_total == L * mha_per_layer
assert gqa_total == L * gqa_per_layer

ratio = gqa_total / mha_total if mha_total > 0 else 0.0

print("Raw KV-cache bytes from the lab model dimensions\n")
print(f"  L (layers)     = {L}")
print(f"  H_q (Q heads)  = {mha_dims['H_q']}")
print(f"  D_HEAD         = {D_HEAD}")
print(f"  S (cached tok) = {S}")
print(f"  B (batch)      = {B}")
print(f"  b (bytes/elem) = {b} (FP16)")
print()
print(
    f"{'Variant':>10} {'H_kv':>6} {'Per-layer bytes':>18} "
    f"{'Total (' + str(L) + ' layers)':>18} {'Ratio vs MHA':>14}"
)
print("-" * 70)
print(f"{'MHA':>10} {mha_dims['H_kv']:>6} {mha_per_layer:>18,} {mha_total:>18,} {'1.00':>14}")
print(f"{'GQA':>10} {gqa_dims['H_kv']:>6} {gqa_per_layer:>18,} {gqa_total:>18,} {ratio:>14.2f}")
print("-" * 70)
print()
print(f"MHA: {mha_per_layer:,} bytes per layer; {mha_total:,} bytes total across {L} layers")
print(f"GQA: {gqa_per_layer:,} bytes per layer; {gqa_total:,} bytes total across {L} layers")
print("GQA figures are calculated from H_kv=1, not measured from an instantiated cache.")
print(f"GQA / MHA ratio:     {ratio:.2f} ({int((1 - ratio) * 100)}% reduction)")
print()
print("Note: cache-size ratios only. Total model memory is dominated by weights,")
print("and latency depends on kernel fusion, bandwidth, and compute vs memory bound.")
print("This is raw tensor bytes, not GPU allocator memory.")
print()
print("PASS: per-layer and total byte figures are consistent.")