"""Benchmark con diverse configurazioni di temperatura e repetition penalty.

Uso:
    python benchmark_v3.py --model-pt checkpoints/finetune/model.pt --label "V1"
    python benchmark_v3.py --model-pt checkpoints/finetune_p3/model.pt --label "V3"
"""
from __future__ import annotations

import argparse

import torch
from tokenizers import Tokenizer

from src.model.config import BoccaccioConfig
from src.model.transformer import BoccaccioForCausalLM

QUESTIONS = [
    "Chi e' Silvio Berlusconi?",
    "Qual e' la capitale d'Italia?",
    "Spiega la fotosintesi in modo semplice.",
    "Scrivi una funzione Python che calcola il fattoriale di un numero.",
    "Quali sono le cause della Prima Guerra Mondiale?",
]

CONFIGS = [
    ("temp=0.3, rep=1.2", 0.3, 1.2),
    ("temp=0.1, rep=1.2", 0.1, 1.2),
    ("temp=0.3, rep=1.0", 0.3, 1.0),
    ("temp=0.7, rep=1.2", 0.7, 1.2),
]


def generate(model, tokenizer, prompt, max_new_tokens, temperature, rep_penalty, device="cuda"):
    full = f"### Utente: {prompt}\n### Assistente:"
    ids = tokenizer.encode(full).ids
    x = torch.tensor([ids], device=device)

    with torch.no_grad():
        for _ in range(max_new_tokens):
            out = model(x)
            logits = out["logits"][0, -1, :]

            # Repetition penalty: penalize tokens already in the sequence
            if rep_penalty != 1.0:
                for tok_id in set(x[0].tolist()):
                    logits[tok_id] /= rep_penalty

            if temperature > 0:
                logits = logits / temperature
                v, _ = torch.topk(logits, 40)
                logits[logits < v[-1]] = float("-inf")
                probs = torch.softmax(logits, dim=-1)
                tok_id = torch.multinomial(probs, num_samples=1).item()
            else:
                tok_id = torch.argmax(logits).item()

            if tok_id == 2:  # </s>
                break
            x = torch.cat([x, torch.tensor([[tok_id]], device=device)], dim=1)

    new_tokens = x[0].cpu().tolist()[len(ids):]
    return tokenizer.decode(new_tokens)


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--model-pt", type=str, default="checkpoints/finetune_p3/model.pt")
    parser.add_argument("--label", type=str, default="V3")
    args = parser.parse_args()

    print(f"Caricamento modello {args.label} da {args.model_pt}...")
    config = BoccaccioConfig.from_yaml("configs/model.yaml", variant="model")
    model = BoccaccioForCausalLM(config).to("cuda")

    sd = torch.load(args.model_pt, map_location="cuda", weights_only=True)
    cleaned = {}
    for k, v in sd.items():
        nk = k[len("_orig_mod."):] if k.startswith("_orig_mod.") else k
        cleaned[nk] = v
    model.load_state_dict(cleaned)
    model.eval()
    print(f"Modello {args.label} caricato.\n")

    tokenizer = Tokenizer.from_file("tokenizer/boccaccio-32k.json")

    for config_name, temp, rep in CONFIGS:
        print("=" * 80)
        print(f"[{args.label}] CONFIG: {config_name}")
        print("=" * 80)
        for i, q in enumerate(QUESTIONS, 1):
            print(f"\n--- Domanda {i}/5: {q} ---")
            resp = generate(model, tokenizer, q, max_new_tokens=150,
                            temperature=temp, rep_penalty=rep)
            print(f"R: {resp}")
        print()


if __name__ == "__main__":
    main()