"""causalscale LowRankGNN benchmark: measures VRAM scaling (the key 100M claim).
Uses train_lowrank_gnn's actual correlation-reconstruction objective.
"""
import time, json, numpy as np, sys
sys.stdout.reconfigure(encoding='utf-8')
# causalscale installed via pip install -e .
import torch, causalscale as cs

print(f"GPU: {torch.cuda.get_device_name(0)}, VRAM: {torch.cuda.get_device_properties(0).total_memory/1e9:.1f} GB")
print("=== causalscale LowRankGNN memory/timing benchmark ===\n")

results = []
for d, r in [(500, 32), (1000, 32), (2000, 16), (5000, 8), (10000, 4)]:
    n = min(2000, 2*d)
    print(f"d={d}, r={r}, n={n}... ", end='', flush=True)
    np.random.seed(42)
    X = np.random.randn(n, d).astype(np.float32)
    X = (X - X.mean(0)) / (X.std(0) + 1e-8)
    
    torch.cuda.empty_cache(); torch.cuda.reset_peak_memory_stats()
    t0 = time.time()
    try:
        result_d = cs.train_lowrank_gnn(X, rank=r, epochs=200, lr=3e-4, device='cuda', verbose=False)
        t1 = time.time()
        mem = torch.cuda.max_memory_allocated()/1e9
        f1 = result_d['f1']
        nedges = result_d['gnn_edges']
        gt_n = result_d['gt_edges']
        rv = {'d':d,'r':r,'n':n,'time_s':round(t1-t0,1),'vram_gb':round(mem,2),
              'edges':nedges,'f1':f1,'gt_edges':gt_n}
        results.append(rv)
        print(f"F1={f1:.3f}, {nedges}/{gt_n} edges, {t1-t0:.1f}s, {mem:.2f}GB")
    except torch.cuda.OutOfMemoryError:
        print("OOM"); results.append({'d':d,'r':r,'vram_gb':'OOM'}); break

out = r'results/memory_scaling.json'
json.dump(results, open(out,'w'), indent=2)

valid = [r for r in results if isinstance(r['vram_gb'],(int,float)) and r['vram_gb']>0]
if len(valid)>=2:
    # Model params: U (d*r) + V (d*r) = 2dr float32 params = 8dr bytes
    # Adam states: 2x params = 16dr bytes
    # Total model: ~24dr bytes + X data (n*d*4 bytes) + intermediates
    # Dominant: model params = 24*d*r bytes
    ds=[r['d'] for r in valid]; rs=[r['r'] for r in valid]
    model_bytes = [24*ds[i]*rs[i]/1e9 for i in range(len(valid))]
    data_bytes = [valid[i]['n']*ds[i]*4/1e9 for i in range(len(valid))]
    vrams = [r['vram_gb'] for r in valid]
    
    for i in range(len(valid)):
        print(f"  d={ds[i]}: model={model_bytes[i]:.2f}GB, data={data_bytes[i]:.2f}GB, actual={vrams[i]:.2f}GB")
    
    # VRAM = model + data + overhead. Overhead ~ 20-40% for intermediates
    # Extrapolation: d=1e8, r=4
    est_model = 24*1e8*4/1e9  # 9.6 GB
    est_data_1k = 1000*1e8*4/1e9  # 400 GB for n=1000
    est_data_500 = 500*1e8*4/1e9  # 200 GB
    
    print(f"\n=== Extrapolation to d=1e8 ===")
    print(f"d=1e8, r=4: model params = {est_model:.1f} GB")
    print(f"d=1e8, n=500: data = {est_data_500:.1f} GB")
    print(f"d=1e8, n=1000: data = {est_data_1000:.1f} GB")
    print(f"FP16 model: {est_model/2:.1f} GB")
    print(f"FP16 with gradient checkpoint: {est_model/4:.1f} GB")
    print(f"\nKey insight: data (n*d) dominates at d=1e8, not model params (d*r)")
    print(f"  With n=500 (streaming batches): feasible on A100 80GB")
    print(f"  With n=100: model+data ~ {est_model+100*1e8*4/1e9:.0f} GB on FP16")
    print(f"  Consumer 8GB GPU: requires batched processing + CPU offloading")
    
    time_slope = np.polyfit(ds, [r['time_s'] for r in valid], 1)[0] if len(valid)>=2 else 0
    if time_slope > 0:
        print(f"Time ~ {time_slope*1000:.1f} ms per 1000d")
        print(f"d=1e8: ~{time_slope*1e8/3600:.0f} hours (linear scaling)")

print(f"\nSaved {out}")