"""Ship the reversible CA's block rule as a ternary matrix tile,
variants/neural_ca.safetensors. The 2x2 block rule compiles to a stack of
ternary matrices with a Heaviside step; the rule is a bijection, so the tile is
a permutation matrix product with a 0.5 analog margin, and the same tile applied
to every block of a lattice is one lattice step."""
from __future__ import annotations
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 ca
from matrix8 import Net, compile_net, MatrixMachine

OUT = os.path.join(ROOT, "variants", "neural_ca.safetensors")


def block_net():
    net = Net()
    ntl, ntr = net.NOT("ntl", "tl"), net.NOT("ntr", "tr")
    nbl, nbr = net.NOT("nbl", "bl"), net.NOT("nbr", "br")
    d1 = net.AND("d1", ["tl", ntr, nbl, "br"])
    d2 = net.AND("d2", [ntl, "tr", "bl", nbr])
    dg = net.OR("dg", [d1, d2])                       # is_diag
    TL = net.XOR("oTL", "br", dg)
    TR = net.XOR("oTR", "bl", dg)
    BL = net.XOR("oBL", "tr", dg)
    BR = net.XOR("oBR", "tl", dg)
    return net, ["tl", "tr", "bl", "br"], [TL, TR, BL, BR]


def matrix_step(mm, grid, phase):
    """One Margolus update driven entirely by the matrix tile."""
    H, W = len(grid), len(grid[0])
    out = [row[:] for row in grid]
    for r0 in range(phase, phase + H, 2):
        for c0 in range(phase, phase + W, 2):
            r, r1, c, c1 = r0 % H, (r0 + 1) % H, c0 % W, (c0 + 1) % W
            v = torch.tensor([[float(grid[r][c]), float(grid[r][c1]),
                               float(grid[r1][c]), float(grid[r1][c1])]])
            o = mm.step(v)[0]
            out[r][c], out[r][c1], out[r1][c], out[r1][c1] = (int(o[0]), int(o[1]),
                                                              int(o[2]), int(o[3]))
    return out


def main() -> int:
    net, inp, outp = block_net()
    layers, info = compile_net(net, inp, outp)
    mm = MatrixMachine(layers)

    seen, bad, vecs = set(), 0, []
    for s in range(16):
        b = tuple((s >> k) & 1 for k in (3, 2, 1, 0))     # tl,tr,bl,br
        v = torch.tensor([[float(x) for x in b]])
        got = tuple(int(x) for x in mm.step(v)[0])
        if got != ca.rule(b):
            bad += 1
        seen.add(got)
        vecs.append(v[0])
    perm = len(seen) == 16
    margin = mm.min_margin(torch.stack(vecs))

    tensors = {}
    for k, (W, B) in enumerate(layers):
        tensors[f"matrix.layer{k:03d}.weight"] = W.to(torch.int8)
        tensors[f"matrix.layer{k:03d}.bias"] = B.to(torch.int8)
    meta = {"machine": "ca",
            "rule": "Margolus reversible: rotate 180 except diagonal pair swaps (BBM class)",
            "inputs": "tl,tr,bl,br", "outputs": "TL,TR,BL,BR", "layers": str(info["layers"])}
    save_file(tensors, OUT, metadata=meta)
    print(f"Built {os.path.relpath(OUT, ROOT)}: reversible CA block rule as a ternary matrix tile")
    print(f"  layers={info['layers']}  gates={info['gates']}  size={os.path.getsize(OUT)} bytes")
    print(f"  every weight ternary: "
          f"{'OK' if all(((W == -1) | (W == 0) | (W == 1)).all() for W, _ in layers) else 'FAIL'}")
    print(f"  tile matches the block rule over all 16 states: {'OK' if bad == 0 else f'FAIL({bad})'}")
    print(f"  tile is a permutation (16 distinct outputs): {'OK' if perm else 'FAIL'}")
    print(f"  analog noise margin: {margin:.3f} (guarantee 0.5)")

    # analog robustness: the tile must reproduce the rule under read noise and
    # static conductance mismatch, as neural_matrix8 measures
    states = torch.stack(vecs)
    ref = [ca.rule(tuple((s >> k) & 1 for k in (3, 2, 1, 0))) for s in range(16)]

    def outs(machine, **kw):
        v = machine.step(states, **kw)
        return [tuple(int(v[i][j]) for j in range(4)) for i in range(v.shape[0])]

    print("  read-noise sweep (exact = tile matches the rule on all 16 states):")
    for sigma in (0.05, 0.10, 0.15, 0.20):
        okn = all(outs(mm, analog=True, noise_sigma=sigma,
                       gen=torch.Generator().manual_seed(s)) == ref for s in range(4))
        print(f"    sigma={sigma:.2f}: {'exact' if okn else 'errors appear'}")
    print("  conductance-mismatch sweep:")
    for sg in (0.05, 0.10, 0.15):
        okg = outs(mm.perturbed(sg, seed=0), analog=True) == ref
        print(f"    sigma_G={sg:.2f}: {'exact' if okg else 'errors appear'}")

    # the loaded tile, applied to every block, is one whole-lattice CA step
    t = load_file(OUT)
    n = 0
    lyr = []
    while f"matrix.layer{n:03d}.weight" in t:
        lyr.append((t[f"matrix.layer{n:03d}.weight"], t[f"matrix.layer{n:03d}.bias"]))
        n += 1
    mm2 = MatrixMachine(lyr)
    g = ca._rand_grid(8, 8, 3)
    gmat = matrix_step(mm2, g, 0)
    gref = ca.step(g, 0)
    print(f"  loaded tile drives a full lattice step (matches ca.step): "
          f"{'OK' if gmat == gref else 'FAIL'}")

    ok = bad == 0 and perm and abs(margin - 0.5) < 1e-6 and gmat == gref
    print("PASS" if ok else "FAIL")
    return 0 if ok else 1


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