#!/usr/bin/env python3
"""field_verify.py: the second implementation of integer neural fields
(extensions/neural-fields.md), written from the extension's text and sharing no code with
wai-rs.

  python3 -I field_verify.py <corpus-dir>

For the corpus volumetric-conformance/fields/ it parses every scene's WIG1 payload and its
WIL1 levels with §8's checks in their order, and for every prefix of each scene's levels
(§7.3: each level applied only as the field's next level) computes the field's values at every
point vectors.json lists (§5: the levels' direct index or spatial hash, the trilinear weights in
fixed point, the network with its rounding and clamps) and the SHA-256 of their bytes (§9);
renders the field at its own view and at every listed view (§6: the camera, a ray a pixel, the
pinned samples, the transmittance, the stop and the background) and compares each render's
SHA-256; requires every fault payload to be refused with the code vectors.json states, every
level applied out of order or of the wrong shape to be refused, and each staged sequence's
final stage hash to be the render of the base and the levels its final stage holds. Then it
plants each fault of §11 in its own decode (an interpolation weight off by one, a hash
multiplier changed, a layer's rounding flipped to floor, one ray sample too many, a level
applied whatever its index) and requires the corpus to catch each one. Standard library only.
Exit 0 when everything holds.
"""
import hashlib
import json
import math
import os
import struct
import sys

ONE = 1 << 16
PRIMES = (1, 2654435761, 805459861)
FAULTS = ("weight_off_by_one", "hash_constant", "activation_rounding", "ray_step", "level_order")


class Refused(Exception):
    def __init__(self, code, why):
        super().__init__(f"{code}: {why}")
        self.code = code


# ------------------------------------------------------------------ integer-operators §2


def rhu(p, s):
    """R_half_up(p, s): floor((p + 2^(s-1)) / 2^s), and p when s = 0."""
    return p if s == 0 else (p + (1 << (s - 1))) >> s


def rdiv(a, b):
    """floor((2a + b) / 2b) for b > 0 (splat-exact §5.1)."""
    return (2 * a + b) // (2 * b)


def exp_neg_q16(x):
    u = -x
    n = u // 45426
    r = x + 45426 * n
    t = r + 88670
    qq = (t * t) >> 16
    y = ((23495 * qq) >> 16) + 22544
    return 0 if n >= 63 else y >> n


# ------------------------------------------------------------------ §3, §4: parsing


def u16(b, at):
    return b[at] | (b[at + 1] << 8)


def i32(b, at):
    return struct.unpack_from("<i", b, at)[0]


def read_view(b):
    v = {"eye": [i32(b, 0), i32(b, 4), i32(b, 8)], "target": [i32(b, 12), i32(b, 16), i32(b, 20)],
         "up": [i32(b, 24), i32(b, 28), i32(b, 32)], "focal": u16(b, 36), "width": u16(b, 38),
         "height": u16(b, 40), "steps": u16(b, 42), "t_near": i32(b, 44), "t_far": i32(b, 48),
         "background": [b[52], b[53], b[54]]}
    return v


def cross(a, b):
    return [a[1] * b[2] - a[2] * b[1], a[2] * b[0] - a[0] * b[2], a[0] * b[1] - a[1] * b[0]]


def dot(a, b):
    return a[0] * b[0] + a[1] * b[1] + a[2] * b[2]


def camera(v):
    """splat-exact §5.3: the right, up and forward rows at 2^14, or None if degenerate."""
    fw = [v["target"][i] - v["eye"][i] for i in range(3)]
    nf = math.isqrt(dot(fw, fw))
    if nf == 0:
        return None
    f = [rdiv(c << 14, nf) for c in fw]
    rr = cross(fw, v["up"])
    nr = math.isqrt(dot(rr, rr))
    if nr == 0:
        return None
    r = [rdiv(c << 14, nr) for c in rr]
    u = [rhu(c, 14) for c in cross(r, f)]
    return r, u, f


def check_view(v):
    for k in ("eye", "target", "up"):
        if any(abs(c) > 1 << 24 for c in v[k]):
            raise Refused("field_range", "a coordinate past 2^24")
    if not 1 <= v["focal"] <= 4096:
        raise Refused("field_range", "focal")
    if not 1 <= v["width"] <= 1024:
        raise Refused("field_range", "width")
    if not 1 <= v["height"] <= 1024:
        raise Refused("field_range", "height")
    if not 1 <= v["steps"] <= 1024:
        raise Refused("field_range", "steps")
    if not 0 <= v["t_near"] <= 1 << 22:
        raise Refused("field_range", "t_near")
    if not v["t_near"] < v["t_far"] <= 1 << 22:
        raise Refused("field_range", "t_far")
    if camera(v) is None:
        raise Refused("field_range", "a degenerate pose")


def parse(p):
    """§3 in §8's order: the field, or Refused."""
    if p[:4] == b"WIL1":
        raise Refused("format_capability_mismatch", "a level")
    if len(p) < 4 or p[:4] != b"WIG1":
        raise Refused("bad_magic", "not WIG1")
    if len(p) < 5:
        raise Refused("stream_length", "no version")
    if p[4] != 1:
        raise Refused("version", "not 1")
    if len(p) < 16:
        raise Refused("stream_length", "fixed header")
    L, K, F, T, ffrac, wfrac, N = p[5], p[6], p[7], p[8], p[9], p[10], p[11]
    vat = 16 + 2 * L + 2 * (N + 1)
    end = vat + 56
    if len(p) < end:
        raise Refused("stream_length", "header")
    if p[12:16] != b"\0\0\0\0" or p[end - 1] != 0:
        raise Refused("reserved_nonzero", "reserved")
    sides = [u16(p, 16 + 2 * i) for i in range(L)]
    widths = [u16(p, 16 + 2 * L + 2 * j) for j in range(N + 1)]
    view = read_view(p[vat:end])
    rules = [
        1 <= L <= 16, 1 <= K <= L, 1 <= F <= 8, 4 <= T <= 19, 8 <= ffrac <= 16, wfrac <= 15, 1 <= N <= 8,
    ]
    for ok in rules:
        if not ok:
            raise Refused("field_range", "a header field")
    for i, s in enumerate(sides):
        if s < 1 or (i and s <= sides[i - 1]):
            raise Refused("field_range", "sides")
    if any(not 1 <= w <= 256 for w in widths):
        raise Refused("field_range", "a width")
    if widths[0] != L * F:
        raise Refused("field_range", "input width")
    if widths[-1] != 4:
        raise Refused("field_range", "output width")
    check_view(view)
    if sum(widths[j] * widths[j + 1] for j in range(N)) > 1 << 16:
        raise Refused("too_large", "network")
    at = end
    layers = []

    def take(n):
        nonlocal at
        if len(p) - at < n:
            raise Refused("stream_length", "a section")
        s = p[at:at + n]
        at += n
        return s

    for j in range(N):
        ni, no = widths[j], widths[j + 1]
        w = list(struct.unpack(f"<{ni * no}h", take(2 * ni * no)))
        b = list(struct.unpack(f"<{no}i", take(4 * no)))
        layers.append(([w[o * ni:(o + 1) * ni] for o in range(no)], b))
    grids = []
    for s in sides[:K]:
        e = min((s + 1) ** 3, 1 << T)
        flat = struct.unpack(f"<{e * F}h", take(2 * e * F))
        grids.append([flat[i * F:(i + 1) * F] for i in range(e)])
    if at != len(p):
        raise Refused("trailing", "after the last level")
    return {"L": L, "F": F, "T": T, "ffrac": ffrac, "wfrac": wfrac, "sides": sides, "layers": layers,
            "view": view, "grids": grids}


def parse_level(p):
    """§7.2 in §8's order: the level, or Refused."""
    if p[:4] == b"WIG1":
        raise Refused("format_capability_mismatch", "a field")
    if len(p) < 4 or p[:4] != b"WIL1":
        raise Refused("bad_magic", "not WIL1")
    if len(p) < 5:
        raise Refused("stream_length", "no version")
    if p[4] != 1:
        raise Refused("version", "not 1")
    if len(p) < 12:
        raise Refused("stream_length", "header")
    if p[10] or p[11]:
        raise Refused("reserved_nonzero", "reserved")
    index, F, T, R = p[5], p[6], p[7], u16(p, 8)
    if not 1 <= index <= 15 or not 1 <= F <= 8 or not 4 <= T <= 19 or R < 1:
        raise Refused("field_range", "a header field")
    n = min((R + 1) ** 3, 1 << T) * F
    if len(p) - 12 < 2 * n:
        raise Refused("stream_length", "table")
    if len(p) - 12 > 2 * n:
        raise Refused("trailing", "after the table")
    flat = struct.unpack(f"<{n}h", p[12:])
    return {"index": index, "F": F, "T": T, "R": R, "table": [flat[i * F:(i + 1) * F] for i in range(n // F)]}


def apply_level(fd, lv, fault=None):
    """§7.3: add the field's next level, or return geometry_mismatch and change nothing."""
    at = lv["index"] if fault == "level_order" else len(fd["grids"])
    if lv["index"] != at or at >= fd["L"] or len(fd["grids"]) >= fd["L"]:
        return "geometry_mismatch"
    if lv["F"] != fd["F"] or lv["T"] != fd["T"] or lv["R"] != fd["sides"][at]:
        return "geometry_mismatch"
    fd["grids"].append(lv["table"])
    return "ok"


def load(d, rel):
    with open(os.path.join(d, rel), "rb") as fh:
        return fh.read()


# ------------------------------------------------------------------ §5: values at a point


def lookup(side, T, x, y, z, fault):
    n = side + 1
    if n * n * n <= 1 << T:
        return x + n * y + n * n * z
    p3 = PRIMES[2] + 2 if fault == "hash_constant" else PRIMES[2]
    h = ((x * PRIMES[0]) % 2 ** 32) ^ ((y * PRIMES[1]) % 2 ** 32) ^ ((z * p3) % 2 ** 32)
    return h % (1 << T)


def values(fd, pt, fault=None):
    if any(c < 0 or c > ONE for c in pt):
        return [0, 0, 0, 0]
    F = fd["F"]
    x = [0] * (fd["L"] * F)
    for lvl, grid in enumerate(fd["grids"]):
        side = fd["sides"][lvl]
        cell, frac = [], []
        for c in pt:
            s = c * side
            k = min(s // ONE, side - 1)
            cell.append(k)
            frac.append(s - k * ONE + (1 if fault == "weight_off_by_one" else 0))
        sums = [0] * F
        for dz in (0, 1):
            for dy in (0, 1):
                for dx in (0, 1):
                    wgt = (frac[0] if dx else ONE - frac[0]) * (frac[1] if dy else ONE - frac[1]) * (frac[2] if dz else ONE - frac[2])
                    e = grid[lookup(side, fd["T"], cell[0] + dx, cell[1] + dy, cell[2] + dz, fault)]
                    for c in range(F):
                        sums[c] += wgt * e[c]
        for c in range(F):
            x[lvl * F + c] = rhu(sums[c], 48)
    a = x
    last = len(fd["layers"]) - 1
    for j, (rows, biases) in enumerate(fd["layers"]):
        out = []
        for row, b in zip(rows, biases):
            acc = b + sum(w * v for w, v in zip(row, a))
            y = acc >> fd["wfrac"] if fault == "activation_rounding" else rhu(acc, fd["wfrac"])
            y = max(-(1 << 31), min((1 << 31) - 1, y))
            if j < last:
                y = max(0, y)
            out.append(y)
        a = out
    sh, top = 16 - fd["ffrac"], 1 << fd["ffrac"]
    return [min(max(a[0], 0) << sh, 1 << 24)] + [max(0, min(top, a[c])) << sh for c in (1, 2, 3)]


# ------------------------------------------------------------------ §6: the render


def render(fd, v, fault=None):
    r, u, f = camera(v)
    W, H, fl = v["width"], v["height"], v["focal"]
    S = v["steps"] + (1 if fault == "ray_step" else 0)
    near, span = v["t_near"], v["t_far"] - v["t_near"]
    edges = [near + (k * span) // S for k in range(S + 1)]
    out = bytearray()
    for j in range(H):
        for i in range(W):
            dx, dy = 256 * i + 128 - 128 * W, 128 * H - 256 * j - 128
            d = [256 * fl * f[c] + dx * r[c] + dy * u[c] for c in range(3)]
            nd = math.isqrt(dot(d, d))
            tr, acc = ONE, [0, 0, 0]
            if nd:
                dirn = [rdiv(c << 16, nd) for c in d]
                for k in range(S):
                    if tr < 7:
                        break
                    delta = edges[k + 1] - edges[k]
                    m = edges[k] + delta // 2
                    pt = [v["eye"][c] + rhu(dirn[c] * m, 16) for c in range(3)]
                    q = values(fd, pt, fault)
                    tau = rhu(q[0] * delta, 16)
                    alpha = ONE - min(ONE, exp_neg_q16(-tau))
                    c = rhu(alpha * tr, 16)
                    tr -= c
                    for ch in range(3):
                        acc[ch] += rhu(q[ch + 1] * c, 16)
            for ch in range(3):
                out.append(min(255, rhu(255 * acc[ch] + tr * v["background"][ch], 16)))
    return bytes(out)


# ------------------------------------------------------------------ §9 and the corpus


def values_digest(vals):
    return hashlib.sha256(b"".join(struct.pack("<4q", *v) for v in vals)).hexdigest()


def check_corpus(d, vec, fault=None):
    """Every statement of vectors.json this decode does not reproduce (with `fault` planted,
    it stops at the first)."""
    fails = []
    for s in vec["scenes"]:
        fd = parse(load(d, s["base"]))
        for k, st in enumerate(s["stages"]):
            if k and apply_level(fd, parse_level(load(d, s["levels"][k - 1])), fault) != "ok":
                fails.append(f"{s['name']} stage {k}: its level does not apply")
                return fails
            if len(fd["grids"]) != st["held"]:
                fails.append(f"{s['name']} stage {k}: holds {len(fd['grids'])} levels, not {st['held']}")
            vals = [values(fd, p, fault) for p in s["points"]]
            if vals != st["values"] or values_digest(vals) != st["values_sha256"]:
                fails.append(f"{s['name']} stage {k}: values")
                if fault:
                    return fails
            for rv in st["renders"]:
                view = fd["view"] if rv["view"] == "own" else s["views"][rv["view"]]
                if hashlib.sha256(render(fd, view, fault)).hexdigest() != rv["rgb_sha256"]:
                    fails.append(f"{s['name']} stage {k}: render at view {rv['view']}")
                    if fault:
                        return fails
    for g in vec["geometry"]:
        fd = parse(load(d, g["base"]))
        got = [apply_level(fd, parse_level(load(d, rel)), fault) for rel in g["apply"]]
        if got != g["expect"]:
            fails.append(f"geometry {g['name']}: {got}")
            if fault:
                return fails
    orb = next(s for s in vec["scenes"] if s["name"] == "orb")
    for q in vec["staged"]:
        fd = parse(load(d, orb["base"]))
        for rel in orb["levels"][:q["final"]["stage"]]:
            apply_level(fd, parse_level(load(d, rel)))
        if hashlib.sha256(render(fd, fd["view"], fault)).hexdigest() != q["final"]["stage_sha256"]:
            fails.append(f"staged {q['name']}: the final stage hash")
            if fault:
                return fails
    return fails


def main(argv):
    if len(argv) != 2:
        print(__doc__)
        return 2
    d = argv[1]
    with open(os.path.join(d, "vectors.json"), encoding="utf-8") as fh:
        vec = json.load(fh)
    fails = check_corpus(d, vec)
    n_values = sum(len(s["points"]) * len(s["stages"]) for s in vec["scenes"])
    n_renders = sum(len(st["renders"]) for s in vec["scenes"] for st in s["stages"])
    for rf in vec["refusals"]:
        b = load(d, rf["file"])
        try:
            (parse_level if rf["format"] == "WIL1" else parse)(b)
            got = None
        except Refused as e:
            got = e.code
        if got != rf["code"]:
            fails.append(f"{rf['file']}: refused with {got}, not {rf['code']}")
    caught = 0
    for fault in FAULTS:
        if check_corpus(d, vec, fault):
            caught += 1
            print(f"planted {fault}: caught")
        else:
            fails.append(f"planted {fault}: not caught")
    for f in fails:
        print("FAIL " + f)
    if fails:
        print(f"{len(fails)} failed")
        return 1
    n_apply = sum(len(g["apply"]) for g in vec["geometry"])
    print(f"field_verify: {n_values} values, {n_renders} renders, {len(vec['refusals'])} refusals, "
          f"{n_apply} level applications and {len(vec['staged'])} stage hashes reproduced; "
          f"{caught} of {len(FAULTS)} planted faults caught")
    return 0


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