from common import *

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.'
    Cpd, kappa, g = 1004.5, 0.2857143, 9.80616
    N, Teq, peq = 0.01, 300., 1e5
    N2, g2 = N*N, g*g
    G = g2/(N2*Cpd)

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

    theta = T*((peq/p)**kappa)
    Thetab = Teq*np.exp(N2*z/g)
    Tb = G + (Teq-G)*np.exp(N2*z/g) # background temperature
    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')

gridfile = 'netcdf/output_dcmip2016_regular.nc'
nc = cdf.Dataset(gridfile, "r")
llm, nlon, nlat, ntime = getdims(nc, 'lev','lon','lat','time_counter')
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)