"""
Byte-nivå LLM med färdiga Qwen3-arkitekturen. Kräver: pip install transformers torch numpy
Förberedelse: se tidigare — train.bin/val.bin som råa uint8-bytes.
"""
import math, os, time
import numpy as np
import torch
from transformers import Qwen3Config, Qwen3ForCausalLM

device = "cuda" if torch.cuda.is_available() else "cpu"
print(f"device: {device}")
os.makedirs("models", exist_ok=True)

# --- modell: Qwen3, rakt av, fast liten och med vocab=256 -------------
config = Qwen3Config(
    vocab_size=256,
    hidden_size=160,
    head_dim=24,
    num_attention_heads=4,
    num_key_value_heads=4,
    num_hidden_layers=5,
    intermediate_size=320,
    max_position_embeddings=1024,
    tie_word_embeddings=True,
    attention_dropout=0.0,
)

raw_model = Qwen3ForCausalLM(config).to(device)
raw_model.set_attn_implementation("sdpa")
print(f"parametrar: {sum(p.numel() for p in raw_model.parameters())/1e3:.1f}K")
print(f"attention: {raw_model.config._attn_implementation}")

# kompilerad variant för träning; raw_model för generate/checkpoint
model = raw_model

# --- sekventiell loader ------------------------------------------------
class SequentialLoader:
    def __init__(self, path, B, T):
        self.data = np.memmap(path, dtype=np.uint8, mode="r")
        self.B, self.T = B, T
        self.seg = (len(self.data) - 1) // B
        self.pos, self.epoch = 0, 0
    def next_batch(self):
        starts = np.arange(self.B) * self.seg + self.pos
        x = np.stack([self.data[s:s+self.T] for s in starts]).astype(np.int64)
        y = np.stack([self.data[s+1:s+1+self.T] for s in starts]).astype(np.int64)
        self.pos += self.T
        if self.pos + self.T + 1 > self.seg:
            self.pos, self.epoch = 0, self.epoch + 1
        return torch.from_numpy(x), torch.from_numpy(y)

B, T = 128, 1024
accum_steps = 2
base_lr = 3e-3
epochs = 10
train = SequentialLoader("data/train.bin", B, T)
steps_per_epoch = (train.seg // T) // accum_steps   # räknar OPTIMIZER-steg, inte mikrosteg
max_steps = epochs * steps_per_epoch
print(f"steg per epok: {steps_per_epoch:,} (effektiv batch: {B*accum_steps})")

# --- optimizer + lr-schema ---------------------------------------------
opt = torch.optim.AdamW(raw_model.parameters(), lr=base_lr,
                        betas=(0.9, 0.95), weight_decay=0.1,
                        fused=(device == "cuda"))
warmup = 1000
def lr_at(step):
    if step < warmup:
        return base_lr * (step + 1) / warmup
    t = (step - warmup) / max(1, max_steps - warmup)
    return base_lr / 10 + 0.5 * (base_lr - base_lr / 10) * (1 + math.cos(math.pi * min(t, 1.0)))


import json
import matplotlib.pyplot as plt

def save_losses(losses, path="models/losses"):
    # rådatan till fil
    with open(f"{path}.json", "w") as f:
        json.dump(losses, f)


# --- träningsloop --------------------------------------------------------
model.train()
t0 = time.time()
losses = []
opt.zero_grad(set_to_none=True)

for step in range(max_steps):
    for g in opt.param_groups:
        g["lr"] = lr_at(step)

    # --- ackumulera gradienter över accum_steps mikrobatcher ---
    for micro in range(accum_steps):
        x, y = train.next_batch()
        x = x.pin_memory().to(device, non_blocking=True)
        y = y.pin_memory().to(device, non_blocking=True)

        with torch.autocast(device_type=device, dtype=torch.bfloat16, enabled=device == "cuda"):
            out = model(input_ids=x, labels=x)   # HF räknar cross-entropy åt oss
            loss = out.loss / accum_steps        # skala ner så gradienterna summeras rätt

        loss.backward()
        losses.append(out.loss.item())           # logga oskalad loss

    torch.nn.utils.clip_grad_norm_(raw_model.parameters(), 1.0)
    opt.step()
    opt.zero_grad(set_to_none=True)

    if step % 50 == 0:
        tok_s = 50 * accum_steps * B * T / max(time.time() - t0, 1e-9)
        print(f"steg {step:6d} | epok {train.epoch} | loss {losses[-1]:.4f} "
              f"| lr {lr_at(step):.2e} | {tok_s/1e3:.0f}k tok/s")
        t0 = time.time()

    if step % 1000 == 0 and step > 0:
        save_losses(losses)
        torch.save({"model": raw_model.state_dict(), "step": step,
                    "pos": train.pos, "epoch": train.epoch}, "models/ckpt.pt")
        # smakprov — på raw_model, inte den kompilerade
        seed = torch.tensor(list("Det var en gång".encode()), dtype=torch.long)[None].to(device)
        with torch.no_grad():
            gen = raw_model.generate(seed, max_new_tokens=200,
                                     do_sample=True, temperature=0.8)
        print(bytes(gen[0].tolist()).decode("utf-8", errors="replace"))
        t0 = time.time()   # räkna inte in generate-tiden i tok/s

print("klart.")