# Source code for pyFTS.models.ensemble.multiseasonal

    """
    Silva, P. C. L et al. Probabilistic Forecasting with Seasonal Ensemble Fuzzy Time-Series
    XIII Brazilian Congress on Computational Intelligence, 2017. Rio de Janeiro, Brazil.
    """

    import numpy as np
    from pyFTS.common import Util as cUtil
    from pyFTS.models.ensemble import ensemble
    from pyFTS.models.seasonal import cmsfts
    from pyFTS.probabilistic import ProbabilityDistribution
    from copy import deepcopy
    from joblib import Parallel, delayed
    import multiprocessing

    def train_individual_model(partitioner, train_data, indexer):
        pttr = str(partitioner.__module__).split('.')[-1]
        diff = "_diff" if partitioner.transformation is not None else ""
        _key = "msfts_" + pttr + str(partitioner.partitions) + diff + "_" + indexer.name

        print(_key)

        model = cmsfts.ContextualMultiSeasonalFTS(_key, indexer=indexer)
        model.append_transformation(partitioner.transformation)
        model.train(train_data, partitioner.sets, order=1)

        cUtil.persist_obj(model, "models/"+_key+".pkl")

        return model

    class SeasonalEnsembleFTS(ensemble.EnsembleFTS):
        def __init__(self, name, **kwargs):
            super(SeasonalEnsembleFTS, self).__init__(name="Seasonal Ensemble FTS", **kwargs)
            self.min_order = 1
            self.indexers = []
            self.partitioners = []
            self.is_multivariate = True
            self.has_seasonality = True
            self.has_probability_forecasting = True

        def update_uod(self, data):
            self.original_max = max(self.indexer.get_data(data))
            self.original_min = min(self.indexer.get_data(data))

        def train(self, data, **kwargs):
            self.original_max = max(self.indexer.get_data(data))
            self.original_min = min(self.indexer.get_data(data))

            num_cores = multiprocessing.cpu_count()

            pool = {}
            count = 0
            for ix in self.indexers:
                for pt in self.partitioners:
                    pool[count] = {'ix': ix, 'pt': pt}
                    count += 1

            results = Parallel(n_jobs=num_cores)(
                delayed(train_individual_model)(deepcopy(pool[m]['pt']), data, deepcopy(pool[m]['ix']))
                for m in pool.keys())

            for tmp in results:
                self.append_model(tmp)

            cUtil.persist_obj(self, "models/"+self.name+".pkl")

        def forecast_distribution(self, data, **kwargs):

            ret = []

            smooth = kwargs.get("smooth", "KDE")
            alpha = kwargs.get("alpha", None)

            uod = self.get_UoD()

            for k in data.index:

                tmp = self.get_models_forecasts(data.ix[k])

                if alpha is None:
                    tmp = np.ravel(tmp).tolist()
                else:
                    tmp = self.get_distribution_interquantile( np.ravel(tmp).tolist(), alpha)

                name = str(self.indexer.get_index(data.ix[k]))

                dist = ProbabilityDistribution.ProbabilityDistribution(smooth, uod=uod, data=tmp,
                                                                       name=name, **kwargs)

                ret.append(dist)

            return ret