# Standard-ish packages
import requests
import re
import numpy as np
import urllib
import io
import matplotlib.pyplot as plt
import pandas as pd

# Less standard, but all of the following are pip- or conda-installable
import rasterio

# pip install azure-storage-blob
from azure.storage.blob import ContainerClient

from osgeo import gdal,osr

# Storage locations are documented at http://aka.ms/ai4edata-hls
hls_container_name = 'hls'
hls_account_name = 'hlssa'
hls_account_url = 'https://' + hls_account_name + '.blob.core.windows.net/'
hls_blob_root = hls_account_url + hls_container_name

# This file is provided by NASA; it indicates the lat/lon extents of each
# hls tile.
#
# The file originally comes from:
#
# https://hls.gsfc.nasa.gov/wp-content/uploads/2016/10/S2_TilingSystem2-1.txt
#
# ...but as of 8/2019, there is a bug with the column names in the original file, so we
# access a copy with corrected column names.
hls_tile_extents_url = 'https://ai4edatasetspublicassets.blob.core.windows.net/assets/S2_TilingSystem2-1.txt?st=2019-08-23T03%3A25%3A57Z&se=2028-08-24T03%3A25%3A00Z&sp=rl&sv=2018-03-28&sr=b&sig=KHNZHIJuVG2KqwpnlsJ8truIT5saih8KrVj3f45ABKY%3D'

# Load this file into a table, where each row is:
#
# Tile ID, Xstart, Ystart, UZ, EPSG, MinLon, MaxLon, MinLon, MaxLon
print('Reading tile extents...')
s = requests.get(hls_tile_extents_url).content
hls_tile_extents = pd.read_csv(io.StringIO(s.decode('utf-8')),delimiter=r'\s+')
print('Read tile extents for {} tiles'.format(len(hls_tile_extents)))

# Read-only shared access signature (SAS) URL for the hls container
hls_sas_token = 'st=2019-08-07T14%3A54%3A43Z&se=2050-08-08T14%3A54%3A00Z&sp=rl&sv=2018-03-28&sr=c&sig=EYNJCexDl5yxb1TxNH%2FzILznc3TiAnJq%2FPvCumkuV5U%3D'

hls_container_client = ContainerClient(account_url=hls_account_url, 
                                         container_name=hls_container_name,
                                         credential=None)
                                

%matplotlib inline

def get_hls_tile(blob_url):
    """
    Given a URL pointing to an HLS image in blob storage, load that image via GDAL
    and return both data and metadata.
    """    
    
    formatted_gdal_bloburl='/{}/{}'.format('vsicurl',blob_url)
    
    tile_open = gdal.Open(formatted_gdal_bloburl)
    data = tile_open.GetRasterBand(1)
    ndv,xsize,ysize = data.GetNoDataValue(),tile_open.RasterXSize,tile_open.RasterYSize
    
    projection = osr.SpatialReference()
    projection.ImportFromWkt(tile_open.GetProjectionRef())
    
    datatype = data.DataType
    datatype = gdal.GetDataTypeName(datatype)  
    data_array = data.ReadAsArray()

    return ndv,xsize,ysize,projection,data_array


def list_available_tiles(prefix):
    """
    List all blobs in an Azure blob container matching a prefix.  
    
    We'll use this to query tiles by location and year.
    """
    
    files = []
    generator = hls_container_client.list_blobs(name_starts_with=prefix)
    for blob in generator:
        files.append(blob.name)
    return files

    
def lat_lon_to_hls_tile_id(lat,lon):
    """
    Get the hls tile ID for a given lat/lon coordinate pair
    """  
    
    found_matching_tile = False

    for i_row,row in hls_tile_extents.iterrows():
        found_matching_tile = lat >= row.MinLat and lat <= row.MaxLat \
        and lon >= row.MinLon and lon <= row.MaxLon
        if found_matching_tile:
            break
    
    if not found_matching_tile:
        return None
    else:
        return row.TilID

# Specify a location and year of interest
lat = 47.6101; lon = -122.2015 # Bellevue, WA

year = '2019'
daynum = '109'    # 1-indexed day-of-year
folder = 'S309'   # 'S309' for Sentinel, 'L309' for Landsat
product = 'S30'   # 'S30' for Sentinel, 'L30' for Landsat
year = '2019'

tile_id = lat_lon_to_hls_tile_id(lat,lon)
assert tile_id is not None, 'Invalid lat/lon'
prefix = folder + '/HLS.' + product + '.T' + tile_id + '.' + year

print('Finding tiles with prefix {}'.format(prefix))
matches = list_available_tiles(prefix)
assert len(matches) > 0, 'No matching tiles'

blob_name = matches[0]
print('Found {} matching tiles, using file {}'.format(len(matches),blob_name))

lat = 47.6101; lon = -122.2015 # Bellevue, WA

year    = '2019'
daynum  = '001'   # 1-indexed day-of-year
folder  = 'S309'  # 'S309' for Sentinel, 'L309' for Landsat
product = 'S30'   # 'S30' for Sentinel, 'L30' for Landsat
band    = '01'
tile_id = '10TET' # See hls.gsfc.nasa.gov/wp-content/uploads/2016/10/S2_TilingSystem2-1.txt
version = 'v1.4'  # Currently always v1.4

blob_name = folder + '/HLS.' + product + '.T' + tile_id + '.' + year + daynum + '.' + version \
    + '_' + band + '.tif'

print('Using file {}'.format(blob_name))

gdal.SetConfigOption('GDAL_HTTP_UNSAFESSL', 'YES')
blob_url = hls_blob_root + '/' + blob_name
print('Reading tile from {}'.format(blob_url))
ndv,xsize,ysize,projection,data_array = get_hls_tile(blob_url)

print('No-data value: {}'.format(ndv))
print('\nSize: {},{}'.format(xsize,ysize))
print('\nProjection:\n{}'.format(projection))

# Bands 2, 3, and 4 are B, G, and R in Sentinel-2 HLS images

base_url = '/vsicurl/' + hls_blob_root + '/' + blob_name
band2_url = re.sub('_(\d+).tif','_02.tif',base_url)
band3_url = re.sub('_(\d+).tif','_03.tif',base_url)
band4_url = re.sub('_(\d+).tif','_04.tif',base_url)
print('Reading bands from:\n{}\n{}\n{}'.format(band2_url,band3_url,band4_url))

band2 = rasterio.open(band2_url)
band3 = rasterio.open(band3_url)
band4 = rasterio.open(band4_url)

norm_value = 2000
image_data = []
for band in [band4,band3,band2]:
    band_array = band.read(1)
    band_array = band_array / norm_value
    image_data.append(band_array)
    band.close()

rgb = np.dstack((image_data[0],image_data[1],image_data[2]))
np.clip(rgb,0,1,rgb)
plt.imshow(rgb)