import geopandas as gpd
import planetary_computer
import pystac_client
import odc.stac
import sys

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)
search = catalog.search(collections=["sentinel-2-l2a"], bbox=bbox, datetime=time_range, query={"eo:cloud_cover": {"lt": 30}})
items = list(search.items())
items = sorted(items, key=lambda x: x.properties.get("eo:cloud_cover", 100))[:4]
items = [planetary_computer.sign(item) for item in items]

gdf = gpd.read_file("train/ST_training_data_updated_1130points_new.shp")
gdf = gdf.to_crs("EPSG:32648")

row = gdf.iloc[0]
x, y_coord = row.geometry.x, row.geometry.y
point_bbox = [x - 80, y_coord - 80, x + 80, y_coord + 80]

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

print("patch_s2 vars:", patch_s2.data_vars)
if patch_s2.dims['x'] < 16 or patch_s2.dims['y'] < 16:
    print("Too small:", patch_s2.dims)
else:
    print("Success dimension:", patch_s2.dims)