#This script analyzes NCEP monthly mean data, produces
#data on the seasonal cycle, and does selected plots.
#It uses the Nio package (distributed as PyNio from the same
#place as PyNgl) for reading the data files.

Chapter = "7"
Figure = "**FIG"

from ClimateUtilities import *
import Nio

NCEPpath = '/Users/rtp1/bigdata/ncep_mm_surface/' #Edit this to appropriate data directory
#Note that the NCEP netcdf files are not currently in the WorkbookDatasets.
#They should be moved there.

#This is a specialized plotting routine for contour plots
#on map projections, with present Earth continent outlines
#It is not very extensively used in the courseware, so it
#hasn't been put in ClimateUtilities.  If you want to
#make plots of various aspects of present Earth climate
#based on netCDF data, you can use this routine as a model
#This doesn't have the features built in to allow saving the map
#to a file. Needs a lot of prettying up.
def map(v,**kwargs):
    rw = Dummy()
    rw.wkColorMap = "temp1" #"gsdtol"# "temp1"
    w = Ngl.open_wks('x11','Climate Workbook',rw)
    #
    r = Dummy()
    r.cnFillOn = True
    r.cnLinesOn = True
    
    r.pmLabelBarDisplayMode = "Always"
    r.mpProjection      = "Mercator"
    r.mpGridAndLimbOn   = True
    r.mpGridLatSpacingF = 30.
    r.mpGridLonSpacingF = 30.
    r.mpFillOn = False

    #Uncomment the following to tweak the contour levels
    r.cnLevelSelectionMode = "ManualLevels"
    r.cnMinLevelValF       = -50
    r.cnMaxLevelValF       = 50
    r.cnLevelSpacingF      = 5
    
    #I think the following only apply for the Mollweide case
    #r.vpXF = .1
    #r.vpYF = .9
    #r.vpWidthF = .8
    #r.vpHeightF = .8

    try:
        #Look for axis data in the array (just for cdms variables)
        r.sfXArray= v.getLongitude()[:]
        r.sfYArray = v.getLatitude()[:]
    except:
        #If it's not there, it has to be put in as an argument
        r.sfXArray  = kwargs['x']
        r.sfYArray  = kwargs['y']

    #Now make the plot
    r.nglDraw = False
    r.nglFrame = False
    plot = Ngl.contour_map(w,v,r)
    #Now draw the plot
    r.nglDraw = True
    Ngl.panel(w,[plot],[1,1],r)
    return plotObj(w,plot,rw) #So user can delete or save plot

#Open the file and get the temperature data
f = Nio.open_file(NCEPpath + 'air.mon.mean.nc')

#f.listvariables() will return a list of the variables in
#the file if you don't already know.  From previous experience,
#I know that the temperature in this file is called 'air'

T = f.variables['air'][:] #It's stupid but the [:] is needed so arithmetic
                          #can be done on the object. An Nio design flaw
print f.variables.keys()
# Convert to degrees K from degrees C
T = T + 273.15 

#Find the index of the starting year
yearHours = 8765.81277
t = f.variables['time'][:]/yearHours

startYear = 1970
nyears = 30

start = Numeric.searchsorted(t,startYear)

#Make a variable to hold one year of composited data for each month,
# **ToDo: Install axes
Tbar = numpy.zeros((12,T.shape[1],T.shape[2]),numpy.float)

#Compute the averages
for month in range(12):
    Tbar[month] = numpy.average(T[start+month:start+nyears*12:12],axis=0)
TbarZM = numpy.average(Tbar,axis=2) #Zonal mean. Axis 2 is longitude

#Save the data to a netcdf file. **ToDo

#Plot a subset of the zonal mean data, and save to a text file
c = Curve()
c.addCurve(f.variables['lat'][:],'lat','Latitude')
c.addCurve(TbarZM[0],'Jan','January')
c.addCurve(TbarZM[3],'Apr','April')
c.addCurve(TbarZM[6],'Jul','July')
c.addCurve(TbarZM[9],'Oct','October')
w = plot(c)
#c.dump('ZMSeasCyc.txt')

#Now make a map of the seasonal cycle
diff = Tbar[6]-Tbar[0]
lat = f.variables['lat'][:]
lon = f.variables['lon'][:].copy() #Make a copy so we can change it
#w = contour(diff,x=lon,y=lat) #Use simple contour plotter
lon[-1]=360. #Kludge to eliminate the bar at the greenwich meridian
w = map(diff,x=lon,y=lat) #Use fancy mapmaker