#!/usr/bin/env python3
"""Little Fig GPU Training Benchmark — FP16 vs NF4 vs FigQuant on TinyLlama"""
import os, sys, subprocess, json, time, gc, traceback
import numpy as np
subprocess.check_call([sys.executable, "-m", "pip", "install", "-q", "transformers", "accelerate", "peft", "bitsandbytes", "datasets", "sentencepiece", "protobuf", "psutil", "numpy"])
if not os.path.exists("/app/littlefig"):
    subprocess.check_call(["git", "clone", "https://github.com/ticketguy/littlefig.git", "/app/littlefig"])
subprocess.check_call([sys.executable, "-m", "pip", "install", "-q", "-e", "/app/littlefig[train]"])
sys.path.insert(0, "/app/littlefig/src")
import torch
import torch.nn.functional as F

MODEL = "TinyLlama/TinyLlama-1.1B-Chat-v1.0"
LORA_R = 16; LORA_ALPHA = 32; LORA_TARGETS = ["q_proj","k_proj","v_proj","o_proj"]
MAX_SEQ = 512; TRAIN_STEPS = 100; BATCH_SIZE = 4; GRAD_ACCUM = 4; LR = 2e-4
RESULTS = {}

def log(msg): print(f"[BENCH] {msg}", flush=True)
def gpu_mb(): return torch.cuda.max_memory_allocated()/1e6 if torch.cuda.is_available() else 0
def reset(): gc.collect(); torch.cuda.empty_cache(); torch.cuda.reset_peak_memory_stats()

log(f"PyTorch {torch.__version__}, CUDA={torch.cuda.is_available()}")
if torch.cuda.is_available():
    log(f"GPU: {torch.cuda.get_device_name()} ({torch.cuda.get_device_properties(0).total_memory/1e9:.1f}GB)")

from datasets import load_dataset
ds = load_dataset("tatsu-lab/alpaca", split="train").select(range(1000))
log(f"Dataset: {len(ds)} examples")

def hf_loop(model, tokenizer, name):
    dev = next(model.parameters()).device
    def tok_fn(ex):
        inst=ex.get("instruction",""); inp=ex.get("input","").strip(); out=ex.get("output","")
        txt = f"### Instruction:\n{inst}\n\n### Input:\n{inp}\n\n### Response:\n{out}" if inp else f"### Instruction:\n{inst}\n\n### Response:\n{out}"
        e = tokenizer(txt, truncation=True, max_length=MAX_SEQ, padding="max_length")
        e["labels"] = e["input_ids"].copy(); return e
    td = ds.map(tok_fn, remove_columns=ds.column_names); td.set_format("torch")
    from torch.utils.data import DataLoader
    dl = DataLoader(td, batch_size=BATCH_SIZE, shuffle=True, collate_fn=lambda b: {k:torch.stack([x[k] for x in b]) for k in b[0] if isinstance(b[0][k], torch.Tensor)}, drop_last=True)
    opt = torch.optim.AdamW([p for p in model.parameters() if p.requires_grad], lr=LR, weight_decay=0.01)
    model.train(); losses=[]; gs=0; al=0.0; reset(); t0=time.time()
    for batch in dl:
        if gs>=TRAIN_STEPS*GRAD_ACCUM: break
        batch = {k:v.to(dev) for k,v in batch.items()}
        with torch.autocast("cuda", dtype=torch.bfloat16):
            loss = model(**batch).loss / GRAD_ACCUM
        loss.backward(); al+=loss.item(); gs+=1
        if gs%GRAD_ACCUM==0:
            torch.nn.utils.clip_grad_norm_([p for p in model.parameters() if p.requires_grad], 1.0)
            opt.step(); opt.zero_grad()
            s=gs//GRAD_ACCUM; losses.append(al); al=0.0
            if s%20==0: log(f"  [{name}] step={s} loss={losses[-1]:.4f}")
    tt=time.time()-t0; pm=gpu_mb()
    del model, opt; gc.collect(); torch.cuda.empty_cache()
    RESULTS[name] = {"final":float(losses[-1]) if losses else None, "time_s":tt, "steps":len(losses), "gpu_mb":pm, "losses":[float(l) for l in losses]}
    log(f"  [{name}] DONE: loss={losses[-1]:.4f} time={tt:.0f}s gpu={pm:.0f}MB")

# ── FP16 LoRA ──
log("\n=== FP16 LoRA ===")
from transformers import AutoModelForCausalLM, AutoTokenizer
from peft import LoraConfig, get_peft_model
reset()
m = AutoModelForCausalLM.from_pretrained(MODEL, torch_dtype=torch.float16, device_map="auto")
t = AutoTokenizer.from_pretrained(MODEL); t.pad_token=t.eos_token
m.gradient_checkpointing_enable(gradient_checkpointing_kwargs={"use_reentrant":False})
m = get_peft_model(m, LoraConfig(r=LORA_R, lora_alpha=LORA_ALPHA, target_modules=LORA_TARGETS, bias="none", task_type="CAUSAL_LM"))
hf_loop(m, t, "fp16_lora")

# ── BnB NF4 ──
log("\n=== BnB NF4 QLoRA ===")
from transformers import BitsAndBytesConfig
from peft import prepare_model_for_kbit_training
reset()
m = AutoModelForCausalLM.from_pretrained(MODEL, quantization_config=BitsAndBytesConfig(load_in_4bit=True, bnb_4bit_quant_type="nf4", bnb_4bit_use_double_quant=True, bnb_4bit_compute_dtype=torch.bfloat16), device_map="auto")
t = AutoTokenizer.from_pretrained(MODEL); t.pad_token=t.eos_token
m = prepare_model_for_kbit_training(m)
m = get_peft_model(m, LoraConfig(r=LORA_R, lora_alpha=LORA_ALPHA, target_modules=LORA_TARGETS, bias="none", task_type="CAUSAL_LM"))
hf_loop(m, t, "bnb_nf4")

# ── FigQuant ──
log("\n=== FigQuant LoRA ===")
from little_fig.engine import FigModel
from little_fig.engine.tier import TrainingTier
reset()
model = FigModel.from_pretrained(MODEL, lora_r=LORA_R, lora_alpha=LORA_ALPHA, tier=TrainingTier.STREAMING_LORA, target_modules=LORA_TARGETS)
tok = model.tokenizer
examples = [dict(r) for r in ds]
def tok_fn(ex):
    inst=ex.get("instruction",""); inp=ex.get("input","").strip(); out=ex.get("output","")
    txt = f"### Instruction:\n{inst}\n\n### Input:\n{inp}\n\n### Response:\n{out}" if inp else f"### Instruction:\n{inst}\n\n### Response:\n{out}"
    e = tok(txt, truncation=True, max_length=MAX_SEQ, padding="max_length")
    return {"input_ids": e["input_ids"], "labels": e["input_ids"].copy(), "attention_mask": e["attention_mask"]}
tokenized = [tok_fn(ex) for ex in examples]
class DS(torch.utils.data.Dataset):
    def __init__(s, d): s.d=d
    def __len__(s): return len(s.d)
    def __getitem__(s, i): return {k:torch.tensor(v, dtype=torch.long) for k,v in s.d[i].items()}
from torch.utils.data import DataLoader
dl = DataLoader(DS(tokenized), batch_size=BATCH_SIZE, shuffle=True, drop_last=True)
dev = torch.device("cuda"); model = model.to(dev)
params = model.get_trainable_parameters()
opt = torch.optim.AdamW(params, lr=LR, weight_decay=0.01)
model.model.train(); losses=[]; gs=0; al=0.0; reset(); t0=time.time()
for batch in dl:
    if gs>=TRAIN_STEPS*GRAD_ACCUM: break
    batch = {k:v.to(dev) for k,v in batch.items()}
    with torch.autocast("cuda", dtype=torch.bfloat16):
        loss = model(input_ids=batch["input_ids"], attention_mask=batch["attention_mask"], labels=batch["labels"]).loss / GRAD_ACCUM
    loss.backward(); al+=loss.item(); gs+=1
    if gs%GRAD_ACCUM==0:
        torch.nn.utils.clip_grad_norm_(params, 1.0); opt.step(); opt.zero_grad()
        s=gs//GRAD_ACCUM; losses.append(al); al=0.0
        if s%20==0: log(f"  [figquant] step={s} loss={losses[-1]:.4f}")
tt=time.time()-t0; pm=gpu_mb()
RESULTS["figquant"] = {"final":float(losses[-1]) if losses else None, "time_s":tt, "steps":len(losses), "gpu_mb":pm}
log(f"  [figquant] DONE: loss={losses[-1]:.4f} time={tt:.0f}s gpu={pm:.0f}MB")
del model, opt; gc.collect(); torch.cuda.empty_cache()

# ── Summary ──
log("\n" + "="*60)
log("  RESULTS: TinyLlama 1.1B, 100 steps, batch=4x4, LoRA r=16")
log("="*60)
log(f"  {'Method':>12} {'Loss':>8} {'Time':>7} {'GPU MB':>8}")
log(f"  {'─'*40}")
for k in ["fp16_lora","bnb_nf4","figquant"]:
    if k in RESULTS:
        r=RESULTS[k]
        log(f"  {k:>12} {r['final']:.4f} {r['time_s']:.0f}s {r['gpu_mb']:.0f}")
log("="*60)
with open("/app/results.json","w") as f: json.dump(RESULTS, f, indent=2, default=str)