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

# In[1]:


get_ipython().run_line_magic('matplotlib', 'inline')

import json

import numpy as np
import rioxarray
import stackstac
from matplotlib import pyplot as plt

item = json.load(open("test-tl-stac.geojson", "rt"))
xx = stackstac.stack(
    [item], dtype="float32", bounds_latlon=(-180, -90, 180, +90), xy_coords="topleft"
)
yy = stackstac.stack(
    [item], dtype="float32", bounds_latlon=(-180, -90, 180, +90), xy_coords="center"
)


# ## Check pixel data is the same
# 
# We expect image data to be identical between the two modes

# In[2]:


assert (xx.data == yy.data).all().compute()


# In[3]:


fig, axs = plt.subplots(2, 1, figsize=(8, 7))
for idx in range(2):
    xx.isel(band=idx).isel(time=0).plot.imshow(ax=axs[idx])


# ## Check Coordinates
# 
# Data spans whole globe in `EPSG:4326` projection with 1 degree pixels.

# In[4]:


_, _, NY, NX = xx.shape
expect_x_left = np.linspace(-180, 180, NX + 1)[:-1]  # -180   -> +179
expect_x_center = expect_x_left + 0.5  # -179.5 -> 179.5

expect_y_left = np.linspace(90, -90, NY + 1)[:-1]  #  +90   -> -89
expect_y_center = expect_y_left - 0.5  #   89.5 -> -89.5


# ### Check topleft
# 
# Top-left coordinates are done correctly

# In[5]:


np.testing.assert_allclose(xx.x, expect_x_left)
np.testing.assert_allclose(xx.y, expect_y_left)


# ### Check center
# 
# Y coordinate is incorrectly computed when `xy_coord="center"`. As a result `rioxarray` reports incorrect transform and bounds.

# In[6]:


(yy.rio.transform(), yy.attrs["transform"], yy.rio.bounds(), yy.attrs["spec"].bounds)


# In[7]:


np.testing.assert_allclose(yy.x, expect_x_center)


# In[8]:


np.testing.assert_allclose(yy.y, expect_y_center)


# -------------------------------------