"""Sentinel-2 cloud and cloud-shadow masking in Python: SCL vs s2cloudless vs OmniCloudMask vs a
simple threshold baseline over the Lake District, plus a cloud-free median composite.

TechnicalGIS tutorial: https://www.technicalgis.com/tutorials/sentinel-2-cloud-masking-python/
Data: Copernicus Sentinel-2 L2A (Microsoft Planetary Computer) and L1C (Google Cloud public bucket).
"""
import os
import sys
import time
import warnings
from importlib.metadata import version
from pathlib import Path

for var in ("PROJ_LIB", "PROJ_DATA", "GDAL_DATA"):   # stray installs (e.g. PostGIS) break the wheels' own copies
    os.environ.pop(var, None)

import numpy as np
import pandas as pd
import requests
import rasterio
from rasterio.enums import Resampling
from rasterio.transform import from_origin
from rasterio.windows import from_bounds
from scipy import ndimage
import pystac_client
import planetary_computer as pc
import odc.stac
from odc.geo.geobox import GeoBox
import matplotlib
matplotlib.use("Agg")   # render to files only
import matplotlib.pyplot as plt
from matplotlib.colors import ListedColormap
from matplotlib.patches import Patch
from s2cloudless import S2PixelCloudDetector
from omnicloudmask import predict_from_array

warnings.filterwarnings("ignore")
plt.rcParams.update({"figure.dpi": 150, "savefig.bbox": "tight", "font.size": 9})
OUT = Path("output"); OUT.mkdir(exist_ok=True)
CHIPS = OUT / "_chips"; CHIPS.mkdir(exist_ok=True)

STAC_URL = "https://planetarycomputer.microsoft.com/api/stac/v1"
GCS = "https://storage.googleapis.com"
L1C_BUCKET = "gcp-public-data-sentinel-2"
CRS = "EPSG:32630"                                # UTM 30N, the native grid of tile 30UVF
AOI = (478000, 6030000, 498000, 6050000)          # 20 x 20 km: Keswick, Derwentwater, Borrowdale, Buttermere
RES = 20                                          # analysis grid in metres (SCL and s2cloudless are 20 m products)
H = W = int((AOI[2] - AOI[0]) / RES)
TILE = "30UVF"
BBOX_LL = [-3.35, 54.41, -3.02, 54.60]
EVAL_DATES = ["2026-05-26", "2026-08-08", "2026-09-23"]   # three cloudy-but-not-overcast dates
COMPOSITE_PERIOD = "2026-05-01/2026-09-30"
CLEAR_DATE = "2026-07-12"                          # the clearest summer date, used to check the composites
N_POINTS = 80                                      # reference points per evaluation date
SEED = 42

# Every mask uses the same codes: 0 clear, 1 cloud, 2 cloud shadow
CLEAR, CLOUD, SHADOW = 0, 1, 2
MASK_CMAP = ListedColormap(["#00000000", "#ffffbf", "#5e4fa2"])
METHODS = ["SCL", "s2cloudless", "OmniCloudMask", "Threshold"]
COLOURS = {"SCL": "#3288bd", "s2cloudless": "#66c2a5", "OmniCloudMask": "#d53e4f", "Threshold": "#fdae61"}


def section(title):
    print(f"\n=== {title} ===")


def read_window(href, out_shape=(H, W), resampling=Resampling.average):
    """Read the AOI from a COG or JP2 over HTTP, resampled onto the analysis grid."""
    with rasterio.Env(GDAL_DISABLE_READDIR_ON_OPEN="EMPTY_DIR", GDAL_HTTP_MAX_RETRY="4", GDAL_HTTP_RETRY_DELAY="2"):
        with rasterio.open(href) as src:
            win = from_bounds(*AOI, src.transform)
            return src.read(1, window=win, out_shape=out_shape, resampling=resampling)


def to_reflectance(dn):
    """Processing baseline 04.00+ adds 1000 to every DN, so reflectance = (DN - 1000) / 10000. DN 0 is no data."""
    return np.where(dn == 0, np.nan, (dn.astype("float32") - 1000) / 10000)


def rgb(stack, gain=3.2):
    """True-colour display image with a fixed gain, so that dates and composites can be compared."""
    return np.clip(np.nan_to_num(np.dstack(stack)) * gain, 0, 1) ** 0.85


def kappa(a, b):
    po = np.mean(a == b)
    pe = np.mean(a) * np.mean(b) + np.mean(~a) * np.mean(~b)
    return (po - pe) / (1 - pe)


section("versions")
print(f"python {sys.version.split()[0]} | rasterio {version('rasterio')} | odc-stac {version('odc-stac')} | "
      f"s2cloudless {version('s2cloudless')} | omnicloudmask {version('omnicloudmask')} | torch {version('torch')}")

# ---------------------------------------------------------------- Step 1: find the scenes
section("Step 1: Sentinel-2 L2A scenes for tile 30UVF")
catalog = pystac_client.Client.open(STAC_URL, modifier=pc.sign_inplace)


def l2a_item(date):
    items = catalog.search(collections=["sentinel-2-l2a"], bbox=BBOX_LL, datetime=date,
                           query={"s2:mgrs_tile": {"eq": TILE}}).item_collection()
    return sorted(items, key=lambda i: i.properties["s2:generation_time"])[-1]   # latest processing


eval_items = {d: l2a_item(d) for d in EVAL_DATES}
for d, it in eval_items.items():
    p = it.properties
    print(f"{d}  {it.id}  tile cloud {p['eo:cloud_cover']:.1f}%  sun zenith {p['s2:mean_solar_zenith']:.1f}  "
          f"azimuth {p['s2:mean_solar_azimuth']:.1f}")

# ---------------------------------------------------------------- Step 2: load the bands
section("Step 2: load L2A bands on a 20 m grid")
scenes = {}
for d, it in eval_items.items():
    t0 = time.time()
    b = {k: to_reflectance(read_window(it.assets[k].href)) for k in ["B02", "B03", "B04", "B08", "B11"]}
    b["SCL"] = read_window(it.assets["SCL"].href, resampling=Resampling.nearest)
    # 10 m true colour for the figures and the reference chips
    b["rgb10"] = [to_reflectance(read_window(it.assets[k].href, out_shape=(2 * H, 2 * W))) for k in ["B04", "B03", "B02"]]
    b["item"] = it
    scenes[d] = b
    print(f"{d}: {H} x {W} px, median B02 {np.nanmedian(b['B02']):.3f}, B08 {np.nanmedian(b['B08']):.3f} "
          f"({time.time() - t0:.0f} s)")

# ---------------------------------------------------------------- Step 3: SCL mask
section("Step 3: Scene Classification Layer (SCL)")
SCL_NAMES = {0: "no data", 1: "saturated/defective", 2: "dark area", 3: "cloud shadow", 4: "vegetation",
             5: "not vegetated", 6: "water", 7: "unclassified", 8: "cloud medium prob.", 9: "cloud high prob.",
             10: "thin cirrus", 11: "snow/ice"}


def scl_mask(scl):
    m = np.zeros(scl.shape, "uint8")
    m[np.isin(scl, [8, 9, 10])] = CLOUD
    m[scl == 3] = SHADOW
    return m


for d, b in scenes.items():
    vals, counts = np.unique(b["SCL"], return_counts=True)
    print(d, " | ".join(f"{SCL_NAMES[v]} {c / b['SCL'].size * 100:.1f}%" for v, c in zip(vals, counts)
                        if c / b["SCL"].size >= 0.005))
    b["SCL_mask"] = scl_mask(b["SCL"])

# ---------------------------------------------------------------- Step 4: s2cloudless on L1C top-of-atmosphere bands
section("Step 4: s2cloudless (L1C from the Google Cloud public bucket)")
S2C_BANDS = ["B01", "B02", "B04", "B05", "B08", "B8A", "B09", "B10", "B11", "B12"]


def l1c_band_urls(item):
    """Find the L1C product with the same sensing time and orbit in the public Google Cloud bucket."""
    _, _, sensing, orbit, _, _ = item.id.split("_")       # e.g. S2C_MSIL2A_20260526T113309_R080_T30UVF_...
    prefix = f"tiles/{TILE[:2]}/{TILE[2]}/{TILE[3:]}/{item.id[:3]}_MSIL1C_{sensing}"
    api = f"{GCS}/storage/v1/b/{L1C_BUCKET}/o"
    r = requests.get(api, params={"prefix": prefix, "delimiter": "/"}, timeout=60)
    safe = sorted(s for s in r.json().get("prefixes", []) if f"_{orbit}_" in s)[-1]
    r = requests.get(api, params={"prefix": safe + "GRANULE/", "fields": "items(name)"}, timeout=60)
    names = [i["name"] for i in r.json()["items"] if "/IMG_DATA/" in i["name"] and i["name"].endswith(".jp2")]
    urls = {band: f"/vsicurl/{GCS}/{L1C_BUCKET}/{n}" for band in S2C_BANDS for n in names if n.endswith(f"_{band}.jp2")}
    return safe.split("/")[-2], urls


def project_shadows(cloud, dark, sun_azimuth, sun_zenith, heights_m=range(200, 3001, 200)):
    """Cast the cloud mask away from the sun for cloud heights of 0.2-3 km; keep the dark pixels it lands on."""
    az = np.radians(sun_azimuth)
    projected = np.zeros_like(cloud)
    for h in heights_m:
        dist = h * np.tan(np.radians(sun_zenith)) / RES          # horizontal offset in pixels
        dx, dy = -np.sin(az) * dist, np.cos(az) * dist            # image rows increase southwards
        projected |= ndimage.shift(cloud.astype("uint8"), (dy, dx), order=0, cval=0).astype(bool)
    return projected & dark & ~cloud


detector = S2PixelCloudDetector(threshold=0.4, average_over=4, dilation_size=2, all_bands=False)
for d, b in scenes.items():
    t0 = time.time()
    safe, urls = l1c_band_urls(b["item"])
    toa = np.dstack([to_reflectance(read_window(urls[k])) for k in S2C_BANDS])
    prob = detector.get_cloud_probability_maps(np.nan_to_num(toa)[np.newaxis])[0]
    cloud = detector.get_mask_from_prob(prob[np.newaxis])[0].astype(bool)
    p = b["item"].properties
    mndwi = (b["B03"] - b["B11"]) / (b["B03"] + b["B11"])
    dark = (b["B08"] < 0.15) & (mndwi < 0)                        # dark, but not water
    shadow = project_shadows(cloud, dark, p["s2:mean_solar_azimuth"], p["s2:mean_solar_zenith"])
    m = np.zeros((H, W), "uint8"); m[shadow] = SHADOW; m[cloud] = CLOUD
    b["s2c_prob"], b["s2cloudless_mask"] = prob, m
    print(f"{d}: {safe}\n    cloud {cloud.mean() * 100:.1f}%  projected shadow {shadow.mean() * 100:.1f}%  "
          f"({time.time() - t0:.0f} s)")

# ---------------------------------------------------------------- Step 5: OmniCloudMask
section("Step 5: OmniCloudMask (red, green and NIR)")


def omnicloudmask(red, green, nir):
    arr = np.nan_to_num(np.stack([red, green, nir]), nan=0.0).astype("float32")
    return predict_from_array(arr, inference_device="cpu", mosaic_device="cpu", no_data_value=-9999)[0]


for d, b in scenes.items():
    t0 = time.time()
    pred = omnicloudmask(b["B04"], b["B03"], b["B08"])
    # OmniCloudMask classes: 0 clear, 1 thick cloud, 2 thin cloud, 3 cloud shadow
    m = np.zeros((H, W), "uint8"); m[np.isin(pred, [1, 2])] = CLOUD; m[pred == 3] = SHADOW
    b["ocm_raw"], b["OmniCloudMask_mask"] = pred, m
    print(f"{d}: thick {np.mean(pred == 1) * 100:.1f}%  thin {np.mean(pred == 2) * 100:.1f}%  "
          f"shadow {np.mean(pred == 3) * 100:.1f}%  ({time.time() - t0:.0f} s on CPU)")

# ---------------------------------------------------------------- Step 6: simple threshold baseline
section("Step 6: threshold baseline")
BLUE_CLOUD = 0.18      # clear land in summer has B02 reflectance of about 0.02-0.08
NIR_SHADOW = 0.12      # sunlit vegetation has B08 of 0.2-0.4; in shadow it drops below about 0.12


def threshold_mask(b):
    mndwi = (b["B03"] - b["B11"]) / (b["B03"] + b["B11"])
    cloud = b["B02"] > BLUE_CLOUD
    shadow = (b["B08"] < NIR_SHADOW) & (mndwi < 0) & ~cloud
    m = np.zeros((H, W), "uint8"); m[shadow] = SHADOW; m[cloud] = CLOUD
    return m


for d, b in scenes.items():
    b["Threshold_mask"] = threshold_mask(b)
    print(f"{d}: cloud {np.mean(b['Threshold_mask'] == CLOUD) * 100:.1f}%  "
          f"shadow {np.mean(b['Threshold_mask'] == SHADOW) * 100:.1f}%")

# ---------------------------------------------------------------- Step 7: how much do the methods agree?
section("Step 7: agreement between methods")
rows = []
for d, b in scenes.items():
    for m in METHODS:
        mk = b[f"{m}_mask"]
        rows.append({"date": d, "method": m, "cloud_%": np.mean(mk == CLOUD) * 100,
                     "shadow_%": np.mean(mk == SHADOW) * 100, "masked_%": np.mean(mk > 0) * 100})
cover = pd.DataFrame(rows)
for col in ["masked_%", "cloud_%", "shadow_%"]:
    print(f"\n{col.replace('_%', '')} (% of the area):")
    print(cover.pivot(index="method", columns="date", values=col).round(1).loc[METHODS].to_string())

agree = np.zeros((len(METHODS), len(METHODS)))
kap = np.zeros_like(agree)
for i, mi in enumerate(METHODS):
    for j, mj in enumerate(METHODS):
        a = np.concatenate([(scenes[d][f"{mi}_mask"] > 0).ravel() for d in EVAL_DATES])
        c = np.concatenate([(scenes[d][f"{mj}_mask"] > 0).ravel() for d in EVAL_DATES])
        agree[i, j] = np.mean(a == c) * 100
        kap[i, j] = kappa(a, c)
print("\npixel agreement, masked vs not masked, all three dates (%):")
print(pd.DataFrame(agree, index=METHODS, columns=METHODS).round(1).to_string())
print("\nCohen's kappa:")
print(pd.DataFrame(kap, index=METHODS, columns=METHODS).round(2).to_string())
all_agree = np.concatenate([np.all([scenes[d][f"{m}_mask"] > 0 for m in METHODS], axis=0).ravel() |
                            np.all([scenes[d][f"{m}_mask"] == 0 for m in METHODS], axis=0).ravel()
                            for d in EVAL_DATES])
print(f"\nall four methods agree on {all_agree.mean() * 100:.1f}% of pixels")

# ---------------------------------------------------------------- Step 8: reference points
section("Step 8: check against visually interpreted reference points")
rng = np.random.default_rng(SEED)
pts = []
for d in EVAL_DATES:
    rr, cc = rng.integers(16, H - 16, N_POINTS), rng.integers(16, W - 16, N_POINTS)
    pts += [{"point_id": f"{d[5:7]}{d[8:10]}-{k:02d}", "date": d, "row": int(r), "col": int(c)}
            for k, (r, c) in enumerate(zip(rr, cc))]
pts = pd.DataFrame(pts)

# Contact sheets of 10 m true-colour chips (640 m across, centre 20 m pixel boxed), blind to every mask
HALF = 32
for d in EVAL_DATES:
    img = rgb(scenes[d]["rgb10"])
    sub = pts[pts.date == d].reset_index(drop=True)
    for s in range(0, len(sub), 20):
        fig, axes = plt.subplots(4, 5, figsize=(10, 8.4))
        for ax, p in zip(axes.ravel(), sub.iloc[s:s + 20].itertuples()):
            r10, c10 = 2 * p.row, 2 * p.col
            ax.imshow(img[r10 - HALF:r10 + HALF, c10 - HALF:c10 + HALF], interpolation="nearest")
            ax.add_patch(plt.Rectangle((HALF - 0.5, HALF - 0.5), 2, 2, fill=False, ec="#ff00ff", lw=1.2))
            ax.set_title(p.point_id, fontsize=8); ax.set_xticks([]); ax.set_yticks([])
        fig.savefig(CHIPS / f"chips_{d}_{s // 20 + 1}.png", dpi=110)
        plt.close(fig)

labels_csv = Path("reference_points.csv")
if not labels_csv.exists():
    pts.assign(label="").to_csv(labels_csv, index=False)
    sys.exit(f"Label the points in {labels_csv} from the contact sheets in {CHIPS}, then run again.")
ref = pts.merge(pd.read_csv(labels_csv)[["point_id", "label"]], on="point_id")
print("reference labels:", ref.label.value_counts().to_dict())
ref = ref[ref.label.isin(["clear", "cloud", "shadow"])].copy()
ref["truth"] = ref.label.map({"clear": CLEAR, "cloud": CLOUD, "shadow": SHADOW})

scores = []
for m in METHODS:
    pred = np.array([scenes[p.date][f"{m}_mask"][p.row, p.col] for p in ref.itertuples()])
    ref[m] = pred
    t_bad, p_bad = ref.truth.values > 0, pred > 0
    tp, fp, fn = np.sum(t_bad & p_bad), np.sum(~t_bad & p_bad), np.sum(t_bad & ~p_bad)
    scores.append({"method": m, "overall_acc_%": np.mean(t_bad == p_bad) * 100, "kappa": kappa(t_bad, p_bad),
                   "omission_%": fn / (tp + fn) * 100,        # cloud or shadow the method left in
                   "commission_%": fp / max(tp + fp, 1) * 100,  # masked pixels that were really clear
                   "cloud_found_%": np.mean(pred[ref.truth == CLOUD] > 0) * 100,
                   "shadow_found_%": np.mean(pred[ref.truth == SHADOW] > 0) * 100,
                   "clear_kept_%": np.mean(pred[ref.truth == CLEAR] == 0) * 100})
scores = pd.DataFrame(scores).set_index("method")
print(f"\n{len(ref)} usable points (clear {np.sum(ref.truth == CLEAR)}, cloud {np.sum(ref.truth == CLOUD)}, "
      f"shadow {np.sum(ref.truth == SHADOW)}); masked = cloud or shadow")
print(scores.round(2).to_string())

print("\nconfusion matrices (rows: reference, columns: method)")
for m in METHODS:
    cm = pd.crosstab(ref.label, ref[m].map({CLEAR: "clear", CLOUD: "cloud", SHADOW: "shadow"}))
    print(f"-- {m}\n{cm.to_string()}")
ref.to_csv(OUT / "reference_point_results.csv", index=False)

# ---------------------------------------------------------------- figures for the evaluation dates
section("figures: masks")
EXTENT = [0, (AOI[2] - AOI[0]) / 1000, 0, (AOI[3] - AOI[1]) / 1000]
d0 = EVAL_DATES[0]

scl_cols = {0: "#000000", 1: "#ff0000", 2: "#404040", 3: "#5e4fa2", 4: "#66bd63", 5: "#d9b365", 6: "#3288bd",
            7: "#999999", 8: "#bdbdbd", 9: "#ffffbf", 10: "#9ecae1", 11: "#e0f3f8"}
fig, axes = plt.subplots(1, 2, figsize=(11.5, 5.4))
axes[0].imshow(rgb(scenes[d0]["rgb10"]), extent=EXTENT)
axes[0].set_title(f"True colour, {d0}")
present = [v for v in range(12) if np.mean(scenes[d0]["SCL"] == v) > 0.0005]
axes[1].imshow(scenes[d0]["SCL"], cmap=ListedColormap([scl_cols[v] for v in range(12)]), vmin=-0.5, vmax=11.5,
               extent=EXTENT, interpolation="nearest")
axes[1].set_title("Scene Classification Layer (SCL)")
axes[1].legend(handles=[Patch(fc=scl_cols[v], ec="#444444", lw=0.5, label=f"{v}  {SCL_NAMES[v]}") for v in present],
               loc="upper left", bbox_to_anchor=(1.02, 1), frameon=False, fontsize=8)
for ax in axes:
    ax.set_xlabel("km east"); ax.set_ylabel("km north")
fig.savefig(OUT / "01-scene-and-scl.png")
plt.close(fig)

for d in EVAL_DATES:
    fig, axes = plt.subplots(1, 5, figsize=(16, 3.9), layout="constrained")
    axes[0].imshow(rgb(scenes[d]["rgb10"]), extent=EXTENT)
    axes[0].set_title(f"True colour {d}")
    grey = np.nanmean(np.dstack([scenes[d]["B04"], scenes[d]["B03"], scenes[d]["B02"]]), axis=2)
    for ax, m in zip(axes[1:], METHODS):
        ax.imshow(grey, cmap="gray", vmin=0, vmax=0.35, extent=EXTENT)
        ax.imshow(scenes[d][f"{m}_mask"], cmap=MASK_CMAP, vmin=0, vmax=2, extent=EXTENT, interpolation="nearest")
        ax.set_title(f"{m}: {np.mean(scenes[d][f'{m}_mask'] > 0) * 100:.0f}% masked")
    for ax in axes:
        ax.set_xticks([]); ax.set_yticks([])
    fig.legend(handles=[Patch(fc="#ffffbf", ec="#444444", label="cloud"),
                        Patch(fc="#5e4fa2", label="cloud shadow")],
               loc="lower center", ncol=2, frameon=False, bbox_to_anchor=(0.5, -0.05))
    fig.savefig(OUT / ("02-four-masks.png" if d == d0 else f"_masks-{d}.png"))
    plt.close(fig)

fig, axes = plt.subplots(1, 2, figsize=(12.5, 4.6), layout="constrained", gridspec_kw={"width_ratios": [1.15, 1]})
x = np.arange(len(EVAL_DATES)); wbar = 0.2
for k, m in enumerate(METHODS):
    sub = cover[cover.method == m]
    axes[0].bar(x + (k - 1.5) * wbar, sub["cloud_%"], wbar, color=COLOURS[m], label=m)
    axes[0].bar(x + (k - 1.5) * wbar, sub["shadow_%"], wbar, bottom=sub["cloud_%"], color=COLOURS[m], alpha=0.5,
                hatch="///", edgecolor="white", lw=0)
axes[0].set_xticks(x, EVAL_DATES); axes[0].set_ylabel("% of the 20 x 20 km area")
axes[0].set_title("Masked area: cloud (solid) and shadow (hatched)")
axes[0].legend(frameon=False, fontsize=8)
im = axes[1].imshow(kap, cmap="Spectral", vmin=0, vmax=1)
axes[1].set_xticks(range(4), METHODS, rotation=20); axes[1].set_yticks(range(4), METHODS)
for i in range(4):
    for j in range(4):
        axes[1].text(j, i, f"{kap[i, j]:.2f}\n{agree[i, j]:.0f}%", ha="center", va="center", fontsize=8)
axes[1].set_title("Pairwise Cohen's kappa and % agreement")
fig.colorbar(im, ax=axes[1], shrink=0.8, label="kappa")
fig.savefig(OUT / "03-agreement.png")
plt.close(fig)

fig, ax = plt.subplots(figsize=(8.5, 3.8))
metrics = ["overall_acc_%", "clear_kept_%", "cloud_found_%", "shadow_found_%"]
names = ["Overall accuracy", "Clear points kept", "Cloud points masked", "Shadow points masked"]
x = np.arange(len(metrics))
for k, m in enumerate(METHODS):
    bars = ax.bar(x + (k - 1.5) * wbar, scores.loc[m, metrics], wbar, color=COLOURS[m], label=m)
    ax.bar_label(bars, fmt="%.0f", fontsize=7, padding=1)
ax.set_xticks(x, names); ax.set_ylim(0, 110); ax.set_ylabel("%")
ax.set_title(f"Against {len(ref)} visually interpreted reference points on three dates")
ax.legend(frameon=False, ncol=4, fontsize=8, loc="lower center", bbox_to_anchor=(0.5, -0.3))
fig.savefig(OUT / "04-reference-accuracy.png")
plt.close(fig)

# ---------------------------------------------------------------- Step 9: cloud-free median composite
section("Step 9: cloud-free median composite, May-September 2026")
items = catalog.search(collections=["sentinel-2-l2a"], bbox=BBOX_LL, datetime=COMPOSITE_PERIOD,
                       query={"s2:mgrs_tile": {"eq": TILE}, "eo:cloud_cover": {"lt": 90}}).item_collection()
latest = {}                                          # one item per sensing time: the latest processing
for it in items:
    key = it.id.split("_")[2]
    if key not in latest or it.properties["s2:generation_time"] > latest[key].properties["s2:generation_time"]:
        latest[key] = it
items = sorted(latest.values(), key=lambda i: i.datetime)
print(f"{len(items)} scenes with tile cloud cover < 90%")
t0 = time.time()
geobox = GeoBox.from_bbox(AOI, crs=CRS, resolution=RES)
ds = odc.stac.load(items, bands=["B02", "B03", "B04", "B08", "SCL"], geobox=geobox, groupby="solar_day",
                   resampling={"SCL": "nearest", "*": "average"})
print(f"loaded {ds.sizes['time']} dates x {ds.sizes['y']} x {ds.sizes['x']} px in {time.time() - t0:.0f} s")
refl = {k: to_reflectance(ds[k].values) for k in ["B02", "B03", "B04", "B08"]}
nodata = ds["SCL"].values == 0
scl_bad = np.isin(ds["SCL"].values, [3, 8, 9, 10])

t0 = time.time()
ocm_bad = np.zeros_like(scl_bad)
for i in range(ds.sizes["time"]):
    if not nodata[i].all():
        ocm_bad[i] = omnicloudmask(refl["B04"][i], refl["B03"][i], refl["B08"][i]) > 0
print(f"OmniCloudMask on {ds.sizes['time']} dates: {time.time() - t0:.0f} s")

valid = (~nodata).sum(axis=(1, 2))
frac_scl = (scl_bad & ~nodata).sum(axis=(1, 2)) / np.maximum(valid, 1)
frac_ocm = (ocm_bad & ~nodata).sum(axis=(1, 2)) / np.maximum(valid, 1)
dates = pd.to_datetime(ds.time.values).strftime("%Y-%m-%d")
tab = pd.DataFrame({"date": dates, "SCL_%": frac_scl * 100, "OCM_%": frac_ocm * 100}).round(1)
print(tab.to_string(index=False))
print(f"dates at least 90% clear: SCL {np.sum(frac_scl < 0.1)}, OmniCloudMask {np.sum(frac_ocm < 0.1)}")


def median_composite(bad):
    out = {}
    for k, v in refl.items():
        a = v.copy()
        a[bad | nodata] = np.nan
        out[k] = np.nanmedian(a, axis=0)
    out["n"] = np.sum(~(bad | nodata), axis=0)
    return out


comp = {"No mask": median_composite(np.zeros_like(scl_bad)),
        "SCL": median_composite(scl_bad),
        "OmniCloudMask": median_composite(ocm_bad),
        "SCL + OmniCloudMask": median_composite(scl_bad | ocm_bad)}

ci = list(dates).index(CLEAR_DATE)
clear = {k: v[ci] for k, v in refl.items()}
print(f"\ncomposites compared with the clear date {CLEAR_DATE} "
      f"(SCL masks {frac_scl[ci] * 100:.1f}% of it, OmniCloudMask {frac_ocm[ci] * 100:.1f}%):")
rows = []
for name, c in comp.items():
    ndvi = (c["B08"] - c["B04"]) / (c["B08"] + c["B04"])
    rows.append({"composite": name, "median_obs": np.median(c["n"]), "min_obs": int(c["n"].min()),
                 "no_data_%": np.mean(c["n"] == 0) * 100,
                 "brighter_blue_%": np.nanmean((c["B02"] - clear["B02"]) > 0.03) * 100,
                 "darker_nir_%": np.nanmean((c["B08"] - clear["B08"]) < -0.08) * 100,
                 "median_ndvi": np.nanmedian(ndvi)})
print(pd.DataFrame(rows).round(3).to_string(index=False))

# Before and after
pick = int(np.argmin(np.abs(frac_scl - 0.5)))      # a typical half-cloudy date
fig, axes = plt.subplots(1, 3, figsize=(15, 5.3))
axes[0].imshow(rgb([refl[k][pick] for k in ["B04", "B03", "B02"]]), extent=EXTENT)
axes[0].set_title(f"Before: one scene, {dates[pick]}")
axes[1].imshow(rgb([comp["No mask"][k] for k in ["B04", "B03", "B02"]]), extent=EXTENT)
axes[1].set_title(f"Median of {ds.sizes['time']} dates, no mask")
axes[2].imshow(rgb([comp["SCL + OmniCloudMask"][k] for k in ["B04", "B03", "B02"]]), extent=EXTENT)
axes[2].set_title("After: median, SCL + OmniCloudMask masked")
for ax in axes:
    ax.set_xticks([]); ax.set_yticks([])
fig.savefig(OUT / "05-composite-before-after.png")
plt.close(fig)

fig, axes = plt.subplots(1, 2, figsize=(11.5, 4.8))
vmax = np.percentile(comp["SCL"]["n"], 99)
for ax, name in zip(axes, ["SCL", "SCL + OmniCloudMask"]):
    im = ax.imshow(comp[name]["n"], cmap="Spectral", vmin=0, vmax=vmax, extent=EXTENT, interpolation="nearest")
    ax.set_title(f"Clear observations per pixel: {name}")
    ax.set_xlabel("km east"); ax.set_ylabel("km north")
    fig.colorbar(im, ax=ax, shrink=0.8, label="dates")
fig.savefig(OUT / "06-clear-observations.png")
plt.close(fig)

# Save the composite as a Cloud Optimised GeoTIFF for reuse
profile = dict(driver="COG", dtype="float32", count=4, width=W, height=H, crs=CRS,
               transform=from_origin(AOI[0], AOI[3], RES, RES), nodata=np.nan, compress="deflate")
with rasterio.open(OUT / "lake_district_median_2026.tif", "w", **profile) as dst:
    for i, k in enumerate(["B02", "B03", "B04", "B08"], start=1):
        dst.write(comp["SCL + OmniCloudMask"][k].astype("float32"), i)
        dst.set_band_description(i, k)
print(f"wrote {OUT / 'lake_district_median_2026.tif'}")