import skopt
from skopt.space import Categorical, Integer, Real

from src.bnn.cross_val import cross_val
from src.bnn.dropout.hyper_opt import opt

space = (
    # number of dense layers followed by a dropout layer
    Integer(1, 5, name="n_layers"),
    # learning rate of the adam optimizer
    Real(1e-5, 1e-4, "log-uniform", name="learning_rate"),
)


def objective(
    model,
    data,
    n_calls=(10, 10),
    method="gp",
    n_splits=3,
    acq_func="gp_hedge",
    random_state=0,
):
    """
    model: Keras Model instance
    data: instance of Data class
    n_calls: 2-tuple of ints, number of iterations for the (outer, inner)
        minimization loop
    method: tuple of strings containing one or more of gp, dummy, forest
        specifies which of skopt's minimizers to try
    n_splits: int how many folds of cross validation (cv) to perform, set to 0 for no cv
    random_state (int): random state
    """
    # iteration counter for the current outer and inner minimization loop
    call_counts = [1, 1]

    # if n_splits is non-zero, cross validation (cv) (or shuffled split) will be used,
    # in that case, only split data into train and test set (0.9, 0, 0.1) at this stage,
    #   since train will be split into train/val by cv
    # else split into train/val/test (0.8, 0.1, 0.1)
    # train, val, test = data.split((0.9, 0, 0.1) if n_splits else (0.8, 0.1, 0.1))

    def curried_inner_objective(n_layers, learning_rate):
        def inner_objective(inner_params):
            model.log(
                f"""outer loop: {call_counts[0]}/{n_calls[0]},
                inner loop: {call_counts[1]}/{n_calls[0]}"""
            )
            call_counts[1] += 1

            nodes, dropouts, activations = [
                inner_params[i : i + n_layers]
                for i in range(0, len(inner_params), n_layers)
            ]
            params = opt.HyperParams(
                n_layers, learning_rate, nodes, dropouts, activations
            )
            if n_splits:
                loss = cross_val(
                    model,
                    data[0],
                    params,
                    n_splits=n_splits,
                    cbs=[opt.tb_log(model, params)],
                )
            else:
                loss = model.run(data, params, cbs=[opt.tb_log(model, params)])

            model.losses.append(loss)
            model.params.append(params)

            if loss < model.min_loss:
                model.log(f"Found new min loss {round(loss, 4)} with {params}")
                model.save()
                model.min_loss = loss
                model.best_model = model.model
                model.best_params = params

            model.log(f"{params}")
            return loss

        return inner_objective

    def outer_objective(outer_params):

        n_layers, learning_rate = outer_params
        # number of nodes in each dense layer
        nodes = (Integer(10, 100, name=f"n_nodes_{i+1}") for i in range(n_layers))
        # fraction of dropped nodes in each dropout layer
        dropouts = (Real(0, 0.5, name=f"dropout_rate_{i+1}") for i in range(n_layers))
        # list of activation functions for each dense layer
        activations = (
            Categorical(["tanh", "relu"], name=f"activation_{i+1}")
            for i in range(n_layers)
        )
        inner_space = (*nodes, *dropouts, *activations)
        res = getattr(skopt, method + "_minimize")(
            curried_inner_objective(n_layers, learning_rate),
            inner_space,
            acq_func=acq_func,
            n_calls=n_calls[1],
            random_state=random_state,
        )
        model.log(
            f"lowest loss obtained in outer loop {call_counts[0]}/{n_calls[0]}: {res.fun}"
        )
        call_counts[0] += 1
        call_counts[1] = 1
        return res.fun

    return outer_objective
