import argparse
import json
from pathlib import Path

import torch
import tiktoken
from safetensors.torch import load_file

from modeling_ys import GPTModel


TOKENIZER_VOCAB_SIZE = 50257
END_OF_TEXT_ID = 50256


def main():
    parser = argparse.ArgumentParser(description="Génère du texte avec Ys.")
    parser.add_argument("--prompt", default="Il était une fois")
    parser.add_argument("--max-new-tokens", type=int, default=120)
    parser.add_argument("--temperature", type=float, default=0.8,
                        help="0 pour choisir le token le plus probable")
    parser.add_argument("--top-k", type=int, default=40,
                        help="0 pour désactiver le filtrage top-k")
    parser.add_argument("--device", choices=("auto", "cpu", "cuda"), default="auto")
    args = parser.parse_args()

    repo_dir = Path(__file__).resolve().parent
    config = json.loads((repo_dir / "config.json").read_text(encoding="utf-8"))
    device = args.device
    if device == "auto":
        device = "cuda" if torch.cuda.is_available() else "cpu"
    if device == "cuda" and not torch.cuda.is_available():
        raise SystemExit("CUDA n'est pas disponible. Relance avec --device cpu.")
    if args.temperature < 0:
        raise SystemExit("--temperature doit être positif ou égal à 0.")

    model = GPTModel(config)
    state_dict = load_file(str(repo_dir / "model.safetensors"), device="cpu")
    model.load_state_dict(state_dict, strict=True)
    model.to(device).eval()

    tokenizer = tiktoken.get_encoding("gpt2")
    token_ids = torch.tensor([tokenizer.encode(args.prompt)], dtype=torch.long, device=device)
    if token_ids.shape[1] == 0:
        token_ids = torch.tensor([[END_OF_TEXT_ID]], dtype=torch.long, device=device)

    for _ in range(args.max_new_tokens):
        context = token_ids[:, -config["context_length"]:]
        with torch.inference_mode():
            logits = model(context)[:, -1, :TOKENIZER_VOCAB_SIZE]

        if args.temperature == 0:
            next_id = torch.argmax(logits, dim=-1, keepdim=True)
        else:
            logits = logits / args.temperature
            if 0 < args.top_k < logits.shape[-1]:
                threshold = torch.topk(logits, args.top_k, dim=-1).values[:, -1, None]
                logits = logits.masked_fill(logits < threshold, float("-inf"))
            probabilities = torch.softmax(logits, dim=-1)
            next_id = torch.multinomial(probabilities, num_samples=1)

        if next_id.item() == END_OF_TEXT_ID:
            break
        token_ids = torch.cat((token_ids, next_id), dim=1)

    print(tokenizer.decode(token_ids[0].tolist()))


if __name__ == "__main__":
    main()