"""Çarpan-LoRA — AUBIN eğitiminde GPU verim yığınının ölçümü (T4, Gemma-4 12B 4-bit, LoRA r16).

Aynı veri (Kev decision-v7 train), aynı model ve LoRA; yalnız eğitim döngüsünün kurulumu değişir. Ölçü: saniyede YARARLI
(dolgu olmayan) eğitim tokenı ve adım süresi. Kurulumlar:
  naive   : bs=1, rastgele sıra, TÜM konumların logiti (262k sözlük), gradyan kontrol noktası
  +keep   : yalnız son konumun logiti (logits_to_keep=1) — cevap tek harf, diğer konumların logiti boşa hesap
  +bucket : uzunluk kovalı batch (dolgu ~0), bs=4
  +bs8    : aynı, bs=8 (keep'in açtığı bellekle)
    python carpan_lora.py --steps 12 --out carpan_lora.json
"""
import argparse, json, random, time
import torch
import torch.nn.functional as F
import kevdata
import kev_llm as KL


def run(S, pool, bs, keep, bucket, steps, warm, rng):
    plen = lambda it: len(it["state"]) + sum(len(t) for t in it["texts"]) + len(it["q"])
    if bucket:
        srt = sorted(pool, key=plen); B = max(bs * 8, 32)
        buckets = [srt[i:i + B] for i in range(0, len(srt), B)]
    opt = torch.optim.AdamW([p for p in S.m.parameters() if p.requires_grad], lr=1e-6)
    useful = tot = 0; dt = 0.0
    torch.cuda.reset_peak_memory_stats()
    for s in range(warm + steps):
        if bucket:
            bk = buckets[rng.randrange(len(buckets))]; batch = [bk[rng.randrange(len(bk))] for _ in range(bs)]
        else:
            batch = [pool[rng.randrange(len(pool))] for _ in range(bs)]
        ps = [S.prompt(it, list(range(len(it["keys"])))) for it in batch]
        enc = S.tok(ps, return_tensors="pt", padding=True, add_special_tokens=False).to(S.dev)
        torch.cuda.synchronize(); t0 = time.time()
        out = S.m(**enc, logits_to_keep=1) if keep else S.m(**enc)
        last = out.logits[:, -1].float()
        lg = torch.stack([last[:, ids].logsumexp(-1) for ids in S.letter_ids], -1)
        loss = sum(F.cross_entropy(lg[b, :len(it["keys"])].unsqueeze(0), torch.tensor([it["y"]], device=lg.device))
                   for b, it in enumerate(batch)) / bs
        opt.zero_grad(set_to_none=True); loss.backward(); opt.step()
        torch.cuda.synchronize()
        if s >= warm:
            dt += time.time() - t0; useful += int(enc["attention_mask"].sum()); tot += enc["attention_mask"].numel()
        del out, last, lg, loss
    return {"bs": bs, "keep": keep, "bucket": bucket, "steps": steps, "sec_per_step": round(dt / steps, 3),
            "useful_tok_per_s": round(useful / dt, 1), "pad_frac": round(1 - useful / tot, 3),
            "peak_gb": round(torch.cuda.max_memory_allocated() / 2 ** 30, 2)}


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--model", default="google/gemma-4-12B-it"); ap.add_argument("--work", default="/tmp/kev")
    ap.add_argument("--steps", type=int, default=12); ap.add_argument("--warm", type=int, default=2)
    ap.add_argument("--out", default="carpan_lora.json")
    a = ap.parse_args()
    K = kevdata.load(a.work)
    S = KL.Scorer(a.model, four_bit=True)
    S.tok.padding_side = "left"
    if S.tok.pad_token is None:
        S.tok.pad_token = S.tok.eos_token
    from peft import LoraConfig, get_peft_model, prepare_model_for_kbit_training
    S.m = prepare_model_for_kbit_training(S.m, use_gradient_checkpointing=True)
    S.m = get_peft_model(S.m, LoraConfig(r=16, lora_alpha=32, lora_dropout=0.05, task_type="CAUSAL_LM",
                                         target_modules=r"^(?!.*(vision|audio|siglip|embed_vision|embed_audio)).*\.(q_proj|k_proj|v_proj|o_proj|gate_proj|up_proj|down_proj)$"))
    S.m.train()
    pool = [x for x in KL.items(K["kev_train"]) if len(x["keys"]) <= 26]
    R = {"model": a.model, "gpu": torch.cuda.get_device_name(0), "configs": []}
    for name, bs, keep, bucket in [("naive", 1, False, False), ("+keep", 1, True, False),
                                   ("+keep+bucket bs4", 4, True, True), ("+keep+bucket bs8", 8, True, True),
                                   ("naive bs4 (pad, all logits)", 4, False, False)]:
        try:
            r = run(S, pool, bs, keep, bucket, a.steps, a.warm, random.Random(0)); r["name"] = name
        except torch.cuda.OutOfMemoryError as e:
            r = {"name": name, "oom": True}
        torch.cuda.empty_cache()
        R["configs"].append(r); print(json.dumps(r), flush=True)
        json.dump(R, open(a.out, "w"), indent=1)
    base = next(c for c in R["configs"] if c["name"] == "naive")["useful_tok_per_s"]
    for c in R["configs"]:
        if "useful_tok_per_s" in c:
            c["x_vs_naive"] = round(c["useful_tok_per_s"] / base, 2)
    json.dump(R, open(a.out, "w"), indent=1)
    print("CARPAN_LORA", json.dumps(R), flush=True)


if __name__ == "__main__":
    main()