"""Sentinel-1 SAR flood mapping in Python: the 2022 Pakistan floods around Manchar Lake and Dadu.

Searches the Microsoft Planetary Computer STAC API for radiometrically terrain-corrected
(RTC) Sentinel-1 scenes, streams VV backscatter from Cloud Optimised GeoTIFFs onto a fixed
grid, converts to decibels, removes speckle with a Lee filter, picks a water threshold
with Otsu's method and maps the flood extent and area through the 2022 monsoon.
Run with:  python sar_flood_mapping.py
"""
import platform
from importlib.metadata import version
from pathlib import Path

import geopandas as gpd
import matplotlib

matplotlib.use("Agg")
import matplotlib.dates as mdates
import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
import planetary_computer
import pystac_client
import rasterio
import rasterio.features
from matplotlib.colors import ListedColormap
from pyproj import Transformer
from rasterio.enums import Resampling
from rasterio.transform import from_origin
from rasterio.vrt import WarpedVRT
from rasterio.windows import from_bounds
from scipy import ndimage
from skimage.filters import threshold_otsu

HERE = Path(__file__).parent
OUT = HERE / "output"
OUT.mkdir(exist_ok=True)
plt.rcParams.update({"figure.dpi": 150, "savefig.bbox": "tight", "font.size": 9})
pd.set_option("display.width", 120)

print("=== versions ===")
print(f"python {platform.python_version()} | pystac-client {pystac_client.__version__} | "
      f"planetary-computer {version('planetary-computer')} | rasterio {rasterio.__version__} | "
      f"numpy {np.__version__} | scikit-image {version('scikit-image')}")

STAC_URL = "https://planetarycomputer.microsoft.com/api/stac/v1"
BBOX = (67.45, 26.15, 68.25, 27.05)  # lon/lat: Manchar Lake, Dadu and the Indus in Sindh, Pakistan
SEASON = "2022-06-01/2022-12-31"
ORBIT = 42                            # one relative orbit = identical viewing geometry on every date
BEFORE, AFTER = "2022-06-16", "2022-09-08"
RES = 40                              # analysis resolution in metres
CRS = "EPSG:32642"                    # UTM zone 42N, the native CRS of these scenes

to_utm = Transformer.from_crs("EPSG:4326", CRS, always_xy=True)
xs, ys = to_utm.transform([BBOX[0], BBOX[2], BBOX[0], BBOX[2]], [BBOX[1], BBOX[1], BBOX[3], BBOX[3]])
x0, x1 = np.floor(min(xs) / RES) * RES, np.ceil(max(xs) / RES) * RES
y0, y1 = np.floor(min(ys) / RES) * RES, np.ceil(max(ys) / RES) * RES
WIDTH, HEIGHT = int((x1 - x0) / RES), int((y1 - y0) / RES)
TRANSFORM = from_origin(x0, y1, RES, RES)
PIXEL_KM2 = RES * RES / 1e6


def section(title: str) -> None:
    print(f"\n=== {title} ===")


def read_gamma0(href: str) -> np.ndarray:
    """Warp a remote RTC COG onto the analysis grid (averaging linear power), NaN where empty."""
    with rasterio.open(href) as src, WarpedVRT(src, crs=CRS, transform=TRANSFORM, width=WIDTH, height=HEIGHT,
                                               resampling=Resampling.average, src_nodata=src.nodata,
                                               nodata=np.nan, dtype="float32") as vrt:
        return vrt.read(1)


def to_db(power: np.ndarray) -> np.ndarray:
    with np.errstate(divide="ignore", invalid="ignore"):
        return 10 * np.log10(np.where(power > 0, power, np.nan))


def lee_filter(img: np.ndarray, size: int = 7, enl: float = 4.4) -> np.ndarray:
    """Lee (1980) speckle filter on linear power. enl = equivalent number of looks."""
    nan = np.isnan(img)
    img = np.where(nan, np.nanmedian(img), img)
    mean = ndimage.uniform_filter(img, size)
    var = np.maximum(ndimage.uniform_filter(img * img, size) - mean * mean, 0)
    ci2 = var / np.maximum(mean * mean, 1e-12)        # local coefficient of variation, squared
    w = np.clip(1 - (1 / enl) / np.maximum(ci2, 1e-12), 0, 1)
    out = mean + w * (img - mean)                      # smooth flat areas, keep edges
    out[nan] = np.nan
    return out


# --- Step 1: search the catalogue ----------------------------------------------
catalog = pystac_client.Client.open(STAC_URL, modifier=planetary_computer.sign_inplace)
search = catalog.search(collections=["sentinel-1-rtc"], bbox=BBOX, datetime=SEASON,
                        query={"sat:relative_orbit": {"eq": ORBIT}, "sat:orbit_state": {"eq": "ascending"}})
items = sorted(search.items(), key=lambda i: i.datetime)
section("STAC search")
print(f"{len(items)} scenes from relative orbit {ORBIT} (ascending), {items[0].datetime:%d %b} to "
      f"{items[-1].datetime:%d %b %Y}")
example = next(i for i in items if i.datetime.strftime("%Y-%m-%d") == AFTER)
p = example.properties
print(example.id)
print(f"assets: {', '.join(example.assets)}")
print(f"polarisations: {p['sar:polarizations']} | mode {p['sar:instrument_mode']} | "
      f"pixel {p['sar:pixel_spacing_range']} m | ENL {p['sar:looks_equivalent_number']} | {p['proj:code']}")
print("vv raster:bands:", example.assets["vv"].extra_fields["raster:bands"])
ENL = p["sar:looks_equivalent_number"]

# --- Step 2: speckle at full resolution -------------------------------------------
CHIP = (67.75, 26.70, 67.81, 26.76)  # 6 km box around Dadu town and the floodwater surrounding it
cx, cy = to_utm.transform([CHIP[0], CHIP[2]], [CHIP[1], CHIP[3]])
with rasterio.open(example.assets["vv"].href) as src:
    chip = src.read(1, window=from_bounds(cx[0], cy[0], cx[1], cy[1], src.transform)).astype("float32")
chip[chip == -32768] = np.nan
chip_lee = lee_filter(chip, size=7, enl=ENL)
section("speckle (10 m chip)")
print(f"chip {chip.shape[1]} x {chip.shape[0]} pixels; linear gamma0 {np.nanmin(chip):.4f}..{np.nanmax(chip):.2f}")


def local_std(db: np.ndarray, size: int = 5) -> float:
    """Median of the 5 x 5 moving-window standard deviation: pixel-to-pixel noise, not scene contrast."""
    m = ndimage.uniform_filter(db, size)
    return float(np.nanmedian(np.sqrt(np.maximum(ndimage.uniform_filter(db * db, size) - m * m, 0))))


print(f"pixel-to-pixel noise (median 5x5 std): raw {local_std(to_db(chip)):.2f} dB, Lee {local_std(to_db(chip_lee)):.2f} dB")

# --- Step 3: read every date on the 40 m grid ----------------------------------------
section("reading VV on a 40 m grid")
vv_db, dates = {}, []
for it in items:
    d = it.datetime.strftime("%Y-%m-%d")
    power = read_gamma0(it.assets["vv"].href)
    # Each 40 m pixel averages 16 pixels (~4 independent resolution cells), so ENL rises ~4x.
    vv_db[d] = to_db(lee_filter(power, size=5, enl=ENL * 4))
    dates.append(d)
    print(f"{d}: {np.isfinite(vv_db[d]).mean():.1%} valid, median {np.nanmedian(vv_db[d]):.1f} dB")

# --- Step 4: Otsu threshold on the flood image -------------------------------------------
before, after = vv_db[BEFORE], vv_db[AFTER]
valid = np.isfinite(before) & np.isfinite(after)
thr = threshold_otsu(after[valid])
section("Otsu threshold")
print(f"VV threshold from {AFTER}: {thr:.1f} dB")
print(f"median VV: water-side {np.median(after[valid & (after < thr)]):.1f} dB, "
      f"land-side {np.median(after[valid & (after >= thr)]):.1f} dB")


def water_mask(db: np.ndarray) -> np.ndarray:
    w = np.isfinite(db) & (db < thr)
    w = ndimage.binary_opening(w, iterations=1)                  # drop isolated dark pixels
    lab, n = ndimage.label(w)
    sizes = ndimage.sum(w, lab, index=np.arange(1, n + 1))
    return np.isin(lab, 1 + np.flatnonzero(sizes >= 7))          # keep patches >= ~1 ha


water_before, water_after = water_mask(before) & valid, water_mask(after) & valid
flood = water_after & ~water_before
receded = water_before & ~water_after
section("flood extent")
print(f"valid area: {valid.sum() * PIXEL_KM2:,.0f} km²")
print(f"water on {BEFORE}: {water_before.sum() * PIXEL_KM2:,.0f} km²")
print(f"water on {AFTER}: {water_after.sum() * PIXEL_KM2:,.0f} km²")
print(f"newly flooded: {flood.sum() * PIXEL_KM2:,.0f} km² ({flood.sum() / valid.sum():.1%} of the area)")
print(f"water before but not after: {receded.sum() * PIXEL_KM2:,.0f} km²")

# Change detection without a threshold: how much darker did each pixel get?
diff = after - before
print(f"pixels at least 3 dB darker: {((diff < -3) & valid).sum() * PIXEL_KM2:,.0f} km²; "
      f"of the Otsu flood pixels, {(diff[flood] < -3).mean():.0%} dropped by 3 dB or more")

# --- Step 5: time series -------------------------------------------------------------
rows = []
for d in dates:
    v = np.isfinite(vv_db[d]) & valid
    w = water_mask(vv_db[d]) & v
    rows.append({"date": pd.Timestamp(d), "water_km2": round(w.sum() * PIXEL_KM2),
                 "flooded_km2": round((w & ~water_before).sum() * PIXEL_KM2), "valid_pct": round(v.mean() * 100, 1)})
ts = pd.DataFrame(rows)
section("water area through the 2022 monsoon")
print(ts.assign(date=ts.date.dt.strftime("%Y-%m-%d")).to_string(index=False))
ts.to_csv(OUT / "flood_area_timeseries.csv", index=False)
peak = ts.loc[ts.flooded_km2.idxmax()]
print(f"peak: {peak.date:%d %B %Y}, {peak.flooded_km2:,} km² newly flooded")

# --- Step 6: save the flood extent -----------------------------------------------------
classes = np.zeros((HEIGHT, WIDTH), "uint8")
classes[water_before & water_after] = 1
classes[flood] = 2
classes[~valid] = 255
with rasterio.open(OUT / "flood_extent_2022-09-08.tif", "w", driver="GTiff", height=HEIGHT, width=WIDTH,
                   count=1, dtype="uint8", crs=CRS, transform=TRANSFORM, nodata=255, compress="deflate") as dst:
    dst.write(classes, 1)
shapes = rasterio.features.shapes(flood.astype("uint8"), mask=flood, transform=TRANSFORM)
polys = gpd.GeoDataFrame.from_features([{"geometry": g, "properties": {"class": "flood"}} for g, _ in shapes],
                                       crs=CRS)
polys["area_km2"] = polys.area / 1e6
polys.to_file(OUT / "flood_extent_2022-09-08.gpkg", driver="GPKG")
section("vector output")
print(f"{len(polys):,} flood polygons, total {polys.area_km2.sum():,.0f} km², largest {polys.area_km2.max():,.0f} km²")

# --- Step 7: figures --------------------------------------------------------------------
extent = [0, (x1 - x0) / 1000, 0, (y1 - y0) / 1000]
pretty = {d: pd.Timestamp(d).strftime("%d %b %Y") for d in dates}


def km_axes(ax):
    ax.set_xlabel("km east")
    ax.set_ylabel("km north")


def stretch(db, lo=-25, hi=0):
    return np.clip((db - lo) / (hi - lo), 0, 1)


rgb = np.dstack([stretch(before), stretch(after), stretch(after)])
rgb[~valid] = 1
fig, axes = plt.subplots(1, 3, figsize=(12, 4.6), constrained_layout=True)
for ax, arr, title in ((axes[0], before, f"VV before, {pretty[BEFORE]}"), (axes[1], after, f"VV after, {pretty[AFTER]}")):
    im = ax.imshow(arr, cmap="gray", vmin=-25, vmax=0, extent=extent)
    ax.set_title(title)
    km_axes(ax)
axes[2].imshow(rgb, extent=extent)
axes[2].set_title("Change composite: R = before, G/B = after")
km_axes(axes[2])
fig.colorbar(im, ax=axes[:2], shrink=0.8, label="VV backscatter (dB)", location="bottom", aspect=40)
fig.savefig(OUT / "01-vv-before-after.png")
plt.close(fig)

cext = [0, (cx[1] - cx[0]) / 1000, 0, (cy[1] - cy[0]) / 1000]
fig, axes = plt.subplots(1, 2, figsize=(10, 4.9), constrained_layout=True)
for ax, arr, title in ((axes[0], chip, "Raw VV, 10 m: speckle"), (axes[1], chip_lee, "Lee filter, 7 × 7 window")):
    im = ax.imshow(to_db(arr), cmap="gray", vmin=-25, vmax=0, extent=cext)
    ax.set_title(f"{title} ({pretty[AFTER]})")
    km_axes(ax)
fig.colorbar(im, ax=axes, shrink=0.85, label="VV backscatter (dB)")
fig.savefig(OUT / "02-speckle-lee-filter.png")
plt.close(fig)

fig, ax = plt.subplots(figsize=(8.5, 3.8))
bins = np.linspace(-30, 5, 141)
ax.hist(before[valid], bins=bins, color="#adb5bd", alpha=0.85, label=f"before ({pretty[BEFORE]})")
ax.hist(after[valid], bins=bins, color="#1d6fa3", alpha=0.6, label=f"after ({pretty[AFTER]})")
ax.axvline(thr, color="#e8590c", lw=2, label=f"Otsu threshold {thr:.1f} dB")
ax.set_xlabel("VV backscatter (dB), Lee-filtered, 40 m")
ax.set_ylabel("pixels")
ax.set_title("Flood water adds a second, dark peak to the histogram")
ax.legend(fontsize=8)
ax.grid(axis="y", alpha=0.3)
fig.savefig(OUT / "03-otsu-histogram.png")
plt.close(fig)

fig, ax = plt.subplots(figsize=(8, 8.4), constrained_layout=True)
ax.imshow(before, cmap="gray", vmin=-25, vmax=5, extent=extent)
overlay = np.ma.masked_equal(np.where(classes == 255, 0, classes), 0)
ax.imshow(overlay, cmap=ListedColormap(["#1d6fa3", "#e8590c"]), vmin=1, vmax=2, extent=extent,
          interpolation="nearest")
handles = [plt.Rectangle((0, 0), 1, 1, color=c) for c in ("#1d6fa3", "#e8590c")]
ax.legend(handles, [f"Water before and after ({(classes == 1).sum() * PIXEL_KM2:,.0f} km²)",
                    f"Flooded by {pretty[AFTER]} ({flood.sum() * PIXEL_KM2:,.0f} km²)"], loc="lower right", fontsize=8)
for name, lon, lat in (("Manchar Lake", 67.67, 26.42), ("Dadu", 67.78, 26.73), ("Sehwan", 67.86, 26.42)):
    ex, ey = to_utm.transform(lon, lat)
    ax.annotate(name, ((ex - x0) / 1000, (ey - y0) / 1000), xytext=(6, 6), textcoords="offset points",
                color="white", fontsize=8, fontweight="bold")
    ax.plot((ex - x0) / 1000, (ey - y0) / 1000, "o", ms=3, color="white")
ax.set_title(f"Sentinel-1 flood extent, {pretty[AFTER]} vs {pretty[BEFORE]}")
km_axes(ax)
fig.savefig(OUT / "04-flood-extent-map.png")
plt.close(fig)

fig, ax = plt.subplots(figsize=(8.5, 3.6))
ax.fill_between(ts.date, ts.water_km2, color="#74c0fc", alpha=0.5, label="all open water")
ax.plot(ts.date, ts.flooded_km2, "o-", color="#e8590c", ms=4, label=f"newly flooded vs {pretty[BEFORE]}")
for d, v in zip(ts.date, ts.flooded_km2):
    if d == peak.date:
        ax.annotate(f"{v:,} km²", (d, v), xytext=(8, 4), textcoords="offset points", fontsize=8)
ax.set_ylabel("area (km²)")
ax.set_title(f"Water in the {extent[1]:.0f} × {extent[3]:.0f} km study area, Sentinel-1 orbit {ORBIT}, 2022")
ax.legend(fontsize=8, loc="upper right")
ax.grid(alpha=0.3)
ax.xaxis.set_major_formatter(mdates.DateFormatter("%b"))
fig.savefig(OUT / "05-flood-area-timeseries.png")
plt.close(fig)