"""SAR speckle filtering in Python: boxcar, Lee, refined Lee, enhanced Lee, Frost and Gamma-MAP on Sentinel-1.

Streams a 13 x 11 km chip of Sentinel-1 RTC (VV, gamma0) over Rutland Water, Oakham and the surrounding
farmland in England from the Microsoft Planetary Computer (no account or key), checks the multiplicative
speckle model against the gamma distribution, implements six speckle filters in NumPy/SciPy and scores them
with the equivalent number of looks (ENL) in homogeneous areas, an edge-preservation index (EPI) against a
29-date multi-temporal reference, the ratio image and runtime.
Run with:  python speckle_filters.py
"""
import platform
import time
from pathlib import Path

import matplotlib

matplotlib.use("Agg")
import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
import planetary_computer
import pystac_client
import rasterio
import scipy
from matplotlib.patches import Rectangle
from pyproj import Transformer
from rasterio.windows import from_bounds
from scipy import ndimage, special, stats

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

print("=== versions ===")
print(f"python {platform.python_version()} | numpy {np.__version__} | scipy {scipy.__version__} | "
      f"rasterio {rasterio.__version__} | pystac-client {pystac_client.__version__}")


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


# --- Step 1: Sentinel-1 RTC chips from the Planetary Computer -------------------------
BBOX = (-0.76, 52.61, -0.56, 52.71)   # lon/lat: Rutland Water, Oakham and farmland, East Midlands, England
ORBIT = 154                           # one relative orbit (descending): identical geometry on every date
SEASON = "2025-10-01/2026-03-31"
TARGET = "2026-01-10"                 # the single date we filter

section("Step 1: search and read Sentinel-1 RTC")
catalog = pystac_client.Client.open("https://planetarycomputer.microsoft.com/api/stac/v1",
                                    modifier=planetary_computer.sign_inplace)
items = sorted(catalog.search(collections=["sentinel-1-rtc"], bbox=BBOX, datetime=SEASON,
                              query={"sat:relative_orbit": {"eq": ORBIT}}).items(), key=lambda i: i.datetime)
print(f"{len(items)} scenes from relative orbit {ORBIT}, {items[0].datetime:%d %b %Y} to {items[-1].datetime:%d %b %Y}")
target = next(i for i in items if i.datetime.strftime("%Y-%m-%d") == TARGET)
ENL_NOMINAL = target.properties["sar:looks_equivalent_number"]
print(f"target: {target.id} | {target.properties['platform']} | looks {target.properties['sar:looks_range']} x "
      f"{target.properties['sar:looks_azimuth']} | nominal ENL {ENL_NOMINAL} | {target.properties['proj:code']}")
to_utm = Transformer.from_crs(4326, target.properties["proj:code"], always_xy=True)
xs, ys = to_utm.transform([BBOX[0], BBOX[2]], [BBOX[1], BBOX[3]])


def read_vv(item) -> np.ndarray:
    """Read the chip of linear gamma0 VV at the native 10 m grid (only the needed COG tiles are fetched)."""
    cache = DATA / f"{item.datetime:%Y-%m-%d}_vv.npy"
    if not cache.exists():
        with rasterio.open(item.assets["vv"].href) as src:
            win = from_bounds(xs[0], ys[0], xs[1], ys[1], src.transform).round_offsets().round_lengths()
            a = src.read(1, window=win).astype("float32")
            a[a == src.nodata] = np.nan
        np.save(cache, a)
    return np.load(cache)


t0 = time.perf_counter()
stack = np.stack([read_vv(i) for i in items])
img = stack[[i.id for i in items].index(target.id)]
ref = np.nanmean(stack, axis=0)        # multi-temporal multilook: nearly speckle-free, edges kept
print(f"read {len(items)} chips of {img.shape[0]} x {img.shape[1]} pixels (10 m) in {time.perf_counter() - t0:.0f} s; "
      f"NaN pixels in target: {int(np.isnan(img).sum())}")
print(f"target gamma0 VV: min {np.nanmin(img):.4f}, median {np.nanmedian(img):.4f}, max {np.nanmax(img):.2f} (linear)")


def db(x):
    with np.errstate(divide="ignore", invalid="ignore"):
        return 10 * np.log10(x)


# --- Step 2: find homogeneous areas in the reference and test the speckle model ---------
section("Step 2: the multiplicative speckle model")
W = 21                                      # 210 m homogeneous windows
m = ndimage.uniform_filter(ref, W)
cv_ref = np.sqrt(np.maximum(ndimage.uniform_filter(ref * ref, W) - m * m, 0)) / m
cv_ref[: W, :] = cv_ref[-W:, :] = cv_ref[:, :W] = cv_ref[:, -W:] = np.inf
m_db = db(m)
windows, kinds = [], []
taken = np.zeros_like(img, bool)
# the four most uniform windows on open water (wind-roughened, well above the noise floor) and four on land
for kind, ok in (("water", m_db < -13), ("land", m_db > -11)):
    n = 0
    for flat in np.argsort(np.where(ok, cv_ref, np.inf), axis=None):
        r, c = np.unravel_index(flat, img.shape)
        if taken[r - W:r + W + 1, c - W:c + W + 1].any():
            continue
        windows.append((r - W // 2, c - W // 2))
        kinds.append(kind)
        taken[r - W // 2:r + W // 2 + 1, c - W // 2:c + W // 2 + 1] = True
        n += 1
        if n == 4:
            break
sl = [np.s_[r:r + W, c:c + W] for r, c in windows]


def enl(x: np.ndarray, kind: str | None = None) -> float:
    """Equivalent number of looks: mean^2 / variance of intensity, averaged over the homogeneous windows."""
    return float(np.mean([np.nanmean(x[s]) ** 2 / np.nanvar(x[s]) for s, k in zip(sl, kinds) if kind in (None, k)]))


for k, s in enumerate(sl):
    print(f"window {k + 1} ({kinds[k]:5s}): rows {s[0].start}-{s[0].stop - 1}, cols {s[1].start}-{s[1].stop - 1} | reference CV "
          f"{cv_ref[windows[k][0] + W // 2, windows[k][1] + W // 2]:.3f} | mean {db(np.nanmean(img[s])):.1f} dB | "
          f"ENL {np.nanmean(img[s]) ** 2 / np.nanvar(img[s]):.2f}")
enl_raw = enl(img)
print(f"ENL of the raw target image: {enl_raw:.2f} (nominal {ENL_NOMINAL}); speckle CV {1 / np.sqrt(enl_raw):.3f}; "
      f"reference ENL {enl(ref):.0f}")
norm = np.concatenate([(img[s] / np.nanmean(img[s])).ravel() for s in sl])
L_fit, _, scale_fit = stats.gamma.fit(norm, floc=0)
print(f"normalised intensity in windows: mean {norm.mean():.3f}, variance {norm.var():.3f}; "
      f"gamma fit shape L = {L_fit:.2f}, scale {scale_fit:.3f}")
ks_fit = stats.kstest(norm, "gamma", args=(L_fit, 0, 1 / L_fit)).statistic
ks_exp = stats.kstest(norm, "expon").statistic
print(f"Kolmogorov-Smirnov distance: gamma(L={L_fit:.2f}) {ks_fit:.3f}; single-look exponential {ks_exp:.3f}")
mean_db_of_px = np.mean([np.nanmean(db(img[s])) for s in sl])
db_of_mean = np.mean([db(np.nanmean(img[s])) for s in sl])
theory = 10 / np.log(10) * (special.digamma(L_fit) - np.log(L_fit))
print(f"mean of dB pixels {mean_db_of_px:.2f} dB vs dB of mean power {db_of_mean:.2f} dB: "
      f"bias {mean_db_of_px - db_of_mean:+.2f} dB (theory for L={L_fit:.2f}: {theory:+.2f} dB)")


# --- Step 3: the filters ----------------------------------------------------------------
section("Step 3: speckle filters")


def moments(x, size):
    mean = ndimage.uniform_filter(x, size)
    var = np.maximum(ndimage.uniform_filter(x * x, size) - mean * mean, 0)
    return mean, var


def fill_nan(x):
    return np.where(np.isnan(x), np.nanmedian(x), x)


def boxcar(x, size=7, enl=None):
    return ndimage.uniform_filter(fill_nan(x), size)


def lee(x, size=7, enl=4.4):
    """Lee (1980): out = mean + k (x - mean), k from the local variance in excess of pure speckle."""
    x = fill_nan(x)
    mean, var = moments(x, size)
    cu2 = 1 / enl                                   # speckle variance (coefficient of variation squared)
    var_x = np.maximum((var - mean ** 2 * cu2) / (1 + cu2), 0)
    k = var_x / np.maximum(var, 1e-12)
    return mean + k * (x - mean)


def enhanced_lee(x, size=7, enl=4.4, damping=1.0):
    """Enhanced Lee (Lopes et al. 1990): full smoothing below Cu, none above Cmax, exponential in between."""
    x = fill_nan(x)
    mean, var = moments(x, size)
    ci = np.sqrt(var) / np.maximum(mean, 1e-12)
    cu, cmax = 1 / np.sqrt(enl), np.sqrt(1 + 2 / enl)
    c = np.clip(ci, cu, cmax - 1e-6)
    w = np.exp(-damping * (c - cu) / (cmax - c))     # 1 in homogeneous areas, 0 at point targets
    out = mean * w + x * (1 - w)
    return np.where(ci >= cmax, x, out)


def gamma_map(x, size=7, enl=4.4):
    """Gamma-MAP (Lopes et al. 1990): maximum a posteriori estimate with a gamma-distributed scene."""
    x = fill_nan(x)
    mean, var = moments(x, size)
    ci2 = var / np.maximum(mean ** 2, 1e-12)
    cu2, cmax2 = 1 / enl, 1 + 2 / enl
    with np.errstate(divide="ignore", invalid="ignore"):
        alpha = (1 + cu2) / (ci2 - cu2)
        b = alpha - enl - 1
        out = (b * mean + np.sqrt(np.maximum(mean ** 2 * b ** 2 + 4 * alpha * enl * mean * x, 0))) / (2 * alpha)
    return np.where(ci2 <= cu2, mean, np.where(ci2 >= cmax2, x, out))


def frost(x, size=7, enl=None, damping=2.0):
    """Frost (1982): exponentially weighted mean, with a decay that grows with local heterogeneity."""
    x = fill_nan(x)
    mean, var = moments(x, size)
    a = damping * var / np.maximum(mean ** 2, 1e-12)
    h = size // 2
    pad = np.pad(x, h, mode="reflect")
    num, den = np.zeros_like(x), np.zeros_like(x)
    for dy in range(-h, h + 1):
        for dx in range(-h, h + 1):
            w = np.exp(-a * np.hypot(dy, dx))
            num += w * pad[h + dy:h + dy + x.shape[0], h + dx:h + dx + x.shape[1]]
            den += w
    return num / den


def _refined_masks():
    i, j = np.mgrid[0:7, 0:7]
    return {"N": i <= 3, "S": i >= 3, "W": j <= 3, "E": j >= 3,
            "NW": i + j <= 6, "SE": i + j >= 6, "NE": j >= i, "SW": j <= i}


def refined_lee(x, enl=4.4, size=7):
    """Refined Lee (Lee 1981): Lee statistics from the edge-aligned half of a 7 x 7 window."""
    x = fill_nan(x)
    m3 = ndimage.uniform_filter(x, 3)
    pad = np.pad(m3, 2, mode="reflect")
    s = [[pad[2 + 2 * (r - 1):2 + 2 * (r - 1) + x.shape[0], 2 + 2 * (c - 1):2 + 2 * (c - 1) + x.shape[1]]
          for c in range(3)] for r in range(3)]                           # 3 x 3 grid of 3 x 3 sub-window means
    grads = np.stack([abs(s[0][1] - s[2][1]), abs(s[1][0] - s[1][2]),   # vertical, horizontal
                      abs(s[0][0] - s[2][2]), abs(s[0][2] - s[2][0])])  # two diagonals
    pairs = [("N", "S", s[0][1], s[2][1]), ("W", "E", s[1][0], s[1][2]),
             ("NW", "SE", s[0][0], s[2][2]), ("NE", "SW", s[0][2], s[2][0])]
    best = np.argmax(grads, axis=0)
    masks = _refined_masks()
    mean, var = np.empty_like(x), np.empty_like(x)
    for d, (a, b, sa, sb) in enumerate(pairs):
        side_a = abs(s[1][1] - sa) <= abs(s[1][1] - sb)                 # the centre belongs to the closer side
        for name, sel in ((a, side_a), (b, ~side_a)):
            k = masks[name].astype("float32") / masks[name].sum()
            mu = ndimage.correlate(x, k, mode="reflect")
            v = np.maximum(ndimage.correlate(x * x, k, mode="reflect") - mu * mu, 0)
            pick = (best == d) & sel
            mean[pick], var[pick] = mu[pick], v[pick]
    cu2 = 1 / enl
    var_x = np.maximum((var - mean ** 2 * cu2) / (1 + cu2), 0)
    return mean + var_x / np.maximum(var, 1e-12) * (x - mean)


FILTERS = {"Boxcar 7×7": boxcar, "Lee": lee, "Refined Lee": refined_lee, "Enhanced Lee": enhanced_lee,
           "Frost": frost, "Gamma-MAP": gamma_map}
ENL_IN = round(enl_raw, 1)    # use the measured ENL of this image, not the nominal value
print(f"filtering with a 7 x 7 window and ENL = {ENL_IN}")
results, times = {}, {}
for name, f in FILTERS.items():
    f(img[:64, :64], enl=ENL_IN)                                         # warm-up
    best_t = np.inf
    for _ in range(3):
        t0 = time.perf_counter()
        out = f(img, enl=ENL_IN)
        best_t = min(best_t, time.perf_counter() - t0)
    out[np.isnan(img)] = np.nan
    results[name], times[name] = out.astype("float32"), best_t
    print(f"{name:13s} {best_t * 1000:7.0f} ms  ({img.size / best_t / 1e6:5.1f} Mpixel/s)")

# --- Step 4: score the filters ---------------------------------------------------------
section("Step 4: ENL, edge preservation and the ratio image")
ref_db = db(ref)
gy, gx = ndimage.sobel(ref_db, 0), ndimage.sobel(ref_db, 1)
edges = np.hypot(gx, gy) > np.nanpercentile(np.hypot(gx, gy), 90)     # strongest 10% of reference edges
print(f"edge pixels (top 10% Sobel gradient of the reference): {int(edges.sum()):,}")


def highpass(x_db: np.ndarray) -> np.ndarray:
    x_db = np.nan_to_num(x_db)
    return x_db - ndimage.gaussian_filter(x_db, 3)


def epi(x: np.ndarray) -> float:
    """Edge-preservation index (after Sattar et al. 1997): correlation of the high-pass detail of the image
    and of the speckle-free reference, over the edge pixels."""
    a, b = highpass(db(x))[edges], highpass(ref_db)[edges]
    a, b = a - a.mean(), b - b.mean()
    return float((a * b).sum() / np.sqrt((a * a).sum() * (b * b).sum()))


rows = [{"filter": "None (raw)", "ENL": enl_raw, "ENL water": enl(img, "water"), "ENL land": enl(img, "land"), "EPI": epi(img), "ratio mean": 1.0, "ratio std": 0.0,
         "RMSE vs ref (dB)": float(np.sqrt(np.nanmean((db(img) - ref_db) ** 2))), "ms": 0.0}]
for name, out in results.items():
    ratio = img / out
    valid = np.isfinite(ratio)
    rows.append({"filter": name, "ENL": enl(out), "ENL water": enl(out, "water"), "ENL land": enl(out, "land"), "EPI": epi(out), "ratio mean": float(np.mean(ratio[valid])),
                 "ratio std": float(np.std(ratio[valid])),
                 "RMSE vs ref (dB)": float(np.sqrt(np.nanmean((db(out) - ref_db) ** 2))), "ms": times[name] * 1000})
score = pd.DataFrame(rows).set_index("filter")
print(score.round({"ENL": 1, "ENL water": 1, "ENL land": 1, "EPI": 3, "ratio mean": 3, "ratio std": 3, "RMSE vs ref (dB)": 2, "ms": 0}).to_string())
print(f"(pure speckle would give a ratio image with mean 1 and std {1 / np.sqrt(enl_raw):.3f})")

# Filtering in dB instead of linear power: the same boxcar, two answers
box_db = 10 ** (boxcar(db(img)) / 10)
for lab, x in [("raw image", img), ("boxcar in linear", results["Boxcar 7×7"]), ("boxcar in dB", box_db)]:
    print(f"{lab:18s}: mean over homogeneous windows {np.mean([db(np.nanmean(x[s])) for s in sl]):6.2f} dB; "
          f"whole-chip mean power {db(np.nanmean(x)):6.2f} dB")

# --- Figures --------------------------------------------------------------------------
section("Figures")
VMIN, VMAX = -22, 2
ZOOM = np.s_[300:620, 40:560]                    # Oakham, the western lagoons and the reservoir shore
TRANSECT_ROW, TRANSECT_COLS = 545, (330, 700)    # west to east across a lagoon, the reservoir edge and open water


def show(ax, x, title):
    ax.imshow(db(x), cmap="gray", vmin=VMIN, vmax=VMAX, interpolation="nearest")
    ax.set_title(title, fontsize=9)
    ax.set_axis_off()


# 1. Study area
fig, axes = plt.subplots(1, 2, figsize=(11, 4.9))
show(axes[0], img, f"Sentinel-1 VV γ⁰, {target.datetime:%d %b %Y} (raw, 10 m)")
show(axes[1], ref, f"Mean of {len(items)} dates, Oct 2025 – Mar 2026 (reference)")
for ax in axes:
    for k, (r, c) in enumerate(windows):
        ax.add_patch(Rectangle((c, r), W, W, fill=False, ec="#fdae61", lw=1.2))
        ax.text(c + W + 3, r + W // 2, str(k + 1), color="#fdae61", fontsize=7, va="center", fontweight="bold")
    z = ZOOM
    ax.add_patch(Rectangle((z[1].start, z[0].start), z[1].stop - z[1].start, z[0].stop - z[0].start,
                           fill=False, ec="#2b83ba", lw=1.2, ls="--"))
    ax.plot(TRANSECT_COLS, [TRANSECT_ROW] * 2, color="#d7191c", lw=1.2)
axes[0].text(0.01, 0.01, "orange: homogeneous windows | blue: zoom | red: transect", transform=axes[0].transAxes,
             color="white", fontsize=7, va="bottom", bbox=dict(fc="black", alpha=0.5, lw=0))
fig.savefig(OUT / "01-study-area-raw-and-reference.png")
plt.close(fig)

# 2. Speckle statistics
fig, axes = plt.subplots(1, 2, figsize=(10, 3.8))
ax = axes[0]
xg = np.linspace(0, 4, 400)
ax.hist(norm, bins=120, range=(0, 4), density=True, color="#abd9e9", label="pixels / window mean")
ax.plot(xg, stats.gamma.pdf(xg, L_fit, scale=1 / L_fit), color="#2c7bb6", lw=1.8, label=f"gamma, L = {L_fit:.2f} (fit)")
ax.plot(xg, stats.expon.pdf(xg), color="#d7191c", lw=1.2, ls="--", label="exponential (single look)")
ax.set(xlabel="Normalised intensity I / mean", ylabel="Density", title="Speckle in homogeneous windows: linear power")
ax.legend(fontsize=8)
ax = axes[1]
ndb = db(norm)
ax.hist(ndb, bins=100, range=(-8, 6), density=True, color="#fdae61")
ax.axvline(0, color="black", lw=1, label="mean power (0 dB)")
ax.axvline(ndb.mean(), color="#d7191c", lw=1.5, ls="--", label=f"mean of dB values ({ndb.mean():+.2f} dB)")
ax.set(xlabel="Normalised intensity (dB)", ylabel="Density", title="The same pixels in decibels: skewed and biased")
ax.legend(fontsize=8)
fig.savefig(OUT / "02-speckle-statistics.png")
plt.close(fig)

# 3. Filter comparison (zoom)
fig, axes = plt.subplots(2, 4, figsize=(12, 4.6))
panels = [("Raw", img)] + list(results.items()) + [("29-date reference", ref)]
for ax, (name, x) in zip(axes.ravel(), panels):
    sub = f"ENL {score.loc[name, 'ENL']:.0f}, EPI {score.loc[name, 'EPI']:.2f}" if name in score.index else ""
    show(ax, x[ZOOM], f"{name}\n{sub}" if sub else f"{name}\n(ENL {enl(x):.0f})")
fig.subplots_adjust(wspace=0.03, hspace=0.25)
fig.savefig(OUT / "03-filter-comparison.png")
plt.close(fig)

# 4. Transect across edges
fig, ax = plt.subplots(figsize=(9, 3.8))
cols = np.arange(*TRANSECT_COLS)
dist = (cols - cols[0]) * 10 / 1000
ax.plot(dist, db(img[TRANSECT_ROW, cols]), color="#bbb", lw=0.8, label="Raw")
colours = {"Boxcar 7×7": "#fdae61", "Lee": "#abdda4", "Refined Lee": "#2b83ba", "Gamma-MAP": "#d7191c"}
for name, col in colours.items():
    ax.plot(dist, db(results[name][TRANSECT_ROW, cols]), color=col, lw=1.4, label=name)
ax.plot(dist, ref_db[TRANSECT_ROW, cols], color="black", lw=1.4, ls="--", label="29-date reference")
ax.set(xlabel="Distance along transect, west to east (km)", ylabel="VV γ⁰ (dB)",
       title="A transect across a lagoon, the reservoir shore and open water")
ax.set_ylim(top=0)
ax.legend(fontsize=7, ncol=3, loc="upper right")
ax.grid(alpha=0.3)
fig.savefig(OUT / "04-edge-transect.png")
plt.close(fig)

# 5. Trade-off: ENL vs EPI, marker size = runtime
fig, ax = plt.subplots(figsize=(6.5, 4.4))
pal = ["#555555", "#fdae61", "#abdda4", "#2b83ba", "#5e3c99", "#1a9641", "#d7191c"]
for (name, r), col in zip(score.iterrows(), pal):
    ax.scatter(r["ENL"], r["EPI"], s=30 + r["ms"] / 4, color=col, edgecolor="black", lw=0.5, zorder=3)
    off = {"Frost": (-75, -40), "Gamma-MAP": (-35, -70), "Lee": (-60, 22), "Enhanced Lee": (5, -100)}.get(name)
    ax.annotate(f"{name} ({r['ms']:.0f} ms)" if r["ms"] else name, (r["ENL"], r["EPI"]), xytext=off or (6, 4),
                textcoords="offset points", fontsize=8,
                arrowprops=dict(arrowstyle="-", color="#888", lw=0.6) if off else None)
ax.set(xlabel="ENL in homogeneous windows (higher = smoother)", ylabel="Edge-preservation index (higher = sharper)",
       title="Speckle reduction against edge preservation (7 × 7 window)")
ax.set_xlim(0, score["ENL"].max() * 1.35)
ax.set_ylim(score["EPI"].min() - 0.05, score["EPI"].max() + 0.06)
ax.grid(alpha=0.3)
fig.savefig(OUT / "05-enl-vs-epi.png")
plt.close(fig)

# 6. Ratio images: structure that a filter removed shows up as visible edges
fig, axes = plt.subplots(2, 3, figsize=(11, 5))
for ax, (name, out) in zip(axes.ravel(), results.items()):
    ax.imshow((img / out)[ZOOM], cmap="gray", vmin=0, vmax=2.5, interpolation="nearest")
    ax.set_title(f"{name}: raw / filtered (std {score.loc[name, 'ratio std']:.2f})", fontsize=9)
    ax.set_axis_off()
fig.subplots_adjust(wspace=0.03, hspace=0.18)
fig.savefig(OUT / "06-ratio-images.png")
plt.close(fig)
print("figures written:", ", ".join(sorted(p.name for p in OUT.glob("*.png"))))