from functools import partial

from configs.config_utils import *
from configs.feature_config import feature_config
from experiment_setup import get_fertility_based_eval_functions

from utils.data_utils import *
from utils.training import *

from models.transformer import *
from models.utils import *
from models.collation import *

base_transformer_model_config = {
    "model_name": "transformer",
    "version": "1.0.0",
    "model_class": TransformerModel,
    "feature_config": feature_config | {"ignored_features":
        [
            "fertility_probability",
            "ov_over_probability",
            # "days_relative_to_ov",
            # "is_biphasic",
            "ov_day",
        ]},
    "batch_fn": produce_window_batches,
    "model_creation_fn": simple_model_creation,
    "collate_fn": simple_x_y_collate,
    "model_save_fn": simple_model_save,
    "model_load_fn": partial(simple_model_load, model_creation_fn=simple_model_creation),
    "batch_loss_fn": get_model_loss,
    "actual_fn": simple_get_y,
    "predict_fn": simple_x_y_predict,
}

base_transformer_training_config = {
    "batch_size": 128,
    "model_class": TransformerModel,
    "learning_parameters": {
        "max_lr": 1e-5,
        "epochs": 30,
        "patience": 5,
    },
    "loss_functions": [
        nn.MSELoss(),
        nn.BCEWithLogitsLoss(),
    ],
    "max_grad_norm": 1.0,
    "train_size": 0.7,
    "val_size": 0.15,
    "test_size": 0.15,
}

fertility_complexity_run = get_run_config(
    run_name="transformer_fertility_complexity_run",
    base_model_configuration=base_transformer_model_config | get_data_config_set(
        values_per_day=12,
        shift_in_hours=12,
        input_window_length_in_days=160,
        output_window_length=1,
        output_window_offset=0
    ) | {
                                 "feature_config": feature_config | {"ignored_features":
                                     [
                                         # "fertility_probability",
                                         # "ov_over_probability",
                                         "days_relative_to_ov",
                                         "is_biphasic",
                                         "ov_day",
                                     ]},
                                 # "model_parameters": {
                                 #     "embed_dim": 256,
                                 #     "num_enc_layers": 2,
                                 #     "num_heads": 2,
                                 # },
                             },
    base_training_configuration=base_transformer_training_config,
    evaluation_configuration={
        "eval_functions_getter": get_fertility_based_eval_functions
    },
    model_config_kwargs_list=[
        {
            "model_parameters": {
                "embed_dim": 16,
                "num_enc_layers": 1,
                "num_heads": 1,
            },
        },
        {
            "model_parameters": {
                "embed_dim": 32,
                "num_enc_layers": 1,
                "num_heads": 1,
            },
        },
        {
            "model_parameters": {
                "embed_dim": 64,
                "num_enc_layers": 1,
                "num_heads": 1,
            },
        },
        {
            "model_parameters": {
                "embed_dim": 64,
                "num_enc_layers": 2,
                "num_heads": 2,
            },
        },
        {
            "model_parameters": {
                "embed_dim": 128,
                "num_enc_layers": 2,
                "num_heads": 2,
            },
        },
        {
            "model_parameters": {
                "embed_dim": 128,
                "num_enc_layers": 4,
                "num_heads": 4,
            },
        },
        {
            "model_parameters": {
                "embed_dim": 256,
                "num_enc_layers": 4,
                "num_heads": 4,
            },
        },
        {
            "model_parameters": {
                "embed_dim": 512,
                "num_enc_layers": 4,
                "num_heads": 4,
            },
        },
        {
            "model_parameters": {
                "embed_dim": 512,
                "num_enc_layers": 8,
                "num_heads": 8,
            },
        },
    ],
    # add model config provider function to set seq_len in model parameters
    additional_model_config_provider_fns={
        "model_parameters": {
            "seq_len": "eval(lambda x: x['input_window_length'])"
        }
    }
)

fertility_input_run = get_run_config(
    run_name="transformer_fertility_input_run",
    base_model_configuration=base_transformer_model_config | {
        "feature_config": feature_config | {"ignored_features":
            [
                # "fertility_probability",
                # "ov_over_probability",
                "days_relative_to_ov",
                "is_biphasic",
                "ov_day",
            ]},
        "model_parameters": {
            "embed_dim": 256,
            "num_enc_layers": 2,
            "num_heads": 2,
        },
    },
    base_training_configuration=base_transformer_training_config,
    evaluation_configuration={
        "eval_functions_getter": get_fertility_based_eval_functions
    },
    model_config_kwargs_list=input_run_kwargs_list,
    # add model config provider function to set seq_len in model parameters
    additional_model_config_provider_fns={
        "model_parameters": {
            "seq_len": "eval(lambda x: x['input_window_length'])"
        }
    }
)
best_config_run = get_run_config(
    run_name="transformer_best_config_run",
    base_model_configuration=base_transformer_model_config | get_data_config_set(
        values_per_day=12,
        shift_in_hours=12,
        input_window_length_in_days=160,
        output_window_length=1,
        output_window_offset=0
    ) |
                             {
                                 "feature_config": feature_config | {"ignored_features":
                                     [
                                         # "fertility_probability",
                                         # "ov_over_probability",
                                         "days_relative_to_ov",
                                         "is_biphasic",
                                         "ov_day",
                                     ]},
                                 "model_parameters": {
                                     "embed_dim": 512,
                                     "num_enc_layers": 4,
                                     "num_heads": 4,
                                 },
                             },
    base_training_configuration=base_transformer_training_config | {
        "learning_parameters": {
            "max_lr": 1e-4,
            # "learning_rate": base_lr * (batch_size / 4),
            "epochs": 30,
            "patience": 5,
        },
    },
    evaluation_configuration={
        "eval_functions_getter": get_fertility_based_eval_functions
    },
    model_config_kwargs_list=[],
    # add model config provider function to set seq_len in model parameters
    additional_model_config_provider_fns={
        "model_parameters": {
            "seq_len": "eval(lambda x: x['input_window_length'])"
        }
    }
)