import glob
import numpy as np
import xarray as xr
import hvplot.xarray
import hvplot.pandas
import holoviews as hv
import planetary_computer
from holoviews import opts
from holoviews.element import tiles
import cartopy.crs as ccrs
import cartopy.crs as ccrs
import cartopy.feature as cfeature
import matplotlib.pyplot as plt
from pystac_client import Client
import odc.stac
from shapely.geometry import Polygon, MultiPolygon, box
from shapely.ops import unary_union
import rioxarray as rxr
import pandas as pd
import geopandas as gpd
from sklearn.metrics import classification_report
from sklearn.ensemble import RandomForestClassifier
from sklearn.metrics import accuracy_score, ConfusionMatrixDisplay
from sklearn.model_selection import train_test_split
#from cuml import RandomForestClassifier
#import cupy as cp
#from cupyx.scipy.ndimage import convolve
#from dask import compute, delayed
#from dask_ml.model_selection import train_test_split
from rasterio.enums import Resampling

xr.set_options(keep_attrs=True)

admin_boundaries_gdf = gpd.read_file("./2016_phc_vut_pid_4326.geojson")
admin_boundaries_gdf = admin_boundaries_gdf.set_index(keys="pname")  # set province name as the index

# Download the data (subset_vbos_smaller.zip)
!gdown "https://drive.google.com/uc?id=1rLNwGQf8C93CXpdQ293zKaKiQejDdJnp"

# Unzip the compressed data
!mkdir -p flood/
!unzip subset_vbos_smaller.zip -d flood/

ls flood/subset_vbos_smaller/

# Get list of all labels
mask_paths = glob.glob('flood/subset_vbos_smaller/tst-ground-truth-flood-masks/*.tif')

len(mask_paths)

df = pd.read_csv('flood/subset_vbos_smaller/metadata/metadata.csv')
# Filter for specific countries
# Ensure the column is in datetime format
df["event_date"] = pd.to_datetime(df["Event Date"], format="%m/%d/%Y", errors="coerce")

# Filter for rows where the year is 2020
#df_2020 = df[df["event_date"].dt.year == 2020]
#print(len(df_2020))

selected_countries = {"Vanuatu"} #, "Timor-Leste" , "Philippines"} #, "Viet Nam", "Australia"}
df = df[df["Country"].isin(selected_countries)]

print("len(df):", len(df))

df

ls flood/subset_vbos_smaller/tst-ground-truth-flood-masks/EMSR434_AOI05_GRA_PRODUCT_ground_truth.tif

# Read the flood mask as a DataArray
flood_mask = rxr.open_rasterio(
    "flood/subset_vbos_smaller/tst-ground-truth-flood-masks/EMSR434_AOI05_GRA_PRODUCT_ground_truth.tif",
    masked=True # masks no data values
).squeeze("band")  # Remove band dimension if only 1 band

np.unique(flood_mask.data)

flood_mask.plot.imshow(cmap="Blues")

flood_mask.rio.crs

# Reproject vector data to match flood_mask CRS
if hasattr(flood_mask, 'rio') and flood_mask.rio.crs is not None:
    admin_boundaries_gdf = admin_boundaries_gdf.to_crs(flood_mask.rio.crs)

fig, ax = plt.subplots(figsize=(10, 8))

# MALAMPA boundary (the province which overlaps witht he flood mask)
malampa_geom = admin_boundaries_gdf.loc["MALAMPA"]
malampa_gdf = gpd.GeoDataFrame(geometry=[malampa_geom.geometry], crs=admin_boundaries_gdf.crs)
malampa_gdf.boundary.plot(ax=ax, edgecolor='green', linewidth=2, label="Boundary of Malampa")

# Plot the bounding box of Malampa
minx, miny, maxx, maxy = malampa_gdf.total_bounds
bbox_geom = box(minx, miny, maxx, maxy)
bbox_gdf = gpd.GeoDataFrame(geometry=[bbox_geom], crs=malampa_gdf.crs)
minx, miny, maxx, maxy = bbox_gdf.total_bounds
ax.set_xlim(minx - 0.1, maxx + 0.1)
ax.set_ylim(miny - 0.1, maxy + 0.1)
bbox_gdf.boundary.plot(ax=ax, edgecolor='red', linewidth=2, label="Bounding Box of Malampa")


# Plot flood mask on top
flood_mask.plot.imshow(ax=ax, cmap="Blues", alpha=0.6, add_colorbar=True)

plt.title("Flood Mask on Top of Bounding Box and MALAMPA Boundary")
plt.legend()
plt.tight_layout()
plt.show()

minx, miny, maxx, maxy = malampa_gdf.total_bounds

# Latitude decreases as you go south, so maxy is the northern bound and miny the southern
lat_slice = slice(maxy, miny)   # from north to south
lon_slice = slice(minx, maxx)   # from west to east

# Bounding box for Malampa largest island
print("lat_slice:", lat_slice)
print("lon_slice:", lon_slice)

PAT = "edh_pat_..."

ds = xr.open_dataset(
    f"https://edh:{PAT}@data.earthdatahub.destine.eu/era5/reanalysis-era5-land-no-antartica-v0.zarr",
    chunks={},
    engine="zarr",
)
ds

ds.data_vars

# Load ERA5 dataset
# Variables: ds.tp, ds.swvl1, ds.swvl2, ds.ro, ds.msl

# Convert units and subset
tp_mm = ds.tp * 1000  # m → mm
tp_mm.attrs["units"] = "mm"
tp_mm_vanu = tp_mm.sel(latitude=lat_slice, longitude=lon_slice)

swvl1_mm = ds.swvl1 * 1000  # m³/m³ → mm
swvl1_mm.attrs["units"] = "mm"
swvl1_vanu = swvl1_mm.sel(latitude=lat_slice, longitude=lon_slice)

swvl2_mm = ds.swvl2 * 1000
swvl2_mm.attrs["units"] = "mm"
swvl2_vanu = swvl2_mm.sel(latitude=lat_slice, longitude=lon_slice)

ro_mm = ds.ro * 1000  # m → mm
ro_mm.attrs["units"] = "mm"
ro_vanu = ro_mm.sel(latitude=lat_slice, longitude=lon_slice)

pev_mm = ds.pev * 1000  # m/day → mm/day
pev_mm.attrs["units"] = "mm/day"
pev_vanu = pev_mm.sel(latitude=lat_slice, longitude=lon_slice)

# Combine into a composite xarray Dataset
flood_composite = xr.Dataset({
    "precip_mm": tp_mm_vanu,
    "soil_moisture_l1_mm": swvl1_vanu,
    "soil_moisture_l2_mm": swvl2_vanu,
    "runoff_mm": ro_vanu,
    "potential_evaporation_mm_day": pev_vanu,
})

flood_composite

flood_composite_04072020 = flood_composite.sel(
    valid_time="2020-04-07"
)

flood_composite_04072020

flood_composite_04072020

# Select only the year 2020
flood_composite_2020 = flood_composite.sel(
    valid_time=slice('2020', '2020')
)

flood_composite_2020

# Select only the year 2021
flood_composite_2021 = flood_composite.sel(
    valid_time=slice('2021', '2021')
)

flood_composite_2021

hv.extension('bokeh')

proj = ccrs.Mercator.GOOGLE
precip_plot = flood_composite_04072020["precip_mm"].hvplot(
    x='longitude',
    y='latitude',
    geo=True,
    cmap='Blues_r',
    clim=(0, 400),
    title="Total precipitation 04/07/2020",
    frame_width=600,
    colorbar=True,
)

# Add basemap and set projection globally
basemap = tiles.OSM().opts(alpha=0.5)

combined = (precip_plot * basemap).opts(
    opts.Overlay(
        projection=proj,
        global_extent=False,
        frame_width=600
    )
)

combined

def plot_xarray_maps(datasets, titles=None, projection=ccrs.UTM(zone=59, southern_hemisphere=True), vmax=None, cmap='Blues'):
    """
    Plot multiple xarray 2D datasets.
    Zooms into Vanuatu region.
    """
    n = len(datasets)
    fig, axs = plt.subplots(1, n, figsize=(6 * n, 6),
                            subplot_kw={'projection': projection})

    if n == 1:
        axs = [axs]

    for i, (data, ax) in enumerate(zip(datasets, axs)):
        im = data.plot.imshow(
            ax=ax,
            transform=ccrs.PlateCarree(),
            cmap=cmap,
            vmax=vmax,
            add_colorbar=True
        )

        # Try to set extent from data bounds if available
        try:
            lon_min = float(data.lon.min())
            lon_max = float(data.lon.max())
            lat_min = float(data.lat.min())
            lat_max = float(data.lat.max())
            ax.set_extent([lon_min, lon_max, lat_min, lat_max], crs=ccrs.PlateCarree())
        except Exception:
            # Fallback to manual Vanuatu bounding box
            ax.set_extent([166, 170, -17.5, -12.5], crs=ccrs.PlateCarree())

        ax.coastlines(resolution='10m', linewidth=0.5)
        ax.add_feature(cfeature.BORDERS, linewidth=0.3)
        ax.set_title(titles[i] if titles else f"Map {i+1}")
        ax.gridlines(draw_labels=True, alpha=0.5)

    plt.tight_layout()
    plt.show()

plot_xarray_maps(
    datasets=[flood_composite_2020["precip_mm"].sum("valid_time"), flood_composite_2021["precip_mm"].sum("valid_time")],
    titles=[
        "Total precipitation, 2020",
        "Total precipitation, 2021"
    ],
    projection=ccrs.UTM(zone=59, southern_hemisphere=True),

    vmax=400,
    cmap="Blues"
)

flood_composite_2020["precip_mm"].sum("valid_time").values.max(), flood_composite_2021["precip_mm"].sum("valid_time").values.max()

# Spatial average over lat and lon dims
flood_composite_2020_mean = flood_composite_2020["precip_mm"].mean(dim=['latitude', 'longitude'])
flood_composite_2021_mean = flood_composite_2021["precip_mm"].mean(dim=['latitude', 'longitude'])

plt.style.use('fivethirtyeight')
fig, ax = plt.subplots(figsize=(10,6))

ax.plot(flood_composite_2020_mean['valid_time'], flood_composite_2020_mean.values, label='2020 Accumulated Precipitation', linewidth=2)
ax.plot(flood_composite_2021_mean['valid_time'], flood_composite_2021_mean.values, label='2021 Accumulated Precipitation', linewidth=2)

ax.set_ylim([-100, 1750])
ax.set_title('Accumulated precipitation [mm], Vanuatu (Approximate bbox)', loc='left', fontsize=14)
ax.set_xlabel('Time')
ax.set_ylabel('Precipitation (mm)')
ax.legend(prop={'size': 10})

plt.show()

flood_composite_04072020 = flood_composite_04072020.rio.write_crs("EPSG:4326")
flood_composite_2020 = flood_composite_2020.rio.write_crs("EPSG:4326")
flood_composite_2021 = flood_composite_2021.rio.write_crs("EPSG:4326")

# Connect to Microsoft Planetary Computer STAC
stac_client = Client.open("https://planetarycomputer.microsoft.com/api/stac/v1")

# Search NASA DEM
search = stac_client.search(
    collections=["nasadem"],
    bbox=malampa_gdf.bounds,
    max_items=1
)

dem_item = planetary_computer.sign(search.item_collection()[0])

dem_item

dem = odc.stac.load(
    items=[dem_item],
    crs="EPSG:4326",  # Match with other layers
    bbox=malampa_gdf.bounds,  # Same bounding box as the composite
)

dem

dem = dem.rename({'time': 'valid_time'})

# Check output ranges
vals = dem["elevation"].isel(valid_time=0).values
vals_clean = vals[np.isfinite(vals)]

print("Unique:", np.unique(vals_clean))
print("Min:", vals_clean.min())
print("Max:", vals_clean.max())

flood_composite_2020.dims, dem.dims

# Calculate mean elevation ignoring NaNs
mean_elev = dem["elevation"].isel(valid_time=0).mean(skipna=True)

# Fill NaNs with mean elevation
dem_filled = dem["elevation"].isel(valid_time=0).fillna(mean_elev)

# Reproject DEM to match the flood dataset grid
dem_interp = dem_filled.interp_like(flood_composite_2020, method="nearest")

# Expand dims to add valid_time dimension matching flood_composite_2020's time dimension
dem_expanded = dem_interp.expand_dims(valid_time=flood_composite_2020.valid_time)

# Now assign this expanded DEM to the composite dataset
flood_composite_2020["elevation"] = dem_expanded

# Check shapes to verify
print(flood_composite_2020["elevation"].shape)  # Should be matching the flood composite (valid_time, lat, lon)

flood_composite_2020

# Expand dims to add valid_time dimension matching flood_composite_2021's time dimension
dem_expanded = dem_interp.expand_dims(valid_time=flood_composite_2021.valid_time)

# Now assign this expanded DEM to the composite dataset
flood_composite_2021["elevation"] = dem_expanded

# Check shapes to verify
print(flood_composite_2021["elevation"].shape)  # Should be matching the flood composite (valid_time, lat, lon)

# Expand dims to add valid_time dimension matching flood_composite_2021's time dimension
dem_expanded = dem_interp.expand_dims(valid_time=flood_composite_04072020.valid_time)

# Now assign this expanded DEM to the composite dataset
flood_composite_04072020["elevation"] = dem_expanded

# Check shapes to verify
print(flood_composite_04072020["elevation"].shape)  # Should be matching the flood composite (valid_time, lat, lon)

# Check output ranges
vals = flood_composite_2020["elevation"].isel(valid_time=0).values
vals_clean = vals[np.isfinite(vals)]
print("Min:", vals_clean.min())
print("Max:", vals_clean.max())

# Check output ranges
vals = flood_composite_2021["elevation"].isel(valid_time=0).values
vals_clean = vals[np.isfinite(vals)]
print("Min:", vals_clean.min())
print("Max:", vals_clean.max())

# Check output ranges
vals = flood_composite_04072020["elevation"].isel(valid_time=0).values
vals_clean = vals[np.isfinite(vals)]
print("Min:", vals_clean.min())
print("Max:", vals_clean.max())

def normalize_zscore(ds: xr.Dataset, vars_to_normalize: list, time_dim="valid_time"):
    """Z-score normalize selected variables over a time dimension, ignoring NaNs and handling divide-by-zero."""
    norm_vars = {}
    for var in vars_to_normalize:
        v = ds[var]
        v_mean = v.mean(dim=time_dim, skipna=True)
        v_std = v.std(dim=time_dim, skipna=True)

        # Prevent division by zero
        v_std = xr.where(v_std == 0, np.nan, v_std)

        norm = (v - v_mean) / v_std

        # Preserve original NaN structure
        norm_vars[var] = norm.where(~np.isnan(v))

    return ds.assign(norm_vars)

flood_composite_2020.data_vars

vars_to_norm = ["runoff_mm", "soil_moisture_l1_mm", "precip_mm", "soil_moisture_l2_mm", "potential_evaporation_mm_day"]
flood_composite_2020_zscore = normalize_zscore(flood_composite_2020, vars_to_norm)
flood_composite_2021_zscore = normalize_zscore(flood_composite_2021, vars_to_norm)
flood_composite_04072020_zscore = normalize_zscore(flood_composite_04072020, vars_to_norm)

# Select elevation DataArray
elev = flood_composite_2020["elevation"]
# Compute spatial mean and std (over latitude and longitude dims)
mean_spatial = elev.mean(dim=["latitude", "longitude"])
std_spatial = elev.std(dim=["latitude", "longitude"])
# Broadcast mean and std to match dims of elevation (including time)
# Then apply spatial z-score normalization
elev_norm = (elev - mean_spatial) / std_spatial
# Assign back to dataset or save as new variable
flood_composite_2020_zscore["elevation"] = elev_norm

elev = flood_composite_2021["elevation"]
mean_spatial = elev.mean(dim=["latitude", "longitude"])
std_spatial = elev.std(dim=["latitude", "longitude"])
elev_norm = (elev - mean_spatial) / std_spatial
flood_composite_2021_zscore["elevation"] = elev_norm

elev = flood_composite_04072020["elevation"]
mean_spatial = elev.mean(dim=["latitude", "longitude"])
std_spatial = elev.std(dim=["latitude", "longitude"])
elev_norm = (elev - mean_spatial) / std_spatial
flood_composite_04072020_zscore["elevation"] = elev_norm

# Loop through all data variables with 3 dimensions
for var_name, da in flood_composite_2020.data_vars.items():
    if "valid_time" in da.dims:
        print(f"\nVariable: {var_name}")
        vals = da.values
        vals_clean = vals[np.isfinite(vals)]
        print("  Min:", vals_clean.min())
        print("  Max:", vals_clean.max())
        print("  Mean:", vals_clean.mean())

# Loop through all data variables with 3 dimensions
for var_name, da in flood_composite_2020_zscore.data_vars.items():
    if "valid_time" in da.dims:
        print(f"\nVariable: {var_name}")
        vals = da.values
        vals_clean = vals[np.isfinite(vals)]
        print("  Min:", vals_clean.min())
        print("  Max:", vals_clean.max())
        print("  Mean:", vals_clean.mean())

precip_thresh = 20.0
sm1_thresh = 350.0
sm2_thresh = 350.0
runoff_thresh = 10.0
pe_thresh = -40.0
elev_thresh = 50.0

precip_thresh_zscore = 2.0
sm1_thresh_zscore = 1.5
sm2_thresh_zscore = 1.5
runoff_thresh_zscore = 2.0
pe_thresh_zscore = -1.0
elev_thresh_zscore = 1.0

flood_mask_2020 = (
    (flood_composite_2020["runoff_mm"] > runoff_thresh) &
    (flood_composite_2020["soil_moisture_l1_mm"] > sm1_thresh) &
    (flood_composite_2020["precip_mm"] > precip_thresh) &
    (flood_composite_2020["soil_moisture_l2_mm"] > sm2_thresh) &
    (flood_composite_2020["potential_evaporation_mm_day"] < pe_thresh)
)

flood_mask_2020_zscore = (
    (flood_composite_2020_zscore["runoff_mm"] > runoff_thresh_zscore) &
    (flood_composite_2020_zscore["soil_moisture_l1_mm"] > sm1_thresh_zscore) &
    (flood_composite_2020_zscore["precip_mm"] > precip_thresh_zscore) &
    (flood_composite_2020_zscore["soil_moisture_l2_mm"] > sm2_thresh_zscore) &
    (flood_composite_2020_zscore["potential_evaporation_mm_day"] < pe_thresh_zscore)
)

flood_mask_2021 = (
    (flood_composite_2021["runoff_mm"] > runoff_thresh) &
    (flood_composite_2021["soil_moisture_l1_mm"] > sm1_thresh) &
    (flood_composite_2021["precip_mm"] > precip_thresh) &
    (flood_composite_2021["soil_moisture_l2_mm"] > sm2_thresh) &
    (flood_composite_2021["potential_evaporation_mm_day"] < pe_thresh)
)

flood_mask_2021_zscore = (
    (flood_composite_2021_zscore["runoff_mm"] > runoff_thresh_zscore) &
    (flood_composite_2021_zscore["soil_moisture_l1_mm"] > sm1_thresh_zscore) &
    (flood_composite_2021_zscore["precip_mm"] > precip_thresh_zscore) &
    (flood_composite_2021_zscore["soil_moisture_l2_mm"] > sm2_thresh_zscore) &
    (flood_composite_2021_zscore["potential_evaporation_mm_day"] < pe_thresh_zscore)
)

flood_mask_2020_zscore = flood_mask_2020_zscore.astype(int)
flood_mask_2021_zscore = flood_mask_2021_zscore.astype(int)

flood_composite_2020_zscore = flood_composite_2020_zscore.compute()
flood_composite_2021_zscore = flood_composite_2021_zscore.compute()
flood_composite_04072020_zscore = flood_composite_04072020_zscore.compute()
flood_mask_2020_zscore = flood_mask_2020_zscore.compute()
flood_mask_2021_zscore = flood_mask_2021_zscore.compute()

# Stack features into (time, lat, lon, features)
stacked_fts_2020 = flood_composite_2020_zscore.to_array().transpose("valid_time", "latitude", "longitude", "variable").fillna(-9999)
stacked_fts_2021 = flood_composite_2021_zscore.to_array().transpose("valid_time", "latitude", "longitude", "variable").fillna(-9999)
# Concatenate along the valid_time dimension
stacked_fts = xr.concat([stacked_fts_2020, stacked_fts_2021], dim="valid_time")
stacked_fts_pred = flood_composite_04072020_zscore.to_array().transpose("valid_time", "latitude", "longitude", "variable").fillna(-9999)

X = stacked_fts.data.reshape(-1, stacked_fts.shape[-1])  # shape: (num_pixels, num_features)
X_pred = stacked_fts_pred.data.reshape(-1, stacked_fts_pred.shape[-1])  # shape: (num_pixels, num_features)

# Flatten labels
y_2020 = flood_mask_2020_zscore.data.flatten()
y_2021 = flood_mask_2021_zscore.data.flatten()
# Concatenate into one array
y = np.concatenate([y_2020, y_2021])

#X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=42)
X_train, X_test, y_train, y_test = train_test_split(
    X, y, test_size=0.2, random_state=42, shuffle=True
)

np.unique(y_train), np.unique(y_test)

clf = RandomForestClassifier(n_estimators=100, random_state=42)
clf.fit(X_train, y_train)

y_pred = clf.predict(X_test)
print(classification_report(y_test, y_pred))

X_pred_ = clf.predict(X_pred)

X_pred_.shape

flood_mask_2020_zscore.sel(valid_time="2020-04-07").shape 

output_shape = flood_mask_2020_zscore.sel(valid_time="2020-04-07").shape  # (time, lat, lon)
full_pred = X_pred_.reshape(output_shape)

full_pred.shape

np.unique(full_pred)

full_pred_xr = xr.DataArray(
    data=full_pred,
    coords={
        "valid_time": flood_composite_04072020_zscore.valid_time,
        "latitude": flood_composite_04072020_zscore.latitude,
        "longitude": flood_composite_04072020_zscore.longitude,
    },
    dims=["valid_time", "latitude", "longitude"],
    name="predicted_class"
)

full_pred_ds = xr.Dataset({"predicted_class": full_pred_xr})

full_pred_ds

full_pred_ds["predicted_class"].hvplot(
    x='longitude',
    y='latitude',
    groupby='valid_time',      # time slider
    cmap='Blues',        
    colorbar=True,
    clim=(0, full_pred_ds["predicted_class"].max().item()), 
    title="Predicted Class over Time",
    width=600,
    height=500
)

# Extract the mask for pixels with value 2 (flood)
flood_mask_2 = (flood_mask == 2)

# Find indices of those pixels
y_idx, x_idx = np.where(flood_mask_2)

# Get lat/lon for those pixels
lats = flood_mask.y[y_idx].values
lons = flood_mask.x[x_idx].values

# Calculate pixel resolution (assuming regular grid)
res_lat = np.abs(flood_mask.y[1] - flood_mask.y[0]).item()
res_lon = np.abs(flood_mask.x[1] - flood_mask.x[0]).item()

# Build polygons for each pixel
geoms = [
    box(lon - res_lon/2, lat - res_lat/2, lon + res_lon/2, lat + res_lat/2)
    for lon, lat in zip(lons, lats)
]

# Merge polygons into one boundary
merged_geom = unary_union(geoms)
flood_boundary_gdf = gpd.GeoDataFrame(geometry=[merged_geom], crs="EPSG:4326")

# Plot raster prediction using matplotlib
fig, ax = plt.subplots(figsize=(10, 8))
full_pred_ds["predicted_class"].isel(valid_time=0).plot(ax=ax, cmap="Blues", alpha=0.7)

malampa_gdf.boundary.plot(ax=ax, edgecolor='green', linewidth=2, alpha=0.8, label='Malampa Boundary')

# Plot the flood boundary polygon on top in red
flood_boundary_gdf.boundary.plot(ax=ax, edgecolor='red', linewidth=2, alpha=0.7)


plt.title("Predicted Class with Flood Mask Boundary")
plt.xlabel("Longitude")
plt.ylabel("Latitude")
plt.show()