"""
eval_benchmark.py — Fixed benchmark suite for Phase 5 monotonic improvement proof.

50 tasks across 3 domains:
  - Factual QA (retrieve-needed): 20 tasks
  - Math (no-retrieve): 20 tasks
  - Commonsense (no-retrieve): 10 tasks

Metric: token-level F1 between generated tokens and ground truth.
Evaluation points: 0, 100, 250, 500, 750, 1000 tasks.

Token F1 = 2 * precision * recall / (precision + recall)
  where precision = |gen ∩ ref| / |gen|
        recall    = |gen ∩ ref| / |ref|
"""
from typing import List, Optional
from collections import Counter
import torch

# ── Fixed benchmark tasks ──────────────────────────────────────────────────
# (question, answer, task_type)
# These are fixed at definition time — never shuffled.
BENCHMARK_TASKS = [
    # Factual QA (20)
    ("What is the capital of France?", "Paris", "factual"),
    ("Who wrote Romeo and Juliet?", "William Shakespeare", "factual"),
    ("What is the chemical symbol for gold?", "Au", "factual"),
    ("In what year did World War II end?", "1945", "factual"),
    ("What is the speed of light in km/s?", "299792", "factual"),
    ("Who painted the Mona Lisa?", "Leonardo da Vinci", "factual"),
    ("What planet is closest to the Sun?", "Mercury", "factual"),
    ("What is the largest ocean on Earth?", "Pacific Ocean", "factual"),
    ("Who invented the telephone?", "Alexander Graham Bell", "factual"),
    ("What is the atomic number of carbon?", "6", "factual"),
    ("What country has the largest population?", "China", "factual"),
    ("Who discovered penicillin?", "Alexander Fleming", "factual"),
    ("What is the hardest natural substance?", "diamond", "factual"),
    ("What is the currency of Japan?", "yen", "factual"),
    ("How many bones are in the adult human body?", "206", "factual"),
    ("What is the tallest mountain on Earth?", "Mount Everest", "factual"),
    ("In what year did the Berlin Wall fall?", "1989", "factual"),
    ("What element has the symbol Na?", "sodium", "factual"),
    ("Who wrote 'Pride and Prejudice'?", "Jane Austen", "factual"),
    ("What is the longest river in the world?", "Nile", "factual"),
    # Math (20)
    ("What is 17 multiplied by 13?", "221", "math"),
    ("What is the square root of 144?", "12", "math"),
    ("What is 15% of 200?", "30", "math"),
    ("What is 2 to the power of 10?", "1024", "math"),
    ("What is 99 plus 73?", "172", "math"),
    ("What is 1000 divided by 8?", "125", "math"),
    ("What is the area of a circle with radius 7? (use pi=3.14, round to nearest integer)", "154", "math"),
    ("What is 456 minus 289?", "167", "math"),
    ("What is 7 factorial?", "5040", "math"),
    ("What is 3 cubed?", "27", "math"),
    ("What is 48 divided by 6?", "8", "math"),
    ("What is the sum of angles in a triangle in degrees?", "180", "math"),
    ("What is 12 times 12?", "144", "math"),
    ("What is 25 percent of 80?", "20", "math"),
    ("What is the next prime after 17?", "19", "math"),
    ("What is 1000 minus 437?", "563", "math"),
    ("What is 9 squared?", "81", "math"),
    ("What is 6 times 7?", "42", "math"),
    ("How many seconds in one hour?", "3600", "math"),
    ("What is 11 times 11?", "121", "math"),
    # Commonsense (10)
    ("What do you use to cut bread?", "knife", "commonsense"),
    ("What color is the sky on a clear day?", "blue", "commonsense"),
    ("How many days are in a week?", "7", "commonsense"),
    ("What language is spoken in Brazil?", "Portuguese", "commonsense"),
    ("What do bees produce?", "honey", "commonsense"),
    ("What is the opposite of hot?", "cold", "commonsense"),
    ("How many sides does a hexagon have?", "6", "commonsense"),
    ("What season comes after summer?", "autumn", "commonsense"),
    ("What do plants need to make food?", "sunlight", "commonsense"),
    ("What is water made of?", "hydrogen and oxygen", "commonsense"),
]

assert len(BENCHMARK_TASKS) == 50

# ── Special token IDs ───────────────────────────────────────────────────────
BOS_ID       = 2
USER_ID      = 11
ASSISTANT_ID = 12
EOS_ID       = 3
PAD_ID       = 0
BLOCK_SIZE   = 1024
MAX_GEN      = 80   # max tokens to generate per benchmark answer


# ── Canonical prompt/format helpers ─────────────────────────────────────────
# The tokenizer's post-processor auto-wraps every encode() with BOS/EOS, so
# raw text must always be cleaned with add_special_tokens=False before any
# manual special-token IDs are spliced in. This is the SINGLE canonical
# format used by eval_benchmark.evaluate, online_learner (embed/generate/
# train), and TaskMemory retrieval keys — they MUST stay identical or
# train/inference will diverge silently (see Phase 5 root-cause writeup).
def clean_ids(tokenizer, text: str) -> List[int]:
    """Tokenize text WITHOUT the tokenizer's auto BOS/EOS wrapping."""
    return tokenizer.encode(text, add_special_tokens=False).ids


def build_prompt_ids(tokenizer, question: str) -> List[int]:
    """Canonical prompt: [BOS, USER] + clean(question) + [ASSISTANT].

    Generation/embedding happens immediately AFTER this sequence.
    """
    return [BOS_ID, USER_ID] + clean_ids(tokenizer, question) + [ASSISTANT_ID]


# ── Token F1 ───────────────────────────────────────────────────────────────
def token_f1(prediction_ids: List[int], reference_ids: List[int]) -> float:
    """Token-level F1 between two token ID sequences."""
    if not prediction_ids or not reference_ids:
        return 0.0
    pred_count = Counter(prediction_ids)
    ref_count  = Counter(reference_ids)
    common     = sum((pred_count & ref_count).values())
    if common == 0:
        return 0.0
    precision = common / len(prediction_ids)
    recall    = common / len(reference_ids)
    return 2 * precision * recall / (precision + recall)


# ── Greedy generation ───────────────────────────────────────────────────────
@torch.no_grad()
def generate(model, tokenizer, prompt_ids: List[int],
             max_new: int = MAX_GEN,
             adapter=None) -> List[int]:
    """Greedy decode up to max_new tokens or EOS."""
    ids  = torch.tensor([prompt_ids], dtype=torch.long)
    gen  = []
    for _ in range(max_new):
        if ids.shape[1] >= BLOCK_SIZE:
            break
        logits = model(ids, adapter=adapter)   # [1, T, V]
        next_id = logits[0, -1].argmax().item()
        if next_id == EOS_ID:
            break
        gen.append(next_id)
        ids = torch.cat([ids, torch.tensor([[next_id]])], dim=1)
    return gen


# ── Benchmark evaluation ────────────────────────────────────────────────────
def evaluate(model, tokenizer, memory=None, device='cpu') -> dict:
    """
    Run all 50 benchmark tasks. Returns per-domain and overall F1.
    memory: TaskMemory (used to retrieve adapters); None → base model only.
    """
    model.eval()
    results = {'factual': [], 'math': [], 'commonsense': [], 'all': []}

    for question, answer, task_type in BENCHMARK_TASKS:
        prompt_ids = build_prompt_ids(tokenizer, question)[:BLOCK_SIZE - MAX_GEN]

        adapter = None
        if memory is not None and len(memory) > 0:
            emb     = model.embed_task(torch.tensor([prompt_ids], dtype=torch.long),
                                       adapter=None)
            adapter = memory.retrieve_merged(emb)

        gen_ids = generate(model, tokenizer, prompt_ids,
                           adapter=adapter)
        ref_ids = clean_ids(tokenizer, answer)

        f1 = token_f1(gen_ids, ref_ids)
        results[task_type].append(f1)
        results['all'].append(f1)

    return {
        'overall_f1':     sum(results['all'])         / len(results['all']),
        'factual_f1':     sum(results['factual'])      / len(results['factual']),
        'math_f1':        sum(results['math'])         / len(results['math']),
        'commonsense_f1': sum(results['commonsense'])  / len(results['commonsense']),
        'n_tasks':        len(results['all']),
    }