import abc
from typing import Callable
import numpy as np
from ..util import tsv_to_df
from .. import C
import cloudpickle
import os
import sys


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

    Properties
    ----------
    sum_stat_calculator: Callable
        A list of Callable to the summary statistics functions
    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.
    loc: str
        The morpheus simulation folder.
    output_file: srt
        A name of the file containing the simulation output.
    parameter_list: list
        A list of parameters/arguments to be passed to the summary
        statistics function
    ignore: list
        A list of keys to ignore for the simulation output. This is helpful
        to ignore information that cannot be avoided written by morpheus.
    read_data: Callable
        A Callable define by user to handle reading simulated data.
    """

    def __init__(
        self, sum_stat_calculator: Callable = None,
            # loc: str,
            output_file: str = C.OUTPUT_FILE,
            name: str = None,
            parameter_list: list = None,
            ignore: list = None,
            read_data: Callable = None,
    ):
        self.sum_stat_calculator = sum_stat_calculator
        self.name = name
        # self.loc = loc
        self.output_file = output_file
        self.parameter_list = parameter_list
        self.ignore = ignore
        self.read_data = read_data
        if self.ignore is None:
            self.ignore = []

    @abc.abstractmethod
    def __call__(self, loc):
        """
        Parameters
        ----------
        loc: str
            The morpheus simulation folder.
        """
        sumstat_edited = None
        if self.read_data is None:
            sumstat_edited = self.read_sumstat(loc)
        else:
            sumstat_edited = self.read_data(loc)
        if self.sum_stat_calculator is None:

            # sumstat_edited = self.read_sumstat(loc)
            #     tmp_sumstat = self.read_sumstat(loc)
            #     self.parameter_list[0] = tsv_to_df(loc, self.output_file)
            #     func = getattr(self.class_name, self.name)
            #     return self.sum_stat_calculator(*self.parameter_list)
            return sumstat_edited
        elif isinstance(self.sum_stat_calculator, dict):
            return self.call_sum_stat_dict(sumstat_edited)
        if isinstance(self.sum_stat_calculator, Callable):
            return self.sum_stat_calculator(sumstat_edited)

    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

    def read_sumstat(self, loc):
        df = tsv_to_df(loc, self.output_file)
        dct = df.to_dict(orient='list')
        new_dct = {}
        for key, val in dct.items():
            if key in self.ignore:
                continue
            if self.name is not None:
                _key = self.name + '__' + str(key)
            else:
                _key = str(key)
            new_dct[_key] = np.asarray(val)
        return new_dct

    # def call_sum_stat_dict(self, sumstats, function_name='main'):
    #     if len(self.sum_stat_calculator) == 0:
    #         return sumstats
    #
    #     if isinstance(self.sum_stat_calculator, Callable):
    #         sumstats = self.sum_stat_calculator(sumstats)
    #     elif isinstance(self.sum_stat_calculator, dict):
    #         if not set(self.sum_stat_calculator.keys()).issubset(
    #             sumstats.keys()
    #         ):
    #             raise ValueError(
    #                 "the keys on the 'ss_post_processing' does not "
    #                 "match the one on the summary statistics names."
    #             )
    #         if isinstance(list(self.sum_stat_calculator.values())[0],
    #         Callable):
    #             sumstats = self.
    #             _call_summary_statistics_use_function(sumstats)
    #         elif isinstance(
    #             list(self.sum_stat_calculator.values())[0], ModuleType
    #         ):
    #             sumstats = self._call_post_processing_ss_use_module(
    #                 sumstats, function_name
    #             )
    #     else:
    #         raise ValueError(
    #             f"the type of 'post_processing_ss' should be str or dict."
    #             f" However, {type(self.sum_stat_calculator)} was given."
    #         )
    #     return sumstats

    def call_sum_stat_dict(self, sumstats, function_name='main'):
        if len(self.sum_stat_calculator) == 0:
            return sumstats
        # generic_key_list = [key.split("__", 1)[1] for key in
        #                     list(self.sum_stat_calculator.keys())]
        generic_key_list = list(self.sum_stat_calculator.keys())
        if not set(generic_key_list).issubset(
                sumstats.keys()
        ):
            raise ValueError(
                "the keys on the 'sumstats' does not "
                "match the one on the summary statistics names."
            )
        # sumstat_dict = {}
        # for key, val in self.sum_stat_calculator.items():
        #
        #     if isinstance(val, Callable):
        #         sumstats = self.
        #         _call_summary_statistics_use_function(sumstats)
        #     elif isinstance(val, ModuleType):
        #         sumstats = self._call_post_processing_ss_use_module(
        #             sumstats, function_name
        #         )
        #     sumstat_dict[key] = sumstats
        sumstat_dict = {}
        for key, function in self.sum_stat_calculator.items():
            try:
                if type(function) == bytes:
                    function = unserialise_function(function)
                post_summary_statistics = function(sumstats)
                sumstat_dict[key] = post_summary_statistics[key]
            except Exception as e:
                raise RuntimeError(
                    f"the selected ss_post_processing function "
                    f"can not be called. `{e}"
                )

        return sumstat_dict

    def _call_summary_statistics_use_module(self, sumstats, function_name):
        sumstat_pp = {}
        for key, _module in self.sum_stat_calculator.items():
            try:
                func = getattr(
                    self.sum_stat_calculator[key], function_name, None
                )
                sumstat_pp[key] = func({key: sumstats[key]})
            except Exception as e:
                raise RuntimeError(
                    f"the selected ss_post_processing function "
                    f"can not be called. Be sure that the main "
                    f"function called `main(). `{e}"
                )
        return sumstat_pp

    def _call_summary_statistics_use_dict(self, sumstats):
        sumstat_pp = {}
        for key, function in self.sum_stat_calculator.items():
            try:
                post_summary_statistics = function({key: sumstats[key]})
                sumstat_pp[key] = post_summary_statistics
            except Exception as e:
                raise RuntimeError(
                    f"the selected ss_post_processing function "
                    f"can not be called. `{e}"
                )
        return sumstat_pp

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


# class NoSummaryStatistics(SummaryStatistics):
#     # def __init__(self, name: str = "IdSumstat"):
#     #     if name != 'IdSumstat':
#     #         name = name + '__' + 'IdSumstat'
#     #     super().__init__(name)
#
#     def __call__(self, loc, output_file):
#         """
#         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 = {}
#         for key, val in dct.items():
#             if key in self.ignore:
#                 continue
#             if self.name is not None:
#                 _key = self.name + '__' + str(key)
#             else:
#                 _key = str(key)
#             new_dct[key] = np.asarray(val)
#         return new_dct


def unserialise_function(uns_func):
    path = os.getcwd()
    sys.path.append(path)
    func = cloudpickle.loads(uns_func)
    return func
