"""The lazy dask pipeline, the eager NumPy equivalent and a benchmark runner.

dask_raster.py imports the functions; it also runs this file as a fresh subprocess for each
benchmark, so that memory left over from one run can never inflate the next:

    python raster_pipeline.py dask 8 2048 8 0      # mode, dates, chunk px, threads, lock
    python raster_pipeline.py eager 4 0 0 0
    python raster_pipeline.py distributed 8 2048 2 0   # LocalCluster: 4 processes x 2 threads
"""
import json
import os
import sys
import threading
import time
from pathlib import Path

import numpy as np

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

HERE = Path(__file__).resolve().parent
DATA = HERE / "data"
TILE = "30UXC"
DATES = ["2025-03-06", "2025-04-02", "2025-04-05", "2025-04-30",
         "2025-05-02", "2025-06-19", "2025-07-11", "2025-11-21"]
OFFSET = 1000          # processing baseline 04.00+ adds 1000 to every DN


def band_path(date, band):
    return DATA / f"{TILE}_{date}_{band}.tif"


def open_stack(dates, chunk, lock=False):
    """Open B04 and B08 for every date as lazy, chunked DataArrays: nothing is read yet."""
    import rioxarray
    import xarray as xr

    chunks = {"band": 1, "y": chunk, "x": chunk}
    red, nir = [], []
    for d in dates:
        red.append(rioxarray.open_rasterio(band_path(d, "B04"), chunks=chunks, lock=lock).squeeze("band", drop=True))
        nir.append(rioxarray.open_rasterio(band_path(d, "B08"), chunks=chunks, lock=lock).squeeze("band", drop=True))
    time_index = xr.Variable("time", np.array(dates, dtype="datetime64[ns]"))
    return xr.concat(red, dim=time_index), xr.concat(nir, dim=time_index)


def median_ndvi(red, nir):
    """NDVI for every date, then the median through time. Still lazy."""
    red = red.where(red > 0).astype("float32") - OFFSET          # DN 0 is no data
    nir = nir.where(nir > 0).astype("float32") - OFFSET
    ndvi = (nir - red) / (nir + red)
    ndvi = ndvi.where(np.isfinite(ndvi))
    ndvi = ndvi.chunk({"time": -1})                              # a median needs every date in one chunk
    out = ndvi.median("time", skipna=True).astype("float32")
    out.attrs = {}                                               # drop the uint16 _FillValue/scale attributes
    return out


def eager_median_ndvi(dates):
    """The same calculation in plain NumPy: every pixel of every date in memory at once."""
    import rasterio

    def read(d, b):
        with rasterio.open(band_path(d, b)) as src:
            return src.read(1)

    red = np.stack([read(d, "B04") for d in dates]).astype("float32")
    nir = np.stack([read(d, "B08") for d in dates]).astype("float32")
    red[red == 0] = np.nan
    nir[nir == 0] = np.nan
    red -= OFFSET
    nir -= OFFSET
    ndvi = (nir - red) / (nir + red)
    del red, nir
    ndvi[~np.isfinite(ndvi)] = np.nan
    return np.nanmedian(ndvi, axis=0).astype("float32")


class RSSSampler:
    """Sample this process's resident memory every 50 ms in a background thread."""

    def __init__(self, interval=0.05):
        import psutil

        self.proc = psutil.Process()
        self.interval = interval
        self.trace = []
        self.cpu = []                                     # % of one core, so 800% = 8 busy cores
        self._stop = threading.Event()

    def _run(self):
        while not self._stop.is_set():
            self.trace.append((time.perf_counter() - self.t0, self.proc.memory_info().rss))
            self.cpu.append(self.proc.cpu_percent())
            time.sleep(self.interval)

    def __enter__(self):
        self.t0 = time.perf_counter()
        self.base = self.proc.memory_info().rss
        self.proc.cpu_percent()
        self.th = threading.Thread(target=self._run, daemon=True)
        self.th.start()
        return self

    def __exit__(self, *exc):
        self._stop.set()
        self.th.join()

    @property
    def peak_gb(self):
        return (max(r for _, r in self.trace) - self.base) / 1e9


def bench(mode, n_dates, chunk, workers, lock):
    import dask

    dates = DATES[:n_dates]
    with RSSSampler() as s:
        t0 = time.perf_counter()
        if mode == "eager":
            result = eager_median_ndvi(dates)
        elif mode == "dask":
            red, nir = open_stack(dates, chunk, lock=lock)
            with dask.config.set(scheduler="threads", num_workers=workers):
                result = median_ndvi(red, nir).compute().values
        else:
            from dask.distributed import Client, LocalCluster

            with LocalCluster(n_workers=4, threads_per_worker=workers, memory_limit="6GB",
                              dashboard_address=":0") as cluster, Client(cluster):
                t0 = time.perf_counter()                 # exclude cluster start-up
                red, nir = open_stack(dates, chunk, lock=lock)
                result = median_ndvi(red, nir).compute().values
        elapsed = time.perf_counter() - t0
    trace = [(round(t, 2), round((r - s.base) / 1e9, 3)) for t, r in s.trace[::4]]
    return {"mode": mode, "n_dates": n_dates, "chunk": chunk, "workers": workers, "lock": lock,
            "seconds": round(elapsed, 2), "peak_gb": round(s.peak_gb, 3),
            "result_mean": round(float(np.nanmean(result)), 4), "trace": trace}


if __name__ == "__main__":
    m, n, c, w, lk = sys.argv[1:6]
    print(json.dumps(bench(m, int(n), int(c), int(w), bool(int(lk)))))