#!/usr/bin/env python3
import collections
import glob
import importlib.machinery
import lzma
import os
import sys
import time


def parse(path):
    opener = lzma.open if path.endswith(".xz") else open
    nv = 0
    clauses = []
    with opener(path, "rt") as f:
        for line in f:
            if not line or line[0] in "c%":
                continue
            if line.startswith("p "):
                nv = int(line.split()[2])
            else:
                clauses.append(tuple(int(x) for x in line.split()[:-1]))
    return nv, clauses


def detect_n(clauses):
    pairs = set(c for c in clauses if len(c) == 2)
    best = 0
    for cand in range(1, 300):
        ok = True
        for i in range(1, cand + 1):
            if (-(cand + i), i) not in pairs or (cand + i, -i) not in pairs:
                ok = False
                break
        if ok:
            best = cand
    return best


def unit_prop(nv, clauses, assigns):
    val = [0] * (nv + 1)
    def assign(lit):
        v = abs(lit)
        s = 1 if lit > 0 else -1
        if val[v] == -s:
            return False
        if val[v] == 0:
            val[v] = s
        return True

    for lit in assigns:
        if not assign(lit):
            return None
    changed = True
    while changed:
        changed = False
        for clause in clauses:
            sat = False
            last = 0
            un = 0
            for lit in clause:
                vv = val[abs(lit)]
                if vv and ((lit > 0 and vv > 0) or (lit < 0 and vv < 0)):
                    sat = True
                    break
                if vv == 0:
                    un += 1
                    last = lit
            if sat:
                continue
            if un == 0:
                return None
            if un == 1:
                if not assign(last):
                    return None
                changed = True
    return val


def recover(nv, clauses):
    n = detect_n(clauses)
    negidx = [[] for _ in range(nv + 1)]
    for ci, c in enumerate(clauses):
        for lit in c:
            if lit < 0:
                negidx[-lit].append(ci)

    memo = {}

    def fanin(var, lo, hi):
        if lo <= var <= hi:
            return {var - lo + 1}
        if var in memo:
            return memo[var]
        acc = set()
        for ci in negidx[var]:
            c = clauses[ci]
            ok = True
            others = []
            for lit in c:
                if lit == -var:
                    continue
                if lit <= 0 or lit >= var:
                    ok = False
                    break
                others.append(lit)
            if ok and others:
                for lit in others:
                    acc |= fanin(lit, lo, hi)
        memo[var] = acc
        return acc

    def local_rank(var, lo, hi, max_k=None):
        support = sorted(fanin(var, lo, hi))
        if not support:
            return None, support
        if max_k is None:
            max_k = len(support)
        local = [c for c in clauses if len(c) > 1 and max(abs(x) for x in c) <= var]
        for k in range(0, min(max_k, len(support)) + 1):
            assigns = []
            for idx, p in enumerate(support):
                assigns.append((lo + p - 1) if idx < k else -(lo + p - 1))
            val = unit_prop(nv, local, assigns)
            if val and val[var] > 0:
                return k, support
        return None, support

    preds = []
    theta = []
    for j in range(1, n + 1):
        a = 2 * n + j
        seed = j
        tr = None
        for ci in negidx[a]:
            c = clauses[ci]
            if seed in c:
                tr = [x for x in c if x not in (-a, seed) and x > 0]
                break
        if tr is None:
            raise ValueError(f"missing trigger for {j}")
        ps = set()
        for x in tr:
            ps |= fanin(x, n + 1, 2 * n)
        ps = sorted(ps)
        if ps:
            rank, _ = local_rank(max(tr), n + 1, 2 * n, len(ps))
            if rank is None:
                rank = 1
        else:
            rank = 10 ** 9
        preds.append([p - 1 for p in ps])
        theta.append(rank)

    units = [c[0] for c in clauses if len(c) == 1]
    neg_units = [-u for u in units if u < 0]
    pos_units = [u for u in units if u > 0]
    kr, _ = local_rank(neg_units[0], 1, n, min(n, 100))
    target = pos_units[0] - (nv - n)
    return n, preds, theta, kr - 1, target


def closure(preds, theta, seeds):
    n = len(preds)
    active = seeds
    changed = True
    while changed:
        changed = False
        for v in range(n):
            if (active >> v) & 1:
                continue
            cnt = 0
            for p in preds[v]:
                if (active >> p) & 1:
                    cnt += 1
            if cnt >= theta[v]:
                active |= 1 << v
                changed = True
    return active


if __name__ == "__main__":
    paths = []
    for arg in sys.argv[1:] or ["/data/train/*.cnf.xz"]:
        paths += glob.glob(arg)
    for path in sorted(paths):
        t0 = time.time()
        nv, clauses = parse(path)
        n, preds, theta, k, target = recover(nv, clauses)
        print(
            os.path.basename(path),
            "vars", nv,
            "N", n,
            "K", k,
            "target", target,
            "theta", dict(sorted(collections.Counter(theta).items())),
            "indeg", dict(sorted(collections.Counter(len(p) for p in preds).items())),
            "time", round(time.time() - t0, 3),
        )
