from __future__ import annotations

from dataclasses import dataclass
from typing import Any, Optional, Sequence
import numpy as np


@dataclass
class BOResult:
    """Result returned by ask()."""

    indices: np.ndarray
    scores: Optional[np.ndarray] = None
    X: Optional[np.ndarray] = None


class BaseSurrogate:
    """Base class for surrogate models.

    Surrogates provide posterior mean/std.  They do not decide the next point.
    """

    supports_std: bool = False

    def fit(self, X: np.ndarray, y: np.ndarray) -> "BaseSurrogate":
        raise NotImplementedError

    def predict(self, X: np.ndarray, return_std: bool = False):
        raise NotImplementedError

    def get_params(self, deep: bool = True) -> dict[str, Any]:
        return {}

    def set_params(self, **params: Any) -> "BaseSurrogate":
        for key, val in params.items():
            setattr(self, key, val)
        return self


class BaseAcquisition:
    """Base class for acquisition functions."""

    name: str = "base"

    def __call__(
        self,
        X: np.ndarray,
        mean: np.ndarray,
        std: np.ndarray,
        y_observed: Optional[np.ndarray] = None,
        observed_mask: Optional[np.ndarray] = None,
        maximize: bool = True,
        **kwargs: Any,
    ) -> np.ndarray:
        raise NotImplementedError


class BaseOptimizer:
    """Base ask/tell optimizer for discrete candidate sets."""

    def initialize(
        self,
        X_candidates: np.ndarray,
        y: Optional[Sequence[float]] = None,
        observed_mask: Optional[Sequence[bool]] = None,
        observed_indices: Optional[Sequence[int]] = None,
    ) -> "BaseOptimizer":
        raise NotImplementedError

    def ask(self, n_points: int = 1) -> BOResult:
        raise NotImplementedError

    def tell(self, indices: Sequence[int] | int, y_new: Sequence[float] | float) -> "BaseOptimizer":
        raise NotImplementedError

    def predict(self, X: Optional[np.ndarray] = None, return_std: bool = True):
        raise NotImplementedError


def as_2d_array(X: Any, name: str = "X") -> np.ndarray:
    arr = np.asarray(X, dtype=float)
    if arr.ndim == 1:
        arr = arr.reshape(-1, 1)
    if arr.ndim != 2:
        raise ValueError(f"{name} must be 1D or 2D array-like, got shape={arr.shape}")
    return arr


def as_1d_array(y: Any, name: str = "y") -> np.ndarray:
    arr = np.asarray(y, dtype=float).reshape(-1)
    if arr.ndim != 1:
        raise ValueError(f"{name} must be 1D array-like")
    return arr


def make_observed_mask(
    n: int,
    y: Optional[Sequence[float]] = None,
    observed_mask: Optional[Sequence[bool]] = None,
    observed_indices: Optional[Sequence[int]] = None,
) -> np.ndarray:
    if observed_mask is not None:
        mask = np.asarray(observed_mask, dtype=bool).reshape(-1)
        if len(mask) != n:
            raise ValueError(f"observed_mask length must be {n}, got {len(mask)}")
        return mask

    mask = np.zeros(n, dtype=bool)
    if observed_indices is not None:
        idx = np.asarray(observed_indices, dtype=int).reshape(-1)
        mask[idx] = True
        return mask

    if y is not None:
        yy = np.asarray(y, dtype=float).reshape(-1)
        if len(yy) != n:
            raise ValueError(
                "When observed_mask/observed_indices is omitted, y must have one value per candidate "
                f"(nan for unobserved). Expected {n}, got {len(yy)}"
            )
        mask = ~np.isnan(yy)
        return mask

    return mask


def top_k_from_score(score: np.ndarray, k: int, observed_mask: Optional[np.ndarray] = None) -> np.ndarray:
    score = np.asarray(score, dtype=float).copy().reshape(-1)
    if observed_mask is not None:
        score[np.asarray(observed_mask, dtype=bool)] = -np.inf
    if k < 1:
        raise ValueError("n_points must be >= 1")
    k = min(k, len(score))
    return np.argsort(score)[::-1][:k]