import numpy as np

from astropy import units as u


class RegressorClassifierBase:
    """
    This class collects one model for every camera type -- given by `cam_id_list` --
    to get an estimate for the energy of an air-shower.
    The class interfaces with `scikit-learn` in that it relays all function calls not
    exeplicitly defined here through `.__getattr__` to any model in its collection.
    This gives us the full power and convenience of scikit-learn and we only have to
    bother about the high-level interface.

    The fitting approach is to train on single images; then separately predict an image's
    energy and combine the different predictions for all the telscopes in each event in
    one single energy  estimate.

    TODO: come up with a better merging approach to combine the estimators of the various
          camera types

    Parameters
    ----------
    model : scikit-learn model, optional (default: RandomForestRegressor)
        the model you want to use to estimate the shower energies
    cam_id_list : list of strings
        list of identifiers to differentiate the various sources of the images; could be
        the camera IDs or even telescope IDs. We will train one model for each of the
        identifiers.
    unit : 1 or astropy unit, optional (default: 1)
        scikit-learn regressors don't work with astropy unit. so, tell in advance in which
        unit we want to deal here in case we need one.
    kwargs
        arguments to be passed on to the constructor of the regressors
    """
    def __init__(self, model, cam_id_list, unit=1, **kwargs):

        self.model_dict = {}
        self.unit = unit
        for cam_id in cam_id_list or []:
            self.model_dict[cam_id] = model(**kwargs)

    def __getattr__(self, attr):
        """
        We interface this class with the "first" model in `.model_dict` and relay all
        function calls over to it. This gives access to all the fancy implementations in
        the scikit-learn classes right from this class and we only have to overwrite the
        high-level interface functions we want to adapt to our use-case.
        """
        return getattr(next(iter(self.model_dict.values())), attr)

    def __str__(self):
        """ return the class name of the used model """
        return str(next(iter(self.model_dict.values()))).split("(")[0]

    def reshuffle_event_list(self, X, y):
        """
        I collect the data event-wise as dictionaries of the different camera identifiers
        but for the training it is more convenient to have a dictionary containing flat
        arrays of the images. This function flattens the `X` and `y` arrays accordingly.

        This is only for convenience during the training and testing phase, later on, we
        won't collect data first event- and then telescope-wise. That's why this is a
        separate function and not implemented in `.fit`.

        Parameters
        ----------
        X : list of dictionaries of lists
            collection of training features see Notes section for a sketch of
            the container's expected shape
        y : list of astropy quantities
            list of the training targets (i.e. shower energies) for all events

        Returns
        -------
        trainFeatures, trainTarget : dictionaries of lists
            flattened containers to be handed over to `.fit`

        Raises
        ------
        KeyError
            in case `X` contains keys that were not provided with `cam_id_list` during
            `.__init__` or `.load`.

        Notes
        -----
        `X` and `y` are expected to be lists of events:
        `X = [event_1, event_2, ..., event_n]`
        `y = [energy_1, energy_2, ..., energy_n]`

        with `event_i` being a dictionary:
        `event_i = {tel_type_1: list_of_tels_1, ...}`
        and `energy_i` is the energy of `event_i`

        `list_of_tels_i` being the list of all telescope of type `tel_type_1` containing
        the features to fit and predict the shower energy:
        `list_of_tels_i = [[tel_1_feature_1, tel_1_feature_2, ...],
                           [tel_2_feature_1, tel_2_feature_2, ...], ...]

        after reshuffling, the resulting lists will look like this
        `trainFeatures = {tel_type_1: [features_event_1_tel_1,
                                       features_event_1_tel_2,
                                       features_event_2_tel_1, ...],
                          tel_type_2: [features_event_1_tel_3,
                                       features_event_2_tel_3,
                                       features_event_2_tel_4, ...],
                          ...}`
        `trainTarget` will be a dictionary with the same keys and a lists of energies
        corresponding to the features in the `trainFeatures` lists.
        """

        trainFeatures = {a: [] for a in self.model_dict}
        trainTarget = {a: [] for a in self.model_dict}

        for evt, target in zip(X, y):
            for cam_id, tels in evt.items():
                try:
                    # append the features-lists of the current event and telescope type to
                    # the flat list of features-lists for this telescope type
                    trainFeatures[cam_id] += tels
                except KeyError:
                    raise KeyError("cam_id '{}' in X but no model defined: {}"
                                   .format(cam_id, [k for k in self.model_dict]))

                try:
                    # add an energy-entry for every feature-list
                    trainTarget[cam_id] += [target.to(self.unit).value]*len(tels)
                except AttributeError:
                    # in case the energy is not given as an astropy quantity
                    # let's hope that the user keeps proper track of the unit themself
                    trainTarget[cam_id] += [target]*len(tels)
        return trainFeatures, trainTarget

    def fit(self, X, y):
        """
        This function fits a model against the collected features; separately for every
        telescope identifier.

        Parameters
        ----------
        X : dictionary of lists of lists
            Dictionary that maps the telescope identifiers to lists of feature-lists.
            The values of the dictionary are the lists `scikit-learn` regressors train on
            and are supposed to comply to their format requirements
            e.g. each featuer-list has to contain the same features at the same position
        y : dictionary of lists
            the energies corresponding to all the feature-lists of `X`

        Returns
        -------
        self

        Raises
        ------
        KeyError
            in case `X` contains keys that are either not in `y` or were not provided
            before with `cam_id_list`.

        """
        for cam_id in X:
            if cam_id not in y:
                raise KeyError("cam_id '{}' in X but not in y: {}"
                               .format(cam_id, [k for k in y]))

            if cam_id not in self.model_dict:
                raise KeyError("cam_id '{}' in X but no model defined: {}"
                               .format(cam_id, [k for k in self.model_dict]))

            # for every `cam_id` train one model
            self.model_dict[cam_id].fit(X[cam_id], y[cam_id])

        return self

    # def predict(self, X, cam_id=None):
    #     """
    #     In the tradition of scikit-learn, `.predict` takes a "list of feature-lists" and
    #     returns an estimate for targeted quantity for every  set of features.
    #
    #     Parameters
    #     ----------
    #     X : list of lists of floats
    #         the list of feature-lists
    #     cam_id : any (e.g. string, int), optional
    #         identifier of the singular camera type to consider here
    #         if not set, `.reg_dict` is assumed to have a single key which is used in place
    #
    #     Returns
    #     -------
    #     predict : list of floats
    #         predictions for the target quantity for every set of features given in `X`
    #
    #     Raises
    #     ------
    #     ValueError
    #         if `cam_id is None` and the number of registered models is not 1
    #     """
    #
    #     if cam_id is None:
    #         if len(self.model_dict) == 1:
    #             cam_id = next(iter(self.model_dict.keys()))
    #         else:
    #             raise ValueError("you need to provide a cam_id")
    #
    #     return self.model_dict[cam_id].predict(X)*self.energy_unit

    def save(self, path):
        """
        saves the models in `.reg_dict` each in a separate pickle to disk

        TODO: investigate more stable containers to write out models than joblib dumps

        Parameters
        ----------
        path : string
            Path to store the different models.
            Expects to contain `{cam_id}` or at least an empty `{}` to replace it with the
            keys in `.reg_dict`.

        """

        from sklearn.externals import joblib
        for cam_id, model in self.model_dict.items():
            try:
                # assume that there is a `{cam_id}` keyword to replace in the string
                joblib.dump(model, path.format(cam_id=cam_id))
            except IndexError:
                # if not, assume there is a naked `{}` somewhere left
                # if not, format won't do anything, so it doesn't break but will overwrite
                # every pickle with the following one
                joblib.dump(model, path.format(cam_id))

    @classmethod
    def load(cls, path, cam_id_list, unit=1):
        """
        Load the pickled dictionary of model from disk, create a husk
        `cls` instance and fill the model dictionary.

        Parameters
        ----------
        path : string
            the path where the pre-trained, pickled regressors are stored
            `path` is assumed to contain a `{cam_id}` keyword to be replaced by each
            camera identifier in `cam_id_list` (or at least a naked `{}`).
        cam_id_list : list
            list of camera identifiers like telescope ID or camera ID and the assumed
            distinguishing feature in the filenames of the various pickled regressors.
        unit : 1 or astropy unit, optional (default: 1)
            scikit-learn regressor/classifier do not work with units. so append this one
            to the predictions in case you deal with unified targets (like energy).
            assuming that the models where trained with consistent units.
clf
        self : RegressorClassifierBase
            in derived classes, this will return a ready-to-use instance of that class to
            predict any problem you have trained for
        """
        from sklearn.externals import joblib

        # need to get an instance of this class
        # `cam_id_list=None` prevents `.__init__` to initialise `.model_dict` itself,
        # since we are going to set it with the pickled models manually
        self = cls(cam_id_list=None, unit=unit)
        for key in cam_id_list:
            try:
                # assume that there is a `{cam_id}` keyword to replace in the string
                self.model_dict[key] = joblib.load(path.format(cam_id=key))
            except IndexError:
                # if not, assume there is a naked `{}` somewhere left
                # if not, format won't do anything, so it doesn't matter
                # though this will load the same model for every `key`
                self.model_dict[key] = joblib.load(path.format(key))

        return self

    def show_importances(self, feature_labels=None):
        """
        Creates a matplotlib figure that shows the importances of the different features
        for the various trained models in a grid of barplots.
        The features are sorted by descending importance.

        Parameters
        ----------
        feature_labels : list of strings, optional
            a list of the feature names in proper order
            to be used as x-axis tick-labels

        Returns
        -------
        fig : matplotlib figure
            the figure holding the different bar plots
        """
        import matplotlib.pyplot as plt
        n_tel_types = len(self.model_dict)
        n_cols = np.ceil(np.sqrt(n_tel_types)).astype(int)
        n_rows = np.ceil(n_tel_types / n_cols).astype(int)

        fig, axs = plt.subplots(nrows=n_rows, ncols=n_cols, squeeze=False)
        plt.suptitle("Feature Importances")
        for i, (cam_id, model) in enumerate(self.model_dict.items()):
            plt.sca(axs.ravel()[i])
            importances = model.feature_importances_
            bins = range(importances.shape[0])
            plt.title(cam_id)
            if feature_labels:
                importances, s_feature_labels = \
                    zip(*sorted(zip(importances, feature_labels), reverse=True))
                plt.xticks(bins, s_feature_labels, rotation=17)
            plt.bar(bins, importances,
                    color='r', align='center')
        return fig
