from __future__ import annotations

import numpy as np

from ..base import BaseAcquisition


class SteinLiteAcquisition(BaseAcquisition):
    """Practical Stein-like exploration score.

    This is not a strict SVGD implementation.  It is an extensible heuristic:

        score = std_weight * std + grad_weight * std * ||grad mean||

    The gradient is evaluated by finite difference using surrogate.predict().
    This class expects ``model`` to be passed in kwargs by CustomOptimizer.
    """

    name = "stein"

    def __init__(self, eps: float = 1.0e-5, std_weight: float = 1.0, grad_weight: float = 1.0):
        self.eps = eps
        self.std_weight = std_weight
        self.grad_weight = grad_weight

    def _gradient_mean(self, model, X: np.ndarray) -> np.ndarray:
        X = np.asarray(X, dtype=float)
        grad = np.zeros_like(X, dtype=float)
        for j in range(X.shape[1]):
            Xp = X.copy()
            Xm = X.copy()
            Xp[:, j] += self.eps
            Xm[:, j] -= self.eps
            mp = np.asarray(model.predict(Xp), dtype=float).reshape(-1)
            mm = np.asarray(model.predict(Xm), dtype=float).reshape(-1)
            grad[:, j] = (mp - mm) / (2.0 * self.eps)
        return grad

    def __call__(self, X, mean, std, y_observed=None, observed_mask=None, maximize: bool = True, **kwargs):
        model = kwargs.get("model", None)
        X = np.asarray(X, dtype=float)
        mean = np.asarray(mean, dtype=float).reshape(-1)
        std = np.asarray(std, dtype=float).reshape(-1)

        if model is None:
            # Fallback: pure uncertainty.
            return std

        grad = self._gradient_mean(model, X)
        grad_norm = np.linalg.norm(grad, axis=1)
        return self.std_weight * std + self.grad_weight * std * grad_norm