#Copyright (c) 2008 Erik Tollerud (etolleru@uci.edu)

import numpy as np
from math import pi

class Pca(object):
    """
    A basic class for Principal Component Analysis (PCA).

    p is the number of dimensions, while N is the number of data points
    """
    _colors=('r','g','b','c','y','m','k') #defaults

    def __calc(self):
        A = self.A
        M=A-np.mean(A,axis=0)
        N=M/np.std(M,axis=0)

        self.M = M
        self.N = N
        self._eig = None

    def __init__(self,data,names=None):
        """
        p X N matrix input
        """
        from warnings import warn
        A = np.array(data).T
        n,p = A.shape
        self.n,self.p = n,p
        if p > n:
            warn('p > n - intentional?')
        self.A = A
        self._origA=A.copy()

        self.__calc()

        self._colors= np.tile(self._colors,int((p-1)/len(self._colors))+1)[:p]
        if names is not None and len(names) != p:
            raise ValueError('names must match data dimension')
        self.names = None if names is None else tuple([str(n) for n in names])


    def getCovarianceMatrix(self):
        """
        returns the covariance matrix for the dataset
        """
        return np.cov(self.N.T)

    def getEigensystem(self):
        """
        returns a tuple of (eigenvalues,eigenvectors) for the data set.
        """
        if self._eig is None:
            res = np.linalg.eig(self.getCovarianceMatrix())
            sorti=np.argsort(res[0])[::-1]
            res=(res[0][sorti],res[1][:,sorti])
            self._eig=res
        return self._eig

    def getEigenvalues(self):
        return self.getEigensystem()[0]

    def getEigenvectors(self):
        return self.getEigensystem()[1]

    def getEnergies(self):
        """
        "energies" are just normalized eigenvectors
        """
        v=self.getEigenvalues()
        return v/np.sum(v)

    def plot2d(self,ix=0,iy=1,clf=True):
        """
        Generates a 2-dimensional plot of the data set and principle components
        using matplotlib.

        ix specifies which p-dimension to put on the x-axis of the plot
        and iy specifies which to put on the y-axis (0-indexed)
        """
        import matplotlib.pyplot as plt
        x,y=self.N[:,ix],self.N[:,iy]
        if clf:
            plt.clf()
        plt.scatter(x,y)
        vals,evs=self.getEigensystem()
        #evx,evy=evs[:,ix],evs[:,iy]
        xl,xu=plt.xlim()
        yl,yu=plt.ylim()
        dx,dy=(xu-xl),(yu-yl)
        for val,vec,c in zip(vals,evs.T,self._colors):
            plt.arrow(0,0,val*vec[ix],val*vec[iy],head_width=0.05*(dx*dy/4)**0.5,fc=c,ec=c)
        #plt.arrow(0,0,vals[ix]*evs[ix,ix],vals[ix]*evs[iy,ix],head_width=0.05*(dx*dy/4)**0.5,fc='g',ec='g')
        #plt.arrow(0,0,vals[iy]*evs[ix,iy],vals[iy]*evs[iy,iy],head_width=0.05*(dx*dy/4)**0.5,fc='r',ec='r')
        if self.names is not None:
            plt.xlabel('$'+self.names[ix]+'/\\sigma$')
            plt.ylabel('$'+self.names[iy]+'/\\sigma$')

    def plot3d(self,ix=0,iy=1,iz=2,clf=True):
        """
        Generates a 3-dimensional plot of the data set and principle components
        using mayavi.

        ix, iy, and iz specify which of the input p-dimensions to place on each of
        the x,y,z axes, respectively (0-indexed).
        """
        import enthought.mayavi.mlab as M
        if clf:
            M.clf()
        z3=np.zeros(3)
        v=(self.getEigenvectors()*self.getEigenvalues())
        M.quiver3d(z3,z3,z3,v[ix],v[iy],v[iz],scale_factor=5)
        M.points3d(self.N[:,ix],self.N[:,iy],self.N[:,iz],scale_factor=0.3)
        if self.names:
            M.axes(xlabel=self.names[ix]+'/sigma',ylabel=self.names[iy]+'/sigma',zlabel=self.names[iz]+'/sigma')
        else:
            M.axes()

    def sigclip(self,sigs):
        """
        clips out all data points that are more than a certain number
        of standard deviations from the mean.

        sigs can be either a single value or a length-p sequence that
        specifies the number of standard deviations along each of the
        p dimensions.
        """
        if np.isscalar(sigs):
            sigs=sigs*np.ones(self.N.shape[1])
        sigs = sigs*np.std(self.N,axis=1)
        n = self.N.shape[0]
        m = np.all(np.abs(self.N) < sigs,axis=1)
        self.A=self.A[m]
        self.__calc()
        return n-sum(m)

    def reset(self):
        self.A = self._origA.copy()
        self.__calc()


    def project(self,vals=None,enthresh=None,nPCs=None,cumen=None):
        """
        projects the normalized values onto the components

        enthresh, nPCs, and cumen determine how many PCs to use

        if vals is None, the normalized data vectors are the values to project.
        Otherwise, it should be convertable to a p x N array

        returns n,p(>threshold) dimension array
        """
        nonnones = sum([e != None for e in (enthresh,nPCs,cumen)])
        if nonnones == 0:
            m = slice(None)
        elif nonnones > 1:
            raise ValueError("can't specify more than one threshold")
        else:
            if enthresh is not None:
                m = self.energies() > enthresh
            elif nPCs is not None:
                m = slice(None,nPCs)
            elif cumen is not None:
                m = np.cumsum(self.energies()) <  cumen
            else:
                raise RuntimeError('Should be unreachable')

        if vals is None:
            vals = self.N.T
        else:
            vals = np.array(vals,copy=False)
            if self.N.T.shape[0] != vals.shape[0]:
                raise ValueError("shape for vals doesn't match")
        proj = np.matrix(self.getEigenvectors()).T*vals
        return proj[m].T

    def deproject(self,A,normed=True):
        """
        input is an n X q array, where q <= p

        output is p X n
        """
        A=np.atleast_2d(A)
        n,q = A.shape
        p = self.A.shape[1]
        if q > p :
            raise ValueError("q > p")

        evinv=np.linalg.inv(np.matrix(self.getEigenvectors()).T)

        zs = np.zeros((n,p))
        zs[:,:q]=A

        proj = evinv*zs.T

        if normed:
            return np.array(proj.T).T
        else:
            mns=np.mean(self.A,axis=0)
            sds=np.std(self.M,axis=0)
            return (np.array(proj.T)*sds+mns).T

    def subtractPC(self,pc,vals=None):
        """
        pc can be a scalar or any sequence of pc indecies

        if vals is None, the source data is self.A, else whatever is in vals
        (which must be p x m)
        """
        if vals is None:
            vals = self.A
        else:
            vals = vals.T
            if vals.shape[1]!= self.A.shape[1]:
                raise ValueError("vals don't have the correct number of components")

        pcs=self.project()
        zpcs=np.zeros_like(pcs)
        zpcs[:,pc]=pcs[:,pc]
        upc=self.deproject(zpcs,False)

        A = vals.T-upc
        B = A.T*np.std(self.M,axis=0)
        return B+np.mean(self.A,axis=0)