# Source code for pyFTS.common.Util

    """
    Common facilities for pyFTS
    """

    import time
    import matplotlib.pyplot as plt
    import numpy as np
    import pandas as pd
    import matplotlib.cm as cmx
    import matplotlib.colors as pltcolors
    from pyFTS.probabilistic import ProbabilityDistribution
    from pyFTS.common import Transformations

    def plot_compared_intervals_ahead(original, models, colors, distributions, time_from, time_to, intervals = True,
                                   save=False, file=None, tam=[20, 5], resolution=None,
                                   cmap='Blues', linewidth=1.5):
        """
        Plot the forecasts of several one step ahead models, by point or by interval

        :param original: Original time series data (list)
        :param models: List of models to compare
        :param colors: List of models colors
        :param distributions: True to plot a distribution
        :param time_from: index of data poit to start the ahead forecasting
        :param time_to: number of steps ahead to forecast
        :param interpol: Fill space between distribution plots
        :param save: Save the picture on file
        :param file: Filename to save the picture
        :param tam: Size of the picture
        :param resolution:
        :param cmap: Color map to be used on distribution plot
        :param option: Distribution type to be passed for models
        :return:
        """
        fig = plt.figure(figsize=tam)
        ax = fig.add_subplot(111)

        cm = plt.get_cmap(cmap)
        cNorm = pltcolors.Normalize(vmin=0, vmax=1)
        scalarMap = cmx.ScalarMappable(norm=cNorm, cmap=cm)

        if resolution is None: resolution = (max(original) - min(original)) / 100

        mi = []
        ma = []

        for count, fts in enumerate(models, start=0):
            if fts.has_probability_forecasting and distributions[count]:
                density = fts.forecast_ahead_distribution(original[time_from - fts.order:time_from], time_to,
                                                          resolution=resolution)

                #plot_density_scatter(ax, cmap, density, fig, resolution, time_from, time_to)
                plot_density_rectange(ax, cm, density, fig, resolution, time_from, time_to)

            if fts.has_interval_forecasting and intervals:
                forecasts = fts.forecast_ahead_interval(original[time_from - fts.order:time_from], time_to)
                lower = [kk[0] for kk in forecasts]
                upper = [kk[1] for kk in forecasts]
                mi.append(min(lower))
                ma.append(max(upper))
                for k in np.arange(0, time_from - fts.order):
                    lower.insert(0, None)
                    upper.insert(0, None)
                ax.plot(lower, color=colors[count], label=fts.shortname, linewidth=linewidth)
                ax.plot(upper, color=colors[count], linewidth=linewidth*1.5)

        ax.plot(original, color='black', label="Original", linewidth=linewidth*1.5)
        handles0, labels0 = ax.get_legend_handles_labels()
        if True in distributions:
            lgd = ax.legend(handles0, labels0, loc=2)
        else:
            lgd = ax.legend(handles0, labels0, loc=2, bbox_to_anchor=(1, 1))
        _mi = min(mi)
        if _mi < 0:
            _mi *= 1.1
        else:
            _mi *= 0.9
        _ma = max(ma)
        if _ma < 0:
            _ma *= 0.9
        else:
            _ma *= 1.1

        ax.set_ylim([_mi, _ma])
        ax.set_ylabel('F(T)')
        ax.set_xlabel('T')
        ax.set_xlim([0, len(original)])

        show_and_save_image(fig, file, save, lgd=lgd)

    def plot_density_rectange(ax, cmap, density, fig, resolution, time_from, time_to):
        """
        Auxiliar function to plot_compared_intervals_ahead
        """
        from matplotlib.patches import Rectangle
        from matplotlib.collections import PatchCollection
        patches = []
        colors = []
        for x in density.index:
            for y in density.columns:
                s = Rectangle((time_from + x, y), 1, resolution, fill=True, lw = 0)
                patches.append(s)
                colors.append(density[y][x]*5)
        pc = PatchCollection(patches=patches, match_original=True)
        pc.set_clim([0, 1])
        pc.set_cmap(cmap)
        pc.set_array(np.array(colors))
        ax.add_collection(pc)
        cb = fig.colorbar(pc, ax=ax)
        cb.set_label('Density')

    def plot_probability_distributions(pmfs, lcolors, tam=[15, 7]):
        fig = plt.figure(figsize=tam)
        ax = fig.add_subplot(111)

        for k,m in enumerate(pmfs,start=0):
            m.plot(ax, color=lcolors[k])

        handles0, labels0 = ax.get_legend_handles_labels()
        ax.legend(handles0, labels0)

    def plot_distribution(ax, cmap, probabilitydist, fig, time_from, reference_data=None):
        """
        Plot forecasted ProbabilityDistribution objects on a matplotlib axis

        :param ax: matplotlib axis
        :param cmap: matplotlib colormap name
        :param probabilitydist: list of ProbabilityDistribution objects
        :param fig: matplotlib figure
        :param time_from: starting time (on x axis) to begin the plots
        :param reference_data:
        :return:
        """
        from matplotlib.patches import Rectangle
        from matplotlib.collections import PatchCollection
        patches = []
        colors = []
        for ct, dt in enumerate(probabilitydist):
            disp = 0.0
            if reference_data is not None:
                disp = reference_data[time_from+ct]

            for y in dt.bins:
                s = Rectangle((time_from+ct, y+disp), 1, dt.resolution, fill=True, lw = 0)
                patches.append(s)
                colors.append(dt.density(y))
        scale = Transformations.Scale()
        colors = scale.apply(colors)
        pc = PatchCollection(patches=patches, match_original=True)
        pc.set_clim([0, 1])
        pc.set_cmap(cmap)
        pc.set_array(np.array(colors))
        ax.add_collection(pc)
        cb = fig.colorbar(pc, ax=ax)
        cb.set_label('Density')

    def plot_distribution2(probabilitydist, data, **kwargs):
        """
        Plot distributions in y-axis over the time (x-axis)

        :param probabilitydist: the forecasted probability distributions to plot
        :param data: the original test sample
        :keyword start_at: the time index (inside of data) to start to plot the probability distributions
        :keyword ax: a matplotlib axis. If no value was provided a new axis is created.
        :keyword cmap: a matplotlib colormap name, the default value is 'Blues'
        :keyword quantiles: the list of quantiles intervals to plot, e. g. [.05, .25, .75, .95]
        :keyword median: a boolean value indicating if the median value will be plot.
        """
        import matplotlib.colorbar as cbar
        import matplotlib.cm as cm

        ax = kwargs.get('ax',None)
        if ax is None:
            fig, ax = plt.subplots(nrows=1, ncols=1, figsize=[15, 5])

        l = len(probabilitydist)

        cmap = kwargs.get('cmap','Blues')
        cmap = plt.get_cmap(cmap)

        start_at = kwargs.get('start_at',0)

        x = [k + start_at for k in range(l + 1)]

        qt = kwargs.get('quantiles',None)

        if qt is None:
            qt = [round(k, 2) for k in np.arange(.05, 1., .05)]
            qt.insert(0, .01)
            qt.append(.99)

        lq = len(qt)

        normal = plt.Normalize(min(qt), max(qt))
        scalarMap = cm.ScalarMappable(norm=normal, cmap=cmap)

        for ct in np.arange(1, int(lq / 2) + 1):
            try:
                y = [[data[start_at], data[start_at]]]
                for pd in probabilitydist:
                    qts = pd.quantile([qt[ct - 1], qt[-ct]])
                    y.append(qts)

                ax.fill_between(x, [k[0] for k in y], [k[1] for k in y],
                                facecolor=scalarMap.to_rgba(ct / lq))
            except:
                pass

        if kwargs.get('median',True):
            y = [data[start_at]]
            for pd in probabilitydist:
                qts = pd.quantile([.5])
                y.append(qts[0])

            ax.plot(x, y, color='red', label='Median')

        cax, _ = cbar.make_axes(ax)
        cb = cbar.ColorbarBase(cax, cmap=cmap, norm=normal)
        cb.set_label('Density')

    def plot_distribution_tiled(distributions,data=None,rows=5,cols=5,index=None,axis=None,size=[10,20]):
        """
        Plot one distribution individually in each axis, with probability in y-axis and UoD on x-axis

        :param distributions:
        :param data:
        :param rows:
        :param cols:
        :param index:
        :param axis:
        :param size:
        :return:
        """

        if axis is None:
            fig, axis = plt.subplots(nrows=rows, ncols=cols, figsize=size)

        for ct in range(rows*cols):
            col = ct % cols
            row = ct // cols
            if index is None:
                ix = ct
            else:
                ix =index[ct]
            forecast = distributions[ix]
            forecast.plot(axis=axis[row][col])
            if data is not None:
                axis[row][col].axvline(data[ix])
            axis[row][col].set_title('t+{}'.format(ix))
            axis[row][col].set_xlabel(None)

        plt.tight_layout()

    def plot_interval(axis, intervals, order, label, color='red', typeonlegend=False, ls='-', linewidth=1):
        """
        Plot forecasted intervals on matplotlib

        :param axis: matplotlib axis
        :param intervals: list of forecasted intervals
        :param order: order of the model that create the forecasts
        :param label: figure label
        :param color: matplotlib color name
        :param typeonlegend:
        :param ls: matplotlib line style
        :param linewidth: matplotlib width
        :return:
        """
        lower = [kk[0] for kk in intervals]
        upper = [kk[1] for kk in intervals]
        mi = min(lower) * 0.95
        ma = max(upper) * 1.05
        for k in np.arange(0, order+1):
            lower.insert(0, None)
            upper.insert(0, None)
        if typeonlegend: label += " (Interval)"
        axis.plot(lower, color=color, label=label, ls=ls,linewidth=linewidth)
        axis.plot(upper, color=color, ls=ls,linewidth=linewidth)
        return [mi, ma]

    def plot_interval2(intervals, data, **kwargs):
        """
        Plot forecasted intervals on matplotlib

        :param intervals: list of forecasted intervals
        :param data: the original test sample
        :keyword start_at: the time index (inside of data) to start to plot the intervals
        :keyword label: figure label
        :keyword color: matplotlib color name
        :keyword typeonlegend:
        :keyword ls: matplotlib line style
        :keyword linewidth: matplotlib width
        """

        l = len(intervals)

        nintervals = intervals

        start_at = kwargs.get('start_at', 1)

        ax = kwargs.get('ax', None)
        if ax is None:
            fig, ax = plt.subplots(nrows=1, ncols=1, figsize=[15, 5])

        for k in np.arange(0, start_at):
            nintervals.insert(0, [None,None])

        nintervals.insert(start_at, [data[start_at], data[start_at]])

        lower = [kk[0] for kk in nintervals]
        upper = [kk[1] for kk in nintervals]

        typeonlegend = kwargs.get('typeonlegend', False)
        color = kwargs.get('color', 'red')
        label = kwargs.get('label','')
        linewidth = kwargs.get('linewidth', 1)

        ls = kwargs.get('ls','-')

        if typeonlegend: label += " (Interval)"
        ax.plot(lower, color=color, label=label, ls=ls,linewidth=linewidth)
        ax.plot(upper, color=color, ls=ls,linewidth=linewidth)

    def plot_rules(model, size=[5, 5], axis=None, rules_by_axis=None, columns=1):
        """
        Plot the FLRG rules of a FTS model on a matplotlib axis

        :param model: FTS model
        :param size: figure size
        :param axis: matplotlib axis
        :param rules_by_axis: number of rules plotted by column
        :param columns: number of columns
        :return:
        """
        if axis is None and rules_by_axis is None:
            rows = 1
        elif axis is None and rules_by_axis is not None:
            rows = (((len(model.flrgs.keys())//rules_by_axis)) // columns)+1

        fig, axis = plt.subplots(nrows=rows, ncols=columns, figsize=size)

        if rules_by_axis is None:
            draw_sets_on_axis(axis, model, size)

        _lhs = model.partitioner.ordered_sets if not model.is_high_order else model.flrgs.keys()

        for ct, key in enumerate(_lhs):

            xticks = []
            xtickslabels = []

            if rules_by_axis is None:
                ax = axis
            else:
                colcount = (ct // rules_by_axis) % columns
                rowcount = (ct // rules_by_axis) // columns

                if rows > 1 and columns > 1:
                    ax = axis[rowcount, colcount]
                elif columns > 1:
                    ax = axis[rowcount]
                else:
                    ax = axis

                if ct % rules_by_axis == 0:
                    draw_sets_on_axis(ax, model, size)

            if not model.is_high_order:
                if key in model.flrgs:
                    x = (ct % rules_by_axis) + 1
                    flrg = model.flrgs[key]
                    y = model.sets[key].centroid
                    ax.plot([x],[y],'o')
                    xticks.append(x)
                    xtickslabels.append(key)
                    for rhs in flrg.RHS:
                        dest = model.sets[rhs].centroid
                        ax.arrow(x+.1, y, 0.8, dest - y, #length_includes_head=True,
                                   head_width=0.1, head_length=0.1, shape='full', overhang=0,
                                   fc='k', ec='k')
            else:
                flrg = model.flrgs[key]
                x = (ct%rules_by_axis)*model.order + 1
                for ct2, lhs in enumerate(flrg.LHS):
                    y = model.sets[lhs].centroid
                    ax.plot([x+ct2], [y], 'o')
                    xticks.append(x+ct2)
                    xtickslabels.append(lhs)
                for ct2 in range(1, model.order):
                    fs1 = flrg.LHS[ct2-1]
                    fs2 = flrg.LHS[ct2]
                    y = model.sets[fs1].centroid
                    dest = model.sets[fs2].centroid
                    ax.plot([x+ct2-1,x+ct2], [y,dest],'-')

                y = model.sets[flrg.LHS[-1]].centroid
                for rhs in flrg.RHS:
                    dest = model.sets[rhs].centroid
                    ax.arrow(x + model.order -1 + .1, y, 0.8, dest - y,  # length_includes_head=True,
                               head_width=0.1, head_length=0.1, shape='full', overhang=0,
                               fc='k', ec='k')

            ax.set_xticks(xticks)
            ax.set_xticklabels(xtickslabels)
            ax.set_xlim([0,rules_by_axis*model.order+1])

        plt.tight_layout()
        plt.show()

    def draw_sets_on_axis(axis, model, size):
        if axis is None:
            fig, axis = plt.subplots(nrows=1, ncols=1, figsize=size)
        for ct, key in enumerate(model.partitioner.ordered_sets):
            fs = model.sets[key]
            axis.plot([0, 1, 0], fs.parameters, label=fs.name)
            axis.axhline(fs.centroid, c="lightgray", alpha=0.5)
        axis.set_xlim([0, len(model.partitioner.ordered_sets)])
        axis.set_xticks(range(0, len(model.partitioner.ordered_sets)))
        tmp = ['']
        tmp.extend(model.partitioner.ordered_sets)
        axis.set_xticklabels(tmp)
        axis.set_ylim([model.partitioner.min, model.partitioner.max])
        axis.set_yticks([model.sets[k].centroid for k in model.partitioner.ordered_sets])
        axis.set_yticklabels([str(round(model.sets[k].centroid, 1)) + " - " + k
                              for k in model.partitioner.ordered_sets])

    current_milli_time = lambda: int(round(time.time() * 1000))

    def uniquefilename(name):
        if '.' in name:
            tmp = name.split('.')
            return  tmp[0] + str(current_milli_time()) + '.' + tmp[1]
        else:
            return name + str(current_milli_time())

    def show_and_save_image(fig, file, flag, lgd=None):
        """
        Show and image and save on file

        :param fig: Matplotlib Figure object
        :param file: filename to save the picture
        :param flag: if True the image will be saved
        :param lgd: legend
        """
        plt.show()
        if flag:
            if lgd is not None:
                fig.savefig(file, additional_artists=lgd,bbox_inches='tight')  #bbox_extra_artists=(lgd,), )
            else:
                fig.savefig(file)
            plt.close(fig)

    def enumerate2(xs, start=0, step=1):
        for x in xs:
            yield (start, x)
            start += step

    def sliding_window(data, windowsize, train=0.8, inc=0.1, **kwargs):
        """
        Sliding window method of cross validation for time series

        :param data: the entire dataset
        :param windowsize: window size
        :param train: percentual of the window size will be used for training the models
        :param inc: percentual of data used for slide the window
        :return: window count, training set, test set
        """

        multivariate = True if isinstance(data, pd.DataFrame) else False

        l = len(data) if not multivariate else len(data.index)
        ttrain = int(round(windowsize * train, 0))
        ic = int(round(windowsize * inc, 0))

        progressbar = kwargs.get('progress', None)

        rng = np.arange(0,l-windowsize+ic,ic)

        if progressbar:
            from tqdm import tqdm
            rng = tqdm(rng)

        for count in rng:
            if count + windowsize > l:
                _end = l
            else:
                _end = count + windowsize
            if multivariate:
                yield (count, data.iloc[count: count + ttrain], data.iloc[count + ttrain: _end])
            else:
                yield (count,  data[count : count + ttrain], data[count + ttrain : _end]  )

    def persist_obj(obj, file):
        """
        Persist an object on filesystem. This function depends on Dill package

        :param obj: object on memory
        :param file: file name to store the object
        """
        import dill
        try:
            with open(file, 'wb') as _file:
                dill.dump(obj, _file)
        except Exception as ex:
            print("File {} could not be saved due exception {}".format(file, ex))

    def load_obj(file):
        """
        Load to memory an object stored filesystem. This function depends on Dill package

        :param file: file name where the object is stored
        :return: object
        """
        import dill
        with open(file, 'rb') as _file:
            obj = dill.load(_file)
        return obj

    def persist_env(file):
        """
        Persist an entire environment on file. This function depends on Dill package

        :param file: file name to store the environment
        """
        import dill
        dill.dump_session(file)

    def load_env(file):
        import dill
        dill.load_session(file)