"""Load GUIDO-small 200M (bugfix) e genera testo.
Run:  python inference_example.py
"""
import sys
from pathlib import Path
import torch
from transformers import AutoTokenizer

# Aggiungi questa cartella al sys.path così Vathos viene trovato
HERE = Path(__file__).resolve().parent
sys.path.insert(0, str(HERE / "vathos"))      # bundled Vathos package
sys.path.insert(0, str(HERE))                  # modeling_guido.py

# Il modeling è in modeling_guido.py (fork di train_guido_small.py — usa la classe PiCOFormerLM).
# `import modeling_guido as m` registra Vathos in production mode + carica le classi.
import modeling_guido as m


def load_model(ckpt_path: str, device="cuda"):
    """Carica il ckpt fp32, infera l'architettura dai pesi, restituisce model+tokenizer."""
    state = torch.load(ckpt_path, map_location="cpu", weights_only=False)
    sd = state["model"] if "model" in state else state
    # strip eventuali prefissi
    sd = {k[len("_orig_mod."):] if k.startswith("_orig_mod.") else k: v for k, v in sd.items()}

    # Infer arch dal state_dict (così funziona anche se il config.json non c'è)
    hp = m.HP()
    hp.head_dim = 64
    hp.d_model = sd["backbone.embedder.embedding.weight"].shape[1]
    hp.n_heads = hp.d_model // hp.head_dim
    hp.d_ff = sd["backbone.blocks.0.channel_mixer.expand.weight"].shape[0]
    hp.n_layers = max(int(k.split("blocks.")[1].split(".")[0])
                      for k in sd if "blocks." in k) + 1
    vocab = sd["backbone.embedder.embedding.weight"].shape[0]
    print(f"[load] arch inferred: {hp.n_layers}L × {hp.d_model} × {hp.n_heads}h × {hp.d_ff}ff  vocab={vocab}")

    model = m.PiCOFormerLM(hp, vocab_size=vocab).to(device).bfloat16()
    missing, unexpected = model.load_state_dict(sd, strict=False)
    if missing:    print(f"[load] missing keys ({len(missing)}): {missing[:3]}...")
    if unexpected: print(f"[load] unexpected ({len(unexpected)}): {unexpected[:3]}...")
    model.eval()

    tok = AutoTokenizer.from_pretrained("mistralai/Mathstral-7B-v0.1")
    return model, tok


@torch.no_grad()
def generate(model, tok, prompt: str, max_new=200, temperature=0.6, top_p=0.95, top_k=40, device="cuda"):
    """Sampling semplice (T, top-p, top-k)."""
    import torch.nn.functional as F
    ids = tok(prompt, return_tensors="pt").input_ids.to(device)
    eos = tok.eos_token_id or 2
    softcap = model.logit_softcap
    for _ in range(max_new):
        h = model._hidden(ids)
        logits = F.linear(h[:, -1, :], model.classifier_weight)
        if softcap > 0:
            logits = softcap * torch.tanh(logits / softcap)
        logits = logits.float() / temperature
        if top_k and top_k < logits.size(-1):
            kth = logits.topk(top_k, dim=-1).values[:, -1:]
            logits = logits.masked_fill(logits < kth, float("-inf"))
        probs = F.softmax(logits, dim=-1)
        if top_p < 1.0:
            sp, si = probs.sort(dim=-1, descending=True)
            csum = sp.cumsum(dim=-1)
            mask = csum - sp > top_p
            sp = sp.masked_fill(mask, 0.0); sp /= sp.sum(dim=-1, keepdim=True)
            choice = torch.multinomial(sp, 1)
            nxt = si.gather(-1, choice)
        else:
            nxt = torch.multinomial(probs, 1)
        ids = torch.cat([ids, nxt], dim=1)
        if nxt.item() == eos:
            break
    return tok.decode(ids[0], skip_special_tokens=False)


if __name__ == "__main__":
    # Adatta il path al checkpoint
    CKPT = HERE / "pytorch_model.bin"
    device = "cuda" if torch.cuda.is_available() else "cpu"
    model, tok = load_model(str(CKPT), device=device)

    prompts = [
        "Question: Sarah has 5 apples and gives 2 to Tom. How many does she have left?\nAnswer: ",
        "Solve: $x^2 + 3x + 2 = 0$. Step 1:",
        "The largest planet in the solar system is",
    ]
    for p in prompts:
        print("=" * 60)
        print("PROMPT:", p)
        out = generate(model, tok, p, max_new=120, device=device)
        print("OUT:   ", out)