import re
import os

with open('/home/x79/remote-sensing/new_import_ODC.py', 'r', encoding='utf-8') as f:
    content = f.read()

# 1. load_data
load_data_replacement = """def load_data(dc, date_range, longtitude_range, latitude_range):
    import os, hashlib
    cache_dir = "dataset_cache"
    os.makedirs(cache_dir, exist_ok=True)
    key_str = f"s2_{date_range}_{longtitude_range}_{latitude_range}"
    cache_key = hashlib.md5(key_str.encode()).hexdigest() + ".nc"
    cache_path = os.path.join(cache_dir, cache_key)
    
    if os.path.exists(cache_path):
        print(f"✅ Loading cached S2 data from {cache_path}")
        return xr.open_dataset(cache_path)

    product = 's2_l2a'
    bbox = [longtitude_range[0], latitude_range[0], longtitude_range[1], latitude_range[1]]
    
    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-2-l2a"],
        bbox=bbox,
        datetime=f"{date_range[0]}/{date_range[1]}",
    )
    items = list(search.items())
    
    data = odc.stac.load(
        items,
        bands=["red", "nir", "SCL"],
        bbox=bbox,
        crs="EPSG:32648",
        resolution=RESOLUTION,
        chunks={"x": 2048, "y": 2048, "time": 1},
        groupby="solar_day"
    )
    if "SCL" in data.data_vars:
        data = data.rename({"SCL": "scl"})
        
    print(f"💾 Caching S2 data to {cache_path}")
    data = data.compute()
    data.to_netcdf(cache_path)
    
    return data"""
content = re.sub(r'def load_data\(dc, date_range, longtitude_range, latitude_range\):.*?return data', load_data_replacement, content, flags=re.DOTALL)


# 2. load_sen1
load_sen1_replacement = """def load_sen1(bbox, time_range):
    import os, hashlib
    cache_dir = "dataset_cache"
    os.makedirs(cache_dir, exist_ok=True)
    key_str = f"s1_vh_vv_{bbox}_{time_range}"
    cache_key_vh = hashlib.md5((key_str + "vh").encode()).hexdigest() + ".nc"
    cache_key_vv = hashlib.md5((key_str + "vv").encode()).hexdigest() + ".nc"
    cache_path_vh = os.path.join(cache_dir, cache_key_vh)
    cache_path_vv = os.path.join(cache_dir, cache_key_vv)
    
    if os.path.exists(cache_path_vh) and os.path.exists(cache_path_vv):
        print(f"✅ Loading cached S1 data from {cache_path_vh} and {cache_path_vv}")
        return xr.open_dataarray(cache_path_vh), xr.open_dataarray(cache_path_vv)

    import pystac_client
    import planetary_computer
    import odc.stac
    
    # Kết nối STAC Client
    catalog = pystac_client.Client.open(
        "https://planetarycomputer.microsoft.com/api/stac/v1",
        modifier=planetary_computer.sign_inplace,
    )

    # Tìm kiếm Items
    search = catalog.search(
        collections=["sentinel-1-rtc"],
        bbox=bbox,
        datetime=time_range,
    )
    items = list(search.items())

    # Tải dữ liệu thành xarray Dataset
    ds_s1 = odc.stac.load(
        items,
        bands=["vv", "vh"],
        bbox=bbox,
        crs="EPSG:32648",
        resolution=RESOLUTION,
        chunks={"x": 2048, "y": 2048, "time": 1}
    )

    # Tính giá trị trung vị theo thời gian
    ds_median = ds_s1.median(dim="time").compute()
    vv = ds_median["vv"]
    vh = ds_median["vh"]

    # Thêm chiều 'band' để giống hệt rioxarray
    vv = vv.expand_dims(dim="band")
    vh = vh.expand_dims(dim="band")

    # Phục hồi metadata về toạ độ
    vv = vv.rio.write_crs("EPSG:32648")
    vh = vh.rio.write_crs("EPSG:32648")

    print(f"💾 Caching S1 data to {cache_dir}")
    vh.to_netcdf(cache_path_vh)
    vv.to_netcdf(cache_path_vv)

    return vh, vv"""
content = re.sub(r'def load_sen1\(bbox, time_range\):.*?return vh, vv', load_sen1_replacement, content, flags=re.DOTALL)


# 3. load_data_sen1
load_data_sen1_replacement = """def load_data_sen1(dc, date_range, coordinates):
    import os, hashlib
    longtitude_range, latitude_range = coordinates
    bbox = [longtitude_range[0], latitude_range[0], longtitude_range[1], latitude_range[1]]
    
    cache_dir = "dataset_cache"
    os.makedirs(cache_dir, exist_ok=True)
    key_str = f"data_sen1_{date_range}_{bbox}"
    cache_key_vh = hashlib.md5((key_str + "vh").encode()).hexdigest() + ".nc"
    cache_key_vv = hashlib.md5((key_str + "vv").encode()).hexdigest() + ".nc"
    cache_path_vh = os.path.join(cache_dir, cache_key_vh)
    cache_path_vv = os.path.join(cache_dir, cache_key_vv)
    
    if os.path.exists(cache_path_vh) and os.path.exists(cache_path_vv):
        print(f"✅ Loading cached S1 (coord) data")
        return xr.open_dataarray(cache_path_vh), xr.open_dataarray(cache_path_vv)
    
    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=f"{date_range[0]}/{date_range[1]}",
    )
    items = list(search.items())
    
    data_sen1 = odc.stac.load(
        items,
        bands=["vv", "vh"],
        bbox=bbox,
        crs="EPSG:32648",
        resolution=RESOLUTION,
        chunks={"x": 2048, "y": 2048, "time": 1},
        groupby="solar_day"
    )
    
    data_sen1 = data_sen1.compute()
    dsvh = data_sen1.vh
    dsvv = data_sen1.vv
    
    print(f"💾 Caching S1 (coord) data")
    dsvh.to_netcdf(cache_path_vh)
    dsvv.to_netcdf(cache_path_vv)
    
    return dsvh, dsvv"""
content = re.sub(r'def load_data_sen1\(dc, date_range, coordinates\):.*?return dsvh, dsvv', load_data_sen1_replacement, content, flags=re.DOTALL)


# 4. load_data_sen2
load_data_sen2_replacement = """def load_data_sen2(dc, date_range, coordinates):
    import os, hashlib
    longtitude_range, latitude_range = coordinates
    bbox = [longtitude_range[0], latitude_range[0], longtitude_range[1], latitude_range[1]]
    
    cache_dir = "dataset_cache"
    os.makedirs(cache_dir, exist_ok=True)
    key_str = f"data_sen2_{date_range}_{bbox}"
    cache_key = hashlib.md5(key_str.encode()).hexdigest() + ".nc"
    cache_path = os.path.join(cache_dir, cache_key)
    
    if os.path.exists(cache_path):
        print(f"✅ Loading cached S2 (coord) data from {cache_path}")
        return xr.open_dataset(cache_path)
    
    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-2-l2a"],
        bbox=bbox,
        datetime=f"{date_range[0]}/{date_range[1]}",
    )
    items = list(search.items())
    
    data = odc.stac.load(
        items,
        bands=["red", "nir", "SCL"],
        bbox=bbox,
        crs="EPSG:32648",
        resolution=RESOLUTION,
        chunks={"x": 2048, "y": 2048, "time": 1},
        groupby="solar_day"
    )
    if "SCL" in data.data_vars:
        data = data.rename({"SCL": "scl"})
        
    data = data.compute()
    print(f"💾 Caching S2 (coord) data to {cache_path}")
    data.to_netcdf(cache_path)
    
    return data"""
content = re.sub(r'def load_data_sen2\(dc, date_range, coordinates\):.*?return data', load_data_sen2_replacement, content, flags=re.DOTALL)


with open('/home/x79/remote-sensing/new_import_ODC.py', 'w', encoding='utf-8') as f:
    f.write(content)

print("Patching successful.")