# Import analysis modules
import numpy as np
import xarray as xr

# Import plotting modules
import matplotlib.pyplot as plt
from matplotlib import rcParams

#------------------------------------------------------------------------------ 
# Parameters
#------------------------------------------------------------------------------

# Set path to dataset here
data_file = "../NCEP2_time_mean.nc"


# Specify useful constants
Rd = 287.04			# Dry gas constant (J/K/kg)


#------------------------------------------------------------------------------ 
# Read data
#------------------------------------------------------------------------------
print('Reading dataset.')

# Load data
NCEP2 = xr.open_dataset(data_file)

# Calculate  zonal means 
NCEP2_zonal_mean = NCEP2.mean(dim='lon', skipna=True)

# Calculate the zonal-mean density
rho = NCEP2_zonal_mean['p'] / ( Rd * NCEP2_zonal_mean['T'] )
rho = rho.to_dataset(name='rho')

# Add theta to the original NCEP dataset
NCEP2_zonal_mean = xr.merge([NCEP2_zonal_mean, rho])

# Display some info about the variables in the dataset
print(NCEP2_zonal_mean)

#------------------------------------------------------------------------------ 
# Plot
#------------------------------------------------------------------------------
print('Plotting.')

# Close existing figures
plt.close("all")


# Turn off interactive mode
plt.ioff()

# Plot zonal-mean relative humidity
fig = plt.figure()
ax = plt.axes()



# Use either contour or contourf
contourfPlot = xr.plot.contourf(NCEP2_zonal_mean['RH'],
    x = 'lat', y = 'p', 
    ax=ax,
    add_colorbar=False,
    levels=[0,5,10,15,20,25,30,35,40,45,50,55,60,65,70,75,80,85,90,95,100],
    )

# Plotting options
plt.xlabel('Latitude [degrees]')
plt.ylabel('Pressure [hPa]')
cBar=plt.colorbar(contourfPlot, label='relative humudity [\%]')
ax.invert_yaxis()
ax.set_title('NCEP2 2008-2017: zonal- and time-mean relative humidity')
plt.show()
fig.savefig("RH_zonal_mean.png", dpi=300,  bbox_inches='tight')


# Plot Temperatue at 500 hPa and 850 hPa
fig = plt.figure()
ax = plt.axes()

T = NCEP2_zonal_mean['T']
linea = xr.plot.line(T[3,:], 'k', linewidth=1)
lineb = xr.plot.line(T[5,:], 'r', linewidth=1)
plt.xlabel('Latitude [degrees]')
plt.ylabel('temperature [K]')
ax.set_title('NCEP2 2008-2017: zonal- and time-mean temperature')
plt.legend(['850 hPa','500 hPa'])
plt.show()
fig.savefig("T850_zonal_mean.png", dpi=300,  bbox_inches='tight')