import numpy as np
import netCDF4 as cdf
# select non-interactive backend, cf http://stackoverflow.com/questions/4931376/generating-matplotlib-graphs-without-a-running-x-server
import matplotlib
matplotlib.use('Agg') 
import matplotlib.pyplot as plt

def getdims(nc, *names): return [len(nc.dimensions[name]) for name in names]
def getvars(nc, *names): return [nc.variables[name] for name in names]

def axis_longitude():
    plt.xlim((0,360))
    plt.xlabel('longitude (degrees)')
    plt.xticks(np.arange(0,361,30))

def axis_latitude():
    plt.ylim((-90,90))
    plt.ylabel('latitude (degrees)')
    plt.yticks(np.arange(-90,91,30))

def slice_lon(nlon,llm,lon,Phi):
    lon2, z = np.zeros((llm,nlon)), np.zeros((llm,nlon))
    for lev in range(llm):
        # average from interfaces to full levels
        z[lev,:] = (.5/g)*(Phi[lev,:]+Phi[lev+1,:]) 
        lon2[lev,:] = lon[:]
    return lon2, z

#--------------------------- DCMIP21 ---------------------------

def post_DCMIP21():
    def plot_var(nlon,nlat,llm, lon,Phi,var,varname):
        # vertical slice at final time
        print 'Reading data ...'
        var, Phi = var[-1,:,nlat/2,:], Phi[-1,:,nlat/2,:]
        print '... done.'
        lon2, z = slice_lon(nlon,llm,lon,Phi)
        plt.figure(figsize=(12,6))
        plt.contourf(lon2,z,var)
        plt.colorbar() 
        plt.title('%s at final time'%varname)
        axis_longitude()
        plt.ylabel('z (m)')
        #        plt.yticks(np.arange(0, 10001, 1000))
        plt.savefig('%s.png'%varname)
    lon, lat, Omega,T,u,Phi = getvars(nc, 'lon','lat','OMEGA', 'T', 'U', 'PHI')
    plot_var(nlon,nlat,llm, lon,Phi,Omega,'Omega')
    plot_var(nlon,nlat,llm, lon,Phi,u,'u')
    plot_var(nlon,nlat,llm, lon,Phi,T,'T')

#--------------------------- DCMIP31 ---------------------------

def post_DCMIP31():
    
    def plot_T850(lon,lat,T850):  # T850 at final time
        print 'Reading data ...'
        lon, lat, T850 = lon[:], lat[:], T850[-1, :, :]
        print '... done.'
        plt.figure(figsize=(12,6))
        plt.contourf(lon,lat,T850)
        plt.colorbar() 
        plt.title('T850')
        axis_longitude()
        axis_latitude()
        plt.savefig('T850.png')
        
    def plot_dT(nlon,nlat,llm, lon,T,p,Phi): # perturbation temp, final time
        # vertical slice at final time
        print 'Reading data ...'
        T,p, Phi = T[-1,:,nlat/2,:], p[-1,:,nlat/2,:], Phi[-1,:,nlat/2,:]
        print '... done.'
        N, Teq, peq = 0.01, 300., 1e5
        N2, g2 = N*N, g*g
        G = g2/(N2*Cpd)

        lon2, z = slice_lon(nlon,llm,lon,Phi)
        theta = T*((peq/p)**kappa)
        Thetab = Teq*np.exp(N2*z/g)
        plt.figure(figsize=(12,6))
        plt.contourf(lon2,z,theta-Thetab, levels=np.arange(-0.12,0.12,0.02) )
        plt.colorbar() 
        plt.title('$\\Theta\'$')
        axis_longitude()
        plt.ylabel('z (m)')
        plt.yticks(np.arange(0, 10001, 1000))
        plt.savefig('dT.png')

    lon, lat, T850, T, Phi, p = getvars(nc, 'lon','lat','T850', 'T', 'PHI','P')
    plot_dT(nlon,nlat,llm, lon,T,p,Phi)
    plot_T850(lon,lat,T850)

#--------------------------- DCMIP41 ---------------------------

def post_DCMIP41():
    def plot_T850(day, lon,lat,T850):  # T850 at final time
        print 'Reading data ...'
        lon, lat, T850 = lon[:], lat[:], T850[day-1, :, :]
        print '... done.'
        plt.figure(figsize=(12,5))
        plt.contourf(lon,lat,T850)
        plt.colorbar() 
        plt.title('T850 at day %d'%day)
        axis_longitude()
        axis_latitude()
        plt.savefig('T850_day%02d.png'%day)

    lon, lat, T850, T, Phi = getvars(nc, 'lon','lat','T850', 'T', 'PHI')    
    plot_T850(7,lon,lat,T850)
    plot_T850(10,lon,lat,T850)
    plot_T850(30,lon,lat,T850)

#------------------------ Held & Suarez ----------------------

post_held_suarez = post_DCMIP41
    
#--------------------------- MAIN ----------------------------

gridfile = 'netcdf/output_dcmip2016_regular.nc'
nc = cdf.Dataset(gridfile, "r")
llm, nlon, nlat, ntime = getdims(nc, 'lev','lon','lat','time_counter')
Cpd, kappa, g = 1004.5, 0.2857143, 9.80616

# Now call a routine post_XXX()