from typing import List

import connectome.model as model_library
import connectome.noise as noise_library
from connectome import analysis as network_analysis
from connectome.analysis import AnalysisBuilder
from modelflow import DynamicGraph
from util.random_variables import Distribution


class InvalidNoiseType(Exception):
    pass


class InvalidModelType(Exception):
    pass


class ParameterDict(dict):
    def fixed_parameters(self):
        result = ParameterDict()
        for parameter_name, parameter_value in self.items():
            try:
                parameter_value.keys()
            except AttributeError:
                result[parameter_name] = parameter_value
        return result

    def variable_parameters(self):
        result = ParameterDict()
        for parameter_name, parameter_value in self.items():
            try:
                parameter_value.keys()
                result[parameter_name] = parameter_value
            except AttributeError:
                pass
        return result

    def prefixed_parameters(self, prefix):
        return ParameterDict({prefix + key: value for key, value in self.items()})


class PipelineBuilder:
    """
    Build models, priors, noise, analysis
    and plug them together to a pipeline.

    Parameters
    ----------

    model_config_dict: dict
        Dictionary with model definitions.
    """
    network_prefix = "network_"
    noise_prefix = "noise_"
    analysis_prefix = "analysis_"

    def __init__(self, model_config_dict):
        self.model_config_dict = model_config_dict

    def __repr__(self):
        return "<{} models={}>".format(self.__class__.__name__, self.defined_models())

    def defined_models(self) -> List[str]:
        """
        Return available models.

        Returns
        -------

        models_list: list
            Sorted list of models.
        """
        return list(sorted(list(self.model_config_dict['models']['model_specific_parameters'].keys())))

    def build_pipeline(self, model_name):
        """
        Build a pipeline consisting of:

          1. model,
          2. noise,
          3. analysis.

        Parameters
        ----------

        model_name: str
            Name of the model.

        Returns
        -------

        pipeline: Node
            The built pipeline.
        """
        if model_name not in self.defined_models():
            raise InvalidModelType("Model", model_name, "not defined in configuration dictionary.")

        model = self.build_model(model_name)
        noise = self._make_noise()
        analysis = self.build_analysis()
        pipeline = concatenate_model_noise_analysis(model, noise, analysis)
        return pipeline

    def build_model_with_noise(self, model_name: str):
        """

        Parameters
        ----------
        model_name: str
            Name of the model.

        Returns
        -------

        model_with_noise: Node
            The noisy model.
        """
        if model_name not in self.defined_models():
            raise InvalidModelType("Model", model_name, "not defined in configuration dictionary.")

        model = self.build_model(model_name)
        noise = self._make_noise()
        model_with_noise = concatenate_model_noise_analysis(model, noise)
        return model_with_noise

    def build_prior(self, model_name):
        """
        Build the prior for a model.

        Parameters
        ----------

        model_name: str
            Name of the model for which to build the prior.


        Returns
        -------

        prior: Distribution
            The built prior.
        """
        noise_parameters = self._noise_parameters().variable_parameters().prefixed_parameters(self.noise_prefix)
        model_parameters = self._model_parameters(model_name).variable_parameters().prefixed_parameters(self.network_prefix)

        combined_parameters = noise_parameters.copy()
        combined_parameters.update(model_parameters)
        return Distribution.from_dictionary_of_dictionaries(combined_parameters)

    def _make_noise(self):
        try:
            noise = getattr(noise_library, self._noise_type())(**self._noise_parameters().fixed_parameters())
        except AttributeError:
            raise InvalidNoiseType(self._noise_type())
        return noise

    def build_model(self, model_name):
        """
        Build the model.

        Parameters
        ----------
        model_name: str
            Name of the model.

        Returns
        -------

        model: Node
            The built model.
        """
        try:
            model = getattr(model_library, model_name)(**self._model_parameters(model_name).fixed_parameters())
        except AttributeError:
            raise InvalidModelType(model_name)
        return model

    def _noise_type(self):
        return self.model_config_dict['noise']['type']

    def _noise_parameters(self):
        return ParameterDict(self.model_config_dict['noise']['parameters'])

    def _model_parameters(self, model_name):
        joint_parameters = self.model_config_dict['models']['joint_parameters']
        model_specific_parameters = self.model_config_dict['models']['model_specific_parameters'][model_name]
        parameters = joint_parameters.copy()
        parameters.update(model_specific_parameters)
        return ParameterDict(parameters)

    def build_analysis(self):
        """
        Build the analysis.

        Returns
        -------

        analysis: DynamicGraph
            The analysis.
        """
        analysis_builder = AnalysisBuilder()
        analysis_parameters = self.model_config_dict['analysis']
        for analysis, kwargs in analysis_parameters.items():
            analysis_builder.add(getattr(network_analysis, analysis), **kwargs)
        analysis = analysis_builder.get_result()
        return analysis


def get_open_inputs(node):
    return [inp for inp in node.input_names if inp not in node.inputs_set_at_initialization and inp != "network"]


def prefix_list(prefix, lst):
    return [prefix + item for item in lst]


def concatenate_model_noise_analysis(model, noise, analysis=None):
    """
    Concatenate model, noise and analysis to a pipeline.

    Build a chain of model, noise and analysis.
    Resulting computational node has as input parameters
    all parameters which were *not* already set in model, noise and analysis
    prefixed by "model\_", "noise\_" and "analysis\_" respectively.

    Parameters
    ----------

    model: Node
        The network model.

    noise:  Node
        The network noise.

    analysis: Node
        The network analysis.
    """
    input_mapping = [(model, PipelineBuilder.network_prefix, get_open_inputs(model)),
                     (noise, PipelineBuilder.noise_prefix, get_open_inputs(noise))]
    if analysis is not None:
        input_mapping.append((analysis, PipelineBuilder.analysis_prefix, get_open_inputs(analysis)))

    result = DynamicGraph()
    result.add_inputs(sum(map(lambda x: prefix_list(x[1], x[2]), input_mapping), []))
    for node, prefix, open_inputs in input_mapping:
        for inp in open_inputs:
            result.connect_input_to_node(prefix + inp, node, inp)
    result.connect_node_to_node(model, noise, "network")
    if analysis is not None:
        result.connect_node_to_node(noise, analysis, "network")
        result.set_output(analysis)
    else:
        result.set_output(noise)
    return result
