import h5py
import numpy as np
import datetime as dt
from datetime import timezone as tz
from matplotlib import pyplot as plt
import pyproj
from numpy import ma
from a301utils.a301_readfile import download
from a301lib.cloudsat import get_geo

z_file='2008082060027_10105_CS_2B-GEOPROF_GRANULE_P_R04_E02.h5'
download(z_file)

meters2km=1.e3

lats,lons,date_times,prof_times,dem_elevation=get_geo(z_file)
    
with h5py.File(z_file,'r') as zin:
    zvals=zin['2B-GEOPROF']['Data Fields']['Radar_Reflectivity'][...]
    factor=zin['2B-GEOPROF']['Data Fields']['Radar_Reflectivity'].attrs['factor']
    missing=zin['2B-GEOPROF']['Data Fields']['Radar_Reflectivity'].attrs['missing']
    height=zin['2B-GEOPROF']['Geolocation Fields']['Height'][...]

"""
Explanation: using pcolormesh to plot an x-z cross section of the cloudsat radar reflectivity
This notebook shows how to read in the reflectivity, convert it to dbZe (dbZ equivalent,
which means the dbZ that the measured reflectivity would have it the cloud was made
of liquid water drops)
1. Read in the height and reflectivity fields
End of explanation
"""


hit=(zvals == missing)
zvals = zvals/factor
zvals[hit]=np.nan
zvals=ma.masked_invalid(zvals)

"""
Explanation: 2. Make a masked array of the reflectivity so that pcolormesh will plot it
note that I need to find the missing data before I divide by factor=100 to
convert from int16 to float
End of explanation
"""


first_time=date_times[0]
print('orbit start: {}'.format(first_time))
start_hour=6
start_minute=45
storm_start=starttime=dt.datetime(first_time.year,first_time.month,first_time.day,
                                        start_hour,start_minute,0,tzinfo=tz.utc)
#
# get 3 minutes of data from the storm_start
#
storm_stop=storm_start + dt.timedelta(minutes=3)
print('storm start: {}'.format(storm_start))
hit = np.logical_and(date_times > storm_start,date_times < storm_stop)
lats = lats[hit]
lons=lons[hit]
prof_times=prof_times[hit]
zvals=zvals[hit,:]
height=height[hit,:]
date_times=date_times[hit]

"""
Explanation: 3. Find the part of the orbing that corresponds to the 3 minutes containing the storm
You need to enter the start_hour and start_minute for the start time of your cyclone in the granule
End of explanation
"""


great_circle=pyproj.Geod(ellps='WGS84')
distance=[0]
start=(lons[0],lats[0])
for index in np.arange(1,len(lons)):
    azi12,azi21,step= great_circle.inv(lons[index-1],lats[index-1],lons[index],lats[index])
    distance.append(distance[index-1] + step)
distance=np.array(distance)/meters2km

"""
Explanation: 4. convert time to distance by using pyproj to get the greatcircle distance between shots
End of explanation
"""


%matplotlib inline
plt.close('all')
fig,ax=plt.subplots(1,1,figsize=(40,4))

from matplotlib import cm
from matplotlib.colors import Normalize
vmin=-30
vmax=20
the_norm=Normalize(vmin=vmin,vmax=vmax,clip=False)
cmap=cm.jet
cmap.set_over('w')
cmap.set_under('b',alpha=0.2)
cmap.set_bad('0.75') #75% grey

col=ax.pcolormesh(distance,height[0,:]/meters2km,zvals.T,cmap=cmap,
                  norm=the_norm)
ax.set(ylim=[0,17],xlim=(0,1200))
fig.colorbar(col,extend='both')
ax.set(xlabel='distance (km)',ylabel='height (km)')
fig.savefig('cloudsat.png')
plt.show()    

!open cloudsat.png

"""
Explanation: 5. Make the plot assuming that height is the same for every shot
i.e. assume that height[0,:] = height[1,:] = ...
in reality, the bin heights are depend on the details of the radar returns, so
we would need to historgram the heights into a uniform set of bins -- ignore that for this qualitative picture
End of explanation
"""


