import time
import matplotlib.pyplot as plt
from sliderule import gedi
from sliderule import earthdata
import sliderule

# initialize client
gedi.init("slideruleearth.io", verbose=False)

# Specify region of interest from geojson
poly_fn = 'grandmesa.geojson'
region = sliderule.toregion(poly_fn)

granules = earthdata.cmr(short_name="GEDI02_A", polygon=region["poly"])

# Build GEDI L4A Request Parameters
parms = {
    "poly": region["poly"],
    "degrade_flag": 0,
    "l2_quality_flag": 1,
    "beam": 0
}

# Latch Start Time
perf_start = time.perf_counter()

# Request GEDI Data
gedi02a = gedi.gedi02ap(parms)
                        
# Latch Stop Time
perf_stop = time.perf_counter()

# Display Statistics
perf_duration = perf_stop - perf_start
print("Completed in {:.3f} seconds of wall-clock time".format(perf_duration))
print("Received {} footprints".format(gedi02a.shape[0]))
if len(gedi02a) > 0:
    print("Beams: {}".format(gedi02a["beam"].unique()))

gedi02a

# plot elevations
f, ax = plt.subplots(1, 2, figsize=[12,8])
ax[0].set_title("Elevations Highest Return")
ax[0].set_aspect('equal')
vmin_hr, vmax_hr = gedi02a['elevation_hr'].quantile((0.2, 0.8))
gedi02a.plot(ax=ax[0], column='elevation_hr', cmap='inferno', s=0.1, vmin=vmin_hr, vmax=vmax_hr)
ax[1].set_title("Elevations Lowest Mode")
ax[1].set_aspect('equal')
vmin_lm, vmax_lm = gedi02a['elevation_lm'].quantile((0.2, 0.8))
gedi02a.plot(ax=ax[1], column='elevation_lm', cmap='inferno', s=0.1, vmin=vmin_lm, vmax=vmax_lm)