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

# # Writing xarray -> COGs
# 
# Hi all,
# 
# I'm looking for guidance / best practices on writing an xarray object to (a collection of) COGs. Let's start with a common case of a DataArray that's indexed by `(time, band, y, x)`. Let's also assume that it's a chunked DataArray, with a chunksize of 1 for `time` and `band`, and it might be chunked along `y` and `x` as well.
# 
# My high-level questions:
# 
# 1. Does rioxarray's `.rio.to_raster(path, driver="COG")` have the right defaults? Anything special we should do to make sure we write "good" COGs for a single chunk?
# 2. Is there an established convention for organizing a directory of COG files that represent a 4-d datacube?
# 
# I'm particularly interested in item 2. My proposed naming convention is
# 
# ```
# <prefix>/time=<time>/band=<band>-y=<y-coord>-x=<x-coord>.tif
# ```
# 
# This works well for xarray: we have coordinate information available when writing the chunk, so we can safely generate a unique name for a chunk using the `(time, band, y, x)` coordinates of, say, the top-left value in the chunk.
# 
# Here's a small example:

# In[ ]:


get_ipython().system('pip install -q -U --no-deps git+https://github.com/TomAugspurger/xcog')


# ## Data generation
# 
# We'll mock up some data that has the right structure for pystac / rioxarray to do their thing.

# In[4]:


import xarray as xr
import numpy as np
import dask.array as da
import stackstac
import rioxarray
import pystac
import pandas as pd

values = da.random.uniform(size=(2, 3, 10980, 10980), chunks=(1, 1, 5490, 5490))

x = np.arange(399960, 509751, step=10.)
y = np.arange(4800000, 4690210 - 1, step=-10.)
band = np.array(["B02", "B03", "B04"])
time = pd.to_datetime(["2021-01-01T17:07:19.024000000", "2021-01-04T17:17:19.024000000"])

data = xr.DataArray(
    values,
    dims=("time", "band", "y", "x"),
    coords={
        "time": xr.DataArray(time, name="time", dims="time"),
        "band": xr.DataArray(band, name="band", dims="band"),
        "y": xr.DataArray(y, name="y", dims="y"),
        "x": xr.DataArray(x, name="x", dims="x"),
        "common_name": xr.DataArray(['blue', 'green', 'red'], dims="band", name="common_name"),
        "center_wavelength":xr.DataArray([0.49 , 0.56 , 0.665], dims="band", name="center_wavelength"),
        "full_width_half_max": xr.DataArray([0.098, 0.045, 0.038], dims="band", name="full_width_half_max"),
    },
    attrs={
        "crs": "epsg:32615",
    },
)
data


# ## Data writing
# 
# We're using `xcog` here, a simple little library with some utilities for writing out chunks. We'll write to local disk, but we should be able to use any fsspec-compatible file-system.

# In[5]:


from pathlib import Path
import xcog


dst = Path("/tmp/cogs/")
dst.mkdir(parents=True, exist_ok=True)

template = xcog.make_template(data)
r = data.map_blocks(
    xcog.write_block,
    kwargs=dict(
        prefix=str(dst),
        storage_options=dict(auto_mkdir=True),
    ),
    template=template
)
r


# In[6]:


get_ipython().run_line_magic('time', 'result = r.compute()')


# Here's paths of the COGs we wrote out:

# In[7]:


get_ipython().system('tree /tmp/cogs/')


# ## Read back STAC + COGs
# 
# Our `result` DataArray is a bunch of STAC items (one per original chunk).

# In[12]:


result[0, 0, 0, 0].item()


# We can group those together (all the assets with the same ID are merged into a single item)

# In[14]:


new_items = xcog.collate(result)
new_items[:5]


# And those can be fed back to stackstac, so kind of a round-trip from DataArray -> {STAC + COG} -> DataArray

# In[15]:


stackstac.stack([x.to_dict() for x in new_items], chunksize=5490).groupby("time").apply(stackstac.mosaic)


# (We'll ignore the fact that we're off by one on the size).