"""Deep learning on Sentinel-2 satellite image time series for habitat classification, in PyTorch on a CPU.

Part 4 of the "AI habitat classification with Python" series (Dark Peak, Peak District, UK).
Classifies each pixel's full, irregular and cloudy 2023 Sentinel-2 time series (50 dates, 10 bands) into
nine Living England habitat classes with two temporal neural networks:

  * TempCNN (Pelletier et al. 2019): a 1D convolutional network over time;
  * L-TAE (Garnot and Landrieu 2020): a lightweight temporal attention encoder with day-of-year encoding;

each fed two ways: (1) the raw irregular dates with a cloud mask, or (2) linearly gap-filled onto a regular
10-day grid. LightGBM on the same inputs is the tabular baseline. Training uses the training blocks only
(spatial folds for early stopping), three fixed seeds, and the shared test blocks for the headline numbers.
Then: which dates matter (attention weights and month occlusion), robustness to missing dates, and a
wall-to-wall 10 m habitat map predicted from the full per-date stack, strip by strip.

Needs the shared per-sample time series first (about 4 minutes, once):
    python ../ai-habitat-classification-common/prepare_data.py --timeseries
Run with:  python habitat_time_series_dl.py
"""
import json
import os
import platform
import sys
import time
from importlib.metadata import version
from pathlib import Path

import matplotlib

matplotlib.use("Agg")
import lightgbm as lgb
import matplotlib.patheffects
import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
import torch
from matplotlib.patches import Patch
from torch import nn

T_START = time.time()
HERE = Path(__file__).resolve().parent
for p in (HERE, HERE.parent / "ai-habitat-classification-common"):   # habitat_common.py: shared series module
    sys.path.insert(0, str(p))
import habitat_common as hc  # noqa: E402

OUT = HERE / "output"
DATA = HERE / "data"                     # the full-AOI per-date stack (gitignored, about 5.6 GB)
OUT.mkdir(exist_ok=True)
DATA.mkdir(exist_ok=True)
plt.rcParams.update({"figure.dpi": 150, "savefig.bbox": "tight", "font.size": 9})
pd.set_option("display.width", 160)
pd.set_option("display.max_columns", 20)
KEYS = [c["key"] for c in hc.CLASSES]
NAMES = {c["key"]: c["name"] for c in hc.CLASSES}
COLOURS = {c["key"]: c["colour"] for c in hc.CLASSES}
BANDS = hc.S2_BANDS
N_THREADS = 8                            # a typical laptop; results were produced with 8 threads
RUNS = 3                                 # seeds 42, 43, 44; run r stops early on spatial fold r
P_DROP = 0.2                             # training augmentation: hide 20% of clear observations per batch
GRID = np.arange(1, 366, 10)             # regular 10-day day-of-year grid for gap-filling (37 steps)
torch.set_num_threads(N_THREADS)
torch.use_deterministic_algorithms(True)


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


print("=== versions ===")
print(f"python {platform.python_version()} | torch {torch.__version__} | lightgbm {lgb.__version__} | "
      f"numpy {np.__version__} | pandas {pd.__version__} | odc-stac {version('odc-stac')} | "
      f"matplotlib {matplotlib.__version__} | {N_THREADS} CPU threads, no GPU")

# --- Step 1: the irregular, cloudy time series at every sample ----------------------------------------
section("Step 1: per-date Sentinel-2 reflectance at the samples")
samples = hc.load_samples().sort_values("sample_id").reset_index(drop=True)
ts = hc.s2_sample_timeseries().sort_values(["sample_id", "date"])
dates = pd.DatetimeIndex(np.sort(ts["date"].unique()))
DOY = dates.dayofyear.to_numpy().astype("float32")
N, T, C = len(samples), len(dates), len(BANDS)
assert (ts["sample_id"].to_numpy().reshape(N, T)[:, 0] == samples["sample_id"].to_numpy()).all()
X = ts[BANDS].to_numpy("float32").reshape(N, T, C)      # reflectance, NaN where cloudy
M = ts["clear"].to_numpy(bool).reshape(N, T)            # True where the pixel was clear on that date
y = samples["class_id"].to_numpy() - 1                  # 0..8
is_train = (samples["split"] == "train").to_numpy()
is_test = ~is_train
fold = samples["fold"].to_numpy()
n_clear = M.sum(1)
print(f"{N:,} samples x {T} dates x {C} bands; {M.mean():.1%} of observations clear")
print(f"clear dates per sample: median {np.median(n_clear):.0f}, 10th percentile {np.percentile(n_clear, 10):.0f}, "
      f"min {n_clear.min()}, max {n_clear.max()}; {(n_clear < 5).sum()} samples have fewer than 5")
gaps = [np.diff(DOY[m]).max() if m.sum() > 1 else 365 for m in M]
print(f"longest gap between clear views: median {np.median(gaps):.0f} days, 90th percentile "
      f"{np.percentile(gaps, 90):.0f} days")
by_month = pd.Series(M.mean(0), index=dates).groupby(dates.month).mean()
print("share of sample observations clear, by month: " +
      ", ".join(f"{pd.Timestamp(2023, m, 1):%b} {v:.0%}" for m, v in by_month.items()))
print("median clear dates per sample by class: " +
      ", ".join(f"{k} {int(np.median(n_clear[y == i]))}" for i, k in enumerate(KEYS)))

# Standardise each band with the mean and standard deviation of clear TRAINING observations only
mu = np.nanmean(X[is_train].reshape(-1, C), axis=0)
sd = np.nanstd(X[is_train].reshape(-1, C), axis=0)
XN = np.nan_to_num((X - mu) / sd, nan=0.0).astype("float32")   # cloudy -> 0 (the training mean)


def gapfill(xn, m, t_obs=DOY, t_grid=GRID):
    """Linear interpolation of each pixel's clear observations onto a regular day-of-year grid.

    xn: (n, T, C) standardised values, m: (n, T) clear mask. Before the first or after the last clear view
    the nearest clear value is repeated; pixels with no clear view at all get 0 (the training mean).
    Fully vectorised: no Python loop over pixels.
    """
    n, nt = m.shape
    ar = np.arange(nt)
    prev = np.maximum.accumulate(np.where(m, ar, -1), axis=1)                       # last clear index <= t
    nxt = np.minimum.accumulate(np.where(m, ar, nt)[:, ::-1], axis=1)[:, ::-1]      # first clear index >= t
    j = np.searchsorted(t_obs, t_grid, side="right")            # number of dates on or before each grid day
    p = np.concatenate([np.full((n, 1), -1), prev], 1)[:, j]    # last clear view on or before the grid day
    q = np.concatenate([nxt, np.full((n, 1), nt)], 1)[:, j]     # first clear view after the grid day
    has_p, has_q = p >= 0, q < nt
    pc, qc = np.clip(p, 0, nt - 1), np.clip(q, 0, nt - 1)
    tp, tq = t_obs[pc], t_obs[qc]
    w = np.where(has_p & has_q, (t_grid - tp) / np.maximum(tq - tp, 1), np.where(has_q, 1.0, 0.0))
    xp = np.take_along_axis(xn, pc[..., None], 1)
    xq = np.take_along_axis(xn, qc[..., None], 1)
    out = (1 - w[..., None]) * xp + w[..., None] * xq
    out[~(has_p | has_q)] = 0.0
    return out.astype("float32")


def make_input(rep, xn, m):
    """(values, day of year, mask) for one input representation."""
    if rep == "gapfill":
        g = gapfill(xn, m)
        return g, np.broadcast_to(GRID.astype("float32"), g.shape[:2]), np.ones(g.shape[:2], bool)
    return xn * m[..., None], np.broadcast_to(DOY, m.shape), m


# --- Step 2: the models ------------------------------------------------------------------------------
class TempCNN(nn.Module):
    """TempCNN (Pelletier et al. 2019): three 1D convolutions over time, then a dense layer.

    With the raw irregular dates the clear mask is added as an extra input channel; the network has no
    notion of the actual dates, only of their order.
    """

    def __init__(self, c_in, t_len, n_cls, width=64, kernel=5, fc=256, drop=0.3):
        super().__init__()
        layers = []
        for i in range(3):
            layers += [nn.Conv1d(c_in if i == 0 else width, width, kernel, padding=kernel // 2),
                       nn.BatchNorm1d(width), nn.ReLU(), nn.Dropout(drop)]
        self.conv = nn.Sequential(*layers)
        self.head = nn.Sequential(nn.Flatten(), nn.Linear(width * t_len, fc), nn.BatchNorm1d(fc), nn.ReLU(),
                                  nn.Dropout(drop), nn.Linear(fc, n_cls))
        self.use_mask = None

    def forward(self, x, doy, mask):
        if self.use_mask:
            x = torch.cat([x, mask.unsqueeze(-1).float()], -1)
        return self.head(self.conv(x.transpose(1, 2)))


def doy_encoding(doy, d, tau=1000.0):
    """Sinusoidal positional encoding of the day of year (Vaswani et al. 2017, with tau = 1000 as in L-TAE)."""
    i = torch.arange(d // 2, dtype=torch.float32, device=doy.device)
    ang = doy.unsqueeze(-1) / tau ** (2 * i / d)
    return torch.cat([torch.sin(ang), torch.cos(ang)], -1)


class LTAE(nn.Module):
    """Lightweight Temporal Attention Encoder (Garnot and Landrieu 2020), for single pixels.

    Each date is embedded (linear layer + layer norm) and its day-of-year encoding added. Every head h has
    one learned "master query" q_h; its attention over dates is softmax(q_h . k_t / sqrt(d_k)), with cloudy
    dates masked out, and it pools its own slice of the embedding channels (channel grouping).
    """

    def __init__(self, c_in, n_cls, d_model=128, n_head=16, d_k=8, drop=0.2):
        super().__init__()
        self.d_model, self.n_head, self.d_k = d_model, n_head, d_k
        self.embed = nn.Sequential(nn.Linear(c_in, d_model), nn.LayerNorm(d_model))
        self.key = nn.Linear(d_model, n_head * d_k)
        self.query = nn.Parameter(torch.randn(n_head, d_k) * d_k ** -0.5)
        self.attn_drop = nn.Dropout(0.1)
        self.head = nn.Sequential(nn.LayerNorm(d_model), nn.Linear(d_model, 128), nn.BatchNorm1d(128), nn.ReLU(),
                                  nn.Dropout(drop), nn.Linear(128, 64), nn.BatchNorm1d(64), nn.ReLU(),
                                  nn.Linear(64, n_cls))

    def forward(self, x, doy, mask, return_attention=False):
        b, t, _ = x.shape
        e = self.embed(x) + doy_encoding(doy, self.d_model)                  # (b, t, d_model)
        k = self.key(e).view(b, t, self.n_head, self.d_k)
        score = torch.einsum("bthd,hd->bht", k, self.query) / self.d_k ** 0.5
        score = score.masked_fill(~mask.unsqueeze(1), -1e9)                  # cloudy dates get no weight
        a = torch.softmax(score, dim=-1)                                     # (b, heads, t)
        v = e.view(b, t, self.n_head, self.d_model // self.n_head)           # channel grouping
        o = torch.einsum("bht,bthd->bhd", self.attn_drop(a), v).reshape(b, self.d_model)
        out = self.head(o)
        return (out, a) if return_attention else out


def build(model, rep):
    t_len = len(GRID) if rep == "gapfill" else T
    if model == "tempcnn":
        net = TempCNN(C + (rep == "mask"), t_len, len(KEYS))
        net.use_mask = rep == "mask"
        return net
    return LTAE(C, len(KEYS))


def to_t(*arrs):
    return [torch.from_numpy(np.ascontiguousarray(a)) for a in arrs]


@torch.no_grad()
def predict_proba(net, rep, xn, m, batch=8192, attention=False):
    net.eval()
    probs, attn = [], []
    for i in range(0, len(xn), batch):
        x, d, k = to_t(*make_input(rep, xn[i:i + batch], m[i:i + batch]))
        if attention:
            out, a = net(x, d, k, return_attention=True)
            attn.append(a.mean(1).numpy())
        else:
            out = net(x, d, k)
        probs.append(torch.softmax(out, -1).numpy())
    p = np.concatenate(probs)
    return (p, np.concatenate(attn)) if attention else p


def macro_f1(yt, yp, k=len(KEYS)):
    cm = np.bincount(yt * k + yp, minlength=k * k).reshape(k, k).astype(float)
    with np.errstate(invalid="ignore", divide="ignore"):
        pa, ua = np.diag(cm) / cm.sum(1), np.diag(cm) / cm.sum(0)
        return float(np.nanmean(np.nan_to_num(2 * pa * ua / (pa + ua))))


def train_net(model, rep, run, max_epochs=150, patience=20, batch=128):
    """Train on the training blocks of four folds; stop early on the macro F1 of spatial fold `run`."""
    seed = hc.SEED + run
    torch.manual_seed(seed)
    rng = np.random.default_rng(seed)
    tr = np.flatnonzero(is_train & (fold != run))
    va = np.flatnonzero(is_train & (fold == run))
    net = build(model, rep)
    opt = torch.optim.AdamW(net.parameters(), lr=1e-3, weight_decay=1e-4)
    loss_fn = nn.CrossEntropyLoss(label_smoothing=0.05)
    best, best_state, wait, hist = -1.0, None, 0, []
    t0 = time.time()
    for epoch in range(max_epochs):
        net.train()
        order = rng.permutation(tr)
        tot = 0.0
        for i in range(0, len(order) - batch + 1, batch):       # drop the last partial batch (batch norm)
            idx = order[i:i + batch]
            mk = M[idx] & (rng.random((len(idx), T)) > P_DROP)  # augmentation: hide some clear views
            x, d, k = to_t(*make_input(rep, XN[idx], mk))
            loss = loss_fn(net(x, d, k), torch.from_numpy(y[idx]))
            opt.zero_grad()
            loss.backward()
            opt.step()
            tot += loss.item() * len(idx)
        f1 = macro_f1(y[va], predict_proba(net, rep, XN[va], M[va]).argmax(1))
        hist.append((epoch + 1, tot / len(order), f1))
        if f1 > best:
            best, wait = f1, 0
            best_state = {kk: vv.clone() for kk, vv in net.state_dict().items()}
        else:
            wait += 1
            if wait >= patience:
                break
    net.load_state_dict(best_state)
    best_epoch = max(hist, key=lambda h: h[2])[0]
    return net, {"val_f1": best, "epochs": len(hist), "best_epoch": best_epoch, "seconds": time.time() - t0,
                 "history": hist}


def lgb_features(rep, xn, m):
    if rep == "gapfill":
        return gapfill(xn, m).reshape(len(xn), -1)
    return np.where(m[..., None], xn, np.nan).reshape(len(xn), -1)   # LightGBM handles missing values itself


def train_lgb(rep, run):
    seed = hc.SEED + run
    tr = np.flatnonzero(is_train & (fold != run))
    va = np.flatnonzero(is_train & (fold == run))
    t0 = time.time()
    clf = lgb.LGBMClassifier(n_estimators=3000, learning_rate=0.03, num_leaves=31, min_child_samples=20,
                             subsample=0.8, subsample_freq=1, colsample_bytree=0.5, reg_lambda=1.0,
                             random_state=seed, n_jobs=N_THREADS, deterministic=True, force_row_wise=True,
                             verbose=-1)
    clf.fit(lgb_features(rep, XN[tr], M[tr]), y[tr], eval_X=(lgb_features(rep, XN[va], M[va]),), eval_y=(y[va],),
            callbacks=[lgb.early_stopping(100, verbose=False)])
    f1 = macro_f1(y[va], clf.predict(lgb_features(rep, XN[va], M[va])))
    return clf, {"val_f1": f1, "epochs": clf.best_iteration_, "best_epoch": clf.best_iteration_,
                 "seconds": time.time() - t0, "history": []}


class LGBWrap:
    """Gives a LightGBM model the same predict_proba(xn, m) interface as the networks."""

    def __init__(self, clf, rep):
        self.clf, self.rep = clf, rep

    def proba(self, xn, m):
        return self.clf.predict_proba(lgb_features(self.rep, xn, m))


def ens_proba(cfg, xn, m):
    """Average class probabilities of the three runs (seeds) of one configuration."""
    if cfg[0] == "lightgbm":
        return np.mean([w.proba(xn, m) for w in MODELS[cfg]], 0)
    return np.mean([predict_proba(net, cfg[1], xn, m) for net in MODELS[cfg]], 0)


section("Step 2: models")
for model in ("tempcnn", "ltae"):
    for rep in ("mask", "gapfill"):
        n_par = sum(p.numel() for p in build(model, rep).parameters())
        print(f"{model:<8} {rep:<8} {n_par:>9,} parameters")

# --- Step 3: train on the training blocks, three seeds, spatial early stopping ----------------------
section("Step 3: training (3 runs each: seed 42+r, early stopping on spatial fold r)")
CONFIGS = [("tempcnn", "mask"), ("tempcnn", "gapfill"), ("ltae", "mask"), ("ltae", "gapfill"),
           ("lightgbm", "mask"), ("lightgbm", "gapfill")]
LABEL = {("tempcnn", "mask"): "TempCNN, raw dates + mask", ("tempcnn", "gapfill"): "TempCNN, gap-filled",
         ("ltae", "mask"): "L-TAE, raw dates + mask", ("ltae", "gapfill"): "L-TAE, gap-filled",
         ("lightgbm", "mask"): "LightGBM, raw dates (NaN)", ("lightgbm", "gapfill"): "LightGBM, gap-filled"}
MODELS, INFO = {}, {}
for cfg in CONFIGS:
    MODELS[cfg], INFO[cfg] = [], []
    for run in range(RUNS):
        m_, info = train_lgb(cfg[1], run) if cfg[0] == "lightgbm" else train_net(*cfg, run)
        MODELS[cfg].append(LGBWrap(m_, cfg[1]) if cfg[0] == "lightgbm" else m_)
        INFO[cfg].append(info)
        unit = "trees" if cfg[0] == "lightgbm" else "epochs"
        print(f"{LABEL[cfg]:<27} run {run} (val fold {run}): val macro F1 {info['val_f1']:.3f}, "
              f"best at {info['best_epoch']} {unit}, {info['seconds']:.0f} s", flush=True)
t_train = sum(i["seconds"] for v in INFO.values() for i in v)
print(f"total training time: {t_train / 60:.1f} min")


# --- Step 4: accuracy on the test blocks ---------------------------------------------------------------
def summarise(yt, yp):
    k = len(KEYS)
    m = np.bincount(yt * k + yp, minlength=k * k).reshape(k, k).astype(float)
    n = m.sum()
    oa = np.trace(m) / n
    pe = (m.sum(0) * m.sum(1)).sum() / n ** 2
    pa = np.diag(m) / m.sum(1)
    ua = np.diag(m) / np.where(m.sum(0) == 0, np.nan, m.sum(0))
    f1 = np.nan_to_num(2 * pa * ua / (pa + ua))
    return {"cm": m.astype(int), "oa": oa, "kappa": (oa - pe) / (1 - pe), "pa": pa, "ua": ua, "f1": f1,
            "macro_f1": f1.mean()}


section("Step 4: accuracy on the 27 test blocks")
te = np.flatnonzero(is_test)
yt = y[te]
tb = samples["block_id"].to_numpy()[te]
rows, PROBA = [], {}
for cfg in CONFIGS:
    run_oa, run_f1 = [], []
    for i in range(RUNS):
        p = (MODELS[cfg][i].proba(XN[te], M[te]) if cfg[0] == "lightgbm"
             else predict_proba(MODELS[cfg][i], cfg[1], XN[te], M[te]))
        s = summarise(yt, p.argmax(1))
        run_oa.append(s["oa"])
        run_f1.append(s["macro_f1"])
    PROBA[cfg] = ens_proba(cfg, XN[te], M[te])
    s = summarise(yt, PROBA[cfg].argmax(1))
    rows.append({"model": LABEL[cfg], "val_macro_f1": np.mean([i["val_f1"] for i in INFO[cfg]]),
                 "test_oa_mean": np.mean(run_oa), "test_oa_sd": np.std(run_oa, ddof=1),
                 "test_f1_mean": np.mean(run_f1), "test_f1_sd": np.std(run_f1, ddof=1),
                 "ensemble_oa": s["oa"], "ensemble_macro_f1": s["macro_f1"], "ensemble_kappa": s["kappa"],
                 "runs_oa": run_oa})
res_df = pd.DataFrame(rows).set_index("model")
print("three runs per model: mean and s.d. of single runs, and the 3-run ensemble (averaged probabilities)")
print(res_df.drop(columns="runs_oa").round(3).to_string())
print("Part 1 spectral rules on the same test samples: OA 0.600, macro F1 0.630, kappa 0.547")

# Headline: the deep model with the best mean VALIDATION macro F1 (the test set plays no part in the choice)
deep = [c for c in CONFIGS if c[0] != "lightgbm"]
BEST = max(deep, key=lambda c: res_df.loc[LABEL[c], "val_macro_f1"])
BEST_LGB = max([c for c in CONFIGS if c[0] == "lightgbm"], key=lambda c: res_df.loc[LABEL[c], "val_macro_f1"])
print(f"chosen on validation folds: {LABEL[BEST]} (deep), {LABEL[BEST_LGB]} (LightGBM)")
res = summarise(yt, PROBA[BEST].argmax(1))
ub = np.unique(tb)
block_idx = [np.flatnonzero(tb == b) for b in ub]


def block_boot(correct, n_boot=2000, seed=hc.SEED):
    g = np.random.default_rng(seed)
    out = []
    for _ in range(n_boot):                             # resample whole test blocks, not pixels
        idx = np.concatenate([block_idx[i] for i in g.integers(0, len(ub), len(ub))])
        out.append(correct[idx].mean())
    return np.percentile(out, [2.5, 97.5])


lo, hi = block_boot(PROBA[BEST].argmax(1) == yt)
lgb_pred = PROBA[BEST_LGB].argmax(1)
lo_l, hi_l = block_boot(lgb_pred == yt)
# Paired block bootstrap of the OA difference between the deep model and LightGBM
g = np.random.default_rng(hc.SEED)
diff = (PROBA[BEST].argmax(1) == yt).astype(float) - (lgb_pred == yt)
dboot = [diff[np.concatenate([block_idx[i] for i in g.integers(0, len(ub), len(ub))])].mean() for _ in range(2000)]
print(f"\n{LABEL[BEST]} (3-run ensemble), n = {len(te):,} test samples in {len(ub)} blocks")
cm_show = pd.DataFrame(res["cm"], index=[f"ref {k}" for k in KEYS], columns=[k.split("_")[0][:5] for k in KEYS])
print("confusion matrix (rows = Living England reference, columns = prediction):")
print(cm_show.to_string())
per_class = pd.DataFrame({"n_test": res["cm"].sum(1), "producers_acc": res["pa"], "users_acc": res["ua"],
                          "f1": res["f1"]}, index=[NAMES[k] for k in KEYS])
lgb_s = summarise(yt, lgb_pred)
per_class["f1_lightgbm"] = lgb_s["f1"]
print(per_class.round(3).to_string())
print(f"overall accuracy {res['oa']:.3f} (block-bootstrap 95% CI {lo:.3f} to {hi:.3f}); "
      f"macro F1 {res['macro_f1']:.3f}; kappa {res['kappa']:.3f}")
print(f"{LABEL[BEST_LGB]}: OA {lgb_s['oa']:.3f} (95% CI {lo_l:.3f} to {hi_l:.3f}), macro F1 {lgb_s['macro_f1']:.3f}")
print(f"OA difference, deep minus LightGBM: {np.mean(diff):+.3f} (paired block-bootstrap 95% CI "
      f"{np.percentile(dboot, 2.5):+.3f} to {np.percentile(dboot, 97.5):+.3f})")
both = summarise(yt, (PROBA[BEST] + PROBA[BEST_LGB]).argmax(1))
print(f"averaging the deep and LightGBM probabilities: OA {both['oa']:.3f}, macro F1 {both['macro_f1']:.3f}")

# --- Step 5: which dates matter? ------------------------------------------------------------------------
section("Step 5: which dates matter")
LTAE_CFG = ("ltae", "mask")
att = np.mean([predict_proba(net, "mask", XN[te], M[te], attention=True)[1] for net in MODELS[LTAE_CFG]], 0)
mt = M[te]
rel = att * mt.sum(1, keepdims=True)          # attention relative to an equal share of the clear dates (=1)
rel_cls = np.array([[np.nanmean(np.where(mt[yt == i][:, j], rel[yt == i][:, j], np.nan)) if mt[yt == i][:, j].any()
                     else np.nan for j in range(T)] for i in range(len(KEYS))])
top = pd.Series(np.nanmean(np.where(mt, rel, np.nan), 0), index=dates).sort_values(ascending=False)
print("L-TAE attention relative to an equal share (1.0), averaged over test pixels where the date is clear")
print("top 6 dates: " + ", ".join(f"{d:%d %b} {v:.2f}" for d, v in top.head(6).items()))
print("bottom 3 dates: " + ", ".join(f"{d:%d %b} {v:.2f}" for d, v in top.tail(3).items()))
for i, k in enumerate(KEYS):
    j = np.nanargmax(rel_cls[i])
    print(f"  {NAMES[k]:<24} most attended date {dates[j]:%d %b} ({rel_cls[i, j]:.2f})")

OCC_CFGS = list(dict.fromkeys([("ltae", "mask"), ("tempcnn", "gapfill"), BEST_LGB]))
months = np.arange(1, 13)
occ = {}
for cfg in OCC_CFGS:
    base = (PROBA[cfg].argmax(1) == yt).mean()
    drops = []
    for mo in months:
        mo_mask = mt & (dates.month.to_numpy() != mo)[None, :]   # pretend every view in this month was cloudy
        drops.append(base - (ens_proba(cfg, XN[te], mo_mask).argmax(1) == yt).mean())
    occ[cfg] = np.array(drops)
occ_df = pd.DataFrame({LABEL[c]: v for c, v in occ.items()}, index=[f"{pd.Timestamp(2023, m, 1):%b}" for m in months])
print("drop in test OA when all views in one month are hidden (percentage points):")
print((100 * occ_df).round(1).T.to_string())

section("Step 5b: robustness to fewer clear views (test samples)")
keep_levels = [1.0, 0.75, 0.5, 0.25]
rob = {}
for cfg in CONFIGS:
    rob[cfg] = []
    for keep in keep_levels:
        g = np.random.default_rng(7)
        mk = mt & (g.random(mt.shape) < keep)
        rob[cfg].append((ens_proba(cfg, XN[te], mk).argmax(1) == yt).mean())
rob_df = pd.DataFrame({LABEL[c]: v for c, v in rob.items()}, index=[f"keep {k:.0%}" for k in keep_levels]).T
print("test OA when a random share of the clear views is kept:")
print(rob_df.round(3).to_string())

# --- Step 6: wall-to-wall map from the full per-date stack ----------------------------------------------
section("Step 6: wall-to-wall prediction")
R, W = hc.shape()
STACK = DATA / f"s2_{hc.YEAR}_stack_int16.dat"     # (dates, bands, rows, cols) reflectance x 10000
DONE = DATA / f"s2_{hc.YEAR}_stack_done.json"
t6 = time.time()
if not (STACK.exists() and DONE.exists()):
    items = hc.search_s2(f"{hc.YEAR}-01-01", f"{hc.YEAR}-12-31")
    off = hc.s2_offsets(items)
    ds = hc.load_s2(items, bands=BANDS).sel(time=dates)
    mm = np.lib.format.open_memmap(STACK, mode="w+", dtype="int16", shape=(T, C, R, W))
    for i, t in enumerate(dates):
        day = ds.sel(time=t).compute()
        for b, band in enumerate(BANDS):
            mm[i, b] = np.clip(day[band].values.astype("int32") - off[t], -32768, 32767).astype("int16")
        if (i + 1) % 10 == 0:
            print(f"  stack {i + 1}/{T} dates ({time.time() - t6:.0f} s)", flush=True)
    mm.flush()
    del mm
    DONE.write_text(json.dumps({"dates": [str(d.date()) for d in dates], "bands": BANDS}))
    print(f"downloaded the {T}-date x {C}-band stack in {time.time() - t6:.0f} s")
stack = np.load(STACK, mmap_mode="r")
print(f"stack: {stack.shape} int16 on disk = {STACK.stat().st_size / 1e9:.1f} GB (memory-mapped, read in strips)")
clear_all, _ = hc.s2_clear_masks()
clear = clear_all.sel(time=dates).values                       # (dates, rows, cols) bool
t7 = time.time()
pred = np.zeros((R, W), "uint8")
conf = np.zeros((R, W), "float32")
STRIP = 50
for r0 in range(0, R, STRIP):
    r1 = min(R, r0 + STRIP)
    xs = np.asarray(stack[:, :, r0:r1, :], dtype="float32") / 10000          # (T, C, h, W)
    xs = xs.transpose(2, 3, 0, 1).reshape(-1, T, C)
    ms = clear[:, r0:r1, :].transpose(1, 2, 0).reshape(-1, T)
    xn = np.where(ms[..., None], (xs - mu) / sd, 0.0).astype("float32")
    p = ens_proba(BEST, xn, ms)
    lab = (p.argmax(1) + 1).astype("uint8")
    lab[ms.sum(1) == 0] = 0                                     # never clear in 2023: no prediction
    pred[r0:r1] = lab.reshape(r1 - r0, W)
    conf[r0:r1] = p.max(1).reshape(r1 - r0, W)
print(f"predicted {R * W:,} pixels with {LABEL[BEST]} (3-run ensemble) in {time.time() - t7:.0f} s; "
      f"{(pred == 0).mean():.3%} never clear")
hc.write_geotiff(OUT / f"habitat_{BEST[0]}_{BEST[1]}_2023.tif", pred, nodata=0)
labels = hc.label_raster()
both_px = (labels > 0) & (pred > 0)
xc = hc.AOI_BOUNDS[0] + (np.arange(W) + 0.5) * hc.RES
yc = hc.AOI_BOUNDS[3] - (np.arange(R) + 0.5) * hc.RES
blocks = hc.load_blocks()
in_test = np.isin(hc.block_id(*np.meshgrid(xc, yc)), blocks.loc[blocks.split == "test", "block_id"].to_numpy())
print(f"wall-to-wall agreement with Living England: test blocks {(labels == pred)[both_px & in_test].mean():.3f}, "
      f"whole study area {(labels == pred)[both_px].mean():.3f} (Part 1 rules: 0.483 and 0.490)")
area = pd.DataFrame({"model_km2": [(pred == i + 1).sum() / 1e4 for i in range(len(KEYS))],
                     "living_england_km2": [(labels == i + 1).sum() / 1e4 for i in range(len(KEYS))]},
                    index=[NAMES[k] for k in KEYS])
print(area.round(1).to_string())
print(f"median top-class probability {np.median(conf[pred > 0]):.2f}; {(conf[pred > 0] < 0.5).mean():.1%} of pixels below 0.5")

# --- Metrics for the series comparison ------------------------------------------------------------------
t_figs = time.time()
metrics = {
    "oa": round(float(res["oa"]), 4), "oa_ci95": [round(float(lo), 4), round(float(hi), 4)],
    "macro_f1": round(float(res["macro_f1"]), 4), "kappa": round(float(res["kappa"]), 4),
    "per_class": {k: {"pa": round(float(res["pa"][i]), 4), "ua": round(float(res["ua"][i]), 4),
                      "f1": round(float(res["f1"][i]), 4)} for i, k in enumerate(KEYS)},
    "n_test": int(len(te)), "runtime_s": None,
    "method": f"{LABEL[BEST]} on the 2023 Sentinel-2 time series (50 dates, 10 bands), PyTorch on CPU, "
              f"ensemble of 3 seeds",
    "oa_runs": [round(float(v), 4) for v in res_df.loc[LABEL[BEST], "runs_oa"]],
    "lightgbm_baseline": {"method": LABEL[BEST_LGB], "oa": round(float(lgb_s["oa"]), 4),
                          "oa_ci95": [round(float(lo_l), 4), round(float(hi_l), 4)],
                          "macro_f1": round(float(lgb_s["macro_f1"]), 4), "kappa": round(float(lgb_s["kappa"]), 4)},
}

# --- Figures ---------------------------------------------------------------------------------------------
ext = (hc.AOI_BOUNDS[0], hc.AOI_BOUNDS[2], hc.AOI_BOUNDS[1], hc.AOI_BOUNDS[3])
cmap, norm = hc.class_cmap()
legend = [Patch(color=COLOURS[k], label=NAMES[k]) for k in KEYS]
MODEL_COL = {("tempcnn", "mask"): "#fdae61", ("tempcnn", "gapfill"): "#d7191c", ("ltae", "mask"): "#2b83ba",
             ("ltae", "gapfill"): "#abdda4", ("lightgbm", "mask"): "#bababa", ("lightgbm", "gapfill"): "#404040"}
month_ticks = [pd.Timestamp(2023, m, 1).dayofyear for m in months]
month_lab = [f"{pd.Timestamp(2023, m, 1):%b}" for m in months]
B04, B08 = BANDS.index("B04"), BANDS.index("B08")
with np.errstate(invalid="ignore", divide="ignore"):
    NDVI = (X[..., B08] - X[..., B04]) / (X[..., B08] + X[..., B04])

# Figure 1: the irregular time series and gap-filling
fig = plt.figure(figsize=(12, 7.6))
gs = fig.add_gridspec(2, 2, height_ratios=[1, 1.5], hspace=0.38, wspace=0.18)
ax = fig.add_subplot(gs[0, :])
ax.bar(DOY, 100 * M.mean(0), width=4, color="#2b83ba")
ax.set_xlim(0, 366)
ax.set_xticks(month_ticks, month_lab)
ax.set_ylabel("Sample pixels clear (%)")
ax.set_title(f"{T} Sentinel-2 dates in 2023 with at least 5% of the study area clear: on average only "
             f"{M.mean():.0%} of the 9,000 sample pixels are clear on a date")
ax = fig.add_subplot(gs[1, 0])
g = np.random.default_rng(3)
egs = []
for k in ("bracken", "improved_grass", "bog"):
    i = KEYS.index(k)
    cand = np.flatnonzero((y == i) & is_train & (n_clear >= np.median(n_clear) - 2) & (n_clear <= np.median(n_clear) + 2)
                    & (np.where(M, NDVI, 1).min(1) > 0.1))  # skip pixels with a residual snow or shadow view
    egs.append((k, g.choice(cand)))
for k, s_ in egs:
    gf = gapfill(XN[s_:s_ + 1], M[s_:s_ + 1])[0] * sd + mu
    nd_gf = (gf[:, B08] - gf[:, B04]) / (gf[:, B08] + gf[:, B04])
    col = COLOURS[k] if k != "improved_grass" else "#4d9221"
    ax.plot(GRID, nd_gf, "-", color=col, lw=1.4, label=f"{NAMES[k]} (pixel {s_}): gap-filled, 10-day grid")
    ax.plot(DOY[M[s_]], NDVI[s_][M[s_]], "o", color=col, mec="black", mew=0.6, ms=5)
ax.set_xlim(0, 366)
ax.set_xticks(month_ticks, month_lab)
ax.set_ylabel("NDVI")
ax.set_title("Three training pixels: clear views (dots) and linear gap-filling (lines)", fontsize=9)
ax.legend(fontsize=7, loc="upper center", bbox_to_anchor=(0.5, -0.08), frameon=False)
ax = fig.add_subplot(gs[1, 1])
gfa = gapfill(XN[is_train], M[is_train]) * sd + mu
nd_all = (gfa[..., B08] - gfa[..., B04]) / (gfa[..., B08] + gfa[..., B04])
for i, k in enumerate(KEYS):
    if k == "water":
        continue
    med = np.median(nd_all[y[is_train] == i], 0)
    ax.plot(GRID, med, lw=1.6, color=COLOURS[k], label=NAMES[k],
            path_effects=[matplotlib.patheffects.withStroke(linewidth=2.6, foreground="#555555")]
            if k in ("seminat_grass", "heath") else None)
ax.set_xlim(0, 366)
ax.set_xticks(month_ticks, month_lab)
ax.set_title("Median gap-filled NDVI by class (training pixels; water omitted)", fontsize=9)
ax.legend(fontsize=7, ncol=3, loc="upper center", bbox_to_anchor=(0.5, -0.08), frameon=False)
fig.savefig(OUT / "01-irregular-time-series.png")
plt.close(fig)

# Figure 2: model comparison and robustness to missing dates
fig, axes = plt.subplots(1, 2, figsize=(12.5, 4.8), gridspec_kw={"width_ratios": [1.25, 1]})
ax = axes[0]
pos = np.arange(len(CONFIGS))
for i, cfg in enumerate(CONFIGS):
    r = res_df.loc[LABEL[cfg]]
    ax.bar(i, r["ensemble_oa"], 0.62, color=MODEL_COL[cfg], edgecolor="#333333", lw=0.6)
    ax.plot([i] * RUNS, r["runs_oa"], "o", color="white", mec="black", ms=4.5, zorder=3)
    ax.text(i, max(r["ensemble_oa"], max(r["runs_oa"])) + 0.012, f"{r['ensemble_oa']:.1%}", ha="center", fontsize=8)
ax.axhline(0.600, color="#7b3294", ls="--", lw=1.3)
ax.text(-0.45, 0.594, "Part 1 spectral rules 60.0%", ha="left", va="top", fontsize=8, color="#7b3294",
        bbox={"facecolor": "white", "edgecolor": "none", "alpha": 0.85, "pad": 1.5}, zorder=4)
ax.set_xticks(pos, [LABEL[c].replace(", ", "\n") for c in CONFIGS], fontsize=7.5)
ax.set_ylim(0.5, max(res_df["ensemble_oa"]) + 0.06)
ax.set_ylabel("Overall accuracy, 2,456 test samples")
ax.set_title("Bars: 3-seed ensemble; dots: single runs", fontsize=9)
ax = axes[1]
for cfg in CONFIGS:
    ax.plot([100 * k for k in keep_levels], rob[cfg], "o-", color=MODEL_COL[cfg], label=LABEL[cfg],
            lw=2 if cfg[0] != "lightgbm" else 1.4, ls="-" if cfg[1] == "mask" else "--")
ax.invert_xaxis()
ax.set_xticks([100, 75, 50, 25])
ax.set_xlabel("Share of each test pixel's clear views kept (%)")
ax.set_ylabel("Overall accuracy")
ax.set_title("Robustness: randomly hide clear views at test time", fontsize=9)
ax.legend(fontsize=7, loc="lower left")
fig.tight_layout()
fig.savefig(OUT / "02-model-comparison.png")
plt.close(fig)

# Figure 3: L-TAE attention by date and class
fig, axes = plt.subplots(2, 1, figsize=(12, 6.4), height_ratios=[1, 3.2], sharex=True)
axes[0].bar(np.arange(T), top.reindex(dates).to_numpy(), color="#2b83ba", width=0.8)
axes[0].axhline(1, color="#555555", lw=0.8, ls=":")
axes[0].set_ylabel("All classes")
axes[0].set_title("L-TAE attention relative to an equal share of a pixel's clear dates (1 = equal share), "
                  "test pixels where the date is clear", fontsize=9)
im = axes[1].imshow(rel_cls, aspect="auto", cmap="magma", vmin=0, vmax=np.nanpercentile(rel_cls, 98),
                    interpolation="nearest")
axes[1].set_yticks(range(len(KEYS)), [NAMES[k] for k in KEYS], fontsize=8)
step = 2
axes[1].set_xticks(range(0, T, step), [f"{d:%d %b}" for d in dates[::step]], rotation=90, fontsize=7)
cb = fig.colorbar(im, ax=axes, shrink=0.75, pad=0.01)
cb.set_label("Attention / equal share")
fig.savefig(OUT / "03-attention-by-date.png")
plt.close(fig)

# Figure 4: month occlusion
fig, ax = plt.subplots(figsize=(10, 3.8))
wbar = 0.8 / len(OCC_CFGS)
for i, cfg in enumerate(OCC_CFGS):
    ax.bar(months + (i - (len(OCC_CFGS) - 1) / 2) * wbar, 100 * occ[cfg], wbar, color=MODEL_COL[cfg],
           edgecolor="#333333", lw=0.5, label=LABEL[cfg])
ax.axhline(0, color="#555555", lw=0.8)
ax.set_xticks(months, month_lab)
ax.set_ylabel("Drop in test OA (percentage points)")
ax.set_title("Occlusion: hide every clear view in one month and re-predict the 2,456 test samples", fontsize=9)
ax.legend(fontsize=8)
fig.tight_layout()
fig.savefig(OUT / "04-month-occlusion.png")
plt.close(fig)


# Figure 5: wall-to-wall map against Living England
def km_axes(ax):
    ax.set_xticks(np.arange(400000, 425001, 5000))
    ax.set_yticks(np.arange(380000, 400001, 5000))
    ax.set_xticklabels([f"{v / 1000:.0f}" for v in ax.get_xticks()])
    ax.set_yticklabels([f"{v / 1000:.0f}" for v in ax.get_yticks()])
    ax.set_xlabel("Easting (km, British National Grid)")
    ax.set_ylabel("Northing (km)")


fig, axes = plt.subplots(1, 2, figsize=(13, 5.6), sharey=True)
axes[0].imshow(np.ma.masked_equal(pred, 0), cmap=cmap, norm=norm, extent=ext, interpolation="nearest")
axes[0].set_title(f"{LABEL[BEST]}: 2023 Sentinel-2 time series, 10 m")
axes[1].imshow(np.ma.masked_equal(labels, 0), cmap=cmap, norm=norm, extent=ext, interpolation="nearest")
axes[1].set_title(f"Living England 2022-23 (pixel agreement: {(labels == pred)[both_px].mean():.0%})")
for ax in axes:
    km_axes(ax)
axes[1].set_ylabel("")
fig.legend(handles=legend, loc="lower center", ncol=9, fontsize=8, frameon=False, bbox_to_anchor=(0.5, -0.04))
fig.savefig(OUT / "05-habitat-map.png")
plt.close(fig)

# Figure 6: confusion matrix
m = res["cm"]
rown = m / m.sum(1, keepdims=True)
fig, ax = plt.subplots(figsize=(8.6, 7.4))
ax.imshow(rown, cmap="YlGnBu", vmin=0, vmax=1)
for i in range(len(KEYS)):
    for j in range(len(KEYS)):
        if m[i, j]:
            ax.text(j, i, m[i, j], ha="center", va="center", fontsize=8, color="white" if rown[i, j] > 0.55 else "black")
short = [NAMES[k].replace(" woodland", "").replace(" grassland", " grass") for k in KEYS]
ax.set_xticks(range(len(KEYS)), [f"{s}\nUA {u:.0%}" for s, u in zip(short, res["ua"])], rotation=45, ha="right",
              fontsize=7.5)
ax.set_yticks(range(len(KEYS)), [f"{s}  PA {p:.0%}" for s, p in zip(short, res["pa"])], fontsize=7.5)
ax.set_xlabel(f"{LABEL[BEST]} (UA = user's accuracy)")
ax.set_ylabel("Living England reference (PA = producer's accuracy)")
ax.set_title(f"Test blocks: overall accuracy {res['oa']:.1%} (block-bootstrap 95% CI {lo:.1%} to {hi:.1%}), "
             f"kappa {res['kappa']:.2f}", fontsize=9)
fig.savefig(OUT / "06-confusion-matrix.png")
plt.close(fig)

runtime = time.time() - T_START
metrics["runtime_s"] = round(runtime, 1)
(OUT / "metrics.json").write_text(json.dumps(metrics, indent=2))
print(f"\nsaved 6 figures, metrics.json and the GeoTIFF to {OUT.name}/")
print(f"total runtime {runtime / 60:.1f} min (training {t_train / 60:.1f} min, stack and wall-to-wall "
      f"{(t_figs - t6) / 60:.1f} min)")