"""Dataset-native trainer that compiles router checkpoints."""

from __future__ import annotations

from dataclasses import dataclass
from typing import Literal

import numpy as np

from nordlys.checkpoint import build_checkpoint
from nordlys.clustering import Clusterer, KMeansClusterer, compute_cluster_metrics
from nordlys.dataset import Dataset
from nordlys.embeddings import Embedder, SentenceTransformers
from nordlys.reduction import Reducer, restore_reducer
from nordlys.reduction.base import ReductionPayload
from nordlys.router import Router


@dataclass(frozen=True)
class FittedStructure:
    """Result of learning task structure from data.

    Attributes:
        cluster_centers: Centroid vectors of shape (n_clusters, n_features)
        n_clusters: Number of clusters
        embedding_config: Embedding model configuration
        clustering_config: Clustering algorithm configuration
        reduction: Optional dimensionality reduction payload
        metrics: Clustering quality metrics
    """

    cluster_centers: np.ndarray
    n_clusters: int
    embedding_config: dict
    clustering_config: dict
    reduction: ReductionPayload | None
    metrics: dict

    @property
    def centroids(self) -> np.ndarray:
        """Alias for cluster_centers for backwards compatibility."""
        return self.cluster_centers


@dataclass(frozen=True)
class RoutingPolicy:
    """Result of calibrating model quality per cluster.

    Attributes:
        scores: Dict mapping model_id to list of per-cluster scores (0-1, higher is better)
        sample_counts: Number of calibration samples per cluster
    """

    scores: dict[str, list[float]]
    sample_counts: list[int]

    def model_scores(self, model_id: str) -> list[float]:
        """Get scores for a specific model."""
        return self.scores.get(model_id, [])

    def best_model_per_cluster(self) -> list[str]:
        """Get the best model (highest score) for each cluster."""
        best = []
        for cluster_idx in range(len(self.sample_counts)):
            best_model = None
            best_score = -1.0
            for model_id, model_scores in self.scores.items():
                if cluster_idx < len(model_scores):
                    if model_scores[cluster_idx] > best_score:
                        best_score = model_scores[cluster_idx]
                        best_model = model_id
            best.append(best_model if best_model is not None else "")
        return best


@dataclass(frozen=True)
class EvaluationReport:
    """Result of evaluating a router on held-out data.

    Attributes:
        routing_decisions: List of route results for each evaluation sample
        metrics: Evaluation metrics such as accuracy
    """

    routing_decisions: list
    metrics: dict


@dataclass(frozen=True)
class Trainer:
    """Train router checkpoints from Dataset objects.

    Canonical path: ``Trainer -> fit_clusters() / score_clusters() / build() -> Router``.

    Example:
        >>> from nordlys import Dataset, Trainer
        >>> from nordlys.clustering import KMeansClusterer
        >>> train_ds = Dataset.from_list([
        ...     {"id": "1", "input": "Fix the parser"},
        ... ])
        >>> val_ds = Dataset.from_list([
        ...     {"id": "2", "input": "Write tests", "targets": {"gpt-4": 1}},
        ... ])
        >>> # Pass your own clusterer
        >>> trainer = Trainer(
        ...     models=["gpt-4"],
        ...     clusterer=KMeansClusterer(n_clusters=10),
        ... )
        >>> fitted = trainer.fit_clusters(train_ds)
        >>> scored = trainer.score_clusters(fitted, val_ds)
        >>> router = trainer.build(fitted, scored)
        >>> result = router.route("Explain quantum computing")
    """

    models: list[str]
    input_col: str = "input"
    target_col: str = "targets"

    embedding_model: str = "sentence-transformers/all-MiniLM-L6-v2"
    embedding_batch_size: int = 64
    embedding_normalize: bool = True
    embedder: Embedder | None = None

    clusterer: Clusterer | None = None
    n_clusters: int = 20

    reducer: Reducer | None = None

    random_state: int = 42
    allow_trust_remote_code: bool = False
    device: Literal["cpu", "cuda"] = "cpu"

    def fit_structure(self, dataset: Dataset) -> FittedStructure:
        """Fit clustering on training data.

        This is the first step in a split workflow where you cluster on one dataset
        and score models on another.

        Args:
            dataset: Training dataset with 'id' and 'input' columns

        Returns:
            FittedClustering object containing cluster state
        """
        self._validate_fit_clusters(dataset)

        embedder = self._make_embedder()
        embeddings = embedder.encode(dataset.column(self.input_col))
        cluster_input, reduction_payload = self._reduce_or_pass(embeddings)

        clusterer = self._make_clusterer()
        clusterer.fit(cluster_input)

        centroids = clusterer.cluster_centers_
        labels = clusterer.labels_
        n_clusters = len(centroids)

        inertia: float | None = None
        if isinstance(clusterer, KMeansClusterer):
            inertia = clusterer.inertia_

        metrics = compute_cluster_metrics(
            cluster_input,
            labels,
            inertia,
        )

        return FittedStructure(
            cluster_centers=np.asarray(centroids, dtype=np.float32),
            n_clusters=n_clusters,
            embedding_config=embedder.checkpoint_config(),
            clustering_config={
                "n_clusters": n_clusters,
                "random_state": self.random_state,
                "max_iter": 300,
                "n_init": 10,
                "algorithm": "lloyd",
                "normalization": "l2" if self.embedding_normalize else "none",
            },
            reduction=reduction_payload,
            metrics={
                "n_samples": metrics.n_samples,
                "cluster_sizes": metrics.cluster_sizes,
                "silhouette_score": metrics.silhouette_score,
                "inertia": metrics.inertia,
            },
        )

    def assign_clusters(self, fitted: FittedStructure, dataset: Dataset) -> np.ndarray:
        """Assign dataset instances to clusters using fitted centroids.

        Uses nearest-centroid assignment.

        Args:
            fitted: Fitted clustering from fit_clusters()
            dataset: Dataset with 'id' and 'input' columns

        Returns:
            Array of cluster assignments of shape (n_samples,)
        """
        embedder = self._make_embedder()
        embeddings = embedder.encode(dataset.column(self.input_col))

        # Apply the same reduction as used during fitting
        reducer = restore_reducer(fitted.reduction)
        if reducer is not None:
            embeddings = reducer.transform(embeddings)

        # Compute distances to centroids
        # embeddings: (n_samples, n_features), centroids: (n_clusters, n_features)
        # result: (n_samples, n_clusters)
        centroids = fitted.cluster_centers
        distances = np.sqrt(
            np.sum(
                (embeddings[:, np.newaxis, :] - centroids[np.newaxis, :, :]) ** 2,
                axis=2,
            )
        )
        labels = np.argmin(distances, axis=1)

        return labels

    def calibrate(self, fitted: FittedStructure, dataset: Dataset) -> RoutingPolicy:
        """Calibrate routing policy using fitted structure and validation data.

        This computes per-cluster model scores from calibration data.

        Args:
            fitted: Fitted structure from fit_structure()
            dataset: Calibration dataset with 'id', 'input', and 'targets' columns

        Returns:
            RoutingPolicy object with per-cluster model scores
        """
        self._validate_score_clusters(dataset)

        # Assign calibration data to clusters
        labels = self.assign_clusters(fitted, dataset)

        # Compute scores from assigned labels
        targets = dataset.column(self.target_col)
        n_clusters = fitted.n_clusters

        scores = self._calc_scores(
            labels,
            targets,
            self.models,
            n_clusters,
        )

        # Compute sample counts per cluster
        sample_counts = [0] * n_clusters
        for label in labels:
            if 0 <= label < n_clusters:
                sample_counts[label] += 1

        return RoutingPolicy(
            scores=scores,
            sample_counts=sample_counts,
        )

    def compile(self, fitted: FittedStructure, policy: RoutingPolicy) -> Router:
        """Compile a router from structure and policy.

        Args:
            fitted: Fitted structure from fit_structure()
            policy: Routing policy from calibrate()

        Returns:
            Router ready for routing
        """
        models_payload = [
            {
                "model_id": model_id,
                "scores": policy.scores[model_id],
            }
            for model_id in self.models
        ]

        checkpoint = build_checkpoint(
            cluster_centers=fitted.cluster_centers,
            models=models_payload,
            embedding=fitted.embedding_config,
            clustering=fitted.clustering_config,
            reduction=fitted.reduction,
            metrics=fitted.metrics,
        )

        return Router(checkpoint=checkpoint, device=self.device)

    def fit_router(
        self, structure_dataset: Dataset, calibration_dataset: Dataset
    ) -> Router:
        """Convenience method: fit structure, calibrate, compile.

        Args:
            structure_dataset: Dataset for learning task structure
            calibration_dataset: Dataset for calibrating model quality per cluster

        Returns:
            Router ready for routing
        """
        fitted = self.fit_structure(structure_dataset)
        policy = self.calibrate(fitted, calibration_dataset)
        return self.compile(fitted, policy)

    def _validate_fit_clusters(self, dataset: Dataset) -> None:
        """Validate dataset for fit_clusters step."""
        if not self.models:
            raise ValueError("At least one model is required")

        errors = dataset.validate_schema(
            required_columns=["id", self.input_col],
            check_unique_ids=True,
        )
        if errors:
            raise ValueError(f"Dataset validation failed: {errors}")

    def _validate_score_clusters(self, dataset: Dataset) -> None:
        """Validate dataset for score_clusters step."""
        allowed_model_ids = set(self.models)

        errors = dataset.validate_schema(
            required_columns=["id", self.input_col, self.target_col],
            check_unique_ids=True,
        )
        if errors:
            raise ValueError(f"Dataset validation failed: {errors}")

        targets = dataset.column(self.target_col)
        target_errors: list[str] = []
        for idx, value in enumerate(targets):
            if value is None:
                continue
            if not isinstance(value, dict):
                target_errors.append(f"Row {idx}: targets must be a dict")
                continue
            for model_id, score in value.items():
                if model_id not in allowed_model_ids:
                    target_errors.append(
                        f"Row {idx}: unknown model_id '{model_id}'. "
                        f"Allowed model IDs: {sorted(allowed_model_ids)}"
                    )
                if not isinstance(model_id, str):
                    target_errors.append(f"Row {idx}: model_id must be str")
                if not isinstance(score, (int, float)):
                    target_errors.append(
                        f"Row {idx}: target for {model_id} must be numeric"
                    )

        if target_errors:
            raise ValueError(f"Dataset validation failed: {target_errors}")

    def _validate(self, dataset: Dataset) -> None:
        if not self.models:
            raise ValueError("At least one model is required")

        allowed_model_ids = set(self.models)

        errors = dataset.validate_schema(
            required_columns=["id", self.input_col, self.target_col],
            check_unique_ids=True,
        )
        if errors:
            raise ValueError(f"Dataset validation failed: {errors}")

        targets = dataset.column(self.target_col)
        errors.extend(
            [
                f"Row {idx}: targets must be a dict"
                for idx, value in enumerate(targets)
                if value is not None and not isinstance(value, dict)
            ]
        )
        if errors:
            raise ValueError(f"Dataset validation failed: {errors}")

        target_errors: list[str] = []
        for idx, value in enumerate(targets):
            if value is None:
                continue
            if not isinstance(value, dict):
                continue  # Already caught above
            for model_id, score in value.items():
                if model_id not in allowed_model_ids:
                    target_errors.append(
                        f"Row {idx}: unknown model_id '{model_id}'. "
                        f"Allowed model IDs: {sorted(allowed_model_ids)}"
                    )
                if not isinstance(model_id, str):
                    target_errors.append(f"Row {idx}: model_id must be str")
                if score not in (0, 1):
                    target_errors.append(
                        f"Row {idx}: target for {model_id} must be 0 or 1"
                    )

        if target_errors:
            raise ValueError(f"Dataset validation failed: {target_errors}")

    def _make_embedder(self) -> Embedder:
        if self.embedder is not None:
            return self.embedder

        return SentenceTransformers(
            model=self.embedding_model,
            batch_size=self.embedding_batch_size,
            normalize=self.embedding_normalize,
            trust_remote_code=self.allow_trust_remote_code,
            device=self.device,
        )

    def _reduce_or_pass(
        self, embeddings: np.ndarray
    ) -> tuple[np.ndarray, ReductionPayload | None]:
        if self.reducer is None:
            return embeddings, None

        reduced = self.reducer.fit_transform(embeddings)
        return reduced, self.reducer.checkpoint_payload()

    def _make_clusterer(self) -> Clusterer:
        if self.clusterer is None:
            return KMeansClusterer(
                n_clusters=self.n_clusters,
                random_state=self.random_state,
            )

        return self.clusterer

    def _calc_scores(
        self,
        labels: np.ndarray,
        targets: list[dict[str, int] | None],
        model_ids: list[str],
        n_clusters: int,
    ) -> dict[str, list[float]]:
        allowed_model_ids = set(model_ids)

        # Validate all targets are from allowed model IDs
        for row_targets in targets:
            if row_targets is None:
                continue
            for model_id in row_targets.keys():
                if model_id not in allowed_model_ids:
                    raise ValueError(
                        f"Unknown model_id '{model_id}' in targets. "
                        f"Allowed: {sorted(allowed_model_ids)}"
                    )

        # Use numpy for faster computation
        valid_mask = labels >= 0
        valid_labels = np.unique(labels[valid_mask])

        label_map = (
            {int(label): i for i, label in enumerate(valid_labels)}
            if len(valid_labels) == n_clusters
            else {i: i for i in range(n_clusters)}
        )

        # Pre-allocate accumulator arrays
        n_models = len(model_ids)
        model_to_idx = {mid: i for i, mid in enumerate(model_ids)}
        sums = np.zeros((n_models, n_clusters), dtype=np.float64)
        counts = np.zeros((n_models, n_clusters), dtype=np.int32)

        # Vectorized accumulation
        for row_idx, (label, row_targets) in enumerate(zip(labels, targets)):
            if label < 0:
                continue
            cluster_idx = label_map.get(int(label))
            if cluster_idx is None:
                continue
            if row_targets is None:
                continue
            for model_id, score in row_targets.items():
                model_idx = model_to_idx.get(model_id)
                if model_idx is not None:
                    sums[model_idx, cluster_idx] += float(score)
                    counts[model_idx, cluster_idx] += 1

        # Raise if any cluster has no data
        if np.any(counts == 0):
            zero_count = int(np.sum(counts == 0))
            raise ValueError(
                f"No training data for {zero_count} cluster-model combination(s). "
                "Consider increasing cluster size or reducing n_clusters."
            )

        # Return scores (higher is better) - already in 0-1 range
        scores = sums / counts

        return {mid: scores[i].tolist() for i, mid in enumerate(model_ids)}
