"""125M/110M pretrening LM na ROCm/7900XTX (bf16). Resumable, checkpointowany, logowany.
Uruchamiaj z `python -u` w środowisku z PyTorch (trenowane na ROCm).

Spec uzgodniony (forum 2026-08-10): d768/12L/12H, vocab=32000, block=512, bf16, wd=0.1, 1 epoka.
Dane: --data = .bin uint16 tokenow (mmap), najlepiej na szybkim dysku lokalnym.
Logi: <log-dir>/<run-id>/{train.log, metrics.jsonl}.
Checkpointy: ./ckpt-<run-id>/ (duze pliki).

Przyklad (Arek odpala live):
  python -u train_125m.py --data <data-dir>/corpus_32k.bin
"""
import argparse, json, math, os, time, glob, datetime
import numpy as np, torch, torch.nn as nn, torch.nn.functional as F

class Block(nn.Module):
    def __init__(s, d, nh):
        super().__init__(); s.nh = nh
        s.ln1 = nn.LayerNorm(d); s.ln2 = nn.LayerNorm(d)
        s.qkv = nn.Linear(d, 3*d); s.proj = nn.Linear(d, d)
        s.fc = nn.Linear(d, 4*d); s.fc2 = nn.Linear(4*d, d)
    def forward(s, x):
        B, T, D = x.shape
        q, k, v = s.qkv(s.ln1(x)).split(D, 2)
        q = q.view(B, T, s.nh, D//s.nh).transpose(1, 2)
        k = k.view(B, T, s.nh, D//s.nh).transpose(1, 2)
        v = v.view(B, T, s.nh, D//s.nh).transpose(1, 2)
        y = F.scaled_dot_product_attention(q, k, v, is_causal=True)
        y = y.transpose(1, 2).contiguous().view(B, T, D)
        x = x + s.proj(y)
        x = x + s.fc2(F.gelu(s.fc(s.ln2(x))))
        return x

class GPT(nn.Module):
    def __init__(s, vocab, d, nl, nh, block):
        super().__init__()
        s.tok = nn.Embedding(vocab, d); s.pos = nn.Embedding(block, d)
        s.blocks = nn.ModuleList([Block(d, nh) for _ in range(nl)])
        s.lnf = nn.LayerNorm(d); s.head = nn.Linear(d, vocab, bias=False)
        s.tok.weight = s.head.weight  # weight tying (GPT-2 standard, 110M)
    def forward(s, idx, tgt=None):
        B, T = idx.shape
        x = s.tok(idx) + s.pos(torch.arange(T, device=idx.device))[None]
        for b in s.blocks: x = b(x)
        logits = s.head(s.lnf(x))
        loss = None if tgt is None else F.cross_entropy(logits.view(-1, logits.size(-1)), tgt.view(-1))
        return logits, loss

def get_batch(data, block, B, dev):
    ix = torch.randint(len(data) - block - 1, (B,))
    x = torch.stack([torch.from_numpy(data[i:i+block].astype(np.int64)) for i in ix])
    y = torch.stack([torch.from_numpy(data[i+1:i+1+block].astype(np.int64)) for i in ix])
    return x.to(dev, non_blocking=True), y.to(dev, non_blocking=True)

def lr_at(step, warmup, total, lr):
    if step < warmup: return lr * (step+1) / warmup
    r = (step - warmup) / max(1, total - warmup)
    return 0.1*lr + 0.5*(0.9*lr) * (1 + math.cos(math.pi * min(1.0, r)))

def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--data", required=True); ap.add_argument("--vocab", type=int, default=32000)
    ap.add_argument("--d", type=int, default=768); ap.add_argument("--nl", type=int, default=12); ap.add_argument("--nh", type=int, default=12)
    ap.add_argument("--block", type=int, default=512); ap.add_argument("--batch", type=int, default=64)
    ap.add_argument("--lr", type=float, default=3e-4); ap.add_argument("--wd", type=float, default=0.1)
    ap.add_argument("--warmup-frac", type=float, default=0.01); ap.add_argument("--grad-clip", type=float, default=1.0)
    ap.add_argument("--epochs", type=float, default=1.0); ap.add_argument("--total-tokens", type=float, default=0.0)
    ap.add_argument("--log-dir", default="./runs")
    ap.add_argument("--ckpt-every", type=int, default=2000); ap.add_argument("--log-every", type=int, default=20)
    ap.add_argument("--accum-steps", type=int, default=1)  # grad-accum: microbatch = batch//accum; --batch = EFEKTYWNY batch (LR/total_steps/licznik liczone z batch, nie micro)
    ap.add_argument("--run-id", default=datetime.datetime.now().strftime("%Y%m%d_%H%M%S"))
    a = ap.parse_args()
    dev = "cuda"; assert torch.cuda.is_available(), "ROCm/cuda niedostepny — source env.sh?"
    rundir = os.path.join(a.log_dir, a.run_id); os.makedirs(rundir, exist_ok=True)
    ckpt_dir = f"./ckpt-{a.run_id}"; os.makedirs(ckpt_dir, exist_ok=True)
    logf = open(os.path.join(rundir, "train.log"), "a", buffering=1)
    metf = open(os.path.join(rundir, "metrics.jsonl"), "a", buffering=1)
    def log(msg):
        line = f"[{datetime.datetime.now().strftime('%H:%M:%S')}] {msg}"
        print(line, flush=True); logf.write(line + "\n")

    data = np.memmap(a.data, dtype=np.uint16, mode="r")
    tok_per_step = a.batch * a.block
    total_tokens = a.total_tokens if a.total_tokens > 0 else a.epochs * len(data)
    total_steps = int(total_tokens / tok_per_step)
    warmup = max(1, int(a.warmup_frac * total_steps))
    log(f"run={a.run_id} dev={torch.cuda.get_device_name(0)} data={len(data):,}tok "
        f"cfg=d{a.d}/{a.nl}L/{a.nh}H/vocab{a.vocab}/block{a.block}/B{a.batch} "
        f"steps={total_steps:,} tok/step={tok_per_step} wd={a.wd} warmup={warmup} ckpt={ckpt_dir}")

    model = GPT(a.vocab, a.d, a.nl, a.nh, a.block).to(dev)
    P = sum(p.numel() for p in model.parameters()); log(f"params={P:,}")
    opt = torch.optim.AdamW(model.parameters(), lr=a.lr, weight_decay=a.wd, betas=(0.9, 0.95), foreach=True)
    step0 = 0
    cks = sorted(glob.glob(os.path.join(ckpt_dir, "ckpt_*.pt")))
    if cks:
        ck = torch.load(cks[-1], map_location=dev, weights_only=True)
        model.load_state_dict(ck["model"]); opt.load_state_dict(ck["opt"]); step0 = ck["step"]
        log(f"RESUME z {cks[-1]} @ step {step0}")

    model.train(); t0 = time.time(); seen = 0
    for step in range(step0, total_steps):
        for g in opt.param_groups: g["lr"] = lr_at(step, warmup, total_steps, a.lr)
        micro = max(1, a.batch // a.accum_steps)
        opt.zero_grad(set_to_none=True)
        for _ in range(a.accum_steps):
            xb, yb = get_batch(data, a.block, micro, dev)
            with torch.autocast("cuda", dtype=torch.bfloat16):
                _, loss = model(xb, yb)
            (loss / a.accum_steps).backward()
        gnorm = nn.utils.clip_grad_norm_(model.parameters(), a.grad_clip); opt.step()
        seen += tok_per_step
        if (step+1) % a.log_every == 0:
            torch.cuda.synchronize(); dt = time.time() - t0; tps = seen / dt
            eta = (total_steps - step - 1) * tok_per_step / tps / 3600
            lv, lr_now = loss.item(), opt.param_groups[0]["lr"]
            log(f"step {step+1:>8,}/{total_steps:,} loss {lv:.4f} lr {lr_now:.2e} "
                f"grad {float(gnorm):.2f} tok/s {tps:,.0f} elapsed {dt/3600:.2f}h ETA {eta:.1f}h")
            metf.write(json.dumps({"step": step+1, "loss": lv, "lr": lr_now, "grad_norm": float(gnorm),
                                   "tok_s": round(tps), "elapsed_s": round(dt), "eta_h": round(eta, 2)}) + "\n")
        if (step+1) % a.ckpt_every == 0:
            p = os.path.join(ckpt_dir, f"ckpt_{step+1:08d}.pt")
            torch.save({"model": model.state_dict(), "opt": opt.state_dict(), "step": step+1, "cfg": vars(a)}, p)
            log(f"[ckpt] {p}")
    fp = os.path.join(ckpt_dir, f"ckpt_{total_steps:08d}.pt")
    if not os.path.exists(fp):  # final ckpt gdy total_steps nie trafia w ckpt_every (np. 41348 % 2000 != 0)
        torch.save({"model": model.state_dict(), "opt": opt.state_dict(), "step": total_steps, "cfg": vars(a)}, fp)
        log(f"[ckpt-final] {fp}")
    log("DONE")

if __name__ == "__main__": main()