"""Exercise the attractor computer: exact forward evaluation, the canonical
whole-network energy relaxation, backward inversion (factoring), and SAT
solving (universality of the solve direction)."""
from __future__ import annotations
import os
import random
import sys

sys.path.insert(0, os.path.join(os.path.dirname(os.path.dirname(os.path.abspath(__file__))), "src"))
from attractor import Circuit, adder, multiplier, cnf


def test_forward():
    ok = True
    for bits in (4, 8):
        c, io = adder(bits)
        rng = random.Random(bits)
        bad = 0
        for _ in range(300):
            a, b = rng.randint(0, (1 << bits) - 1), rng.randint(0, (1 << bits) - 1)
            clamp = {io["cin"]: 0}
            for k in range(bits):
                clamp[io["xs"][k]] = (a >> k) & 1
                clamp[io["ys"][k]] = (b >> k) & 1
            s = c.forward_eval(clamp)
            got = sum(s[w] << k for k, w in enumerate(io["sum"]))
            if got != a + b or c.energy(s) != 0:
                bad += 1
        print(f"  forward adder {bits}-bit: {'OK' if bad == 0 else f'FAIL({bad})'}")
        ok &= bad == 0
    for bits in (3, 5):
        c, io = multiplier(bits)
        rng = random.Random(100 + bits)
        bad = 0
        for _ in range(300):
            a, b = rng.randint(0, (1 << bits) - 1), rng.randint(0, (1 << bits) - 1)
            clamp = {io["zero"]: 0}
            for k in range(bits):
                clamp[io["xs"][k]] = (a >> k) & 1
                clamp[io["ys"][k]] = (b >> k) & 1
            s = c.forward_eval(clamp)
            got = sum(s[w] << k for k, w in enumerate(io["prod"]))
            if got != a * b or c.energy(s) != 0:
                bad += 1
        print(f"  forward multiplier {bits}-bit: {'OK' if bad == 0 else f'FAIL({bad})'}")
        ok &= bad == 0
    return ok


def test_energy_relax():
    """The canonical form: anneal the whole network (no propagation shortcut)."""
    c, io = adder(4)
    rng = random.Random(3)
    bad = 0
    for _ in range(20):
        a, b = rng.randint(0, 15), rng.randint(0, 15)
        clamp = {io["cin"]: 0}
        for k in range(4):
            clamp[io["xs"][k]] = (a >> k) & 1
            clamp[io["ys"][k]] = (b >> k) & 1
        conv = False
        for attempt in range(4):                 # annealers restart
            s, conv = c.relax_energy(clamp, sweeps=6000, seed=rng.randint(0, 1 << 30))
            got = sum(s[w] << k for k, w in enumerate(io["sum"]))
            if conv and got == a + b:
                break
        if not conv:
            bad += 1
    print(f"  whole-network energy relaxation (4-bit adder, 20 cases): "
          f"{'OK' if bad == 0 else f'reached ground state in {20 - bad}/20'}")
    return bad == 0


def test_factor():
    ok = True
    for bits, targets in ((4, [15, 35, 143]), (5, [21, 55, 91])):
        c, io = multiplier(bits)
        for N in targets:
            target = {io["prod"][k]: (N >> k) & 1 for k in range(2 * bits)}
            s = c.solve(io["xs"] + io["ys"], {io["zero"]: 0}, target, seed=N)
            if s is None:
                print(f"  factor {N}: not found")
                ok = False
                continue
            a = sum(s[io["xs"][k]] << k for k in range(bits))
            b = sum(s[io["ys"][k]] << k for k in range(bits))
            good = a * b == N and 1 < a < N and 1 < b < N
            print(f"  factor {N} ({bits}x{bits}): {a} x {b} {'OK' if a * b == N else 'WRONG'}")
            ok &= a * b == N
    return ok


def test_sat():
    # (x1 | x2 | ~x3) & (~x1 | x3) & (x2 | x3) & (~x2 | ~x3), a satisfiable 3-SAT.
    clauses = [[1, 2, -3], [-1, 3], [2, 3], [-2, -3]]
    c, io = cnf(clauses, 3)
    s = c.solve(list(io["vars"].values()), {}, {io["sat"]: 1}, seed=1)
    if s is None:
        print("  SAT solve: no model found")
        return False
    assign = {v: s[w] for v, w in io["vars"].items()}
    sat = all(any((assign[abs(l)] == 1) if l > 0 else (assign[abs(l)] == 0) for l in cl)
              for cl in clauses)
    print(f"  SAT solve: model {assign} {'satisfies' if sat else 'FAILS'} the formula")
    return sat


if __name__ == "__main__":
    print("Attractor computer\n" + "=" * 40)
    print("Forward evaluation (exact, energy 0):")
    a = test_forward()
    print("Canonical relaxation:")
    b = test_energy_relax()
    print("Backward inversion (factoring by relaxation):")
    c_ = test_factor()
    print("SAT (clamp output to 1, relax to a model):")
    d = test_sat()
    print("=" * 40)
    print("ALL PASS" if (a and b and c_ and d) else "FAILURES")
    sys.exit(0 if (a and b and c_ and d) else 1)