import geopandas as gpd
import planetary_computer
import pystac_client
import odc.stac
import numpy as np
import time
from shapely.geometry import Point, box, shape

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())

x = 561609
y = 1024183

start = time.time()
# Filter items by spatial intersection
from pyproj import Transformer
# The items geometry are in EPSG:4326 (lon, lat)
# Our x, y are in EPSG:32648
transformer = Transformer.from_crs("epsg:32648", "epsg:4326", always_xy=True)
lon, lat = transformer.transform(x, y)
point = Point(lon, lat)

filtered_items = []
for item in items:
    geom = shape(item.geometry)
    if geom.contains(point):
        filtered_items.append(item)

filtered_items = sorted(filtered_items, key=lambda x: x.properties["eo:cloud_cover"])

print("Original items:", len(items))
print("Filtered items:", len(filtered_items))
print("Time to filter:", time.time() - start)

start = time.time()
filtered_items = [planetary_computer.sign(item) for item in filtered_items]
ds = odc.stac.load(filtered_items[:4], bands=["B02"], 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("Time to load 4 items:", time.time() - start)
print(ds["B02"].shape)