"""Crop monitoring with Sentinel-1 SAR time series: winter wheat, oilseed rape and more in South Cambridgeshire.

Gets field boundaries from OpenStreetMap (landuse=farmland), labels them with the Rural Payments Agency's
Crop Map of England (CROME) 2024, streams a full season of Sentinel-1 RTC backscatter (VV and VH) from the
Microsoft Planetary Computer, averages it per field to beat speckle, and plots the VV, VH and VH/VV ratio
signatures of each crop. It then counts how many cloud-free Sentinel-2 views the same fields got, compares
NDVI with the radar ratio, and tests how well the radar time series alone separates the crops.
Run with:  python s1_crop_monitoring.py
"""
import platform
import time
from concurrent.futures import ThreadPoolExecutor
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 osmnx as ox
import pandas as pd
import planetary_computer
import pystac_client
import rasterio
import rasterio.features
import requests
from matplotlib.patches import Patch
from rasterio.enums import Resampling
from rasterio.transform import from_origin
from rasterio.vrt import WarpedVRT
from scipy import ndimage
from shapely.geometry import box

HERE = Path(__file__).parent
OUT, DATA, CACHE = HERE / "output", HERE / "data", HERE / "cache"
for d in (OUT, DATA, CACHE):
    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()} | pystac-client {pystac_client.__version__} | "
      f"planetary-computer {version('planetary-computer')} | rasterio {rasterio.__version__} | "
      f"osmnx {ox.__version__} | numpy {np.__version__} | pandas {pd.__version__}")


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


# --- Step 1: field boundaries from OpenStreetMap ----------------------------------------------
BBOX = (-0.12, 52.08, 0.02, 52.17)    # lon/lat: arable South Cambridgeshire, between Royston and Cambridge
CRS = "EPSG:32630"                    # UTM zone 30N: the native grid of these Sentinel-1 and Sentinel-2 scenes
INSET_M = 20                          # shrink each field by 20 m to drop hedges, tracks and mixed edge pixels

fields_file = DATA / "osm_farmland.gpkg"
if not fields_file.exists():
    ox.settings.cache_folder = CACHE / "osmnx"
    osm = None
    for attempt in range(5):   # public Overpass servers are often busy: retry with backoff and a mirror
        for url in ("https://overpass-api.de/api", "https://overpass.kumi.systems/api"):
            ox.settings.overpass_url = url
            try:
                osm = ox.features_from_bbox(BBOX, {"landuse": "farmland"})
                break
            except Exception as exc:
                print(f"{url} failed: {str(exc).splitlines()[0][:70]}")
        if osm is not None:
            break
        time.sleep(30 * (attempt + 1))
    osm = osm[osm.geometry.type.isin(["Polygon", "MultiPolygon"])].reset_index()
    osm[["id", "geometry"]].to_file(fields_file, driver="GPKG")
fields = gpd.read_file(fields_file).to_crs(CRS)
aoi_utm = gpd.GeoSeries([box(*BBOX)], crs=4326).to_crs(CRS).iloc[0]
fields = fields[fields.within(aoi_utm)].copy()
fields["area_ha"] = fields.area / 1e4
fields["core"] = fields.buffer(-INSET_M)
fields = fields[(fields["area_ha"] >= 3) & (fields["area_ha"] <= 100) & ~fields["core"].is_empty].reset_index(drop=True)
section("OpenStreetMap farmland polygons")
print(f"{len(fields)} polygons of 3-100 ha inside the study area, median {fields['area_ha'].median():.1f} ha, "
      f"total {fields['area_ha'].sum():,.0f} ha")

# --- Step 2: crop labels from the Crop Map of England (CROME) 2024 ----------------------------
CROME = ("https://environment.data.gov.uk/spatialdata/crop-map-of-england-2024/ogc/features/v1/"
         "collections/Crop_Map_of_England_2024_Cambridgeshire/items")
CROPS = {"AC66": "Winter wheat", "AC67": "Winter oilseed rape", "AC63": "Winter barley", "AC01": "Spring barley",
         "AC17": "Maize", "AC03": "Sugar beet", "LG20": "Winter field beans", "PG01": "Grass",
         "TG01": "Temporary grass", "FA01": "Fallow"}
MIN_HEX, MIN_PURITY = 10, 0.8          # a field needs >= 10 CROME hexagons, and >= 80% of them must agree

crome_file = DATA / "crome_2024.gpkg"
if not crome_file.exists():
    feats, url = [], CROME
    params = {"f": "json", "limit": 10000, "bbox": ",".join(map(str, BBOX))}
    while url:                                  # OGC API Features: follow the "next" links page by page
        r = requests.get(url, params=params, timeout=300)
        r.raise_for_status()
        page = r.json()
        feats += page["features"]
        url, params = next((lk["href"] for lk in page.get("links", []) if lk["rel"] == "next"), None), None
    gpd.GeoDataFrame.from_features(feats, crs=4326)[["cromeid", "lucode", "prob", "geometry"]].to_file(crome_file)
hexes = gpd.read_file(crome_file).to_crs(CRS)
hexes["geometry"] = hexes.centroid
joined = gpd.sjoin(hexes, fields.set_geometry("core")[["core"]], predicate="within")
stats = joined.groupby("index_right")["lucode"].agg(n_hex="size", lucode=lambda s: s.mode().iloc[0],
                                                     purity=lambda s: s.value_counts(normalize=True).iloc[0])
fields = fields.join(stats)
fields["crop"] = fields["lucode"].map(CROPS).fillna("Other")
labelled = fields[(fields["n_hex"] >= MIN_HEX) & (fields["purity"] >= MIN_PURITY)].reset_index(drop=True)
labelled["fid"] = np.arange(1, len(labelled) + 1)
section("CROME 2024 labels")
print(f"{len(hexes):,} CROME hexagons; {fields['n_hex'].notna().sum()} polygons contain hexagons")
print(f"{len(labelled)} fields are >= {MIN_PURITY:.0%} one crop (the rest are mixed, or the OSM polygon spans several fields)")
print(labelled.groupby("crop")["area_ha"].agg(fields="size", median_ha="median").sort_values("fields", ascending=False)
      .round(1).to_string())

# --- Step 3: one analysis grid at 10 m -------------------------------------------------------
RES = 10
x0, y0, x1, y1 = gpd.GeoSeries([box(*BBOX)], crs=4326).to_crs(CRS).total_bounds
x0, y0, x1, y1 = (np.floor(x0 / RES) * RES, np.floor(y0 / RES) * RES, np.ceil(x1 / RES) * RES, np.ceil(y1 / RES) * RES)
WIDTH, HEIGHT = int((x1 - x0) / RES), int((y1 - y0) / RES)
TRANSFORM = from_origin(x0, y1, RES, RES)
labels = rasterio.features.rasterize(zip(labelled["core"], labelled["fid"]), out_shape=(HEIGHT, WIDTH),
                                     transform=TRANSFORM, fill=0, dtype="int32")
FIDS = labelled["fid"].to_numpy()
npix = ndimage.sum(np.ones_like(labels), labels, FIDS)
section("analysis grid")
print(f"{WIDTH} x {HEIGHT} pixels of {RES} m ({(x1 - x0) / 1000:.1f} x {(y1 - y0) / 1000:.1f} km), {CRS}")
print(f"field cores: median {np.median(npix):.0f} pixels per field (min {npix.min():.0f}, max {npix.max():.0f})")

# --- Step 4: a season of Sentinel-1 RTC from one orbit ----------------------------------------
STAC_URL = "https://planetarycomputer.microsoft.com/api/stac/v1"
SEASON = "2023-09-01/2024-09-30"      # the 2024 harvest year: winter crops are sown in autumn 2023
ORBIT = 59                            # one relative orbit (ascending, ~18:00 UTC): identical geometry every 12 days
GDAL_ENV = dict(GDAL_DISABLE_READDIR_ON_OPEN="EMPTY_DIR", GDAL_HTTP_MAX_RETRY=5, GDAL_HTTP_RETRY_DELAY=2,
                GDAL_HTTP_MERGE_CONSECUTIVE_RANGES="YES")

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}})
s1_items = sorted(search.items(), key=lambda i: i.datetime)
section("Sentinel-1 STAC search")
print(f"{len(s1_items)} scenes from relative orbit {ORBIT} ({s1_items[0].properties['sat:orbit_state']}), "
      f"{s1_items[0].datetime:%d %b %Y} to {s1_items[-1].datetime:%d %b %Y}")
print(f"platforms: {pd.Series([i.properties['platform'] for i in s1_items]).value_counts().to_dict()}")
print(f"example: {s1_items[0].id} | polarisations {s1_items[0].properties['sar:polarizations']}")


def read_warped(href: str, resampling=Resampling.average, nodata=None, dtype="float32") -> np.ndarray:
    """Stream a remote COG straight onto the analysis grid."""
    for attempt in range(4):
        try:
            with rasterio.Env(**GDAL_ENV), rasterio.open(planetary_computer.sign(href)) as src, \
                    WarpedVRT(src, crs=CRS, transform=TRANSFORM, width=WIDTH, height=HEIGHT, resampling=resampling,
                              src_nodata=src.nodata if nodata is None else nodata, nodata=0, dtype=dtype) as vrt:
                return vrt.read(1)
        except rasterio.errors.RasterioIOError:
            time.sleep(5 * (attempt + 1))
    raise RuntimeError(f"could not read {href}")


def field_means(img: np.ndarray) -> np.ndarray:
    """Mean of the valid (non-NaN) pixels in each field core. Use it on linear power, never on dB."""
    ok = np.isfinite(img)
    s = ndimage.sum(np.where(ok, img, 0), labels, FIDS)
    n = ndimage.sum(ok, labels, FIDS)
    with np.errstate(invalid="ignore", divide="ignore"):
        return np.where(n >= 0.8 * npix, s / n, np.nan)   # need 80% of the field core covered


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


# Keep three full images (late autumn, early spring, early summer) for the maps.
s1_days = pd.DatetimeIndex([i.datetime for i in s1_items]).tz_localize(None)
KEEP_DATES = sorted({s1_days[np.abs(s1_days - pd.Timestamp(t)).argmin()].strftime("%Y-%m-%d")
                     for t in ("2023-11-15", "2024-03-25", "2024-06-15")})


def s1_scene(item) -> dict:
    out = {}
    cache = CACHE / f"{item.id}.npz"
    if cache.exists():
        arr = np.load(cache)
        return {"date": item.datetime, **{k: arr[k] for k in arr.files}}
    for pol in ("vv", "vh"):
        g0 = read_warped(item.assets[pol].href)
        g0[g0 <= 0] = np.nan                     # nodata (and the odd non-positive value) -> NaN
        out[pol] = field_means(g0)
        if item.datetime.strftime("%Y-%m-%d") in KEEP_DATES:
            out[f"{pol}_img"] = g0.astype("float32")
    np.savez_compressed(cache, **out)
    return {"date": item.datetime, **out}


t0 = time.time()
with ThreadPoolExecutor(max_workers=8) as pool:
    scenes = list(pool.map(s1_scene, s1_items))
print(f"read and field-averaged {len(scenes)} scenes x 2 polarisations in {time.time() - t0:.0f} s")
S1_DATES = pd.DatetimeIndex([s["date"] for s in scenes]).tz_localize(None).normalize()
VV = np.array([s["vv"] for s in scenes])             # dates x fields, linear power
VH = np.array([s["vh"] for s in scenes])

# --- Step 5: speckle, and why we average per field --------------------------------------------
section("speckle: single pixels versus field averages")
imgs = {s["date"].strftime("%Y-%m-%d"): s["vh_img"] for s in scenes if "vh_img" in s}
MID = KEEP_DATES[1]
vh_img = imgs[MID]
inside = labels > 0
pix_db = db(vh_img[inside])
lab_in = labels[inside]
within_std = pd.Series(pix_db).groupby(lab_in).std()
rng = np.random.default_rng(0)
half = rng.random(inside.sum()) < 0.5          # split each field's pixels at random into two halves
lin = vh_img[inside]
mean_a = pd.Series(lin[half]).groupby(lab_in[half]).mean()
mean_b = pd.Series(lin[~half]).groupby(lab_in[~half]).mean()
half_diff = (db(mean_a) - db(mean_b)).dropna()
print(f"VH on {MID}: median within-field pixel std {within_std.median():.2f} dB")
print(f"split-half test: std of (half A - half B) = {half_diff.std():.3f} dB, "
      f"so a whole-field mean is good to about {half_diff.std() / 2:.3f} dB")

# --- Step 6: crop signatures ----------------------------------------------------------------
vv_db, vh_db = db(VV), db(VH)
ratio_db = vh_db - vv_db                       # VH/VV ratio in dB = VH(dB) - VV(dB)
long = pd.DataFrame({
    "date": np.repeat(S1_DATES, len(FIDS)), "fid": np.tile(FIDS, len(S1_DATES)),
    "vv_db": vv_db.ravel(), "vh_db": vh_db.ravel(), "ratio_db": ratio_db.ravel()})
long = long.merge(labelled[["fid", "id", "crop", "area_ha"]].rename(columns={"id": "osm_id"}), on="fid")
long.round(3).to_csv(OUT / "field_backscatter_timeseries.csv", index=False)
labelled[["fid", "id", "crop", "lucode", "n_hex", "purity", "area_ha", "geometry"]].to_file(
    OUT / "fields_crome_2024.gpkg", driver="GPKG")
print(f"wrote {len(long):,} field-date rows to field_backscatter_timeseries.csv and {len(labelled)} fields to "
      "fields_crome_2024.gpkg")

SHOW = ["Winter wheat", "Winter oilseed rape", "Spring barley", "Maize"]
COLOURS = {"Winter wheat": "#c77c02", "Winter oilseed rape": "#2b8a3e", "Spring barley": "#1c7ed6",
           "Maize": "#862e9c", "Winter barley": "#e8590c", "Sugar beet": "#c2255c", "Winter field beans": "#5c940d",
           "Grass": "#868e96", "Temporary grass": "#adb5bd", "Fallow": "#ced4da", "Other": "#e9ecef"}
med = long.groupby(["crop", "date"])[["vv_db", "vh_db", "ratio_db"]].median()


def window_mean(crop, var, start, end):
    s = med.loc[crop, var]
    return s[(s.index >= start) & (s.index <= end)].mean()


section("crop signatures (median of field means, dB)")
rows = []
for crop in [c for c in COLOURS if c in med.index.get_level_values(0)]:
    n = (labelled["crop"] == crop).sum()
    if n < 5:
        continue
    rows.append({"crop": crop, "fields": n,
                 "VH Nov-Dec": window_mean(crop, "vh_db", "2023-11-01", "2023-12-31"),
                 "VH Mar": window_mean(crop, "vh_db", "2024-03-01", "2024-03-31"),
                 "VH May-Jun": window_mean(crop, "vh_db", "2024-05-01", "2024-06-30"),
                 "VV May-Jun": window_mean(crop, "vv_db", "2024-05-01", "2024-06-30"),
                 "VH/VV May-Jun": window_mean(crop, "ratio_db", "2024-05-01", "2024-06-30"),
                 "VH/VV Sep": window_mean(crop, "ratio_db", "2024-09-01", "2024-09-30")})
print(pd.DataFrame(rows).round(1).to_string(index=False))
all_vv = np.nanmedian(vv_db, axis=1)          # every field on each date: weather moves them all together
jump = pd.Series(np.diff(all_vv), index=S1_DATES[1:])
print(f"largest all-field VV jumps: " + ", ".join(f"{d:%d %b %Y} {v:+.1f} dB" for d, v in
                                                  jump.reindex(jump.abs().sort_values(ascending=False).index[:3]).items()))
for crop in ["Winter wheat", "Winter oilseed rape"]:
    r, vv = med.loc[crop, "ratio_db"], med.loc[crop, "vv_db"]
    print(f"{crop}: VH/VV highest {r.max():.1f} dB on {r.idxmax():%d %b %Y}, lowest {r.min():.1f} dB on "
          f"{r.idxmin():%d %b %Y}; VV lowest {vv.min():.1f} dB on {vv.idxmin():%d %b %Y}")

# --- Step 7: how often could Sentinel-2 see the same fields? --------------------------------
section("Sentinel-2 over the same fields")
s2_all = list(catalog.search(collections=["sentinel-2-l2a"], bbox=BBOX, datetime=SEASON).items())
aoi = box(*BBOX)
tiles = pd.Series([i.properties["s2:mgrs_tile"] for i in s2_all]).value_counts()
TILE = tiles.index[0]
best = {}
for it in s2_all:
    if it.properties["s2:mgrs_tile"] != TILE:
        continue
    key = it.datetime.date()
    if key not in best or it.properties["s2:processing_baseline"] > best[key].properties["s2:processing_baseline"]:
        best[key] = it
s2_items = sorted(best.values(), key=lambda i: i.datetime)
print(f"{len(s2_items)} Sentinel-2 L2A acquisitions of tile {TILE} (any cloud cover), "
      f"{s2_items[0].datetime:%d %b %Y} to {s2_items[-1].datetime:%d %b %Y}")
CLEAR_SCL = [4, 5]            # SCL vegetation, not vegetated. Cloud, shadow, cirrus, unclassified etc. are unusable


def s2_scene(item) -> dict:
    cache = CACHE / f"{item.id}.npz"
    if cache.exists():
        arr = np.load(cache)
        return {"date": item.datetime, **{k: arr[k] for k in arr.files}}
    scl = read_warped(item.assets["SCL"].href, Resampling.nearest, nodata=0, dtype="uint8")
    clear = np.isin(scl, CLEAR_SCL)
    covered = float((scl[inside] > 0).mean())
    n_clear = ndimage.sum(clear, labels, FIDS)
    out = {"covered": np.array(covered), "clear_frac": np.array(float(clear[inside].mean())),
           "ndvi": np.full(len(FIDS), np.nan)}
    if out["clear_frac"] >= 0.2:               # only read the 10 m bands when there is something to see
        off = -1000 if float(item.properties["s2:processing_baseline"]) >= 4.0 else 0
        red = read_warped(item.assets["B04"].href, nodata=0) + off
        nir = read_warped(item.assets["B08"].href, nodata=0) + off
        with np.errstate(invalid="ignore", divide="ignore"):
            ndvi = np.where(clear, (nir - red) / (nir + red), np.nan)
        s = ndimage.sum(np.nan_to_num(ndvi), labels, FIDS)
        with np.errstate(invalid="ignore", divide="ignore"):
            out["ndvi"] = np.where(n_clear >= 0.8 * npix, s / n_clear, np.nan)
    np.savez_compressed(cache, **out)
    return {"date": item.datetime, **out}


t0 = time.time()
with ThreadPoolExecutor(max_workers=8) as pool:
    s2 = list(pool.map(s2_scene, s2_items))
s2 = [s for s in s2 if float(s["covered"]) >= 0.95]   # drop slivers at the swath edge
print(f"read SCL (and red/NIR where useful) for {len(s2_items)} acquisitions in {time.time() - t0:.0f} s; "
      f"{len(s2)} cover the fields fully")
S2_DATES = pd.DatetimeIndex([s["date"] for s in s2]).tz_localize(None).normalize()
clear_frac = pd.Series([float(s["clear_frac"]) for s in s2], index=S2_DATES)
NDVI = np.array([s["ndvi"] for s in s2])
usable = clear_frac >= 0.8
print(f"acquisitions with >= 80% of field pixels clear: {usable.sum()} of {len(s2)} ({usable.mean():.0%})")
print(f"acquisitions where at least one field is clear: {(np.isfinite(NDVI).any(axis=1)).sum()}")
per_field = np.isfinite(NDVI).sum(axis=0)
print(f"clear NDVI observations per field over 13 months: median {np.median(per_field):.0f}, "
      f"min {per_field.min()}, max {per_field.max()} (Sentinel-1: {len(S1_DATES)} on every field)")
field_gap = [np.diff(S2_DATES[np.isfinite(NDVI[:, k])].to_numpy()).astype("timedelta64[D]").astype(int).max()
             for k in range(len(FIDS))]
spring = (S2_DATES >= "2024-03-01") & (S2_DATES < "2024-05-01")
print(f"longest gap between clear views of the same field: median {np.median(field_gap):.0f} days "
      f"(range {min(field_gap)}-{max(field_gap)}); clear views per field in March-April 2024: "
      f"median {np.median(np.isfinite(NDVI[spring]).sum(axis=0)):.0f}")
gaps = np.diff(S2_DATES[usable].to_numpy()).astype("timedelta64[D]").astype(int)
s1_gaps = np.diff(S1_DATES.to_numpy()).astype("timedelta64[D]").astype(int)
g = int(np.argmax(gaps))
print(f"longest gap between >= 80%-clear Sentinel-2 dates: {gaps.max()} days "
      f"({S2_DATES[usable][g]:%d %b %Y} to {S2_DATES[usable][g + 1]:%d %b %Y}); longest Sentinel-1 gap: {s1_gaps.max()} days")
monthly = pd.DataFrame({
    "S1 scenes (orbit 59)": pd.Series(1, index=S1_DATES).resample("MS").sum(),
    "S2 acquisitions": pd.Series(1, index=S2_DATES).resample("MS").sum(),
    "S2 >= 80% clear": usable.astype(int).resample("MS").sum()}).fillna(0).astype(int)
monthly.index = monthly.index.strftime("%b %Y")
print(monthly.to_string())

section("NDVI versus VH/VV ratio (crop medians, dates within 3 days)")
ndvi_long = pd.DataFrame({"date": np.repeat(S2_DATES, len(FIDS)), "fid": np.tile(FIDS, len(S2_DATES)),
                          "ndvi": NDVI.ravel()}).dropna().merge(labelled[["fid", "crop"]], on="fid")
ndvi_med = ndvi_long.groupby(["crop", "date"])["ndvi"].agg(["median", "size"])
ndvi_med = ndvi_med[ndvi_med["size"] >= 5]["median"]
for crop in ["Winter wheat", "Winter oilseed rape"]:
    a, b = ndvi_med.loc[crop], med.loc[crop, "ratio_db"]
    pairs = pd.merge_asof(a.rename("ndvi").reset_index().sort_values("date"),
                          b.rename("ratio").reset_index().sort_values("date"),
                          on="date", direction="nearest", tolerance=pd.Timedelta("3D")).dropna()
    r = np.corrcoef(pairs["ndvi"], pairs["ratio"])[0, 1]
    print(f"{crop}: {len(a)} NDVI dates with >= 5 clear fields; {len(pairs)} within 3 days of a radar date; "
          f"Pearson r = {r:.2f}")
    odd = a[(a - a.rolling(5, center=True, min_periods=3).median()).abs() > 0.25]
    print(f"  NDVI outliers (> 0.25 from the 5-date running median): "
          + (", ".join(f"{d:%d %b %Y} ({v:.2f})" for d, v in odd.items()) or "none"))

# --- Step 8: can radar alone tell the crops apart? -------------------------------------------
section("nearest-centroid crop classification from Sentinel-1 alone (leave-one-out)")
CLASSES = [c for c in ["Winter wheat", "Winter oilseed rape", "Winter barley", "Spring barley", "Maize",
                       "Winter field beans", "Sugar beet"] if (labelled["crop"] == c).sum() >= 8]
sel = labelled["crop"].isin(CLASSES).to_numpy()
X = np.hstack([vv_db.T, vh_db.T])[sel]        # fields x (dates x 2) feature matrix
ok = np.isfinite(X).all(axis=0)
X = X[:, ok]
X = (X - X.mean(axis=0)) / X.std(axis=0)
y = labelled.loc[sel, "crop"].to_numpy()
pred = []
for i in range(len(y)):                        # leave one field out, classify it by the nearest class centroid
    keep = np.arange(len(y)) != i
    cents = {c: X[keep][y[keep] == c].mean(axis=0) for c in CLASSES}
    pred.append(min(cents, key=lambda c: np.linalg.norm(X[i] - cents[c])))
pred = np.array(pred)
cm = pd.crosstab(pd.Series(y, name="CROME label"), pd.Series(pred, name="predicted")).reindex(
    index=CLASSES, columns=CLASSES, fill_value=0)
print(f"{len(y)} fields, {X.shape[1]} features ({ok.sum() // 2} dates x VV, VH); "
      f"overall accuracy {np.mean(pred == y):.0%}")
print(cm.to_string())
for c in CLASSES:
    print(f"  {c:<20} recall {np.mean(pred[y == c] == c):.0%}  ({(y == c).sum()} fields)")

# --- Figures ---------------------------------------------------------------------------------
extent = (x0 / 1000, x1 / 1000, y0 / 1000, y1 / 1000)
outline = labelled.copy()
outline["geometry"] = outline.geometry.scale(1e-3, 1e-3, origin=(0, 0))


def vh_stretch(img):
    return np.clip((db(ndimage.uniform_filter(np.nan_to_num(img, nan=1e-4), 3)) + 25) / 13, 0, 1)


# Figure 1: multi-temporal VH composite and the CROME labels
rgb = np.dstack([vh_stretch(imgs[d]) for d in KEEP_DATES])
fig, axes = plt.subplots(1, 2, figsize=(11.5, 4.6), sharex=True, sharey=True)
axes[0].imshow(rgb, extent=extent)
axes[0].set_title(f"Sentinel-1 VH: R {KEEP_DATES[0]}, G {KEEP_DATES[1]}, B {KEEP_DATES[2]}")
axes[1].imshow(vh_stretch(imgs[MID]), cmap="gray", extent=extent, alpha=0.5)
present = labelled["crop"].value_counts()
for crop in present.index:
    outline[outline["crop"] == crop].plot(ax=axes[1], color=COLOURS.get(crop, "#e9ecef"), edgecolor="black", lw=0.3)
axes[1].legend(handles=[Patch(facecolor=COLOURS.get(c, "#e9ecef"), edgecolor="black", lw=0.3, label=f"{c} ({n})")
                        for c, n in present.items()],
               loc="upper left", bbox_to_anchor=(1.02, 1), frameon=False, fontsize=8, title="CROME 2024 label (fields)")
axes[1].set_title(f"{len(labelled)} single-crop fields")
for ax in axes:
    ax.set_xlabel("Easting (km, UTM 30N)")
axes[0].set_ylabel("Northing (km)")
axes[1].set_ylabel("")
fig.tight_layout()
fig.savefig(OUT / "01-vh-composite-crop-map.png")
plt.close(fig)

# Figure 2: speckle at pixel level versus the field average
big = labelled.sort_values("area_ha", ascending=False).iloc[0].geometry.centroid
zx0, zy0 = np.floor((big.x - 1500 - x0) / RES) * RES + x0, np.floor((big.y - 1500 - y0) / RES) * RES + y0
c0, r0 = int((zx0 - x0) / RES), int((y1 - (zy0 + 3000)) / RES)
win = (slice(max(r0, 0), max(r0, 0) + 300), slice(max(c0, 0), max(c0, 0) + 300))
lut = np.full(labels.max() + 1, np.nan)
lut[FIDS] = db(field_means(vh_img))
fm = np.where(labels > 0, lut[labels], np.nan)[win]
fig, axes = plt.subplots(1, 3, figsize=(11.5, 3.7), gridspec_kw={"width_ratios": [1, 1, 1.25]})
for ax, im, t in [(axes[0], db(vh_img[win]), "Single 10 m pixels"), (axes[1], fm, "Field-core means")]:
    h = ax.imshow(im, cmap="viridis", vmin=-26, vmax=-12)
    ax.set_title(f"{t}\nVH, {MID}, 3 x 3 km")
    ax.set_xticks([]), ax.set_yticks([])
fig.colorbar(h, ax=axes[:2], shrink=0.85, label="VH backscatter (dB)")
for crop in ["Winter wheat", "Winter oilseed rape"]:
    f = labelled.loc[labelled["crop"] == crop, "fid"].to_numpy()
    vals = pix_db[np.isin(lab_in, f)]
    axes[2].hist(vals[np.isfinite(vals)], bins=np.arange(-34, -6, 0.5), histtype="step", lw=1.5, density=True,
                 color=COLOURS[crop], label=f"{crop}: all pixels")
    fmeans = db(field_means(vh_img))[np.isin(FIDS, f)]
    axes[2].plot(fmeans, np.full(len(fmeans), -0.012 if crop == "Winter wheat" else -0.03), "|",
                 color=COLOURS[crop], ms=8, label=f"{crop}: field means")
axes[2].set_xlabel("VH (dB)")
axes[2].set_ylabel("Density")
axes[2].set_title("All pixels versus field means")
axes[2].set_ylim(-0.045, 0.34)
axes[2].legend(fontsize=7, frameon=False, loc="upper left")
fig.savefig(OUT / "02-speckle-field-averaging.png")
plt.close(fig)

# Figure 3: crop signatures in VV, VH and VH/VV
fig, axes = plt.subplots(3, 1, figsize=(9, 8.2), sharex=True)
for ax, var, lab in zip(axes, ["vv_db", "vh_db", "ratio_db"], ["VV (dB)", "VH (dB)", "VH/VV ratio (dB)"]):
    for crop in SHOW:
        grp = long[long["crop"] == crop].groupby("date")[var]
        q25, q50, q75 = grp.quantile(0.25), grp.median(), grp.quantile(0.75)
        ax.fill_between(q50.index, q25, q75, color=COLOURS[crop], alpha=0.18, lw=0)
        ax.plot(q50.index, q50, color=COLOURS[crop], lw=1.8, marker="o", ms=2.5,
                label=f"{crop} ({(labelled['crop'] == crop).sum()} fields)")
    ax.set_ylabel(lab)
    ax.grid(alpha=0.3)
axes[0].legend(ncol=2, fontsize=8, frameon=False, loc="lower left")
axes[2].xaxis.set_major_locator(mdates.MonthLocator(bymonth=[9, 11, 1, 3, 5, 7, 9]))
axes[2].xaxis.set_major_formatter(mdates.DateFormatter("%b\n%Y"))
axes[0].set_title(f"Sentinel-1 backscatter per crop (orbit {ORBIT}): median field and interquartile range")
fig.tight_layout()
fig.savefig(OUT / "03-crop-signatures-vv-vh-ratio.png")
plt.close(fig)

# Figure 4: Sentinel-1 versus Sentinel-2 availability
partly = (clear_frac >= 0.2) & ~usable
cols = np.where(usable, "#2b8a3e", np.where(partly, "#f59f00", "#adb5bd"))
fig, (ax, ax2) = plt.subplots(2, 1, figsize=(9, 4.8), sharex=True, gridspec_kw={"height_ratios": [1.1, 1]})
ax.vlines(S1_DATES, 1.6, 2.4, color="#1543b0", lw=1.5)
ax.vlines(S2_DATES, 0.1, 0.9, color=cols, lw=1.2)
ax.set_yticks([0.5, 2.0], ["Sentinel-2", "Sentinel-1"])
ax.set_ylim(0, 2.5)
ax.legend(handles=[Patch(color="#1543b0", label=f"Sentinel-1 scene, orbit {ORBIT} ({len(S1_DATES)})"),
                   Patch(color="#2b8a3e", label=f"Sentinel-2, >= 80% of field pixels clear ({usable.sum()})"),
                   Patch(color="#f59f00", label=f"Sentinel-2, 20-80% clear ({partly.sum()})"),
                   Patch(color="#adb5bd", label=f"Sentinel-2, < 20% clear ({(clear_frac < 0.2).sum()})")],
          ncol=2, fontsize=7.5, frameon=False, loc="lower center", bbox_to_anchor=(0.5, 1.0))
ax2.bar(clear_frac.index, clear_frac * 100, width=2.5, color=cols)
ax2.set_ylabel("Sentinel-2 field\npixels clear (%)")
ax2.axhline(80, color="black", lw=0.6, ls="--")
ax2.xaxis.set_major_formatter(mdates.DateFormatter("%b\n%Y"))
fig.tight_layout()
fig.savefig(OUT / "04-s1-vs-s2-availability.png")
plt.close(fig)

# Figure 5: NDVI versus VH/VV for wheat and oilseed rape
fig, axes = plt.subplots(2, 1, figsize=(9, 6), sharex=True)
for ax, crop in zip(axes, ["Winter wheat", "Winter oilseed rape"]):
    r = med.loc[crop, "ratio_db"]
    ax.plot(r.index, r, color="#1543b0", lw=1.8, marker="o", ms=3, label="Sentinel-1 VH/VV (dB, left axis)")
    ax.set_ylabel("VH/VV (dB)", color="#1543b0")
    axr = ax.twinx()
    nd = ndvi_med.loc[crop]
    axr.plot(nd.index, nd, color=COLOURS[crop], lw=0, marker="s", ms=5, label="Sentinel-2 NDVI, clear dates (right axis)")
    axr.set_ylabel("NDVI", color=COLOURS[crop])
    axr.set_ylim(0, 1)
    ax.set_title(f"{crop}: {len(r)} radar dates, {len(nd)} usable NDVI dates")
    ax.grid(alpha=0.3)
    h1, l1 = ax.get_legend_handles_labels()
    h2, l2 = axr.get_legend_handles_labels()
    ax.legend(h1 + h2, l1 + l2, fontsize=7.5, frameon=False, loc="lower left")
axes[1].xaxis.set_major_formatter(mdates.DateFormatter("%b\n%Y"))
fig.tight_layout()
fig.savefig(OUT / "05-ndvi-vs-vhvv-ratio.png")
plt.close(fig)

# Figure 6: confusion matrix
frac = cm.div(cm.sum(axis=1), axis=0)
fig, ax = plt.subplots(figsize=(6.6, 5.2))
ax.imshow(frac, cmap="Blues", vmin=0, vmax=1)
for i in range(len(CLASSES)):
    for j in range(len(CLASSES)):
        ax.text(j, i, f"{cm.iloc[i, j]}", ha="center", va="center", fontsize=9,
                color="white" if frac.iloc[i, j] > 0.55 else "black")
ax.set_xticks(range(len(CLASSES)), [c.replace(" ", "\n", 1) for c in CLASSES], fontsize=8)
ax.set_yticks(range(len(CLASSES)), CLASSES, fontsize=8)
ax.set_xlabel("Predicted from the Sentinel-1 VV + VH time series")
ax.set_ylabel("CROME 2024 label")
ax.set_title(f"Leave-one-out nearest centroid: {np.mean(pred == y):.0%} of {len(y)} fields correct")
fig.tight_layout()
fig.savefig(OUT / "06-confusion-matrix.png")
plt.close(fig)
print(f"\nfigures and outputs written to {OUT.name}/")