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 PhysBONativeOptimizer(BaseOptimizer):
    """PHYSBO-native discrete optimizer.

    Uses PHYSBO's own bayes_search and score modes (typically EI/PI/TS).
    For arbitrary custom scores, use CustomOptimizer with PhysBOSurrogate or
    SklearnGPRSurrogate instead.
    """

    def __init__(
        self,
        score_mode: str = "EI",
        num_rand_basis: int = 200,
        interval: int = 0,
        random_seed: int | None = None,
        maximize: bool = True,
    ):
        self.score_mode = score_mode
        self.num_rand_basis = num_rand_basis
        self.interval = interval
        self.random_seed = random_seed
        self.maximize = maximize
        self.physbo = None
        self.policy = None

    def _import_physbo(self):
        if self.physbo is None:
            try:
                import physbo
            except ImportError:
                print("IMPORT ERROR: physbo")
                print("Install example: pip install physbo")
                raise
            self.physbo = physbo
        return self.physbo

    def initialize(
        self,
        X_candidates,
        y: Optional[Sequence[float]] = None,
        observed_mask: Optional[Sequence[bool]] = None,
        observed_indices: Optional[Sequence[int]] = None,
    ):
        physbo = self._import_physbo()
        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 observed_indices is not None and len(yy) == len(observed_indices):
                self.y_all[np.asarray(observed_indices, dtype=int)] = yy
            elif len(yy) == n:
                self.y_all[:] = yy
            else:
                raise ValueError("y length must be n_candidates or len(observed_indices).")

        self.observed_mask = make_observed_mask(
            n,
            y=self.y_all,
            observed_mask=observed_mask,
            observed_indices=observed_indices,
        )
        idx_train = np.where(self.observed_mask)[0]
        y_train = self.y_all[idx_train]
        if not self.maximize:
            y_train = -y_train

        self.policy = physbo.search.discrete.policy(
            test_X=self.X_candidates,
            initial_data=(idx_train, y_train),
        )
        if self.random_seed is not None:
            self.policy.set_seed(int(self.random_seed))
        return self

    def ask(self, n_points: int = 1) -> BOResult:
        if self.policy is None:
            raise RuntimeError("Optimizer is not initialized.")

        actions = self.policy.bayes_search(
            max_num_probes=n_points,
            simulator=None,
            score=self.score_mode,
            interval=self.interval,
            num_rand_basis=self.num_rand_basis,
        )
        indices = np.asarray(actions, dtype=int).reshape(-1)
        if len(indices) > n_points:
            indices = indices[-n_points:]

        score = None
        try:
            score = self.policy.get_score(mode=self.score_mode, xs=self.X_candidates)
        except Exception:
            pass
        return BOResult(indices=indices, scores=None if score is None else score[indices], X=self.X_candidates[indices])

    def tell(self, indices, y_new):
        # PHYSBO policy update API is less transparent than simple refit for custom loops.
        # Rebuild the policy from accumulated observations for clarity.
        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
        return self.initialize(self.X_candidates, y=self.y_all, observed_mask=self.observed_mask)

    def predict(self, X=None, return_std: bool = True):
        if self.policy is None:
            raise RuntimeError("Optimizer is not initialized.")
        if X is None:
            X = self.X_candidates
        X = as_2d_array(X)
        mean = self.policy.get_post_fmean(X)
        if not self.maximize:
            mean = -mean
        if return_std:
            var = self.policy.get_post_fcov(X)
            std = np.sqrt(np.maximum(var, 0.0))
            return mean, std
        return mean