#!/usr/bin/env python
# coding: utf-8

# ## Accessing GOES data with the Planetary Computer STAC API
# 
# The [Cloud and Moisture Imagery](https://planetarycomputer.microsoft.com/dataset/goes-cmi) product contains one or more Earth-view images with pixel values identifying brightness values that are scaled to support visual analysis.
# 
# ### Data Access
# 
# The datasets hosted by the Planetary Computer are available from [Azure Blob Storage](https://docs.microsoft.com/en-us/azure/storage/blobs/). We'll use [pystac-client](https://pystac-client.readthedocs.io/) to search the Planetary Computer's [STAC API](https://planetarycomputer.microsoft.com/api/stac/v1/docs) for the subset of the data that we care about, and then we'll load the data directly from Azure Blob Storage. We'll specify a `modifier` so that we can access the data stored in the Planetary Computer's private Blob Storage Containers. See [Reading from the STAC API](https://planetarycomputer.microsoft.com/docs/quickstarts/reading-stac/) and [Using tokens for data access](https://planetarycomputer.microsoft.com/docs/concepts/sas/) for more.
# 
# This example will load data from the `goes-cmi` collection and plot it. We'll select a date and time that captures images of Hurricane Florence.

# In[1]:


import pystac_client
import planetary_computer

catalog = pystac_client.Client.open(
    "https://planetarycomputer.microsoft.com/api/stac/v1",
    modifier=planetary_computer.sign_inplace,
)

bbox = [-67.2729, 25.6000, -61.7999, 27.5423]
search = catalog.search(
    collections=["goes-cmi"],
    bbox=bbox,
    datetime=["2018-09-11T13:00:00Z", "2018-09-11T15:40:00Z"],
    limit=1,
    query={"goes:image-type": {"eq": "MESOSCALE"}},
)

item = next(search.items())


# The GOES-CMI items have many assets available for the various bands. You can view the full list at https://planetarycomputer.microsoft.com/dataset/goes-cmi. We'll use the blue, red, and near-infrared bands.

# In[2]:


import rioxarray
import xarray as xr

bands = ["C01_2km", "C02_2km", "C03_2km"]
common_names = [
    item.assets[band].extra_fields["eo:bands"][0]["common_name"] for band in bands
]
ds = xr.concat(
    [rioxarray.open_rasterio(item.assets[band].href) for band in bands], dim="band"
).assign_coords(band=common_names)
ds


# GOES doesn't have a true green band, which we need for our true color image. We'll simulate it with a linear combination of the other bands (See [Bah et. al (2018)](https://agupubs.onlinelibrary.wiley.com/doi/10.1029/2018EA000379) for more on this technique).

# In[3]:


green = (
    0.45 * ds.sel(band="red") + 0.1 * ds.sel(band="nir09") + 0.45 * ds.sel(band="blue")
).assign_coords(band="green")
green


# Finally, we'll apply a color correction.

# In[4]:


import numpy as np

γ = 2.2

rgb = xr.concat([ds, green], dim="band").sel(band=["red", "green", "blue"])
rgb = rgb / rgb.max(dim=["band", "y", "x"])
rgb = np.clip(rgb ** (1 / γ), 0, 1)


# In[5]:


import matplotlib.pyplot as plt

fig, ax = plt.subplots(figsize=(12, 12))

rgb.plot.imshow(rgb="band", ax=ax)
ax.set(title="GOES true color image");