#!/usr/bin/env python3
"""wai_int_field_export.py: write integer neural fields (extensions/neural-fields.md).

The encoder side of wai.field.int_grid and wai.field.int_refine. A field is a multi-resolution
feature grid and a small network; in its JSON form (neural-fields.md §10.1) every feature,
weight and position is a float. This tool quantizes such a model as §10.2 states and writes the
field's WIG1 payload, carrying its first `carried` levels, and one WIL1 payload for each level
after them (§7.2). What a sink decodes from those payloads is fixed by the payloads alone; how
the float model was made is the encoder's choice.

  wai_int_field_export.py export MODEL.json OUTDIR [--name NAME]
      write OUTDIR/NAME.wig1 and OUTDIR/NAME.l<k>.wil1 for every level k past `carried`
  wai_int_field_export.py demo OUTDIR [--check]
      build the scenes of volumetric-conformance/fields/scenes/ (orb, mlp, edge) and write
      them, or with --check require OUTDIR to hold exactly those bytes
  wai_int_field_export.py model NAME
      print a demo scene's float model in the JSON form

The scenes are made here, with no download. `orb` is fitted: a procedural scene (a ball and
a box of absorbing medium, coloured by position) sampled at seeded points, each level's table
fitted coarse to fine by stochastic gradient descent to what the levels below it leave,
through the level's own interpolation and hash, and quantized before the next level is
fitted. `mlp` is a seeded random field, the network's general case. `edge` holds the
saturating cases. Every float operation here is +, -, *, / or a comparison, whose results IEEE
754 fixes, so the scenes are the same bytes on every machine. Standard library only.
"""
import json
import math
import os
import struct
import sys

FORM = "wai-int-field/1"
HASH_PRIMES = (1, 2654435761, 805459861)
MASK32 = 0xFFFFFFFF


# ---------------------------------------------------------------- quantization (§10.2)


def q(v, frac, lo, hi):
    """v * 2^frac, rounded half up, clamped to [lo, hi]."""
    return max(lo, min(hi, math.floor(v * (1 << frac) + 0.5)))


def q16(v):
    return q(v, 16, -(1 << 24), 1 << 24)


def entries(side, log2_table):
    return min((side + 1) ** 3, 1 << log2_table)


def table_bytes(lv, ffrac):
    flat = [x for e in lv["table"] for x in e]
    return struct.pack(f"<{len(flat)}h", *(q(x, ffrac, -32768, 32767) for x in flat))


def export(model):
    """The WIG1 payload of a model in the JSON form, and the WIL1 payload of each level it
    does not carry, as (index, bytes)."""
    if model.get("form") != FORM:
        raise ValueError(f"not a {FORM} model")
    ffrac, wfrac, t = model["ffrac"], model["wfrac"], model["log2_table"]
    levels, net, view = model["levels"], model["network"], model["view"]
    carried = model.get("carried", len(levels))
    nfeat = len(levels[0]["table"][0])
    for lv in levels:
        if len(lv["table"]) != entries(lv["side"], t) or any(len(e) != nfeat for e in lv["table"]):
            raise ValueError("a level does not hold min((side + 1)^3, 2^T) entries of F features")
    out = bytearray(b"WIG1")
    out += bytes([1, len(levels), carried, nfeat, t, ffrac, wfrac, len(net)]) + bytes(4)
    for lv in levels:
        out += struct.pack("<H", lv["side"])
    widths = [len(levels) * nfeat] + [len(layer["weights"]) for layer in net]
    out += struct.pack(f"<{len(widths)}H", *widths)
    out += struct.pack("<9i", *(q16(c) for k in ("eye", "target", "up") for c in view[k]))
    out += struct.pack("<4H", view["focal"], view["width"], view["height"], view["steps"])
    out += struct.pack("<2i", q16(view["t_near"]), q16(view["t_far"]))
    out += bytes(view["background"]) + b"\0"
    for layer, n_in in zip(net, widths):
        for row in layer["weights"]:
            if len(row) != n_in:
                raise ValueError("a layer's row does not match its input")
            out += struct.pack(f"<{n_in}h", *(q(w, wfrac, -32768, 32767) for w in row))
        out += struct.pack(f"<{len(layer['biases'])}i",
                           *(q(b, ffrac + wfrac, -(1 << 31), (1 << 31) - 1) for b in layer["biases"]))
    for lv in levels[:carried]:
        out += table_bytes(lv, ffrac)
    wil1 = []
    for k, lv in enumerate(levels[carried:], start=carried):
        head = b"WIL1" + bytes([1, k, nfeat, t]) + struct.pack("<HH", lv["side"], 0)
        wil1.append((k, head + table_bytes(lv, ffrac)))
    return bytes(out), wil1


def files_of(name, model):
    wig1, wil1 = export(model)
    out = {f"{name}.wig1": wig1}
    for k, b in wil1:
        out[f"{name}.l{k}.wil1"] = b
    return out


# ---------------------------------------------------------------- the grid, in float


def entry(side, log2_table, x, y, z):
    """The entry vertex (x, y, z) of a level reads (§5.2)."""
    n = side + 1
    if n * n * n <= 1 << log2_table:
        return x + n * (y + n * z)
    h = ((x * HASH_PRIMES[0]) & MASK32) ^ ((y * HASH_PRIMES[1]) & MASK32) ^ ((z * HASH_PRIMES[2]) & MASK32)
    return h & ((1 << log2_table) - 1)


def corners(side, log2_table, p):
    """The eight entries a point of [0, 1]^3 reads at a level, each with its trilinear weight."""
    cell, w = [], []
    for c in p:
        s = c * side
        i = min(int(math.floor(s)), side - 1)
        cell.append(i)
        w.append(s - i)
    out = []
    for k in range(8):
        d = (k & 1, (k >> 1) & 1, (k >> 2) & 1)
        wt = 1.0
        for a in range(3):
            wt *= w[a] if d[a] else 1.0 - w[a]
        out.append((entry(side, log2_table, cell[0] + d[0], cell[1] + d[1], cell[2] + d[2]), wt))
    return out


def lerp(table, cs, nfeat):
    out = [0.0] * nfeat
    for e, wt in cs:
        v = table[e]
        for c in range(nfeat):
            out[c] += wt * v[c]
    return out


def mix(seed):
    """splitmix64, as floats in [0, 1) of 53 bits."""
    s = seed
    while True:
        s = (s + 0x9E3779B97F4A7C15) & 0xFFFFFFFFFFFFFFFF
        z = s
        z = ((z ^ (z >> 30)) * 0xBF58476D1CE4E5B9) & 0xFFFFFFFFFFFFFFFF
        z = ((z ^ (z >> 27)) * 0x94D049BB133111EB) & 0xFFFFFFFFFFFFFFFF
        yield ((z ^ (z >> 31)) >> 11) / float(1 << 53)


# ---------------------------------------------------------------- the scenes


def orb_target(x, y, z):
    """A ball of density 8 with a box of density 6 beside it; colour by position."""
    dx, dy, dz = x - 0.45, y - 0.5, z - 0.5
    r2 = dx * dx + dy * dy + dz * dz
    sigma = 8.0 * max(0.0, 1.0 - r2 / 0.09)
    if 0.62 <= x <= 0.82 and 0.15 <= y <= 0.35 and 0.3 <= z <= 0.5:
        sigma += 6.0
    return [sigma, min(1.0, 0.15 + 0.85 * x), 0.25 + 0.6 * y, 0.9 - 0.7 * z]


def orb():
    """Four levels, the two finer hashed, fitted coarse to fine. The network sums the levels'
    features, so each level holds what the levels below it leave of the scene, and the payload
    carries the two coarser levels: the two finer are its refinements."""
    ffrac, t, sides, nfeat = 11, 9, [3, 7, 15, 31], 4
    g = mix(0x5741492D4649454C)
    pts = [(next(g), next(g), next(g)) for _ in range(3000)]
    targets = [orb_target(*p) for p in pts]
    below = [[0.0] * nfeat for _ in pts]
    scale = float(1 << ffrac)
    levels = []
    for side in sides:
        table = [[0.0] * nfeat for _ in range(entries(side, t))]
        cs = [corners(side, t, p) for p in pts]
        for _ in range(6):
            for i in range(len(pts)):
                pred = lerp(table, cs[i], nfeat)
                r = [targets[i][c] - below[i][c] - pred[c] for c in range(nfeat)]
                for e, wt in cs[i]:
                    v = table[e]
                    for c in range(nfeat):
                        v[c] += 0.35 * wt * r[c]
        table = [[q(x, ffrac, -32768, 32767) / scale for x in e] for e in table]
        for i in range(len(pts)):
            add = lerp(table, cs[i], nfeat)
            below[i] = [below[i][c] + add[c] for c in range(nfeat)]
        levels.append({"side": side, "table": table})
    n_in = len(sides) * nfeat
    w0 = [[0.0] * n_in for _ in range(2 * nfeat)]
    for c in range(nfeat):
        for lv in range(len(sides)):
            w0[c][lv * nfeat + c] = 1.0
            w0[nfeat + c][lv * nfeat + c] = -1.0
    w1 = [[0.0] * (2 * nfeat) for _ in range(nfeat)]
    for c in range(nfeat):
        w1[c][c], w1[c][nfeat + c] = 1.0, -1.0
    return {
        "form": FORM, "ffrac": ffrac, "wfrac": 8, "log2_table": t, "carried": 2, "levels": levels,
        "network": [{"weights": w0, "biases": [0.0] * (2 * nfeat)},
                    {"weights": w1, "biases": [0.0] * nfeat}],
        "view": {"eye": [0.5, 0.6, -1.0], "target": [0.5, 0.5, 0.5], "up": [0.0, 1.0, 0.0],
                 "focal": 22, "width": 20, "height": 20, "steps": 32,
                 "t_near": 0.9, "t_far": 2.4, "background": [16, 16, 24]},
    }


def mlp():
    """Three levels of seeded random features, two of them hashed, and a three-layer network
    of seeded random weights: the network's general case, every level carried."""
    g = mix(0x5741494649454C44)
    t, sides, nfeat, widths = 8, [2, 9, 40], 2, [16, 12, 4]
    levels = [{"side": s, "table": [[4.0 * next(g) - 2.0 for _ in range(nfeat)] for _ in range(entries(s, t))]}
              for s in sides]
    net, n_in = [], len(sides) * nfeat
    for w in widths:
        net.append({"weights": [[2.0 * next(g) - 1.0 for _ in range(n_in)] for _ in range(w)],
                    "biases": [next(g) - 0.5 for _ in range(w)]})
        n_in = w
    net[-1]["biases"][0] += 1.5
    return {
        "form": FORM, "ffrac": 10, "wfrac": 10, "log2_table": t, "carried": 3, "levels": levels,
        "network": net,
        "view": {"eye": [1.6, 1.3, -0.7], "target": [0.5, 0.5, 0.5], "up": [0.0, 1.0, 0.0],
                 "focal": 15, "width": 16, "height": 14, "steps": 24,
                 "t_near": 0.6, "t_far": 2.8, "background": [255, 255, 255]},
    }


def edge():
    """One level of side 1 and one layer with weights of 0 fraction bits: a density past the
    32-bit clamp that makes the medium opaque within a sample, colours past both ends of their
    range, and a view partly outside the domain."""
    s = 1.0 / 65536.0
    table = []
    for k in range(2):
        for j in range(2):
            for i in range(2):
                x0 = 32767 if (i + j + k) % 2 == 0 else -32768
                x1 = 20000 if k else -32768
                x2 = 32767 if i else -32768
                x3 = 32767 if j else -20000
                table.append([x0 * s, x1 * s, x2 * s, x3 * s])
    return {
        "form": FORM, "ffrac": 16, "wfrac": 0, "log2_table": 4, "carried": 1,
        "levels": [{"side": 1, "table": table}],
        "network": [{"weights": [[32767.0, 32767.0, 0.0, 0.0], [0.0, 0.0, 1.0, 0.0],
                                 [0.0, 0.0, 0.0, 1.0], [0.0, 0.0, -1.0, -1.0]],
                     "biases": [1073741824 * s, 20000 * s, -5 * s, 0.0]}],
        "view": {"eye": [-0.5, 0.5, 0.5], "target": [1.0, 0.5, 0.5], "up": [0.0, 1.0, 0.0],
                 "focal": 9, "width": 12, "height": 10, "steps": 8,
                 "t_near": 0.0, "t_far": 2.0, "background": [0, 128, 255]},
    }


DEMO = {"orb": orb, "mlp": mlp, "edge": edge}


def main(argv):
    if len(argv) >= 4 and argv[1] == "export":
        name = argv[argv.index("--name") + 1] if "--name" in argv else os.path.splitext(os.path.basename(argv[2]))[0]
        with open(argv[2], encoding="utf-8") as f:
            model = json.load(f)
        os.makedirs(argv[3], exist_ok=True)
        for fn, b in files_of(name, model).items():
            with open(os.path.join(argv[3], fn), "wb") as f:
                f.write(b)
            print(f"{fn}: {len(b)} bytes")
        return 0
    if len(argv) >= 3 and argv[1] == "demo":
        out = {}
        for name, build in DEMO.items():
            out.update(files_of(name, build()))
        if "--check" in argv:
            bad = []
            have = sorted(os.listdir(argv[2])) if os.path.isdir(argv[2]) else []
            if have != sorted(out):
                bad.append(f"{argv[2]} holds {have}, the demo writes {sorted(out)}")
            for fn, b in out.items():
                p = os.path.join(argv[2], fn)
                if os.path.exists(p):
                    with open(p, "rb") as f:
                        if f.read() != b:
                            bad.append(f"{fn} differs from what the demo writes")
            for x in bad:
                print("FAIL " + x)
            print("demo scenes: " + ("FAIL" if bad else f"{len(out)} files, byte for byte"))
            return 1 if bad else 0
        os.makedirs(argv[2], exist_ok=True)
        for fn, b in out.items():
            with open(os.path.join(argv[2], fn), "wb") as f:
                f.write(b)
            print(f"{fn}: {len(b)} bytes")
        return 0
    if len(argv) == 3 and argv[1] == "model" and argv[2] in DEMO:
        print(json.dumps(DEMO[argv[2]](), indent=1))
        return 0
    print(__doc__)
    return 2


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