import glob
import os
from pathlib import Path
import sys
import json
import numpy as np
import h5py
from matplotlib import pyplot as plt
from mpl_toolkits.basemap import Basemap
from datetime import datetime, timezone


cloudsat_file='MYD021KM.A2006303.2220.006.groundtrack.h5'
with h5py.File(cloudsat_file,'r') as groundfile:
    reproject_name=groundfile.attrs['reproject_filename']
    cloudsat_lats=groundfile['cloudsat_lats'][...]
    cloudsat_lons=groundfile['cloudsat_lons'][...]
    cloudsat_times=groundfile['cloudsat_times'][...]
    #
    # we know we had original datetimes in the Greenwich (UTC) timezone
    # so assign than timezone when we read them in
    #
    cloudsat_dates=[datetime.fromtimestamp(item,timezone.utc) for item in cloudsat_times]

"""
Explanation: Locating the groundtrack on the channel1 raster
This notebook shows how to use the geotiff transform to edit the cloudsat groundtrack so it only covers
the modis granule
Approach:  we need to find the [row,col] coordinates on the chan1 image for each (lat,lon) point
in the cloudsat groundtrack and see whether it has a valid radiance or not.
End of explanation
"""


with h5py.File(reproject_name,'r') as h5_file:
    basemap_args=json.loads(h5_file.attrs['basemap_args'])
    chan1=h5_file['channels']['1'][...]
    geo_string = h5_file.attrs['geotiff_args']
    geotiff_args = json.loads(geo_string)

"""
Explanation: Read in the resampled channel 1 data
End of explanation
"""


´

"""
Explanation: Prune the cloudsat data so it only covers the channel 1 image
Our ground_track.ipynb notebook showed that the lats and lons from cloudsat  run outside the 
channel 1 image.  We need a way to discard those cloudsat values that are outside the modis granule.
We can do that if we can convert between cloudsat x/y and channel1 row,column, and test to
see whether the chan1[row,col] value is missing (throw cloudsat out) or valid (keep cloudsat).
These two functions do that conversion:
Remember that we introduced this transform in the resample2 notebook, and that modis_to_h5 writes
the transform out as part of geotiff_args, specifically geotiff_args['adfgeotransform'].  This gives
us the x,y coordinates of the upper left corner and the width and height of a pixel.  From
there it's easy to figure out the col,row by counting pixels.
End of explanation
"""


bmap=Basemap(**basemap_args)
print(bmap.projparams)

"""
Explanation: First convert the lon/lat values to x,y with basemap
End of explanation
"""


x,y=bmap(cloudsat_lons,cloudsat_lats)
print('beginning and end before ',x[0],x[-1],y[0],y[-1])
x_centered = x - bmap.projparams['x_0']
y_centered = y - bmap.projparams['y_0']
print('after removing offest: ',x_centered[0],x_centered[-1],y_centered[0],y_centered[-1])


"""
Explanation: Note that basemap adds two extra items: x_0=1282439.1275676694  called "the easting" and 
y_0: 1185390.3964149079, called "the northing".  These offsets are so that basemap doesn't have
to work with negative x,y values (no idea why that is a problem for anyone, but whatever).
We need to subtract off the easting and northing before finding the row and column from the geotransform.
End of explanation
"""


print(geotiff_args)

col, row = xy_to_col_row(x_centered,y_centered,geotiff_args['adfgeotransform'])
print('chan1 height and width are {}'.format(chan1.shape))
print('row and col for beginning and end are: {} {} {} {}'.format(row[0],row[-1],col[0],col[-1]))

"""
Explanation: get the transform from geotiff_args['adfgeotransform']:
End of explanation
"""


hit=[]
for the_row,the_col in zip(row,col):
    if np.isnan(chan1[the_row,the_col]):
        #
        # missing data, don't plot
        #
        hit.append(False)
    else:
        #
        # good data, plot
        #
        hit.append(True)
hit=np.array(hit)
cloudsat_lons=cloudsat_lons[hit]
cloudsat_lats=cloudsat_lats[hit]

"""
Explanation: Now prune the cloudsat groundtrack so that it fits the chan1 granule
Now that we know the col and row for every groundtrack point, we can test to
see whether they lie within the granule (good data) or outside of it
(chan1 is np.nan)
End of explanation
"""


with h5py.File(cloudsat_file,'a') as groundfile:
    dset=groundfile.create_dataset('trimmed_lons',cloudsat_lons.shape,cloudsat_lons.dtype)
    dset[...] = cloudsat_lons[...]
    dset.attrs['long_name']='cloudsat lons that fit modis granule'
    dset=groundfile.create_dataset('trimmed_lats',cloudsat_lats.shape,cloudsat_lats.dtype)
    dset[...] = cloudsat_lats[...]
    dset.attrs['long_name']='cloudsat lats that fit modis granule'
   

"""
Explanation: overwrite the cloudsat values in the h5 file with the new data
we can't modify existing data, but we can add new vectors to the original h5 file by
opening as 'a' (append) instead of 'w' (write)
End of explanation
"""


%matplotlib inline
from matplotlib import cm
from matplotlib.colors import Normalize

cmap=cm.autumn  #see http://wiki.scipy.org/Cookbook/Matplotlib/Show_colormaps
cmap.set_over('w')
cmap.set_under('b',alpha=0.2)
cmap.set_bad('0.75') #75% grey

plt.close('all')
fig,ax = plt.subplots(1,1,figsize=(14,14))
#
# set up the Basemap object
#
basemap_args['ax']=ax
basemap_args['resolution']='c'
bmap = Basemap(**basemap_args)
#
# transform the ground track lons/lats to x/y
#
cloudsatx,cloudsaty=bmap(cloudsat_lons,cloudsat_lats)
#
# plot as blue circles
#
bmap.plot(cloudsatx,cloudsaty,'bo')
#
# now plot channel 1
#
num_meridians=180
num_parallels = 90
col = bmap.imshow(chan1, origin='upper',cmap=cmap, vmin=0, vmax=0.4)

lon_sep, lat_sep = 5,5
parallels = np.arange(-90, 90, lat_sep)
meridians = np.arange(0, 360, lon_sep)
bmap.drawparallels(parallels, labels=[1, 0, 0, 0],
                       fontsize=10, latmax=90)
bmap.drawmeridians(meridians, labels=[0, 0, 0, 1],
                       fontsize=10, latmax=90)
bmap.drawcoastlines()
colorbar=fig.colorbar(col, shrink=0.5, pad=0.05,extend='both')
colorbar.set_label('channel1 reflectivity',rotation=-90,verticalalignment='bottom')
_=ax.set(title='trimmed cloudsat groundtrack')

"""
Explanation: Repeat run of 
python -m a301utils.hf5dump MYD021KM.A2006303.2220.006.groundtrack.h5
shows the trimmed lat/lons have been added
++++++++++++++++++++
found the following top-level items: 
cloudsat_lats: &lt;HDF5 dataset "cloudsat_lats": shape (2305,), type "&lt;f4"&gt;
cloudsat_lons: &lt;HDF5 dataset "cloudsat_lons": shape (2305,), type "&lt;f4"&gt;
cloudsat_times: &lt;HDF5 dataset "cloudsat_times": shape (2305,), type "&lt;f8"&gt;
trimmed_lats: &lt;HDF5 dataset "trimmed_lats": shape (1882,), type "&lt;f4"&gt;
trimmed_lons: &lt;HDF5 dataset "trimmed_lons": shape (1882,), type "&lt;f4"&gt;
++++++++++++++++++++
_______________
root group object &lt;HDF5 dataset "cloudsat_lats": shape (2305,), type "&lt;f4"&gt;
_______________
    long_name: cloudsat latitude
    units: degrees North
_______________
root group object &lt;HDF5 dataset "cloudsat_lons": shape (2305,), type "&lt;f4"&gt;
_______________
    long_name: cloudsat longitude
    units: degrees East
_______________
root group object &lt;HDF5 dataset "cloudsat_times": shape (2305,), type "&lt;f8"&gt;
_______________
    long_name: cloudsat UTC datetime timestamp
    units: seconds since Jan. 1, 1970
_______________
root group object &lt;HDF5 dataset "trimmed_lats": shape (1882,), type "&lt;f4"&gt;
_______________
    long_name: cloudsat lats that fit modis granule
_______________
root group object &lt;HDF5 dataset "trimmed_lons": shape (1882,), type "&lt;f4"&gt;
_______________
    long_name: cloudsat lons that fit modis granule
-------------------
attributes for the root file
-------------------
attribute name:  cloudsat_filename --- value:  2006303212128_02702_CS_2B-GEOPROF-LIDAR_GRANULE_P2_R04_E02.h5
attribute name:  modis_filename --- value:  MYD021KM.A2006303.2220.006.2012078143305.h5
attribute name:  reproject_filename --- value:  MYD021KM.A2006303.2220.006.reproject.h5

now plot this
End of explanation
"""


