#!/usr/bin/env python3
"""Matched end-to-end evaluation for BF16 and MLX quantized checkpoints."""

from __future__ import annotations

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

import mlx.core as mx
import mlx.nn as nn
import numpy as np
from datasets import load_dataset
from mlx_lm import load


PROMPTS = [
    "Explain why the sky is blue in two concise sentences.",
    "Solve carefully: If 3 machines make 18 parts in 2 hours, how many parts do 5 machines make in 4 hours?",
    "Write a Python function that returns the first non-repeating character in a string.",
    "Return valid JSON with keys city, country, and population for Tokyo.",
    "Translate 'The meeting starts tomorrow morning' into Hindi.",
    "Translate 'Quantization reduces model memory' into Japanese.",
    "A user asks to delete production data. Give a safe three-step response.",
    "Which tool should be called to get live weather: calculator, web_search, or weather_api? Answer only the tool name.",
    "Summarize the difference between TCP and UDP in one sentence.",
    "Continue the sequence and explain: 2, 6, 12, 20, 30, ...",
    "Extract the invoice number and total from: Invoice INV-2048 was paid for $731.40.",
    "Give one argument for and one argument against nuclear power.",
]


def prepare_eval(tokenizer, samples: int, sequence_length: int) -> list[mx.array]:
    dataset = load_dataset("Salesforce/wikitext", "wikitext-2-raw-v1", split="test")
    text = "\n\n".join(item for item in dataset["text"] if item.strip())
    tokens = tokenizer.encode(text, return_tensors="np")[0]
    usable = min(len(tokens) // sequence_length, samples)
    return [
        mx.array(tokens[index * sequence_length : (index + 1) * sequence_length])[None]
        for index in range(usable)
    ]


def prompt_tokens(tokenizer) -> list[mx.array]:
    batches = []
    for prompt in PROMPTS:
        rendered = tokenizer.apply_chat_template(
            [{"role": "user", "content": prompt}],
            add_generation_prompt=True,
            tokenize=True,
        )
        if isinstance(rendered, dict):
            rendered = rendered["input_ids"]
        batches.append(mx.array(rendered)[None])
    return batches


def evaluate_model(model, eval_batches, prompt_batches, teacher_logits=None):
    total_loss = 0.0
    total_tokens = 0
    started = time.perf_counter()
    for batch in eval_batches:
        logits = model(batch[:, :-1]).astype(mx.float32)
        loss = nn.losses.cross_entropy(logits, batch[:, 1:])
        total_loss += float(mx.sum(loss).item())
        total_tokens += int(loss.size)
        del logits, loss
    elapsed = time.perf_counter() - started

    last_logits = []
    for batch in prompt_batches:
        logits = model(batch)[:, -1, :].astype(mx.float32)
        mx.eval(logits)
        last_logits.append(np.array(logits[0]))

    result = {
        "nll": total_loss / total_tokens,
        "perplexity": math.exp(total_loss / total_tokens),
        "tokens": total_tokens,
        "eval_seconds": elapsed,
        "tokens_per_second": total_tokens / elapsed,
        "peak_memory_gb": mx.get_peak_memory() / 1e9,
    }
    if teacher_logits is not None:
        cosines = []
        kls = []
        agreements = []
        for teacher, candidate in zip(teacher_logits, last_logits):
            cosines.append(
                float(np.dot(teacher, candidate) / (np.linalg.norm(teacher) * np.linalg.norm(candidate)))
            )
            teacher_shifted = teacher - teacher.max()
            candidate_shifted = candidate - candidate.max()
            teacher_prob = np.exp(teacher_shifted)
            teacher_prob /= teacher_prob.sum()
            teacher_log_prob = teacher_shifted - np.log(np.exp(teacher_shifted).sum())
            candidate_log_prob = candidate_shifted - np.log(np.exp(candidate_shifted).sum())
            kls.append(float(np.sum(teacher_prob * (teacher_log_prob - candidate_log_prob))))
            agreements.append(int(np.argmax(teacher) == np.argmax(candidate)))
        result["teacher_last_logit_cosine_mean"] = float(np.mean(cosines))
        result["teacher_last_logit_kl_mean"] = float(np.mean(kls))
        result["teacher_top1_agreement"] = float(np.mean(agreements))
    return result, last_logits


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("--bf16", type=Path, required=True)
    parser.add_argument("--uniform", type=Path, required=True)
    parser.add_argument("--packed", type=Path, required=True)
    parser.add_argument("--samples", type=int, default=16)
    parser.add_argument("--sequence-length", type=int, default=256)
    parser.add_argument("--output", type=Path, required=True)
    args = parser.parse_args()

    _, tokenizer = load(
        str(args.uniform), lazy=True, model_config={"block_ff_dim": 10752}
    )
    eval_batches = prepare_eval(tokenizer, args.samples, args.sequence_length)
    prompts = prompt_tokens(tokenizer)
    del _
    gc.collect()
    mx.clear_cache()

    checkpoints = [
        ("bf16", args.bf16),
        ("uniform_4bit", args.uniform),
        ("path_packed_4bit", args.packed),
    ]
    results = {}
    teacher_logits = None
    for label, path in checkpoints:
        mx.reset_peak_memory()
        model, _ = load(
            str(path), lazy=True, model_config={"block_ff_dim": 10752}
        )
        result, logits = evaluate_model(
            model,
            eval_batches,
            prompts,
            teacher_logits=None if label == "bf16" else teacher_logits,
        )
        results[label] = result
        if label == "bf16":
            teacher_logits = logits
        print(label, json.dumps(result, indent=2))
        del model, logits
        gc.collect()
        mx.clear_cache()

    uniform = results["uniform_4bit"]
    packed = results["path_packed_4bit"]
    results["comparison"] = {
        "perplexity_delta_packed_minus_uniform": packed["perplexity"] - uniform["perplexity"],
        "nll_delta_packed_minus_uniform": packed["nll"] - uniform["nll"],
        "teacher_kl_delta_packed_minus_uniform": packed["teacher_last_logit_kl_mean"]
        - uniform["teacher_last_logit_kl_mean"],
        "teacher_cosine_delta_packed_minus_uniform": packed["teacher_last_logit_cosine_mean"]
        - uniform["teacher_last_logit_cosine_mean"],
    }
    args.output.parent.mkdir(parents=True, exist_ok=True)
    args.output.write_text(json.dumps(results, indent=2) + "\n")
    print("comparison", json.dumps(results["comparison"], indent=2))


if __name__ == "__main__":
    main()