"""Ship the self-assembling tile computer as variants/neural_tile.safetensors: a
tile set (the binary counter) stored as its glue tables, together with the
binding gate that governs growth. A tile binds at a site when the summed
strength of its matching glues meets tau, which is the Heaviside gate
H(strength . match - tau) with per-tile weights = glue strengths and bias = -tau.
Round-trips the file, regrows the counter, and confirms row y encodes y."""
from __future__ import annotations
import json
import os
import sys

import torch
from safetensors.torch import save_file, load_file
from safetensors import safe_open

ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
sys.path.insert(0, os.path.join(ROOT, "src"))
import tile as T

OUT = os.path.join(ROOT, "variants", "neural_tile.safetensors")
NBITS = 8
TAU = 2


def main() -> int:
    ts = T.counter_tileset()
    strength = {"edge": 2}
    glues = sorted({g for t in ts for g in (t.N, t.E, t.S, t.W) if g})
    gid = {g: i for i, g in enumerate(glues)}
    tile_glues = torch.tensor([[gid.get(t.N, -1), gid.get(t.E, -1),
                                gid.get(t.S, -1), gid.get(t.W, -1)] for t in ts],
                              dtype=torch.long)
    glue_strength = torch.tensor([strength.get(g, 1) for g in glues], dtype=torch.long)
    # per-tile binding-gate weights = strengths of the tile's four glues (0 = null)
    bind_w = torch.tensor([[strength.get(g, 1) if g else 0 for g in (t.N, t.E, t.S, t.W)]
                           for t in ts], dtype=torch.long)
    tensors = {"tile_glues": tile_glues, "glue_strength": glue_strength,
               "binding_weight": bind_w, "binding_bias": torch.tensor(-TAU)}
    meta = {"machine": "tile", "tau": str(TAU), "glues": json.dumps(glues),
            "tile_names": json.dumps([t.name for t in ts]), "program": "binary counter"}
    save_file(tensors, OUT, metadata=meta)
    print(f"Built {os.path.relpath(OUT, ROOT)}: binary-counter tile set")
    print(f"  tiles={len(ts)}  glues={len(glues)}  tau={TAU}  size={os.path.getsize(OUT)} bytes")

    # round-trip: reconstruct the tiles from the file and regrow the counter
    t = load_file(OUT)
    with safe_open(OUT, framework="pt") as f:
        m = f.metadata()
    gl = json.loads(m["glues"])
    strg = {gl[i]: int(s) for i, s in enumerate(t["glue_strength"].tolist())}
    tiles = []
    for row, name in zip(t["tile_glues"].tolist(), json.loads(m["tile_names"])):
        sides = [gl[i] if i >= 0 else "" for i in row]
        tiles.append(T.Tile(N=sides[0], E=sides[1], S=sides[2], W=sides[3], name=name))

    rows = (1 << NBITS) - 1
    A, det = T.grow(tiles, T.counter_seed(NBITS), int(m["tau"]), strg,
                    (0, 0, NBITS, rows))
    bad = filled = 0
    for y in range(1, rows + 1):
        cells = [A.get((x, y)) for x in range(NBITS)]
        if any(c is None for c in cells):
            continue
        filled += 1
        v = sum((1 if c.N == "b1" else 0) << (NBITS - 1 - x) for x, c in enumerate(cells))
        if v != (y & ((1 << NBITS) - 1)):
            bad += 1
    print(f"  round-trip regrow {NBITS}-bit counter: directed={det}  rows={filled}  "
          f"row y == y {'OK' if bad == 0 else f'FAIL({bad})'}")

    # the stored binding gate reproduces the model's binding decision
    gate_ok = True
    Atest = {(1, 0): T.Tile(N="b0"), (2, 0): T.Tile(N="b0")}
    for ti, tt in enumerate(tiles):
        for site in [(1, 1), (2, 1)]:
            w = t["binding_weight"][ti].tolist()
            match = [1 if tt.glue(side) and Atest.get((site[0] + dx, site[1] + dy))
                     and tt.glue(side) == Atest[(site[0] + dx, site[1] + dy)].glue(opp) else 0
                     for dx, dy, side, opp in T._SIDES]
            gate = 1 if sum(wi * mi for wi, mi in zip(w, match)) + int(t["binding_bias"]) >= 0 else 0
            if gate != int(T.binds(Atest, site[0], site[1], tt, TAU, strg)):
                gate_ok = False
    print(f"  stored binding gate H(weight.match - tau) matches growth rule: "
          f"{'OK' if gate_ok else 'FAIL'}")
    return 0 if (bad == 0 and det and gate_ok) else 1


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