"""Offline dataset management for threshold calibration.

This module provides utilities for loading, validating, and splitting
offline evaluation data for threshold calibration experiments.
"""

from __future__ import annotations

import logging
from dataclasses import dataclass, field
from pathlib import Path
from typing import Iterator

import numpy as np
import pandas as pd

logger = logging.getLogger(__name__)


@dataclass
class CalibrationSample:
    """Single sample for calibration.

    Attributes:
        query_id: Unique identifier for the query.
        query_text: Original query text.
        task_id: Task/dataset identifier.
        mean_logit: GNN mean logit prediction.
        max_prob: GNN max probability prediction.
        logit_margin: GNN logit margin (top1 - top2).
        logit_std: GNN logit standard deviation.
        entropy: GNN prediction entropy.
        top1_model: Model with highest prediction.
        top1_prob: Probability of top model.
        ground_truth_effects: Dict mapping model -> ground truth effect.
        oracle_model: Best model according to ground truth.
        oracle_effect: Best effect value.
    """

    query_id: str
    query_text: str
    task_id: str
    mean_logit: float
    max_prob: float
    logit_margin: float
    logit_std: float
    entropy: float
    top1_model: str
    top1_prob: float
    ground_truth_effects: dict[str, float]
    oracle_model: str
    oracle_effect: float

    @property
    def is_trivial_ground_truth(self) -> bool:
        """Query is trivial if >90% of models succeed."""
        threshold = 0.5
        n_correct = sum(1 for e in self.ground_truth_effects.values() if e >= threshold)
        return n_correct >= 0.9 * len(self.ground_truth_effects)

    @property
    def is_hard_ground_truth(self) -> bool:
        """Query is hard if <30% of models succeed."""
        threshold = 0.5
        n_correct = sum(1 for e in self.ground_truth_effects.values() if e >= threshold)
        return n_correct <= 0.3 * len(self.ground_truth_effects)

    @property
    def perf_range(self) -> float:
        """Performance range (max - min effect)."""
        effects = list(self.ground_truth_effects.values())
        return max(effects) - min(effects) if effects else 0.0


@dataclass
class DatasetSplit:
    """A split of the calibration dataset (train/val/test).

    Attributes:
        name: Split name (e.g., "train", "val", "test", "fold_0").
        samples: List of calibration samples.
        metadata: Optional metadata about the split.
    """

    name: str
    samples: list[CalibrationSample]
    metadata: dict = field(default_factory=dict)

    def __len__(self) -> int:
        return len(self.samples)

    def __iter__(self) -> Iterator[CalibrationSample]:
        return iter(self.samples)

    @property
    def query_ids(self) -> list[str]:
        return [s.query_id for s in self.samples]

    @property
    def mean_logits(self) -> np.ndarray:
        return np.array([s.mean_logit for s in self.samples])

    @property
    def max_probs(self) -> np.ndarray:
        return np.array([s.max_prob for s in self.samples])

    @property
    def logit_margins(self) -> np.ndarray:
        return np.array([s.logit_margin for s in self.samples])

    @property
    def oracle_effects(self) -> np.ndarray:
        return np.array([s.oracle_effect for s in self.samples])

    def filter_by_task(self, task_id: str) -> "DatasetSplit":
        """Return subset for a specific task."""
        filtered = [s for s in self.samples if s.task_id == task_id]
        return DatasetSplit(
            name=f"{self.name}_{task_id}",
            samples=filtered,
            metadata={**self.metadata, "task_filter": task_id},
        )

    def sample(self, n: int, random_state: int = 42) -> "DatasetSplit":
        """Random sample of n samples."""
        rng = np.random.RandomState(random_state)
        indices = rng.choice(len(self.samples), size=min(n, len(self.samples)), replace=False)
        sampled = [self.samples[i] for i in indices]
        return DatasetSplit(
            name=f"{self.name}_sample_{n}",
            samples=sampled,
            metadata={**self.metadata, "sampled": True, "sample_size": n},
        )


@dataclass
class OfflineDataset:
    """Complete offline dataset for calibration.

    This class manages the full dataset and provides utilities for
    splitting, cross-validation, and analysis.

    Attributes:
        samples: All calibration samples.
        model_names: List of model names in the dataset.
        task_ids: List of unique task IDs.
        metadata: Dataset metadata.
    """

    samples: list[CalibrationSample]
    model_names: list[str]
    task_ids: list[str]
    metadata: dict = field(default_factory=dict)

    def __len__(self) -> int:
        return len(self.samples)

    def __iter__(self) -> Iterator[CalibrationSample]:
        return iter(self.samples)

    @classmethod
    def from_csv(
        cls,
        path: str | Path,
        signal_columns: dict[str, str] | None = None,
        effect_prefix: str = "effect_",
    ) -> "OfflineDataset":
        """Load dataset from CSV file.

        Expected CSV format:
        - query_id: Unique query identifier
        - query_text: Query text
        - task_id: Task/dataset identifier
        - mean_logit, max_prob, logit_margin, logit_std, entropy: GNN signals
        - top1_model, top1_prob: Top prediction
        - effect_<model_name>: Ground truth effect for each model

        Args:
            path: Path to CSV file.
            signal_columns: Mapping from expected column names to actual names.
            effect_prefix: Prefix for effect columns.

        Returns:
            OfflineDataset instance.
        """
        df = pd.read_csv(path)
        return cls._from_dataframe(df, signal_columns, effect_prefix)

    @classmethod
    def from_parquet(
        cls,
        path: str | Path,
        signal_columns: dict[str, str] | None = None,
        effect_prefix: str = "effect_",
    ) -> "OfflineDataset":
        """Load dataset from Parquet file."""
        df = pd.read_parquet(path)
        return cls._from_dataframe(df, signal_columns, effect_prefix)

    @classmethod
    def from_router_data(
        cls,
        router_data_path: str | Path,
        gnn_predictions_path: str | Path,
    ) -> "OfflineDataset":
        """Load from router_data.csv and GNN predictions.

        This is the standard format for EmberGraphRouter training data.

        Args:
            router_data_path: Path to router_data.csv with ground truth effects.
            gnn_predictions_path: Path to CSV with GNN signal predictions.

        Returns:
            OfflineDataset instance.
        """
        # Load router data (ground truth)
        router_df = pd.read_csv(router_data_path)

        # Pivot to wide format: query -> {model: effect}
        pivot_df = router_df.pivot_table(
            index=["query", "task_id"],
            columns="llm",
            values="effect",
            aggfunc="first",
        ).reset_index()

        # Load GNN predictions
        predictions_df = pd.read_csv(gnn_predictions_path)

        # Merge
        merged_df = pivot_df.merge(predictions_df, on="query", how="inner")

        # Extract model names from pivot columns
        model_names = [c for c in pivot_df.columns if c not in ["query", "task_id"]]

        # Rename effect columns
        for model in model_names:
            if model in merged_df.columns:
                merged_df = merged_df.rename(columns={model: f"effect_{model}"})

        return cls._from_dataframe(merged_df, effect_prefix="effect_")

    @classmethod
    def _from_dataframe(
        cls,
        df: pd.DataFrame,
        signal_columns: dict[str, str] | None = None,
        effect_prefix: str = "effect_",
    ) -> "OfflineDataset":
        """Internal method to construct dataset from DataFrame."""
        # Default column mapping
        col_map = {
            "query_id": "query_id",
            "query_text": "query_text",
            "task_id": "task_id",
            "mean_logit": "mean_logit",
            "max_prob": "max_prob",
            "logit_margin": "logit_margin",
            "logit_std": "logit_std",
            "entropy": "entropy",
            "top1_model": "top1_model",
            "top1_prob": "top1_prob",
        }
        if signal_columns:
            col_map.update(signal_columns)

        # Handle alternative column names
        if "query" in df.columns and "query_id" not in df.columns:
            df = df.rename(columns={"query": "query_id"})
        if "query_id" in df.columns and "query_text" not in df.columns:
            df["query_text"] = df["query_id"]

        # Extract effect columns
        effect_cols = [c for c in df.columns if c.startswith(effect_prefix)]
        model_names = [c.replace(effect_prefix, "") for c in effect_cols]

        if not model_names:
            raise ValueError(f"No effect columns found with prefix '{effect_prefix}'")

        logger.info(
            "Loading dataset: %d samples, %d models: %s",
            len(df),
            len(model_names),
            model_names,
        )

        samples = []
        for _, row in df.iterrows():
            # Extract ground truth effects
            effects = {}
            for model, col in zip(model_names, effect_cols, strict=False):
                if pd.notna(row.get(col)):
                    effects[model] = float(row[col])

            if not effects:
                continue

            # Find oracle
            oracle_model = max(effects, key=effects.get)
            oracle_effect = effects[oracle_model]

            sample = CalibrationSample(
                query_id=str(row.get(col_map["query_id"], "")),
                query_text=str(row.get(col_map["query_text"], "")),
                task_id=str(row.get(col_map["task_id"], "unknown")),
                mean_logit=float(row.get(col_map["mean_logit"], 0.0)),
                max_prob=float(row.get(col_map["max_prob"], 0.5)),
                logit_margin=float(row.get(col_map["logit_margin"], 0.5)),
                logit_std=float(row.get(col_map["logit_std"], 0.5)),
                entropy=float(row.get(col_map["entropy"], 0.5)),
                top1_model=str(row.get(col_map["top1_model"], oracle_model)),
                top1_prob=float(row.get(col_map["top1_prob"], 0.5)),
                ground_truth_effects=effects,
                oracle_model=oracle_model,
                oracle_effect=oracle_effect,
            )
            samples.append(sample)

        if col_map["task_id"] in df.columns:
            task_ids = list(df[col_map["task_id"]].unique())
        else:
            task_ids = ["unknown"]

        return cls(
            samples=samples,
            model_names=model_names,
            task_ids=task_ids,
            metadata={
                "source": "dataframe",
                "n_samples": len(samples),
                "n_models": len(model_names),
            },
        )

    def train_test_split(
        self,
        test_fraction: float = 0.2,
        stratify_by_task: bool = True,
        random_state: int = 42,
    ) -> tuple[DatasetSplit, DatasetSplit]:
        """Split dataset into train and test sets.

        Args:
            test_fraction: Fraction of data for test set.
            stratify_by_task: Whether to stratify by task_id.
            random_state: Random seed.

        Returns:
            Tuple of (train_split, test_split).
        """
        rng = np.random.RandomState(random_state)

        if stratify_by_task:
            train_samples = []
            test_samples = []

            for task_id in self.task_ids:
                task_samples = [s for s in self.samples if s.task_id == task_id]
                rng.shuffle(task_samples)

                n_test = max(1, int(len(task_samples) * test_fraction))
                test_samples.extend(task_samples[:n_test])
                train_samples.extend(task_samples[n_test:])
        else:
            indices = np.arange(len(self.samples))
            rng.shuffle(indices)

            n_test = int(len(indices) * test_fraction)
            test_indices = indices[:n_test]
            train_indices = indices[n_test:]

            train_samples = [self.samples[i] for i in train_indices]
            test_samples = [self.samples[i] for i in test_indices]

        return (
            DatasetSplit(
                name="train",
                samples=train_samples,
                metadata={"split": "train", "fraction": 1 - test_fraction},
            ),
            DatasetSplit(
                name="test",
                samples=test_samples,
                metadata={"split": "test", "fraction": test_fraction},
            ),
        )

    def cross_validation_splits(
        self,
        n_folds: int = 5,
        stratify_by_task: bool = True,
        random_state: int = 42,
    ) -> list[tuple[DatasetSplit, DatasetSplit]]:
        """Generate cross-validation splits.

        Args:
            n_folds: Number of CV folds.
            stratify_by_task: Whether to stratify by task_id.
            random_state: Random seed.

        Returns:
            List of (train_split, val_split) tuples.
        """
        rng = np.random.RandomState(random_state)

        if stratify_by_task:
            # Assign fold per task
            task_to_samples: dict[str, list[CalibrationSample]] = {}
            for sample in self.samples:
                task_to_samples.setdefault(sample.task_id, []).append(sample)

            # Shuffle within tasks
            for task_samples in task_to_samples.values():
                rng.shuffle(task_samples)

            # Create folds
            folds: list[list[CalibrationSample]] = [[] for _ in range(n_folds)]
            for task_samples in task_to_samples.values():
                for i, sample in enumerate(task_samples):
                    folds[i % n_folds].append(sample)
        else:
            indices = np.arange(len(self.samples))
            rng.shuffle(indices)

            fold_size = len(indices) // n_folds
            folds = []
            for i in range(n_folds):
                start = i * fold_size
                end = start + fold_size if i < n_folds - 1 else len(indices)
                fold_indices = indices[start:end]
                folds.append([self.samples[j] for j in fold_indices])

        # Create train/val pairs
        splits = []
        for i in range(n_folds):
            val_samples = folds[i]
            train_samples = []
            for j in range(n_folds):
                if j != i:
                    train_samples.extend(folds[j])

            splits.append((
                DatasetSplit(
                    name=f"train_fold_{i}",
                    samples=train_samples,
                    metadata={"fold": i, "split": "train"},
                ),
                DatasetSplit(
                    name=f"val_fold_{i}",
                    samples=val_samples,
                    metadata={"fold": i, "split": "val"},
                ),
            ))

        return splits

    def summary(self) -> dict:
        """Return summary statistics about the dataset."""
        mean_logits = np.array([s.mean_logit for s in self.samples])
        max_probs = np.array([s.max_prob for s in self.samples])
        oracle_effects = np.array([s.oracle_effect for s in self.samples])

        n_trivial = sum(1 for s in self.samples if s.is_trivial_ground_truth)
        n_hard = sum(1 for s in self.samples if s.is_hard_ground_truth)

        return {
            "n_samples": len(self.samples),
            "n_models": len(self.model_names),
            "n_tasks": len(self.task_ids),
            "model_names": self.model_names,
            "task_ids": self.task_ids,
            "signal_stats": {
                "mean_logit": {
                    "mean": float(np.mean(mean_logits)),
                    "std": float(np.std(mean_logits)),
                    "min": float(np.min(mean_logits)),
                    "max": float(np.max(mean_logits)),
                },
                "max_prob": {
                    "mean": float(np.mean(max_probs)),
                    "std": float(np.std(max_probs)),
                    "min": float(np.min(max_probs)),
                    "max": float(np.max(max_probs)),
                },
                "oracle_effect": {
                    "mean": float(np.mean(oracle_effects)),
                    "std": float(np.std(oracle_effects)),
                },
            },
            "difficulty_distribution": {
                "trivial_pct": n_trivial / len(self.samples) * 100,
                "hard_pct": n_hard / len(self.samples) * 100,
                "medium_pct": (len(self.samples) - n_trivial - n_hard) / len(self.samples) * 100,
            },
        }
