"""Witness search for the dependence hypothesis of the size proposition.

For every coordinate q of the record region, exhibit a state x and a data-region
coordinate s with U(x)_s != U(x xor e_q)_s. Random states almost never witness
this because a random output address lands in the 106-bit data region with
probability about one in ten; the search below forces the record under
evaluation to write into the data region and randomizes everything else.
"""
import os, random, sys
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
import torch
from reflect import (Cfg, build_net, Leveled, encode_gate, pad, state_to_vec,
                     vec_to_state, slot_off, field_off)


def targeted(cfg, U, rng, q, g, r, D0, D1, budget=4000):
    """Second pass for address bits, index flags and weight codes: make the
    field under test decide the output, and place both the read and the write
    where they can be observed. Free coordinates are the data region and the
    inactive record region, neither of which the evaluated record occupies."""
    A, SLOT = cfg.A, cfg.SLOT
    free = list(range(D0, D1)) + list(range(cfg.NET1, cfg.S))
    freeset = set(free)
    Q0 = cfg.NET0
    base = Q0 + g * cfg.R
    for _ in range(budget):
        slot = 0 if r < SLOT else (1 if r < 2 * SLOT else None)
        # the tested slot decides the output: out = x at that slot's address
        if slot is None:
            w = (1, 1)
            bias = rng.randrange(-(1 << (cfg.BB - 1)), 1 << (cfg.BB - 1))
        else:
            w = (1, 0) if slot == 0 else (0, 1)
            bias = -1
        a = [rng.choice(free), rng.choice(free)]
        i = [rng.randint(0, 1), rng.randint(0, 1)]
        ao = rng.randrange(D0, D1)
        io = 0
        ptr = rng.randrange(0, 1 << A)
        rec = encode_gate(cfg, pad(cfg, [(a[0], w[0], i[0]), (a[1], w[1], i[1])]),
                          bias, (ao, io))
        sig = [0] * cfg.S
        for d in free:
            sig[d] = rng.randint(0, 1)
        for k in range(A):
            sig[cfg.PTR_BASE + k] = (ptr >> (A - 1 - k)) & 1
        sig[base:base + cfg.R] = rec
        s0, s1 = list(sig), list(sig)
        s1[q] ^= 1
        x = U(s0, g)
        y = U(s1, g)
        if any(x[d] != y[d] for d in range(D0, D1)):
            return budget
    return 0


def main() -> int:
    cfg = Cfg()
    net, inputs, outputs = build_net(cfg)
    lev = Leveled(net, inputs, outputs, device="cpu")

    def U(sig, gp):
        v = state_to_vec(cfg, {"sig": sig, "gp": gp, "halt": 0}).unsqueeze(0)
        return vec_to_state(cfg, lev.step(v[:, :len(inputs)])[0])["sig"]

    D0, D1 = cfg.WORK_BASE, cfg.NET0          # data region [D0, D1)
    rng = random.Random(20260910)
    Q0 = cfg.NET0
    Qbits = list(range(Q0, Q0 + cfg.G * cfg.R))
    witnessed, tries_used = [], []
    for q in Qbits:
        g, r = divmod(q - Q0, cfg.R)
        found = 0
        for attempt in range(160):
            # a record that writes into the data region, everything else random
            a0 = rng.randrange(D0, D1 - 1)
            a1 = rng.randrange(D0, D1 - 1)
            ao = rng.randrange(D0, D1 - 1)
            w0 = rng.choice((-1, 0, 1))
            w1 = rng.choice((-1, 0, 1))
            bias = rng.randrange(-(1 << (cfg.BB - 1)), 1 << (cfg.BB - 1))
            i0, i1, io = rng.randint(0, 1), rng.randint(0, 1), rng.randint(0, 1)
            rec = encode_gate(cfg, pad(cfg, [(a0, w0, i0), (a1, w1, i1)]),
                              bias, (ao, io))
            sig = [0] * cfg.S
            for d in range(D0, D1):
                sig[d] = rng.randint(0, 1)
            ptr = rng.choice((0, 1, 2, 3))
            for k in range(cfg.A):
                sig[cfg.PTR_BASE + k] = (ptr >> (cfg.A - 1 - k)) & 1
            base = cfg.bank_base(0) + g * cfg.R
            sig[base:base + cfg.R] = rec
            s0 = list(sig)
            s1 = list(sig)
            s1[q] ^= 1
            a = U(s0, g)
            b = U(s1, g)
            if any(a[d] != b[d] for d in range(D0, D1)):
                found = attempt + 1
                break
        if not found:
            found = targeted(cfg, U, rng, q, g, r, D0, D1)
        if found:
            witnessed.append(q)
            tries_used.append(found)
    ok = len(witnessed) == len(Qbits)
    print(f"record-region bits: {len(Qbits)}")
    print(f"bits with an exhibited witness: {len(witnessed)}"
          f"{'  (hypothesis holds for every bit)' if ok else '  MISSING: '
             + str([q - Q0 for q in Qbits if q not in witnessed][:20])}")
    if tries_used:
        print(f"attempts needed: max {max(tries_used)}, mean "
              f"{sum(tries_used)/len(tries_used):.1f}")
    print("every data coordinate of U is a multiplexer output depending on the "
          "write-select line and the previous value, hence on at least two "
          "coordinates")
    return 0 if ok else 1


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