import abc
from typing import Callable
import numpy as np
from ..util import tsv_to_df


class SumstatFun:
    """
    Base class for summary statistic functions working on
    morpheus models.

    Properties
    ----------

    name: str
        An identifier for his summary statistic. Used by pyabc
        to generate keys. The names of all summary statistics
        used jointly should be unique.
    """

    def __init__(self,
                 name: str,
                 class_name: str = None,
                 parameter_list: list = None):
        self.name = name
        self.class_name = class_name
        self.parameter_list = parameter_list

    @abc.abstractmethod
    def __call__(self, loc: str, output_file):
        """
        Parameters
        ----------

        loc: str
            The morpheus simulation folder.
        """
        if self.class_name and self.parameter_list is not None:

            self.parameter_list[0] = tsv_to_df(loc, output_file)
            func = getattr(self.class_name, self.name)
            result = func(*self.parameter_list)
            return result
        else:
            raise NotImplementedError()

    def create_own_sumstat_extractor(self):
        """
        Extract from the whole `sumstat` dictionary those entries
        that were generated by the present sumstat function.

        This is necessary since before the calculation of distances
        all sumstats are aggregated in a format which can also be
        stored.

        Returns
        -------

        own_sumstat_extractor: Callable[dict, dict]
            A function that extract the sumstats generated by the given
            sumstat fun based on the name attribute.
        """
        # extract all dependencies to avoid pickling issues
        name = self.name

        def own_sumstat_extractor(sumstat: dict):
            """
            Parameters
            ----------
            sumstat: dict
                The aggregated sumstat dictionary.

            Returns
            -------
            own_sumstat: dict
                The extracted values.
            """
            own_sumstat = {
                key: sumstat[key] for key in sumstat if
                key == name or key.startswith(name + "__")
            }
            return own_sumstat

        return own_sumstat_extractor

    def preprocess_distance_function(
            self, distance: Callable[[dict, dict], float]):
        """
        Preprocess the summary statistics to extract those values that
        were generated by the present summary statistics function.

        Parameters
        ----------
        distance: Callable[[dict, dict], float]
            The distance to apply to the given summary statistics.
            Note that subclasses may like to have a default distance to
            use for their data, such that the user does not need to
            specify anything at all.

        Returns
        -------
        preprocessed_distance: Callable[[dict, dict], float]
            A callable applying the `distance` to the own summary statistics
            only.
        """
        own_sumstat_extractor = self.create_own_sumstat_extractor()

        def preprocessed_distance(x, y):
            x = own_sumstat_extractor(x)
            y = own_sumstat_extractor(y)
            return distance(x, y)

        return preprocessed_distance


class IdSumstatFun(SumstatFun):

    def __init__(self, name: str = "IdSumstat"):
        if name != 'IdSumstat':
            name = name + '__' + 'IdSumstat'
        super().__init__(name)

    def __call__(self, loc: str, ignore: list,
                 output_file: str):
        """
        Parameters
        ----------
        loc: str
            location of the simulation output.
        ignore: list
            a list of keys to ignore for the simulation output. This is helpful
            to ignore information that cannot be avoid written by morpheus.
        """

        df = tsv_to_df(loc, output_file)
        dct = df.to_dict(orient='list')
        new_dct = dict()
        for key, val in dct.items():
            if key in ignore:
                continue
            _key = self.name + '__' + str(key)
            new_dct[_key] = np.array(val)
        return new_dct
