"""Attractor computer: an energy-based threshold network.

A Boolean circuit is compiled to a quadratic pseudo-Boolean energy
E(s) = sum_i L[i] s_i + sum_{i<j} Q[i,j] s_i s_j over binary wire variables, with
per-gate gadgets that are non-negative and zero exactly on the gate truth table:

    AND z=x&y : 3z + xy - 2xz - 2yz
    OR  z=x|y : x + y + z + xy - 2xz - 2yz
    NOT z=~x  : 1 - x - z + 2xz

AND/OR/NOT are functionally complete, so any circuit compiles and its consistent
assignment is the global minimum. There is no program counter or clock: the
coupling matrix Q (with linear terms L) is the program, and execution is
relaxation toward the minimum. The relaxation step is a threshold neuron over the
same integer weights, s_i <- H(-(L[i] + sum_j Q[i,j] s_j)).

The clamped wire subset selects the mode. Clamp inputs to evaluate forward (exact
via topological propagation to the energy-0 fixed point); clamp outputs to invert
(a multiplier run backward returns factors); clamp a CNF output to 1 to solve
SAT. Forward evaluation is exact and linear in gate count; inversion and
open-constraint solving are annealed ground-state search, NP-hard in general,
with a zero-energy state certifying a correct assignment.
"""
from __future__ import annotations
import math
import random
from collections import defaultdict
from typing import Dict, List, Optional, Tuple


class Circuit:
    """Wire allocator and energy accumulator. Gates append exact QUBO gadgets
    and record the relation for topological forward evaluation."""

    def __init__(self) -> None:
        self.n = 0
        self.L: Dict[int, int] = defaultdict(int)
        self.Q: Dict[Tuple[int, int], int] = defaultdict(int)
        self.const = 0
        self.gates: List[Tuple[str, int, Tuple[int, ...]]] = []

    def wire(self) -> int:
        i = self.n
        self.n += 1
        return i

    def wires(self, k: int) -> List[int]:
        return [self.wire() for _ in range(k)]

    def _q(self, i: int, j: int, c: int) -> None:
        if i == j:
            self.L[i] += c
        else:
            self.Q[(min(i, j), max(i, j))] += c

    def AND(self, x: int, y: int) -> int:
        z = self.wire()
        self.L[z] += 3
        self._q(x, y, 1); self._q(x, z, -2); self._q(y, z, -2)
        self.gates.append(("AND", z, (x, y)))
        return z

    def OR(self, x: int, y: int) -> int:
        z = self.wire()
        self.L[x] += 1; self.L[y] += 1; self.L[z] += 1
        self._q(x, y, 1); self._q(x, z, -2); self._q(y, z, -2)
        self.gates.append(("OR", z, (x, y)))
        return z

    def NOT(self, x: int) -> int:
        z = self.wire()
        self.const += 1
        self.L[x] += -1; self.L[z] += -1
        self._q(x, z, 2)
        self.gates.append(("NOT", z, (x,)))
        return z

    def XOR(self, x: int, y: int) -> int:
        return self.OR(self.AND(x, self.NOT(y)), self.AND(self.NOT(x), y))

    def full_adder(self, x: int, y: int, cin: int) -> Tuple[int, int]:
        axy = self.XOR(x, y)
        s = self.XOR(axy, cin)
        cout = self.OR(self.AND(x, y), self.AND(cin, axy))
        return s, cout

    # ---- energy + couplings ------------------------------------------------
    def energy(self, s: List[int]) -> int:
        e = self.const
        for i, c in self.L.items():
            e += c * s[i]
        for (i, j), c in self.Q.items():
            e += c * s[i] * s[j]
        return e

    def neighbors(self) -> Dict[int, List[Tuple[int, int]]]:
        nbr: Dict[int, List[Tuple[int, int]]] = defaultdict(list)
        for (i, j), c in self.Q.items():
            nbr[i].append((j, c))
            nbr[j].append((i, c))
        return nbr

    # ---- relaxation modes --------------------------------------------------
    def forward_eval(self, clamp: Dict[int, int]) -> List[int]:
        """Exact forward relaxation: propagate clamped inputs through the gate
        relations in topological order onto the energy-0 fixed point."""
        s = [0] * self.n
        for w, v in clamp.items():
            s[w] = v
        for op, z, ins in self.gates:
            if op == "AND":
                s[z] = s[ins[0]] & s[ins[1]]
            elif op == "OR":
                s[z] = s[ins[0]] | s[ins[1]]
            else:
                s[z] = 1 - s[ins[0]]
        return s

    def relax_energy(self, clamp: Dict[int, int], sweeps: int = 4000,
                     t0: float = 4.0, t1: float = 0.02, seed: int = 0
                     ) -> Tuple[List[int], bool]:
        """Canonical relaxation: anneal the full threshold network (every free
        wire), tracking the lowest-energy state. Universal but hard; the exact
        gadgets keep the target at energy 0."""
        nbr = self.neighbors()
        rng = random.Random(seed)
        s = [rng.randint(0, 1) for _ in range(self.n)]
        for w, v in clamp.items():
            s[w] = v
        free = [i for i in range(self.n) if i not in clamp]
        best, best_e = list(s), self.energy(s)
        for step in range(sweeps):
            T = t0 * (t1 / t0) ** (step / max(1, sweeps - 1))
            for _ in range(len(free)):
                i = free[rng.randrange(len(free))]
                field = self.L[i] + sum(c * s[j] for j, c in nbr[i])
                dE = (1 - 2 * s[i]) * field
                if dE <= 0 or rng.random() < math.exp(-dE / T):
                    s[i] ^= 1
            e = self.energy(s)
            if e < best_e:
                best, best_e = list(s), e
                if best_e == 0:
                    return best, True
        return best, best_e == 0

    def solve(self, free_inputs: List[int], fixed: Dict[int, int],
              target: Dict[int, int], sweeps: int = 3000, restarts: int = 80,
              seed: int = 0) -> Optional[List[int]]:
        """Open-constraint relaxation over a chosen set of driver wires, with
        the rest slaved through the circuit; anneal the output Hamming mismatch
        to zero. Clamp outputs and pass the inputs here to run backward."""
        rng = random.Random(seed)

        def mism(vals: Dict[int, int]) -> int:
            s = self.forward_eval({**fixed, **vals})
            return sum(1 for w, v in target.items() if s[w] != v)

        for _ in range(restarts):
            vals = {w: rng.randint(0, 1) for w in free_inputs}
            m = mism(vals)
            if m == 0:
                return self.forward_eval({**fixed, **vals})
            for step in range(sweeps):
                T = 2.0 * (0.02 / 2.0) ** (step / sweeps)
                w = free_inputs[rng.randrange(len(free_inputs))]
                vals[w] ^= 1
                m2 = mism(vals)
                if m2 <= m or rng.random() < math.exp(-(m2 - m) / T):
                    m = m2
                    if m == 0:
                        return self.forward_eval({**fixed, **vals})
                else:
                    vals[w] ^= 1
        return None


# ---------------------------------------------------------------------------
# Circuit builders
# ---------------------------------------------------------------------------
def adder(bits: int) -> Tuple[Circuit, dict]:
    c = Circuit()
    xs, ys = c.wires(bits), c.wires(bits)
    cin = c.wire()
    outs, carry = [], cin
    for k in range(bits):
        s, carry = c.full_adder(xs[k], ys[k], carry)
        outs.append(s)
    return c, {"xs": xs, "ys": ys, "cin": cin, "sum": outs + [carry]}


def multiplier(bits: int) -> Tuple[Circuit, dict]:
    c = Circuit()
    xs, ys = c.wires(bits), c.wires(bits)
    zero = c.wire()
    acc = [zero] * (2 * bits)
    for i in range(bits):
        carry = zero
        for j in range(bits):
            acc[i + j], carry = c.full_adder(acc[i + j], c.AND(xs[i], ys[j]), carry)
        acc[i + bits] = carry
    return c, {"xs": xs, "ys": ys, "zero": zero, "prod": acc}


_OPCODE = {"AND": 0, "OR": 1, "NOT": 2}
_OPNAME = {v: k for k, v in _OPCODE.items()}


def to_tensors(circ: Circuit, io: dict):
    """Serialize the coupling matrix (the program) and the gate list to tensors.
    Q is stored sparsely as index pairs and integer values."""
    import torch
    qi = sorted(circ.Q)
    q_idx = torch.tensor(qi if qi else [], dtype=torch.long).reshape(-1, 2)
    q_val = torch.tensor([circ.Q[k] for k in qi], dtype=torch.long)
    li = sorted(circ.L)
    l_idx = torch.tensor(li, dtype=torch.long)
    l_val = torch.tensor([circ.L[i] for i in li], dtype=torch.long)
    g_op = torch.tensor([_OPCODE[op] for op, _, _ in circ.gates], dtype=torch.long)
    g_out = torch.tensor([o for _, o, _ in circ.gates], dtype=torch.long)
    g_in = torch.tensor([[ins[0], ins[1] if len(ins) > 1 else -1]
                         for _, _, ins in circ.gates], dtype=torch.long).reshape(-1, 2)
    t = {"Q_idx": q_idx, "Q_val": q_val, "L_idx": l_idx, "L_val": l_val,
         "gate_op": g_op, "gate_out": g_out, "gate_in": g_in}
    import json
    meta = {"machine": "attractor", "n": str(circ.n), "const": str(circ.const),
            "io": json.dumps({k: v for k, v in io.items()})}
    return t, meta


def from_tensors(t: dict, meta: dict) -> Tuple[Circuit, dict]:
    import json
    c = Circuit()
    c.n = int(meta["n"])
    c.const = int(meta["const"])
    for (i, j), v in zip(t["Q_idx"].tolist(), t["Q_val"].tolist()):
        c.Q[(i, j)] = v
    for i, v in zip(t["L_idx"].tolist(), t["L_val"].tolist()):
        c.L[i] = v
    for op, out, ins in zip(t["gate_op"].tolist(), t["gate_out"].tolist(), t["gate_in"].tolist()):
        c.gates.append((_OPNAME[op], out, tuple(x for x in ins if x >= 0)))
    return c, json.loads(meta["io"])


def cnf(clauses: List[List[int]], n_vars: int) -> Tuple[Circuit, dict]:
    """Compile a CNF formula. Literals are +v (var v) or -v (negation), v>=1.
    Returns the circuit, the variable wires, and the wire that is 1 iff the
    formula is satisfied. Clamp that wire to 1 and relax to find a model."""
    c = Circuit()
    var = {v: c.wire() for v in range(1, n_vars + 1)}
    clause_ws = []
    for cl in clauses:
        lits = [var[abs(l)] if l > 0 else c.NOT(var[abs(l)]) for l in cl]
        acc = lits[0]
        for w in lits[1:]:
            acc = c.OR(acc, w)
        clause_ws.append(acc)
    sat = clause_ws[0]
    for w in clause_ws[1:]:
        sat = c.AND(sat, w)
    return c, {"vars": var, "sat": sat}