#!/usr/bin/env python3
"""Evaluate a causal language model on the finance-related MMLU subjects."""

from __future__ import annotations

import argparse
import json
import math
import statistics
import time
import urllib.error
import urllib.request
from dataclasses import dataclass
from datetime import datetime, timezone
from pathlib import Path
from typing import Any, Iterable

import pyarrow.parquet as parquet
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer


DATASET_ID = "cais/mmlu"
DATASET_REVISION = "c30699e8356da336a370243923dbaf21066bb9fe"
DEFAULT_SUBJECTS = (
    "business_ethics",
    "econometrics",
    "high_school_macroeconomics",
    "high_school_microeconomics",
    "management",
    "marketing",
    "professional_accounting",
)
ANSWER_LABELS = ("A", "B", "C", "D")


@dataclass(frozen=True)
class MmluExample:
    question: str
    choices: tuple[str, ...]
    answer: int


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser()
    parser.add_argument("--model", required=True)
    parser.add_argument("--output-dir", required=True)
    parser.add_argument("--cache-dir", default=".artifacts/mmlu-finance")
    parser.add_argument("--revision", default=None)
    parser.add_argument("--subjects", nargs="+", default=list(DEFAULT_SUBJECTS))
    parser.add_argument("--few-shot", type=int, default=5)
    parser.add_argument("--batch-size", type=int, default=2)
    parser.add_argument("--max-samples-per-subject", type=int, default=None)
    parser.add_argument("--device", choices=("auto", "cpu", "cuda", "mps"), default="auto")
    return parser.parse_args()


def resolve_device(requested: str) -> str:
    if requested != "auto":
        return requested
    if torch.cuda.is_available():
        return "cuda"
    if torch.backends.mps.is_available():
        return "mps"
    return "cpu"


def dataset_url(subject: str, split: str) -> str:
    filename = f"{split}-00000-of-00001.parquet"
    return (
        f"https://huggingface.co/datasets/{DATASET_ID}/resolve/"
        f"{DATASET_REVISION}/{subject}/{filename}"
    )


def download_file(url: str, destination: Path, attempts: int = 6) -> None:
    if destination.exists() and destination.stat().st_size > 0:
        return
    destination.parent.mkdir(parents=True, exist_ok=True)
    temporary = destination.with_suffix(destination.suffix + ".part")
    last_error: BaseException | None = None
    for attempt in range(1, attempts + 1):
        try:
            with (
                urllib.request.urlopen(url, timeout=120) as response,
                temporary.open("wb") as output,
            ):
                while chunk := response.read(1024 * 1024):
                    output.write(chunk)
            temporary.replace(destination)
            return
        except (OSError, urllib.error.URLError) as exc:
            last_error = exc
            temporary.unlink(missing_ok=True)
            if attempt == attempts:
                break
            delay = min(2 ** (attempt - 1), 30)
            print(
                json.dumps(
                    {
                        "download_retry": attempt,
                        "delay_seconds": delay,
                        "url": url,
                        "error": str(exc),
                    }
                ),
                flush=True,
            )
            time.sleep(delay)
    raise RuntimeError(f"Unable to download {url} after {attempts} attempts") from last_error


def load_split(cache_dir: Path, subject: str, split: str) -> list[MmluExample]:
    path = cache_dir / subject / f"{split}.parquet"
    download_file(dataset_url(subject, split), path)
    records = parquet.read_table(path).to_pylist()
    return [
        MmluExample(
            question=str(record["question"]),
            choices=tuple(str(choice) for choice in record["choices"]),
            answer=int(record["answer"]),
        )
        for record in records
    ]


def format_example(example: MmluExample, include_answer: bool) -> str:
    lines = [example.question]
    lines.extend(
        f"{label}. {choice}"
        for label, choice in zip(ANSWER_LABELS, example.choices, strict=True)
    )
    if include_answer:
        lines.append(f"Answer: {ANSWER_LABELS[example.answer]}")
    else:
        lines.append("Answer:")
    return "\n".join(lines)


def build_prompt(subject: str, few_shot: Iterable[MmluExample], test: MmluExample) -> str:
    readable_subject = subject.replace("_", " ")
    header = (
        "The following are multiple choice questions (with answers) "
        f"about {readable_subject}.\n\n"
    )
    demonstrations = "\n\n".join(
        format_example(example, include_answer=True) for example in few_shot
    )
    return f"{header}{demonstrations}\n\n{format_example(test, include_answer=False)}"


def answer_token_ids(tokenizer: Any) -> list[int]:
    result: list[int] = []
    for label in ANSWER_LABELS:
        encoded = tokenizer.encode(f" {label}", add_special_tokens=False)
        if len(encoded) != 1:
            raise ValueError(f"Answer label {label!r} is not a single token: {encoded}")
        result.append(encoded[0])
    return result


def chunks(values: list[Any], size: int) -> Iterable[list[Any]]:
    for start in range(0, len(values), size):
        yield values[start : start + size]


def write_json_atomic(path: Path, payload: dict[str, Any]) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    temporary = path.with_suffix(path.suffix + ".part")
    temporary.write_text(
        json.dumps(payload, indent=2, sort_keys=True) + "\n",
        encoding="utf-8",
    )
    temporary.replace(path)


def write_jsonl_atomic(path: Path, rows: list[dict[str, Any]]) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    temporary = path.with_suffix(path.suffix + ".part")
    with temporary.open("w", encoding="utf-8") as output:
        for row in rows:
            output.write(json.dumps(row, sort_keys=True) + "\n")
    temporary.replace(path)


def load_subject_checkpoint(
    output_dir: Path,
    subject: str,
    expected_config: dict[str, Any],
) -> tuple[dict[str, Any], list[dict[str, Any]]] | None:
    subject_dir = output_dir / "subjects" / subject
    result_path = subject_dir / "result.json"
    predictions_path = subject_dir / "predictions.jsonl"
    if not result_path.exists() or not predictions_path.exists():
        return None
    result = json.loads(result_path.read_text(encoding="utf-8"))
    predictions = [
        json.loads(line)
        for line in predictions_path.read_text(encoding="utf-8").splitlines()
        if line.strip()
    ]
    if result.get("sample_size") != len(predictions):
        return None
    if result.get("evaluation_config") != expected_config:
        return None
    return result, predictions


def write_subject_checkpoint(
    output_dir: Path,
    subject: str,
    result: dict[str, Any],
    predictions: list[dict[str, Any]],
) -> None:
    subject_dir = output_dir / "subjects" / subject
    write_json_atomic(subject_dir / "result.json", result)
    write_jsonl_atomic(subject_dir / "predictions.jsonl", predictions)


def wilson_interval(correct: int, total: int, z: float = 1.959963984540054) -> tuple[float, float]:
    if total == 0:
        return 0.0, 0.0
    proportion = correct / total
    denominator = 1 + (z * z / total)
    centre = proportion + z * z / (2 * total)
    margin = z * math.sqrt((proportion * (1 - proportion) + z * z / (4 * total)) / total)
    return (centre - margin) / denominator, (centre + margin) / denominator


def evaluate_subject(
    *,
    model: Any,
    tokenizer: Any,
    device: str,
    subject: str,
    dev: list[MmluExample],
    test: list[MmluExample],
    few_shot_count: int,
    batch_size: int,
) -> tuple[dict[str, Any], list[dict[str, Any]]]:
    few_shot = dev[:few_shot_count]
    answer_ids = torch.tensor(answer_token_ids(tokenizer), device=device)
    predictions: list[dict[str, Any]] = []
    started = time.perf_counter()

    for batch in chunks(test, batch_size):
        prompts = [build_prompt(subject, few_shot, example) for example in batch]
        encoded = tokenizer(
            prompts,
            padding=True,
            return_tensors="pt",
            add_special_tokens=True,
        ).to(device)
        with torch.inference_mode():
            logits = model(**encoded).logits
        # Left padding keeps every final prompt token at the final sequence index.
        final_positions = torch.full(
            (len(batch),),
            encoded["input_ids"].shape[1] - 1,
            device=device,
            dtype=torch.long,
        )
        row_indices = torch.arange(len(batch), device=device)
        final_logits = logits[row_indices, final_positions]
        choice_logits = final_logits.index_select(dim=1, index=answer_ids)
        probabilities = torch.softmax(choice_logits.float(), dim=1).cpu()
        predicted = choice_logits.argmax(dim=1).cpu().tolist()

        for example, prediction, probability in zip(
            batch, predicted, probabilities.tolist(), strict=True
        ):
            predictions.append(
                {
                    "subject": subject,
                    "expected": ANSWER_LABELS[example.answer],
                    "predicted": ANSWER_LABELS[prediction],
                    "correct": prediction == example.answer,
                    "choice_probabilities": {
                        label: round(value, 8)
                        for label, value in zip(ANSWER_LABELS, probability, strict=True)
                    },
                }
            )

    correct = sum(int(row["correct"]) for row in predictions)
    total = len(predictions)
    lower, upper = wilson_interval(correct, total)
    return (
        {
            "subject": subject,
            "correct": correct,
            "sample_size": total,
            "accuracy": round(correct / total, 8) if total else None,
            "confidence_interval_95": [round(lower, 8), round(upper, 8)],
            "duration_seconds": round(time.perf_counter() - started, 3),
        },
        predictions,
    )


def main() -> None:
    args = parse_args()
    if args.few_shot < 0 or args.batch_size < 1:
        raise ValueError("few-shot must be non-negative and batch-size must be positive")

    output_dir = Path(args.output_dir)
    cache_dir = Path(args.cache_dir)
    output_dir.mkdir(parents=True, exist_ok=True)
    device = resolve_device(args.device)
    dtype = torch.bfloat16 if device in {"cuda", "mps"} else torch.float32

    # Resolve every dataset dependency before allocating model memory. A transient
    # CDN failure must not invalidate hours of completed inference.
    splits: dict[str, tuple[list[MmluExample], list[MmluExample]]] = {}
    for subject in args.subjects:
        dev = load_split(cache_dir, subject, "dev")
        test = load_split(cache_dir, subject, "test")
        if args.max_samples_per_subject is not None:
            test = test[: args.max_samples_per_subject]
        splits[subject] = (dev, test)

    tokenizer = AutoTokenizer.from_pretrained(args.model, revision=args.revision)
    if tokenizer.pad_token_id is None:
        tokenizer.pad_token = tokenizer.eos_token
    tokenizer.padding_side = "left"
    model = AutoModelForCausalLM.from_pretrained(
        args.model,
        revision=args.revision,
        dtype=dtype,
        low_cpu_mem_usage=True,
    ).to(device)
    model.eval()

    subject_results: list[dict[str, Any]] = []
    predictions: list[dict[str, Any]] = []
    for subject in args.subjects:
        dev, test = splits[subject]
        evaluation_config = {
            "dataset_revision": DATASET_REVISION,
            "few_shot": args.few_shot,
            "model_revision": args.revision,
            "sample_size": len(test),
        }
        checkpoint = load_subject_checkpoint(
            output_dir,
            subject,
            evaluation_config,
        )
        if checkpoint is None:
            result, subject_predictions = evaluate_subject(
                model=model,
                tokenizer=tokenizer,
                device=device,
                subject=subject,
                dev=dev,
                test=test,
                few_shot_count=args.few_shot,
                batch_size=args.batch_size,
            )
            result["evaluation_config"] = evaluation_config
            write_subject_checkpoint(
                output_dir,
                subject,
                result,
                subject_predictions,
            )
        else:
            result, subject_predictions = checkpoint
            print(
                json.dumps(
                    {
                        "subject": subject,
                        "status": "resumed_from_checkpoint",
                        "sample_size": result["sample_size"],
                    }
                ),
                flush=True,
            )
        subject_results.append(result)
        predictions.extend(subject_predictions)
        print(json.dumps(result, sort_keys=True), flush=True)

    total = sum(row["sample_size"] for row in subject_results)
    correct = sum(row["correct"] for row in subject_results)
    lower, upper = wilson_interval(correct, total)
    accuracies = [
        row["accuracy"] for row in subject_results if row["accuracy"] is not None
    ]
    summary = {
        "benchmark": "MMLU finance and business subset",
        "dataset": DATASET_ID,
        "dataset_revision": DATASET_REVISION,
        "model": args.model,
        "model_revision": args.revision,
        "evaluated_at": datetime.now(timezone.utc).isoformat(),
        "device": device,
        "dtype": str(dtype).replace("torch.", ""),
        "few_shot": args.few_shot,
        "sample_limit_per_subject": args.max_samples_per_subject,
        "subjects": subject_results,
        "sample_size": total,
        "micro_accuracy": round(correct / total, 8) if total else None,
        "macro_accuracy": round(statistics.mean(accuracies), 8) if accuracies else None,
        "confidence_interval_95": [round(lower, 8), round(upper, 8)],
        "status": "community_evaluation",
        "official_leaderboard_result": False,
    }
    write_json_atomic(output_dir / "results.json", summary)
    write_jsonl_atomic(output_dir / "predictions.jsonl", predictions)
    print(json.dumps(summary, indent=2, sort_keys=True))


if __name__ == "__main__":
    main()