import pystac_client
import dinamis_sdk

api = pystac_client.Client.open(
    'https://stacapi-dinamis.apps.okd.crocc.meso.umontpellier.fr',
    modifier=dinamis_sdk.sign_inplace
)

"""
L'exemple suivant montre comment calculer une différence de NDVI 
sur une zone donnée entre deux années consécutives
"""
import pyotb

def mosa(year):
    """Return a pyotb application that perform a mosaic."""
    res = api.search(
        bbox=[4, 42.99, 5, 44.05],
        datetime=[f"{year}-01-01", f"{year}-12-25"],
        collections=["spot-6-7-drs"]
    )
    urls = [f"/vsicurl/{r.assets['src_xs'].href}" for r in res.items()]
    return pyotb.Mosaic({"il": urls})


def ndvi(xs):
    """Return a pyotb application that computes NDVI."""
    return pyotb.BandMath({"il": [xs], "exp": "(im1b4-im1b1)/(im1b4+im1b1)"})

mosa_22 = mosa("2022")  # mosaique XS 2022
mosa_21 = mosa("2021")  # mosaique XS 2021
ndvi_22 = ndvi(mosa_22)  # NDVI 2022
ndvi_21 = ndvi(mosa_21)  # NDVI 2021
delta_ndvi = ndvi_22 - pyotb.Superimpose({"inr": ndvi_22, "inm": ndvi_21})
#delta_ndvi.write("ndvi_loss.tif?&box=5000:5000:4096:4096")

"""
L'exemple suivant montre comment on peut extraire des patches de 
128 x 128   pixels dans les images `delta_ndvi` et `mosa_22` 
et afficher ces patches dans le notebook avec un widget
"""
import matplotlib.pyplot as plt
import numpy as np
def nice_rgb(img): return np.minimum(np.maximum(0, (img - np.mean(img))/(3 * np.std(img))), 1)
mosa_subset = mosa_22[1000:1600, 1000:1800, 0:3]
mosa_arr = mosa_subset.to_numpy()
plt.imshow(nice_rgb(mosa_arr))
plt.show()

plt.imshow(delta_ndvi[1000:1600, 1000:1800, :].to_numpy())
plt.show()

vec = "patches_pos.gpkg"
inp = delta_ndvi[1000:5000, 1000:5000]

pyotb.PatchesSelection({
    "in": inp,
    "grid.step": 512,
    "grid.psize": 128,
    "strategy": "all",
    "outtrain": vec
})

patches = pyotb.PatchesExtraction({
    "source1.il": [inp],
    "source1.patchsizex": 128,
    "source1.patchsizey": 128,
    "source2.il": [mosa_22],
    "source2.patchsizex": 128,
    "source2.patchsizey": 128,
    "vec": vec,
    "field": "id"
}, n_sources=2)

from ipywidgets import interact, widgets
from IPython.display import display
import matplotlib.pyplot as plt
import numpy as np

x_arr = patches["source1.out"].to_numpy()
y_arr = patches["source2.out"].to_numpy()
patches_sz = x_arr.shape[1]
patches_nb = int(x_arr.shape[0] / patches_sz)
x_arr = np.reshape(x_arr, (patches_nb, patches_sz, patches_sz, 1))
y_arr = np.reshape(y_arr, (patches_nb, patches_sz, patches_sz, 4))

fig = plt.figure(figsize=(6, 4))
ndvi = fig.add_subplot(121) 
xs = fig.add_subplot(122) 

def show(idx):
    ndvi.imshow(x_arr[idx, :])
    xs.imshow(nice_rgb(y_arr[idx, :, :, 0:3]))
    
def f(idx):
    show(idx)
    display(fig)

show(0)
interact(f, idx=widgets.IntSlider(min=0, max=patches_nb-1, step=1, value=0));