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

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 = sorted(items, key=lambda x: x.properties["eo:cloud_cover"])

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

# Find a point that fails. Let's just test a few points.
for idx, row in gdf.head(20).iterrows():
    x_coord = row['geometry'].x
    y_coord = row['geometry'].y
    
    transformer = Transformer.from_crs("epsg:32648", "epsg:4326", always_xy=True)
    lon, lat = transformer.transform(x_coord, y_coord)
    point = Point(lon, lat)
    
    filtered = []
    for item in items:
        if shape(item.geometry).contains(point):
            filtered.append(item)
            
    filtered = [planetary_computer.sign(item) for item in filtered]
    if not filtered:
        print(f"Point {idx}: NO ITEMS CONTAINS POINT!")
        continue
        
    ds = odc.stac.load(
        filtered,
        bands=["B02"],
        x=(x_coord - 80, x_coord + 80),
        y=(y_coord - 80, y_coord + 80),
        crs="EPSG:32648",
        resolution=10,
        patch_url=planetary_computer.sign,
        fail_on_error=False
    ).compute()
    
    sums = ds["B02"].sum(dim=["x", "y"]).values
    non_zero = (sums > 0).sum()
    print(f"Point {idx}: {len(filtered)} items, {non_zero} non-zero time steps")