### Obtain Heights for specified dates ###
### Create composites of the specified dates ###
### Created by Michael Fischer on 5/10/16 ###
### Edited by Dylan Card on 5/25/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('/home11/ugrad/2014/dc187197/MHC/warm_enhance.csv',delimiter=',')
warm_pure=np.genfromtxt('/home11/ugrad/2014/dc187197/MHC/warm_pure.csv',delimiter=',')
warm=np.genfromtxt('/home11/ugrad/2014/dc187197/MHC/warm.csv',delimiter=',')
cold=np.genfromtxt('/home11/ugrad/2014/dc187197/MHC/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,warm_enhance,warm_pure):
    year=var[0] 
    month=var[1]
    day=var[2]
    hour=var[3]

    print
    print
    print "############ Starting Process Heights ###################"
    print


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

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

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

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

    
        time = g_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(g_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

        g[sn,:,:,:] = g_fin.variables['g'][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
        g_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_g = np.mean(g,axis=0)
    comp_g = np.array(comp_g)
    comp_g = np.multiply(comp_g,9.81) # Convert from geopotetial meters to geopotetial height
    comp_g = np.divide(comp_g,100) # Change to Decameters


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

### Plot Composite Winds ###

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

#Plot Cold Cases
    if len(var[0])==len(cold[0]):
        plt.figure(figsize=(11,7))
        plt.title('Composite 500 hPa Heights- Cold Cases (dm) (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])
        b.drawmeridians(np.arange(west,east,10),labels=[1,0,0,1])
        bounds = np.arange(0,np.amax(comp_g[levi]),0.001) #Pick bounds for the colorbar
        x , y = b(lon_lim,lat_lim)
        levels=np.arange(500,600,6)
        CS = plt.contour(x,y,comp_g[levi,::,::],colors='r',levels=levels)
        #To remove trailing zeros 
        class nf(float):
            def __repr__(self):
                str = '%.1f' % (self.__float__(),)
                if str[-1] == '0':
                    return '%.0f' % self.__float__()
                else:
                    return '%.1f' % self.__float__()
        # Recast levels to new class
        CS.levels = [nf(val) for val in CS.levels]
        # Label levels with specially formatted floats
        if plt.rcParams["text.usetex"]:
            fmt = r'%r \%%'
        else:
            fmt = '%r'

        plt.clabel(CS, CS.levels, inline=1, fmt=fmt, fontsize=10)
        plt.savefig('heights_cold.png')


#Plot Warm Cases
    if len(var[0])==len(warm[0]):
        plt.figure(figsize=(11,7))
        plt.title('Composite 500 hPa Heights- Warm Cases (dm) (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])
        b.drawmeridians(np.arange(west,east,10),labels=[1,0,0,1])
        bounds = np.arange(0,np.amax(comp_g[levi]),0.001) #Pick bounds for the colorbar
        x , y = b(lon_lim,lat_lim)
        levels=np.arange(500,600,3)
        CS = plt.contour(x,y,comp_g[levi,::,::],colors='r',levels=levels)
        #To remove trailing zeros 
        class nf(float):
            def __repr__(self):
                str = '%.1f' % (self.__float__(),)
                if str[-1] == '0':
                    return '%.0f' % self.__float__()
                else:
                    return '%.1f' % self.__float__()
        # Recast levels to new class
        CS.levels = [nf(val) for val in CS.levels]
        # Label levels with specially formatted floats
        if plt.rcParams["text.usetex"]:
            fmt = r'%r \%%'
        else:
            fmt = '%r'

        plt.clabel(CS, CS.levels, inline=1, fmt=fmt, fontsize=10)
        plt.savefig('heights_warm.png')



#Plot Enhanced Warm Cases
    if len(var[0])==len(warm_enhance[0]):
        plt.figure(figsize=(11,7))
        plt.title('Composite 500 hPa Heights- Enhanced Warm Cases (dm) (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,np.amax(comp_g[levi]),0.001) #Pick bounds for the colorbar
        x , y = b(lon_lim,lat_lim)
        levels=np.arange(500,600,3)
        CS = plt.contour(x,y,comp_g[levi,::,::],colors='r',levels=levels)
        #To remove trailing zeros 
        class nf(float):
            def __repr__(self):
                str = '%.1f' % (self.__float__(),)
                if str[-1] == '0':
                    return '%.0f' % self.__float__()
                else:
                    return '%.1f' % self.__float__()
        # Recast levels to new class
        CS.levels = [nf(val) for val in CS.levels]
        # Label levels with specially formatted floats
        if plt.rcParams["text.usetex"]:
            fmt = r'%r \%%'
        else:
            fmt = '%r'

        plt.clabel(CS, CS.levels, inline=1, fmt=fmt, fontsize=10)
        plt.savefig('heights_warm_e.png')



#Plot Pure Warm Cases
    if len(var[0])==len(warm_pure[0]):
        plt.figure(figsize=(11,7))
        plt.title('Composite 500 hPa Heights- Pure Warm Cases (dm) (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,np.amax(comp_g[levi]),0.001) #Pick bounds for the colorbar
        x , y = b(lon_lim,lat_lim)
        levels=np.arange(500,600,3)
        CS = plt.contour(x,y,comp_g[levi,::,::],colors='r',levels=levels)
        #To remove trailing zeros 
        class nf(float):
            def __repr__(self):
                str = '%.1f' % (self.__float__(),)
                if str[-1] == '0':
                    return '%.0f' % self.__float__()
                else:
                    return '%.1f' % self.__float__()
        # Recast levels to new class
        CS.levels = [nf(val) for val in CS.levels]
        # Label levels with specially formatted floats
        if plt.rcParams["text.usetex"]:
            fmt = r'%r \%%'
        else:
            fmt = '%r'

        plt.clabel(CS, CS.levels, inline=1, fmt=fmt, fontsize=10)
        plt.savefig('heights_warm_p.png')


### End ###