"""Zero-shot task adapters + metrics for the fetched eval sets.

Two task shapes, both scored purely by LM likelihood:

* Minimal pairs (BLiMP): accuracy = fraction where the grammatical sentence gets
  a higher total log-probability than the ungrammatical one.
* Multiple choice (COPA / XCOPA): accuracy = fraction where the gold choice gets
  the highest (length-normalised) continuation log-probability.

Loaders read the immutable raw parquet directly (pyarrow), so they do not depend
on dataset loading scripts.
"""

from __future__ import annotations

import glob
from dataclasses import dataclass, field
from pathlib import Path

import torch

from strata.eval.scoring import encode_with_continuation, score_sequences

# Cause/effect connectors for building COPA/XCOPA prompts, per language.
_CONNECTORS = {
    "en": {"cause": "because", "effect": "so"},
    "zh": {"cause": "因为", "effect": "所以"},
}


@dataclass(frozen=True, slots=True)
class MinimalPairExample:
    good: str
    bad: str
    tag: str = ""


@dataclass(frozen=True, slots=True)
class MultipleChoiceExample:
    context: str
    choices: tuple[str, ...]
    gold: int


@dataclass(slots=True)
class TaskResult:
    task: str
    metric: str
    value: float
    n_examples: int
    subscores: dict[str, float] = field(default_factory=dict)

    def to_dict(self) -> dict[str, object]:
        return {
            "task": self.task,
            "metric": self.metric,
            "value": self.value,
            "n_examples": self.n_examples,
            "subscores": self.subscores,
        }


# ---------------------------------------------------------------------------
# Evaluation drivers
# ---------------------------------------------------------------------------


def evaluate_minimal_pairs(
    model,
    tokenizer,
    examples: list[MinimalPairExample],
    *,
    task: str,
    device: torch.device,
    pad_id: int,
    batch_size: int = 16,
    precision: str = "bf16",
    predicate_memory_intervention: str = "none",
    predicate_memory_residual_scale: float | None = None,
    graph_object_residual_scale: float | None = None,
) -> TaskResult:
    """Total-log-prob comparison of good vs. bad sentences."""

    if not examples:
        raise ValueError(f"{task}: no examples to evaluate")
    seqs: list[list[int]] = []
    ctx: list[int] = []
    for ex in examples:
        for text in (ex.good, ex.bad):
            ids = list(tokenizer.encode(text, add_bos=True).input_ids)
            seqs.append(ids)
            ctx.append(1)  # score every real token (position 0 is BOS)
    scores = score_sequences(
        model, seqs, ctx, device=device, pad_id=pad_id, batch_size=batch_size, precision=precision,
        predicate_memory_intervention=predicate_memory_intervention,
        predicate_memory_residual_scale=predicate_memory_residual_scale,
        graph_object_residual_scale=graph_object_residual_scale,
    )
    correct = 0
    per_tag: dict[str, list[int]] = {}
    for i, ex in enumerate(examples):
        good_score, bad_score = scores[2 * i], scores[2 * i + 1]
        hit = int(good_score > bad_score)
        correct += hit
        if ex.tag:
            per_tag.setdefault(ex.tag, []).append(hit)
    subscores = {tag: sum(v) / len(v) for tag, v in sorted(per_tag.items())}
    return TaskResult(task, "accuracy", correct / len(examples), len(examples), subscores)


def evaluate_multiple_choice(
    model,
    tokenizer,
    examples: list[MultipleChoiceExample],
    *,
    task: str,
    device: torch.device,
    pad_id: int,
    batch_size: int = 16,
    precision: str = "bf16",
    length_normalize: bool = True,
    predicate_memory_intervention: str = "none",
    predicate_memory_residual_scale: float | None = None,
    graph_object_residual_scale: float | None = None,
) -> TaskResult:
    """Pick the choice with the highest (length-normalised) continuation log-prob."""

    if not examples:
        raise ValueError(f"{task}: no examples to evaluate")
    seqs: list[list[int]] = []
    ctx: list[int] = []
    spans: list[tuple[int, int]] = []  # (start, end) index into seqs per example
    for ex in examples:
        start = len(seqs)
        for choice in ex.choices:
            ids, context_len = encode_with_continuation(tokenizer, ex.context, choice)
            seqs.append(ids)
            ctx.append(context_len)
        spans.append((start, len(seqs)))
    scores = score_sequences(
        model, seqs, ctx, device=device, pad_id=pad_id, batch_size=batch_size,
        precision=precision, length_normalize=length_normalize,
        predicate_memory_intervention=predicate_memory_intervention,
        predicate_memory_residual_scale=predicate_memory_residual_scale,
        graph_object_residual_scale=graph_object_residual_scale,
    )
    correct = 0
    for ex, (start, end) in zip(examples, spans):
        choice_scores = scores[start:end]
        pred = max(range(len(choice_scores)), key=lambda k: choice_scores[k])
        correct += int(pred == ex.gold)
    return TaskResult(task, "accuracy", correct / len(examples), len(examples))


# ---------------------------------------------------------------------------
# Loaders (read raw parquet directly)
# ---------------------------------------------------------------------------


def _read_parquet_rows(path: str) -> list[dict]:
    import pyarrow.parquet as pq

    return pq.ParquetFile(path).read().to_pylist()


def load_blimp(raw_dir: Path, *, configs: list[str] | None = None, max_per_config: int | None = None) -> list[MinimalPairExample]:
    examples: list[MinimalPairExample] = []
    paths = sorted(glob.glob(str(raw_dir / "*" / "*.parquet")))
    if not paths:
        raise FileNotFoundError(f"no BLiMP parquet under {raw_dir}")
    for path in paths:
        uid = Path(path).parent.name
        if configs is not None and uid not in configs:
            continue
        rows = _read_parquet_rows(path)
        if max_per_config is not None:
            rows = rows[:max_per_config]
        for row in rows:
            examples.append(MinimalPairExample(good=row["sentence_good"], bad=row["sentence_bad"], tag=uid))
    return examples


def _copa_context(premise: str, question: str, language: str) -> str:
    premise = premise.rstrip().rstrip(".。").strip()
    connectors = _CONNECTORS.get(language, _CONNECTORS["en"])
    connector = connectors.get(str(question), connectors["effect"])
    joiner = "" if language == "zh" else " "
    return f"{premise}{joiner}{connector}{joiner}"


def load_copa(parquet_path: str, *, language: str = "en", max_examples: int | None = None) -> list[MultipleChoiceExample]:
    rows = _read_parquet_rows(parquet_path)
    rows = [r for r in rows if int(r.get("label", -1)) in (0, 1)]  # drop unlabeled test rows
    if max_examples is not None:
        rows = rows[:max_examples]
    examples: list[MultipleChoiceExample] = []
    for row in rows:
        context = _copa_context(row["premise"], row["question"], language)
        examples.append(
            MultipleChoiceExample(
                context=context,
                choices=(str(row["choice1"]), str(row["choice2"])),
                gold=int(row["label"]),
            )
        )
    return examples