### Obtain tropospheric winds for specified dates ###
### Create composites of the specified dates ###
### Created by Michael Fischer on 5/10/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 ###
#Let's randomly select to composite tropospheric winds for every Nov 27, from 2000-2010 at 12Z

year = np.arange(2000,2011,1) #Randomly chose 2000-2010. Note 2011 won't be chosen due to np.arange format
month = np.tile(11,np.size(year)) #Randomly chose November. Make sure size of array matches size of year array
day = np.tile(27,np.size(year)) #Same reasoning as month.
hour = np.tile(12,np.size(year)) #Same reasoning as month.

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

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

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

bot_lev = np.where(levels==1000)[0][0]
top_lev = np.where(levels==100)[0][0]

levsize = np.size(levels[bot_lev:top_lev+1])

### Create a limited domain ###
def find_closest(array,value):
    idx = (np.abs(array-value)).argmin()
    return idx
    
north = 80
south = 10
west = -150 + 360 #Since the ERA-Interim longitudes are arranged from 0 to 360
east = -50 + 360

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]

u = np.zeros((np.size(year),levsize,np.size(lat_lim),np.size(lon_lim))) #create the u-wind array
v = np.zeros((np.size(year),levsize,np.size(lat_lim),np.size(lon_lim))) #create the v-wind array

for sn in xrange(np.size(year)): #This loops over every storm (sn=storm number)
    year_str = str(year[sn]) #This is the string of the year used to retrieve the data
    print sn
    print year_str
    
    uwind_datafile = 'http://ramadda.atmos.albany.edu:8080/repository/opendap/Top/UAlbany+Repository/ERA-Interim/'+year_str+'/u.'+year_str+'.nc/entry.das'
    uwind_fin = NC4.Dataset(uwind_datafile)
    vwind_datafile = 'http://ramadda.atmos.albany.edu:8080/repository/opendap/Top/UAlbany+Repository/ERA-Interim/'+year_str+'/v.'+year_str+'.nc/entry.das'
    vwind_fin = NC4.Dataset(vwind_datafile)
    
    time = uwind_fin.variables['time']
    dates = NC4.num2date(time[:],time.units).tolist() #Get all dates for the current year
    time_ind = dates.index(DT.datetime(year[sn],month[sn],day[sn],hour[sn],0)) #This is the time index
    
    levs = np.array(uwind_fin.variables['lev']) #The vertical pressure level
    
    bot_lev = np.where(levs==1000)[0][0] #Start at 1000 hPa
    top_lev = np.where(levs==100)[0][0] #End at 100 hPa

    u[sn,:,:,:] = uwind_fin.variables['u'][time_ind,bot_lev:top_lev+1,lat_south:lat_north+1,lon_west:lon_east+1] #retrieve the u wind
    v[sn,:,:,:] = vwind_fin.variables['v'][time_ind,bot_lev:top_lev+1,lat_south:lat_north+1,lon_west:lon_east+1] #retrieve the v wind
    #Close the data files to prevent memory leaks
    uwind_fin.close()
    vwind_fin.close()
    
### Create the composites ###

#Since array is already in a domain around Albany, we can just take the mean across the time dimension of the array (the first axis)

comp_u = np.mean(u,axis=0)
comp_v = np.mean(v,axis=0)

#Now the arrays are only three dimensional (pressure level, lat, lon)

### Plot Composite Winds ###

levi = np.where(levs == 250)[0][0] #Let's just randomly select to plot the 250 hPa flow

skip = 5 #This is variable used for the plotting of the wind barbs, which allows better spacing of the barbs
speed = np.sqrt(comp_u**2. + comp_v**2.) #This is used for colors in the plotting

plt.figure(figsize=(11,7))
plt.title('Composite 250 hPa Winds',fontsize=14)
b = Basemap(projection='cyl', lat_ts=10,llcrnrlon=west, urcrnrlon=east,llcrnrlat=south,urcrnrlat=north,resolution='h') #Recall spatial bounds were defined earlier
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,0])
b.drawmeridians(np.arange(west,east,10),labels=[0,0,0,1])
bounds = np.arange(0,50.1,5) #Pick bounds for the colorbar
cmap=plt.cm.jet
norm = colors.BoundaryNorm(bounds, cmap.N)
barbs = plt.barbs(lon_lim[::skip],lat_lim[::skip],comp_u[levi,::skip,::skip],comp_v[levi,::skip,::skip],speed[levi,::skip,::skip],length=5,sizes=dict(spacing=0.22),cmap=cmap,norm=norm,barb_increments=dict(half=2.5, full=5, flag=25))
plt.colorbar(barbs,norm=norm,boundaries=bounds,ticks=bounds)
plt.show()


### End ###