__author__ = 'wittawat'

from abc import ABCMeta, abstractmethod
import math
import matplotlib.pyplot as plt
import numpy as np
import fsic.util as util
import matplotlib.pyplot as plt
import scipy.stats as stats

class PairedData(object):
    """Class representing paired data for independence testing

    properties:
    X, Y: numpy array. X and Y are paired of the same sample size. The
        dimensions are not necessarily the same.
    """

    def __init__(self, X, Y, label=None):
        """
        :param X: n x d numpy array for dataset X
        :param Y: n x d' numpy array for dataset Y
        """
        self.X = X
        self.Y = Y
        # short description to be used as a plot label
        self.label = label

        nx, dx = X.shape
        ny, dy = Y.shape
        if nx != ny:
            raise ValueError('Data size of the paired sample must be the same.')

        if not np.all(np.isfinite(X)):
            print 'X:'
            print util.fullprint(X)
            raise ValueError('Not all elements in X are finite.')

        if not np.all(np.isfinite(Y)):
            print 'Y:'
            print util.fullprint(Y)
            raise ValueError('Not all elements in Y are finite.')

    def __str__(self):
        mean_x = np.mean(self.X, 0)
        std_x = np.std(self.X, 0) 
        mean_y = np.mean(self.Y, 0)
        std_y = np.std(self.Y, 0) 
        prec = 4
        desc = ''
        desc += 'E[x] = %s \n'%(np.array_str(mean_x, precision=prec ) )
        desc += 'E[y] = %s \n'%(np.array_str(mean_y, precision=prec ) )
        desc += 'Std[x] = %s \n' %(np.array_str(std_x, precision=prec))
        desc += 'Std[y] = %s \n' %(np.array_str(std_y, precision=prec))
        return desc

    def dx(self):
        """Return the dimension of X."""
        dx = self.X.shape[1]
        return dx

    def dy(self):
        """Return the dimension of Y."""
        dy = self.Y.shape[1]
        return dy

    def sample_size(self):
        return self.X.shape[0]

    def xy(self):
        """Return (X, Y) as a tuple"""
        return (self.X, self.Y)

    def split_tr_te(self, tr_proportion=0.5, seed=820):
        """Split the dataset into training and test sets. Assume n is the same 
        for both X, Y. 
        
        Return (PairedData for tr, PairedData for te)"""
        X = self.X
        Y = self.Y
        nx, dx = X.shape
        ny, dy = Y.shape
        if nx != ny:
            raise ValueError('Require nx = ny')
        Itr, Ite = util.tr_te_indices(nx, tr_proportion, seed)
        label = '' if self.label is None else self.label
        tr_data = PairedData(X[Itr, :], Y[Itr, :], 'tr_' + label)
        te_data = PairedData(X[Ite, :], Y[Ite, :], 'te_' + label)
        return (tr_data, te_data)

    def subsample(self, n, seed=87):
        """Subsample without replacement. Return a new PairedData """
        if n > self.X.shape[0] or n > self.Y.shape[0]:
            raise ValueError('n should not be larger than sizes of X, Y.')
        ind_x = util.subsample_ind( self.X.shape[0], n, seed )
        ind_y = util.subsample_ind( self.Y.shape[0], n, seed )
        return PairedData(self.X[ind_x, :], self.Y[ind_y, :], self.label)

    def clone(self):
        """
        Return a new PairedData object with a separate copy of each internal 
        variable, and with the same content.
        """
        nX = np.copy(self.X)
        nY = np.copy(self.Y)
        nlabel = self.label
        return PairedData(nX, nY, nlabel)

    def __add__(self, pdata2):
        """
        Merge the current PairedData with another one.
        Create a new PairedData and create a new copy for all internal variables.
        label is set to None.
        """
        copy = self.clone()
        copy2 = pdata2.clone()
        nX = np.vstack((copy.X, copy2.X))
        nY = np.vstack((copy.Y, copy2.Y))
        return PairedData(nX, nY)



### end PairedData class        

class PairedSource(object):
    """A data source where it is possible to resample. Subclasses may prefix 
    class names with PS. 

    - If possible, prefix with PSInd to indicate that the 
    PairedSource contains two independent samples. 
    - Prefix with PSDep otherwise.
    - Use PS if the PairedSource can be either one depending on the provided 
    paramters."""

    __metaclass__ = ABCMeta

    @abstractmethod
    def sample(self, n, seed):
        """Return a PairedData. Returned result should be deterministic given 
        the input (n, seed)."""
        raise NotImplementedError()

    @abstractmethod
    def dx(self):
        """Return the dimension of X"""
        raise NotImplementedError()
    
    @abstractmethod
    def dy(self):
        """Return the dimension of Y"""
        raise NotImplementedError()
        

class PSResample(PairedSource):
    """
    A PairedSource which subsamples without replacement from the specified
    PairedData.
    """

    def __init__(self, pdata):
        self.pdata = pdata

    def sample(self, n, seed=900):
        pdata_sub = self.pdata.subsample(n, seed)
        return pdata_sub

    def dx(self):
        return self.pdata.dx()

    def dy(self):
        return self.pdata.dy()


class PSStraResample(PairedSource):
    """
    A PairedSource which does a stratified subsampling. without replacement
    from the specified PairedData. 
    The implementation is only approximately correctly.
    """
    def __init__(self, pdata, pivot):
        """
        pivot: a one-dimensional numpy array of the same size as pdata.sample_size()
            indicating the class of each point.
        """
        if len(pivot) != pdata.sample_size():
            raise ValueError('pivot must have the same length as the data.')
        self.pdata = pdata
        self.pivot = pivot
        uniq, counts = np.unique(pivot, return_counts=True)
        self._uniques = uniq
        self._counts = counts

    def sample(self, n, seed=900):
        pdata = self.pdata
        n_sam = pdata.sample_size()
        if n > n_sam:
            raise ValueError('Cannot subsample %d points from %d points.'%(n, n_sam))

        X, Y = pdata.xy()
        import math
        # permute X, Y. Keep pairs 
        I = util.subsample_ind(n_sam, n_sam, seed=seed+3)
        X = X[I, :]
        Y = Y[I, :]
        perm_pivot = self.pivot[I]
        list_chosenI = []
        for ui, v in enumerate(self._uniques):
            Iv  = np.nonzero(np.abs(perm_pivot - v) <= 1e-8)
            Iv = Iv[0]
            niv = self._counts[ui]
            # ceil guarantees that at least 1 instance will be chosen 
            # from each class. 
            n_class = int(math.ceil(niv/float(n_sam)*n))
            chosenI = Iv[:n_class]
            #print chosenI
            list_chosenI.append(chosenI)
        final_chosenI = np.hstack(list_chosenI)
        #print final_chosenI
        reduceI = util.subsample_ind(len(final_chosenI), min(n, len(final_chosenI)), seed+5)
        final_chosenI = final_chosenI[reduceI]
        assert len(final_chosenI) == n, 'final_chosenI has length %d which is not n=%d'%(len(final_chosenI), n)

        Xsam = X[final_chosenI, :]
        Ysam = Y[final_chosenI, :]
        new_label = None if pdata.label is None else pdata.label + '_stra'
        return PairedData(Xsam, Ysam, label=new_label)

    def dx(self):
        return self.pdata.dx()

    def dy(self):
        return self.pdata.dy()

# end PSStraResample

class PSNullShuffle(PairedSource):
    """
    Randomly permute the order of one sample so that the pairs are guaranteed
    to be broken. This is very similar to PSNullResample except it does not 
    sample. The sampling part is delegated to another PairedSource. 
    Essentially, PSNullResample = PSNullShuffle + PSResample.
    Decorator pattern.
    """
    def __init__(self, ps):
        self.ps = ps 

    def sample(self, n, seed=7):
        if n == 1:
            pdata = self.ps.sample(2, seed=seed+27)
            X, Y = pdata.xy()
            nX = X[[0], :]
            nY = Y[[1], :]
        else:
            pdata = self.ps.sample(n, seed=seed+27)
            nX, Y = pdata.xy()
            ind_shift1 = np.roll(range(n), 1)
            nY = Y[ind_shift1, :]

        new_label = 'null_shuffle'
        return PairedData(nX, nY, label=new_label)

    def dx(self):
        return self.ps.dx()

    def dy(self):
        return self.ps.dy()

# end PSNullShuffle



class PSNullResample(PairedSource):
    """
    Randomly permute the order of one sample so that the pairs are guaranteed
    to be broken. This is meant to simulate the case where [H0: X, Y are
    independent] is true.

    A PairedSource which subsamples without replacement from the permuted two
    samples.
    """

    def __init__(self, pdata):
        """
        pdata: A PairedData object
        """
        self.pdata = pdata

    def sample(self, n, seed=981):
        if n > self.pdata.sample_size():
            raise ValueError('cannot sample more points than what the original dataset has')
        X, Y =  self.pdata.xy()
        if n == 1:
            ind = util.subsample_ind(self.pdata.sample_size(), 2, seed=seed)
            nX = X[[ind[0]], :] 
            nY = Y[[ind[1]], :]
        else:
            ind = util.subsample_ind(self.pdata.sample_size(), n, seed=seed)
            nX = X[ind, :]
            ind_shift1 = np.roll(ind, 1)
            nY = Y[ind_shift1, :]
        new_label = None if self.pdata.label is None else self.pdata.label + '_shuf'
        return PairedData(nX, nY, label=new_label)

    def dx(self):
        return self.pdata.dx()

    def dy(self):
        return self.pdata.dy()

# end class PSNullResample

class PSStandardize(PairedSource):
    """
    A PairedSource that standardizes dimensions of X, Y independently so that 
    each has 0 mean and unit variance. Useful with PSResample or PSNullResample 
    when working with real data whose variables do not have the same scaling.

    Decorator pattern.
    """
    def __init__(self, ps):
        """
        ps: a PairedSource
        """
        self.ps = ps 

    def sample(self, n, seed=55):
        ps = self.ps 
        pdata = ps.sample(n, seed=seed)
        X, Y = pdata.xy()

        Zx = util.standardize(X)
        Zy = util.standardize(Y)
        assert np.all(np.isfinite(Zx))
        assert np.all(np.isfinite(Zy))
        new_label = None if pdata.label is None else pdata.label + '_std'
        return PairedData(Zx, Zy, label=new_label)
    
    def dx(self):
        return self.ps.dx()

    def dy(self):
        return self.ps.dy()

# end of class PSStandardize

class PSGaussNoiseDims(PairedSource):
    """
    A PairedSource that adds noise dimensions to X, Y drawn from the specified 
    PairedSource. The noise follows the standard normal distribution.

    Decorator pattern.
    """
    def __init__(self, ps, ndx, ndy):
        """
        ndx: number of noise dimensions for X 
        ndy: number of noise dimensions for Y 
        """
        assert ndx >= 0
        assert ndy >= 0
        self.ps = ps 
        self.ndx = ndx 
        self.ndy = ndy
    
    def sample(self, n, seed=44):
        with util.NumpySeedContext(seed=seed+100):
            NX = np.random.randn(n, self.ndx)
            NY = np.random.randn(n, self.ndy)

            pdata = self.ps.sample(n, seed=seed)
            X, Y = pdata.xy()
            Zx = np.hstack((X, NX))
            Zy = np.hstack((Y, NY))
            new_label = None if pdata.label is None else \
                pdata.label + '_ndx%d'%self.ndx + '_ndy%d'%self.ndy
            return PairedData(Zx, Zy, label=new_label)

    def dx(self):
        return self.ps.dx() + self.ndx

    def dy(self):
        return self.ps.dy() + self.ndy

# end of class PSGaussNoiseDims

class PSFunc(PairedSource):
    """
    A PairedSource that generates data (X, Y) such that Y = f(X) for a 
    specified function f (possibly stochastic), and px where X ~ px.
    """

    def __init__(self, f, px):
        """
        f: function such that Y = f(X). (n x dx)  |-> n x dy
        px: prior on X. Used to generate X. n |-> n x dx
        """
        self.f = f 
        self.px = px
        x = px(2)
        y = f(x)
        self.dx = x.shape[1]
        self.dy = y.shape[1]

    def sample(self, n, seed):
        rstate = np.random.get_state()
        np.random.seed(seed)

        px = self.px 
        X = px(n )
        f = self.f 
        Y = f(X )

        np.random.set_state(rstate)
        return PairedData(X, Y, label='psfunc')

    def dx(self):
        return self.dx

    def dy(self):
        return self.dy

class PSUnifRotateNoise(PairedSource):
    """
    X, Y are dependent in the same way as in PS2DUnifRotate. However, this 
    problem adds more extra noise dimensions.
    - The total number of dimensions is 2+2*noise_dim.
    - Only the first dimensions of X and Y are dependent. Dependency strength 
        depends on the specified angle.
    """

    def __init__(self, angle, xlb=-1, xub=1, ylb=-1, yub=1, noise_dim=0):
        """
        angle: angle in radian
        xlb: lower bound for x (a real number)
        xub: upper bound for x (a real number)
        ylb: lower bound for y (a real number)
        yub: upper bound for y (a real number)
        noise_dim: number of noise dimensions to add to each of X and Y. All
            the extra dimensions follow U(-1, 1)
        """
        ps_2d_unif = PS2DUnifRotate(angle, xlb=xlb, xub=xub, ylb=ylb, yub=yub)
        self.ps_2d_unif = ps_2d_unif
        self.noise_dim = noise_dim

    def sample(self, n, seed=883):
        sample2d = self.ps_2d_unif.sample(n, seed)
        noise_dim = self.noise_dim
        if noise_dim <= 0:
            return sample2d

        rstate = np.random.get_state()
        np.random.seed(seed+1)

        # draw n*noise_dim points from U(-1, 1)
        Xnoise = stats.uniform.rvs(loc=-1, scale=2,
                size=noise_dim*n).reshape(n, noise_dim)
        Ynoise = stats.uniform.rvs(loc=-1, scale=2,
                size=noise_dim*n).reshape(n, noise_dim)

        # concatenate the noise dims to the 2d problem 
        X2d, Y2d = sample2d.xy()
        X = np.hstack((X2d, Xnoise))
        Y = np.hstack((Y2d, Ynoise))

        np.random.set_state(rstate)

        return PairedData(X, Y, label='rot_unif_noisedim%d'%(noise_dim))

    def dx(self):
        return 1+ self.noise_dim

    def dy(self):
        return 1+ self.noise_dim


class PS2DSinFreq(PairedSource):
    """
    X, Y follow the density proportional to 1+sin(w*x)sin(w*y) where 
    w is the frequency. The higher w, the close the density is to a uniform
    distribution on [-pi, pi] x [-pi, pi].

    This dataset was used in Arthur Gretton's lecture notes.
    """
    def __init__(self, freq):
        """
        freq: a nonnegative floating-point number
        """
        self.freq = freq

    def sample(self, n, seed=81):
        ps = PSSinFreq(self.freq, d=1)
        pdata = ps.sample(n, seed=seed)
        X, Y = pdata.xy()

        return PairedData(X, Y, label='sin_freq%.2f'%self.freq)

    def _sample_sequential(self, n, seed=81):
        """
        With a loop, slow.
        """
        rstate = np.random.get_state()
        np.random.seed(seed)

        # rejection sampling
        w = self.freq
        sam = np.zeros((n, 2))
        ind = 0
        #unif_den = 1.0/(4*math.pi**2)
        #ref_bound = 2.0/unif_den
        while ind<n:
            # uniformly randomly draw x, y from U(-pi, pi)
            x = stats.uniform.rvs(loc=-math.pi, scale=2*math.pi, size=1)
            y = stats.uniform.rvs(loc=-math.pi, scale=2*math.pi, size=1)
            if stats.uniform.rvs() < (1+np.sin(w*x)*np.sin(w*y))/2.0:
                # accept 
                sam[ind, :] = [x, y]
                ind = ind + 1

        np.random.set_state(rstate)
        return PairedData(sam[:, [0]], sam[:, [1]], label='sin_freq%.2f'%self.freq)


    def dx(self):
        return 1
    
    def dy(self):
        return 1

# end class PS2DSinFreq

class PSSinFreq(PairedSource):
    """
    X, Y follow the density proportional to 
        1+\prod_{i=1}^{d} [ sin(w*x_i)sin(w*y_i) ]
    w is the frequency. The higher w, the close the density is to a uniform
    distribution on [-pi, pi] x [-pi, pi].
    - This is a generalization of PS2DSinFreq.
    """
    def __init__(self, freq, d):
        """
        freq: a nonnegative floating-point number
        """
        self.freq = freq
        self.d = d

    def sample(self, n, seed=81):
        d = self.d
        Sam = PSSinFreq.sample_d_variates(self.freq, n, 2*self.d, seed)
        X = Sam[:, :d]
        Y = Sam[:, d:]
        return PairedData(X, Y, label='sin_freq%.2f_d%d'%(self.freq, d) )

    def dx(self):
        return self.d

    def dy(self):
        return self.d

    @staticmethod 
    def sample_d_variates(w, n, D, seed=81):
        """
        Return an n x D sample matrix. 
        """
        with util.NumpySeedContext(seed=seed):
            # rejection sampling
            sam = np.zeros((n, D))
            # sample block_size*D at a time.
            block_size = 500
            from_ind = 0
            while from_ind < n:
                # uniformly randomly draw x, y from U(-pi, pi)
                X = stats.uniform.rvs(loc=-math.pi, scale=2*math.pi, size=D*block_size)
                X = np.reshape(X, (block_size, D))
                un_den = 1.0+np.prod(np.sin(w*X), 1)
                I = stats.uniform.rvs(size=block_size) < un_den/2.0

                # accept 
                accepted_count = np.sum(I)
                to_take = min(n - from_ind, accepted_count)
                end_ind = from_ind + to_take

                AX = X[I, :]
                X_take = AX[:to_take, :]
                sam[from_ind:end_ind, :] = X_take
                from_ind = end_ind
        return sam



class PS2DUnifRotate(PairedSource):
    """
    X, Y follow uniform distributions (default to U(-1, 1)). Rotate them by a
    rotation matrix of the specified angle. This can be used to simulate the
    setting of an ICA problem.
    """
    def __init__(self, angle, xlb=-1, xub=1, ylb=-1, yub=1):
        """
        angle: angle in radian
        xlb: lower bound for x (a real number)
        xub: upper bound for x (a real number)
        ylb: lower bound for y (a real number)
        yub: upper bound for y (a real number)
        """
        self.angle = angle
        self.xlb = xlb
        self.xub = xub 
        self.ylb = ylb 
        self.yub = yub

    def sample(self, n, seed=389):
        t = self.angle
        rot = np.array([[np.cos(t), -np.sin(t)], [np.sin(t), np.cos(t)]])

        ps_unif = PSIndUnif(xlb=[self.xlb], xub=[self.xub], ylb=[self.ylb], yub=[self.yub])
        pdata = ps_unif.sample(n, seed)
        X, Y = pdata.xy()
        XY = np.hstack((X, Y))
        rot_XY = XY.dot(rot.T)

        return PairedData(rot_XY[:, [0]], rot_XY[:, [1]], label='rot_unif_a%.2f'%(t))

    def dx(self):
        return 1

    def dy(self):
        return 1


class PSIndUnif(PairedSource):
    """
    Multivariate (or univariate) uniform distributions for both X, Y
    on the specified boundaries
    """

    def __init__(self, xlb, xub, ylb, yub):
        """
        xlb: a numpy array of lower bounds of x
        xub: a numpy array of upper bounds of x
        ylb: a numpy array of lower bounds of y
        yub: a numpy array of upper bounds of y
        """
        convertif = lambda a: np.array(a) if isinstance(a, list) else a 
        xlb, xub, ylb, yub = map(convertif, [xlb, xub, ylb, yub])
        if xlb.shape[0] != xub.shape[0]:
            raise ValueError('lower and upper bounds of X must be of the same length.')

        if ylb.shape[0] != yub.shape[0]:
            raise ValueError('lower and upper bounds of X must be of the same length.')

        if not np.all(xub - xlb > 0):
            raise ValueError('Require upper - lower to be positive. False for x')

        if not np.all(yub - ylb > 0):
            raise ValueError('Require upper - lower to be positive. False for y')

        self.xlb = xlb
        self.xub = xub 
        self.ylb = ylb 
        self.yub = yub

    def sample(self, n, seed):
        rstate = np.random.get_state()
        np.random.seed(seed)

        dx = self.xlb.shape[0]
        dy = self.ylb.shape[0]
        X = np.zeros((n, dx)) 
        Y = np.zeros((n, dy)) 

        pscale = self.xub - self.xlb 
        qscale = self.yub - self.ylb
        for i in range(dx):
            X[:, i] = stats.uniform.rvs(loc=self.xlb[i], scale=pscale[i], size=n)
        for i in range(dy):
            Y[:, i] = stats.uniform.rvs(loc=self.ylb[i], scale=qscale[i], size=n)

        np.random.set_state(rstate)
        return PairedData(X, Y, label='ind_unif_dx%d_dy%d'%(dx, dy))

    def dx(self):
        return self.xlb.shape[0]

    def dy(self):
        return self.ylb.shape[0]


class PSIndSameGauss(PairedSource):
    """Two same standard Gaussians for P, Q.  """
    def __init__(self, dx, dy):
        """
        dx: dimension of X
        dy: dimension of Y
        """
        self.dimx = dx 
        self.dimy = dy

    def sample(self, n, seed):
        rstate = np.random.get_state()
        np.random.seed(seed)

        X = np.random.randn(n, self.dx())
        Y = np.random.randn(n, self.dy()) 
        np.random.set_state(rstate)
        return PairedData(X, Y, label='sg_dx%d_dy%d'%(self.dx(), self.dy()) )

    def dx(self):
        return self.dimx

    def dy(self):
        return self.dimy

# end class PSIndSameGauss


class PSPairwiseSign(PairedSource):
    """
    A toy problem given in section 5.3 of 

    Large-Scale Kernel Methods for Independence Testing
    Qinyi Zhang, Sarah Filippi,  Arthur Gretton, Dino Sejdinovic

    X ~ N(0, I_d)
    Y = \sqrt(2/d) \sum_{j=1}^{d/2} sign(X_{2j-1 * X_{2j}})|Z_j| + Z_{d/2+1}

    where Z ~ N(0, I_{d/2+1})
    """

    def __init__(self, dx):
        """
        dx: the dimension of X
        """
        if dx <= 0 or dx%2 != 0:
            raise ValueError('dx has to be even')
        self.dimx = dx

    def sample(self, n, seed):
        d = self.dimx 
        with util.NumpySeedContext(seed=seed):
            Z = np.random.randn(n, d/2+1)
            X = np.random.randn(n, d)
            Y = np.zeros((n, 1))
            for j in range(d/2):
                Y = Y + np.sign(X[:, [2*j]]*X[:, [2*j+1]])*np.abs(Z[:, [j]])
            Y = np.sqrt(2.0/d)*Y + Z[:, [d/2]]
        return PairedData(X, Y, label='pairwise_sign_dx%d'%self.dimx)


    def dx(self):
        return self.dimx

    def dy(self):
        return 1

# end class PSPairwiseSign


class PSGaussSign(PairedSource):
    """
    A toy problem where X follows the standard multivariate Gaussian, 
    and Y = sign(product(X))*|Z| where Z ~ N(0, 1). 
    """

    def __init__(self, dx):
        """
        dx: the dimension of X 
        """
        if dx <= 0:
            raise ValueError('dx must be > 0')
        self.dimx = dx 

    def sample(self, n, seed):
        d = self.dimx 
        with util.NumpySeedContext(seed=seed):
            Z = np.random.randn(n, 1)
            X = np.random.randn(n, d)
            Xs = np.sign(X)
            Y = np.prod(Xs, 1)[:, np.newaxis]*np.abs(Z)
        return PairedData(X, Y, label='gauss_sign_dx%d'%d)


    def dx(self):
        return self.dimx 

    def dy(self):
        return 1



