"""Find the optimal training config: batch, seq, and fused CE.

The bottleneck is the logits tensor: [B, T, 50257] = 1.6 GB at batch=16, seq=1024.
The cross-entropy + backward through this is expensive.

Optimizations to test:
  1. Larger batch + shorter seq (amortize weight reads)
  2. Longer seq + smaller batch (more tokens per forward)
  3. Fused cross-entropy (avoid materializing full logits)
  4. Gradient checkpointing (trade compute for memory, fit bigger batches)
"""
import time
import torch
import torch.nn.functional as F
from config import Config
from model import build_model

torch.backends.cuda.matmul.allow_tf32 = True
torch.backends.cudnn.allow_tf32 = True

cfg = Config.v5_500m()
model = build_model(cfg, "cuda")
model.train()

# Freeze medusa
for name, param in model.named_parameters():
    if name.startswith("medusa_"):
        param.requires_grad = False

opt = torch.optim.AdamW(
    [p for p in model.parameters() if p.requires_grad],
    lr=3e-4, betas=(0.9, 0.95), weight_decay=0.1, fused=True)

print(f"Model: {sum(p.numel() for p in model.parameters())/1e6:.1f}M params")
print(f"Trainable: {sum(p.numel() for p in model.parameters() if p.requires_grad)/1e6:.1f}M")
print()

# --- Standard CE vs Fused CE ---
def standard_ce(model, ids, cfg):
    h = model(ids)
    logits = model.lm_head(h)
    loss = F.cross_entropy(
        logits[:, :-1].reshape(-1, cfg.vocab_size),
        ids[:, 1:].reshape(-1))
    return loss

def fused_ce(model, ids, cfg):
    """Fused CE: use F.linear + F.cross_entropy without materializing full logits.
    
    Actually PyTorch's F.cross_entropy already fuses the softmax + NLL.
    The issue is the [B*T, V] logits tensor. We can't avoid it with standard CE.
    
    But we CAN use the 'chunked' approach: compute CE in vocab chunks.
    This trades compute for memory, letting us use bigger batches.
    """
    h = model(ids)  # [B, T, d]
    # Tied embeddings: lm_head weight = tok_emb weight = [V, d]
    weight = model.tok_emb.weight  # [V, d]
    # Reshape for CE: h_flat [B*T, d], targets [B*T]
    h_flat = h[:, :-1].reshape(-1, cfg.d_model)
    targets = ids[:, 1:].reshape(-1)
    # F.cross_entropy with weight = compute logits on the fly
    # Actually, F.cross_entropy(input, target) where input is [N, C]
    # still materializes [N, C] logits. 
    # The real fused version: use torch's internal fused kernel
    logits = F.linear(h_flat, weight)  # [B*T, V]
    loss = F.cross_entropy(logits, targets)
    return loss

def chunked_ce(model, ids, cfg, chunk_size=8192):
    """Chunked CE: compute loss in chunks to save memory.
    Allows much larger batch sizes.
    """
    h = model(ids)  # [B, T, d]
    weight = model.tok_emb.weight  # [V, d]
    h_flat = h[:, :-1].reshape(-1, cfg.d_model)  # [B*T, d]
    targets = ids[:, 1:].reshape(-1)  # [B*T]
    
    N = h_flat.shape[0]
    total_loss = 0.0
    for i in range(0, N, chunk_size):
        chunk_h = h_flat[i:i+chunk_size]
        chunk_t = targets[i:i+chunk_size]
        chunk_logits = F.linear(chunk_h, weight)
        total_loss = total_loss + F.cross_entropy(chunk_logits, chunk_t) * chunk_h.shape[0]
    return total_loss / N

# --- Test configs ---
configs = [
    # (batch, seq, ce_fn, label)
    (16, 1024, standard_ce, "standard CE"),
    (16, 1024, fused_ce, "fused CE (same speed)"),
    (16, 1024, chunked_ce, "chunked CE (8k)"),
    (32, 1024, standard_ce, "batch=32"),
    (32, 1024, chunked_ce, "batch=32 chunked"),
    (64, 1024, chunked_ce, "batch=64 chunked"),
    (16, 2048, standard_ce, "seq=2048"),
    (16, 2048, chunked_ce, "seq=2048 chunked"),
    (32, 2048, chunked_ce, "batch=32 seq=2048 chunked"),
    (8, 4096, chunked_ce, "seq=4096 chunked"),
    (16, 4096, chunked_ce, "batch=16 seq=4096 chunked"),
]

print(f"{'config':>30} {'tokens':>8} {'step(ms)':>9} {'tok/s':>10} {'VRAM(GB)':>9}")
print("-" * 70)

for batch, seq, ce_fn, label in configs:
    try:
        ids = torch.randint(0, cfg.vocab_size, (batch, seq), device="cuda")
        
        # Warmup
        for _ in range(3):
            opt.zero_grad(set_to_none=True)
            loss = ce_fn(model, ids, cfg)
            loss.backward()
            opt.step()
        torch.cuda.synchronize()
        
        # Measure
        torch.cuda.reset_peak_memory_stats()
        times = []
        for _ in range(5):
            torch.cuda.synchronize()
            t0 = time.perf_counter()
            opt.zero_grad(set_to_none=True)
            loss = ce_fn(model, ids, cfg)
            loss.backward()
            opt.step()
            torch.cuda.synchronize()
            times.append(time.perf_counter() - t0)
        times.sort()
        step_ms = times[2] * 1000
        tokens = batch * seq
        tok_s = tokens / (step_ms / 1000)
        vram = torch.cuda.max_memory_allocated() / 1e9
        print(f"{label:>30} {tokens:>8} {step_ms:>9.1f} {tok_s:>10,.0f} {vram:>9.2f}")
    except RuntimeError as e:
        if "out of memory" in str(e).lower():
            print(f"{label:>30} {batch*seq:>8} {'OOM':>9}")
            torch.cuda.empty_cache()
        else:
            raise