"""Deep learning for satellite imagery, part 2: train a U-Net to segment Sentinel-2 chips, in PyTorch, on a CPU.

Loads the georeferenced 128 x 128 chips made in part 1 (Nottingham, Sentinel-2 L2A summer 2021, five classes from
Overture buildings and ESA WorldCover), trains a segmentation_models_pytorch U-Net whose ResNet-18 encoder starts
from TorchGeo's Sentinel-2 self-supervised weights (SSL4EO-S12 MoCo), compares it with the same network trained
from scratch and with a leaky random split, and reports loss curves, per-class IoU and F1 and a confusion matrix
on the spatially separate test blocks. Saves the model for part 3 (deep-learning-inference-large-rasters).
Run with:  python train_unet.py      (it runs part 1 first if the chips are missing)
"""
import os
import platform
import subprocess
import sys
import time
import urllib.request
from importlib.metadata import version
from pathlib import Path

import geopandas as gpd
import matplotlib

matplotlib.use("Agg")
import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
import rasterio
import segmentation_models_pytorch as smp
import torch
import torch.nn.functional as F
from matplotlib.colors import ListedColormap
from matplotlib.patches import Patch
from torch.utils.data import DataLoader, Dataset
from torchgeo.models import ResNet18_Weights

HERE = Path(__file__).parent
OUT = HERE / "output"
OUT.mkdir(exist_ok=True)
plt.rcParams.update({"figure.dpi": 150, "savefig.bbox": "tight", "font.size": 9})
pd.set_option("display.width", 140)
torch.set_num_threads(8)                 # a typical laptop; no GPU needed
SEED = 42

print("=== versions ===")
print(f"python {platform.python_version()} | torch {torch.__version__} | segmentation-models-pytorch {smp.__version__} | "
      f"torchgeo {version('torchgeo')} | rasterio {rasterio.__version__} | numpy {np.__version__}")
print(f"device: CPU ({torch.get_num_threads()} threads), CUDA available: {torch.cuda.is_available()}")


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


# --- Step 1: get the chips from part 1 ------------------------------------------------------------
PART1 = HERE.parent / "satellite-imagery-training-data-python"
CHIPS = Path(os.environ.get("CHIPS_DIR", PART1 / "data" / "chips"))
if not (CHIPS / "chips.gpkg").exists():
    script = PART1 / "make_training_chips.py"
    if not script.exists():                       # skipped part 1: fetch its script from the site
        PART1.mkdir(parents=True, exist_ok=True)
        url = os.environ.get("TGIS_CODE_URL", "https://www.technicalgis.com/code")
        urllib.request.urlretrieve(f"{url}/satellite-imagery-training-data-python/make_training_chips.py", script)
    print(f"chips not found: running {script.name} (part 1) first")
    subprocess.run([sys.executable, str(script)], check=True, cwd=PART1)

CLASSES = ["Building", "Other built-up", "Trees", "Grass and crops", "Water"]
COLOURS = ["#c21f3a", "#b0aaa0", "#1f5f3a", "#c9d98a", "#1543b0"]
IGNORE, NC = 255, len(CLASSES)
chips = gpd.read_file(CHIPS / "chips.gpkg")


def read(cid: str) -> tuple[np.ndarray, np.ndarray]:
    with rasterio.open(CHIPS / "images" / f"{cid}.tif") as s, rasterio.open(CHIPS / "masks" / f"{cid}.tif") as m:
        return s.read(), m.read(1)


t0 = time.time()
X, Y = map(np.stack, zip(*(read(c) for c in chips["chip_id"])))
with rasterio.open(CHIPS / "images" / f"{chips['chip_id'][0]}.tif") as s:
    BANDS = list(s.descriptions)
section("chips")
print(f"loaded {len(chips)} chips in {time.time() - t0:.1f} s: X {X.shape} {X.dtype}, Y {Y.shape} {Y.dtype}")
print(f"bands: {', '.join(BANDS)}")
print(chips["split"].value_counts().reindex(["train", "val", "test"]).to_string())


# --- Step 2: a Dataset that matches the pretrained encoder ------------------------------------------
# The SSL4EO-S12 encoder expects all 13 Sentinel-2 bands in order, divided by 10000. L2A has no B10 (cirrus),
# so we insert a channel of zeros in its place.
def to_input(x: np.ndarray) -> torch.Tensor:
    x = torch.from_numpy(x.astype("float32") / 10000)
    return torch.cat([x[..., :10, :, :], torch.zeros_like(x[..., :1, :, :]), x[..., 10:, :, :]], dim=-3)


class ChipDataset(Dataset):
    def __init__(self, idx: np.ndarray, augment: bool = False):
        self.idx, self.augment = idx, augment

    def __len__(self):
        return len(self.idx)

    def __getitem__(self, i):
        x, y = to_input(X[self.idx[i]]), torch.from_numpy(Y[self.idx[i]].astype("int64"))
        if self.augment:                                  # the 8 flips and rotations of a square
            k = int(torch.randint(4, (1,)))
            x, y = torch.rot90(x, k, (1, 2)), torch.rot90(y, k, (0, 1))
            if torch.rand(1) < 0.5:
                x, y = x.flip(2), y.flip(1)
        return x, y


# --- Step 3: a U-Net with a Sentinel-2 pretrained ResNet-18 encoder -----------------------------------
WEIGHTS = ResNet18_Weights.SENTINEL2_ALL_MOCO


def make_model(pretrained: bool, quiet: bool = True) -> torch.nn.Module:
    torch.manual_seed(SEED)
    model = smp.Unet("resnet18", encoder_weights=None, in_channels=13, classes=NC)
    if pretrained:
        sd = {k: v for k, v in WEIGHTS.get_state_dict(progress=False).items() if not k.startswith("fc.")}
        torch.nn.Module.load_state_dict(model.encoder, sd, strict=True)    # strict: every tensor must match
        if not quiet:
            print(f"loaded {WEIGHTS.name} into the encoder: {len(sd)} tensors, classifier head dropped")
    return model


section("model")
m = make_model(True, quiet=False)
n_enc = sum(p.numel() for p in m.encoder.parameters())
print(f"U-Net, ResNet-18 encoder: {sum(p.numel() for p in m.parameters()) / 1e6:.1f} M parameters "
      f"({n_enc / 1e6:.1f} M in the encoder)")


# --- Step 4: loss, metrics and the training loop -----------------------------------------------------
counts = np.bincount(Y[(chips["split"] == "train").to_numpy()].ravel(), minlength=256)[:NC]
class_w = torch.tensor(1 / np.sqrt(counts / counts.sum()), dtype=torch.float32)
class_w /= class_w.mean()
dice = smp.losses.DiceLoss("multiclass", ignore_index=IGNORE)


def loss_fn(logits, y):
    return F.cross_entropy(logits, y, weight=class_w, ignore_index=IGNORE) + dice(logits, y)


def confusion(model, loader) -> tuple[np.ndarray, float]:
    """Pixel confusion matrix (rows = reference, columns = prediction) and mean loss."""
    model.eval()
    cm, losses = np.zeros((NC, NC), "int64"), []
    with torch.no_grad():
        for x, y in loader:
            logits = model(x)
            losses.append(loss_fn(logits, y).item())
            p, y = logits.argmax(1).numpy().ravel(), y.numpy().ravel()
            ok = y != IGNORE
            cm += np.bincount(y[ok] * NC + p[ok], minlength=NC * NC).reshape(NC, NC)
    return cm, float(np.mean(losses))


def scores(cm: np.ndarray) -> pd.DataFrame:
    tp = np.diag(cm)
    fp, fn = cm.sum(0) - tp, cm.sum(1) - tp
    return pd.DataFrame({"IoU": tp / (tp + fp + fn), "F1": 2 * tp / (2 * tp + fp + fn),
                         "precision": tp / (tp + fp), "recall": tp / (tp + fn),
                         "share": cm.sum(1) / cm.sum()}, index=CLASSES)


EPOCHS, BATCH, LR = 30, 16, 1e-3


def train(name: str, pretrained: bool, split: pd.Series) -> tuple[torch.nn.Module, pd.DataFrame]:
    idx = {s: np.flatnonzero(split.to_numpy() == s) for s in ("train", "val")}
    g = torch.Generator().manual_seed(SEED)
    train_dl = DataLoader(ChipDataset(idx["train"], augment=True), batch_size=BATCH, shuffle=True, generator=g)
    val_dl = DataLoader(ChipDataset(idx["val"]), batch_size=32)
    model = make_model(pretrained)
    opt = torch.optim.AdamW(model.parameters(), lr=LR, weight_decay=1e-4)
    sched = torch.optim.lr_scheduler.OneCycleLR(opt, max_lr=LR, total_steps=EPOCHS * len(train_dl), pct_start=0.15)
    history, best, t0 = [], -1.0, time.time()
    for epoch in range(1, EPOCHS + 1):
        model.train()
        tl = []
        for x, y in train_dl:
            opt.zero_grad()
            loss = loss_fn(model(x), y)
            loss.backward()
            opt.step()
            sched.step()
            tl.append(loss.item())
        cm, vl = confusion(model, val_dl)
        miou = scores(cm)["IoU"].mean()
        history.append({"epoch": epoch, "train_loss": np.mean(tl), "val_loss": vl, "val_mIoU": miou,
                        "minutes": (time.time() - t0) / 60})
        if miou > best:
            best, state = miou, {k: v.clone() for k, v in model.state_dict().items()}
        if epoch in (1, 2, 5) or epoch % 5 == 0:
            h = history[-1]
            print(f"{name:<22} epoch {epoch:>2}: train loss {h['train_loss']:.3f} | val loss {vl:.3f} | "
                  f"val mIoU {miou:.3f} | {h['minutes']:.1f} min")
    model.load_state_dict(state)
    hist = pd.DataFrame(history)
    b = hist.loc[hist["val_mIoU"].idxmax()]
    print(f"{name}: {len(idx['train'])} training chips, {(time.time() - t0) / 60:.1f} min; "
          f"best val mIoU {b['val_mIoU']:.3f} at epoch {int(b['epoch'])}")
    return model, hist


section(f"training: {EPOCHS} epochs, batch {BATCH}, AdamW + one-cycle, weighted cross-entropy + Dice")
print("class weights: " + ", ".join(f"{c} {w:.2f}" for c, w in zip(CLASSES, class_w.tolist())))
block = chips["split"]
model, hist = train("pretrained encoder", True, block)
scratch, hist_s = train("from scratch", False, block)
# A random split of the same chips, to show how much it flatters the test score.
random_split = pd.Series(np.random.default_rng(0).choice(["train", "val", "test"], len(chips), p=[0.7, 0.15, 0.15]))
leaky, hist_r = train("random split (leaky)", True, random_split)

# --- Step 5: evaluate on the held-out test blocks ---------------------------------------------------
test_dl = DataLoader(ChipDataset(np.flatnonzero(block == "test")), batch_size=32)
cm, _ = confusion(model, test_dl)
cm_s, _ = confusion(scratch, test_dl)
cm_r, _ = confusion(leaky, DataLoader(ChipDataset(np.flatnonzero(random_split == "test")), batch_size=32))
res = scores(cm)
section(f"test set: {int((block == 'test').sum())} chips in 5 km blocks never seen in training "
        f"({cm.sum() / 1e6:.2f} M pixels)")
print(res.round(3).to_string())
summary = pd.DataFrame({
    n: {"overall accuracy": np.trace(c) / c.sum(), "mean IoU": scores(c)["IoU"].mean(),
        "macro F1": scores(c)["F1"].mean(), "Building IoU": scores(c).loc["Building", "IoU"]}
    for n, c in [("pretrained, block split", cm), ("scratch, block split", cm_s),
                 ("pretrained, random split", cm_r)]}).T
print("\n" + summary.round(3).to_string())
print("\nconfusion matrix (rows: reference label, columns: prediction, % of each reference class):")
print(pd.DataFrame(100 * cm / cm.sum(1, keepdims=True), index=CLASSES, columns=CLASSES).round(1).to_string())

# Is it the model or the labels? How often do the labels agree with themselves across a 10 m shift?
yt = Y[(block == "test").to_numpy()]
a, b2 = yt[:, :, 1:].ravel(), yt[:, :, :-1].ravel()
ok = (a != IGNORE) & (b2 != IGNORE)
print(f"\nneighbouring test pixels with the same label: {(a[ok] == b2[ok]).mean():.1%}")

ckpt = OUT / "unet_resnet18_s2_nottingham.pt"
torch.save({"state_dict": model.state_dict(), "arch": "Unet", "encoder": "resnet18", "in_channels": 13,
            "classes": CLASSES, "colours": COLOURS, "bands": BANDS, "b10_zero_index": 10, "scale": 10000,
            "chip": int(X.shape[-1]), "test_scores": res.to_dict()}, ckpt)
print(f"\nsaved {ckpt.name} ({ckpt.stat().st_size / 1e6:.0f} MB)")

# --- Figures ------------------------------------------------------------------------------------
fig, axes = plt.subplots(1, 2, figsize=(11, 3.8), constrained_layout=True)
for h, n, col in [(hist, "pretrained encoder", "#c21f3a"), (hist_s, "from scratch", "#1543b0"),
                  (hist_r, "pretrained, random split", "#a0a0a0")]:
    axes[0].plot(h["epoch"], h["train_loss"], color=col, ls="--", lw=1, label=f"{n}: train")
    axes[0].plot(h["epoch"], h["val_loss"], color=col, lw=1.6, label=f"{n}: validation")
    axes[1].plot(h["epoch"], h["val_mIoU"], color=col, lw=1.6, label=n)
axes[0].set_title("Loss (weighted cross-entropy + Dice)", loc="left")
axes[1].set_title("Validation mean IoU over five classes", loc="left")
for ax in axes:
    ax.set_xlabel("epoch")
    ax.grid(alpha=0.3)
axes[0].legend(fontsize=7, frameon=False, ncol=1)
axes[1].legend(fontsize=8, frameon=False, loc="lower right")
fig.savefig(OUT / "01-training-curves.png")
plt.close(fig)

fig, ax = plt.subplots(figsize=(6.2, 5), constrained_layout=True)
pct = 100 * cm / cm.sum(1, keepdims=True)
ax.imshow(pct, cmap="Greys", vmin=0, vmax=100)
for i in range(NC):
    for j in range(NC):
        ax.text(j, i, f"{pct[i, j]:.0f}%", ha="center", va="center", fontsize=9,
                color="white" if pct[i, j] > 55 else "#0d1315")
ax.set_xticks(range(NC), CLASSES, rotation=30, ha="right")
ax.set_yticks(range(NC), CLASSES)
ax.set_xlabel("predicted")
ax.set_ylabel("reference label")
ax.set_title("Test confusion matrix (row %: recall on the diagonal)", loc="left")
fig.savefig(OUT / "02-confusion-matrix.png")
plt.close(fig)

fig, ax = plt.subplots(figsize=(8.5, 3.6), constrained_layout=True)
xs = np.arange(NC)
for k, (c, n, col) in enumerate([(cm, "pretrained: IoU", "#c21f3a"), (cm_s, "scratch: IoU", "#1543b0")]):
    s = scores(c)
    ax.bar(xs + (k - 0.5) * 0.36, s["IoU"], 0.36, color=col, label=n)
    for x, v, f1 in zip(xs + (k - 0.5) * 0.36, s["IoU"], s["F1"]):
        ax.text(x, v + 0.01, f"{v:.2f}\nF1 {f1:.2f}", ha="center", va="bottom", fontsize=7)
ax.set_xticks(xs, CLASSES)
ax.set_ylim(0, 1.05)
ax.set_ylabel("score on the test blocks")
ax.set_title("Per-class IoU (and F1) on the spatially held-out test chips", loc="left")
ax.legend(fontsize=8, frameon=False, loc="upper left")
ax.grid(axis="y", alpha=0.3)
fig.savefig(OUT / "03-per-class-iou-f1.png")
plt.close(fig)

# Four test chips: the most built-up, the most watery, a wooded one and a mixed one.
test = chips[block == "test"]
props = test[CLASSES].to_numpy()
ent = -(props * np.log(np.clip(props, 1e-9, 1))).sum(1)
picks = list(dict.fromkeys([test["Building"].idxmax(), test["Water"].idxmax(), test["Trees"].idxmax(),
                            test.index[int(ent.argmax())]]))
cmap = ListedColormap(COLOURS)
fig, axes = plt.subplots(len(picks), 4, figsize=(9, 2.3 * len(picks)), constrained_layout=True)
model.eval()
with torch.no_grad():
    preds = model(to_input(X[picks])).argmax(1).numpy()
for r, (i, p) in enumerate(zip(picks, preds)):
    img, y = X[i], Y[i]
    rgb = np.clip(img[[3, 2, 1]].transpose(1, 2, 0) / 2500, 0, 1) ** 0.8
    err = np.where(y == IGNORE, np.nan, (p != y).astype(float))
    for ax, data, title in [(axes[r, 0], rgb, f"{chips.loc[i, 'chip_id'].split('_', 1)[1]}: image"),
                            (axes[r, 1], np.ma.masked_equal(y, IGNORE), "reference labels"),
                            (axes[r, 2], p, "U-Net prediction"),
                            (axes[r, 3], err, f"disagreement ({np.nanmean(err):.0%} of pixels)")]:
        if data is rgb:
            ax.imshow(data)
        elif data is err:
            ax.imshow(rgb)
            ax.imshow(np.ma.masked_where(~(err == 1), err), cmap=ListedColormap(["#ff00c8"]), vmin=0, vmax=1,
                      interpolation="nearest")
        else:
            ax.imshow(data, cmap=cmap, vmin=-0.5, vmax=4.5, interpolation="nearest")
        ax.set_title(title, fontsize=8, loc="left")
        ax.set_xticks([])
        ax.set_yticks([])
fig.legend(handles=[Patch(color=c, label=n) for c, n in zip(COLOURS, CLASSES)] +
           [Patch(color="#ff00c8", label="disagreement")], loc="lower center", ncol=6, fontsize=8, frameon=False,
           bbox_to_anchor=(0.5, -0.03))
fig.savefig(OUT / "04-test-predictions.png")
plt.close(fig)
print(f"figures written to {OUT.name}/")