from __future__ import annotations

from typing import Optional, Sequence
import numpy as np

from ..base import (
    BaseOptimizer,
    BaseSurrogate,
    BaseAcquisition,
    BOResult,
    as_2d_array,
    as_1d_array,
    make_observed_mask,
    top_k_from_score,
)


class CustomOptimizer(BaseOptimizer):
    """Discrete BO optimizer with pluggable surrogate and acquisition.

    This is the main extension point for custom scores such as Stein-lite.
    """

    def __init__(
        self,
        surrogate: BaseSurrogate,
        acquisition: BaseAcquisition,
        maximize: bool = True,
        refit_each_tell: bool = True,
    ):
        self.surrogate = surrogate
        self.acquisition_func = acquisition
        self.maximize = maximize
        self.refit_each_tell = refit_each_tell

    def initialize(
        self,
        X_candidates,
        y: Optional[Sequence[float]] = None,
        observed_mask: Optional[Sequence[bool]] = None,
        observed_indices: Optional[Sequence[int]] = None,
    ):
        self.X_candidates = as_2d_array(X_candidates, "X_candidates")
        n = len(self.X_candidates)
        self.y_all = np.full(n, np.nan, dtype=float)

        if y is not None:
            yy = np.asarray(y, dtype=float).reshape(-1)
            if len(yy) == n:
                self.y_all[:] = yy
            elif observed_indices is not None and len(yy) == len(observed_indices):
                self.y_all[np.asarray(observed_indices, dtype=int)] = yy
            else:
                raise ValueError(
                    "y must be either length n_candidates with nan for unknown, "
                    "or length observed_indices."
                )

        self.observed_mask = make_observed_mask(
            n,
            y=self.y_all,
            observed_mask=observed_mask,
            observed_indices=observed_indices,
        )
        if y is not None and observed_indices is not None and len(np.asarray(y).reshape(-1)) == len(observed_indices):
            self.y_all[:] = np.nan
            self.y_all[np.asarray(observed_indices, dtype=int)] = np.asarray(y, dtype=float).reshape(-1)

        self._fit_surrogate()
        return self

    @property
    def observed_indices(self) -> np.ndarray:
        return np.where(self.observed_mask)[0]

    @property
    def y_observed(self) -> np.ndarray:
        return self.y_all[self.observed_mask]

    def _fit_surrogate(self):
        idx = self.observed_indices
        if len(idx) == 0:
            self.is_fitted_ = False
            return
        self.surrogate.fit(self.X_candidates[idx], self.y_all[idx])
        self.is_fitted_ = True

    def predict(self, X=None, return_std: bool = True):
        if not getattr(self, "is_fitted_", False):
            n = len(self.X_candidates if X is None else X)
            mean = np.full(n, np.nan)
            std = np.full(n, np.inf)
            return (mean, std) if return_std else mean
        if X is None:
            X = self.X_candidates
        return self.surrogate.predict(X, return_std=return_std)

    def acquisition(self, X=None):
        if X is None:
            X = self.X_candidates
        mean, std = self.predict(X, return_std=True)
        return self.acquisition_func(
            X=np.asarray(X, dtype=float),
            mean=mean,
            std=std,
            y_observed=self.y_observed,
            observed_mask=self.observed_mask if X is self.X_candidates else None,
            maximize=self.maximize,
            model=self.surrogate,
        )

    def ask(self, n_points: int = 1) -> BOResult:
        if not getattr(self, "is_fitted_", False):
            unobserved = np.where(~self.observed_mask)[0]
            idx = unobserved[:n_points]
            return BOResult(indices=idx, X=self.X_candidates[idx])

        score = self.acquisition(self.X_candidates)
        idx = top_k_from_score(score, n_points, observed_mask=self.observed_mask)
        return BOResult(indices=idx, scores=score[idx], X=self.X_candidates[idx])

    def tell(self, indices, y_new):
        idx = np.asarray(indices if np.ndim(indices) else [indices], dtype=int).reshape(-1)
        yy = np.asarray(y_new if np.ndim(y_new) else [y_new], dtype=float).reshape(-1)
        if len(idx) != len(yy):
            raise ValueError(f"indices and y_new length mismatch: {len(idx)} vs {len(yy)}")
        self.y_all[idx] = yy
        self.observed_mask[idx] = True
        if self.refit_each_tell:
            self._fit_surrogate()
        return self