#!/usr/bin/env python3
"""Independent pixel-by-pixel check of a Landscape Program (ЛП) run.

Re-implements every node of the ЛП graph in numpy — 00 z_integrated,
01 u_integrated, 02 z_max, 03 u_max, 10 e_ecological_constraint,
20 c_conflicts, 21 p_max, 30 func_zoning (V1–V3 with the run's thresholds and
tie_priority), 31 c_max — from the 13 leaf input rasters the run actually used,
and compares the result with all 15 output rasters the run published.

The program is READ-ONLY: it issues only HTTP GET requests, so it can be
pointed at any stand, production included:

    python lp_verify.py --base-url http://localhost:18080 --run <run-id>
    python lp_verify.py --base-url https://dssgis.dwh.kz  --run <run-id>

Data flow:
  GET /landscape/runs/{id}          run params, per-node inputs and outputs
  GET /datasets/{id}.cog            every input (13) and output (15) raster
      (or /datasets/{id}/download for inputs with --inputs-from download)

The leaf inputs are taken from the run's node `inputs` (group = slot code), the
outputs from the node `outputs`. Expected values are computed through the whole
chain from the leaf inputs only (E from the re-computed Z/U, C and V from the
re-computed E, ...), so a wrong intermediate output cannot hide a later error.
The run's own chain references (node 10 reading node 00's output, ...) are
checked for consistency as well.

Large rasters (production: 22219 x 14977) are processed in windows of
--block-rows x --block-cols pixels (default 512 x 4096, aligned with the
512-pixel COG tiles), so memory stays bounded (~0.5 GB).

Pixel rules (source of truth: internal/landscape/kernels.go):
  valid(v)       1 <= v <= 5 (NoData and anything else is invalid); class
                 values are used as int16(v) (truncation)
  classify(x)    float32 x: < 1 -> NoData, <= 1.8 -> 1, <= 2.6 -> 2,
                 <= 3.4 -> 3, <= 4.2 -> 4, <= 5.0 -> 5, else NoData
  z_integrated   all Z valid: classify(0.25*Z1 + 0.25*Z2 + 0.25*Z3 + 0.25*Z4)
                 (float32, summed left to right)
  u_integrated   water & soil & (crop | veg): classify(num/den),
                 num = 0.25*water + 0.25*soil [+ 0.25*crop] [+ 0.25*veg],
                 den = 0.5 [+ 0.25] [+ 0.25] (float32)
  z_max          all Z valid: max;  u_max: U domain as above, max of valid
  e              Z_int & U_int valid: max(Z, U)
  c_conflicts    P_j & E valid: min(P_j, E)
  p_max, c_max   max over valid inputs; NoData when none is valid
  func_zoning    E valid & any P_k valid; candidate k: P_k valid & rule(E, P_k)
                 V1: E<=2 & S>=v1_min_s_e12 | E==3 & S>=v1_min_s_e3
                 V2: E<=2 | E==3 & S>=v2_min_s_e3 | E==4 & S>=v2_min_s_e4
                 V3: E<=4
                 no candidate -> 1; else sector with max S (ties by
                 tie_priority) + 1

Exit code: 0 when every output matches pixel for pixel, 1 on any mismatch,
2 on a usage / data error.
"""
import argparse
import json
import os
import shutil
import sys
import tempfile
import time
import urllib.error
import urllib.request

import numpy as np

try:
    import rasterio
    from rasterio.windows import Window
except ImportError:  # pragma: no cover
    sys.exit("lp_verify.py needs numpy and rasterio: pip install numpy rasterio")

NODATA = np.int16(-9999)
SLOTS = [
    "lp-z-biodiversity", "lp-z-lcr", "lp-z-soil-protection", "lp-z-wr",
    "lp-u-surf-water", "lp-u-soil", "lp-u-crop", "lp-u-veg",
    "lp-p-irrigation", "lp-p-snow", "lp-p-crop", "lp-p-pasture", "lp-p-tourism",
]
Z = SLOTS[0:4]
U = SLOTS[4:8]  # water, soil, crop, veg
P = SLOTS[8:13]  # sector order 1..5
C = ["lp-c-irrigation", "lp-c-snow", "lp-c-crop", "lp-c-pasture", "lp-c-tourism"]
OUTPUTS = [
    "lp-z-integrated", "lp-u-integrated", "lp-z-max", "lp-u-max", "lp-e",
    *C, "lp-p-max", "lp-v1-conservative", "lp-v2-balanced", "lp-v3-development", "lp-c-max",
]
OUTPUT_NODE = {
    "lp-z-integrated": "00", "lp-u-integrated": "01", "lp-z-max": "02", "lp-u-max": "03", "lp-e": "10",
    **{c: "20" for c in C}, "lp-p-max": "21",
    "lp-v1-conservative": "30", "lp-v2-balanced": "30", "lp-v3-development": "30", "lp-c-max": "31",
}
DEFAULT_PARAMS = {"tie_priority": [1, 2, 3, 4, 5], "v1_min_s_e12": 3, "v1_min_s_e3": 4,
                  "v2_min_s_e3": 3, "v2_min_s_e4": 5}

F32 = np.float32


# ----------------------------------------------------------------- pixel rules
def valid(v):
    return (v >= 1) & (v <= 5)


def cls(v, ok):
    """int16(v) for valid pixels (Go truncation), 0 elsewhere."""
    return np.where(ok, np.trunc(np.where(ok, v, 0)), 0).astype(np.int16)


def classify(x):
    x = x.astype(np.float32, copy=False)
    return np.select(
        [x < F32(1.0), x <= F32(1.8), x <= F32(2.6), x <= F32(3.4), x <= F32(4.2), x <= F32(5.0)],
        [NODATA, 1, 2, 3, 4, 5], default=NODATA).astype(np.int16)


def quarter(v):
    # float32(float32(0.25) * float32(v)), as in the Go kernels
    return F32(0.25) * v.astype(np.float32)


def z_integrated(z):
    ok = valid(z[0]) & valid(z[1]) & valid(z[2]) & valid(z[3])
    score = ((quarter(z[0]) + quarter(z[1])) + quarter(z[2])) + quarter(z[3])
    return np.where(ok, classify(score), NODATA).astype(np.int16)


def u_domain(u):
    water, soil, crop, veg = u
    crop_ok, veg_ok = valid(crop), valid(veg)
    return valid(water) & valid(soil) & (crop_ok | veg_ok), crop_ok, veg_ok


def u_integrated(u):
    water, soil, crop, veg = u
    ok, crop_ok, veg_ok = u_domain(u)
    num = quarter(water) + quarter(soil)
    den = np.full(num.shape, F32(0.5), dtype=np.float32)
    num = np.where(crop_ok, num + quarter(crop), num).astype(np.float32)
    den = np.where(crop_ok, den + F32(0.25), den).astype(np.float32)
    num = np.where(veg_ok, num + quarter(veg), num).astype(np.float32)
    den = np.where(veg_ok, den + F32(0.25), den).astype(np.float32)
    with np.errstate(invalid="ignore", divide="ignore"):
        score = (num / den).astype(np.float32)
    return np.where(ok, classify(score), NODATA).astype(np.int16)


def z_max(z):
    ok = valid(z[0]) & valid(z[1]) & valid(z[2]) & valid(z[3])
    m = np.maximum.reduce([cls(v, ok) for v in z])
    return np.where(ok, m, NODATA).astype(np.int16)


def u_max(u):
    water, soil, crop, veg = u
    ok, crop_ok, veg_ok = u_domain(u)
    m = np.maximum(cls(water, ok), cls(soil, ok))
    m = np.where(crop_ok, np.maximum(m, cls(crop, crop_ok)), m)
    m = np.where(veg_ok, np.maximum(m, cls(veg, veg_ok)), m)
    return np.where(ok, m, NODATA).astype(np.int16)


def e_constraint(zi, ui):
    ok = valid(zi) & valid(ui)
    return np.where(ok, np.maximum(zi, ui), NODATA).astype(np.int16)


def c_conflict(p, e):
    ok = valid(p) & valid(e)
    return np.where(ok, np.minimum(cls(p, ok), e), NODATA).astype(np.int16)


def max_any_valid(layers):
    oks = [valid(v) for v in layers]
    m = np.maximum.reduce([cls(v, ok) for v, ok in zip(layers, oks)])
    anyok = np.logical_or.reduce(oks)
    return np.where(anyok, m, NODATA).astype(np.int16)


def func_zoning(e, p, params):
    eok = valid(e)
    ev = np.where(eok, e, 0).astype(np.int16)
    base = [valid(v) for v in p]
    s = [cls(v, ok) for v, ok in zip(p, base)]
    any_base = np.logical_or.reduce(base)
    t = params
    rules = [
        lambda e_, s_: ((e_ <= 2) & (s_ >= t["v1_min_s_e12"])) | ((e_ == 3) & (s_ >= t["v1_min_s_e3"])),
        lambda e_, s_: (e_ <= 2) | ((e_ == 3) & (s_ >= t["v2_min_s_e3"])) | ((e_ == 4) & (s_ >= t["v2_min_s_e4"])),
        lambda e_, s_: e_ <= 4,
    ]
    outs = []
    for rule in rules:
        cand = [b & rule(ev, sk) for b, sk in zip(base, s)]
        any_cand = np.logical_or.reduce(cand)
        max_s = np.maximum.reduce([np.where(ck, sk, 0) for ck, sk in zip(cand, s)]).astype(np.int16)
        sector = np.zeros(ev.shape, dtype=np.int16)
        done = np.zeros(ev.shape, dtype=bool)
        for code in t["tie_priority"]:
            k = code - 1
            hit = cand[k] & (s[k] == max_s) & ~done
            sector[hit] = code
            done |= hit
        v = np.where(any_cand, sector + 1, 1).astype(np.int16)
        outs.append(np.where(eok & any_base, v, NODATA).astype(np.int16))
    return outs


def expected_block(inp, params):
    """inp: dict slot -> float64 array. Returns dict output group -> int16."""
    z = [inp[g] for g in Z]
    u = [inp[g] for g in U]
    p = [inp[g] for g in P]
    out = {
        "lp-z-integrated": z_integrated(z),
        "lp-u-integrated": u_integrated(u),
        "lp-z-max": z_max(z),
        "lp-u-max": u_max(u),
    }
    e = e_constraint(out["lp-z-integrated"], out["lp-u-integrated"])
    out["lp-e"] = e
    for cg, pk in zip(C, p):
        out[cg] = c_conflict(pk, e)
    out["lp-p-max"] = max_any_valid(p)
    v1, v2, v3 = func_zoning(e, p, params)
    out["lp-v1-conservative"], out["lp-v2-balanced"], out["lp-v3-development"] = v1, v2, v3
    out["lp-c-max"] = max_any_valid([out[c] for c in C])
    return out


# ------------------------------------------------------------------- HTTP (GET)
def http_get(url, timeout=120, retries=3):
    req = urllib.request.Request(url, method="GET", headers={"User-Agent": "lp_verify/1.0"})
    last = None
    for attempt in range(retries):
        try:
            return urllib.request.urlopen(req, timeout=timeout)
        except urllib.error.HTTPError as e:
            if e.code < 500:
                raise
            last = e
        except (urllib.error.URLError, TimeoutError, ConnectionError) as e:
            last = e
        time.sleep(2 * (attempt + 1))
    raise last


def get_json(url):
    with http_get(url) as r:
        return json.loads(r.read().decode("utf-8"))


def download(url, path, wait=0):
    """GET url into path (reused when already cached). A 409 means the COG is
    still being produced (conversion is asynchronous): retry for up to `wait`
    seconds."""
    if os.path.exists(path) and os.path.getsize(path) > 0:
        return "cached"
    tmp = path + ".part"
    deadline = time.time() + wait
    while True:
        try:
            r = http_get(url, timeout=600)
            break
        except urllib.error.HTTPError as e:
            if e.code != 409 or time.time() >= deadline:
                raise
            time.sleep(2)
    with r, open(tmp, "wb") as f:
        shutil.copyfileobj(r, f, length=4 << 20)
    os.replace(tmp, path)
    return "%.1f MB" % (os.path.getsize(path) / 1e6)


# ----------------------------------------------------------------------- main
def parse_points(s):
    pts = []
    for part in filter(None, (x.strip() for x in s.split(";"))):
        c, r = part.split(",")
        pts.append((int(c), int(r)))
    return pts


def main():
    ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    ap.add_argument("--base-url", default=os.environ.get("BASE_URL", "http://localhost:18080"))
    ap.add_argument("--run", required=True, help="landscape run id")
    ap.add_argument("--cache", help="directory for downloaded rasters (kept; existing files are reused)")
    ap.add_argument("--inputs-from", choices=["cog", "download"], default="cog",
                    help="read inputs from their COG (default) or the original upload (what the node read)")
    ap.add_argument("--block-rows", type=int, default=512)
    ap.add_argument("--block-cols", type=int, default=4096)
    ap.add_argument("--wait", type=int, default=180,
                    help="seconds to keep retrying a COG that is not ready yet (HTTP 409), default 180")
    ap.add_argument("--points", help='spot-check pixels "col,row;col,row;..." and print their values')
    args = ap.parse_args()
    base = args.base_url.rstrip("/")

    run = get_json("%s/landscape/runs/%s" % (base, args.run))
    params = dict(DEFAULT_PARAMS)
    raw_params = run.get("params") or {}
    for k in DEFAULT_PARAMS:
        if raw_params.get(k):
            params[k] = raw_params[k]
    print("run        %s  status=%s  from_stage=%s" % (run["id"], run["status"], run.get("from_stage")))
    print("params     %s" % json.dumps(params))
    nodes = sorted(run.get("nodes") or [], key=lambda n: n["node_code"])
    print("nodes      %s" % " ".join("%s:%s" % (n["node_code"], n["status"]) for n in nodes))
    if run["status"] != "succeeded":
        print("ERROR: run is not succeeded; nothing to verify")
        return 2

    errors = []
    slot_ids, out_ids = {}, {}
    for n in nodes:
        for ref in n.get("inputs") or []:
            g, did = ref["group"], ref["dataset_id"]
            if g in SLOTS:
                if slot_ids.setdefault(g, did) != did:
                    errors.append("slot %s resolved to two datasets (%s, %s)" % (g, slot_ids[g], did))
        for ref in n.get("outputs") or []:
            out_ids[ref["group"]] = ref["dataset_id"]
            if OUTPUT_NODE.get(ref["group"]) != n["node_code"]:
                errors.append("output %s reported by node %s" % (ref["group"], n["node_code"]))
    # chain consistency: an intermediate input must be what this run produced
    for n in nodes:
        for ref in n.get("inputs") or []:
            g = ref["group"]
            if g in OUTPUTS and out_ids.get(g) != ref["dataset_id"]:
                errors.append("node %s read %s=%s but the run's output is %s" % (n["node_code"], g, ref["dataset_id"], out_ids.get(g)))
    missing = [g for g in SLOTS if g not in slot_ids] + [g for g in OUTPUTS if g not in out_ids]
    if missing:
        errors.append("run does not reference: %s" % ", ".join(missing))
    if errors:
        for e in errors:
            print("ERROR:", e)
        return 2

    cache = args.cache or tempfile.mkdtemp(prefix="lp_verify_")
    os.makedirs(cache, exist_ok=True)
    files = {}
    t0 = time.time()
    for g in SLOTS + OUTPUTS:
        did = slot_ids.get(g) or out_ids[g]
        if g in SLOTS and args.inputs_from == "download":
            url, path = "%s/datasets/%s/download" % (base, did), os.path.join(cache, did + ".orig.tif")
        else:
            url, path = "%s/datasets/%s.cog" % (base, did), os.path.join(cache, did + ".cog.tif")
        try:
            info = download(url, path, args.wait)
        except urllib.error.HTTPError as e:
            print("ERROR: GET %s -> HTTP %d (output not ready or missing?)" % (url, e.code))
            return 2
        files[g] = path
        print("fetch      %-22s %s  %s" % (g, did, info))
    print("downloaded in %.1f s to %s" % (time.time() - t0, cache))

    with rasterio.Env(GDAL_CACHEMAX=512):
        ds = {g: rasterio.open(files[g]) for g in SLOTS + OUTPUTS}
        ref = ds[SLOTS[0]]
        for g, d in ds.items():
            if (d.width, d.height) != (ref.width, ref.height) or not d.transform.almost_equals(ref.transform) or d.crs != ref.crs:
                errors.append("%s grid %dx%d %s differs from %s %dx%d %s" % (
                    g, d.width, d.height, d.crs, SLOTS[0], ref.width, ref.height, ref.crs))
        for g in OUTPUTS:
            if ds[g].dtypes[0] != "int16" or ds[g].nodata != -9999:
                errors.append("output %s: dtype %s nodata %s (want int16, -9999)" % (g, ds[g].dtypes[0], ds[g].nodata))
        if errors:
            for e in errors:
                print("ERROR:", e)
            return 2
        W, H = ref.width, ref.height
        print("grid       %dx%d  %s  pixel %.3f" % (W, H, ref.crs, ref.transform.a))

        if args.points:
            return spot_check(ds, parse_points(args.points), params)

        mism = {g: 0 for g in OUTPUTS}
        validn = {g: 0 for g in OUTPUTS}
        first = {}
        t0 = time.time()
        for r0 in range(0, H, args.block_rows):
            h = min(args.block_rows, H - r0)
            for c0 in range(0, W, args.block_cols):
                w = min(args.block_cols, W - c0)
                win = Window(c0, r0, w, h)
                inp = {g: ds[g].read(1, window=win, masked=False).astype(np.float64) for g in SLOTS}
                exp = expected_block(inp, params)
                for g in OUTPUTS:
                    act = ds[g].read(1, window=win, masked=False)
                    bad = exp[g] != act
                    nbad = int(bad.sum())
                    validn[g] += int((exp[g] != NODATA).sum())
                    if nbad:
                        mism[g] += nbad
                        if g not in first:
                            rr, cc = np.argwhere(bad)[0]
                            first[g] = (c0 + int(cc), r0 + int(rr), int(exp[g][rr, cc]), int(act[rr, cc]))
        print("compared in %.1f s" % (time.time() - t0))
        for d in ds.values():
            d.close()

    print()
    print("%-20s %-4s %-38s %12s %12s %10s" % ("output", "node", "dataset_id", "pixels", "valid(exp)", "mismatch"))
    total = 0
    for g in OUTPUTS:
        total += mism[g]
        line = "%-20s %-4s %-38s %12d %12d %10d" % (g, OUTPUT_NODE[g], out_ids[g], W * H, validn[g], mism[g])
        if g in first:
            line += "   first at col=%d row=%d expected=%d actual=%d" % first[g]
        print(line)
    print()
    print("TOTAL mismatching pixels: %d -> %s" % (total, "PASS" if total == 0 else "FAIL"))
    return 0 if total == 0 else 1


def spot_check(ds, points, params):
    bad = 0
    shown = ["lp-z-integrated", "lp-u-integrated", "lp-e", *C, "lp-v1-conservative", "lp-v2-balanced", "lp-v3-development"]
    for col, row in points:
        win = Window(col, row, 1, 1)
        inp = {g: ds[g].read(1, window=win, masked=False).astype(np.float64) for g in SLOTS}
        exp = expected_block(inp, params)
        vals = " ".join("%s=%d" % (g.replace("lp-", ""), int(inp[g][0, 0])) for g in SLOTS)
        print("point (%d,%d) inputs: %s" % (col, row, vals))
        for g in shown:
            act = int(ds[g].read(1, window=win, masked=False)[0, 0])
            e = int(exp[g][0, 0])
            ok = act == e
            bad += 0 if ok else 1
            print("   %-20s expected=%6d actual=%6d %s" % (g, e, act, "OK" if ok else "MISMATCH"))
    print("SPOT CHECK: %d points, %d mismatches -> %s" % (len(points), bad, "PASS" if bad == 0 else "FAIL"))
    return 0 if bad == 0 else 1


if __name__ == "__main__":
    sys.exit(main())