### Obtain Specific Humidity for specified dates ###
### Create composites of the specified dates ###
### Created by Michael Fischer on 5/10/16 ###
### Edited by Dylan Card on 5/23/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
import matplotlib as mpl

### 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=',')

print
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 array lengths ###
#
for var in (warm,cold):
    year=var[0] 
    month=var[1]
    day=var[2]
    hour=var[3]


    print
    print "############ Starting Process Specific Humidity ###################"
    print


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

    datafile = 'http://ramadda.atmos.albany.edu:8080/repository/opendap/Top/UAlbany+Repository/CFSR/2005/q.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 = 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]

    q = np.zeros((np.size(year),levsize,np.size(lat_lim),np.size(lon_lim))) #create the q 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(day[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

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

    
        time = q_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
    
        levs = np.array(q_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

        q[sn,:,:,:] = q_fin.variables['q'][time_ind,bot_lev:top_lev+1,lat_south:lat_north+1,lon_west:lon_east+1] #retrieve the specific humidity

    #Close the data files to prevent memory leaks
        q_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_q = np.mean(q,axis=0)
    comp_q = np.array(comp_q)


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

### Plot Composite Winds ###

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

#Plot Warm Cases
    if len(var[0])==len(warm[0]):
        plt.figure(figsize=(11,7))
        #plt.title('Composite 1000 hPa Specific Humidity- Warm Cases (kg/kg) (N=36)',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
        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],fontsize=20)
        b.drawmeridians(np.arange(west,east,10),labels=[1,0,0,1],fontsize=20)
        bounds = np.arange(0,0.017,0.001) #Pick bounds for the colorbar
        cmap=plt.cm.YlGn
        x , y = b(lon_lim,lat_lim)
        norm = colors.BoundaryNorm(bounds, cmap.N)
        cont= plt.contourf(x,y,comp_q[levi,::,::],cmap=cmap, vmin=0, vmax=0.016)
        clbar =  plt.colorbar(cont,norm=norm,boundaries=bounds)
        clbar.ax.tick_params(labelsize=20)
        plt.savefig('humidity_warm.png')

#Plot Enhanced Warm Cases
    if len(var[0])==len(warm_enhance[0]):
        plt.figure(figsize=(11,7))
        plt.title('Composite 1000 hPa Specific Humidity- Enhanced Warm Cases (kg/kg) (N=17)',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
        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(0,0.017,0.001) #Pick bounds for the colorbar
        cmap=plt.cm.YlGn
        norm = colors.BoundaryNorm(bounds, cmap.N)
        x , y = b(lon_lim,lat_lim)
        cont= plt.contourf(x,y,comp_q[levi,::,::],cmap=cmap, vmin=0, vmax=0.016)
        plt.colorbar(cont,norm=norm,boundaries=bounds, ticks=bounds)
        plt.savefig('humidity_warm_e.png')

#Plot Pure Warm Cases
    if len(var[0])==len(warm_pure[0]):
        plt.figure(figsize=(11,7))
        plt.title('Composite 1000 hPa Specific Humidity- Pure Warm Cases (kg/kg) (N=19)',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
        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(0,0.017,0.001) #Pick bounds for the colorbar
        cmap=plt.cm.YlGn
        norm = colors.BoundaryNorm(bounds, cmap.N)
        x , y = b(lon_lim,lat_lim)
        cont= plt.contourf(x,y,comp_q[levi,::,::],cmap=cmap, vmin=0, vmax=0.016)
        plt.colorbar(cont,norm=norm,boundaries=bounds, ticks=bounds)
        plt.savefig('humidity_warm_p.png')

#Plot Cold Cases
    if len(var[0])==len(cold[0]):
        plt.figure(figsize=(11,7))
        #plt.title('Composite 1000 hPa Specific Humidity- Cold Cases (kg/kg) (N=12)',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
        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],fontsize=20)
        b.drawmeridians(np.arange(west,east,10),labels=[1,0,0,1],fontsize=20)
        bounds = np.arange(0,0.017,0.001) #Pick bounds for the colorbar
        cmap=plt.cm.YlGn
        x , y = b(lon_lim,lat_lim)
        norm = colors.BoundaryNorm(bounds, cmap.N)
        cont= plt.contourf(x,y,comp_q[levi,::,::],cmap=cmap, vmin=0, vmax=0.016)
        clbar =  plt.colorbar(cont,norm=norm,boundaries=bounds)
        clbar.ax.tick_params(labelsize=20)
        plt.savefig('humidity_cold.png')

### End ###