### Obtain Surface Pressure for specified dates ###
### Create composites of the specified dates ###
### Created by Michael Fischer on 5/10/16 ###
### Edited by Dylan Card on 5/24/16 ###

### Import Libraries ###
import numpy as np
import netCDF4 as NC4
import datetime as DT
import matplotlib.pyplot as plt
from mpl_toolkits.basemap import Basemap
import matplotlib.colors as colors

### Establish dates to composite ###

warm_enhance=np.genfromtxt('./warm_enhance.csv',delimiter=',')
warm_pure=np.genfromtxt('./warm_pure.csv',delimiter=',')
warm=np.genfromtxt('./warm.csv',delimiter=',')
cold=np.genfromtxt('./cold.csv',delimiter=',')

warm_year=warm[0]
warm_month=warm[1]
warm_day=warm[2]
warm_hour=warm[3]

cold_year=cold[0]
cold_month=cold[1]
cold_day=cold[2]
cold_hour=cold[3]

print 'Total Warm Enhanced Cases: ' 
print len(warm_enhance[0])
print 'Total Warm Pure Cases: ' 
print len(warm_pure[0])
print 'Total Cold Cases: ' 
print len(cold[0])
print

#set up length of year array based on cold or warm
##########################################
x=cold_year 
y=cold_month
z=cold_day
h=cold_hour
##########################################

year=x
month=y
day=z
hour=h


### Load a sample year to see dimensions ###

datafile = 'http://ramadda.atmos.albany.edu:8080/repository/opendap/Top/UAlbany+Repository/CFSR/2005/pmsl.2005.0p5.anl.nc/entry.das'
fin = NC4.Dataset(datafile)

lat = np.array(fin.variables['lat'])
lon = np.array(fin.variables['lon'])


### Create a limited domain ###
def find_closest(array,value):
    idx = (np.abs(array-value)).argmin()
    return idx
    
north = 50
south = 35
west = -90  
east = -60

lat_north = find_closest(lat[:], north)
lat_south = find_closest(lat[:], south)
lon_west = find_closest(lon[:], west)
lon_east = find_closest(lon[:], east)

lat_lim = lat[lat_south:lat_north+1]
lon_lim = lon[lon_west:lon_east+1]

pmsl = np.zeros((np.size(year),np.size(lat_lim),np.size(lon_lim))) #create the surface pressur array


for sn in xrange(np.size(year)): #This loops over every storm (sn=storm number)
    year_int = int(year[sn]) 
    month_int = int(month[sn])
    day_int = int(month[sn])
    hour_int = int(hour[sn])
    year_str = str(year_int)#This is the string of the year used to retrieve the data
    print sn
    print 'Year: ' + year_str

    
    pmsl_datafile = 'http://ramadda.atmos.albany.edu:8080/repository/opendap/Top/UAlbany+Repository/CFSR/'+year_str+'/pmsl.'+year_str+'.0p5.anl.nc/entry.das'
    pmsl_fin = NC4.Dataset(pmsl_datafile)

    
    time = pmsl_fin.variables['time']
    dates = NC4.num2date(time[:],time.units).tolist() #Get all dates for the current year
    time_ind = dates.index(DT.datetime(year_int,month_int,day_int,hour_int,0)) #This is the time index
    
#retrieve the surface pressure
    pmsl[sn,:,:] = pmsl_fin.variables['pmsl'][time_ind,lat_south:lat_north+1,lon_west:lon_east+1] 
    #Close the data files to prevent memory leaks
    pmsl_fin.close()
    
### Create the composites ###

pmsl = np.divide(pmsl,100) #pressure converted to hPa
pmsl = list(pmsl)
comp_pmsl = np.mean(pmsl,axis=0)
comp_pmsl = np.array(comp_pmsl)

### Plot Composite Pressure ###

plt.figure(figsize=(11,7))
plt.title('Composite Mean Sea Level Pressure',fontsize=14)
b = Basemap(projection='cyl', lat_ts=42 ,llcrnrlon=west, urcrnrlon=east,llcrnrlat=south,urcrnrlat=north,resolution='h') #Recall spatial bounds were defined earlier
x , y = b(lon_lim,lat_lim)
b.drawcoastlines(color='k')
b.drawcountries(color='k')
b.drawstates(color='k')
b.drawcounties(color='k')
b.drawparallels(np.arange(south,north,10),labels=[1,0,0,1])
b.drawmeridians(np.arange(west,east,10),labels=[1,0,0,1])
bounds = np.arange(960,1030.1,1) #Pick bounds for the colorbar
cmap=plt.cm.jet
norm = colors.BoundaryNorm(bounds, cmap.N)
cont = plt.contourf(x,y,comp_pmsl)
plt.colorbar(cont,norm=norm,boundaries=bounds,ticks=bounds)
plt.show()




### End ###