from __future__ import annotations

import argparse
import json
import math
import time
from pathlib import Path

import torch
import torch.nn.functional as F
import yaml
from datasets import get_dataset_config_names, load_dataset

from litgpt import Tokenizer
from litgpt.config import Config
from litgpt.model import GPT


def load_model(checkpoint_dir: Path, device: torch.device) -> GPT:
    config = Config.from_file(checkpoint_dir / "model_config.yaml")
    model = GPT(config)
    hyperparameters_path = checkpoint_dir / "hyperparameters.yaml"
    if hyperparameters_path.is_file():
        hyperparameters = yaml.safe_load(hyperparameters_path.read_text(encoding="utf-8")) or {}
        if hyperparameters.get("train", {}).get("tie_embeddings"):
            model.transformer.wte.weight = model.lm_head.weight
    checkpoint = torch.load(checkpoint_dir / "lit_model.pth", map_location="cpu")
    state_dict = checkpoint["model"] if "model" in checkpoint else checkpoint
    model.load_state_dict(state_dict)
    model.to(device)
    model.eval()
    return model


def encode_pair(tokenizer: Tokenizer, prompt: str, continuation: str, block_size: int) -> tuple[torch.Tensor, int]:
    prompt_ids = tokenizer.encode(prompt, bos=True, eos=False).long()
    continuation_ids = tokenizer.encode(continuation, bos=False, eos=False).long()
    ids = torch.cat([prompt_ids, continuation_ids])
    if ids.numel() > block_size:
        keep = min(block_size, continuation_ids.numel() + min(32, prompt_ids.numel()))
        ids = ids[-keep:]
        continuation_len = min(continuation_ids.numel(), ids.numel() - 1)
    else:
        continuation_len = continuation_ids.numel()
    return ids, continuation_len


@torch.no_grad()
def continuation_nll(
    model: GPT,
    tokenizer: Tokenizer,
    prompt: str,
    continuation: str,
    device: torch.device,
) -> float:
    ids, continuation_len = encode_pair(tokenizer, prompt, continuation, model.max_seq_length)
    if continuation_len <= 0 or ids.numel() <= 1:
        return float("inf")
    inputs = ids[:-1].unsqueeze(0).to(device)
    targets = ids[1:].to(device)
    use_autocast = device.type == "cuda"
    with torch.autocast(device_type="cuda", dtype=torch.bfloat16, enabled=use_autocast):
        logits = model(inputs)[0].float()
    start = max(0, targets.numel() - continuation_len)
    loss = F.cross_entropy(logits[start:], targets[start:], reduction="sum")
    return float(loss.cpu()) / continuation_len


def normalize_answer_key(answer_key: str, labels: list[str], texts: list[str]) -> int | None:
    if answer_key in labels:
        return labels.index(answer_key)
    for index, text in enumerate(texts):
        if answer_key.strip().lower() == text.strip().lower():
            return index
    return None


def benchmark_arc_easy(model: GPT, tokenizer: Tokenizer, device: torch.device, limit: int) -> dict[str, float | int]:
    dataset = load_dataset("ai2_arc", "ARC-Easy", split=f"validation[:{limit}]")
    correct = 0
    total = 0
    start = time.perf_counter()
    for row in dataset:
        question = row["question"].strip()
        labels = [str(label) for label in row["choices"]["label"]]
        texts = [str(text) for text in row["choices"]["text"]]
        target = normalize_answer_key(str(row["answerKey"]), labels, texts)
        if target is None:
            continue
        prompt = f"Question: {question}\nAnswer:"
        scores = [continuation_nll(model, tokenizer, prompt, " " + choice, device) for choice in texts]
        prediction = min(range(len(scores)), key=scores.__getitem__)
        correct += int(prediction == target)
        total += 1
    elapsed = time.perf_counter() - start
    return {
        "arc_easy_validation_examples": total,
        "arc_easy_accuracy": correct / total if total else math.nan,
        "arc_easy_seconds": elapsed,
    }


def benchmark_blimp(
    model: GPT,
    tokenizer: Tokenizer,
    device: torch.device,
    configs: list[str],
    examples_per_config: int,
) -> dict[str, float | int | dict[str, float]]:
    per_config: dict[str, float] = {}
    correct = 0
    total = 0
    start = time.perf_counter()
    for config in configs:
        dataset = load_dataset("nyu-mll/blimp", config, split=f"train[:{examples_per_config}]")
        config_correct = 0
        config_total = 0
        for row in dataset:
            good = str(row["sentence_good"]).strip()
            bad = str(row["sentence_bad"]).strip()
            good_score = continuation_nll(model, tokenizer, "", good, device)
            bad_score = continuation_nll(model, tokenizer, "", bad, device)
            config_correct += int(good_score < bad_score)
            config_total += 1
        per_config[config] = config_correct / config_total if config_total else math.nan
        correct += config_correct
        total += config_total
    elapsed = time.perf_counter() - start
    return {
        "blimp_examples": total,
        "blimp_configs": len(configs),
        "blimp_accuracy": correct / total if total else math.nan,
        "blimp_per_config": per_config,
        "blimp_seconds": elapsed,
    }


def default_blimp_configs(limit: int) -> list[str]:
    preferred = [
        "adjunct_island",
        "anaphor_number_agreement",
        "determiner_noun_agreement_1",
        "irregular_past_participle_adjectives",
        "subject_verb_agreement_simple",
    ]
    available = set(get_dataset_config_names("nyu-mll/blimp"))
    configs = [name for name in preferred if name in available]
    if len(configs) < limit:
        configs.extend(name for name in sorted(available) if name not in configs)
    return configs[:limit]


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("--checkpoint-dir", type=Path, required=True)
    parser.add_argument("--tokenizer-dir", type=Path, required=True)
    parser.add_argument("--arc-limit", type=int, default=100)
    parser.add_argument("--blimp-configs", type=int, default=5)
    parser.add_argument("--blimp-examples", type=int, default=50)
    parser.add_argument("--out", type=Path)
    args = parser.parse_args()

    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    model = load_model(args.checkpoint_dir, device)
    tokenizer = Tokenizer(args.tokenizer_dir)
    metrics: dict[str, object] = {
        "checkpoint_dir": str(args.checkpoint_dir),
        "tokenizer_dir": str(args.tokenizer_dir),
        "device": str(device),
        "model_name": model.config.name,
        "parameters": sum(parameter.numel() for parameter in model.parameters()),
    }
    metrics.update(benchmark_arc_easy(model, tokenizer, device, args.arc_limit))
    metrics.update(
        benchmark_blimp(
            model,
            tokenizer,
            device,
            default_blimp_configs(args.blimp_configs),
            args.blimp_examples,
        )
    )

    text = json.dumps(metrics, indent=2, sort_keys=True)
    print(text)
    if args.out:
        args.out.parent.mkdir(parents=True, exist_ok=True)
        args.out.write_text(text + "\n", encoding="utf-8")


if __name__ == "__main__":
    main()