#!/usr/bin/env python3
"""
Erdos 835 small-case settlement -- ROBUST, self-contained.

For each k we decide whether J(2k,k) is (k+1)-colourable, which is exactly
whether 835(k) is YES.  Three independent oracles must agree:

  (O1) OWN exhaustive backtracking prover (MRV + forward checking, sound
       symmetry break fixing one (k+1)-clique to the identity colouring).
       For UNSAT this enumerates the ENTIRE search tree and reports 0
       completions -- a proof that trusts no external solver.
  (O2) CaDiCaL (via python-sat).
  (O3) Glucose4 and MiniSat22 (via python-sat).

For SAT we independently re-verify the returned colouring is proper AND makes
every (k+1)-subset rainbow.  DIMACS CNF is written for third-party replay.
"""
import sys, itertools, json, time
from math import comb
from pysat.formula import CNF
from pysat.solvers import Cadical153, Glucose4, Minisat22

# ----------------------------- Johnson graph -----------------------------
def build(n, k):
    V = list(itertools.combinations(range(n), k))
    idx = {s: i for i, s in enumerate(V)}
    S = [frozenset(s) for s in V]
    adj = [set() for _ in V]
    for i in range(len(V)):
        for j in range(i+1, len(V)):
            if len(S[i] & S[j]) == k-1:
                adj[i].add(j); adj[j].add(i)
    cliques = []
    for A in itertools.combinations(range(n), k+1):
        fac = tuple(idx[tuple(sorted(set(A) - {x}))] for x in A)
        cliques.append(fac)
    return V, idx, S, adj, cliques

# --------------------- O1: own exhaustive prover -------------------------
def own_prover(nV, adj, ncol, fixed):
    """
    fixed: dict vertex->colour (the symmetry-broken clique).
    Returns (colourable: bool, witness: list|None, nodes: int).
    Sound & complete: MRV order, forward checking, full backtracking.
    """
    colour = [None]*nV
    domain = [set(range(ncol)) for _ in range(nV)]
    nodes = 0

    # apply fixed assignments with propagation
    def assign(v, c, undo):
        colour[v] = c
        for w in adj[v]:
            if colour[w] is None and c in domain[w]:
                domain[w].discard(c); undo.append(w)
        return True
    def unassign(v, undo):
        colour[v] = None
        c = None
        for w in undo:
            pass
    # We manage domains via explicit stacks per assignment.
    order_fixed = list(fixed.items())
    def propagate_fixed(i):
        nonlocal nodes
        if i == len(order_fixed):
            return backtrack()
        v, c = order_fixed[i]
        if c not in domain[v]:
            return False
        removed = []
        colour[v] = c
        for w in adj[v]:
            if colour[w] is None and c in domain[w]:
                domain[w].discard(c); removed.append(w)
        ok = propagate_fixed(i+1)
        if ok: return True
        colour[v] = None
        for w in removed: domain[w].add(c)
        return False

    def backtrack():
        nonlocal nodes
        nodes += 1
        # pick unassigned vertex with min remaining values (MRV)
        best = -1; bestsize = ncol+1
        for v in range(nV):
            if colour[v] is None:
                d = len(domain[v])
                if d == 0:
                    return False          # dead end
                if d < bestsize:
                    bestsize = d; best = v
                    if d == 1: break
        if best == -1:
            return True                   # complete colouring found
        v = best
        for c in list(domain[v]):
            removed = []
            colour[v] = c
            dead = False
            for w in adj[v]:
                if colour[w] is None and c in domain[w]:
                    domain[w].discard(c); removed.append(w)
                    if not domain[w]:
                        dead = True
            if not dead and backtrack():
                return True
            colour[v] = None
            for w in removed: domain[w].add(c)
        return False

    ok = propagate_fixed(0)
    witness = list(colour) if ok else None
    return ok, witness, nodes

# ----------------------------- CNF encoding ------------------------------
def var(v, c, ncol): return v*ncol + c + 1
def make_cnf(nV, adj, ncol, fixed):
    cnf = CNF()
    for v in range(nV):
        cnf.append([var(v, c, ncol) for c in range(ncol)])
        for c1 in range(ncol):
            for c2 in range(c1+1, ncol):
                cnf.append([-var(v, c1, ncol), -var(v, c2, ncol)])
    seen = set()
    for v in range(nV):
        for w in adj[v]:
            e = (v, w) if v < w else (w, v)
            if e in seen: continue
            seen.add(e)
            for c in range(ncol):
                cnf.append([-var(v, c, ncol), -var(w, c, ncol)])
    for v, c in fixed.items():
        cnf.append([var(v, c, ncol)])
    return cnf

def sat_solve(SolverCls, cnf):
    s = SolverCls(bootstrap_with=cnf.clauses)
    r = s.solve()
    model = set(l for l in (s.get_model() or []) if l > 0) if r else None
    s.delete()
    return r, model

# ------------------------------- driver ----------------------------------
def settle(k, do_sat=True):
    n, ncol = 2*k, k+1
    V, idx, S, adj, cliques = build(n, k)
    nV = len(V)
    # symmetry break: fix first anti-star clique to identity colouring
    fixed = {v: pos for pos, v in enumerate(cliques[0])}
    assert len(fixed) == ncol

    t0 = time.time()
    ok1, wit, nodes = own_prover(nV, adj, ncol, fixed)
    t1 = time.time()

    res = {"k": k, "n": n, "colours": ncol, "num_vertices": nV,
           "num_edges": sum(len(a) for a in adj)//2,
           "num_anti_star_cliques": len(cliques),
           "own_prover_colourable": ok1, "own_prover_nodes": nodes,
           "own_prover_seconds": round(t1-t0, 3)}

    # independent verification of a witness
    if ok1:
        proper = all(wit[v] != wit[w] for v in range(nV) for w in adj[v])
        rainbow = all(len({wit[f] for f in fac}) == ncol for fac in cliques)
        res["verified_proper"] = proper
        res["verified_all_rainbow"] = rainbow
        with open(f"witness_k{k}.json","w") as f:
            json.dump({"colour_by_vertex_index": wit,
                       "vertices":[list(s) for s in V]}, f)

    # SAT cross-checks
    if do_sat:
        cnf = make_cnf(nV, adj, ncol, fixed)
        cnf.to_file(f"J_{n}_{k}__{ncol}col.cnf")
        res["cnf_clauses"] = len(cnf.clauses); res["cnf_vars"] = nV*ncol
        agree = {}
        for name, Cls in [("cadical153",Cadical153),("glucose4",Glucose4),("minisat22",Minisat22)]:
            r, _ = sat_solve(Cls, cnf)
            agree[name] = bool(r)
        res["sat_solvers_colourable"] = agree
        res["all_oracles_agree"] = (len(set([ok1]+list(agree.values()))) == 1)

    res["conclusion"] = ("835(%d) is YES" % k) if ok1 else ("835(%d) is NO" % k)
    print(json.dumps(res, indent=1))
    print()
    return res

if __name__ == "__main__":
    ks = [int(x) for x in sys.argv[1:]] or [2,3,4]
    out = [settle(k) for k in ks]
    with open("settle_835_results.json","w") as f:
        json.dump(out, f, indent=1)
    print("Wrote settle_835_results.json")
