import warnings
warnings.filterwarnings('ignore')

def load_sen1(bbox, time_range):
    import pystac_client
    import planetary_computer
    import odc.stac
    
    catalog = pystac_client.Client.open(
        "https://planetarycomputer.microsoft.com/api/stac/v1",
        modifier=planetary_computer.sign_inplace,
    )

    search = catalog.search(
        collections=["sentinel-1-rtc"],
        bbox=bbox,
        datetime=time_range,
    )
    items = list(search.items())
    print("Found items:", len(items))

    ds_s1 = odc.stac.load(
        items,
        bands=["vv", "vh"],
        bbox=bbox,
        crs="EPSG:32648",
        resolution=10,
        chunks={"x": 2048, "y": 2048, "time": 1}
    )

    ds_median = ds_s1.median(dim="time").compute()
    vv = ds_median["vv"]
    vh = ds_median["vh"]

    vv = vv.expand_dims(dim="band")
    vh = vh.expand_dims(dim="band")

    vv = vv.rio.write_crs("EPSG:32648")
    vh = vh.rio.write_crs("EPSG:32648")

    return vh, vv

print("Testing load_sen1...")
bbox = [105.5, 9.2, 106.4, 10.0]
time_range = "2022-09-01/2023-10-01"
vh, vv = load_sen1(bbox, time_range)
print("VH shape:", vh.shape)
print("VV shape:", vv.shape)
print("Success!")