#!/usr/bin/env python3
import sys
from collections import defaultdict, deque


def parse_dimacs(path):
    clauses = []
    n = m = None
    with open(path) as f:
        for line in f:
            if not line or line[0] == "c":
                continue
            if line[0] == "p":
                parts = line.split()
                n, m = int(parts[2]), int(parts[3])
                continue
            lits = [int(x) for x in line.split() if x != "0"]
            if lits:
                clauses.append(lits)
    return n, clauses


def litstr(l):
    return ("x%d" % l) if l > 0 else ("~x%d" % (-l))


def terms(lits):
    return " ".join("+1 " + litstr(l) for l in lits)


class Detector:
    def __init__(self, n, clauses):
        self.n = n
        self.clauses = clauses
        self.m = len(clauses)
        self.ids = {tuple(c): i + 1 for i, c in enumerate(clauses)}

    def fail(self, msg):
        raise RuntimeError(msg)

    def detect(self):
        n, clauses = self.n, self.clauses
        in4 = [False] * (n + 1)
        for c in clauses:
            if len(c) == 4:
                for l in c:
                    in4[abs(l)] = True
        pvars = {v for v in range(1, n + 1) if not in4[v]}
        pclause_ids = []
        pclause_lits = []
        for idx, c in enumerate(clauses, 1):
            if len(c) > 2 and all(abs(l) in pvars for l in c):
                pclause_ids.append(idx)
                pclause_lits.append(c)
        if not pclause_lits:
            self.fail("no P clauses")
        p = len(pclause_lits)
        r = len(pclause_lits[0])
        if any(len(c) != r for c in pclause_lits):
            self.fail("P clauses have mixed lengths")
        if len(pvars) != p * r or n != 2 * p * r:
            self.fail(f"bad counts p={p} r={r} pvars={len(pvars)} n={n}")
        canonP = {}
        pgroup = {}
        group_pvars = []
        for g, c in enumerate(pclause_lits):
            arr = []
            for l in c:
                v = abs(l)
                if v in canonP:
                    self.fail("P var in two P clauses")
                canonP[v] = l
                pgroup[v] = g
                arr.append(v)
            group_pvars.append(arr)

        p_adj = defaultdict(list)
        p_edge_id = {}
        p_imp = {}
        rlit_for_p = {}
        for idx, c in enumerate(clauses, 1):
            if len(c) != 2:
                continue
            a, b = c
            ap, bp = abs(a) in pvars, abs(b) in pvars
            if ap and bp:
                if a != -canonP[abs(a)] or b != -canonP[abs(b)]:
                    self.fail("bad P-P signs")
                va, vb = abs(a), abs(b)
                p_adj[va].append(vb)
                p_adj[vb].append(va)
                p_edge_id[tuple(sorted((va, vb)))] = idx
            elif ap ^ bp:
                pl, rl = (a, b) if ap else (b, a)
                pv = abs(pl)
                if pl != -canonP[pv]:
                    self.fail("bad P implication sign")
                if pv in p_imp:
                    self.fail("duplicate P implication")
                p_imp[pv] = idx
                rlit_for_p[pv] = rl
            else:
                self.fail("binary over non-P variables")

        seen = set()
        domains = []
        var_domain = {}
        rlit_to_domain = {}
        for v in sorted(pvars):
            if v in seen:
                continue
            comp = []
            q = deque([v])
            seen.add(v)
            while q:
                x = q.popleft()
                comp.append(x)
                for y in p_adj[x]:
                    if y not in seen:
                        seen.add(y)
                        q.append(y)
            if len(comp) != p:
                self.fail(f"P component size {len(comp)} != p={p}")
            if len({pgroup[x] for x in comp}) != p:
                self.fail("P component does not hit all groups")
            rlits = {rlit_for_p[x] for x in comp}
            if len(rlits) != 1:
                self.fail("P component has multiple R literals")
            rlit = next(iter(rlits))
            d = len(domains)
            domains.append({"pvars": comp, "rlit": rlit, "q_lits": [], "r_clause": None})
            for x in comp:
                var_domain[x] = d
            if abs(rlit) in rlit_to_domain and rlit_to_domain[abs(rlit)] != d:
                self.fail("R variable reused")
            rlit_to_domain[abs(rlit)] = d
        if len(domains) != r:
            self.fail("bad domain count")
        rvars = {abs(d["rlit"]) for d in domains}
        qvars = set(range(1, n + 1)) - pvars - rvars
        if len(qvars) != r * (p - 1):
            self.fail("bad Q count")

        q_domain = {}
        q_canon = {}
        for idx, c in enumerate(clauses, 1):
            if len(c) != p:
                continue
            if any(abs(l) in pvars for l in c):
                continue
            r_lits = [l for l in c if abs(l) in rvars]
            if len(r_lits) != 1:
                continue
            rd = rlit_to_domain[abs(r_lits[0])]
            if r_lits[0] != -domains[rd]["rlit"]:
                continue
            qs = [l for l in c if abs(l) in qvars]
            if len(qs) != p - 1:
                continue
            if domains[rd]["q_lits"]:
                self.fail("duplicate R->Q clause")
            domains[rd]["q_lits"] = qs
            domains[rd]["r_clause"] = idx
            for ql in qs:
                qv = abs(ql)
                if qv in q_domain:
                    self.fail("Q reused in R clauses")
                q_domain[qv] = rd
                q_canon[qv] = ql
        if any(not d["q_lits"] for d in domains):
            self.fail("missing R->Q clause")
        if len(q_domain) != len(qvars):
            self.fail("not all Q vars assigned")

        q_adj = defaultdict(list)
        q_edge_id = {}
        for idx, c in enumerate(clauses, 1):
            if len(c) != 4:
                continue
            if any(abs(l) in pvars for l in c):
                continue
            rs = [l for l in c if abs(l) in rvars]
            qs = [l for l in c if abs(l) in qvars]
            if len(rs) != 2 or len(qs) != 2:
                continue
            d1 = rlit_to_domain[abs(rs[0])]
            d2 = rlit_to_domain[abs(rs[1])]
            if d1 == d2:
                self.fail("injection clause same R domain")
            if rs[0] != -domains[d1]["rlit"] or rs[1] != -domains[d2]["rlit"]:
                self.fail("bad R sign in injection")
            qv1, qv2 = abs(qs[0]), abs(qs[1])
            if qs[0] != -q_canon[qv1] or qs[1] != -q_canon[qv2]:
                self.fail("bad Q sign in injection")
            if {q_domain[qv1], q_domain[qv2]} != {d1, d2}:
                self.fail("Q domains do not match R domains")
            q_adj[qv1].append(qv2)
            q_adj[qv2].append(qv1)
            q_edge_id[tuple(sorted((qv1, qv2)))] = idx

        seen = set()
        color_comps = []
        q_color = {}
        for qv in sorted(qvars):
            if qv in seen:
                continue
            comp = []
            dq = deque([qv])
            seen.add(qv)
            while dq:
                x = dq.popleft()
                comp.append(x)
                for y in q_adj[x]:
                    if y not in seen:
                        seen.add(y)
                        dq.append(y)
            color_comps.append(comp)
        if len(color_comps) != p - 1:
            self.fail(f"bad color count {len(color_comps)} != {p-1}")
        q_by_domain_color = [[None] * (p - 1) for _ in range(r)]
        for color, comp in enumerate(color_comps):
            if len(comp) != r:
                self.fail(f"color size {len(comp)} != r={r}")
            ds = set()
            for qv in comp:
                d = q_domain[qv]
                if d in ds:
                    self.fail("color repeats domain")
                ds.add(d)
                q_color[qv] = color
                q_by_domain_color[d][color] = qv
        missing = []
        present_inj = {}
        for j in range(r):
            for k in range(j + 1, r):
                for c in range(p - 1):
                    a = q_by_domain_color[j][c]
                    b = q_by_domain_color[k][c]
                    eid = q_edge_id.get(tuple(sorted((a, b))))
                    if eid is None:
                        missing.append((j, k, c))
                    else:
                        present_inj[(j, k, c)] = eid
        self.p = p
        self.r = r
        self.pvars = pvars
        self.canonP = canonP
        self.pgroup = pgroup
        self.domains = domains
        self.group_pvars = group_pvars
        self.pclause_ids = pclause_ids
        self.p_edge_id = p_edge_id
        self.p_imp = p_imp
        self.rvars = rvars
        self.qvars = qvars
        self.q_canon = q_canon
        self.q_by_domain_color = q_by_domain_color
        self.q_edge_id = q_edge_id
        self.present_inj = present_inj
        self.missing = missing
        return self


class Proof:
    def __init__(self, det):
        self.d = det
        self.lines = ["pseudo-Boolean proof version 3.0", "", f"f {det.m} ;"]
        self.maxid = det.m
        self.xlit = {}

    def add_line(self, s):
        self.lines.append(s)
        self.maxid += 1
        return self.maxid

    def red(self, lits, witness_var, witness_val):
        cid = self.add_line(f"red {terms(lits)} >= 1 : x{witness_var} -> {witness_val} ;")
        return cid

    def rup_clause(self, lits):
        return self.add_line(f"rup {terms(lits)} >= 1 ;")

    def rup_pb(self, lits, rhs):
        return self.add_line(f"rup {terms(lits)} >= {rhs} ;")

    def pol(self, ids):
        if not ids:
            raise RuntimeError("empty pol")
        expr = str(ids[0])
        for x in ids[1:]:
            expr += f" {x} +"
        return self.add_line(f"pol {expr} ;")

    def pbc_clique(self, atoms, edge_id):
        # atoms are canonical literals that represent true for each item.
        if len(atoms) == 1:
            return None
        if len(atoms) == 2:
            return edge_id[tuple(sorted((abs(atoms[0]), abs(atoms[1]))))]

        groups = []
        i = 0
        while i + 1 < len(atoms):
            a, b = atoms[i], atoms[i + 1]
            eid = edge_id[tuple(sorted((abs(a), abs(b))))]
            groups.append({"atoms": [a, b], "amo": eid})
            i += 2
        if i < len(atoms):
            groups.append({"atoms": [atoms[i]], "amo": None})

        while len(groups) > 1:
            groups.sort(key=lambda g: len(g["atoms"]))
            A = groups.pop(0)
            B = groups.pop(0)
            merged_atoms = A["atoms"] + B["atoms"]
            amo = self.merge_cliques(A, B, edge_id)
            groups.append({"atoms": merged_atoms, "amo": amo})
        return groups[0]["amo"]

    def merge_cliques(self, A, B, edge_id):
        # Ensure singleton special case has A singleton if any.
        if len(B["atoms"]) == 1 and len(A["atoms"]) > 1:
            A, B = B, A
        target_lits = [-x for x in (A["atoms"] + B["atoms"])]
        self.lines.append(f"pbc {terms(target_lits)} >= {len(target_lits) - 1} : subproof")
        old = self.maxid
        neg_id = old + 1
        local = old + 1

        def sub_pol(tokens):
            nonlocal local
            self.lines.append("pol " + " ".join(str(t) for t in tokens) + " ;")
            local += 1
            return local

        if len(A["atoms"]) == 1:
            a = A["atoms"][0]
            if B["amo"] is None:
                raise RuntimeError("bad singleton merge")
            apos = sub_pol([neg_id, B["amo"], "+"])
            unit_ids = []
            for b in B["atoms"]:
                eid = edge_id[tuple(sorted((abs(a), abs(b))))]
                unit_ids.append(sub_pol([apos, eid, "+"]))
            toks = [neg_id]
            for uid in unit_ids:
                toks += [uid, "+"]
            contr = sub_pol(toks)
        else:
            # Derive units for the smaller side.
            if len(B["atoms"]) < len(A["atoms"]):
                A, B = B, A
            bpos = sub_pol([neg_id, A["amo"], "+"])
            apos = sub_pol([neg_id, B["amo"], "+"])
            unit_ids = []
            for a in A["atoms"]:
                toks = [bpos]
                for b in B["atoms"]:
                    eid = edge_id[tuple(sorted((abs(a), abs(b))))]
                    toks += [eid, "+"]
                toks += [len(B["atoms"]), "d"]
                unit_ids.append(sub_pol(toks))
            toks = [apos]
            for uid in unit_ids:
                toks += [uid, "+"]
            contr = sub_pol(toks)
        self.lines.append(f"qed pbc : {contr} ;")
        self.maxid = local + 1
        return self.maxid

    def prove_domain_link(self, dom_idx, p_amo_id):
        d = self.d.domains[dom_idx]
        pvars = d["pvars"]
        lits = [-self.d.canonP[v] for v in pvars] + [d["rlit"]]
        self.lines.append(f"pbc {terms(lits)} >= {self.d.p} : subproof")
        old = self.maxid
        neg = old + 1
        local = old + 1

        def sub_pol(tokens):
            nonlocal local
            self.lines.append("pol " + " ".join(str(t) for t in tokens) + " ;")
            local += 1
            return local

        nr = sub_pol([neg, p_amo_id, "+"])
        unit_ids = []
        for v in pvars:
            unit_ids.append(sub_pol([nr, self.d.p_imp[v], "+"]))
        toks = [neg]
        for uid in unit_ids:
            toks += [uid, "+"]
        contr = sub_pol(toks)
        self.lines.append(f"qed pbc : {contr} ;")
        self.maxid = local + 1
        return self.maxid

    def finish_unsat(self):
        d = self.d
        # First-level domain cardinalities.
        domain_c_ids = []
        for j, dom in enumerate(d.domains):
            atoms = [d.canonP[v] for v in dom["pvars"]]
            eids = {}
            for a in atoms:
                for b in atoms:
                    if abs(a) < abs(b):
                        eids[(abs(a), abs(b))] = d.p_edge_id[(abs(a), abs(b))]
            amo = self.pbc_clique(atoms, eids)
            domain_c_ids.append(self.prove_domain_link(j, amo))
        lower = self.pol(d.pclause_ids + domain_c_ids)

        # Extension variables X[j,c] = R_j & Q_jc.
        xvar = {}
        nextv = d.n + 1
        for j in range(d.r):
            for c in range(d.p - 1):
                xv = nextv
                nextv += 1
                xvar[(j, c)] = xv
                R = d.domains[j]["rlit"]
                Q = d.q_canon[d.q_by_domain_color[j][c]]
                self.red([-xv, R], xv, 0)
                self.red([-xv, Q], xv, 0)
                self.red([-R, -Q, xv], xv, 1)

        # Pairwise at-most-one over X for each color.
        x_edge = {}
        for c in range(d.p - 1):
            for j in range(d.r):
                for k in range(j + 1, d.r):
                    eid = self.rup_clause([-xvar[(j, c)], -xvar[(k, c)]])
                    x_edge[(xvar[(j, c)], xvar[(k, c)])] = eid

        # If R_j is selected, one X_j,c is true.
        t_ids = []
        for j in range(d.r):
            lits = [-d.domains[j]["rlit"]] + [xvar[(j, c)] for c in range(d.p - 1)]
            t_ids.append(self.rup_clause(lits))

        color_amo = []
        for c in range(d.p - 1):
            atoms = [xvar[(j, c)] for j in range(d.r)]
            eids = {}
            for j in range(d.r):
                for k in range(j + 1, d.r):
                    a = xvar[(j, c)]
                    b = xvar[(k, c)]
                    eids[(min(a, b), max(a, b))] = x_edge[(min(a, b), max(a, b))]
            color_amo.append(self.pbc_clique(atoms, eids))
        upper = self.pol(t_ids + color_amo)
        final = self.pol([lower, upper])
        self.lines += ["", "output NONE ;", f"conclusion UNSAT : {final} ;", "end pseudo-Boolean proof ;"]
        return "\n".join(self.lines) + "\n"


def main():
    n, clauses = parse_dimacs(sys.argv[1])
    det = Detector(n, clauses).detect()
    print("detected", det.p, det.r, "missing", len(det.missing), file=sys.stderr)
    if det.missing:
        return 2
    proof = Proof(det).finish_unsat()
    open(sys.argv[2], "w").write(proof)
    return 0


if __name__ == "__main__":
    sys.exit(main())
