"""GSM8K evaluation for quantized MLX LLMs.

Measures grade-school math accuracy to verify quantized models
still reason correctly, not just produce fluent text.
"""

import os
import re
from dataclasses import dataclass

import numpy as np


# 3-shot exemplars for chain-of-thought prompting
FEW_SHOT_EXAMPLES = [
    {
        "question": "There are 15 trees in the grove. Grove workers will plant trees in the grove today. After they are done, there will be 21 trees. How many trees did the grove workers plant today?",
        "answer": "There are 15 trees originally. Then there were 21 trees after some more were planted. So there must have been 21 - 15 = 6 trees planted.\n#### 6",
    },
    {
        "question": "If there are 3 cars in the parking lot and 2 more cars arrive, how many cars are in the parking lot?",
        "answer": "There are originally 3 cars. 2 more cars arrive. 3 + 2 = 5.\n#### 5",
    },
    {
        "question": "Leah had 32 chocolates and her sister had 42. If they ate 35, how many pieces do they have left in total?",
        "answer": "Originally, Leah had 32 chocolates. Her sister had 42. So in total they had 32 + 42 = 74. After eating 35, they had 74 - 35 = 39.\n#### 39",
    },
]


@dataclass
class GSM8KResult:
    n_correct: int
    n_total: int
    accuracy: float
    per_question: list[dict]


def _build_prompt(question: str, n_shots: int = 3) -> str:
    """Build a few-shot prompt for GSM8K."""
    prompt = ""
    for ex in FEW_SHOT_EXAMPLES[:n_shots]:
        prompt += f"Q: {ex['question']}\nA: {ex['answer']}\n\n"
    prompt += f"Q: {question}\nA:"
    return prompt


def _extract_answer(text: str) -> str | None:
    """Extract the numeric answer after #### from model output."""
    # Look for #### marker
    match = re.search(r"####\s*(-?[\d,]+\.?\d*)", text)
    if match:
        return match.group(1).replace(",", "").strip()

    # Fallback: find the last number in the text
    numbers = re.findall(r"-?[\d,]+\.?\d*", text)
    if numbers:
        return numbers[-1].replace(",", "").strip()

    return None


def _extract_ground_truth(answer_text: str) -> str:
    """Extract numeric answer from GSM8K ground truth (after ####)."""
    match = re.search(r"####\s*(-?[\d,]+\.?\d*)", answer_text)
    if match:
        return match.group(1).replace(",", "").strip()
    return ""


def _normalize_number(s: str) -> float | None:
    """Parse a number string to float for comparison."""
    try:
        return float(s.replace(",", ""))
    except (ValueError, TypeError):
        return None


def evaluate_gsm8k(
    model_path: str,
    n_samples: int = 200,
    n_shots: int = 3,
    max_tokens: int = 256,
    seed: int = 42,
) -> GSM8KResult:
    """Evaluate an MLX LLM on GSM8K test set.

    Args:
        model_path: Path to MLX model directory
        n_samples: Number of test questions to evaluate
        n_shots: Number of few-shot examples in prompt
        max_tokens: Max tokens to generate per question
        seed: Random seed for sample selection

    Returns:
        GSM8KResult with accuracy and per-question details
    """
    import mlx.core as mx
    from mlx_lm import load, generate
    from mlx_lm.sample_utils import make_sampler
    from datasets import load_dataset
    from tqdm import tqdm

    # Resolve to absolute path so mlx_lm treats it as local, not HF repo
    if os.path.isdir(model_path):
        model_path = os.path.abspath(model_path)
    print(f"  Loading model from {model_path}...")
    model, tokenizer = load(model_path)

    print(f"  Loading GSM8K test set...")
    ds = load_dataset("openai/gsm8k", "main", split="test")

    # Sample subset
    rng = np.random.RandomState(seed)
    indices = rng.choice(len(ds), size=min(n_samples, len(ds)), replace=False)
    indices.sort()

    n_correct = 0
    per_question = []

    print(f"  Evaluating {len(indices)} questions ({n_shots}-shot)...")
    for idx in tqdm(indices, desc="  GSM8K eval"):
        item = ds[int(idx)]
        question = item["question"]
        gt_answer = _extract_ground_truth(item["answer"])

        prompt = _build_prompt(question, n_shots=n_shots)

        # Generate with greedy decoding for deterministic math
        sampler = make_sampler(temp=0.0)
        output = generate(
            model, tokenizer,
            prompt=prompt,
            max_tokens=max_tokens,
            verbose=False,
            sampler=sampler,
        )

        predicted = _extract_answer(output)

        gt_num = _normalize_number(gt_answer)
        pred_num = _normalize_number(predicted) if predicted else None

        correct = (gt_num is not None and pred_num is not None
                   and abs(gt_num - pred_num) < 1e-3)
        if correct:
            n_correct += 1

        per_question.append({
            "idx": int(idx),
            "question": question[:100] + "...",
            "ground_truth": gt_answer,
            "predicted": predicted,
            "correct": correct,
        })

    accuracy = n_correct / len(indices) if indices.size > 0 else 0.0

    return GSM8KResult(
        n_correct=n_correct,
        n_total=len(indices),
        accuracy=accuracy,
        per_question=per_question,
    )


def print_gsm8k_report(result: GSM8KResult):
    """Print GSM8K evaluation results."""
    print(f"\n  GSM8K Results")
    print(f"  {'=' * 50}")
    print(f"  Accuracy: {result.n_correct}/{result.n_total} "
          f"({result.accuracy:.1%})")

    # Show some examples
    wrong = [q for q in result.per_question if not q["correct"]]
    right = [q for q in result.per_question if q["correct"]]

    if right:
        print(f"\n  Correct examples:")
        for q in right[:3]:
            print(f"    Q: {q['question']}")
            print(f"    GT: {q['ground_truth']}, Pred: {q['predicted']}")

    if wrong:
        print(f"\n  Incorrect examples:")
        for q in wrong[:3]:
            print(f"    Q: {q['question']}")
            print(f"    GT: {q['ground_truth']}, Pred: {q['predicted']}")