"""
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())