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

# In[1]:


from pystac_client import Client

from odc.stac import stac_load, configure_rio


config = {
    "s2_l2a": {
        "assets": {
            "*": {
                "data_type": "uint16",
                "nodata": 0,
                "unit": "1",
            },
            "SCL": {
                "data_type": "uint8",
                "nodata": 0,
                "unit": "1",
            },
        },
        "aliases": {
            "costal_aerosol": "B01",
            "blue": "B02",
            "green": "B03",
            "red": "B04",
            "red_edge_1": "B05",
            "red_edge_2": "B06",
            "red_edge_3": "B07",
            "nir": "B08",
            "nir_narrow": "B08A",
            "water_vapour": "B09",
            "swir_1": "B11",
            "swir_2": "B12",
            "mask": "SCL",
            "aerosol_optical_thickness": "AOT",
            "scene_average_water_vapour": "WVP",
        },
    }
}



configure_rio(
    cloud_defaults=True,
    aws={"aws_unsigned": True},
    AWS_S3_ENDPOINT="s3.af-south-1.amazonaws.com",
)


catalog = Client.open("https://explorer.digitalearth.africa/stac")


# In[2]:


for collection in catalog.get_all_collections():
    print(collection)


# In[3]:


# %% [markdown]
# ## Find STAC Items to Load
#
# ### Define query parameters

# %%
# Set a bounding box
# [xmin, ymin, xmax, ymax] in latitude and longitude
bbox = [37.76, 12.49, 37.77, 12.50]

# Set a start and end date
start_date = "2020-09-01"
end_date = "2020-12-01"

# Set the STAC collections
collections = ["s1_rtc"]


# %% [markdown]
# ### Construct query and get items from catalog

# %%
# Build a query with the set parameters
query = catalog.search(
    bbox=bbox, collections=collections, datetime=f"{start_date}/{end_date}"
)


# In[4]:


# Search the STAC catalog for all items matching the query
items = list(query.get_items())
print(f"Found: {len(items):d} datasets")


# In[5]:


# %% [markdown]
# ## Load the Data
#
# In this step, we specify the desired coordinate system, resolution (here 20m), and bands to load. We also pass the bounding box to the `stac_load` function to only load the requested data. Since the band aliases are contained in the `config` dictionary, bands can be loaded using these aliaes (e.g. `"red"` instead of `"B04"` below).
#
# The data will be lazy-loaded with dask, meaning that is won't be loaded into memory until necessary, such as when it is displayed.

# %%
crs = "EPSG:6933"
resolution = 20

ds = stac_load(
    items,
    bands=("vh", "vv"),
    crs=crs,
    resolution=resolution,
    chunks={},
    groupby="solar_day",
    stac_cfg=config,
    bbox=bbox,
)

# View the Xarray Dataset
ds


# In[6]:


items_dict = query.get_all_items_as_dict()['features']


# In[7]:


len(items_dict)


# In[8]:


items_dict[0]


# In[9]:


import rioxarray as rxr


# In[10]:


url = 's3://deafrica-sentinel-1/s1_rtc/N12E037/2020/09/02/03F841/s1_rtc_03F841_N12E037_2020_09_02_VV.tif'


# In[11]:


da = rxr.open_rasterio(url)


# In[12]:


da


# In[13]:


da.spatial_ref.crs_wkt


# In[14]:


crs = "EPSG:4326"
resolution = 0.002

ds = stac_load(
    items,
    bands=("vh", "vv"),
    crs=crs,
    resolution=resolution,
    chunks={},
    groupby="solar_day",
    stac_cfg=config,
    bbox=bbox
)


# In[ ]: