# coding: utf-8
# coding: utf-8
import numpy as np

from numpy.core.numeric import array

from sklearn.metrics import pairwise

from collections import OrderedDict

from scipy.linalg import eig

from scipy import interpolate as ip
from pyspark.ml.feature import PCA as pcaml
from pyspark.ml.linalg import Vectors as mlvec

from pyspark.mllib.linalg.distributed import IndexedRow, RowMatrix, IndexedRowMatrix, CoordinateMatrix, MatrixEntry
from pyspark.mllib.linalg import DenseVector, Vectors
from pyspark.mllib.feature import PCA as PCAmllib
from pyspark.sql.types import *
from pyspark.mllib.clustering import GaussianMixture
from pyspark.ml.clustering import GaussianMixture as gmmml
from pyspark.mllib.clustering import KMeans
from pyspark.ml.clustering import KMeans as kmml

class LabelPropagationDistributed():
    """
    A Label Propagation semi-supervised clustering algorithm based
    on the paper "Semi-supervised Learning in Gigantic Image Collections"
    by Fergus, Weiss and Torralba with modifications to fit Spark.
    The algorithm begins with discretizing the datapoints into 1-D histograms.
    For each independent dimension, it approximates the density using the histogram
    and solves numerically for eigenfunctions and eigenvalues.
    Then it uses the k eigenfunctions for k smallest eigenvalues to do a 1-D
    interpolation of every data point.
    These interpolated data points are the approximated eigenvectors of the Normalized
    Laplacian of the original data points which are further used to solve a k*k system
    of linear equation which yields alpha.
    The alpha is then used to get approximate functions which are then clustered using
    Gaussian Mixture Model.
    For any new/ unseen data point, the point can be interpolated and using the alpha,
    the approximate function for that point can be obtained whose label can be easily predicted
    using the GMM model learned before making it inductive learning.

    Based on
    U{https://cs.nyu.edu/~fergus/papers/fwt_ssl.pdf}
    Fergus, Weiss and Torralba, Semi-supervised Learning in Gigantic
    Image Collections, Proceedings of the 22nd International Conference
    on Neural Information Processing Systems, p.522-530, December 07-10, 2009,
    Vancouver, British Columbia, Canada

    >>> from LabelPropagationDistributed import LabelPropagationDistributed as LPD
    >>> dataX = array([ 5.76961775,  -6.55673209, 11.30752027,  -1.56316985,
                8.76722337,  -1.54995049, 10.23511359,  -6.20912033,
                3.49161828,  -3.02917744]).reshape(5,2)
    >>> dataY = array([ 1,  -1,  0, -1, -1])
    >>> test = array([2.1159109 ,   6.03520684,   1.04347698,  -4.44740207,
               -8.33902404,   4.20918959,   1.38447488,  -1.50363493]).reshape(4,2)
    >>> lpd = LPD(sc=sc, sqlContext = sqlContext, numBins = 5)
    >>> lpd.fit(sc.parallelize(dataX),sc.parallelize(dataY))
    >>> plabels_ = lpd.predict(sc.parallelize(test))

    """

    global transformer
    global selectValues
    global bc_EdgeMeans
    global bc_newg
    global kb

    def __init__(self, sc = None, sqlContext = None,  k = -1,numBins = -1 ,lagrangian = 1000, gamma = 0.01):

        self.k = k
        self.numBins = numBins
        self.lagrangian = lagrangian
        self.sqlContext = sqlContext
        self.sc = sc
        self.gamma = gamma

    @classmethod
    def print_test(cls):
        print "working"

    def makeDF(self,rotatedData, dimensions):
        """
        Convert data from RDD to a dataframe with every feature
        as a different column.

        :param rotatedData:
            Data points with as independent dimensions as possible

        :param dimensions
            Total number of dimensions of the data

        """
        X_ = rotatedData.map(lambda (ind, vec): (ind,np.array(vec)))
        dataAsDict = X_.map(lambda (ind,x): tuple([ind]+[float(f) for f in x]))
        fields = [StructField("idd",IntegerType(), True)]
        fields += [StructField(str(i),FloatType(), True) for i in range(1,dimensions+1)]
        """
        schemaString = "index "
        for i in range(1,dimensions+1):
            schemaString += str(i) + " "
        schemaString = schemaString.strip()
        fields = [StructField(field_name,FloatType(), True) for field_name in schemaString.split()]
        """
        schema = StructType(fields)
        return self.sqlContext.createDataFrame(dataAsDict, schema)


    def getdataboundaries(self,dictData,k):
        """
        For interpolating the data points, get min and max values for
        the data set.

        :param: dictData:
            Dataframe of dataset

        :param: k:
            The number of eigenvectors to be selected. Default is equal
            to the number of clusters.

        """
        dataBounds = [None for _ in range(k)]
        for i in range(0,k):
            ind = self.jj[i]
            s = str(ind + 1)
            tmp = dictData.select(s).rdd.map(lambda r: r[0]).collect()
            Lower, Upper = np.percentile(tmp, [2.5,97.5])
            dataBounds[i] = (Lower,Upper)
        return dataBounds
        
    def transformer(vec, bounds, bc_EdgeMeans, bc_newg, kb):
        """
        Interpolate and return data points
        Runs on workers

        :param: vec:
            k dimenional data

        :param: bounds:
            Min and Max values for data to be interpolated.

        :param: bc_EdgeMeans:
            Histogram bin edges to build the 1-D interpolator

        :param: bc_newg:
            k smallest eigenvectors to build the 1-D interpolator
        """
        vec1 = np.array(vec)
        tmpArr = np.zeros(vec1.shape)
        edge_means = bc_EdgeMeans.value
        for i in range(0, kb.value):
            inter2 = ip.interp1d(edge_means[:,i], bc_newg.value[:,i], fill_value = 'extrapolate', kind = 'linear')
            (lower, upper) = bounds.value[i]
            val = vec1[i]
            clippedValue = np.clip(val,lower,upper)
            tmpArr[i] = inter2(val)
        return DenseVector(tmpArr)


    @staticmethod
    def indRowMatMaker(irm):
        """
        Converts RDD to indexed Row Matrix

        :param: irm:
            RDD to be converted to IndexedRowMatrix
        """

        return IndexedRowMatrix(irm.zipWithIndex().map(lambda x:IndexedRow(x[1],x[0])))

    def solver(self,functions, alpha):
        """
        f = U * alpha
        Get and return approximate eigenvectors from eigenfunctions using alpha.

        :param: vectors:
            Approximate eigenfunctions
        """

        U = functions
        alpha_bc = self.sc.broadcast(alpha.reshape(self.k,1))
        f = U.rows.map(lambda row: (row.index, np.dot(row.vector.toArray(), alpha_bc.value)))
        return f

    def selectValues(ddict,kb):
        """
        Select and return k eigenvectors from the Dataframe.

        :param: ddict:
            Dataframe Row as a dictionary.

        :param: kb:
            Desired number of eigenvectors to be selected.
            Default equals to number of clusters

        """

        desired_keys = sorted([int(k) for k in ddict.keys()])[0:kb.value]
        newddict = { i: ddict[str(i)] for i in desired_keys }
        return newddict

    def relabel(self,ilabels, km):
        """
        Label the data points in an ordered way depending on the ascending
        order of the gaussian means.

        :param: labels:
            GMM predicted labels

        """

        #gaussians = np.zeros((self.numClasses,self.k))
        #i=0
        #for ls in self.gmm.gaussians:
        #    gaussians[i] = (ls.mu)
        #    i = i +1
        #gaussians = np.array(gmodel.gaussiansDF.select('mean').rdd.map(lambda r: r[0].toArray()).collect()).flatten()
        means = np.array(km.clusterCenters())
        distmeans = self.sc.broadcast(np.argsort(means.flatten()))
        ilabels = ilabels.select(["idd","prediction"]).rdd.map(lambda row: (row.idd, np.where(distmeans.value == row.prediction)[0][0]))
        return ilabels

    def _get_kernel(self, X, y = None,ker=None):
        """
        return pairwise affinity matrix based on the kernel specified.

        :param: X:
            Data array

        :param: y:
            Label array

        :param: ker:
            kernel specified. Currently supporting Euclidean and rbf.

        """

        if ker == "rbf":
            if y is None:
                return pairwise.rbf_kernel(X, X, gamma = self.gamma)
            else:
                return pairwise.rbf_kernel(X, y, gamma = self.gamma)
        elif ker == "linear":
            if y is None:
                return pairwise.euclidean_distances(X, X)
            else:
                return pairwise.euclidean_distances(X, y)
        else:
            raise ValueError("is not a valid kernel. Only rbf and euclidean"
                             " are supported at this time")

    def getParams(self, X, y ):
        """
        Get all necessary parameters like total number of dimensions,
        desired number of dimensions to work on k,
        number of classes based on the labeled data available,
        number of histogram bins and total data points n.

        :param: X
            RDD of index, Data points

        :param: y
            RDD of index, labels

        """

        n = X.cache().count()
        self.dimensions= len(X.take(1)[0][1])
        classes = sorted(y.filter(lambda (i,label): label != -1).map(lambda (i,label): label).distinct().collect())
        if classes[0] == -1:
            classes = np.delete(classes, 0) # remove the -1 from this list

        if self.k == -1:
            self.k = np.size(classes)
        if self.numBins == -1:
                self.numBins = self.k + 1
        if self.k > self.dimensions:
            self.k = self.dimensions
        return (n,classes)

    def rotate(self, X):
        """
        Rotate the data to get independent dimensions.

        :param: X
            RDD of data points to be rotated

        """
        XforPCA = X.map(lambda (ind,vec): (int(ind),mlvec.dense(vec))).toDF(schema = ('idd','features'))
        pca = pcaml(k=self.dimensions,inputCol="features", outputCol="pcaFeatures")
        self.PCA = pca.fit(XforPCA)
        rotatedData = self.PCA.transform(XforPCA)

        rotatedData=rotatedData.select(["idd", "pcaFeatures"]).rdd.map(lambda r: (r.idd,r.pcaFeatures.toArray()))
        return rotatedData

    def approximateDensities(self, s, dictData):
        """
        Discretize data into bins. Returns the new histograms and corresponding bin edges.

        :param: s:
            index to select dimension of the data

        :param: dictData:
            Dataframe of original data

        """
        dimRdd = dictData.select(s)
        binEdges,histograms = dimRdd.rdd.map(lambda x: x[0]).histogram(self.numBins)
        histograms = array(histograms)
        binEdges = np.array(binEdges)
        #db = array(np.diff(binEdges),float)
        #histograms = histograms/db/histograms.sum()
        histograms = histograms + 0.1
        histograms /= histograms.sum()
        return (histograms, binEdges)

    def generalizedEigenSolver(self, binmeans, histograms):
        """
        A generalized Eigen Solver that gives approximate eigenfunctions and eigenvalues.
        Based on Eqn. 2 in the paper.
        NOTE: The first eigenfunction is always going to be a trivial function with maximal smoothness.
              Hence selecting eigenvalue at index 1 instead of 0.

        :params: histograms:
            Discretized data whose eigenfunctions and eigenvalues are to be evaluated.

        """

        Wdis = self._get_kernel(binmeans.reshape(binmeans.shape[0],1),y=None,ker="rbf")
        P = np.diag(histograms).astype('float32')
        PW = np.dot(P,Wdis)
        PWP = np.dot(PW,P)