import numpy as np
import netCDF4 as nc
import os.path
import matplotlib.pyplot as plt
from mpl_toolkits.basemap import Basemap, shiftgrid
#from scipy.interpolate import griddata

class EOF:
    _lat = [-20,20]
    _lon = [122,290]
    def __init__(self, *args, **kwargs):
        if isinstance(args[0], basestring):
            self.initializeFromNetCDF(args[0])
        else:
            self.initializeFromData(*args, **kwargs)
        
    
    def initializeFromData(self, var, filename = None, is99variance = True, detrend=False):
        # Compute EOF
        if var.ndim == 3:
            var = var.reshape(var.shape[0], var.shape[1]*var.shape[2])
        nonMissingIndex = np.where(np.logical_not(np.isnan(var[0])))[0]
#        print nonMissingIndex.shape
        yy      = var[:, nonMissingIndex]
#        yy, Py  = self.removeSeasonal(yy, True, detrend)
        yy      = yy.T
        yy      = yy.astype(np.float64)
#           # using eigenvalue for a*a'
#           c       = np.dot(yy,yy.T)
#           d, v    = np.linalg.eig(c)
#           pc      = np.dot(v.T,yy)
        # using svd for a
        v, s, p = np.linalg.svd(yy)
        norm_yy = (yy**2).sum()
        norm_s  = (s**2).sum()
        if not np.isclose(norm_yy, norm_s):
            print 'norm_yy = {}'.format(norm_yy)
            print 'norm_s = {}'.format(norm_s)
            print 'There might be an error, because the variances are different.'

        d       = s**2/(var.shape[0]*var.shape[1])
#        d       = s**2
        eof       = np.ones([var.shape[1], v.shape[1]], dtype=var.dtype) * np.NaN
        eof       = eof.astype(var.dtype)
        eof[nonMissingIndex,:] = v
        pc      = np.dot(np.diag(s), p)
        n, perc = self.get99Variance(d)

        if is99variance:
            self.d      = d[:n]
            self.eof    = eof[:,:n]
            self.pc     = pc[:n,:]
            self.var99  = n
            self.perc   = perc[:n]
            self.is99variance = True
        else:
            self.d      = d
            self.eof    = eof
            self.pc     = pc
            self.var99  = n
            self.perc   = perc
            self.is99variance = False
#        self.seasonalmeans = Py[:12,:]
        self.nonMissingIndex = nonMissingIndex
        # Save as NetCDF
        if filename:            
            self.saveEOFAsNetCDF(filename, d[:n], eof[:,:n], pc[:n,:], perc[:n])
   
        
    @staticmethod
    def removeSeasonal(var_O, wantPy=False, detrend=False):
        if var_O.ndim == 3:
            var_O = var_O.reshape(var_O.shape[0], var_O.shape[1]*var_O.shape[2])
        nonMissingIndex = np.where(np.logical_not(np.isnan(var_O[0])))[0]
#        print nonMissingIndex.shape
        var      = var_O[:, nonMissingIndex]
#        var[var>900] = 0
#        var[var<-900] = 0
        average = True
        if average == True:
            idmonth = np.eye(12)
            years = np.ceil(var.shape[0]/12.)
            A = np.tile(idmonth, int(years))
            res = var.shape[0]%12
            A = A[:,:].T if res == 0 else A[:,:res-12].T
            AAtInv = np.linalg.inv(np.dot(A.T,A))
            Py = np.dot(np.dot(np.dot(A,AAtInv),A.T),var)
        else:
            years = var.shape[0]/12
            months_trend = np.zeros((years, var.shape[1]))
            x = np.arange(years)
            for m in range (12):
                for i in range(var.shape[1]):
                    poly1d_var = np.poly1d(np.polyfit(x, var[m::12,i], 3))
                    months_trend[m,i] = poly1d_var(x)
            Py = np.tile(months_trend, years, axis=0)
        yy = var-Py
        
        if detrend:
            spatial_length = yy.shape[1]
            time_length = yy.shape[0]
            trend = np.zeros((time_length, spatial_length))
            x = np.arange(time_length)
            for i in range(spatial_length):
#                print var[:,i]
                if np.isnan(var[:,i]).any():
                    print 'index {} has nan values'.format(i)
                    trend[:,i] = np.nan
                else:
                    poly1d_var = np.poly1d(np.polyfit(x, yy[:,i], 1))
                    trend[:,i] = poly1d_var(x)
            yy_t = yy - trend

        result = np.ones([var_O.shape[0], var_O.shape[1]], dtype=var_O.dtype) * np.NaN
        result = result.astype(var_O.dtype)
        result[:,nonMissingIndex] = yy_t if detrend else yy
        if wantPy:
            resultBar = np.ones([var_O.shape[0], var_O.shape[1]], dtype=var_O.dtype) * np.NaN
            resultBar = resultBar.astype(var_O.dtype)
            resultBar[:,nonMissingIndex] = Py + trend if detrend else Py
            return result, resultBar
        else:
            return result
        
    @staticmethod
    def removeSeasonalfromOther(new_var, var):
        if var.ndim == 3:
            var = var.reshape(var.shape[0], var.shape[1]*var.shape[2])
        if new_var.ndim == 3:
            new_var = new_var.reshape(new_var.shape[0], new_var.shape[1]*new_var.shape[2])
        var[var>900] = 0
        var[var<-900] = 0
        idmonth = np.eye(12)
        years = np.ceil(var.shape[0]/12.)
        A = np.tile(idmonth, years)
        res = var.shape[0]%12
        A = A[:,:].T if res == 0 else A[:,:res-12].T
        AAtInv = np.linalg.inv(np.dot(A.T,A))
        Py = np.dot(np.dot(np.dot(A,AAtInv),A.T),var)
        Py = Py[:12,:]
        x  = np.tile(Py.T, years)
        x  = x.T
        yy = new_var - x
        yy[x==0] = 0
        return yy
    
    @staticmethod
    def removeTrend(yy):
        spatial_length = yy.shape[1]
        time_length = yy.shape[0]
        trend = np.zeros((time_length, spatial_length))
        x = np.arange(time_length)
        for i in range(spatial_length):
#                print var[:,i]
            poly1d_var = np.poly1d(np.polyfit(x, yy[:,i], 1))
            trend[:,i] = poly1d_var(x)
        yy_t = yy - trend
        return yy_t
    
    @classmethod
    def removeTrendEvery5Years(cls, var_O):
        if var_O.ndim == 3:
            var_O = var_O.reshape(var_O.shape[0], var_O.shape[1]*var_O.shape[2])
        nonMissingIndex = np.where(np.logical_not(np.isnan(var_O[0])))[0]
#        print nonMissingIndex.shape
        var      = var_O[:, nonMissingIndex]
        
        # Remove seasonal every 5 years
        totalYear = var.shape[0]/12
        totalPeriods = (totalYear-25)/5
        
        detrend = np.zeros(var.shape)
        # first part
        detrend[:15*12] = cls.removeTrend(var[:30*12])[:15*12]
        for p in range(totalPeriods):
            detrend[(15+5*p)*12:(20+5*p)*12] = cls.removeTrend(var[(5*p)*12:(30+5*p)*12])[15*12:20*12]
        detrend[(15+(totalPeriods-1)*5)*12:] = cls.removeTrend(var[var.shape[0]-30*12:])[(15+(totalPeriods-1)*5)*12-var.shape[0]+30*12:]
        
        result = np.ones([var_O.shape[0], var_O.shape[1]], dtype=var_O.dtype) * np.NaN
        result = result.astype(var_O.dtype)
        result[:,nonMissingIndex] = detrend
        return result
        
    # 99% variance
    @staticmethod
    def get99Variance(d):
        perc    = np.absolute(d)/np.sum(np.absolute(d))
        total = 0
        n = 0
        for i in range(perc.shape[0]):
            total = total + perc[i]
            if total > 0.999:
                n = i+1
                break
        return n, perc
        
    # Plot cumulative variance explained by first k modes
    def plotCumulativeVariance(self, k=None):
        plt.close(0)        
        plt.figure(0)
        d = self.d
        if k is not None:
            d = d[:k]
        plt.plot(d.cumsum(), 'x')
        plt.xlabel('Number of modes')
        plt.ylabel('Cumulative explained variance [$C^2$]')
        plt.show()
            
            
    @staticmethod    
    def saveEOFAsNetCDF(fileName, d, v, pc, perc):
#        newlat  = np.arange(20-(-20)+1)-20
#        newlon  = np.arange((285-120)/1.5+1)*1.5+120
#        years   = 16
#        months  = np.zeros(years*12)
#        for y in range(years):
#            for m in range(12):
#                months[y*12+m] = int((1850+y)*100+m+1)
#        months  = months[20:]
        fName = 'ModelData/EOFs/' + fileName
        print 'In EOF:  saveEOFAsNetCDF: {}'.format(fName)
        if os.path.isfile(fName):
            try:
                os.remove(fName)
            except OSError:
                pass
        f   = nc.Dataset(fName, "w", format="NETCDF4")
        f.createDimension("time", pc.shape[1])
        f.createDimension("latlon", v.shape[0])
        f.createDimension("eofs", v.shape[1])
        dd          = f.createVariable('variance',"f4",("eofs"))
        vv          = f.createVariable('eof',"f4",("latlon", "eofs"))
        ppc         = f.createVariable('pc',"f4",("eofs", "time"))
        pperc       = f.createVariable('percentage',"f4",("eofs"))
        dd[:]       = d.real
        vv[:,:]     = v.real
        ppc[:,:]    = pc.real
        pperc[:]    = perc
        f.close()
#        print 'In EOF:  saveEOFAsNetCDF: {}'.format(pc.real.shape)
        
    def initializeFromNetCDF(self, fileName):
        fName = 'ModelData/EOFs/' + fileName
        print 'In EOF:  initializeFromNetCDF: {}'.format(fName)        
        f   = nc.Dataset(fName)
        d   = f.variables['variance'][:]
        eof = f.variables['eof'][:,:]
        pc  = f.variables['pc'][:,:]
        perc= f.variables['percentage'][:]
        self.d      = np.array(d)
        self.eof    = np.array(eof)
        self.pc     = np.array(pc)
        self.var99  = d.shape[0]
        self.perc   = np.array(perc)
        self.is99variance = True
        f.close()
#        print 'In EOF:  initializeFromNetCDF: {}'.format(np.array(pc).shape)
        
    def plotEOF(self, order, title, isWind=False, save=True, grid=None):
        var = self.eof[:,order]
        plt.close(0)
        plt.figure(0)
        print var.shape
        if isWind == True:
            var = var.reshape(41,226)
            v_u = var[:,:113]
            v_v = var[:,113:]
            self.plotVector(v_u, v_v)
        else:
            if var.shape[0] == 41*113:
                var = var.reshape(41,113)
                self.plotContour(var)
            elif var.shape[0] == 31*80:
                var = var.reshape(31, 80)            
                self.plotContour(var, lat=[-15,15], lon=[161,279.5])
            elif var.shape[0] == 41*40:
                var = var.reshape(41,40)
                self.plotContour(var, lat=[-20,20], lon=[221,279.5])
        plt.axes().set_aspect('auto')
        if save:
            plt.savefig('Figure/'+title[:-2]+'/'+title+'_EOF_{}.png'.format(order), bbox_inches='tight', dpi=800)
        plt.show()
        
    def plotContour(self, var, levels=None, ax=None, lat=None, lon=None):
        if lat is None:
            lat = self._lat
        if lon is None:
            lon = self._lon
        m = Basemap(llcrnrlon=lon[0],llcrnrlat=lat[0],urcrnrlat=lat[1],urcrnrlon=lon[1],projection='mill')           
        m.drawcoastlines(linewidth=1.25)
        m.fillcontinents(color='0.8')
        ny = var.shape[0]; nx = var.shape[1]
        lons, lats = m.makegrid(nx, ny) # get lat/lons of ny by nx evenly space grid.
        x, y = m(lons, lats) # compute map proj coordinates.
        #m.contourf(newlon,newlat,x, latlon=true)
        plot = m.contourf(x,y,var, levels, cmap=plt.cm.RdBu_r) if levels is not None else m.contourf(x,y,var)
#        plot.cmap.set_under('b')
#        plot.cmap.set_over('r')
        m.drawparallels(np.arange(-50,70,10),labels=[1,0,0,0])
        m.drawmeridians(np.arange(0,360,40),labels=[0,0,0,1])
        plt.colorbar(plot)
#        plt.axes().set_aspect('auto')
#        plt.savefig('Figure/HA_NC_G/PMM_Pattern_00.png', bbox_inches='tight', dpi=800)
        
    def plotVector(self, u, v, lat=None, lon=None):
        if lat is None:
            lat = self._lat
        if lon is None:
            lon = self._lon
        grid_lat  = np.arange((lat[1]-(lat[0]))/4)*4+lat[0]
        grid_lon  = np.arange((lon[1]-lon[0])/6)*6+lon[0]
        m = Basemap(llcrnrlon=lon[0],llcrnrlat=lat[0],urcrnrlat=lat[1],urcrnrlon=lon[1],projection='mill')   
        m.drawcoastlines(linewidth=1.25)
        m.fillcontinents(color='0.8')
        ny = u.shape[0]/4; nx = u.shape[1]/4
        lons, lats = m.makegrid(nx, ny) # get lat/lons of ny by nx evenly space grid.
        x, y = m(lons, lats) # compute map proj coordinates.
        latNum = np.arange(grid_lat.shape[0])*4
        lonNum = np.arange(grid_lon.shape[0])*4
        # plot wind vectors on projection grid.
        # first, shift grid so it goes from -180 to 180 (instead of 0 to 360
        # in longitude).  Otherwise, interpolation is messed up.
        u_regrid = u[latNum, :]
        u_regrid = u_regrid[:, lonNum]
        v_regrid = v[latNum, :]
        v_regrid = v_regrid[:, lonNum]
        ugrid,newlons = shiftgrid(lon[0], u_regrid.real, grid_lon)
        vgrid,newlons = shiftgrid(lon[0], v_regrid.real, grid_lon)
        # transform vectors to projection grid.
        uproj,vproj,xx,yy = m.transform_vector(ugrid,vgrid,newlons,grid_lat,nx,ny, returnxy=True)
#        uproj,vproj = \
#        m.transform_vector(ugrid,vgrid,newlons,lat,nx,ny, returnxy=True)
        # now plot.
        maximum = max([np.nanmax(u), np.nanmax(v)])
#        N = numpy.sqrt(U**2+V**2)
        m.quiver(x,y,uproj,vproj,scale= maximum*nx/1.5, pivot='mid', units='width')
#        plt.quiverkey(Q, 1.1, 0.5, 1, r'$2 m^2/s^2$', labelpos='S',
#                   fontproperties={'weight': 'bold'})
        #qk = plt.quiverkey(Q, 0.1, 0.1, 20, '20 m/s', labelpos='W')
        m.drawparallels(np.arange(-50,70,10),labels=[1,0,0,0])
        m.drawmeridians(np.arange(0,360,40),labels=[0,0,0,1])
#        return uproj,vproj, xx, yy

    def getFirstEOFs(self, numOfPCs, isWind=False):
        ori_EOFs = np.swapaxes(np.copy(self.eof[:,:numOfPCs]),0,1)
        if isWind == True:
            ori_EOFs = ori_EOFs.reshape(numOfPCs,41,226)
        else:
            ori_EOFs = ori_EOFs.reshape(numOfPCs,41,113)
        ori_fields  = np.zeros((self.pc.shape[1], ori_EOFs.shape[1], ori_EOFs.shape[2]))
        for i in range(numOfPCs):
            ori_fields = ori_fields+np.outer(self.pc[i],ori_EOFs[i,:,:]).reshape(ori_fields.shape)
        return ori_fields
    
    def plotReconstructedVariances(self, numOfPCs, var, text):
        re_var = self.getFirstEOFs(numOfPCs)
        var_re_fields = np.var(re_var, axis=0)
        var_ori_fields = np.var(var, axis=0)
        perc_var_field = 1-(var_ori_fields - var_re_fields)/var_ori_fields

        plt.close(12)
        plt.figure(12)
        f, ax = plt.subplots(1,1)
        levels = np.arange(11)/10.
        m=self.plotContour(perc_var_field, ax=ax, levels=levels)
        ax.set_title(text, fontsize=8)
#        f.subplots_adjust(right=0.925)
#        cbar_ax = f.add_axes([0.95, 0.05, 0.03, 0.83])
#        cbar=plt.colorbar(m, cax=cbar_ax)       
#        cbar.ax.tick_params(labelsize=8)
        percVarSeason = np.zeros(12)
        for i in range(12):
            percVarSeason[i] =1-((np.nanvar(var[i::12])-np.nanvar(re_var[i::12]))/np.nanvar(var[i::12]))
        print percVarSeason
#        plt.savefig('Figure/Temp2/reconstructed.png', bbox_inches='tight', dpi=800)   
        plt.show()
        return percVarSeason