import gpyreg as gpr
import numpy as np
import pytest
from scipy.stats import norm

from pybads.bads import BADS


def get_test_opt_conf(D=3):
    x0 = np.ones((1, D)) * 4
    LB = -100*np.ones(D)                # Lower bound
    UB = 100*np.ones(D)                 # Upper bound
    PLB = -8*np.ones(D)                 # Plausible lower bound
    PUB = 12*np.ones(D)                 # Plausible upper bound
    tol_errs = np.array([0.1, 0.1, 1, 1.])
    return D, x0, LB, UB, PLB, PUB, tol_errs

def run_bads(fun, x0, LB, UB, PLB, PUB, tol_errs, f_min,
            oracle_fun=None, non_box_cons=None, uncertainty_handling=False,
            assert_flag=False, max_fun_evals=None):
    options = {}
    options["display"] = "full"#debug_flag = True

    if uncertainty_handling > 0:
        options["uncertainty_handling"] = True
        options["max_fun_evals"] = 200 if max_fun_evals is None  else max_fun_evals
        if uncertainty_handling > 1:
            options["specify_target_noise"] = True
    else:
        options["max_fun_evals"] = 100 if max_fun_evals is None else max_fun_evals

    optimize_result = BADS(fun=fun, x0=x0, lower_bounds=LB,
                                    upper_bounds=UB, plausible_lower_bounds=PLB,
                                    plausible_upper_bounds=PUB, non_box_cons=non_box_cons,
                                    options=options).optimize()
    x = optimize_result['x']
    fval = optimize_result['fval']

    if oracle_fun is None:
        #print(f"Final value: {fval:.3f} (true value: {f_min}), with {optimize_result['func_count']} fun evals.")
        err = np.abs(fval - f_min)
    else:
        fval_true = oracle_fun(x)
        #print(f"Final value (not-noisy): {fval_true:.3f} (true value: {f_min}) with {optimize_result['func_count']} fun evals.")
        err = np.abs(fval_true - f_min)
    if assert_flag:
        assert np.any(err < tol_errs), f"Error {err} is not smaller than tolerance {tol_errs} when optimizing {fun.__name__}."
    
    return optimize_result, err
    
def test_ellipsoid_opt():
    D, x0, LB, UB, PLB, PUB, tol_errs = get_test_opt_conf()
    fun = lambda x: np.sum((np.atleast_2d(x) / np.arange(1, len(x) + 1) ** 2) ** 2)
    run_bads(fun, x0, LB, UB, PLB, PUB, tol_errs, f_min=0.0, assert_flag=True)
    
def test_univariate_input_and_opt():
    rfn = lambda x: x**2 + 3.2 + np.random.normal(scale=0.1)
    plb = -5
    pub = 5
    x0 = 3
    opt = BADS(rfn, x0, plausible_lower_bounds=plb, plausible_upper_bounds=pub)
    opt.optimize()

def test_1D_opt_ndarray():
    """Test optimization with 1D inputs and nd arrays"""
    fun = lambda x: np.atleast_2d(x)**2
    x0 = np.array([[2.0]])
    LB = np.array([[-10.0]])
    UB = np.array([[10.0]])
    PLB = np.array([[-5.0]])
    PUB = np.array([[5.0]])
    tol_errs = np.array([0.1])
    run_bads(fun, x0, LB, UB, PLB, PUB, tol_errs, f_min=0.0, assert_flag=True)

def test_1D_opt_1darray():
    """Test optimization with 1D inputs and 1D arrays"""
    fun = lambda x: np.atleast_2d(x)**2
    x0 = np.array([2.0])
    LB = np.array([-10.0])
    UB = np.array([10.0])
    PLB = np.array([-5.0])
    PUB = np.array([5.0])
    tol_errs = np.array([0.1])
    run_bads(fun, x0, LB, UB, PLB, PUB, tol_errs, f_min=0.0, assert_flag=True)

def test_1D_opt_scalar():
    """Test optimization with scalar inputs"""
    fun = lambda x: np.atleast_2d(x)**2
    x0 = 2.0
    LB = -10.0
    UB = 10.0
    PLB = -5.0
    PUB = 5.0
    tol_errs = np.array([0.1])
    run_bads(fun, x0, LB, UB, PLB, PUB, tol_errs, f_min=0.0, assert_flag=True)

def test_high_dim_opt():
    D, x0, LB, UB, PLB, PUB, tol_errs = get_test_opt_conf(D=60)
    fun = lambda x: np.sum((np.atleast_2d(x) / np.arange(1, len(x) + 1) ** 2) ** 2)
    run_bads(fun, x0, LB, UB, PLB, PUB, tol_errs, f_min=0.0, assert_flag=False, max_fun_evals=200)

def test_sphere_opt():
    D, x0, LB, UB, PLB, PUB, tol_errs = get_test_opt_conf()
    x0 = np.zeros((1, D))
    fun = lambda x: np.sum(np.atleast_2d(x)**2, axis=1)
    non_box_cons = lambda x: np.atleast_2d(x)[:, 0] + np.atleast_2d(x)[:, 1] >= np.sqrt(2)     # Non-bound constraints
    print(non_box_cons(x0))
    run_bads(fun, x0, LB, UB, PLB, PUB, tol_errs, f_min=1.0, non_box_cons=non_box_cons, assert_flag=True)
    
def test_noisy_sphere_opt():
    D, x0, LB, UB, PLB, PUB, tol_errs = get_test_opt_conf()
    fun = lambda x: np.sum(np.atleast_2d(x)**2, axis=1) + np.random.randn()             # Noisy objective function
    oracle_fun = lambda x: np.sum(np.atleast_2d(x)**2, axis=1)                          # True objective function
    run_bads(fun, x0, LB, UB, PLB, PUB, tol_errs, f_min=0.0, oracle_fun=oracle_fun, assert_flag=True)
    
def test_small_noisy_func():
    np.random.seed(42343)
    def noisy_sphere(x, sigma=1.0):
        """Simple quadratic function with added noise."""
        x_2d = np.atleast_2d(x)
        f = np.sum(x_2d**2, axis=1)
        noise = 1e-4 * sigma * np.random.normal(size=x_2d.shape[0])
        return f + noise
    
    x0 = np.array([-3, -3, -3])
    LB = np.array([-5, -5, -5])
    UB = np.array([5, 5, 5])
    PLB = np.array([-2, -2, -2])
    PUB = np.array([2, 2, 2])
    tol_errs = np.array([0.1, 0.1, 0.1])

    oracle_fun = lambda x: np.sum(np.atleast_2d(x)**2, axis=1)
    run_bads(
        noisy_sphere,
        x0,
        LB,
        UB,
        PLB,
        PUB,
        tol_errs,
        f_min=0.0,
        oracle_fun=oracle_fun,
        uncertainty_handling=1,
        assert_flag=True,
        max_fun_evals=300
    )
    
def he_noisy_sphere(x):
    y = np.sum(np.atleast_2d(x)**2, axis=1)
    s = 2 + 1*np.sqrt(y)
    y = y + s*np.random.randn()
    return y, s

def test_he_noisy_sphere_opt():
    D, x0, LB, UB, PLB, PUB, tol_errs = get_test_opt_conf()
    fun = he_noisy_sphere
    oracle_fun = lambda x: np.sum(np.atleast_2d(x)**2, axis=1)                          # True objective function
    run_bads(fun, x0, LB, UB, PLB, PUB, tol_errs, f_min=0.0, oracle_fun=oracle_fun, uncertainty_handling=2, assert_flag=True)
    