"""VQA-style benchmark evaluation for ArcisVLM.

Supports: VQAv2, GQA, TextVQA, ScienceQA, A-OKVQA, OKVQA
"""

import torch
from typing import Optional


def normalize_answer(answer: str) -> str:
    """Normalize answer for VQA accuracy computation."""
    answer = answer.strip().lower()
    # Remove articles
    for article in ["a ", "an ", "the "]:
        if answer.startswith(article):
            answer = answer[len(article):]
    # Remove punctuation
    answer = answer.rstrip(".")
    return answer


def vqa_accuracy(pred: str, targets: list[str]) -> float:
    """VQA accuracy: min(count(pred in targets) / 3, 1.0)"""
    pred_norm = normalize_answer(pred)
    count = sum(1 for t in targets if normalize_answer(t) == pred_norm)
    return min(count / 3.0, 1.0)


def evaluate_vqa(model, dataset, tokenizer, device="cuda",
                 max_samples: int = None, batch_size: int = 32) -> dict:
    """Evaluate model on a VQA-style dataset.

    Args:
        model: VLJEPAModel instance
        dataset: Dataset yielding (image, question, answers)
        tokenizer: BPE tokenizer
        device: Device to run on
        max_samples: Limit evaluation samples
        batch_size: Batch size for inference

    Returns:
        dict with "accuracy", "num_samples", "predictions"
    """
    model.eval()
    total_acc = 0.0
    num_samples = 0
    predictions = []

    n = min(len(dataset), max_samples) if max_samples else len(dataset)

    for i in range(n):
        sample = dataset[i]
        image = sample["image"].unsqueeze(0).to(device)
        question = sample.get("question", sample.get("instruction", ""))
        answers = sample.get("answers", [sample.get("answer", "")])
        if isinstance(answers, str):
            answers = [answers]

        # Encode question
        query_ids = torch.tensor(
            [tokenizer.encode(question)], dtype=torch.long, device=device
        )

        # Generate
        with torch.no_grad():
            output_ids = model.generate(image, query_ids, max_new_tokens=32, temperature=0.1)

        pred_text = tokenizer.decode(output_ids[0].cpu().tolist())

        acc = vqa_accuracy(pred_text, answers)
        total_acc += acc
        num_samples += 1

        predictions.append({
            "question": question,
            "prediction": pred_text,
            "answers": answers,
            "accuracy": acc,
        })

    return {
        "accuracy": total_acc / max(num_samples, 1) * 100,
        "num_samples": num_samples,
        "predictions": predictions,
    }