#!/usr/bin/env python3
import heapq
import sys
import time

from highlevel import parse_dimacs, recover_high_clauses


class CDCL:
    def __init__(self, clauses, nvars):
        self.nvars = nvars
        self.clauses = [list(c) for c in clauses]
        self.watches = {}
        self.watch_pos = []
        self.assign = [0] * (nvars + 1)
        self.level = [0] * (nvars + 1)
        self.reason = [-1] * (nvars + 1)
        self.trail = []
        self.trail_lim = []
        self.qhead = 0
        self.activity = [0.0] * (nvars + 1)
        self.var_inc = 1.0
        self.heap = []
        self.learned = []
        self.conflicts = 0
        for ci, c in enumerate(self.clauses):
            for lit in c:
                self.activity[abs(lit)] += 1.0
            if len(c) == 1:
                self.watch_pos.append([0, 0])
                self.watches.setdefault(c[0], []).append(ci)
            else:
                self.watch_pos.append([0, 1])
                self.watches.setdefault(c[0], []).append(ci)
                self.watches.setdefault(c[1], []).append(ci)
        for v in range(1, nvars + 1):
            heapq.heappush(self.heap, (-self.activity[v], v))

    def lit_value(self, lit):
        val = self.assign[abs(lit)]
        if val == 0:
            return 0
        return val if lit > 0 else -val

    def enqueue(self, lit, reason):
        v = abs(lit)
        val = 1 if lit > 0 else -1
        if self.assign[v] != 0:
            return self.assign[v] == val
        self.assign[v] = val
        self.level[v] = len(self.trail_lim)
        self.reason[v] = reason
        self.trail.append(lit)
        return True

    def add_clause(self, clause):
        # Remove duplicate literals; tautologies are ignored by returning None.
        seen = {}
        out = []
        for lit in clause:
            v = abs(lit)
            if v in seen:
                if seen[v] != lit:
                    return None
            else:
                seen[v] = lit
                out.append(lit)
        ci = len(self.clauses)
        self.clauses.append(out)
        if len(out) == 0:
            self.watch_pos.append([0, 0])
        elif len(out) == 1:
            self.watch_pos.append([0, 0])
            self.watches.setdefault(out[0], []).append(ci)
        else:
            self.watch_pos.append([0, 1])
            self.watches.setdefault(out[0], []).append(ci)
            self.watches.setdefault(out[1], []).append(ci)
        return ci

    def propagate(self):
        while self.qhead < len(self.trail):
            lit = self.trail[self.qhead]
            self.qhead += 1
            false_lit = -lit
            ws = self.watches.get(false_lit, [])
            i = 0
            while i < len(ws):
                ci = ws[i]
                c = self.clauses[ci]
                if not c:
                    return ci
                wp = self.watch_pos[ci]
                if len(c) == 1:
                    if self.lit_value(c[0]) < 0:
                        return ci
                    i += 1
                    continue
                if c[wp[0]] == false_lit:
                    wi, oi = 0, 1
                elif c[wp[1]] == false_lit:
                    wi, oi = 1, 0
                else:
                    # Stale watch entry left after a move.
                    ws[i] = ws[-1]
                    ws.pop()
                    continue
                other = c[wp[oi]]
                if self.lit_value(other) > 0:
                    i += 1
                    continue
                moved = False
                for k, l2 in enumerate(c):
                    if k == wp[oi] or k == wp[wi]:
                        continue
                    if self.lit_value(l2) >= 0:
                        wp[wi] = k
                        self.watches.setdefault(l2, []).append(ci)
                        ws[i] = ws[-1]
                        ws.pop()
                        moved = True
                        break
                if moved:
                    continue
                if self.lit_value(other) < 0:
                    return ci
                if not self.enqueue(other, ci):
                    return ci
                i += 1
        return None

    def current_level(self):
        return len(self.trail_lim)

    def backtrack(self, lvl):
        if self.current_level() <= lvl:
            return
        keep = self.trail_lim[lvl] if lvl > 0 else 0
        for lit in reversed(self.trail[keep:]):
            v = abs(lit)
            self.assign[v] = 0
            self.reason[v] = -1
            self.level[v] = 0
        del self.trail[keep:]
        del self.trail_lim[lvl:]
        self.qhead = min(self.qhead, len(self.trail))

    def analyze(self, confl):
        self.conflicts += 1
        self.var_inc *= 1.05
        seen = [False] * (self.nvars + 1)
        learnt = []
        path_c = 0
        clause = self.clauses[confl]
        idx = len(self.trail) - 1
        p = 0
        while True:
            for q in clause:
                v = abs(q)
                if not seen[v] and self.level[v] > 0:
                    seen[v] = True
                    self.activity[v] += self.var_inc
                    if self.level[v] == self.current_level():
                        path_c += 1
                    else:
                        learnt.append(q)
                        heapq.heappush(self.heap, (-self.activity[v], v))
            while True:
                p = self.trail[idx]
                idx -= 1
                if seen[abs(p)]:
                    break
            seen[abs(p)] = False
            path_c -= 1
            if path_c <= 0:
                break
            clause = self.clauses[self.reason[abs(p)]]
        learnt.insert(0, -p)
        # Put the highest-level non-asserting literal in position 1.
        back = 0
        best_i = 1
        for i in range(1, len(learnt)):
            lv = self.level[abs(learnt[i])]
            if lv > back:
                back = lv
                best_i = i
        if len(learnt) > 1:
            learnt[1], learnt[best_i] = learnt[best_i], learnt[1]
        return learnt, back

    def pick_branch(self):
        while self.heap:
            negact, v = heapq.heappop(self.heap)
            if self.assign[v] == 0:
                return v
        for v in range(1, self.nvars + 1):
            if self.assign[v] == 0:
                return v
        return 0

    def solve(self, max_conflicts=1000000):
        confl = self.propagate()
        if confl is not None:
            return False
        while True:
            confl = self.propagate()
            if confl is not None:
                if self.current_level() == 0:
                    self.learned.append([])
                    return False
                learnt, back = self.analyze(confl)
                self.learned.append(learnt)
                self.backtrack(back)
                ci = self.add_clause(learnt)
                if ci is not None:
                    self.enqueue(learnt[0], ci)
                if self.conflicts >= max_conflicts:
                    return None
            else:
                v = self.pick_branch()
                if v == 0:
                    return True
                self.trail_lim.append(len(self.trail))
                self.enqueue(v, -1)


def main():
    for path in sys.argv[1:]:
        nvars, clauses = parse_dimacs(path)
        high = recover_high_clauses(clauses)
        high_clauses = [list(k) for k, _ in high]
        t0 = time.time()
        solver = CDCL(high_clauses, nvars // 2)
        res = solver.solve()
        print(path, "res", res, "conflicts", solver.conflicts, "learned", len(solver.learned), "time", time.time() - t0)
        if res is None:
            break


if __name__ == "__main__":
    main()
