"""Evaluator: runs a model on a dataset and computes metrics."""

from __future__ import annotations

import torch
from tqdm import tqdm

from obliteratus.models.loader import ModelHandle


class Evaluator:
    """Evaluate a model handle on a dataset, returning metric results.

    Supports two modes:
      - **perplexity** (default for causal_lm): feeds tokenized text and computes PPL.
      - **classification**: runs forward pass, takes argmax, computes accuracy/F1.
    """

    def __init__(
        self,
        handle: ModelHandle,
        dataset,
        metrics: list[str] | None = None,
        batch_size: int = 8,
        max_length: int = 512,
        max_samples: int | None = None,
        text_column: str = "text",
        label_column: str = "label",
    ):
        self.handle = handle
        self.dataset = dataset
        self.metrics = metrics or (
            ["perplexity"] if handle.task == "causal_lm" else ["accuracy", "f1"]
        )
        self.batch_size = batch_size
        self.max_length = max_length
        self.max_samples = max_samples
        self.text_column = text_column
        self.label_column = label_column

    @torch.no_grad()
    def evaluate(self) -> dict[str, float]:
        """Run evaluation and return a dict of metric_name -> score."""
        if self.handle.task == "causal_lm":
            return self._evaluate_causal_lm()
        elif self.handle.task == "classification":
            return self._evaluate_classification()
        else:
            raise ValueError(f"Unsupported task: {self.handle.task}")

    def _evaluate_causal_lm(self) -> dict[str, float]:
        model = self.handle.model
        tokenizer = self.handle.tokenizer
        device = next(model.parameters()).device

        ds = self.dataset

        if self.max_samples is not None and self.max_samples <= 0:
            raise ValueError("max_samples must be a positive integer or None.")

        # WikiText and similar corpora contain empty separator rows. Walk the
        # map-style dataset row by row so a bounded evaluation stops as soon as
        # it has enough usable texts instead of materializing the whole column.
        valid_indices = []
        for index in range(len(ds)):
            text = ds[index][self.text_column]
            if not isinstance(text, str) or not text.strip():
                continue
            valid_indices.append(index)
            if self.max_samples is not None and len(valid_indices) >= self.max_samples:
                break

        if not valid_indices:
            raise ValueError(
                f"Dataset contains no non-empty text in column "
                f"{self.text_column!r}."
            )

        ds = ds.select(valid_indices)

        total_loss = 0.0
        total_tokens = 0
        skipped_batches = 0

        for i in tqdm(
            range(0, len(ds), self.batch_size),
            desc="Evaluating PPL",
        ):
            batch_texts = ds[i : i + self.batch_size][self.text_column]

            # Defensive filtering in case a custom dataset returns unexpected
            # values after selection or transformation.
            batch_texts = [
                text
                for text in batch_texts
                if isinstance(text, str) and text.strip()
            ]

            if not batch_texts:
                skipped_batches += 1
                continue

            encodings = tokenizer(
                batch_texts,
                return_tensors="pt",
                truncation=True,
                max_length=self.max_length,
                padding=True,
            ).to(device)

            input_ids = encodings["input_ids"]
            attention_mask = encodings["attention_mask"]

            # A causal-LM loss requires at least one prediction target after
            # shifting. A sequence length below two cannot contribute.
            if (
                input_ids.ndim != 2
                or input_ids.numel() == 0
                or input_ids.shape[1] < 2
            ):
                skipped_batches += 1
                continue

            labels = input_ids.clone()

            # Hugging Face causal-LM losses ignore labels set to -100.
            # This prevents padded positions from affecting perplexity.
            labels[attention_mask == 0] = -100

            num_tokens = labels[:, 1:].ne(-100).sum().item()

            if num_tokens <= 0:
                skipped_batches += 1
                continue

            outputs = model(
                input_ids=input_ids,
                attention_mask=attention_mask,
                labels=labels,
            )

            if outputs.loss is None:
                raise RuntimeError(
                    "The causal language model returned no loss value."
                )

            if not torch.isfinite(outputs.loss):
                raise RuntimeError(
                    f"Non-finite evaluation loss encountered at batch "
                    f"offset {i}: {outputs.loss.item()}"
                )

            total_loss += outputs.loss.item() * num_tokens
            total_tokens += num_tokens

        if total_tokens <= 0:
            raise RuntimeError(
                "Perplexity evaluation produced zero valid prediction tokens. "
                "Check the dataset text column and tokenizer configuration."
            )

        if skipped_batches:
            tqdm.write(
                f"Skipped {skipped_batches} empty or too-short batch(es)."
            )

        import math

        avg_loss = total_loss / total_tokens
        return {"perplexity": math.exp(avg_loss)}

    def _evaluate_classification(self) -> dict[str, float]:
        from obliteratus.evaluation.metrics import accuracy as acc_fn
        from obliteratus.evaluation.metrics import f1_score_metric as f1_fn

        model = self.handle.model
        tokenizer = self.handle.tokenizer
        device = next(model.parameters()).device

        ds = self.dataset
        if self.max_samples is not None:
            ds = ds.select(range(min(self.max_samples, len(ds))))

        all_preds = []
        all_labels = []

        for i in tqdm(range(0, len(ds), self.batch_size), desc="Evaluating"):
            batch = ds[i : i + self.batch_size]
            texts = batch[self.text_column]
            labels = batch[self.label_column]

            encodings = tokenizer(
                texts,
                return_tensors="pt",
                truncation=True,
                max_length=self.max_length,
                padding=True,
            ).to(device)

            outputs = model(**encodings)
            preds = outputs.logits.argmax(dim=-1).cpu().tolist()
            all_preds.extend(preds)
            all_labels.extend(labels)

        results = {}
        if "accuracy" in self.metrics:
            results["accuracy"] = acc_fn(all_preds, all_labels)
        if "f1" in self.metrics:
            results["f1"] = f1_fn(all_preds, all_labels)
        return results