"""Search strategies for threshold optimization.

This module implements various strategies for exploring the threshold
search space, including grid search, random search, and Bayesian optimization.
"""

from __future__ import annotations

import logging
from abc import ABC, abstractmethod
from dataclasses import dataclass
from typing import TYPE_CHECKING, Iterator

import numpy as np

if TYPE_CHECKING:
    from .config import CalibrationConfig, ThresholdBounds

from embergraphrouter.inference.tiered_router import TieredThresholds

logger = logging.getLogger(__name__)


@dataclass(frozen=True)
class ThresholdCandidate:
    """A threshold configuration candidate.

    Attributes:
        thresholds: The threshold configuration.
        iteration: Search iteration that produced this candidate.
        source: Strategy that produced this candidate.
    """

    thresholds: TieredThresholds
    iteration: int
    source: str


class SearchStrategy(ABC):
    """Abstract base class for search strategies."""

    @abstractmethod
    def __init__(self, bounds: "ThresholdBounds", random_seed: int = 42):
        """Initialize search strategy.

        Args:
            bounds: Bounds for threshold parameters.
            random_seed: Random seed for reproducibility.
        """
        pass

    @abstractmethod
    def suggest(self, n: int = 1) -> list[ThresholdCandidate]:
        """Suggest next threshold configurations to evaluate.

        Args:
            n: Number of configurations to suggest.

        Returns:
            List of threshold candidates.
        """
        pass

    @abstractmethod
    def update(
        self,
        candidate: ThresholdCandidate,
        score: float,
    ) -> None:
        """Update search state with evaluation result.

        Args:
            candidate: Evaluated candidate.
            score: Combined score (higher is better).
        """
        pass

    @property
    @abstractmethod
    def name(self) -> str:
        """Strategy name for logging."""
        pass


class GridSearch(SearchStrategy):
    """Exhaustive grid search over threshold space.

    This strategy evaluates all combinations of discrete threshold values.
    Best for small search spaces or when thoroughness is required.
    """

    def __init__(
        self,
        bounds: "ThresholdBounds",
        random_seed: int = 42,
        n_points_per_dim: int = 5,
    ):
        """Initialize grid search.

        Args:
            bounds: Bounds for threshold parameters.
            random_seed: Random seed (unused for grid, but kept for interface).
            n_points_per_dim: Number of points to sample per dimension.
        """
        self.bounds = bounds
        self.n_points = n_points_per_dim
        self._iteration = 0
        self._grid_iterator: Iterator[TieredThresholds] | None = None

    @property
    def name(self) -> str:
        return "grid"

    def _create_grid(self) -> Iterator[TieredThresholds]:
        """Generate all grid points.

        Note: Grid search now only explores 3 parameters (trivial_mean_logit,
        uncertain_mean_logit, uncertain_logit_margin) since max_prob-based
        signals were removed due to softmax compression making them non-discriminative.
        """
        trivial_mean_logits = np.linspace(
            self.bounds.trivial_mean_logit_min,
            self.bounds.trivial_mean_logit_max,
            self.n_points,
        )
        uncertain_mean_logits = np.linspace(
            self.bounds.uncertain_mean_logit_min,
            self.bounds.uncertain_mean_logit_max,
            self.n_points,
        )
        uncertain_logit_margins = np.linspace(
            self.bounds.uncertain_logit_margin_min,
            self.bounds.uncertain_logit_margin_max,
            self.n_points,
        )

        for tml in trivial_mean_logits:
            for uml in uncertain_mean_logits:
                for ulm in uncertain_logit_margins:
                    yield TieredThresholds(
                        trivial_mean_logit=float(tml),
                        uncertain_mean_logit=float(uml),
                        uncertain_logit_margin=float(ulm),
                    )

    def suggest(self, n: int = 1) -> list[ThresholdCandidate]:
        """Suggest next grid points."""
        if self._grid_iterator is None:
            self._grid_iterator = self._create_grid()

        candidates = []
        for _ in range(n):
            try:
                thresholds = next(self._grid_iterator)
                self._iteration += 1
                candidates.append(
                    ThresholdCandidate(
                        thresholds=thresholds,
                        iteration=self._iteration,
                        source=self.name,
                    )
                )
            except StopIteration:
                break

        return candidates

    def update(self, candidate: ThresholdCandidate, score: float) -> None:
        """Grid search doesn't need updates."""
        pass


class RandomSearch(SearchStrategy):
    """Random search over threshold space.

    This strategy samples random threshold configurations uniformly
    from the bounded space. Good for exploring diverse configurations.
    """

    def __init__(
        self,
        bounds: "ThresholdBounds",
        random_seed: int = 42,
    ):
        """Initialize random search.

        Args:
            bounds: Bounds for threshold parameters.
            random_seed: Random seed for reproducibility.
        """
        self.bounds = bounds
        self.rng = np.random.RandomState(random_seed)
        self._iteration = 0

    @property
    def name(self) -> str:
        return "random"

    def _sample_thresholds(self) -> TieredThresholds:
        """Sample random thresholds from bounds.

        Note: Only 3 parameters are sampled since max_prob-based signals
        were removed due to softmax compression making them non-discriminative.
        """
        return TieredThresholds(
            trivial_mean_logit=self.rng.uniform(
                self.bounds.trivial_mean_logit_min,
                self.bounds.trivial_mean_logit_max,
            ),
            uncertain_mean_logit=self.rng.uniform(
                self.bounds.uncertain_mean_logit_min,
                self.bounds.uncertain_mean_logit_max,
            ),
            uncertain_logit_margin=self.rng.uniform(
                self.bounds.uncertain_logit_margin_min,
                self.bounds.uncertain_logit_margin_max,
            ),
        )

    def suggest(self, n: int = 1) -> list[ThresholdCandidate]:
        """Suggest random configurations."""
        candidates = []
        for _ in range(n):
            self._iteration += 1
            candidates.append(
                ThresholdCandidate(
                    thresholds=self._sample_thresholds(),
                    iteration=self._iteration,
                    source=self.name,
                )
            )
        return candidates

    def update(self, candidate: ThresholdCandidate, score: float) -> None:
        """Random search doesn't use feedback."""
        pass


class BayesianSearch(SearchStrategy):
    """Bayesian optimization over threshold space.

    Uses Gaussian Process surrogate model to guide search toward
    promising regions of the parameter space. Most sample-efficient
    for expensive evaluations.

    Falls back to random search if scipy is not available.
    """

    def __init__(
        self,
        bounds: "ThresholdBounds",
        random_seed: int = 42,
        n_initial: int = 10,
        acquisition: str = "ei",  # Expected Improvement
    ):
        """Initialize Bayesian optimization.

        Args:
            bounds: Bounds for threshold parameters.
            random_seed: Random seed for reproducibility.
            n_initial: Number of initial random samples.
            acquisition: Acquisition function ("ei", "ucb", "pi").
        """
        self.bounds = bounds
        self.rng = np.random.RandomState(random_seed)
        self.n_initial = n_initial
        self.acquisition = acquisition
        self._iteration = 0

        # Store observations
        self._X: list[np.ndarray] = []
        self._y: list[float] = []

        # Try to import scipy for GP
        try:
            import scipy.stats  # noqa: F401

            self._has_scipy = True
        except ImportError:
            logger.warning(
                "scipy not available, BayesianSearch will use random sampling"
            )
            self._has_scipy = False

        # Parameter bounds as array (3 parameters after removing max_prob-based signals)
        self._bounds_array = np.array([
            [bounds.trivial_mean_logit_min, bounds.trivial_mean_logit_max],
            [bounds.uncertain_mean_logit_min, bounds.uncertain_mean_logit_max],
            [bounds.uncertain_logit_margin_min, bounds.uncertain_logit_margin_max],
        ])

    @property
    def name(self) -> str:
        return "bayesian"

    def _to_array(self, thresholds: TieredThresholds) -> np.ndarray:
        """Convert thresholds to array (3 parameters)."""
        return np.array([
            thresholds.trivial_mean_logit,
            thresholds.uncertain_mean_logit,
            thresholds.uncertain_logit_margin,
        ])

    def _from_array(self, x: np.ndarray) -> TieredThresholds:
        """Convert array to thresholds (3 parameters)."""
        return TieredThresholds(
            trivial_mean_logit=float(x[0]),
            uncertain_mean_logit=float(x[1]),
            uncertain_logit_margin=float(x[2]),
        )

    def _sample_random(self) -> np.ndarray:
        """Sample random point from bounds."""
        return self.rng.uniform(
            self._bounds_array[:, 0],
            self._bounds_array[:, 1],
        )

    def _rbf_kernel(
        self,
        X1: np.ndarray,
        X2: np.ndarray,
        length_scale: float = 1.0,
        variance: float = 1.0,
    ) -> np.ndarray:
        """RBF (Gaussian) kernel."""
        dist_sq = np.sum((X1[:, None, :] - X2[None, :, :]) ** 2, axis=-1)
        return variance * np.exp(-0.5 * dist_sq / length_scale**2)

    def _gp_predict(
        self,
        X_new: np.ndarray,
        X_train: np.ndarray,
        y_train: np.ndarray,
        noise: float = 1e-6,
    ) -> tuple[np.ndarray, np.ndarray]:
        """Simple GP prediction."""
        n = len(X_train)
        if n == 0:
            return np.zeros(len(X_new)), np.ones(len(X_new))

        K = self._rbf_kernel(X_train, X_train)
        K_s = self._rbf_kernel(X_train, X_new)
        K_ss = self._rbf_kernel(X_new, X_new)

        L = np.linalg.cholesky(K + noise * np.eye(n))
        alpha = np.linalg.solve(L.T, np.linalg.solve(L, y_train))
        mu = K_s.T @ alpha

        v = np.linalg.solve(L, K_s)
        var = np.diag(K_ss) - np.sum(v**2, axis=0)
        var = np.maximum(var, 1e-8)

        return mu, np.sqrt(var)

    def _expected_improvement(
        self,
        mu: np.ndarray,
        sigma: np.ndarray,
        y_best: float,
        xi: float = 0.01,
    ) -> np.ndarray:
        """Compute expected improvement acquisition."""
        from scipy.stats import norm

        with np.errstate(divide="ignore", invalid="ignore"):
            imp = mu - y_best - xi
            Z = imp / sigma
            ei = imp * norm.cdf(Z) + sigma * norm.pdf(Z)
            ei[sigma <= 0] = 0.0
        return ei

    def suggest(self, n: int = 1) -> list[ThresholdCandidate]:
        """Suggest next configurations using Bayesian optimization."""
        candidates = []

        for _ in range(n):
            self._iteration += 1

            # Use random sampling for initial points or if scipy unavailable
            if len(self._X) < self.n_initial or not self._has_scipy:
                x = self._sample_random()
            else:
                # Use GP to find next point
                X_train = np.array(self._X)
                y_train = np.array(self._y)
                y_best = np.max(y_train)

                # Sample random candidates and evaluate acquisition
                n_candidates = 1000
                X_cand = np.array([self._sample_random() for _ in range(n_candidates)])
                mu, sigma = self._gp_predict(X_cand, X_train, y_train)
                ei = self._expected_improvement(mu, sigma, y_best)

                # Select best candidate
                best_idx = np.argmax(ei)
                x = X_cand[best_idx]

            candidates.append(
                ThresholdCandidate(
                    thresholds=self._from_array(x),
                    iteration=self._iteration,
                    source=self.name,
                )
            )

        return candidates

    def update(self, candidate: ThresholdCandidate, score: float) -> None:
        """Update GP with new observation."""
        x = self._to_array(candidate.thresholds)
        self._X.append(x)
        self._y.append(score)


def create_search_strategy(
    strategy: str,
    config: "CalibrationConfig",
) -> SearchStrategy:
    """Factory function for search strategies.

    Args:
        strategy: Strategy name ("grid", "random", "bayesian").
        config: Calibration configuration.

    Returns:
        Initialized search strategy.
    """
    strategy = strategy.lower()

    if strategy == "grid":
        return GridSearch(
            bounds=config.threshold_bounds,
            random_seed=config.random_seed,
        )
    elif strategy == "random":
        return RandomSearch(
            bounds=config.threshold_bounds,
            random_seed=config.random_seed,
        )
    elif strategy == "bayesian":
        return BayesianSearch(
            bounds=config.threshold_bounds,
            random_seed=config.random_seed,
        )
    else:
        raise ValueError(f"Unknown search strategy: {strategy}")
