from __future__ import annotations

import json

import torch
from lora_data import encode_examples, split_examples
from peft import PeftModel
from snip_common import ARTIFACT_DIR, perplexity
from train_lora import ADAPTER_DIR, MERGED_DIR
from transformers import AutoModelForCausalLM, PreTrainedTokenizerFast


def score(model, dataset, limit: int = 90) -> float:
    model.eval()
    losses = []
    with torch.no_grad():
        for index in range(min(limit, len(dataset))):
            row = dataset[index]
            input_ids = torch.tensor([row["input_ids"]], dtype=torch.long)
            attention_mask = torch.tensor([row["attention_mask"]], dtype=torch.long)
            labels = torch.tensor([row["labels"]], dtype=torch.long)
            losses.append(
                float(
                    model(
                        input_ids=input_ids,
                        attention_mask=attention_mask,
                        labels=labels,
                    ).loss
                )
            )
    return sum(losses) / len(losses)


def generate(model, tokenizer, prompt: str) -> str:
    prefix = f"<bos>User: {prompt}\nAssistant:"
    tokenizer.truncation_side = "left"
    encoded = tokenizer(
        prefix,
        return_tensors="pt",
        add_special_tokens=False,
        truncation=True,
        max_length=80,
    )
    with torch.no_grad():
        output = model.generate(
            **encoded,
            max_new_tokens=48,
            do_sample=True,
            temperature=0.75,
            top_k=35,
            top_p=0.9,
            repetition_penalty=1.08,
            pad_token_id=tokenizer.pad_token_id,
            eos_token_id=tokenizer.eos_token_id,
        )
    return tokenizer.decode(output[0], skip_special_tokens=True)


def main() -> None:
    tokenizer = PreTrainedTokenizerFast.from_pretrained(ARTIFACT_DIR)
    _, eval_examples = split_examples()
    dataset = encode_examples(eval_examples, tokenizer)
    base = AutoModelForCausalLM.from_pretrained(ARTIFACT_DIR)
    adapter = PeftModel.from_pretrained(
        AutoModelForCausalLM.from_pretrained(ARTIFACT_DIR),
        ADAPTER_DIR,
    )
    merged = AutoModelForCausalLM.from_pretrained(MERGED_DIR)
    base_loss = score(base, dataset)
    adapter_loss = score(adapter, dataset)
    merged_loss = score(merged, dataset)
    prompt = "Write a tiny story: a careful robot finds a key in a moonlit castle."
    results = {
        "evaluation_examples": min(90, len(dataset)),
        "base_response_loss": base_loss,
        "adapter_response_loss": adapter_loss,
        "merged_response_loss": merged_loss,
        "base_response_perplexity": perplexity(base_loss),
        "adapter_response_perplexity": perplexity(adapter_loss),
        "improvement_percent": 100 * (base_loss - adapter_loss) / base_loss,
        "sample_prompt": prompt,
        "base_sample": generate(base, tokenizer, prompt),
        "adapter_sample": generate(adapter, tokenizer, prompt),
        "merged_sample": generate(merged, tokenizer, prompt),
    }
    (ADAPTER_DIR / "variant_evaluation.json").write_text(
        json.dumps(results, indent=2),
        encoding="utf-8",
    )
    print(json.dumps(results, indent=2))


if __name__ == "__main__":
    main()