import geopandas as gpd
import planetary_computer
import pystac_client
import odc.stac
import numpy as np

bbox = [105.5, 9.2, 106.3, 10.0]
time_range = "2023-01-01/2023-04-30"
catalog = pystac_client.Client.open("https://planetarycomputer.microsoft.com/api/stac/v1", modifier=planetary_computer.sign_inplace)
items = list(catalog.search(collections=["sentinel-2-l2a"], bbox=bbox, datetime=time_range, query={"eo:cloud_cover": {"lt": 30}}).items())
items = [planetary_computer.sign(item) for item in items]

x = 561609
y = 1024183
ds = odc.stac.load(items, bands=["B02", "B03", "B04", "B08", "SCL"], x=(x-80, x+80), y=(y-80, y+80), crs="EPSG:32648", resolution=10, patch_url=planetary_computer.sign, fail_on_error=False).compute()

print("Original shape:", ds["B02"].shape)
ds2 = ds.dropna(dim="time", how="all")
print("After dropna time size:", len(ds2.time))
if len(ds2.time) > 0:
    ds2 = ds2.isel(time=slice(0, 4))
    median = ds2["B02"].median(dim="time", skipna=True).values
    print("Median shape:", median.shape)
    print("Zeros in median:", np.sum(median == 0) / median.size)
    print("NaNs in median:", np.sum(np.isnan(median)) / median.size)