"""Exact SAT formulations for zero forcing in P(n,k).

(a) round_sat: the round-based encoding of Proposition P4 with horizon T = 2n - s.
(b) fort_cover: fort-cover formulation (Proposition P2) with lazy fort generation,
    fort minimisation, dihedral images of every fort (sound: automorphisms map
    forts to forts) and first-force normalisation (four directed-edge orbit
    representatives). UNSAT at bound b proves Z > b; the final static CNF is
    re-solved with proof logging and the DRAT proof is checked by drat-trim.
(c) boundary_sat: exists X with |X| = p and |dX| <= b?  (Krishnan-Ma-Zhang style
    boundary formula adapted to k = 5), with rotation normalisation x_{u0} v x_{v0}.

All CNFs are built with python-sat (sequential-counter cardinality encodings).
"""
import json
import os
import subprocess
import time

import common
from pysat.card import CardEnc, EncType
from pysat.formula import CNF
from pysat.solvers import Solver

from closure_oracle import closure, closure_fast
from graph_constructor import adjacency, adjacency_masks, dihedral_maps, label, u, v

DRAT_TRIM = os.path.join(common.ROOT, "code", "tools", "drat-trim", "drat-trim")


# ----------------------------------------------------------------------------
# helpers
# ----------------------------------------------------------------------------
class VarPool:
    def __init__(self):
        self.top = 0

    def new(self):
        self.top += 1
        return self.top


def atmost(lits, bound, pool):
    enc = CardEnc.atmost(lits=lits, bound=bound, top_id=pool.top, encoding=EncType.seqcounter)
    pool.top = max(pool.top, enc.nv)
    return enc.clauses


def equals(lits, bound, pool):
    enc = CardEnc.equals(lits=lits, bound=bound, top_id=pool.top, encoding=EncType.seqcounter)
    pool.top = max(pool.top, enc.nv)
    return enc.clauses


def is_fort(adjmask, F, N):
    """F: bitmask. Fort iff nonempty and every w outside F has 0 or >= 2 neighbours in F."""
    if F == 0:
        return False
    outside = ((1 << N) - 1) & ~F
    while outside:
        low = outside & -outside
        w = low.bit_length() - 1
        outside ^= low
        x = adjmask[w] & F
        if x and (x & (x - 1)) == 0:
            return False
    return True


def minimise_fort(adjmask, F, N):
    """Greedy: remove vertices while the set stays a fort (result is a fort, subset of F)."""
    changed = True
    while changed:
        changed = False
        rest = F
        while rest:
            low = rest & -rest
            rest ^= low
            G = F & ~low
            if G and is_fort(adjmask, G, N):
                F = G
                changed = True
    return F


def mask_to_list(m):
    out = []
    while m:
        low = m & -m
        out.append(low.bit_length() - 1)
        m ^= low
    return out


def solve_with_proof(clauses, nv, solver_name, cnf_path, proof_path, time_limit=None):
    """Solves a static CNF with proof logging, writes DIMACS + DRAT, runs drat-trim.
    Returns dict with sat, proof_checked, checker_exit, hashes, times."""
    cnf = CNF(from_clauses=clauses)
    cnf.nv = max(cnf.nv, nv)
    cnf.to_file(cnf_path)
    t0 = time.time()
    with Solver(name=solver_name, bootstrap_with=clauses, with_proof=True) as s:
        sat = s.solve()
        model = s.get_model() if sat else None
        proof = None if sat else s.get_proof()
    t_solve = time.time() - t0
    out = {"sat": sat, "solver": solver_name, "solve_seconds": round(t_solve, 3), "cnf_path": os.path.relpath(cnf_path, common.ROOT),
           "cnf_sha256": common.sha256_file(cnf_path), "num_vars": cnf.nv, "num_clauses": len(clauses)}
    if sat:
        out["model"] = model
        return out
    with open(proof_path, "w") as f:
        f.write("\n".join(proof) + "\n")
    out["proof_path"] = os.path.relpath(proof_path, common.ROOT)
    out["proof_sha256"] = common.sha256_file(proof_path)
    out["proof_lines"] = len(proof)
    t1 = time.time()
    cmd = [DRAT_TRIM, cnf_path, proof_path]
    r = subprocess.run(cmd, capture_output=True, text=True)
    out["checker_cmd"] = " ".join(os.path.relpath(c, common.ROOT) if c.startswith("/") else c for c in cmd)
    out["checker_exit"] = r.returncode
    out["checker_verified"] = ("s VERIFIED" in r.stdout)
    out["checker_seconds"] = round(time.time() - t1, 3)
    out["checker_stdout_tail"] = r.stdout.strip().splitlines()[-3:]
    return out


# ----------------------------------------------------------------------------
# (a) round-based encoding (P4)
# ----------------------------------------------------------------------------
def round_sat(n, k, s, solver_name="glucose4", T=None, time_limit=None):
    """Is there a zero-forcing set of size <= s?  Horizon T = 2n - s (P4).
    sat is None if the time limit was hit (unknown). Default solver Glucose 4.1
    (independent code base from the CaDiCaL used in the fort-cover / boundary paths)."""
    import threading
    adj = adjacency(n, k)
    N = 2 * n
    if T is None:
        T = 2 * n - s
    pool = VarPool()
    x = [pool.new() for _ in range(N)]
    b = [[pool.new() for _ in range(T + 1)] for _ in range(N)]
    f = {}
    for w in range(N):
        for bb in adj[w]:
            for t in range(1, T + 1):
                f[(bb, w, t)] = pool.new()
    cl = []
    for w in range(N):
        cl.append([-b[w][0], x[w]])
        cl.append([b[w][0], -x[w]])
        for t in range(1, T + 1):
            fs = [f[(bb, w, t)] for bb in adj[w]]
            cl.append([-b[w][t], b[w][t - 1]] + fs)
            cl.append([b[w][t], -b[w][t - 1]])
            for fv in fs:
                cl.append([b[w][t], -fv])
            for bb in adj[w]:
                fv = f[(bb, w, t)]
                cl.append([-fv, b[bb][t - 1]])
                for w2 in adj[bb]:
                    if w2 != w:
                        cl.append([-fv, b[w2][t - 1]])
        cl.append([b[w][T]])
    cl += atmost(x, s, pool)
    t0 = time.time()
    with Solver(name=solver_name, bootstrap_with=cl) as sv:
        if time_limit:
            timer = threading.Timer(time_limit, lambda: sv.interrupt())
            timer.start()
            sat = sv.solve_limited(expect_interrupt=True)
            timer.cancel()
        else:
            sat = sv.solve()
        model = sv.get_model() if sat else None
    S = None
    if sat:
        S = [w for w in range(N) if model[x[w] - 1] > 0]
    return {"sat": sat, "S": S, "num_vars": pool.top, "num_clauses": len(cl), "T": T,
            "solve_seconds": round(time.time() - t0, 3), "solver": solver_name, "timeout": sat is None}


# ----------------------------------------------------------------------------
# (b) fort-cover with lazy forts, symmetry and first-force normalisation
# ----------------------------------------------------------------------------
def first_force_representatives(n, k):
    """Four directed-edge orbit representatives under D_n: (source, target)."""
    return [(u(0, n), v(0, n)), (u(0, n), u(1, n)), (v(0, n), u(0, n)), (v(0, n), v(k, n))]


def fort_cover(n, k, b, solver_name="cadical195", time_limit=1200.0, normalise=True, symmetric=True,
               seed_forts=None, verbose=False):
    """Decides whether P(n,k) has a zero-forcing set of size <= b.
    Returns dict: sat (True/False/None on timeout), witness S (if sat), forts (list of vertex lists),
    base clauses, nv, iteration statistics."""
    adj = adjacency(n, k)
    adjm = adjacency_masks(n, k)
    N = 2 * n
    full = (1 << N) - 1
    pool = VarPool()
    svar = [pool.new() for _ in range(N)]  # svar[w] = w in S
    base = []
    base += atmost(svar, b, pool)
    if normalise:
        sel = []
        for (src, tgt) in first_force_representatives(n, k):
            c = pool.new()
            sel.append(c)
            base.append([-c, svar[src]])
            base.append([-c, -svar[tgt]])
            for w2 in adj[src]:
                if w2 != tgt:
                    base.append([-c, svar[w2]])
        base.append(sel)
    maps = dihedral_maps(n) if symmetric else [list(range(N))]
    fort_clauses = []
    forts = []
    seen = set()

    def add_fort(F):
        """F bitmask; adds clauses for F and its dihedral images."""
        added = 0
        for p in maps:
            img = 0
            for w in mask_to_list(F):
                img |= 1 << p[w]
            if img in seen:
                continue
            seen.add(img)
            forts.append(img)
            fort_clauses.append([svar[w] for w in mask_to_list(img)])
            added += 1
        return added

    if seed_forts:
        for F in seed_forts:
            add_fort(F)
    t0 = time.time()
    iters = 0
    result = None
    witness = None
    with Solver(name=solver_name, bootstrap_with=base + fort_clauses) as sv:
        while True:
            if time.time() - t0 > time_limit:
                result = None
                break
            iters += 1
            sat = sv.solve()
            if not sat:
                result = False
                break
            model = sv.get_model()
            S = 0
            for w in range(N):
                if model[svar[w] - 1] > 0:
                    S |= 1 << w
            B = closure_fast(adjm, S)
            if B == full:
                result = True
                witness = mask_to_list(S)
                break
            W = full & ~B
            assert is_fort(adjm, W, N)
            F = minimise_fort(adjm, W, N)
            assert is_fort(adjm, F, N) and (F & S) == 0
            before = len(fort_clauses)
            add_fort(F)
            for c in fort_clauses[before:]:
                sv.add_clause(c)
            if verbose and iters % 50 == 0:
                print("  iter %d forts %d elapsed %.1fs" % (iters, len(forts), time.time() - t0), flush=True)
    return {"sat": result, "witness": witness, "forts": [mask_to_list(F) for F in forts], "base_clauses": base,
            "fort_clauses": fort_clauses, "nv": pool.top, "iterations": iters, "wall_seconds": round(time.time() - t0, 3),
            "svar": svar, "solver": solver_name, "normalised": normalise, "symmetric": symmetric}


# ----------------------------------------------------------------------------
# (c) boundary formulas
# ----------------------------------------------------------------------------
def boundary_clauses(n, k, p, bmax, normalise=True):
    """CNF: exists X, |X| = p, |dX| <= bmax (z_w >= [w in dX]). Returns (clauses, x, z, nv)."""
    adj = adjacency(n, k)
    N = 2 * n
    pool = VarPool()
    x = [pool.new() for _ in range(N)]
    z = [pool.new() for _ in range(N)]
    cl = []
    for w in range(N):
        for t in adj[w]:
            cl.append([x[w], -x[t], z[w]])  # (not x_w and x_t) -> z_w
        cl.append([-z[w], -x[w]])  # z_w -> not x_w (so z counts only boundary vertices)
    cl += equals(x, p, pool)
    cl += atmost(z, bmax, pool)
    if normalise:
        cl.append([x[u(0, n)], x[v(0, n)]])
    return cl, x, z, pool.top


def boundary_of(n, k, X):
    adj = adjacency(n, k)
    Xs = set(X)
    return sorted({w for x in X for w in adj[x] if w not in Xs})


def boundary_sat(n, k, p, bmax, solver_name="cadical195"):
    cl, x, z, nv = boundary_clauses(n, k, p, bmax)
    t0 = time.time()
    with Solver(name=solver_name, bootstrap_with=cl) as sv:
        sat = sv.solve()
        model = sv.get_model() if sat else None
    X = None
    if sat:
        X = [w for w in range(2 * n) if model[x[w] - 1] > 0]
        assert len(X) == p and len(boundary_of(n, k, X)) <= bmax
    return {"sat": sat, "X": X, "boundary": (boundary_of(n, k, X) if X else None), "num_vars": nv,
            "num_clauses": len(cl), "solve_seconds": round(time.time() - t0, 3)}


def beta(n, k, p, start=11, solver_name="cadical195"):
    """Exact beta_p(n) = min |dX| over |X| = p, if <= start; else returns start+1 meaning '>= start+1'."""
    best = None
    bX = None
    bnd = start
    while True:
        r = boundary_sat(n, k, p, bnd, solver_name)
        if not r["sat"]:
            break
        best = len(r["boundary"])
        bX = r["X"]
        bnd = best - 1
        if bnd < 0:
            break
    return {"n": n, "k": k, "p": p, "beta": best if best is not None else start + 1, "exact": best is not None,
            "X": bX, "boundary": (boundary_of(n, k, bX) if bX else None)}


if __name__ == "__main__":
    # smoke tests on known cases
    for (n, k, expect) in [(5, 2, 5), (11, 5, 6)]:
        r1 = round_sat(n, k, expect)
        r0 = round_sat(n, k, expect - 1)
        f1 = fort_cover(n, k, expect)
        f0 = fort_cover(n, k, expect - 1)
        print("P(%d,%d): round_sat(%d)=%s round_sat(%d)=%s fort(%d)=%s fort(%d)=%s iters=%d/%d" % (
            n, k, expect, r1["sat"], expect - 1, r0["sat"], expect, f1["sat"], expect - 1, f0["sat"], f1["iterations"], f0["iterations"]))
        assert r1["sat"] and not r0["sat"] and f1["sat"] and not f0["sat"]
    print("boundary P(19,4) p=16 b=9:", boundary_sat(19, 4, 16, 9)["sat"], " (Krishnan-Ma-Zhang example: expected True)")
    print("boundary P(20,4) p=16 b=9:", boundary_sat(20, 4, 16, 9)["sat"], " (expected False)")
    print("SAT smoke tests passed")