from pyFTS.common import Membership, FuzzySet as FS
from pyFTS.common.Composite import FuzzySet as Composite
from pyFTS.partitioners import partitioner, Grid
from pyFTS.models.seasonal.common import DateTime, FuzzySet, strip_datepart
import numpy as np
import matplotlib.pylab as plt
from scipy.spatial import KDTree


class TimeGridPartitioner(partitioner.Partitioner):
    """Even Length DateTime Grid Partitioner"""

    def __init__(self, **kwargs):
        """
        Even Length Grid Partitioner
        :param seasonality: Time granularity, from pyFTS.models.seasonal.common.DateTime
        :param data: Training data of which the universe of discourse will be extracted. The universe of discourse is the open interval between the minimum and maximum values of the training data.
        :param npart: The number of universe of discourse partitions, i.e., the number of fuzzy sets that will be created
        :param func: Fuzzy membership function (pyFTS.common.Membership)
        """
        super(TimeGridPartitioner, self).__init__(name="TimeGrid", preprocess=False, **kwargs)

        self.season = kwargs.get('seasonality', DateTime.day_of_year)
        '''Seasonality, a pyFTS.models.seasonal.common.DateTime object'''
        self.mask = kwargs.get('mask', '%Y-%m-%d %H:%M:%S')
        '''A string with datetime formating mask'''

        data = kwargs.get('data', None)
        if self.season == DateTime.year:
            ndata = [strip_datepart(k, self.season) for k in data]
            self.min = min(ndata)
            self.max = max(ndata)
        else:
            tmp = (self.season.value / self.partitions) / 2
            self.min = tmp
            self.max = self.season.value + tmp

        self.type = kwargs.get('type','seasonal')

        self.sets = self.build(None)

        if self.ordered_sets is None and self.setnames is not None:
            self.ordered_sets = self.setnames
        else:
            self.ordered_sets = FS.set_ordered(self.sets)

    def extractor(self,x):
        if self.type == 'seasonal':
            return strip_datepart(x, self.season, self.mask)
        else:
            return x

    def build(self, data):
        sets = {}

        kwargs = {'variable': self.variable, 'type': self.type }

        if self.season == DateTime.year:
            dlen = (self.max - self.min)
            partlen = dlen / self.partitions
        elif self.season == DateTime.day_of_week:
            self.min, self.max, partlen, pl2 = 0, 7, 1, 1
        elif self.season == DateTime.minute:
            self.min, self.max, partlen, pl2 = 0, 60, 1, 1
        elif self.season == DateTime.hour:
            self.min, self.max, partlen, pl2 = 0, 24, 1, 1
        elif self.season == DateTime.month:
            self.min, self.max, partlen, pl2 = 1, 13, 1, 1
        elif self.season  in (DateTime.half, DateTime.third, DateTime.quarter, DateTime.sixth):
            self.min, self.max, partlen, pl2 = 1, self.season.value+1, 1, 1
        else:
            partlen = self.season.value / self.partitions
            pl2 = partlen / 2

        for count, midpoint in enumerate(np.arange(self.min, self.max, partlen)):
            set_name = self.get_name(count)
            if self.membership_function == Membership.trimf:
                if midpoint == self.min or count == 0:
                    tmp = Composite(set_name, superset=True, **kwargs)
                    tmp.append_set(FuzzySet(self.season, set_name, Membership.trimf,
                                            [self.season.value - pl2, self.season.value,
                                             self.season.value + pl2], self.season.value, alpha=1,
                                            **kwargs))
                    tmp.append_set(FuzzySet(self.season, set_name, Membership.trimf,
                                            [midpoint - partlen, midpoint, midpoint + partlen], midpoint,
                                            **kwargs))
                    tmp.centroid = midpoint
                    sets[set_name] = tmp
                elif midpoint == self.max - partlen or count == self.partitions - 1:
                    tmp = Composite(set_name, superset=True, **kwargs)
                    tmp.append_set(FuzzySet(self.season, set_name, Membership.trimf,
                                            [-pl2, 0.0,
                                             pl2], 0.0, alpha=1,
                                            **kwargs))
                    tmp.append_set(FuzzySet(self.season, set_name, Membership.trimf,
                                            [midpoint - partlen, midpoint, midpoint + partlen], midpoint,
                                            **kwargs))
                    tmp.centroid = midpoint
                    sets[set_name] = tmp
                else:
                    sets[set_name] = FuzzySet(self.season, set_name, Membership.trimf,
                                         [midpoint - partlen, midpoint, midpoint + partlen], midpoint,
                                         **kwargs)
            elif self.membership_function == Membership.gaussmf:
                sets[set_name] = FuzzySet(self.season, set_name, Membership.gaussmf, [midpoint, partlen / 3], midpoint,
                                     **kwargs)
            elif self.membership_function == Membership.trapmf:
                q = partlen / 4
                if midpoint == self.min:
                    tmp = Composite(set_name, superset=True)
                    tmp.append_set(FuzzySet(self.season, set_name, Membership.trimf,
                                            [self.season.value - pl2, self.season.value,
                                             self.season.value + 0.0000001], 0,
                                            **kwargs))
                    tmp.append_set(FuzzySet(self.season, set_name, Membership.trapmf,
                                            [midpoint - partlen, midpoint - q, midpoint + q, midpoint + partlen], midpoint,
                                            **kwargs))
                    tmp.centroid = midpoint
                    sets[set_name] = tmp
                else:
                    sets[set_name] = FuzzySet(self.season, set_name, Membership.trapmf,
                                         [midpoint - partlen, midpoint - q, midpoint + q, midpoint + partlen], midpoint,
                                         **kwargs)
            count += 1

        self.min = 0

        return sets

    def build_index(self):
        points = []

        fset = self.sets[self.ordered_sets[0]]
        points.append([fset.sets[1].lower, fset.sets[1].centroid, fset.sets[1].upper])

        for ct, key in enumerate(self.ordered_sets[1:-1]):
            fset = self.sets[key]
            points.append([fset.lower, fset.centroid, fset.upper])

        fset = self.sets[self.ordered_sets[-1]]
        points.append([fset.sets[1].lower, fset.sets[1].centroid, fset.sets[1].upper])

        import sys
        sys.setrecursionlimit(100000)

        self.kdtree = KDTree(points)

        sys.setrecursionlimit(1000)

    def search(self, data, **kwargs):
        '''
        Perform a search for the nearest fuzzy sets of the point 'data'. This function were designed to work with several
        overlapped fuzzy sets.

        :param data: the value to search for the nearest fuzzy sets
        :param type: the return type: 'index' for the fuzzy set indexes or 'name' for fuzzy set names.
        :param results: the number of nearest fuzzy sets to return
        :return: a list with the nearest fuzzy sets
        '''

        type = kwargs.get('type','index')
        results = kwargs.get('results',3)

        if self.kdtree is None:
            self.build_index()

        _, ix = self.kdtree.query([data, data, data], results)

        ix = ix.tolist()

        if 0 in ix:
            ix.insert(0, self.partitions-1)
        elif self.partitions-1 in ix:
            ix.insert(0, 0)

        if type == 'name':
            return [self.ordered_sets[k] for k in sorted(ix)]
        else:
            return sorted(ix)


    def plot(self, ax):
        """
        Plot the
        :param ax:
        :return:
        """
        ax.set_title(self.name)
        ax.set_ylim([0, 1])
        ax.set_xlim([0, self.season.value])
        ticks = []
        x = []
        for key in self.sets.keys():
            s = self.sets[key]
            if s.type == 'composite':
                for ss in s.sets:
                    self.plot_set(ax, ss)
            else:
                self.plot_set(ax, s)
            ticks.append(str(round(s.centroid, 0)) + '\n' + s.name)
            x.append(s.centroid)
        ax.xaxis.set_ticklabels(ticks)
        ax.xaxis.set_ticks(x)
