from functools import partial

from configs.feature_config import feature_config

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

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

take_every_nth = int(288 / 12)
shift_in_hours = 12
input_window_length = (288 // take_every_nth) * 80
# output_window_length = (288 // take_every_nth) * 1
output_window_length = 1
output_window_offset = input_window_length + (288 // take_every_nth) * 0
window_shift = int((288 // take_every_nth) / 24 * shift_in_hours)
min_input_length_fraction_for_padding = ((288 // take_every_nth) * 4) / input_window_length

transformer_batch_size = 128


def get_transformer_run_config(
        run_name: str,
        run_description: str,
        input_window_length: int,
        output_window_length: int,
        take_every_nth: int,
        shift_in_hours: int,
        output_window_offset: int,
        batch_size: int,
        model_parameters: dict,
        max_lr: float,
        num_epochs: int,
        patience: int,
        feature_config: dict):
    """
    Get the configuration for the Transformer model
    Returns:
        run_configuration: configuration for the Transformer model
    """
    base_model_config = {
        "model_name": "transformer_regressor",
        "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",
            ]},
        "preprocessing": {
            "window_shift": int((288 // take_every_nth) / 24 * shift_in_hours),
            "take_every_nth": take_every_nth,
            "min_input_length_fraction_for_padding": min_input_length_fraction_for_padding,
        },
        "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,
        "model_parameters": {
            **model_parameters | {"seq_len": input_window_length},
        },
        "input_window_length": input_window_length,
        "output_window_length": output_window_length,
        "output_window_offset": output_window_offset,
    }
    base_training_config = {
        "batch_size": batch_size,
        "model_class": TransformerModel,
        "learning_parameters": {
            "learning_rate": max_lr * (batch_size / 4),
            # "learning_rate": base_lr * (batch_size / 4),
            "epochs": num_epochs,
            "patience": patience,
        },
        "loss_functions": [
            nn.MSELoss(),
            nn.BCEWithLogitsLoss(),
        ],
        "max_grad_norm": 1.0,
        "train_size": 0.7,
        "val_size": 0.15,
        "test_size": 0.15,
    }
    return {
        "name": run_name,
        "description": run_description,
        "model_configuration": base_model_config,
        "training_configuration": base_training_config
    }


run_configuration = {
    "runs": [
        get_transformer_run_config(
            run_name="transformer_ovulation_regression",
            run_description="Transformer regressor for ovulation prediction",
            input_window_length=input_window_length,
            output_window_length=output_window_length,
            take_every_nth=take_every_nth,
            shift_in_hours=shift_in_hours,
            output_window_offset=output_window_offset,
            batch_size=transformer_batch_size,
            model_parameters={
                "embed_dim": 128,
                "num_heads": 4,
                "num_enc_layers": 4,
            },
            max_lr=1e-5,
            num_epochs=20,
            patience=3,
            feature_config=feature_config,
        )
    ]
}