import tensorflow as tf
from skopt.space import Categorical, Integer, Real

from src.bnn.dropout.hyper_opt import opt


def create_space(n_layers=4):
    learning_rate = Real(1e-5, 1e-4, "log-uniform", name="learning_rate")
    # 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
    dropout_rates = (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)
    )
    space = (learning_rate, *nodes, *dropout_rates, *activations)
    return space


def objective(model, data, n_layers=4):
    def objective_func(params):
        learning_rate = params[0]
        nodes, dropout_rates, activations = [
            params[i : i + n_layers] for i in range(1, len(params), n_layers)
        ]
        params = opt.HyperParams(
            n_layers, learning_rate, nodes, dropout_rates, activations
        )

        loss = model.run(data, params, cbs=[opt.tb_log(model, 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}")
        tf.keras.backend.clear_session()
        return loss

    return objective_func
