from copy import deepcopy
from typing import Callable

from vsm_datascience_common import constants


def recursive_dict_update(base_dict: dict, update_dict: dict) -> dict:
    base_dict = deepcopy(base_dict)
    for key, value in update_dict.items():
        if isinstance(value, dict) and key in base_dict:
            base_dict[key] = recursive_dict_update(base_dict[key], value)
        else:
            base_dict[key] = value

    return base_dict


def recursive_provider_dict_update(base_dict: dict,
                                   provider_fns: dict,
                                   context: dict = None) -> dict:
    """
    Recursively updates a base dictionary with values from provider functions.

    Provider functions get intermediate dictionary as input and return a value.
    Args:
        base_dict: base dictionary to update.
        provider_fns: provider functions to use for updating the dictionary.
        context: base dict to use as context in recursive calls, defaults to None.

    Returns:
        A new dictionary with updated values from provider functions.

    """

    if context is None:
        base_dict = deepcopy(base_dict)
        context = base_dict

    for key, provider_fn in provider_fns.items():
        if isinstance(provider_fn, dict):
            # if the value is a dictionary, recursively update it
            base_dict[key] = recursive_provider_dict_update(base_dict.get(key, {}), provider_fn, context)
        elif callable(provider_fn):
            # call the provider function with the current base_dict as input
            base_dict[key] = provider_fn(context)
        else:
            # if value is string and starts with "eval(", evaluate it
            if isinstance(provider_fn, str) and provider_fn.startswith("eval(") and provider_fn.endswith(")"):
                # evaluate the string as a Python expression
                eval_fn = eval(provider_fn[5:-1])
                base_dict[key] = eval_fn(context)
            else:
                # if the value is not callable, just set it
                base_dict[key] = provider_fn

    return base_dict


def config_generator(config_gen_fn: Callable,
                     fixed_params: dict,
                     variable_param_configs: list,
                     ) -> list:
    configs = list()
    for variable_config in variable_param_configs:
        config = fixed_params.copy()
        for param_name, param_value in variable_config.items():
            # update the config with the variable parameter, values can be None, if default should be used
            if param_value is not None:
                config[param_name] = param_value
        configs.append(config_gen_fn(**config))

    return configs


def get_run_config(
        run_name: str,
        base_model_configuration: dict,
        base_training_configuration: dict,
        evaluation_configuration: dict = None,
        model_config_kwargs_list: list = None,
        training_config_kwargs_list: list = None,
        *args, **kwargs) -> dict:
    """
    Generates a list of training configurations for a given model and training setup.

    Args and kwargs are forwarded to the `get_training_config` function.
    Args:
        run_name: name of the run.
        base_model_configuration: base model configuration dictionary.
        base_training_configuration: base training configuration dictionary.
        evaluation_configuration: evaluation configuration dictionary, defaults to None.
        model_config_kwargs_list: list of dictionaries containing additional keyword arguments for model configuration.
        training_config_kwargs_list: list of dictionaries containing additional keyword arguments for training configuration.
    Returns:
        A dictionary containing the run configuration with a list of training configurations.

    """

    if model_config_kwargs_list is None:
        model_config_kwargs_list = list()
    if training_config_kwargs_list is None:
        training_config_kwargs_list = list()

    training_configs = list()
    num_trainings = max(len(model_config_kwargs_list), len(training_config_kwargs_list), 1)

    num_model_kwargs = len(model_config_kwargs_list)
    num_training_kwargs = len(training_config_kwargs_list)

    # if num_model_kwargs == 0 and num_training_kwargs == 0:
    #     raise ValueError("Both model_config_kwargs_list and training_config_kwargs_list cannot be empty.")

    for i in range(num_trainings):
        # list length can either be equal, one, or zero
        if num_model_kwargs == 0:
            model_config_kwargs = {}
        elif num_model_kwargs == 1:
            model_config_kwargs = model_config_kwargs_list[0]
        else:
            model_config_kwargs = model_config_kwargs_list[i]

        if num_training_kwargs == 0:
            training_config_kwargs = {}
        elif num_training_kwargs == 1:
            training_config_kwargs = training_config_kwargs_list[0]
        else:
            training_config_kwargs = training_config_kwargs_list[i]

        training_config = get_training_config(
            training_name=f"{run_name}_{i}",
            training_description=f"{run_name} training configuration {i}",
            base_model_config=base_model_configuration,
            model_config_kwargs=model_config_kwargs,
            base_training_config=base_training_configuration,
            training_config_kwargs=training_config_kwargs,
            *args,
            **kwargs,
        )
        training_configs.append(training_config)
    return {
        "name": run_name,
        "runs": training_configs,
        "evaluation_configuration": evaluation_configuration if evaluation_configuration is not None else {}
    }


def get_training_config(
        training_name: str,
        training_description: str,
        base_model_config: dict,
        model_config_kwargs: dict,
        base_training_config: dict,
        training_config_kwargs: dict,
        additional_model_config_provider_fns: dict = None,
        additional_training_config_provider_fns: dict = None,
        *args, **kwargs) -> dict:
    """
    Generates a training configuration for a given model and training setup.
    Args:
        training_name: Name of the training.
        training_description: Description of the training.
        base_model_config: Base model configuration dictionary.
        model_config_kwargs: Additional keyword arguments for model configuration.
        base_training_config: Base training configuration dictionary.
        training_config_kwargs: Additional keyword arguments for training configuration.
        additional_model_config_provider_fns: Additional provider functions for parts of the config that depend on the config itself
                                such as input length as model parameter.
        additional_training_config_provider_fns: Additional provider functions for parts of the config that depend on the config itself
                                such as input length as model parameter.
    Returns:
        A dictionary containing the training configuration with model and training parameters.
    """

    # recursively merge the base model configuration with the provided kwargs
    model_configuration = recursive_dict_update(base_model_config, model_config_kwargs)

    # update based on provider functions
    if additional_model_config_provider_fns:
        model_configuration = recursive_provider_dict_update(model_configuration, additional_model_config_provider_fns)

    required_model_configuration_keys = [
        "model_name",
        "version",
        "model_class",
        "feature_config",
        "preprocessing",
        "batch_fn",
        "collate_fn",
        "model_creation_fn",
        "model_save_fn",
        "model_load_fn",
        "batch_loss_fn",
        "actual_fn",
        "predict_fn",
        "model_parameters",
        "input_window_length",
        "output_window_length",
        "output_window_offset"
    ]
    if not all(key in model_configuration for key in required_model_configuration_keys):
        missing_keys = [
            key for key in required_model_configuration_keys if key not in base_model_config
        ]
        raise ValueError(f"Missing required keys in model configuration: {', '.join(missing_keys)}")

    training_configuration = recursive_dict_update(base_training_config, training_config_kwargs)

    # update based on provider functions
    if additional_training_config_provider_fns:
        training_configuration = recursive_provider_dict_update(training_configuration,
                                                                additional_training_config_provider_fns)

    # if "model_class" is not in training_configuration, set it to the model class from the model configuration
    if "model_class" not in training_configuration:
        training_configuration["model_class"] = base_model_config["model_class"]

    required_training_configuration_keys = [
        "model_class",
        "batch_size",
        "learning_parameters",
        "max_grad_norm",
        "train_size",
        "val_size",
        "test_size"
    ]

    if not all(key in training_configuration for key in required_training_configuration_keys):
        missing_keys = [
            key for key in required_training_configuration_keys if key not in training_configuration
        ]
        raise ValueError(f"Missing required keys in training configuration: {', '.join(missing_keys)}")

    return {
        "name": training_name,
        "description": training_description,
        "model_configuration": model_configuration,
        "training_configuration": training_configuration,
    }


def get_data_config_set(
        values_per_day: int,
        shift_in_hours: int,
        input_window_length_in_days: int,
        output_window_length: int,
        output_window_offset: int,
) -> dict:
    measurements_per_day = constants.MEASUREMENTS_PER_DAY
    take_every_nth = int(measurements_per_day / values_per_day)
    input_window_length = (measurements_per_day // take_every_nth) * input_window_length_in_days
    return {
        "input_window_length": input_window_length,
        "output_window_length": output_window_length,
        "output_window_offset": input_window_length + output_window_offset,
        "preprocessing": {
            "window_shift": max(int((measurements_per_day // take_every_nth) / 24 * shift_in_hours), 1),
            "take_every_nth": take_every_nth,
            "min_input_length_fraction_for_padding":
                ((measurements_per_day // take_every_nth) * 4) / input_window_length,
        }
    }


input_run_kwargs_list = [
    get_data_config_set(
        values_per_day=12,
        shift_in_hours=12,
        input_window_length_in_days=10,
        output_window_length=1,
        output_window_offset=0
    ),
    get_data_config_set(
        values_per_day=12,
        shift_in_hours=12,
        input_window_length_in_days=20,
        output_window_length=1,
        output_window_offset=0
    ),
    get_data_config_set(
        values_per_day=12,
        shift_in_hours=12,
        input_window_length_in_days=40,
        output_window_length=1,
        output_window_offset=0
    ),
    get_data_config_set(
        values_per_day=12,
        shift_in_hours=12,
        input_window_length_in_days=80,
        output_window_length=1,
        output_window_offset=0
    ),
    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
    ),
    get_data_config_set(
        values_per_day=1,
        shift_in_hours=12,
        input_window_length_in_days=20,
        output_window_length=1,
        output_window_offset=0
    ),
    get_data_config_set(
        values_per_day=2,
        shift_in_hours=12,
        input_window_length_in_days=20,
        output_window_length=1,
        output_window_offset=0
    ),
    get_data_config_set(
        values_per_day=4,
        shift_in_hours=12,
        input_window_length_in_days=20,
        output_window_length=1,
        output_window_offset=0
    ),
    get_data_config_set(
        values_per_day=12,
        shift_in_hours=12,
        input_window_length_in_days=20,
        output_window_length=1,
        output_window_offset=0
    ),
    get_data_config_set(
        values_per_day=24,
        shift_in_hours=12,
        input_window_length_in_days=20,
        output_window_length=1,
        output_window_offset=0
    ),
    get_data_config_set(
        values_per_day=48,
        shift_in_hours=12,
        input_window_length_in_days=20,
        output_window_length=1,
        output_window_offset=0
    ),
    get_data_config_set(
        values_per_day=72,
        shift_in_hours=12,
        input_window_length_in_days=20,
        output_window_length=1,
        output_window_offset=0
    ),
    get_data_config_set(
        values_per_day=288,
        shift_in_hours=12,
        input_window_length_in_days=20,
        output_window_length=1,
        output_window_offset=0
    )
]