# Mostly-standard imports
import os
import tempfile
import numpy as np
import shutil
import urllib
import matplotlib.pyplot as plt

# Less-common-but-still-pip-installable imports
import xarray
from azure.storage.blob import ContainerClient

# pip install progressbar2, not progressbar
import progressbar

# Storage locations are documented at http://aka.ms/ai4edata-goes16
goes_account_name = 'goes'
goes_container_name = 'noaa-goes16'
goes_account_url = 'https://' + goes_account_name + '.blob.core.windows.net'
goes_blob_root = goes_account_url + '/' + goes_container_name + '/'

# Create a ContainerClient to enumerate blobs
goes_container_client = ContainerClient(account_url=goes_account_url, 
                                         container_name=goes_container_name,
                                         credential=None)

temp_dir = os.path.join(tempfile.gettempdir(),'goes')
os.makedirs(temp_dir,exist_ok=True)

%matplotlib inline

class DownloadProgressBar():
    """
    https://stackoverflow.com/questions/37748105/how-to-use-progressbar-module-with-urlretrieve
    """
    
    def __init__(self):
        self.pbar = None

    def __call__(self, block_num, block_size, total_size):
        if not self.pbar:
            self.pbar = progressbar.ProgressBar(max_value=total_size)
            self.pbar.start()
            
        downloaded = block_num * block_size
        if downloaded < total_size:
            self.pbar.update(downloaded)
        else:
            self.pbar.finish()
            

def download_url(url, destination_filename=None, progress_updater=None, force_download=False):
    """
    Download a URL to a temporary file
    """
    
    # This is not intended to guarantee uniqueness, we just know it happens to guarantee
    # uniqueness for this application.
    if destination_filename is None:
        url_as_filename = url.replace('://', '_').replace('/', '_')    
        destination_filename = \
            os.path.join(temp_dir,url_as_filename)
    if (not force_download) and (os.path.isfile(destination_filename)):
        print('Bypassing download of already-downloaded file {}'.format(
            os.path.basename(url)))
        return destination_filename
    print('Downloading file {} to {}'.format(os.path.basename(url),
                                             destination_filename),end='')
    urllib.request.urlretrieve(url, destination_filename, progress_updater)  
    assert(os.path.isfile(destination_filename))
    nBytes = os.path.getsize(destination_filename)
    print('...done, {} bytes.'.format(nBytes))
    return destination_filename

# Data are stored as product/year/day/hour/filename
product = 'ABI-L2-MCMIPF'
syear = '2020'; sday = '002'; shour = '14';

# There will be several scans this hour, we'll take the first
scan_index = 0

prefix = product + '/' + syear + '/' + sday + '/' + shour + '/'
print('Finding blobs matching prefix: {}'.format(prefix))
generator = goes_container_client.list_blobs(name_starts_with=prefix)
blobs = []
for blob in generator:
    blobs.append(blob.name)
print('Found {} scans'.format(len(blobs)))

scan_url = goes_blob_root + blobs[scan_index]

# GOES-16 MCMIPF files are ~300MB.  Not too big to fit in memory, so sometimes it may be 
# preferable to download to file first, sometimes it will be better to load straight to 
# memory.
download_to_file = True

if download_to_file:
    
    filename = download_url(scan_url,progress_updater=DownloadProgressBar())
    from datetime import datetime
    dataset = xarray.open_dataset(filename)    

else:
    
    import netCDF4
    import requests
    
    # If you know of a good way to show a progress bar with requests.get (i.e., without writing
    # to file), we're all ears, email aiforearthdatasets@microsoft.com!
    print('Downloading {} to memory...'.format(os.path.basename(scan_url)))
    response = requests.get(scan_url)
    print('Finished downloading')
    nc4_ds = netCDF4.Dataset(os.path.basename(scan_url), memory = response.content)
    store = xarray.backends.NetCDF4DataStore(nc4_ds)
    dataset = xarray.open_dataset(store)

print('Scan starts at: {}'.format(dataset.time_coverage_start))
print('Scan ends at: {}'.format(dataset.time_coverage_end))

# Bands are documented at:
#
# https://www.ncdc.noaa.gov/data-access/satellite-data/goes-r-series-satellites/glossary
#
# We'll use the red/"veggie"/blue bands with wavelengths 0.64, 0.86, and 0.47, respectively.
#
# This is close enough to RGB for today, but there's a great tutorial on getting closer to
# true color (and doing other fancy rendering tricks with GOES data!) here:
#
# https://unidata.github.io/python-gallery/examples/mapping_GOES16_TrueColor.html
#
r = dataset['CMI_C02'].data; r = np.clip(r, 0, 1)
g = dataset['CMI_C03'].data; g = np.clip(g, 0, 1)
b = dataset['CMI_C01'].data; b = np.clip(r, 0, 1)

# Brighten the image a bit for to look more stylish
gamma = 2.5; r = np.power(r, 1/gamma); g = np.power(g, 1/gamma); b = np.power(b, 1/gamma)

# Create a single RGB image for plotting
rgb = np.dstack((r, g, b))

fig = plt.figure(figsize=(7.5, 7.5), dpi=100)

# This definitely looks slicker with fancy borders on, at the cost of some extra
# imports.
show_fancy_borders = True

if not show_fancy_borders:
    
    plt.imshow(rgb); ax = plt.gca(); ax.axis('off');

else:
    
    import metpy
    import cartopy.crs as ccrs

    # Pick an arbitrary channel to get the x/y coordinates and projection information 
    # associated with the scan
    dummy_channel = dataset.metpy.parse_cf('CMI_C01')
    x = dummy_channel.x; y = dummy_channel.y

    ax = fig.add_subplot(1, 1, 1, projection=dummy_channel.metpy.cartopy_crs)
    ax.imshow(rgb, origin='upper', extent=(x.min(), x.max(), y.min(), y.max()))
    ax.coastlines(resolution='50m', color='black')
    ax.add_feature(ccrs.cartopy.feature.BORDERS);

shutil.rmtree(temp_dir)