"""
Recurrent Depth × Context Length Interaction at Tiny Scale (v2 - memory fix)
"""
import torch
import torch.nn as nn
import torch.nn.functional as F
import math
import time
import json

# --- Config ---
D_MODEL = 256
N_HEADS = 4
VOCAB = 256
BATCH = 32
STEPS = 2000
LR = 3e-4
WARMUP = 200
DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
SEED = 42

# --- Data: pre-generate tensors (memory efficient) ---
def load_tinystories_bytes(max_train=50_000_000, max_val=2_000_000):
    with open('/work/data/tinystories.txt', 'rb') as f:
        raw = f.read()
    # Split into train/val: 95/5
    split = int(len(raw) * 0.95)
    train_bytes = raw[:split][:max_train]
    val_bytes = raw[split:][:max_val]
    
    # Pre-generate all (x, y) pairs as contiguous tensors
    def make_pairs(data, seq_len=256):
        n = len(data) // seq_len
        arr = torch.tensor(list(data[:n * seq_len + 1]), dtype=torch.long)
        x = arr[:-1].view(n, seq_len)
        y = arr[1:].view(n, seq_len)
        return x, y
    
    train_x, train_y = make_pairs(train_bytes)
    val_x, val_y = make_pairs(val_bytes)
    print(f'  Train: {train_x.shape[0]} sequences, Val: {val_x.shape[0]} sequences')
    return train_x, train_y, val_x, val_y

# --- Model ---
class TransformerBlock(nn.Module):
    def __init__(self, d_model, n_heads, d_ff):
        super().__init__()
        self.ln1 = nn.LayerNorm(d_model)
        self.attn = nn.MultiheadAttention(d_model, n_heads, batch_first=True)
        self.ln2 = nn.LayerNorm(d_model)
        self.ffn = nn.Sequential(
            nn.Linear(d_model, d_ff),
            nn.GELU(),
            nn.Linear(d_ff, d_model)
        )
    
    def forward(self, x):
        h = self.ln1(x)
        causal_mask = torch.ones(h.size(1), h.size(1), device=h.device, dtype=torch.bool).triu(1)
        attn_out, _ = self.attn(h, h, h, attn_mask=causal_mask, need_weights=False)
        x = x + attn_out
        x = x + self.ffn(self.ln2(x))
        return x

class TinyGPT(nn.Module):
    def __init__(self, d_model, n_heads, d_ff, n_blocks, recurrent=False, ctx_len=128):
        super().__init__()
        self.d_model = d_model
        self.ctx_len = ctx_len
        self.recurrent = recurrent
        self.n_blocks = n_blocks
        
        self.tok_emb = nn.Embedding(VOCAB, d_model)
        self.pos_emb = nn.Embedding(ctx_len, d_model)
        
        if recurrent:
            self.block = TransformerBlock(d_model, n_heads, d_ff)
        else:
            self.blocks = nn.ModuleList([
                TransformerBlock(d_model, n_heads, d_ff) for _ in range(n_blocks)
            ])
        
        self.ln_f = nn.LayerNorm(d_model)
        self.head = nn.Linear(d_model, VOCAB, bias=False)
        self.head.weight = self.tok_emb.weight
    
    def forward(self, idx, targets=None):
        B, T = idx.shape
        T = min(T, self.ctx_len)
        idx = idx[:, :T]
        if targets is not None:
            targets = targets[:, :T]
        
        pos = torch.arange(T, device=idx.device).unsqueeze(0)
        x = self.tok_emb(idx) + self.pos_emb(pos)
        
        if self.recurrent:
            for _ in range(self.n_blocks):
                x = self.block(x)
        else:
            for block in self.blocks:
                x = block(x)
        
        x = self.ln_f(x)
        logits = self.head(x)
        
        if targets is not None:
            loss = F.cross_entropy(logits.reshape(-1, VOCAB), targets.reshape(-1))
            return loss
        return logits

def count_params(model):
    return sum(p.numel() for p in model.parameters())

def train(model, train_x, train_y, val_x, val_y, steps=STEPS, batch=BATCH, lr=LR, warmup=WARMUP):
    model = model.to(DEVICE)
    optimizer = torch.optim.AdamW(model.parameters(), lr=lr, betas=(0.9, 0.95), weight_decay=0.01)
    
    def get_lr(step):
        if step < warmup:
            return lr * step / max(warmup, 1)
        return lr * (0.5 + 0.5 * math.cos(math.pi * (step - warmup) / max(steps - warmup, 1)))
    
    n_train = train_x.shape[0]
    model.train()
    losses = []
    val_losses = []
    best_val = float("inf")
    
    for step in range(steps):
        optimizer.param_groups[0]["lr"] = get_lr(step)
        
        idx = torch.randint(0, n_train, (batch,))
        x = train_x[idx].to(DEVICE)
        y = train_y[idx].to(DEVICE)
        
        loss = model(x, y)
        loss.backward()
        torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
        optimizer.step()
        optimizer.zero_grad()
        
        losses.append(loss.item())
        
        if (step + 1) % 200 == 0:
            model.eval()
            val_loss = 0.0
            val_count = 0
            with torch.no_grad():
                for i in range(0, val_x.shape[0], batch):
                    vx = val_x[i:i+batch].to(DEVICE)
                    vy = val_y[i:i+batch].to(DEVICE)
                    vl = model(vx, vy)
                    val_loss += vl.item()
                    val_count += 1
            avg_val = val_loss / max(val_count, 1)
            val_losses.append((step + 1, avg_val))
            if avg_val < best_val:
                best_val = avg_val
            model.train()
    
    # Final validation
    model.eval()
    val_loss = 0.0
    val_count = 0
    with torch.no_grad():
        for i in range(0, val_x.shape[0], batch):
            vx = val_x[i:i+batch].to(DEVICE)
            vy = val_y[i:i+batch].to(DEVICE)
            vl = model(vx, vy)
            val_loss += vl.item()
            val_count += 1
    final_val = val_loss / max(val_count, 1)
    
    return {
        "final_train_loss": losses[-1],
        "avg_train_loss_last_100": sum(losses[-100:]) / 100,
        "best_val_loss": best_val,
        "final_val_loss": final_val,
        "final_val_ppl": math.exp(final_val),
        "val_curve": val_losses,
    }

def main():
    torch.manual_seed(SEED)
    if torch.cuda.is_available():
        torch.cuda.manual_seed_all(SEED)
    
    print("=" * 60)
    print("Recurrent Depth × Context Length Interaction")
    print("=" * 60)
    print(f"Device: {DEVICE}")
    if torch.cuda.is_available():
        free, total = torch.cuda.mem_get_info()
        print(f"GPU: {free/1e9:.1f} GB free / {total/1e9:.1f} GB total")
    
    print("\nLoading data...")
    train_x, train_y, val_x, val_y = load_tinystories_bytes()
    
    configs = [
        ("recurrent_ctx64",  True,  4, 2048, 64),
        ("recurrent_ctx256", True,  4, 2048, 256),
        ("stacked_ctx64",    False, 4, 128,  64),
        ("stacked_ctx256",   False, 4, 128,  256),
    ]
    
    results = {}
    for name, recurrent, n_blocks, d_ff, ctx_len in configs:
        print(f"\n{'='*40}")
        print(f"Config: {name}")
        print(f"  recurrent={recurrent}, blocks={n_blocks}, d_ff={d_ff}, ctx={ctx_len}")
        
        model = TinyGPT(D_MODEL, N_HEADS, d_ff, n_blocks, recurrent=recurrent, ctx_len=ctx_len)
        n_params = count_params(model)
        print(f"  Params: {n_params:,}")
        
        t0 = time.time()
        res = train(model, train_x, train_y, val_x, val_y)
        elapsed = time.time() - t0
        res["params"] = n_params
        res["train_time_s"] = elapsed
        results[name] = res
        
        print(f"  Train loss (final): {res['final_train_loss']:.4f}")
        print(f"  Val PPL: {res['final_val_ppl']:.2f}")
        print(f"  Time: {elapsed:.1f}s")
        
        # Free GPU memory between configs
        del model
        torch.cuda.empty_cache()
    
    # Summary
    print("\n" + "=" * 60)
    print("SUMMARY")
    print("=" * 60)
    print(f"\n{'Config':<25} {'Params':>10} {'Val PPL':>10} {'Time':>8}")
    print("-" * 60)
    for name, res in results.items():
        print(f"{name:<25} {res['params']:>10,} {res['final_val_ppl']:>10.2f} {res['train_time_s']:>7.1f}s")
    
    # Interaction analysis
    print("\n--- Interaction Analysis ---")
    rec_64 = results["recurrent_ctx64"]["final_val_ppl"]
    rec_256 = results["recurrent_ctx256"]["final_val_ppl"]
    stk_64 = results["stacked_ctx64"]["final_val_ppl"]
    stk_256 = results["stacked_ctx256"]["final_val_ppl"]
    
    rec_gain = rec_64 - rec_256
    stk_gain = stk_64 - stk_256
    
    print(f"Recurrent: ctx64={rec_64:.2f} → ctx256={rec_256:.2f} (Δ={rec_gain:+.2f})")
    print(f"Stacked:   ctx64={stk_64:.2f} → ctx256={stk_256:.2f} (Δ={stk_gain:+.2f})")
    print(f"\nRecurrent advantage at ctx64:  {stk_64 - rec_64:+.2f}")
    print(f"Recurrent advantage at ctx256: {stk_256 - rec_256:+.2f}")
    
    if rec_gain > 0.5 and stk_gain < 0.3:
        print("\n>>> FINDING: Longer context disproportionately benefits recurrent models!")
    elif abs(rec_gain) < 0.3 and abs(stk_gain) < 0.3:
        print("\n>>> FINDING: Context length has minimal effect on either architecture.")
    elif rec_gain < 0 and stk_gain < 0:
        print("\n>>> FINDING: Longer context HURTS both architectures (attention dilution).")
    else:
        print(f"\n>>> Mixed: rec_gain={rec_gain:+.2f}, stk_gain={stk_gain:+.2f}")
    
    out = {
        "experiment": "recurrent_ctx_interaction",
        "question": "Does recurrent depth amplify the benefit of longer context at tiny scale?",
        "timestamp": time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime()),
        "config": {
            "d_model": D_MODEL,
            "n_heads": N_HEADS,
            "vocab": VOCAB,
            "batch": BATCH,
            "steps": STEPS,
            "lr": LR,
            "device": DEVICE,
        },
        "results": results,
        "analysis": {
            "recurrent_ctx_gain": rec_gain,
            "stacked_ctx_gain": stk_gain,
            "recurrent_advantage_ctx64": stk_64 - rec_64,
            "recurrent_advantage_ctx256": stk_256 - rec_256,
        }
    }
    
    out_path = "/work/work/recurrent_ctx_interaction_results.json"
    with open(out_path, "w") as f:
        json.dump(out, f, indent=2)
    print(f"\nSaved to {out_path}")

if __name__ == "__main__":
    main()

def preflight_case():
    """CPU-only correctness fixture using the real TinyGPT and shifted TinyStories bytes."""
    from pathlib import Path
    import torch
    source = Path('/work/data/tinystories.txt')
    with source.open('rb') as f:
        raw = f.read(130)
    if len(raw) != 130:
        raise ValueError('Need at least 130 real training bytes for causal preflight')
    tokens = torch.tensor(list(raw), dtype=torch.long)
    x = torch.stack((tokens[0:64], tokens[65:129]))
    shifted_y = torch.stack((tokens[1:65], tokens[66:130]))
    model = TinyGPT(d_model=32, n_heads=4, d_ff=64, n_blocks=2, recurrent=True, ctx_len=64).cpu()
    return model, x, shifted_y, []  # No dedicated memory-write modules; recurrence reuses TransformerBlock weights.