"""
SCUGNIZZ - Hugging Face Jobs edition

NOTE:
- Configurato per training su Hugging Face Jobs.
- Usa FineWeb in streaming.
- Parametri modello aumentati (12L / 768D / 12H).
- TARGET_TOKENS rappresenta un obiettivo logico di training.
- Per usare l'intero FineWeb è consigliabile eliminare il memmap e
  passare a un DataLoader streaming. Questa versione mantiene la
  struttura originale per ridurre le modifiche.
"""

!pip -q install datasets transformers huggingface_hub

import os
import math
import time
import random
import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
from contextlib import nullcontext
from datasets import load_dataset
from transformers import GPT2TokenizerFast

seed = 42
random.seed(seed)
np.random.seed(seed)
torch.manual_seed(seed)

if torch.cuda.is_available():
    torch.cuda.manual_seed_all(seed)

device = "cuda" if torch.cuda.is_available() else "cpu"
print("Device:", device)

if torch.cuda.is_available():
    torch.backends.cuda.matmul.allow_tf32 = True
    torch.backends.cudnn.allow_tf32 = True

# ==========================
# DATASET / TRAINING
# ==========================

TARGET_TOKENS = 10_000_000_000  # logical target for long-running HF Jobs

DATA_FILE = "fineweb_full_uint32.dat"
CKPT_FILE = "pcs_fineweb_checkpoint_last.pt"
BEST_FILE = "pcs_fineweb_checkpoint_best.pt"
FINAL_FILE = "pcs_fineweb_final.pt"

batch_size = 16
block_size = 1024

# circa 1 epoca sui 90M token di training
TOKENS_PER_STEP = batch_size * block_size
TRAIN_TOKENS = int(TARGET_TOKENS * 0.9)
max_iters = (TRAIN_TOKENS + TOKENS_PER_STEP - 1) // TOKENS_PER_STEP

print(f"Training tokens : {TRAIN_TOKENS:,}")
print(f"Token/step      : {TOKENS_PER_STEP:,}")
print(f"Iterazioni      : {max_iters:,} (~1 epoca)")
eval_interval = 500
save_interval = 1000
eval_iters = 20

learning_rate = 3e-4
min_lr = 3e-5
warmup_iters = 1000
weight_decay = 0.1
beta1 = 0.9
beta2 = 0.95
grad_clip = 1.0

n_embd = 768
n_head = 12
n_layer = 12
dropout = 0.1
bias = False

pcs_a = 0.8309193524478643
pcs_b = 0.0

tokenizer = GPT2TokenizerFast.from_pretrained("gpt2")
tokenizer.pad_token = tokenizer.eos_token
vocab_size = tokenizer.vocab_size
eos_id = tokenizer.eos_token_id

print("Vocab size:", vocab_size)

if not os.path.exists(DATA_FILE):
    print("Creating FineWeb memmap...")

    arr = np.memmap(
        DATA_FILE,
        dtype=np.uint32,
        mode="w+",
        shape=(TARGET_TOKENS,)
    )

    ds = load_dataset(
        "HuggingFaceFW/fineweb",
        name="CC-MAIN-2024-10",
        split="train",
        streaming=True
    )

    pos = 0
    last_print_million = -1

    for row in ds:
        txt = row["text"]

        if txt and len(txt) > 100:
            ids = tokenizer.encode(txt + tokenizer.eos_token)
            n = min(len(ids), TARGET_TOKENS - pos)

            if n > 0:
                arr[pos:pos+n] = np.array(ids[:n], dtype=np.uint32)
                pos += n

        cur_million = pos // 1_000_000
        if cur_million != last_print_million:
            print("Saved tokens:", pos)
            last_print_million = cur_million

        if pos >= TARGET_TOKENS:
            break

    arr.flush()
    print("Memmap created. Tokens written:", pos)

else:
    print("Memmap already exists:", DATA_FILE)

data = np.memmap(
    DATA_FILE,
    dtype=np.uint32,
    mode="r",
    shape=(TARGET_TOKENS,)
)

split_idx = int(0.9 * TARGET_TOKENS)
train_len = split_idx
val_len = TARGET_TOKENS - split_idx

print("Train tokens:", train_len)
print("Val tokens:", val_len)

def get_batch(split_name):
    if split_name == "train":
        lo = 0
        hi = train_len - block_size - 1
    else:
        lo = train_len
        hi = TARGET_TOKENS - block_size - 1

    ix = np.random.randint(lo, hi, size=(batch_size,))
    x = np.stack([data[i:i+block_size] for i in ix])
    y = np.stack([data[i+1:i+block_size+1] for i in ix])

    x = torch.tensor(x, dtype=torch.long, device=device)
    y = torch.tensor(y, dtype=torch.long, device=device)
    return x, y

def get_lr(it):
    if it < warmup_iters:
        return learning_rate * (it + 1) / warmup_iters

    if it > max_iters:
        return min_lr

    decay_ratio = (it - warmup_iters) / (max_iters - warmup_iters)
    coeff = 0.5 * (1.0 + math.cos(math.pi * decay_ratio))
    return min_lr + coeff * (learning_rate - min_lr)

class PCS(nn.Module):
    def __init__(self, a=pcs_a, b=pcs_b):
        super().__init__()
        self.a = a
        self.b = b

    def forward(self, x):
        return x * torch.sin(self.a * x) + self.b * torch.cos(x)

class CausalSelfAttention(nn.Module):
    def __init__(self, n_embd, n_head, block_size, dropout, bias=False):
        super().__init__()
        assert n_embd % n_head == 0

        self.n_head = n_head
        self.head_dim = n_embd // n_head

        self.q_proj = nn.Linear(n_embd, n_embd, bias=bias)
        self.k_proj = nn.Linear(n_embd, n_embd, bias=bias)
        self.v_proj = nn.Linear(n_embd, n_embd, bias=bias)
        self.out_proj = nn.Linear(n_embd, n_embd, bias=bias)

        self.attn_dropout = nn.Dropout(dropout)
        self.resid_dropout = nn.Dropout(dropout)

        mask = torch.tril(torch.ones(block_size, block_size))
        self.register_buffer("mask", mask.view(1, 1, block_size, block_size))

    def forward(self, x):
        B, T, C = x.shape

        q = self.q_proj(x).view(B, T, self.n_head, self.head_dim).transpose(1, 2)
        k = self.k_proj(x).view(B, T, self.n_head, self.head_dim).transpose(1, 2)
        v = self.v_proj(x).view(B, T, self.n_head, self.head_dim).transpose(1, 2)

        att = (q @ k.transpose(-2, -1)) / math.sqrt(self.head_dim)
        att = att.masked_fill(self.mask[:, :, :T, :T] == 0, float("-inf"))
        att = F.softmax(att, dim=-1)
        att = self.attn_dropout(att)

        y = att @ v
        y = y.transpose(1, 2).contiguous().view(B, T, C)
        y = self.resid_dropout(self.out_proj(y))
        return y

class MLP(nn.Module):
    def __init__(self, n_embd, dropout, bias=False):
        super().__init__()
        self.fc1 = nn.Linear(n_embd, 4 * n_embd, bias=bias)
        self.act = PCS()
        self.fc2 = nn.Linear(4 * n_embd, n_embd, bias=bias)
        self.dropout = nn.Dropout(dropout)

    def forward(self, x):
        x = self.fc1(x)
        x = self.act(x)
        x = self.fc2(x)
        x = self.dropout(x)
        return x

class Block(nn.Module):
    def __init__(self, n_embd, n_head, block_size, dropout, bias=False):
        super().__init__()
        self.ln1 = nn.LayerNorm(n_embd)
        self.attn = CausalSelfAttention(n_embd, n_head, block_size, dropout, bias=bias)
        self.ln2 = nn.LayerNorm(n_embd)
        self.mlp = MLP(n_embd, dropout, bias=bias)

    def forward(self, x):
        x = x + self.attn(self.ln1(x))
        x = x + self.mlp(self.ln2(x))
        return x

class GPT(nn.Module):
    def __init__(self):
        super().__init__()

        self.tok_emb = nn.Embedding(vocab_size, n_embd)
        self.pos_emb = nn.Embedding(block_size, n_embd)
        self.drop = nn.Dropout(dropout)

        self.blocks = nn.ModuleList([
            Block(n_embd, n_head, block_size, dropout, bias=bias)
            for _ in range(n_layer)
        ])

        self.ln_f = nn.LayerNorm(n_embd)
        self.lm_head = nn.Linear(n_embd, vocab_size, bias=False)

        self.tok_emb.weight = self.lm_head.weight
        self.apply(self._init_weights)

    def _init_weights(self, module):
        if isinstance(module, nn.Linear):
            nn.init.normal_(module.weight, mean=0.0, std=0.02)
            if module.bias is not None:
                nn.init.zeros_(module.bias)
        elif isinstance(module, nn.Embedding):
            nn.init.normal_(module.weight, mean=0.0, std=0.02)

    def forward(self, idx, targets=None):
        B, T = idx.shape
        assert T <= block_size

        pos = torch.arange(0, T, device=idx.device, dtype=torch.long)
        x = self.tok_emb(idx) + self.pos_emb(pos)
        x = self.drop(x)

        for block in self.blocks:
            x = block(x)

        x = self.ln_f(x)
        logits = self.lm_head(x)

        loss = None
        if targets is not None:
            loss = F.cross_entropy(
                logits.reshape(B * T, logits.size(-1)),
                targets.reshape(B * T)
            )

        return logits, loss

model = GPT().to(device)
print("Parameters (M):", sum(p.numel() for p in model.parameters()) / 1e6)

optimizer = torch.optim.AdamW(
    model.parameters(),
    lr=learning_rate,
    betas=(beta1, beta2),
    weight_decay=weight_decay
)

use_amp = (device == "cuda")
scaler = torch.cuda.amp.GradScaler(enabled=use_amp)

start_iter = 0
best_val = float("inf")

if os.path.exists(CKPT_FILE):
    print("Loading checkpoint...")
    ckpt = torch.load(CKPT_FILE, map_location=device)
    model.load_state_dict(ckpt["model"])
    optimizer.load_state_dict(ckpt["optimizer"])
    start_iter = ckpt["iter"] + 1
    best_val = ckpt.get("best_val", float("inf"))
    print("Resume from iter:", start_iter)
    print("Best val:", best_val)

@torch.no_grad()
def estimate_loss():
    out = {}
    model.eval()

    for split in ["train", "val"]:
        losses = torch.zeros(eval_iters)

        for _ in range(eval_iters):
            x, y = get_batch(split)
            ctx = torch.cuda.amp.autocast() if use_amp else nullcontext()

            with ctx:
                _, loss = model(x, y)

            losses[_] = loss.item()

        out[split] = losses.mean().item()

    model.train()
    return out

print("Starting training...")
t0 = time.time()

for it in range(start_iter, max_iters + 1):
    lr = get_lr(it)
    for param_group in optimizer.param_groups:
        param_group["lr"] = lr

    xb, yb = get_batch("train")

    ctx = torch.cuda.amp.autocast() if use_amp else nullcontext()
    with ctx:
        _, loss = model(xb, yb)

    optimizer.zero_grad(set_to_none=True)

    scaler.scale(loss).backward()
    scaler.unscale_(optimizer)
    torch.nn.utils.clip_grad_norm_(model.parameters(), grad_clip)
    scaler.step(optimizer)
    scaler.update()

    if it % eval_interval == 0:
        losses = estimate_loss()
        train_loss = losses["train"]
        val_loss = losses["val"]
        ppl = math.exp(val_loss)

        print("Iter", f"{it:06d}", "|",
              "LR", f"{lr:.6e}", "|",
              "Train", f"{train_loss:.4f}", "|",
              "Val", f"{val_loss:.4f}", "|",
              "PPL", f"{ppl:.2f}")

        if val_loss < best_val:
            best_val = val_loss
            torch.save(
                {
                    "iter": it,
                    "model": model.state_dict(),
                    "optimizer": optimizer.state_dict(),
                    "best_val": best_val
                },
                BEST_FILE
            )
            print("New best checkpoint saved:", BEST_FILE)

    if it % save_interval == 0 and it > 0:
        torch.save(
            {
                "iter": it,
                "model": model.state_dict(),
                "optimizer": optimizer.state_dict(),
                "best_val": best_val
            },
            CKPT_FILE
        )
        print("Checkpoint saved:", CKPT_FILE)

elapsed = (time.time() - t0) / 60
print("Training finished in", round(elapsed, 2), "minutes")

torch.save(model.state_dict(), FINAL_FILE)
print("Final model saved:", FINAL_FILE)

@torch.no_grad()
def generate(prompt, max_new_tokens=150, temperature=0.8, top_k=50, top_p=0.95, repetition_penalty=1.10):
    model.eval()

    ids = tokenizer.encode(prompt)
    x = torch.tensor(ids, dtype=torch.long, device=device).unsqueeze(0)

    for _ in range(max_new_tokens):
        x_cond = x[:, -block_size:]

        ctx = torch.cuda.amp.autocast() if use_amp else nullcontext()
        with ctx:
            logits, _ = model(x_cond)

        logits = logits[:, -1, :]

        if repetition_penalty != 1.0:
            used_tokens = torch.unique(x[0])
            for token_id in used_tokens:
                token_id = token_id.item()
                if logits[0, token_id] < 0:
                    logits[0, token_id] *= repetition_penalty
                else:
                    logits[0, token_id] /= repetition_penalty

        logits = logits / temperature

        if top_k is not None and top_k > 0:
            v, _ = torch.topk(logits, min(top_k, logits.size(-1)))
            logits[logits < v[:, [-1]]] = float("-inf")

        if top_p is not None and top_p < 1.0:
            sorted_logits, sorted_indices = torch.sort(logits, descending=True)
            sorted_probs = F.softmax(sorted_logits, dim=-1)
            cumulative_probs = torch.cumsum(sorted_probs, dim=-1)

            sorted_indices_to_remove = cumulative_probs > top_p
            sorted_indices_to_remove[:, 1:] = sorted_indices_to_remove[:, :-1].clone()
            sorted_indices_to_remove[:, 0] = False

            indices_to_remove = sorted_indices_to_remove.scatter(
                1, sorted_indices, sorted_indices_to_remove
            )
            logits = logits.masked_fill(indices_to_remove, float("-inf"))

        probs = F.softmax(logits, dim=-1)
        next_id = torch.multinomial(probs, num_samples=1)
        x = torch.cat((x, next_id), dim=1)

        if next_id.item() == eos_id:
            break

    return tokenizer.decode(x[0].tolist())

print("============================================================")
print("Quick generation test")
print("============================================================")
print(generate("Artificial intelligence is", max_new_tokens=120))

print("============================================================")
print("PCS GPT - Chat")
print("Scrivi exit per uscire.")
print("============================================================")

while True:
    user_in = input("Tu: ").strip()

    if user_in.lower() == "exit":
        break

    out = generate(
        prompt=user_in,
        max_new_tokens=120,
        temperature=0.8,
        top_k=40,
        top_p=0.9,
        repetition_penalty=1.12
    )

    print("PCS:")
    print(out)