flagship_ce/solver.py

"""leaner_ce — a lean, dependency-free counterexample engine for the Solo protocol.

The FALSE-side engine for a flagship solver: given `E1 ⇒ E2?`, find a finite
magma (Fin n, ◇) where E1 holds for all assignments and E2 fails for some
assignment, and submit it as a Lean `decideFin!` certificate.

Single file, Python stdlib only (no pysat / no LLM), well under 500 KB. Escalating
deterministic ladder:
  1. exhaustive Fin 2-3
  2. structured table families Fin 2-7 + direct products
  3. constraint-propagation backtracking Fin 4-6
  4. value-symmetry targeted backtracking Fin 6 (pin E2 sides to 0/1, propagate)

On the 850 public FALSE problems this clears 848 (only hard2_0027/0051 need a
large magma). It has no true-implication prover — an integrator pairs it with a
proof stage (deterministic miners + LLM) for the full flagship.

Assembled from the tested cebench modules (core/search/backtrack); see
artifacts/ce-backtrack_leaner/ for the development version + benchmarks.
"""
import json
import re
import sys
import time
from itertools import product, permutations

DIAMOND = "◇"


# ── protocol ──

def read_message():
    line = sys.stdin.readline()
    if not line:
        sys.exit(0)
    return json.loads(line.strip())


def send_message(msg):
    print(json.dumps(msg), flush=True)


def call_judge(verdict, code):
    send_message({"call": "judge", "verdict": verdict, "code": code})
    return read_message()


# ── equation parse + codegen evaluator ──

def _normalize(text):
    return text.replace("*", DIAMOND)


def parse_variables(text):
    seen, out = set(), []
    for v in re.findall(r"\b([a-z])\b", text):
        if v not in seen:
            seen.add(v)
            out.append(v)
    return out


def _parse_expr(s, vars_set):
    s = s.strip()
    while len(s) >= 2 and s[0] == "(" and s[-1] == ")":
        depth = 0
        matched = True
        for i, c in enumerate(s):
            if c == "(":
                depth += 1
            elif c == ")":
                depth -= 1
            if depth == 0 and i < len(s) - 1:
                matched = False
                break
        if matched:
            s = s[1:-1].strip()
        else:
            break
    depth = 0
    last_op = -1
    for i, c in enumerate(s):
        if c == "(":
            depth += 1
        elif c == ")":
            depth -= 1
        elif c == DIAMOND and depth == 0:
            last_op = i
    if last_op >= 0:
        return ("op", _parse_expr(s[:last_op], vars_set),
                _parse_expr(s[last_op + 1:], vars_set))
    s = s.strip()
    if len(s) == 1 and s in vars_set:
        return s
    raise ValueError(f"cannot parse: {s!r}")


def _expr_code(ast):
    if isinstance(ast, str):
        return ast
    _, l, r = ast
    return f"t[{_expr_code(l)}][{_expr_code(r)}]"


class Equation:
    __slots__ = ("variables", "lhs", "rhs", "_holds", "_violated")

    def __init__(self, text):
        text = _normalize(text)
        self.variables = parse_variables(text)
        vs = set(self.variables)
        lhs_str, rhs_str = text.split("=", 1)
        self.lhs = _parse_expr(lhs_str, vs)
        self.rhs = _parse_expr(rhs_str, vs)
        k = len(self.variables)
        if k == 0:
            unpack = "_"
        elif k == 1:
            unpack = self.variables[0] + ","
        else:
            unpack = ", ".join(self.variables)
        lc, rc = _expr_code(self.lhs), _expr_code(self.rhs)
        ns = {}
        exec(f"def _h(t, A):\n for {unpack} in A:\n  if {lc}!={rc}: return False\n return True\n", ns)
        exec(f"def _v(t, A):\n for {unpack} in A:\n  if {lc}!={rc}: return True\n return False\n", ns)
        self._holds, self._violated = ns["_h"], ns["_v"]

    def holds(self, t, A):
        return self._holds(t, A)

    def violated(self, t, A):
        return self._violated(t, A)


_AC = {}


def assigns(n, k):
    key = (n, k)
    a = _AC.get(key)
    if a is None:
        a = _AC[key] = list(product(range(n), repeat=k))
    return a


# ── structured table families ──

def structured_tables(n):
    for c in range(n):
        yield [[c] * n for _ in range(n)]
    yield [[i] * n for i in range(n)]
    yield [list(range(n)) for _ in range(n)]
    yield [[(i + j) % n for j in range(n)] for i in range(n)]
    yield [[(i - j) % n for j in range(n)] for i in range(n)]
    yield [[max(i, j) for j in range(n)] for i in range(n)]
    yield [[min(i, j) for j in range(n)] for i in range(n)]
    yield [[i if i != 0 else j for j in range(n)] for i in range(n)]
    yield [[j if j != 0 else i for j in range(n)] for i in range(n)]
    for k in range(1, n):
        yield [[(i + k) % n] * n for i in range(n)]
        yield [[(j + k) % n for j in range(n)] for _ in range(n)]
    if n > 1:
        yield [[(i * j) % n for j in range(n)] for i in range(n)]
    if n in (2, 4):
        yield [[(i ^ j) % n for j in range(n)] for i in range(n)]
    for c in range(n):
        for thresh in range(1, n):
            yield [[i if i < thresh else c for _ in range(n)] for i in range(n)]
            yield [[j if j < thresh else c for j in range(n)] for _ in range(n)]
    yield [[i if i >= j else j for j in range(n)] for i in range(n)]
    yield [[i if i <= j else j for j in range(n)] for i in range(n)]
    yield [[i for _ in range(n)] for i in range(n)]
    yield [[j for j in range(n)] for _ in range(n)]
    if n >= 2:
        yield [[0 if i != j else i for j in range(n)] for i in range(n)]
        yield [[(i + j + 1) % n for j in range(n)] for i in range(n)]
    if n >= 3:
        yield [[i if i == j else (i + j) % n for j in range(n)] for i in range(n)]
        yield [[i if i == j else 0 for j in range(n)] for i in range(n)]
        yield [[i if i == j else n - 1 for j in range(n)] for i in range(n)]
    if n <= 4:
        for perm in permutations(range(n)):
            yield [[perm[j] for j in range(n)] for _ in range(n)]
            yield [[perm[i] for _ in range(n)] for i in range(n)]
    if n >= 4:
        for d in range(2, n):
            if n % d == 0:
                m = n // d
                yield [[(i // m) * m + (j % m) for j in range(n)] for i in range(n)]
                yield [[(i % d) + (j // d) * d for j in range(n)] for i in range(n)]
    if n <= 5:
        for chooser in range(min(2 ** (n * n), 1025)):
            table = [[0] * n for _ in range(n)]
            valid = True
            for i in range(n):
                for j in range(n):
                    bit = (chooser >> (i * n + j)) & 1
                    table[i][j] = i if bit else j
                    if i == j and table[i][j] != i:
                        valid = False
                        break
                if not valid:
                    break
            if valid:
                yield table
    for c in range(n):
        t = [[c] * n for _ in range(n)]
        for i in range(n):
            t[i][i] = i
        yield t
    if n >= 2:
        yield [[(n - 1 - i) for _ in range(n)] for i in range(n)]
        yield [[(n - 1 - j) for j in range(n)] for _ in range(n)]
        yield [[(n - 1 - i + j) % n for j in range(n)] for i in range(n)]
        yield [[(i + n - 1 - j) % n for j in range(n)] for i in range(n)]
    for a in range(n):
        for b in range(n):
            if (a, b) in ((0, 0), (1, 1)):
                continue
            yield [[(a * i + b * j) % n for j in range(n)] for i in range(n)]
    for a in range(1, min(n, 4)):
        for b in range(1, min(n, 4)):
            for c in range(1, min(n, 3)):
                yield [[(a * i + b * j + c) % n for j in range(n)] for i in range(n)]


def product_tables(n):
    for p in range(2, n):
        if n % p:
            continue
        q = n // p
        for a1 in range(p):
            for b1 in range(q):
                for a2 in range(p):
                    for b2 in range(q):
                        if a1 == a2 == b1 == b2 == 0:
                            continue
                        table = [[0] * n for _ in range(n)]
                        for i in range(n):
                            for j in range(n):
                                r, s = i // q, i % q
                                tt, u = j // q, j % q
                                table[i][j] = ((a1 * r + a2 * tt) % p) * q + (b1 * s + b2 * u) % q
                        yield table


# ── searches ──

def search_exhaustive(e1, e2, max_n=3):
    k1, k2 = len(e1.variables), len(e2.variables)
    for n in range(2, max_n + 1):
        A1, A2 = assigns(n, k1), assigns(n, k2)
        for enc in range(n ** (n * n)):
            t = [[(enc // (n ** (i * n + j))) % n for j in range(n)] for i in range(n)]
            if e1.holds(t, A1) and e2.violated(t, A2):
                return n, t
    return None, None


def search_structured(e1, e2, max_n=7):
    k1, k2 = len(e1.variables), len(e2.variables)
    for n in range(2, max_n + 1):
        A1, A2 = assigns(n, k1), assigns(n, k2)
        for t in structured_tables(n):
            if e1.holds(t, A1) and e2.violated(t, A2):
                return n, t
    for n in range(4, 10):
        A1, A2 = assigns(n, k1), assigns(n, k2)
        for t in product_tables(n):
            if e1.holds(t, A1) and e2.violated(t, A2):
                return n, t
    return None, None


# constraint-propagation backtracking (see cebench/backtrack.py)
KNOWN, CELL, OPEN = 0, 1, 2


def _ground(ast, env):
    if isinstance(ast, str):
        return env[ast]
    _, l, r = ast
    return (_ground(l, env), _ground(r, env))


def _gcons(eq, n):
    k = len(eq.variables)
    out = []
    for a in assigns(n, k):
        env = dict(zip(eq.variables, a))
        out.append((_ground(eq.lhs, env), _ground(eq.rhs, env)))
    return out


def _eval(gt, T, n):
    if isinstance(gt, int):
        return (KNOWN, gt)
    l, r = gt
    sl = _eval(l, T, n)
    if sl[0] != KNOWN:
        return (OPEN, sl[1])
    sr = _eval(r, T, n)
    if sr[0] != KNOWN:
        return (OPEN, sr[1])
    cell = sl[1] * n + sr[1]
    v = T[cell]
    return (KNOWN, v) if v != -1 else (CELL, cell)


def _propagate(T, n, cons, trail):
    progress = True
    while progress:
        progress = False
        for lg, rg in cons:
            s1, v1 = _eval(lg, T, n)
            s2, v2 = _eval(rg, T, n)
            if s1 == KNOWN and s2 == KNOWN:
                if v1 != v2:
                    return False
            elif s1 == KNOWN and s2 == CELL:
                T[v2] = v1
                trail.append(v2)
                progress = True
            elif s2 == KNOWN and s1 == CELL:
                T[v1] = v2
                trail.append(v1)
                progress = True
    return True


def _pick(T, n, cons):
    for lg, rg in cons:
        s1, v1 = _eval(lg, T, n)
        s2, v2 = _eval(rg, T, n)
        if s1 == KNOWN and s2 == KNOWN:
            continue
        if s1 in (CELL, OPEN):
            return v1
        if s2 in (CELL, OPEN):
            return v2
    return None


def _e2_complete(T, n, e2):
    A2 = assigns(n, len(e2.variables))
    base = [x if x != -1 else 0 for x in T]
    tbl = [base[i * n:(i + 1) * n] for i in range(n)]
    if e2.violated(tbl, A2):
        return tbl
    free = [i for i in range(n * n) if T[i] == -1]
    if free and n ** len(free) <= 20000:
        for combo in product(range(n), repeat=len(free)):
            for c, val in zip(free, combo):
                base[c] = val
            tbl = [base[i * n:(i + 1) * n] for i in range(n)]
            if e2.violated(tbl, A2):
                return tbl
    return None


def search_backtrack(e1, e2, n, deadline):
    cons = _gcons(e1, n)
    T = [-1] * (n * n)

    def rec():
        if time.time() > deadline:
            return None
        cell = _pick(T, n, cons)
        if cell is None:
            return _e2_complete(T, n, e2)
        for val in range(n):
            trail = [cell]
            T[cell] = val
            if _propagate(T, n, cons, trail):
                res = rec()
                if res is not None:
                    return res
            for c in trail:
                T[c] = -1
        return None

    return rec()


def search_targeted(e1, e2, n, deadline):
    """Pin E2.lhs(a2)=0, E2.rhs(a2)=1 (WLOG by value symmetry) and propagate."""
    cons1 = _gcons(e1, n)
    for a2 in assigns(n, len(e2.variables)):
        if time.time() > deadline:
            return None
        env = dict(zip(e2.variables, a2))
        l2, r2 = _ground(e2.lhs, env), _ground(e2.rhs, env)
        if l2 == r2:
            continue
        cons = cons1 + [(l2, 0), (r2, 1)]
        T = [-1] * (n * n)

        def rec():
            if time.time() > deadline:
                return None
            cell = _pick(T, n, cons)
            if cell is None:
                return [[(T[i * n + j] if T[i * n + j] != -1 else 0)
                         for j in range(n)] for i in range(n)]
            for val in range(n):
                trail = [cell]
                T[cell] = val
                if _propagate(T, n, cons, trail):
                    res = rec()
                    if res is not None:
                        return res
                for c in trail:
                    T[c] = -1
            return None

        res = rec()
        if res is not None:
            return res
    return None


def search_affine(e1, e2, min_n=2, max_n=32, deadline=None):
    """Parametric affine sweep: op(x,y) = (a*x + b*y + c) mod n over a,b,c in Z_n.
    Reaches affine counterexamples at orders that the Fin<=7 structured /
    backtracking / targeted search all miss, at ~ms per candidate — fully
    deterministic and general (not a lookup). E.g. hard2_0051 is 7*(x+y) mod 13,
    invisible to arbitrary-cell search but instant here. (Credit: django-cat.)
    Returns (n, a, b, c) so the caller emits a finOpTable-free direct-Fin cert
    (make_false_affine_code), which is SOUND at every order (unlike finOpTable,
    which mis-parses entries >=10)."""
    k1, k2 = len(e1.variables), len(e2.variables)
    for n in range(min_n, max_n + 1):
        if deadline and time.time() > deadline:
            return None, None, None, None
        A1, A2 = assigns(n, k1), assigns(n, k2)
        for a in range(n):
            for b in range(n):
                for c in range(n):
                    t = [[(a * i + b * j + c) % n for j in range(n)] for i in range(n)]
                    if e1.holds(t, A1) and e2.violated(t, A2):
                        return n, a, b, c
    return None, None, None, None


# ── certificate ──

def make_false_code(n, table):
    return (
        "import JudgeProblem\n"
        "import JudgeDecide.DecideBang\n"
        "import JudgeFinOp.MemoFinOp\n"
        "open MemoFinOp\n\n"
        "set_option maxRecDepth 100000 in\n"
        "set_option maxHeartbeats 4000000 in\n"  # decideFin! at Fin>=11 exceeds the 200k default
        "def submission : Goal := by\n"
        f"  let m : Magma (Fin {n}) := {{\n"
        f"    op := finOpTable \"{json.dumps(table)}\"\n"
        f"  }}\n"
        f"  refine ⟨Fin {n}, m, ?_⟩\n"
        f"  decideFin!\n"
    )


def make_false_affine_code(n, a, b, c):
    """Affine op given DIRECTLY as x ◇ y = (a*x + b*y + c) mod n via Fin.mul /
    Fin.add — no finOpTable, so it is SOUND at every order (finOpTable's
    extractDigits mis-parses table entries >=10). The judge's declaration
    allowlist admits Fin.* (but not the HMul./HAdd. wrappers). (Template credit:
    django-cat's hard2_0051 cert, x ◇ y = 7*(x+y) mod 13.)"""
    op = (f"fun x y => Fin.add (Fin.add (Fin.mul ({a} : Fin {n}) x) "
          f"(Fin.mul ({b} : Fin {n}) y)) ({c} : Fin {n})")
    return (
        "import JudgeProblem\n"
        "import JudgeDecide.DecideBang\n\n"
        "set_option maxRecDepth 100000 in\n"
        "set_option maxHeartbeats 4000000 in\n"  # decideFin! at Fin>=11 exceeds the 200k default
        "def submission : Goal := by\n"
        f"  refine ⟨Fin {n}, ⟨{op}⟩, ?_⟩\n"
        "  decideFin!\n"
    )


def find_counterexample(e1, e2, budget_s):
    """Returns a full Lean FALSE-certificate string, or None if no CE is found."""
    t0 = time.time()
    n, t = search_exhaustive(e1, e2, max_n=3)
    if n:
        return make_false_code(n, t)
    n, t = search_structured(e1, e2, max_n=7)
    if n:
        return make_false_code(n, t)
    # plain propagation backtracking is fast (<=8s) & complete at small n; it
    # owns Fin 4-5 and finds small arbitrary CEs the structured families miss.
    for size in (4, 5):
        t = search_backtrack(e1, e2, size, deadline=min(t0 + budget_s * 0.5, time.time() + 8))
        if t is not None:
            return make_false_code(size, t)
    # Affine Z_n sweep BEFORE the expensive Fin 6-7 targeted search. It is cheap
    # (~seconds to Fin 13) and its cost is machine-independent, so affine-only
    # CEs (hard2_0051 @ Fin13, hard2_0096 @ Fin9) are found in seconds rather
    # than after a ~100s fruitless targeted search — keeping the recovery ROBUST
    # on slow eval hardware instead of racing the deadline. Emitted via the
    # finOpTable-free direct-Fin cert, so it is sound at any order (finOpTable
    # mis-parses table entries >=10). min_n=8 because Fin<=7 (affine included)
    # is already covered by search_structured/backtrack above. Any CE it returns
    # is verified (holds ∧ violated), so it can only add solves, never flip a
    # verdict; running it before targeted (not after) makes 0051/0096 robust
    # while keeping small-n arbitrary CEs (backtrack) as the preferred cert.
    n, a, b, c = search_affine(e1, e2, min_n=8, max_n=13,
                               deadline=min(t0 + budget_s * 0.6, t0 + 15))
    if n:
        return make_false_affine_code(n, a, b, c)
    # targeted search (pin E2 sides) owns Fin 6 — it finds what plain backtrack
    # + all-zero completion misses. Expensive, so it runs last.
    for size in (6, 7):
        t = search_targeted(e1, e2, size, deadline=t0 + budget_s * 0.95)
        if t is not None:
            return make_false_code(size, t)
    return None


def main():
    startup = read_message()
    problem = startup["problem"]
    budget = startup.get("budget", {}) or {}
    budget_s = min(float(budget.get("timeout_seconds", 60)), 120)
    e1 = Equation(problem["equation1"])
    e2 = Equation(problem["equation2"])
    code = find_counterexample(e1, e2, budget_s)
    if code is not None:
        call_judge("false", code)
    # No true-implication prover in this FALSE engine; if no CE, exit (unsolved).


if __name__ == "__main__":
    main()