#  config.py
#
# Created on Mon Aug 07 2023
# Author: Leon K. Martin
#
# Copyright (c) 2023 Charité Universitätsmedizin Berlin
#
"""
TVB-O Model Configuration Module
--------------------------------

This module provides utility functions to configure and manage TVB neural mass models using ontology-based configurations.
These configurations can be applied to model objects to change their default behavior or parameterize them as per
published studies or lab-tested configurations.

### Functions

- `get_default(NMM: str) -> dict`:
    Fetch default values for neural mass model parameters.

- `get_param_config(config_key: str, return_NMM: bool = False, use_classes: bool = False) -> Union[dict, Tuple[dict, str]]`:
    Fetch specific parameter configurations based on a reference key.

- `update_default(model_name: str, config_key: str) -> dict`:
    Update the default parameter configuration of a given model.

- `configure_model(tvb_model: tvb.simulator.models, config_key: str) -> tvb.simulator.models`:
    Apply a specific configuration to a TVB model object, altering its state variables and parameters.

### Usage

To fetch and apply configurations, use the respective functions:

```python
model_defaults = get_default("ExampleModel")
updated_defaults = update_default("ExampleModel", "custom_config")
configured_model = configure_model(model_instance, "custom_config")
```

## Dependencies

This module relies on the `ontology` module for fetching configurations. Ensure it's properly initialized and available.

## See Also

tvbo.ontology:
    For handling and fetching ontology-based configurations.

::: {.callout-note}
- Configurations are either sourced from peer-reviewed publications or have been extensively tested in labs.
- Ensure you've set up the environment and dependencies correctly for smooth operation of these utilities.
:::
"""

from os.path import join

import numpy as np

from tvbo.utils import report
from tvbo.ontology import owl as ontology


class ModelConfiguration:
    """
    A class that handles parameter settings for a given Neural Mass Model (NMM).

    Attributes
    ----------
    model : ontology.NMM
        Neural Mass Model instance or identifier.
    config : dict
        Configuration settings for the NMM.
    config_key : str
        Key identifying the current configuration.
    report : str
        Report generated by the `get_report` method.

    Parameters
    ----------
    NMM : ontology.NMM or str
        An instance or identifier of the Neural Mass Model.

    Methods
    -------
    apply_config(config_key)
        Apply a new configuration based on the provided key.
    get_report(format="latex", decimals=3, **kwargs)
        Generate a report for the current configuration.
    save_report(path=".", format="latex", **kwargs)
        Save the generated report to a file.
    """

    def __init__(self, NMM):
        """
        Parameters
        ----------
        NMM : ontology.NMM or str
            Neural Mass Model instance or identifier.
        """
        if isinstance(NMM, str):
            NMM = ontology.get_model(NMM)
        self.model = NMM
        self.available_configurations = NMM.direct_instances()
        self.config = get_default(NMM)
        self.config_key = "default"

    def apply_config(self, config_key):
        """
        Apply a new configuration based on the provided key.

        Parameters
        ----------
        config_key : str
            Key identifying the new configuration to apply.
        """
        self.config_key = config_key
        new_config = get_param_config(config_key, use_classes=True)
        for k, v in self.config.items():
            if k in new_config.keys():
                self.config[k] = new_config[k]["value"]

    def get_report(self, format="latex", decimals=3, **kwargs):
        """
        Generate a report for the current configuration.

        Parameters
        ----------
        format : str, optional
            Format of the report, by default "latex".
        decimals : int, optional
            Number of decimal places to use in the report, by default 3.
        **kwargs :
            Additional keyword arguments.

        Returns
        -------
        str
            Generated report.
        """
        self.report = report.parameter_report(self, format=format, **kwargs)
        return self.report

    def save_report(self, path=".", format="latex", **kwargs):
        """
        Save the generated report to a file.

        Parameters
        ----------
        path : str, optional
            Path to save the report, by default ".".
        format : str, optional
            Format of the report, by default "latex".
        **kwargs :
            Additional keyword arguments.
        """
        fpath = join(path, self.model.label.first() + "_" + self.config_key + ".tex")
        if format.lower() == "latex":
            with open(fpath, "w") as texfile:
                texfile.write(self.get_report(format="latex", **kwargs))


def get_default(NMM):
    """
    Retrieve the default values for a given Neural Mass Model (NMM).

    If the provided NMM is a string identifier, it attempts to fetch the actual NMM instance.
    It then collects the default values for each descendant class of the NMM.

    Parameters
    ----------
    NMM : ontology.NMM or str
        An instance or identifier of the Neural Mass Model.

    Returns
    -------
    dict
        A dictionary with descendant classes of the NMM as keys and their corresponding default values as values.

    Examples
    --------
    >>> get_default("some_NMM_identifier")
    {<descendant_class1>: 0.5, <descendant_class2>: 1.2, ...}

    """
    if isinstance(NMM, str):
        NMM = ontology.get_model(NMM)

    suffix = ontology.get_model_suffix(NMM)
    values = dict()
    for cls in NMM.descendants():
        v = cls.defaultValue.first()
        if v is not None:
            values[cls] = float(v)
    return values


def get_param_config(
    config_key="default", model=None, return_NMM=False, use_classes=False
):
    """
    Retrieve the parameter configuration for a given configuration key.

    This function searches the ontology for the provided configuration key and
    returns a dictionary with the relevant parameter details. Additionally,
    it checks for ancestors related to Neural Mass Model (NMM) and fetches the
    relevant variables and instances associated with the configuration.

    Parameters
    ----------
    config_key : str
        The key for which the configuration details are to be fetched.
    return_NMM : bool, optional
        If True, returns the Neural Mass Model associated with the configuration. Default is False.
    use_classes : bool, optional
        If True, the labels in the returned dictionary will use class objects instead of string labels. Default is False.

    Returns
    -------
    dict or tuple
        If `return_NMM` is False:
            A dictionary with the configuration details.
        If `return_NMM` is True:
            A tuple where the first element is the configuration details dictionary and the second is the NMM.

    Examples
    --------
    >>> get_param_config("some_config_key")
    {'some_variable': {'category': 'State Variable', 'range': [0.5, 1.2]}}

    >>> get_param_config("some_config_key", return_NMM=True)
    ({'some_variable': {'category': 'State Variable', 'range': [0.5, 1.2]}}, 'NMM_instance')
    """
    if config_key == "default":
        return ontology.get_default_values(model)

    config_dict = dict()
    if not isinstance(config_key, str):
        conf = config_key
    else:
        conf = ontology.onto.search_one(label=config_key)
    if ontology.onto.ModelConfiguration not in conf.is_a:
        return get_default(conf.uses_model.first())

    config_instance_of = conf.is_instance_of

    if isinstance(model, type(None)):
        for ci in config_instance_of:
            if ontology.onto.NeuralMassModel in ci.ancestors():
                NMM = ci
                break
            else:
                return get_default(conf.uses_model.first())

    elif isinstance(model, str):
        NMM = ontology.get_model(model)
    else:
        NMM = model

    default_params = ontology.get_default_values(NMM)
    config_params = {
        k: v.value.first() for k, v in ontology.get_model_parameters(conf).items()
    }
    default_params.update(config_params)
    if return_NMM:
        return default_params, NMM
    return default_params


def update_default(model_name, config_key=None, config_dict=None):
    # TODO: add all params to docstring
    """
    Update the default configuration for a specific model using a given configuration key.

    This function first fetches the default values for the provided model using the ontology.
    If the config_key is "default", it returns the fetched default values.
    Otherwise, it retrieves the parameter configuration associated with the config_key
    and updates the default values accordingly.

    Parameters
    ----------
    model_name : str
        The name of the model for which the default values need to be updated.
    config_key : str
        The reference key for the parameter configuration that will be used to update the default values.

    Returns
    -------
    dict
        The updated default configuration values for the model. If config_key is "default", the original default values are returned.

    Examples
    --------
    >>> update_default("example_model", "default")
    {'param1': 1.2, 'param2': 2.3}

    >>> update_default("example_model", "custom_config")
    {'param1': 1.8, 'param2': 2.3}

    Notes
    -----
    The function makes use of `ontology` for its operations. Ensure that `ontology` is properly initialized and contains the required data.
    """
    default = ontology.get_default_values(model_name)
    if config_key == "default":
        return default

    updated = default.copy()
    if config_key:
        config = get_param_config(config_key)
    else:
        config = config_dict
    for k, v in default.items():
        if k in config.keys():
            updated[k] = config[k]
    return updated


def configure_model(tvb_model, config_key):
    """Apply specific parameter configuration stored in TVB-O. All configurations have been either published in a
    peer-reviewed journal or tested extensively in our lab. Each configuration has a specific reference with further
    information.

    Parameters
    ----------
    tvb_model : tvb.simulator.models
        TVB model object
    config_key : str
        Reference-key for Parameter configuration.

    Returns
    -------
    tvb.simulator.models
        Updated model instance with parameter values according to specific configuration.

    """
    if config_key == "default":
        config_dict = ontology.get_default_values(tvb_model.tvbo_model)
    elif isinstance(config_key, dict):
        config_dict = config_key
    else:
        config_dict, NMM = get_param_config(config_key=config_key, return_NMM=True)
        tvb_model.tvbo_model = NMM.label.first()

    tvb_model.tvbo_config = config_dict

    for label, v in config_dict.items():
        # if v["category"] == "State Variable":
        #     tvb_model.state_variable_range[label] = v["range"]
        # elif v["category"] == "Parameter":
        tvb_model.__dict__[label] = np.array([v])

    tvb_model.configure()

    return tvb_model


def get_coupling_parameters(CF):
    if isinstance(CF, str):
        CF = ontology.get_coupling_function(CF)
    acr = CF.acronym.first()
    param_config = {
        p.label.first().replace(f"_{acr}", ""): (
            p.defaultValue.first() if p.defaultValue else 1
        )
        for p in CF.has_parameter
    }
    return param_config


def get_model_configurations(model):
    if isinstance(model, str):
        model = ontology.get_model(model)

    return list(
        set(
            [
                inst
                for inst in ontology.onto.ModelConfiguration.instances()
                if model in inst.is_a
            ]
        )
    )


############################################
#              LinkML Functions            #
############################################


def getattr_case_insensitive(obj, attr_name, default=None):
    """
    Get the value of an attribute from an object, ignoring case sensitivity.

    This function searches for an attribute with a case-insensitive match to the given `attr_name`
    within the `obj` object. If a match is found, the corresponding value is returned. If no match
    is found, the `default` value is returned.

    :param obj: The object to search for the attribute.
    :param attr_name: The name of the attribute to search for.
    :param default: The default value to return if no matching attribute is found.
    :return: The value of the matching attribute or the default value.
    """
    # Normalize the desired attribute name to lowercase for case-insensitive comparison
    attr_name_lower = attr_name.lower()
    # Iterate through each attribute in the object's dictionary
    for key, value in vars(obj).items():
        # If a case-insensitive match is found, return the corresponding value
        if key.lower() == attr_name_lower:
            return value
    # If no matching attribute is found, return the default value
    return default
