"""

Hyperparameter search-space configuration.

"""

from ray import tune

import numpy as np

import pandas as pd

N_SAMPLES = [256]

DECAY_WEIGHTS = [0.2, 0.4, 0.6, 0.8, 0.9, 1]

TEMPERATURES = [0.5, 1.0, 1.5, 2.0, 2.5, 3.0]

LOG_WEIGHTS = [1e-3, 1e-2, 1e-1, 1e0, 1e1, 1e2]

HIDDEN_DIMS = [64, 128, 256]

N_LAYERS = [3, 4]

N_HEADS = [12, 16, 24]

# Conditional hyper-param functions

def get_dropout_from_n_layers(spec):

    """

    Returns dropout of 0 if n_layers==1

    else random dropout.

    """

    if spec.config.model.n_layers == 1:

        return 0

    else:

        return np.random.choice([0, 0.1, 0.2, 0.4])

# Hyperparameter search-space

config_generic = {

    "lr": tune.grid_search([1e-4, 1e-3, 1e-2]),

    "optimizer": "SGD",  # tune.choice(['Adam','SGD']),

    "momentum": 0.9,  # tune.grid_search(DECAY_WEIGHTS),

    "train_epochs": 100,

    "train_mb_size": tune.grid_search([16, 32, 64, 128]),

}

config_model = {

    "CNN": {

        "hidden_dim": tune.grid_search(HIDDEN_DIMS),

        "n_layers": tune.grid_search(N_LAYERS),

        "kernel_size": tune.grid_search([3, 5, 7]),

        "nonlinearity": tune.grid_search(["tanh", "relu"]),

    },

    "MLP": {

        "hidden_dim": tune.grid_search(HIDDEN_DIMS),

        "n_layers": tune.grid_search(N_LAYERS),

        "dropout": tune.sample_from(get_dropout_from_n_layers),

        "nonlinearity": tune.grid_search(["tanh", "relu"]),

    },

    "Transformer": {

        "hidden_dim": tune.grid_search(HIDDEN_DIMS),

        "n_layers": tune.grid_search(N_LAYERS),

        "n_heads": tune.grid_search(N_HEADS),

        "dropout": tune.sample_from(get_dropout_from_n_layers),

        "nonlinearity": tune.grid_search(["relu", "gelu"]),

    },

    "RNN": {

        "hidden_dim": tune.grid_search(HIDDEN_DIMS),

        "n_layers": tune.grid_search(N_LAYERS),

        "dropout": tune.sample_from(get_dropout_from_n_layers),

        "nonlinearity": tune.grid_search(["tanh", "relu"]),

        "bidirectional": tune.grid_search([True, False]),

    },

    "LSTM": {

        "hidden_dim": tune.grid_search(HIDDEN_DIMS),

        "n_layers": tune.grid_search(N_LAYERS),

        "dropout": tune.sample_from(get_dropout_from_n_layers),

        "bidirectional": tune.grid_search([True, False]),

    },

    "GRU": {

        "hidden_dim": tune.grid_search(HIDDEN_DIMS),

        "n_layers": tune.grid_search(N_LAYERS),

        "dropout": tune.sample_from(get_dropout_from_n_layers),

        "bidirectional": tune.grid_search([True, False]),

    },

}

config_cl = {

    "Replay": {

        "mem_size": tune.grid_search([5 * n for n in N_SAMPLES])

        # JA: Have edited this in replay definition

        #'storage_policy':storage_policy.ClassBalancedStoragePolicy()

    },

    "GDumb": {"mem_size": tune.grid_search(N_SAMPLES)},

    "EWC": {"mode": "separate", "ewc_lambda": tune.grid_search(LOG_WEIGHTS)},

    "OnlineEWC": {

        "mode": "online",

        "ewc_lambda": tune.grid_search(LOG_WEIGHTS),

        "decay_factor": tune.grid_search(DECAY_WEIGHTS),

    },

    "SI": {"si_lambda": tune.grid_search(LOG_WEIGHTS)},

    "LwF": {

        "alpha": tune.grid_search(LOG_WEIGHTS),

        "temperature": tune.grid_search(TEMPERATURES),

    },

    "LFL": {"lambda_e": tune.grid_search([LOG_WEIGHTS])},

    "GEM": {

        "patterns_per_exp": tune.grid_search(N_SAMPLES),

        "memory_strength": tune.grid_search(DECAY_WEIGHTS),

    },

    "AGEM": {

        "patterns_per_exp": tune.grid_search(N_SAMPLES),

        "sample_size": tune.grid_search([i * max(N_SAMPLES) for i in range(1, 3)]),

    },

    #'CoPE':

}

vals = {

    "mem_size": N_SAMPLES,

    "ewc_lambda": LOG_WEIGHTS,

    "decay_factor": DECAY_WEIGHTS,

    "si_lambda": LOG_WEIGHTS,

    "alpha": LOG_WEIGHTS,

    "temperature": TEMPERATURES,

    "lambda_e": LOG_WEIGHTS,

    "patterns_per_exp": N_SAMPLES,

    "memory_strength": DECAY_WEIGHTS,

    "sample_size": [i * max(N_SAMPLES) for i in range(1, 3)],

}

vals2 = {

    "hidden_dim": HIDDEN_DIMS,

    "n_layers": N_LAYERS,

    "nonlinearity": ["tanh", "relu", "gelu*"],

    "n_heads": N_HEADS,

    "bidirectional": ["True", "False"],

}

models = ["MLP", "CNN", "LSTM", "Transformer"]

for k in vals2.keys():

    vals2[k] = (vals2[k], *[k in config_model[model] for model in models])

df_hp = pd.DataFrame(vals.items(), columns=["Hyperparameter", "Values"])

df_hp = df_hp.set_index(["Hyperparameter"])

# print(df_hp.to_latex())

df_hp = pd.DataFrame(

    ((k, *v) for k, v in vals2.items()), columns=["Hyperparameter", "Values", *models]

)

df_hp = df_hp.set_index(["Hyperparameter"])

# print(df_hp.to_latex())