from functools import partial

from configs.feature_config import feature_config

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

from models.third_party.patch_tst.models.PatchTST import Model as PatchTST
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

max_lr = 1e-5
batch_size = 64

run_configuration = {
    "item_limit": 100,
    "runs": [
        {
            "name": "lstm_ovulation_regression",
            "description": "LSTM model for ovulation regression",
            "model_configuration": {
                "model_name": "patch_tst_regressor",
                "version": "1.0.0",
                "model_class": PatchTST,
                "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": ((288 // take_every_nth) * 4) / input_window_length,
                },
                "batch_fn": produce_window_batches,
                "collate_fn": simple_x_y_collate,
                "model_creation_fn": simple_model_creation,
                "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": {
                    "configs": {
                        # core
                        "seq_len": input_window_length,
                        "pred_len": output_window_length,
                        "seq_pred": False,
                        # model
                        "e_layers": 4,
                        "n_heads": 4,
                        "d_model": 128,
                        "d_ff": 128,
                        "dropout": 0.2,
                        "fc_dropout": 0.2,
                        "head_dropout": 0.0,
                        "individual": True,
                        # patch
                        # "patch_len": input_window_length,
                        "patch_len": int(288 / take_every_nth),
                        "stride": int(288 / take_every_nth / 2),
                        "padding_patch": 0,
                        # preprocessing
                        "revin": False,
                        "affine": False,
                        "subtract_last": False,
                        # decomp
                        "decomposition": True,
                        "kernel_size": 3,
                    }
                },
                "input_window_length": input_window_length,
                "output_window_length": output_window_length,
                "output_window_offset": output_window_offset,
            },
            "training_configuration": {
                "batch_size": batch_size,
                "model_class": PatchTST,
                "learning_parameters": {
                    # "learning_rate": base_lr,
                    "learning_rate": max_lr * (batch_size / 4),
                    "epochs": 10,
                    "patience": 3,
                },
                "loss_functions": [
                    nn.MSELoss(),
                    nn.BCEWithLogitsLoss(),
                ],
                "max_grad_norm": 1.0,
                "train_size": 0.7,
                "val_size": 0.15,
                "test_size": 0.15,
            }

        }
    ]
}