"""AI2D evaluation for Vision-Language Models.

Evaluates VQA accuracy on AI2 Diagrams — multiple-choice visual questions
about scientific diagrams (food webs, life cycles, etc.). Uses lmms-lab/ai2d
from HuggingFace Hub.

This is the VLM equivalent of GSM8K: clear right/wrong answers, exercises
genuine visual reasoning, and is sensitive to quantization quality degradation.
"""

import re
from dataclasses import dataclass, field

import numpy as np
from tqdm import tqdm


@dataclass
class AI2DResult:
    n_correct: int
    n_total: int
    accuracy: float
    per_question: list[dict] = field(default_factory=list)


def _build_prompt(question: str, options: list[str]) -> str:
    """Build a multiple-choice VQA prompt."""
    labels = ["A", "B", "C", "D", "E", "F", "G", "H"]
    choices = "\n".join(f"{labels[i]}. {opt}" for i, opt in enumerate(options))
    return (
        f"{question}\n{choices}\n"
        f"Answer with the letter of the correct option."
    )


def _extract_answer_letter(text: str, n_options: int) -> str | None:
    """Extract the answer letter from model output."""
    text = text.strip()
    valid = set("ABCDEFGH"[:n_options])

    # Check if response starts with a valid letter
    if text and text[0].upper() in valid:
        return text[0].upper()

    # Look for patterns like "A.", "A)", "Answer: A", "The answer is A"
    patterns = [
        r"(?:answer|option)\s*(?:is|:)\s*([A-H])",
        r"\b([A-H])\s*[\.\):]",
        r"\b([A-H])\b",
    ]
    for pattern in patterns:
        match = re.search(pattern, text, re.IGNORECASE)
        if match:
            letter = match.group(1).upper()
            if letter in valid:
                return letter

    return None


def evaluate_ai2d_mlx(
    model_path: str,
    n_samples: int = 100,
    max_new_tokens: int = 32,
    max_image_size: int = 512,
    seed: int = 42,
) -> AI2DResult:
    """Evaluate a quantized MLX VLM on AI2D diagram understanding.

    Args:
        model_path: Path to MLX model directory (converted via mlx-vlm).
        n_samples: Number of test questions to evaluate.
        max_new_tokens: Max tokens to generate per question.
        max_image_size: Max image dimension (limits visual token count).
        seed: Random seed for sample selection.

    Returns:
        AI2DResult with accuracy and per-question details.
    """
    import os
    from datasets import load_dataset
    from PIL import Image
    from mlx_vlm import load, generate
    from mlx_vlm.prompt_utils import apply_chat_template
    from mlx_vlm.utils import load_config

    # Resolve to absolute path
    if os.path.isdir(model_path):
        model_path = os.path.abspath(model_path)

    print(f"  Loading MLX VLM from {model_path}...")
    model, processor = load(model_path)
    config = load_config(model_path)

    print("  Loading AI2D test set from HuggingFace...")
    ds = load_dataset("lmms-lab/ai2d", split="test")

    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)} AI2D questions...")
    for idx in tqdm(indices, desc="  AI2D eval"):
        item = ds[int(idx)]
        img = item["image"]
        if not isinstance(img, Image.Image):
            img = Image.open(img)
        img = img.convert("RGB")

        w, h = img.size
        if max(w, h) > max_image_size:
            scale = max_image_size / max(w, h)
            img = img.resize((int(w * scale), int(h * scale)), Image.BILINEAR)

        question = item["question"]
        options = item["options"]
        gt_idx = int(item["answer"])
        labels = ["A", "B", "C", "D", "E", "F", "G", "H"]
        gt_letter = labels[gt_idx]

        prompt_text = _build_prompt(question, options)

        try:
            # Save image to temp file for mlx-vlm generate()
            import tempfile
            with tempfile.NamedTemporaryFile(suffix=".png", delete=False) as f:
                img.save(f, format="PNG")
                temp_path = f.name

            # Build prompt via mlx-vlm chat template
            formatted = apply_chat_template(processor, config, prompt_text, num_images=1)

            result = generate(
                model, processor,
                prompt=formatted,
                image=[temp_path],
                max_tokens=max_new_tokens,
                verbose=False,
                temp=0.0,
            )
            output_text = result.text if hasattr(result, 'text') else str(result)

            os.unlink(temp_path)

            predicted = _extract_answer_letter(output_text, len(options))
            correct = predicted == gt_letter

        except Exception as e:
            output_text = f"ERROR: {e}"
            predicted = None
            correct = False

        if correct:
            n_correct += 1

        per_question.append({
            "idx": int(idx),
            "question": question[:100] + ("..." if len(question) > 100 else ""),
            "ground_truth": gt_letter,
            "predicted": predicted,
            "output": str(output_text)[:200],
            "correct": correct,
        })

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

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


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

    right = [q for q in result.per_question if q["correct"]]
    wrong = [q for q in result.per_question if not 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']}")
            print(f"    Output: {q['output'][:100]}")