from __future__ import annotations

import numpy as np
from scipy.stats import norm

from ..base import BaseAcquisition


def _best_y(y_observed, maximize: bool = True) -> float:
    if y_observed is None:
        return 0.0
    y = np.asarray(y_observed, dtype=float)
    y = y[~np.isnan(y)]
    if len(y) == 0:
        return 0.0
    return float(np.max(y) if maximize else np.min(y))


class ExpectedImprovement(BaseAcquisition):
    """Analytic EI for independent posterior mean/std."""

    name = "ei"

    def __init__(self, xi: float = 0.0):
        self.xi = xi

    def __call__(self, X, mean, std, y_observed=None, observed_mask=None, maximize: bool = True, **kwargs):
        mean = np.asarray(mean, dtype=float).reshape(-1)
        std = np.maximum(np.asarray(std, dtype=float).reshape(-1), 1.0e-300)
        y_best = kwargs.get("y_best", _best_y(y_observed, maximize=maximize))
        if maximize:
            improvement = mean - y_best - self.xi
        else:
            improvement = y_best - mean - self.xi
        z = improvement / std
        return improvement * norm.cdf(z) + std * norm.pdf(z)


class ProbabilityImprovement(BaseAcquisition):
    name = "pi"

    def __init__(self, xi: float = 0.0):
        self.xi = xi

    def __call__(self, X, mean, std, y_observed=None, observed_mask=None, maximize: bool = True, **kwargs):
        mean = np.asarray(mean, dtype=float).reshape(-1)
        std = np.maximum(np.asarray(std, dtype=float).reshape(-1), 1.0e-300)
        y_best = kwargs.get("y_best", _best_y(y_observed, maximize=maximize))
        improvement = (mean - y_best - self.xi) if maximize else (y_best - mean - self.xi)
        return norm.cdf(improvement / std)


class UpperConfidenceBound(BaseAcquisition):
    name = "ucb"

    def __init__(self, kappa: float = 2.0):
        self.kappa = kappa

    def __call__(self, X, mean, std, y_observed=None, observed_mask=None, maximize: bool = True, **kwargs):
        mean = np.asarray(mean, dtype=float).reshape(-1)
        std = np.asarray(std, dtype=float).reshape(-1)
        return mean + self.kappa * std if maximize else -mean + self.kappa * std


class LowerConfidenceBound(BaseAcquisition):
    """LCB score converted so larger is better."""

    name = "lcb"

    def __init__(self, kappa: float = 2.0):
        self.kappa = kappa

    def __call__(self, X, mean, std, y_observed=None, observed_mask=None, maximize: bool = False, **kwargs):
        mean = np.asarray(mean, dtype=float).reshape(-1)
        std = np.asarray(std, dtype=float).reshape(-1)
        # Larger score means more attractive. For minimization, low mean - kappa*std is attractive.
        return -(mean - self.kappa * std)