from a301utils.a301_readfile import download
import numpy as np
import h5py
import sys
import a301lib
from a301lib.geolocate import fast_hist, fast_avg, fast_count,make_plot 
from a301lib.radiation import planckInvert
from matplotlib import pyplot as plt
from a301lib.geolocate import find_corners
#
# use hdfview to see the structure of this file
#
filename = 'vancouver_hires.h5'
download(filename)
h5_file=h5py.File(filename)

def subsample(*datalist,lats=None,lons=None,llcrnr=None,urcrnr=None):
    """
    return a list of satellite scene variables (radiances, reflectivies, ndvi, etc.)
      each cropped to a subset of lats, lons and data where:
      llcnr['lat'] < lats < urcnr['lat'] and
      llcnr['lon'] < lon < urcnr['lon']
      
    Parameters
    ----------
    datalist: list of vectors or arrays of pixel values to be clipped,
        each is same size as lats and lons
    
    lats: vector or array of pixel lats
       units: degrees N
       
    lons: vector or array of pixel lons
        units: degrees E
    
   
      
    llcnr: dictionary
       lower left corner dictionary with keys ['lat','lon']
       containing the latitude and longitude of the lower left corner
       
     urcnr: dictionary
       upper right corner dictionary with keys ['lat','lon']
       containing the latitude and longitude of the upper right corner
       
    Returns
    -------
    
    list: all variables in datalist, followed by lats and lons, cropped
        to the corners

       
    """
    hit_lat = np.logical_and(lats >= llcrnr['lat'], lats <= urcrnr['lat'])
    hit_lon = np.logical_and(lons >= llcrnr['lon'], lons <= urcrnr['lon'])
    selected=np.logical_and(hit_lat,hit_lon)
    outlist=[lats[selected],lons[selected]]
    for item in datalist:
        outlist.append(item[selected])
    return outlist



lat_data=h5_file['latlon']['lat'][...]
lon_data=h5_file['latlon']['lon'][...]
chan1_refl=h5_file['data_fields']['chan1'][...]
chan2_refl=h5_file['data_fields']['chan2'][...]
ndvi = (chan2_refl - chan1_refl)/(chan2_refl + chan1_refl)
llcrnr = dict(lat=49.,lon=-123.5)
urcrnr = dict(lat=49.5,lon=-122.5)
sublats, sublons,subchan1, subchan2, subndvi  = subsample(chan1_refl,chan2_refl,ndvi,
                                                       lats=lat_data,lons=lon_data,
                                                       llcrnr=llcrnr,urcrnr=urcrnr)
corners=find_corners(sublats,sublons)

lon_min= llcrnr['lon']
lon_max = urcrnr['lon']

lat_min = llcrnr['lat']
lat_max = urcrnr['lat']
binsize = 0.006

lon_hist = fast_hist(sublons.ravel(),lon_min,lon_max,binsize=binsize)
lat_hist =  fast_hist(sublats.ravel(),lat_min,lat_max,binsize=binsize)
gridded_ndvi = fast_avg(lat_hist,lon_hist,subndvi.ravel(),bad_neg=False)
gridded_chan1 = fast_avg(lat_hist,lon_hist,subchan1.ravel(),bad_neg=False)

lat_centers=lat_hist['centers_vec']
lon_centers=lon_hist['centers_vec']
lon_array,lat_array=np.meshgrid(lon_centers,lat_centers)
print(lon_array.shape)
masked_reflects = np.ma.masked_invalid(gridded_chan1)
masked_ndvi = np.ma.masked_invalid(gridded_ndvi)

"""
Explanation: Take a look at channel 1 and NDVI at full resolution for Vancouver


Modis channel listing


Band 1 centered at 0.645 microns (red)


data acquired at 250 m resolution, August 11, 2016


I've written the measurements between -125 -> -120 deg lon and 45-50 degrees lat to
  an hdf:  vancouver_hires.h5, download in the cell below


see what channel 1 and the ndvi look like at 250 meter resoluiton
End of explanation
"""


import seaborn as sns
from matplotlib.colors import ListedColormap, LinearSegmentedColormap 
def five_colors():
    colors = ["royal blue", "baby blue", "eggshell", "burnt red", "soft pink"]
    #print([the_color for the_color in colors])
    colors=[sns.xkcd_rgb[the_color] for the_color in colors]
    cmap=ListedColormap(colors,N=5)
    cmap.set_over('w')
    cmap.set_under('k') #black
    cmap.set_bad('0.75') #75% grey
    return cmap


"""
Explanation: Here is an example of a seaborn colormap with five xkcd colors
End of explanation
"""


from matplotlib import cm
from matplotlib.colors import Normalize
def continuous_colors():
    cmap=cm.YlGnBu_r  #see http://wiki.scipy.org/Cookbook/Matplotlib/Show_colormaps
    cmap.set_over('r')
    cmap.set_under('orange',alpha=0.8)
    cmap.set_bad('pink') #75% grey
    return cmap


%matplotlib inline
#
# choose one of the two
#
cmap = continuous_colors()
#cmap = five_colors()

#
# set the range over which the pallette extends so I use
# use all my colors 
#
vmin= 0.0
vmax= 0.10
the_norm=Normalize(vmin=vmin,vmax=vmax,clip=False)
fig,ax = plt.subplots(1,1,figsize=(14,18))
corners['ax'] = ax
corners['resolution']='h'
corners['projection']='lcc'
corners['urcrnrlon'] = urcrnr['lon']
corners['urcrnrlat'] = urcrnr['lat']
corners['llcrnrlat'] = llcrnr['lat']
corners['llcrnrlon'] = llcrnr['lon']
proj = make_plot(corners,lat_sep=0.1,lon_sep=0.25)
lon_array,lat_array=np.meshgrid(lon_centers,lat_centers)
#
# translate every lat,lon pair in the scene to x,y plotting coordinates 
# for th Lambert projection
#
x,y=proj(lon_array,lat_array)
CS=proj.pcolormesh(x, y,masked_reflects, cmap=cmap, norm=the_norm)
CBar=proj.colorbar(CS, 'right', size='5%', pad='5%',extend='both')
CBar.set_label('Channel 1 reflectance',
               rotation=270,verticalalignment='bottom',size=18)
_=ax.set_title('Modis Channel 1, August 11, 2016 Vancouver',size=22)


import seaborn as sns
sns.choose_diverging_palette(as_cmap=True)

cmap=sns.diverging_palette(261, 153,sep=6, s=85, l=66,as_cmap=True)
vmin= -0.9
vmax=  0.9
cmap.set_over('c')
cmap.set_under('k',alpha=0.8)
cmap.set_bad('pink')
the_norm=Normalize(vmin=vmin,vmax=vmax,clip=False)
fig,ax2 = plt.subplots(1,1,figsize=(14,18))
corners['ax']=ax2
proj = make_plot(corners,lat_sep=0.1,lon_sep=0.25)
CS=proj.pcolormesh(x, y,masked_ndvi, cmap=cmap, norm=the_norm)
CBar=proj.colorbar(CS, 'right', size='5%', pad='5%',extend='both')
CBar.set_label('NDVI',
               rotation=270,verticalalignment='bottom',size=18)
_=ax2.set_title('Modis NDVI, August 11, 2016 Vancouver',size=22)


"""
Explanation: and here is our previous colormap
End of explanation
"""


