"""Run the Phase 3 notebook's complete ArithMark-3.0 evaluation protocol."""

from __future__ import annotations

import argparse
import json
import time
from collections import defaultdict
from pathlib import Path

import torch
import torch.nn.functional as F
from datasets import load_dataset
from transformers import AutoModelForCausalLM, AutoTokenizer


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser()
    parser.add_argument("--model", required=True, help="Local model directory or Hub id")
    parser.add_argument("--output", required=True, type=Path)
    parser.add_argument("--max-context", type=int, default=1024)
    parser.add_argument("--progress-every", type=int, default=25)
    parser.add_argument(
        "--dtype",
        choices=("float32", "bfloat16"),
        default="float32",
        help=(
            "Forward precision. bfloat16 reproduces the Phase 3 notebook's "
            "CUDA autocast protocol while keeping the stored model FP32."
        ),
    )
    return parser.parse_args()


def main() -> None:
    args = parse_args()
    if not torch.cuda.is_available():
        raise RuntimeError("This final benchmark run requires CUDA.")

    device = torch.device("cuda")
    tokenizer = AutoTokenizer.from_pretrained(args.model, trust_remote_code=True)
    model = AutoModelForCausalLM.from_pretrained(
        args.model,
        trust_remote_code=True,
        dtype=torch.float32,
    ).to(device)
    model.eval()

    try:
        dataset = load_dataset("AxiomicLabs/Arithmark-3.0", split="train")
    except Exception:
        dataset = load_dataset(
            "AxiomicLabs/Arithmark-3.0",
            data_files="arithmark-3.jsonl",
            split="train",
        )

    @torch.no_grad()
    def score_continuation(context: str, ending: str) -> tuple[float, int]:
        # The dataset owns whitespace: contexts have no trailing space and endings
        # already start with one. Left-truncate the context only, as in Phase 3.
        continuation_ids = tokenizer(
            ending, add_special_tokens=False
        )["input_ids"]
        context_ids = tokenizer(
            context, add_special_tokens=False
        )["input_ids"]
        context_ids = context_ids[-(args.max_context - len(continuation_ids)) :]

        input_ids = torch.tensor(
            [context_ids + continuation_ids], device=device, dtype=torch.long
        )
        with torch.autocast(
            device_type="cuda",
            dtype=torch.bfloat16,
            enabled=args.dtype == "bfloat16",
        ):
            output = model(input_ids)

        logits = output["logits"] if isinstance(output, dict) else output.logits
        if isinstance(logits, (tuple, list)):
            logits = logits[0]
        log_probs = F.log_softmax(
            logits[0, len(context_ids) - 1 : -1].float(), dim=-1
        )
        targets = torch.tensor(continuation_ids, device=device, dtype=torch.long)
        score = log_probs.gather(-1, targets[:, None]).sum().item()
        return score, len(continuation_ids)

    started = time.perf_counter()
    raw_correct = 0
    normalized_correct = 0
    categories: dict[str, list[int]] = defaultdict(lambda: [0, 0, 0])
    samples: list[dict[str, object]] = []

    for index, example in enumerate(dataset):
        label = int(example["label"])
        scored = [
            score_continuation(example["ctx"], ending)
            for ending in example["endings"]
        ]
        raw_scores = [score for score, _ in scored]
        normalized_scores = [score / max(length, 1) for score, length in scored]
        raw_prediction = max(range(len(raw_scores)), key=raw_scores.__getitem__)
        normalized_prediction = max(
            range(len(normalized_scores)), key=normalized_scores.__getitem__
        )
        raw_hit = int(raw_prediction == label)
        normalized_hit = int(normalized_prediction == label)
        raw_correct += raw_hit
        normalized_correct += normalized_hit

        category = example["activity_label"]
        categories[category][0] += raw_hit
        categories[category][1] += normalized_hit
        categories[category][2] += 1
        samples.append(
            {
                "index": index,
                "label": label,
                "raw_prediction": raw_prediction,
                "normalized_prediction": normalized_prediction,
                "raw_scores": raw_scores,
                "normalized_scores": normalized_scores,
                "continuation_lengths": [length for _, length in scored],
                "activity_label": category,
            }
        )

        complete = index + 1
        if complete % args.progress_every == 0 or complete == len(dataset):
            elapsed = time.perf_counter() - started
            print(
                f"ArithMark-3 {complete}/{len(dataset)} "
                f"acc={100 * raw_correct / complete:.2f}% "
                f"acc_norm={100 * normalized_correct / complete:.2f}% "
                f"elapsed={elapsed:.1f}s",
                flush=True,
            )

    total = len(dataset)
    elapsed = time.perf_counter() - started
    result = {
        "benchmark": "ArithMark-3.0",
        "dataset": "AxiomicLabs/Arithmark-3.0",
        "split": "train",
        "model": args.model,
        "stored_model_dtype": "float32",
        "dtype": args.dtype,
        "forward_precision": (
            "CUDA bfloat16 autocast" if args.dtype == "bfloat16" else "float32"
        ),
        "log_softmax_dtype": "float32",
        "max_context": args.max_context,
        "num_examples": total,
        "limit": None,
        "acc": raw_correct / total,
        "acc_norm": normalized_correct / total,
        "raw_correct": raw_correct,
        "normalized_correct": normalized_correct,
        "elapsed_seconds": elapsed,
        "categories": {
            category: {
                "num_examples": counts[2],
                "acc": counts[0] / counts[2],
                "acc_norm": counts[1] / counts[2],
            }
            for category, counts in sorted(categories.items())
        },
        "samples": samples,
    }
    args.output.parent.mkdir(parents=True, exist_ok=True)
    args.output.write_text(json.dumps(result, indent=2), encoding="utf-8")
    print(json.dumps({key: result[key] for key in (
        "num_examples", "acc", "acc_norm", "elapsed_seconds"
    )}, indent=2))


if __name__ == "__main__":
    main()