#!/usr/bin/env python3
import collections
import glob
import json
import lzma
import os
import sys
import time


def parse(path):
    opener = lzma.open if path.endswith(".xz") else open
    clauses = []
    nvars = 0
    sizes = collections.Counter()
    with opener(path, "rt") as fh:
        for line in fh:
            if not line or line[0] == "c":
                continue
            if line.startswith("p "):
                sp = line.split()
                nvars = int(sp[2])
                continue
            lits = [int(x) for x in line.split()]
            if not lits:
                continue
            if lits[-1] == 0:
                lits.pop()
            if not lits:
                continue
            clauses.append(lits)
            sizes[len(lits)] += 1
    return nvars, clauses, sizes


def up(nvars, clauses):
    val = [0] * (nvars + 1)
    watches = [[] for _ in range(2 * nvars + 2)]
    watchpos = []
    trail = []
    head = 0

    def idx(lit):
        v = abs(lit)
        return 2 * v + (1 if lit < 0 else 0)

    def litval(lit):
        x = val[abs(lit)]
        if x == 0:
            return 0
        return 1 if (lit > 0) == (x > 0) else -1

    def enqueue(lit):
        v = abs(lit)
        want = 1 if lit > 0 else -1
        if val[v] == want:
            return True
        if val[v] == -want:
            return False
        val[v] = want
        trail.append(lit)
        return True

    for i, c in enumerate(clauses):
        if len(c) == 1:
            watchpos.append((0, 0))
            watches[idx(c[0])].append(i)
            if not enqueue(c[0]):
                return True, len(trail), i
        else:
            watchpos.append((0, 1))
            watches[idx(c[0])].append(i)
            watches[idx(c[1])].append(i)

    while head < len(trail):
        lit = trail[head]
        head += 1
        wi = idx(-lit)
        j = 0
        ws = watches[wi]
        while j < len(ws):
            ci = ws[j]
            c = clauses[ci]
            a, b = watchpos[ci]
            if c[a] != -lit:
                a, b = b, a
            other = c[b]
            if litval(other) == 1:
                j += 1
                continue
            moved = False
            for k, x in enumerate(c):
                if k == a or k == b:
                    continue
                if litval(x) != -1:
                    watchpos[ci] = (k, b)
                    ws[j] = ws[-1]
                    ws.pop()
                    watches[idx(x)].append(ci)
                    moved = True
                    break
            if moved:
                continue
            if litval(other) == -1:
                return True, len(trail), ci
            if not enqueue(other):
                return True, len(trail), ci
            j += 1
    return False, len(trail), -1


def main():
    files = sys.argv[1:] or sorted(glob.glob("/data/train/*.cnf.xz"))
    base = {}
    bp = "/workspace/baselines/kissat-sc2024.json"
    if os.path.exists(bp):
        for r in json.load(open(bp))["results"]:
            base[r["name"]] = r
    for path in files:
        t = time.time()
        nvars, clauses, sizes = parse(path)
        confl, assigned, ci = up(nvars, clauses)
        name = os.path.basename(path)
        st = base.get(name, {}).get("status")
        print(
            f"{'UP' if confl else 'NOUP'} assigned={assigned} n={nvars} m={len(clauses)} "
            f"sizes={dict(sorted(sizes.items()))} status={st} sec={time.time()-t:.2f} {name}",
            flush=True,
        )


if __name__ == "__main__":
    main()
