"""
babilong_eval.metrics — Exact BABILong scoring logic.

Ported directly from the official BABILong evaluation codebase:
  https://github.com/booydar/babilong

The scoring uses LABEL MATCHING, not exact string match:
  1. preprocess_output() cleans the LLM output (lowercase, first sentence)
  2. compare_answers() extracts known task labels from the output,
     subtracts labels that appear in the question, and checks whether
     the target is the ONLY remaining predicted label.

Per AX64 (Source Text Immutability): scoring logic is reproduced
EXACTLY. No modifications to the algorithm.
"""

from __future__ import annotations

import re
from typing import Dict, List, Set, Tuple


# ---------------------------------------------------------------------------
# Task-specific label vocabularies
# ---------------------------------------------------------------------------

TASK_LABELS: Dict[str, List[str]] = {
    "qa1": [
        "bathroom", "bedroom", "garden", "hallway", "kitchen", "office",
    ],
    "qa2": [
        "bathroom", "bedroom", "garden", "hallway", "kitchen", "office",
    ],
    "qa3": [
        "bathroom", "bedroom", "garden", "hallway", "kitchen", "office",
    ],
    "qa4": [
        "bathroom", "bedroom", "garden", "hallway", "kitchen", "office",
        "north", "south", "east", "west",
    ],
    "qa5": [
        "bathroom", "bedroom", "garden", "hallway", "kitchen", "office",
        "bill", "fred", "jeff", "mary",
    ],
}


# ---------------------------------------------------------------------------
# Output preprocessing
# ---------------------------------------------------------------------------

def preprocess_output(output: str) -> str:
    """Preprocess LLM output for label extraction.

    Steps (matching BABILong exactly):
      1. Lowercase
      2. Take only the first sentence (split on period)
      3. Strip everything after "context", "example", or "Question"
      4. Strip whitespace

    Args:
        output: Raw LLM output string.

    Returns:
        Cleaned, lowercased output ready for label extraction.
    """
    text = output.lower().strip()

    # Take first sentence
    sentences = text.split(".")
    if sentences:
        text = sentences[0]

    # Strip after context/example/question markers
    for marker in ("context", "example", "question"):
        idx = text.find(marker)
        if idx != -1:
            text = text[:idx]

    return text.strip()


# ---------------------------------------------------------------------------
# Answer comparison (label matching)
# ---------------------------------------------------------------------------

def compare_answers(
    output: str,
    target: str,
    question: str,
    task: str,
) -> bool:
    """Compare LLM output against the target using BABILong label matching.

    Algorithm (matching BABILong exactly):
      1. Get the task's label vocabulary
      2. Find all labels present in the preprocessed output → predicted_labels
      3. Find all labels present in the question → question_labels
      4. Subtract question_labels from predicted_labels → net_predicted
      5. Return True if target is in net_predicted AND |net_predicted| == 1

    This is stricter than substring matching: the model must produce
    EXACTLY the target label and no other label (after removing labels
    that appeared in the question itself).

    Args:
        output: Raw LLM output string.
        target: Expected answer (1–2 words, always lowercase).
        question: The question that was asked.
        task: Task name (qa1–qa5) for label vocabulary lookup.

    Returns:
        True if the output matches the target per BABILong scoring.
    """
    task_key = task.lower().strip()
    labels = TASK_LABELS.get(task_key, [])
    if not labels:
        # Fallback: exact substring match
        return target.lower().strip() in output.lower()

    # Preprocess
    processed = preprocess_output(output)
    target_clean = target.lower().strip()
    question_clean = question.lower().strip()

    # Extract labels from output
    predicted: Set[str] = set()
    for label in labels:
        if label in processed:
            predicted.add(label)

    # Extract labels from question (to subtract)
    question_labels: Set[str] = set()
    for label in labels:
        if label in question_clean:
            question_labels.add(label)

    # Net predicted = predicted - question_labels
    net_predicted = predicted - question_labels

    # Target must be the ONLY predicted label
    return net_predicted == {target_clean}


# ---------------------------------------------------------------------------
# Batch scoring
# ---------------------------------------------------------------------------

def compute_accuracy(
    results: List[Tuple[str, str, str, str]],
) -> float:
    """Compute accuracy over a batch of (output, target, question, task) tuples.

    Args:
        results: List of (output, target, question, task) tuples.

    Returns:
        Accuracy as a float in [0.0, 1.0].
    """
    if not results:
        return 0.0

    correct = sum(
        1 for output, target, question, task in results
        if compare_answers(output, target, question, task)
    )
    return correct / len(results)
