from __future__ import annotations

import json
import time

import torch
from snip_common import (
    ARTIFACT_DIR,
    DATA_DIR,
    parameter_count,
    perplexity,
    read_texts,
    texts_to_blocks,
)
from transformers import GPT2LMHeadModel, PreTrainedTokenizerFast

PROMPTS = [
    "Once upon a time",
    "The little robot discovered",
    "Jacob opened the castle door and",
]


def main() -> None:
    tokenizer = PreTrainedTokenizerFast.from_pretrained(ARTIFACT_DIR)
    model = GPT2LMHeadModel.from_pretrained(ARTIFACT_DIR)
    model.eval()
    dataset = texts_to_blocks(read_texts(DATA_DIR / "eval.jsonl"), tokenizer)

    losses: list[float] = []
    started = time.perf_counter()
    with torch.no_grad():
        for index in range(min(len(dataset), 200)):
            row = dataset[index]
            input_ids = torch.tensor([row["input_ids"]], dtype=torch.long)
            outputs = model(input_ids=input_ids, labels=input_ids)
            losses.append(float(outputs.loss))
    elapsed = time.perf_counter() - started
    mean_loss = sum(losses) / len(losses)

    samples = []
    for prompt in PROMPTS:
        encoded = tokenizer(prompt, return_tensors="pt")
        with torch.no_grad():
            output = model.generate(
                **encoded,
                max_new_tokens=64,
                do_sample=True,
                temperature=0.85,
                top_k=40,
                top_p=0.92,
                repetition_penalty=1.08,
                pad_token_id=tokenizer.pad_token_id,
                eos_token_id=tokenizer.eos_token_id,
            )
        samples.append(
            {
                "prompt": prompt,
                "completion": tokenizer.decode(output[0], skip_special_tokens=True),
            }
        )

    results = {
        "model": "SNIP-0.4M",
        "parameters": parameter_count(model),
        "eval_blocks": len(losses),
        "eval_loss": mean_loss,
        "perplexity": perplexity(mean_loss),
        "eval_seconds": elapsed,
        "samples": samples,
    }
    (ARTIFACT_DIR / "evaluation.json").write_text(
        json.dumps(results, indent=2),
        encoding="utf-8",
    )
    print(json.dumps(results, indent=2))


if __name__ == "__main__":
    main()