"""The math eval loop, and the metrics that actually discriminate recipes.

Accuracy alone is not enough. Published measurements on this exact model family
show W4A16 retaining ~99% of MATH-500 while losing ~11-20 points of AIME, and
show quantization roughly doubling the truncation rate on AIME under a fixed
token cap. The mechanism is that low-bit weights perturb high-entropy
"branching" tokens, the model rambles, and it never emits its closing tag.

So every run reports three things:

* ``accuracy``          — did it get the answer right
* ``truncation_rate``   — did it run out of budget instead of stopping
* ``mean_generated_tokens`` / ``think_close_rate`` — is CoT inflating

A recipe that holds accuracy while inflating tokens is not safe; it is a recipe
that will collapse the moment the grader's token cap is tighter than ours.
"""

from __future__ import annotations

import json
import statistics
from dataclasses import asdict, dataclass
from pathlib import Path
from typing import Any, Sequence

from .answers import answers_match, extract_answer
from .data import MathExample, build_prompt
from .generate import GenerationOutput, generate

THINK_CLOSE_TAG = "</think>"


@dataclass
class MathPrediction:
    example_id: str
    source: str
    gold: str
    predicted: str | None
    correct: bool
    finished: bool
    think_closed: bool
    num_generated_tokens: int
    num_prompt_tokens: int
    response: str


def score_generation(
    example: MathExample, output: GenerationOutput
) -> MathPrediction:
    predicted = extract_answer(output.text)
    return MathPrediction(
        example_id=example.example_id,
        source=example.source,
        gold=example.answer,
        predicted=predicted,
        correct=answers_match(predicted, example.answer),
        finished=output.finished,
        # A thinking model that never closes its tag has looped, even if it
        # somehow stopped afterwards.
        think_closed=THINK_CLOSE_TAG in output.text,
        num_generated_tokens=output.num_generated_tokens,
        num_prompt_tokens=output.num_prompt_tokens,
        response=output.text,
    )


def summarize(predictions: Sequence[MathPrediction]) -> dict[str, Any]:
    n = len(predictions)
    if n == 0:
        return {"num_examples": 0}

    lengths = [p.num_generated_tokens for p in predictions]
    finished = [p for p in predictions if p.finished]
    truncated = [p for p in predictions if not p.finished]

    return {
        "num_examples": n,
        "accuracy": sum(p.correct for p in predictions) / n,
        "parse_rate": sum(p.predicted is not None for p in predictions) / n,
        # The headline risk metric: budget exhaustion, not wrong answers.
        "truncation_rate": len(truncated) / n,
        "think_close_rate": sum(p.think_closed for p in predictions) / n,
        "mean_generated_tokens": statistics.mean(lengths),
        "median_generated_tokens": statistics.median(lengths),
        "max_generated_tokens": max(lengths),
        # Splitting accuracy by termination separates "reasoned badly" from
        # "never got to answer" — they need different fixes.
        "accuracy_when_finished": (
            sum(p.correct for p in finished) / len(finished) if finished else None
        ),
        "accuracy_when_truncated": (
            sum(p.correct for p in truncated) / len(truncated) if truncated else None
        ),
    }


def summarize_by_source(predictions: Sequence[MathPrediction]) -> dict[str, Any]:
    sources = sorted({p.source for p in predictions})
    return {s: summarize([p for p in predictions if p.source == s]) for s in sources}


def run_eval(
    model,
    tokenizer,
    examples: Sequence[MathExample],
    *,
    max_new_tokens: int = 65536,
    temperature: float = 0.0,
    top_p: float = 0.95,
    top_k: int = 20,
    presence_penalty: float = 0.0,
    repetition_penalty: float = 1.0,
    batch_size: int = 8,
    enable_thinking: bool | None = None,
) -> list[MathPrediction]:
    outputs = generate(
        model,
        tokenizer,
        [build_prompt(ex) for ex in examples],
        max_new_tokens=max_new_tokens,
        temperature=temperature,
        top_p=top_p,
        top_k=top_k,
        repetition_penalty=repetition_penalty,
        batch_size=batch_size,
        enable_thinking=enable_thinking,
        desc="math eval",
    )
    return [score_generation(ex, out) for ex, out in zip(examples, outputs)]


def run_eval_vllm(
    model_path: str,
    tokenizer,
    examples: Sequence[MathExample],
    *,
    max_new_tokens: int = 65536,
    temperature: float = 0.0,
    top_p: float = 0.95,
    top_k: int = 20,
    presence_penalty: float = 0.0,
    repetition_penalty: float = 1.0,
    enable_thinking: bool | None = None,
    gpu_memory_utilization: float = 0.90,
    allowed_token_ids: Sequence[int] | None = None,
    llm=None,
) -> tuple[list[MathPrediction], object]:
    """Same scoring, vLLM engine. Returns the engine so it can be reused."""
    from .vllm_backend import generate_vllm

    outputs, llm = generate_vllm(
        model_path,
        tokenizer,
        [build_prompt(ex) for ex in examples],
        max_new_tokens=max_new_tokens,
        temperature=temperature,
        top_p=top_p,
        top_k=top_k,
        presence_penalty=presence_penalty,
        repetition_penalty=repetition_penalty,
        enable_thinking=enable_thinking,
        gpu_memory_utilization=gpu_memory_utilization,
        allowed_token_ids=allowed_token_ids,
        llm=llm,
    )
    return [score_generation(ex, out) for ex, out in zip(examples, outputs)], llm


def save_results(
    output_dir: str | Path,
    *,
    run_name: str,
    config: dict[str, Any],
    predictions: Sequence[MathPrediction],
) -> Path:
    """Write ``summary.json`` (tracked) and ``generations.jsonl`` (gitignored)."""
    out = Path(output_dir) / run_name
    out.mkdir(parents=True, exist_ok=True)

    summary = {
        "run_name": run_name,
        "config": config,
        "overall": summarize(predictions),
        "by_source": summarize_by_source(predictions),
    }
    (out / "summary.json").write_text(json.dumps(summary, indent=2))

    with (out / "generations.jsonl").open("w") as fh:
        for prediction in predictions:
            fh.write(json.dumps(asdict(prediction)) + "\n")

    return out


def format_summary(summary: dict[str, Any]) -> str:
    overall = summary["overall"]
    if not overall.get("num_examples"):
        return "no examples evaluated"

    lines = [
        f"  examples            {overall['num_examples']}",
        f"  accuracy            {overall['accuracy']:.3f}",
        f"  parse rate          {overall['parse_rate']:.3f}",
        f"  TRUNCATION RATE     {overall['truncation_rate']:.3f}   <- budget exhaustion",
        f"  think-close rate    {overall['think_close_rate']:.3f}",
        f"  mean gen tokens     {overall['mean_generated_tokens']:.0f}",
        f"  median gen tokens   {overall['median_generated_tokens']:.0f}",
        f"  max gen tokens      {overall['max_generated_tokens']}",
    ]
    if overall.get("accuracy_when_finished") is not None:
        lines.append(f"  acc | finished      {overall['accuracy_when_finished']:.3f}")
    if overall.get("accuracy_when_truncated") is not None:
        lines.append(f"  acc | truncated     {overall['accuracy_when_truncated']:.3f}")

    lines.append("")
    for source, stats in summary["by_source"].items():
        lines.append(
            f"  [{source}] n={stats['num_examples']} acc={stats['accuracy']:.3f} "
            f"trunc={stats['truncation_rate']:.3f} tok={stats['mean_generated_tokens']:.0f}"
        )
    return "\n".join(lines)