# Import required modules.
from osgeo import ogr, osr
from rtree.index import Index
from requests.auth import HTTPBasicAuth
import requests
import os

# Set script parameters, edit the USGS Earthdata username, password, and any file paths as needed.
url_base = 'http://e4ftl01.cr.usgs.gov/MEASURES/SRTMGL1.003/2000.02.11/'
srtm_grid_lyr = r'/home/jovyan/work/data/srtm/srtm30m_bounding_boxes.json'
census_lyr = r'/home/jovyan/work/data/census/lpr_000b16a_e/lpr_000b16a_e.shp'
out_dir = r'/home/jovyan/work/data/srtm'
idx = Index()
geoms = {}
count = 0
grid_cells = set()
username = ''
password = ''

# Open census layer and filter to Ontario.
ds = ogr.Open(census_lyr)
lyr = ds.GetLayer(0)
lyr.SetAttributeFilter("PRNAME = 'Ontario'")

# Read feature geometries and transform to EPSG:4326 (wgs84).
source_sr = lyr.GetSpatialRef()
target_sr = osr.SpatialReference()
target_sr.ImportFromEPSG(4326)
target_sr.SetAxisMappingStrategy(osr.OAMS_TRADITIONAL_GIS_ORDER)
transform = osr.CoordinateTransformation(source_sr, target_sr)
for feat in lyr:
    geom = feat.GetGeometryRef()
    geom.Transform(transform)
    
    # Create a convex hull geometry (minimum bounding polygon) for later use.
    convex_hull = geom.ConvexHull().Clone()
    
    # Iterate over individual polygons in multipolygon, and add bbox envelope to rtree spatial index.
    for g in geom:
        geoms[count] = g.Clone()
        x_min, x_max, y_min, y_max = g.GetEnvelope()
        idx.insert(count, (x_min, y_min, x_max, y_max))
        count += 1
lyr = None
ds = None

# Open SRTM grid layer and filter to convex hull.
ds = ogr.Open(srtm_grid_lyr)
lyr = ds.GetLayer(0)
lyr.SetSpatialFilter(convex_hull)

# Read filtered features and check for intersection with Ontario geometries and save file name if intersecting.
for feat in lyr:
    geom = feat.GetGeometryRef()
    x_min, x_max, y_min, y_max = geom.GetEnvelope()
    for fid in idx.intersection((x_min, y_min, x_max, y_max)):
        if geom.Intersects(geoms[fid]):
            grid_cells.add(feat.GetField('dataFile'))
            break
lyr = None
ds = None

# Set up session for auth and redirect handling, and iterate over grid cells.
with requests.Session() as session:
    session.auth = (username, password)
    for grid_cell in grid_cells:
        url = url_base + grid_cell
        out_file = os.path.join(out_dir, grid_cell)
        
        # Submit request for grid cell file if it has not already been downloaded.
        if not os.path.exists(out_file):
            req = session.request('get', url)
            response = session.get(req.url, stream=True)
            response.raise_for_status()

            # Write response to file.
            with open(out_file, 'wb') as output:
                for chunk in response.iter_content(chunk_size=1024*1024):
                    output.write(chunk)