from __future__ import annotations

from typing import Optional, Sequence
import numpy as np

from ..base import BaseOptimizer, BOResult, as_2d_array, make_observed_mask


class RandomOptimizer(BaseOptimizer):
    def __init__(self, random_seed: int | None = None):
        self.random_seed = random_seed
        self.rng = np.random.default_rng(random_seed)

    def initialize(self, X_candidates, y: Optional[Sequence[float]] = None, observed_mask=None, observed_indices=None):
        self.X_candidates = as_2d_array(X_candidates, "X_candidates")
        n = len(self.X_candidates)
        self.y_all = np.full(n, np.nan)
        if y is not None:
            yy = np.asarray(y, dtype=float).reshape(-1)
            if len(yy) == n:
                self.y_all[:] = yy
        self.observed_mask = make_observed_mask(n, y=self.y_all, observed_mask=observed_mask, observed_indices=observed_indices)
        return self

    def ask(self, n_points: int = 1):
        idx = np.where(~self.observed_mask)[0]
        if len(idx) == 0:
            return BOResult(indices=np.array([], dtype=int), X=np.empty((0, self.X_candidates.shape[1])))
        chosen = self.rng.choice(idx, size=min(n_points, len(idx)), replace=False)
        return BOResult(indices=np.asarray(chosen, dtype=int), X=self.X_candidates[chosen])

    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)
        self.y_all[idx] = yy
        self.observed_mask[idx] = True
        return self

    def predict(self, X=None, return_std: bool = True):
        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


class GridOptimizer(BaseOptimizer):
    """Sequentially returns unobserved candidates in their original order."""

    def initialize(self, X_candidates, y: Optional[Sequence[float]] = None, observed_mask=None, observed_indices=None):
        self.X_candidates = as_2d_array(X_candidates, "X_candidates")
        n = len(self.X_candidates)
        self.y_all = np.full(n, np.nan)
        if y is not None:
            yy = np.asarray(y, dtype=float).reshape(-1)
            if len(yy) == n:
                self.y_all[:] = yy
        self.observed_mask = make_observed_mask(n, y=self.y_all, observed_mask=observed_mask, observed_indices=observed_indices)
        return self

    def ask(self, n_points: int = 1):
        idx = np.where(~self.observed_mask)[0][:n_points]
        return BOResult(indices=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)
        self.y_all[idx] = yy
        self.observed_mask[idx] = True
        return self

    def predict(self, X=None, return_std: bool = True):
        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