"""Process large satellite rasters in parallel, out of core, with dask and rioxarray.

Eight full Sentinel-2 scenes of tile 30UXC (London and the Thames Valley, 2025) are reduced
to one median NDVI map: 16 files of 10980 x 10980 pixels, 1.93 billion pixels in total.

Steps
  1. find eight low-cloud scenes with STAC and download bands B04 and B08 once
  2. open them lazily with rioxarray and dask chunks, and choose a chunk size
  3. build the median NDVI as a lazy task graph and inspect it
  4. run it with dask diagnostics (task stream, CPU and memory)
  5. compare peak memory with eager NumPy
  6. time chunk sizes, thread counts, the read lock and a LocalCluster
  7. write the result to disk chunk by chunk as a tiled GeoTIFF

Run:  python dask_raster.py      (about 25 minutes; downloads 4.6 GB on the first run)
"""
import json
import os
import subprocess
import sys
import threading
import time
from concurrent.futures import ThreadPoolExecutor

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 dask
import distributed
import matplotlib
import networkx as nx
import numpy as np
import planetary_computer as pc
import psutil
import pystac_client
import rasterio
import requests
import rioxarray
import xarray as xr

matplotlib.use("Agg")
import matplotlib.pyplot as plt
from dask.core import get_dependencies
from dask.diagnostics import Profiler, ProgressBar
from matplotlib.patches import Rectangle

from raster_pipeline import (DATA, DATES, HERE, OFFSET, TILE, RSSSampler, band_path,
                             median_ndvi, open_stack)

OUT = HERE / "output"
OUT.mkdir(exist_ok=True)
DATA.mkdir(exist_ok=True)
STAC_URL = "https://planetarycomputer.microsoft.com/api/stac/v1"
BANDS = ["B04", "B08"]
KIND_COLOURS = {"open file": "#5e4fa2", "read chunk": "#3288bd", "elementwise": "#abdda4",
                "concatenate": "#fdae61", "rechunk": "#f46d43", "median": "#9e0142"}
plt.rcParams.update({"savefig.dpi": 150, "font.size": 10, "axes.spines.top": False, "axes.spines.right": False})
T_START = time.perf_counter()


def run_bench(mode, n_dates, chunk=2048, workers=8, lock=False):
    """Run one measurement in a fresh Python process (see raster_pipeline.py)."""
    args = [sys.executable, str(HERE / "raster_pipeline.py"), mode, str(n_dates), str(chunk), str(workers), str(int(lock))]
    out = subprocess.run(args, capture_output=True, text=True, check=True, cwd=HERE).stdout
    return json.loads(out.strip().splitlines()[-1])


def kind(key):
    """Group dask task names into a few readable types."""
    name = str(key[0] if isinstance(key, tuple) else key)
    for word, label in [("original", "open file"), ("open_rasterio", "read chunk"), ("getitem", "read chunk"),
                        ("median", "median"), ("rechunk", "rechunk"), ("concatenate", "concatenate")]:
        if word in name:
            return label
    return "elementwise"


print("=== versions ===")
print(f"python {sys.version.split()[0]} | dask {dask.__version__} | distributed {distributed.__version__} | "
      f"xarray {xr.__version__} | rioxarray {rioxarray.__version__} | rasterio {rasterio.__version__} | "
      f"numpy {np.__version__}")
print(f"machine: {psutil.cpu_count(logical=False)} cores / {psutil.cpu_count()} threads, "
      f"{psutil.virtual_memory().total / 1e9:.0f} GB RAM")

# ------------------------------------------------------------------ Step 1: data
print("\n=== Step 1: eight low-cloud Sentinel-2 scenes of tile 30UXC (2025) ===")
catalog = pystac_client.Client.open(STAC_URL, modifier=pc.sign_inplace)


def find_item(date):
    # Two orbits can cross the tile on one day; keep the one that covers it fully, latest processing first
    items = catalog.search(collections=["sentinel-2-l2a"], datetime=date,
                           query={"s2:mgrs_tile": {"eq": TILE}, "s2:nodata_pixel_percentage": {"lt": 1}}).item_collection()
    return sorted(items, key=lambda i: i.properties["s2:generation_time"])[-1]


def download(href, dest):
    if dest.exists():
        return 0
    tmp = dest.with_suffix(".part")
    with requests.get(href, stream=True, timeout=120) as r:
        r.raise_for_status()
        with open(tmp, "wb") as f:
            for block in r.iter_content(1 << 22):
                f.write(block)
    tmp.rename(dest)
    return dest.stat().st_size


jobs = []
for d in DATES:
    it = find_item(d)
    print(f"{d}  {it.id}  tile cloud {it.properties['eo:cloud_cover']:.1f}%  "
          f"no data {it.properties['s2:nodata_pixel_percentage']:.4f}%")
    jobs += [(it.assets[b].href, band_path(d, b)) for b in BANDS]
t0 = time.perf_counter()
with ThreadPoolExecutor(4) as ex:
    fetched = sum(ex.map(lambda j: download(*j), jobs))
total = sum(p.stat().st_size for _, p in jobs)
print(f"{len(jobs)} files on disk, {total / 1e9:.2f} GB compressed "
      f"({fetched / 1e9:.2f} GB downloaded this run, {time.perf_counter() - t0:.0f} s)")
with rasterio.open(band_path(DATES[0], "B08")) as src:
    print(f"each file: {src.width} x {src.height} px, {src.dtypes[0]}, {src.profile['compress']}, "
          f"internal tiles {src.block_shapes[0]}, overviews {src.overviews(1)}")
    tb = src.bounds
raw_gb = len(DATES) * len(BANDS) * 10980 * 10980 * 2 / 1e9
print(f"uncompressed: {len(DATES)} dates x {len(BANDS)} bands x 10980 x 10980 uint16 = {raw_gb:.2f} GB "
      f"({raw_gb * 2:.2f} GB as float32)")
EXT = [tb.left / 1000, tb.right / 1000, tb.bottom / 1000, tb.top / 1000]

# ------------------------------------------------------------------ Step 2: lazy open, chunks
print("\n=== Step 2: open lazily with rioxarray and choose a chunk size ===")
t0 = time.perf_counter()
red, nir = open_stack(DATES, 2048)
print(f"opened {len(jobs)} files in {time.perf_counter() - t0:.2f} s without reading any pixels")
print(red)
print("\nchunk size options for a 10980 x 10980 band:")
for c in [256, 512, 1024, 2048, 4096]:
    n = int(np.ceil(10980 / c)) ** 2
    print(f"  {c:>4} px: {n:>5} chunks per band, {c * c * 4 / 1e6:5.1f} MB per float32 chunk, "
          f"{c * c * 4 * len(DATES) / 1e6:5.0f} MB per {len(DATES)}-date block for the median")

# Figure 1: the chunk grid over the raster, from the COG overviews (cheap to read)
with rasterio.open(band_path(DATES[2], "B08")) as s8, rasterio.open(band_path(DATES[2], "B04")) as s4:
    ov8 = s8.read(1, out_shape=(686, 686)).astype("float32")
    ov4 = s4.read(1, out_shape=(686, 686)).astype("float32")
    win = rasterio.windows.Window(4096, 4096, 2048, 2048)
    z8 = s8.read(1, window=win).astype("float32")
    z4 = s4.read(1, window=win).astype("float32")
with np.errstate(divide="ignore", invalid="ignore"):
    ndvi_ov = np.where((ov8 > 0) & (ov4 > 0), (ov8 - ov4) / (ov8 + ov4 - 2 * OFFSET), np.nan)
    ndvi_z = (z8 - z4) / (z8 + z4 - 2 * OFFSET)
fig, axes = plt.subplots(1, 2, figsize=(13, 6.3))
ax = axes[0]
im = ax.imshow(ndvi_ov, cmap="Spectral", vmin=-0.2, vmax=0.9, extent=EXT)
for k in range(2048, 10980, 2048):
    ax.axvline(EXT[0] + k / 100, color="black", lw=1.1)
    ax.axhline(EXT[3] - k / 100, color="black", lw=1.1)
for y in range(0, 10980, 2048):
    for x in range(0, 10980, 2048):
        h, w = min(2048, 10980 - y), min(2048, 10980 - x)
        if h < 2048 or w < 2048:
            ax.text(EXT[0] + (x + w / 2) / 100, EXT[3] - (y + h / 2) / 100, f"{h}×{w}", ha="center", va="center",
                    fontsize=6.5, rotation=90 if w < 2048 and h == 2048 else 0,
                    bbox=dict(fc="white", ec="none", alpha=0.75, pad=0.5))
ax.add_patch(Rectangle((EXT[0] + 40.96, EXT[3] - 61.44), 20.48, 20.48, fill=False, ec="#9e0142", lw=2.6))
ax.set_title("Tile 30UXC NDVI (5 April 2025) in 2048 px dask chunks:\n"
             f"36 chunks per band; edge chunks are {10980 - 5 * 2048} px wide", fontsize=10)
ax.set_xlabel("Easting (km, UTM 30N)")
ax.set_ylabel("Northing (km)")
ax = axes[1]
ax.imshow(ndvi_z, cmap="Spectral", vmin=-0.2, vmax=0.9, extent=[0, 2048, 2048, 0])
for k in range(512, 2048, 512):
    ax.axvline(k, color="black", lw=0.9, ls="--")
    ax.axhline(k, color="black", lw=0.9, ls="--")
ax.set_title("The outlined chunk at full 10 m resolution:\n"
             "16 of the file's 512 px internal tiles (dashed)", fontsize=10)
ax.set_xlabel("column in chunk (px)")
ax.set_ylabel("row in chunk (px)")
cb = fig.colorbar(im, ax=axes, shrink=0.8, pad=0.02)
cb.set_label("NDVI")
fig.savefig(OUT / "01-chunk-grid.png", bbox_inches="tight", facecolor="white")
plt.close(fig)

# ------------------------------------------------------------------ Step 3: the task graph
print("\n=== Step 3: build the median NDVI as a lazy task graph ===")
med = median_ndvi(red, nir)
print(med)
print(f"\n2048 px chunks: {len(med.data.__dask_graph__())} tasks in {len(med.data.__dask_graph__().layers)} layers; "
      f"output {med.nbytes / 1e6:.0f} MB in {med.data.npartitions} chunks")
for c in [256, 512, 1024, 2048, 4096]:
    m_ = median_ndvi(*open_stack(DATES, c))
    print(f"  {c:>4} px chunks: {len(m_.data.__dask_graph__()):>6} tasks, {m_.data.npartitions:>4} output chunks")

# A tiny version of the same pipeline (two dates, 2 x 2 chunks) is small enough to draw.
small = median_ndvi(*open_stack(DATES[:2], 5490))
dsk = dict(small.data.__dask_graph__())
(optimised,) = dask.optimize(small.data)
print(f"two dates in 2 x 2 chunks: {len(dsk)} tasks as written, {len(optimised.__dask_graph__())} after dask fuses "
      f"chains of element-wise steps")
G = nx.DiGraph()
for k in dsk:
    G.add_node(k)
    G.add_edges_from((dep, k) for dep in get_dependencies(dsk, k) if dep in dsk)
for i, generation in enumerate(nx.topological_generations(G)):
    for nd in generation:
        G.nodes[nd]["layer"] = i
pos = nx.multipartite_layout(G, subset_key="layer", align="horizontal")
pos = {k: (x, -y) for k, (x, y) in pos.items()}
kinds = {}
for k in G:
    kinds[kind(k)] = kinds.get(kind(k), 0) + 1
print(f"graph as written: {G.number_of_nodes()} tasks, {G.number_of_edges()} dependencies, "
      f"{nx.dag_longest_path_length(G) + 1} tasks on the longest path")
print("task types: " + ", ".join(f"{k} {v}" for k, v in sorted(kinds.items())))
fig, ax = plt.subplots(figsize=(13, 6))
nx.draw_networkx_edges(G, pos, ax=ax, edge_color="#999999", arrows=False, width=0.7)
for kd, col in KIND_COLOURS.items():
    nodes = [k for k in G if kind(k) == kd]
    if nodes:
        nx.draw_networkx_nodes(G, pos, nodelist=nodes, node_color=col, node_size=80, ax=ax,
                               edgecolors="black", linewidths=0.4, label=f"{kd} ({len(nodes)})")
ax.legend(loc="lower left", title="task type (count)", fontsize=9)
ax.set_title(f"Task graph for two dates in 2 × 2 chunks ({G.number_of_nodes()} tasks, read from top to bottom)")
ax.axis("off")
fig.savefig(OUT / "02-task-graph.png", bbox_inches="tight", facecolor="white")
plt.close(fig)

# ------------------------------------------------------------------ Step 4: diagnostics
print("\n=== Step 4: compute with dask diagnostics (8 threads, 2048 px chunks) ===")
with dask.config.set(scheduler="threads", num_workers=8):
    with Profiler() as prof, RSSSampler(interval=0.1) as rprof, ProgressBar(dt=10, out=sys.stdout):
        t0 = time.perf_counter()
        result = med.compute()
        diag_secs = time.perf_counter() - t0
mem = [r / 1e6 for _, r in rprof.trace]   # MB
print(f"computed in {diag_secs:.1f} s: {len(prof.results)} tasks; process memory {min(mem) / 1e3:.2f} GB at start, "
      f"{max(mem) / 1e3:.2f} GB at peak; mean CPU {np.mean(rprof.cpu[1:]):.0f}% (800% = 8 busy threads)")
vals = result.values
print(f"median NDVI: mean {np.nanmean(vals):.3f}; {np.isfinite(vals).mean() * 100:.2f}% of pixels valid; "
      f"{(vals > 0.6).sum() / 1e4:,.0f} km² above 0.6 and {(vals < 0.1).sum() / 1e4:,.0f} km² below 0.1")
del result, vals

t_first = min(r.start_time for r in prof.results)
wmap = {w: i for i, w in enumerate(sorted({r.worker_id for r in prof.results}))}
fig, axes = plt.subplots(2, 1, figsize=(13, 7), sharex=True, gridspec_kw={"height_ratios": [2.2, 1]})
ax = axes[0]
for kd, col in KIND_COLOURS.items():
    rows = [r for r in prof.results if kind(r.key) == kd]
    for w in wmap:
        bars = [(r.start_time - t_first, r.end_time - r.start_time) for r in rows if r.worker_id == w]
        if bars:
            ax.broken_barh(bars, (wmap[w] - 0.4, 0.8), facecolors=col, edgecolor="none")
    if rows:
        ax.bar(0, 0, color=col, label=f"{kd} ({len(rows)} tasks)")
ax.set_yticks(range(len(wmap)))
ax.set_yticklabels([f"thread {i + 1}" for i in range(len(wmap))])
ax.invert_yaxis()
ax.set_title(f"Task stream from dask.diagnostics.Profiler: {len(prof.results)} tasks on 8 threads in {diag_secs:.0f} s")
ax.legend(loc="upper center", bbox_to_anchor=(0.5, -0.04), fontsize=8.5, ncol=5, frameon=False)
ax = axes[1]
t_r = np.array([t for t, _ in rprof.trace])
ax.plot(t_r, np.array(mem) / 1e3, color="#3288bd", lw=2)
ax.set_ylabel("memory (GB)", color="#3288bd")
ax.set_ylim(0, None)
ax2 = ax.twinx()
ax2.plot(t_r, rprof.cpu, color="#f46d43", lw=1)
ax2.set_ylabel("CPU (%)", color="#f46d43")
ax2.spines["right"].set_visible(True)
ax2.set_ylim(0, None)
ax.set_xlabel("seconds after compute()")
ax.set_title("Process memory (blue) and CPU (orange), sampled with psutil every 0.1 s")
fig.tight_layout()
fig.savefig(OUT / "03-dask-diagnostics.png", bbox_inches="tight", facecolor="white")
plt.close(fig)

# ------------------------------------------------------------------ Step 5: memory
print("\n=== Step 5: peak memory, eager NumPy against dask (fresh process per run) ===")
mem_rows = []
for n in [1, 2, 4]:
    r = run_bench("eager", n)
    mem_rows.append(r)
    print(f"eager NumPy, {n} date(s): {r['seconds']:6.1f} s, peak {r['peak_gb']:5.2f} GB  (mean NDVI {r['result_mean']})")
for n in [1, 2, 4, 8]:
    r = run_bench("dask", n)
    mem_rows.append(r)
    print(f"dask,        {n} date(s): {r['seconds']:6.1f} s, peak {r['peak_gb']:5.2f} GB  (mean NDVI {r['result_mean']})")
eager = [r for r in mem_rows if r["mode"] == "eager"]
slope, icpt = np.polyfit([r["n_dates"] for r in eager], [r["peak_gb"] for r in eager], 1)
print(f"eager peak grows by {slope:.2f} GB per date: about {icpt + slope * 8:.0f} GB for 8 dates "
      f"and {icpt + slope * 70:.0f} GB for the 70 or so scenes a year of this tile")

fig, axes = plt.subplots(1, 2, figsize=(13, 4.8))
ax = axes[0]
for r in mem_rows:
    if (r["mode"], r["n_dates"]) in {("eager", 4), ("dask", 4), ("dask", 8)}:
        t, m = zip(*r["trace"])
        col = {"eager": "#d53e4f", 4: "#66c2a5", 8: "#3288bd"}["eager" if r["mode"] == "eager" else r["n_dates"]]
        ax.plot(t, m, color=col, lw=2, label=f"{'eager NumPy' if r['mode'] == 'eager' else 'dask'}, "
                                             f"{r['n_dates']} dates (peak {r['peak_gb']:.1f} GB)")
ax.set_xlabel("seconds")
ax.set_ylabel("memory above the start (GB)")
ax.set_title("Memory during one run")
ax.legend(fontsize=8.5)
ax = axes[1]
ax.plot([r["n_dates"] for r in eager], [r["peak_gb"] for r in eager], "o-", color="#d53e4f", lw=2,
        label="eager NumPy (measured)")
ax.plot([4, 8], [icpt + slope * 4, icpt + slope * 8], ":", color="#d53e4f", lw=1.6,
        label=f"eager NumPy, linear trend ({slope:.1f} GB per date)")
dk = [r for r in mem_rows if r["mode"] == "dask"]
ax.plot([r["n_dates"] for r in dk], [r["peak_gb"] for r in dk], "s-", color="#3288bd", lw=2,
        label="dask, 2048 px chunks, 8 threads (measured)")
ax.set_xticks([1, 2, 4, 8])
ax.set_xlabel("dates in the median")
ax.set_ylabel("peak memory (GB)")
ax.set_title("Peak memory against the size of the stack")
ax.legend(fontsize=8.5)
fig.tight_layout()
fig.savefig(OUT / "04-memory-eager-vs-dask.png", bbox_inches="tight", facecolor="white")
plt.close(fig)

# ------------------------------------------------------------------ Step 6: timings
print("\n=== Step 6: time chunk sizes and thread counts (8 dates, fresh process per run) ===")
CHUNKS = [256, 512, 1024, 2048, 4096]
WORKERS = [1, 2, 4, 8, 16]
timing = []
for c in CHUNKS:
    for w in WORKERS:
        if w == 1 and c not in (512, 2048):
            continue                              # single-threaded runs are slow; two are enough to show scaling
        r = run_bench("dask", 8, chunk=c, workers=w)
        r.pop("trace")
        timing.append(r)
        print(f"{c:>4} px chunks, {w:>2} threads: {r['seconds']:6.1f} s, peak {r['peak_gb']:5.2f} GB")
lock_run = run_bench("dask", 8, chunk=2048, workers=8, lock=True)
print(f"2048 px chunks,  8 threads, lock=True: {lock_run['seconds']:6.1f} s, peak {lock_run['peak_gb']:5.2f} GB")
dist_run = run_bench("distributed", 8, chunk=2048, workers=2)
print(f"LocalCluster, 4 processes x 2 threads, 2048 px chunks: {dist_run['seconds']:6.1f} s "
      f"(memory is in the worker processes)")
best = min(timing, key=lambda r: r["seconds"])
single = [r for r in timing if r["workers"] == 1 and r["chunk"] == 2048][0]
eight = [r for r in timing if r["workers"] == 8 and r["chunk"] == 2048][0]
print(f"fastest: {best['chunk']} px chunks with {best['workers']} threads, {best['seconds']:.1f} s; "
      f"2048 px: 1 thread {single['seconds']:.0f} s -> 8 threads {eight['seconds']:.0f} s "
      f"({single['seconds'] / eight['seconds']:.1f}x)")
(OUT / "metrics.json").write_text(json.dumps({"memory": [{k: v for k, v in r.items() if k != "trace"} for r in mem_rows],
                                              "timing": timing, "lock": {k: v for k, v in lock_run.items() if k != "trace"},
                                              "distributed": {k: v for k, v in dist_run.items() if k != "trace"}}, indent=1))

fig, axes = plt.subplots(1, 2, figsize=(13, 4.8))
cols = dict(zip(CHUNKS, ["#d53e4f", "#fdae61", "#66c2a5", "#3288bd", "#5e4fa2"]))
for ax, field, ylabel, title in [(axes[0], "seconds", "wall-clock seconds (lower is better)", "Time"),
                                 (axes[1], "peak_gb", "peak memory (GB)", "Peak memory")]:
    for c in CHUNKS:
        rows = sorted([r for r in timing if r["chunk"] == c], key=lambda r: r["workers"])
        ax.plot([r["workers"] for r in rows], [r[field] for r in rows], "o-", color=cols[c], lw=2, label=f"{c} px chunks")
    ax.plot([8], [lock_run[field]], "X", color="black", ms=9, label="2048 px, lock=True")
    ax.set_xscale("log", base=2)
    ax.set_xticks(WORKERS)
    ax.set_xticklabels(WORKERS)
    ax.set_xlabel("threads (num_workers)")
    ax.set_ylabel(ylabel)
    ax.set_title(f"{title}: median NDVI of 8 full scenes")
    ax.legend(fontsize=8.5)
fig.tight_layout()
fig.savefig(OUT / "05-timing-chunks-threads.png", bbox_inches="tight", facecolor="white")
plt.close(fig)

# ------------------------------------------------------------------ Step 7: windowed write
print("\n=== Step 7: write the result to disk chunk by chunk ===")
med = median_ndvi(*open_stack(DATES, 2048)).rio.write_nodata(np.nan, encoded=False)
dest = OUT / "ndvi_median_2025_30UXC.tif"
dest.unlink(missing_ok=True)
with RSSSampler() as s, dask.config.set(scheduler="threads", num_workers=8):
    t0 = time.perf_counter()
    med.rio.to_raster(dest, tiled=True, blockxsize=512, blockysize=512, compress="deflate", predictor=3,
                      windowed=True, lock=threading.Lock())
    secs = time.perf_counter() - t0
print(f"wrote {dest.name}: {dest.stat().st_size / 1e6:.0f} MB on disk in {secs:.1f} s; "
      f"memory rose by at most {s.peak_gb:.2f} GB for a {med.nbytes / 1e6:.0f} MB result")
with rasterio.open(dest) as src:
    print(f"check: {src.width} x {src.height} {src.dtypes[0]}, {src.crs}, {src.profile['compress']}, "
          f"{len(list(src.block_windows(1)))} blocks of {src.block_shapes[0]}")
    whole = src.read(1, out_shape=(1098, 1098))
    zoom = src.read(1, window=rasterio.windows.Window(4500, 5000, 1200, 900))   # 12 x 9 km

fig, axes = plt.subplots(1, 2, figsize=(13, 5.8), gridspec_kw={"width_ratios": [1, 1.1]})
im = axes[0].imshow(whole, cmap="Spectral", vmin=0, vmax=0.9, extent=EXT)
zx0, zy1 = EXT[0] + 45, EXT[3] - 50
axes[0].add_patch(Rectangle((zx0, zy1 - 9), 12, 9, fill=False, ec="black", lw=1.8))
axes[0].set_title("Median NDVI, March–November 2025 (8 dates)\nwhole tile, 110 × 110 km")
axes[0].set_xlabel("Easting (km, UTM 30N)")
axes[0].set_ylabel("Northing (km)")
axes[1].imshow(zoom, cmap="Spectral", vmin=0, vmax=0.9, extent=[zx0, zx0 + 12, zy1 - 9, zy1])
axes[1].set_title("The outlined 12 × 9 km box at full 10 m resolution")
axes[1].set_xlabel("Easting (km)")
cb = fig.colorbar(im, ax=axes, shrink=0.85, pad=0.02)
cb.set_label("median NDVI")
fig.savefig(OUT / "06-median-ndvi-result.png", bbox_inches="tight", facecolor="white")
plt.close(fig)

print(f"\nfinished in {(time.perf_counter() - T_START) / 60:.1f} min")