#!/usr/bin/env python3
import argparse
import lzma
import os
import sys
import time
from collections import defaultdict, deque


POSSIBLE_N = [840, 420, 280, 210, 168, 140, 120, 105, 84, 70, 42, 40, 35, 14, 5]


PATTERNS = {
    # Log/binary encoding for 6 states: 000, 001, 010, 011, 100, 101.
    "bin6": {
        "t": 6,
        "bits": 3,
        "patterns": [
            [[(0, 0), (1, 0), (2, 0)]],
            [[(0, 1), (1, 0), (2, 0)]],
            [[(0, 0), (1, 1), (2, 0)], [(0, 0), (1, 1)]],
            [[(0, 1), (1, 1), (2, 0)], [(0, 1), (1, 1)]],
            [[(0, 0), (2, 1)], [(0, 0), (1, 0), (2, 1)]],
            [[(0, 1), (2, 1)], [(0, 1), (1, 0), (2, 1)]],
        ],
        "codes": [0, 1, 2, 3, 4, 5],
    },
    # Log/binary encoding for 3 states: 00, 01, 10.
    "bin3": {
        "t": 3,
        "bits": 2,
        "patterns": [
            [[(0, 0), (1, 0)]],
            [[(0, 1)]],
            [[(0, 0), (1, 1)]],
        ],
        "codes": [0, 1, 2],
    },
    # Twisted-ring encoding for 6 states. The two invalid states are eliminated
    # by local invalid-state clauses, so each two-bit test names one state.
    "tw6": {
        "t": 6,
        "bits": 3,
        "patterns": [
            [[(0, 0), (2, 0)]],
            [[(0, 1), (1, 0)]],
            [[(1, 1), (2, 0)]],
            [[(0, 1), (2, 1)]],
            [[(0, 0), (1, 1)]],
            [[(1, 0), (2, 1)]],
        ],
        "codes": [0, 1, 3, 7, 6, 4],
    },
}


def read_dimacs(path):
    opener = lzma.open if path.endswith(".xz") else open
    nvars = 0
    clauses = []
    with opener(path, "rt") as fp:
        for line in fp:
            if not line or line[0] == "c":
                continue
            if line[0] == "p":
                parts = line.split()
                nvars = int(parts[2])
                continue
            lits = [int(x) for x in line.split()]
            if lits and lits[-1] == 0:
                lits.pop()
            if lits:
                clauses.append(lits)
    return nvars, clauses


class Encoding:
    def __init__(self, kind, n, t, type_bits=0, patterns=None, codes=None, score=0):
        self.kind = kind
        self.n = n
        self.t = t
        self.type_bits = type_bits
        self.patterns = patterns or []
        self.codes = codes or []
        self.score = score

    @property
    def type_start(self):
        return self.n + 1

    @property
    def type_end(self):
        if self.kind.startswith("direct"):
            return self.n + self.n * self.t
        return self.n + self.n * self.type_bits

    def __str__(self):
        return f"{self.kind}: N={self.n} T={self.t} bits={self.type_bits} score={self.score}"


def infer_direct(clauses, nvars, n, t):
    type_start = n + 1
    type_end = n + n * t
    if type_end >= nvars:
        return None
    trans = {}
    for lits in clauses:
        if len(lits) != 3:
            continue
        for src_lit in lits:
            if src_lit >= 0:
                continue
            src_var = -src_lit
            if not (type_start <= src_var <= type_end):
                continue
            src_node = (src_var - type_start) // t + 1
            src_state = (src_var - type_start) % t
            for call_lit in lits:
                if abs(call_lit) != src_node:
                    continue
                for dst_lit in lits:
                    if dst_lit == src_lit or dst_lit == call_lit:
                        continue
                    if type_start <= dst_lit <= type_end:
                        dst_node = (dst_lit - type_start) // t + 1
                        dst_state = (dst_lit - type_start) % t
                        call = 1 if call_lit < 0 else 0
                        trans[(src_node, src_state, call)] = (dst_node, dst_state)
    if not trans:
        return None
    return trans


def infer_pattern(clauses, nvars, n, spec):
    t = spec["t"]
    bits = spec["bits"]
    type_start = n + 1
    type_end = n + n * bits
    if type_end >= nvars:
        return None
    raw = defaultdict(lambda: defaultdict(dict))
    patterns = spec["patterns"]
    for lits in clauses:
        calls = [x for x in lits if abs(x) <= n]
        if len(calls) != 1:
            continue
        call_lit = calls[0]
        src_node = abs(call_lit)
        src_bits = [n + (src_node - 1) * bits + i + 1 for i in range(bits)]
        for state, alternatives in enumerate(patterns):
            for pat in alternatives:
                if len(lits) != len(pat) + 2:
                    continue
                pat_lits = [(-src_bits[b] if val else src_bits[b]) for b, val in pat]
                if not all(x in lits for x in pat_lits):
                    continue
                rem = [x for x in lits if x not in pat_lits and x != call_lit]
                if len(rem) != 1:
                    continue
                dst_lit = rem[0]
                if not (type_start <= abs(dst_lit) <= type_end):
                    continue
                dst_node = (abs(dst_lit) - type_start) // bits + 1
                dst_bit = (abs(dst_lit) - type_start) % bits
                call = 1 if call_lit < 0 else 0
                raw[(src_node, state, call)][dst_node][dst_bit] = dst_lit > 0

    if not raw:
        return None

    trans = {}
    for key, by_dst in raw.items():
        if len(by_dst) != 1:
            continue
        dst_node, bit_values = next(iter(by_dst.items()))
        cands = []
        for state, code in enumerate(spec["codes"]):
            ok = True
            for bit, val in bit_values.items():
                if bool((code >> bit) & 1) != val:
                    ok = False
                    break
            if ok:
                cands.append(state)
        if len(cands) == 1:
            trans[key] = (dst_node, cands[0])
    return trans


def detect_encoding(nvars, clauses):
    candidates = []
    for n in POSSIBLE_N:
        for t in (6, 3):
            tr = infer_direct(clauses, nvars, n, t)
            if tr:
                expected = n * t * 2
                score = len(tr) * 1000 - abs(expected - len(tr))
                candidates.append((score, Encoding(f"direct{t}", n, t, score=score), tr))
        for kind, spec in PATTERNS.items():
            tr = infer_pattern(clauses, nvars, n, spec)
            if tr:
                expected = n * spec["t"] * 2
                # Raw key count is the meaningful score; bin6 clauses can have several
                # destination bit clauses per transition but infer_pattern has grouped them.
                score = len(tr) * 1000 - abs(expected - len(tr))
                candidates.append((score, Encoding(kind, n, spec["t"], spec["bits"], spec["patterns"], spec["codes"], score), tr))
    candidates.sort(key=lambda x: x[0], reverse=True)
    if not candidates:
        return None, None
    return candidates[0][1], candidates[0][2]


def units_to_restrictions(clauses, enc):
    call_fixed = [None] * (enc.n + 1)
    allowed_states = [set(range(enc.t)) for _ in range(enc.n + 1)]
    unit_count = 0
    for lits in clauses:
        if len(lits) != 1:
            continue
        unit_count += 1
        lit = lits[0]
        v = abs(lit)
        val = lit > 0
        if 1 <= v <= enc.n:
            cur = 1 if val else 0
            if call_fixed[v] is not None and call_fixed[v] != cur:
                return None, None, unit_count
            call_fixed[v] = cur
        elif enc.type_start <= v <= enc.type_end:
            if enc.kind.startswith("direct"):
                idx = v - enc.type_start
                node = idx // enc.t + 1
                state = idx % enc.t
                if val:
                    allowed_states[node] &= {state}
                else:
                    allowed_states[node].discard(state)
            else:
                idx = v - enc.type_start
                node = idx // enc.type_bits + 1
                bit = idx % enc.type_bits
                keep = set()
                for state in allowed_states[node]:
                    code = enc.codes[state]
                    if bool((code >> bit) & 1) == val:
                        keep.add(state)
                allowed_states[node] = keep
    return call_fixed, allowed_states, unit_count


class Searcher:
    def __init__(self, enc, trans, call_fixed, allowed_states, time_limit):
        self.enc = enc
        self.trans = trans
        self.call_fixed = call_fixed
        self.allowed_states = allowed_states
        self.time_limit = time_limit
        self.start_time = time.time()
        self.nodes = 0
        self.path = []
        self.call_assign = [None] * (enc.n + 1)
        self.state_assign = [None] * (enc.n + 1)
        self.visited = [False] * (enc.n + 1)
        self.dead_memo = set()
        self.max_depth = 0

    def timed_out(self):
        return time.time() - self.start_time > self.time_limit

    def call_order(self, node, state):
        opts = []
        fixed = self.call_fixed[node]
        for call in (fixed,) if fixed is not None else (1, 0):
            nxt = self.trans.get((node, state, call))
            if nxt:
                opts.append((call, nxt))
        # Prefer constrained forward moves first. This makes fixed-unit
        # instances almost deterministic and still leaves the full search sound.
        def rank(item):
            _call, (dst, dst_state) = item
            return (0 if len(self.allowed_states[dst]) == 1 else 1, dst)
        opts.sort(key=rank)
        return opts

    def solve(self):
        start_states = sorted(self.allowed_states[1])
        # If a unit fixed the start state, this is already singleton.
        for s0 in start_states:
            self.path = []
            self.call_assign = [None] * (self.enc.n + 1)
            self.state_assign = [None] * (self.enc.n + 1)
            self.visited = [False] * (self.enc.n + 1)
            self.state_assign[1] = s0
            self.visited[1] = True
            self.path.append((1, s0))
            if self.dfs(1, s0, 1, s0):
                return True
            if self.timed_out():
                return False
        return False

    def dfs(self, node, state, depth, start_state):
        self.nodes += 1
        if depth > self.max_depth:
            self.max_depth = depth
        if self.nodes % 200000 == 0 and self.timed_out():
            return False
        # The tuple is coarse but catches many repeated traps in these
        # deterministic path searches.
        if depth > 32:
            key = (node, state, depth, tuple(i for i, v in enumerate(self.visited) if v and i % 64 == 0))
            if key in self.dead_memo:
                return False
        for call, (dst, dst_state) in self.call_order(node, state):
            if dst == 1:
                if depth == self.enc.n and dst_state == start_state:
                    self.call_assign[node] = call
                    return True
                continue
            if self.visited[dst]:
                continue
            if dst_state not in self.allowed_states[dst]:
                continue
            old_call = self.call_assign[node]
            old_state = self.state_assign[dst]
            self.call_assign[node] = call
            self.state_assign[dst] = dst_state
            self.visited[dst] = True
            self.path.append((dst, dst_state))
            if self.dfs(dst, dst_state, depth + 1, start_state):
                return True
            self.path.pop()
            self.visited[dst] = False
            self.state_assign[dst] = old_state
            self.call_assign[node] = old_call
            if self.timed_out():
                return False
        return False


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("cnf")
    ap.add_argument("--time", type=float, default=10.0)
    args = ap.parse_args()
    nvars, clauses = read_dimacs(args.cnf)
    enc, trans = detect_encoding(nvars, clauses)
    if not enc:
        print("UNKNOWN no encoding")
        return 1
    print("encoding", enc, "trans", len(trans), "expected", enc.n * enc.t * 2)
    call_fixed, allowed, unit_count = units_to_restrictions(clauses, enc)
    if call_fixed is None:
        print("UNSUPPORTED contradictory units")
        return 1
    print("units", unit_count, "fixed_calls", sum(x is not None for x in call_fixed), "fixed_states", sum(len(s) == 1 for s in allowed[1:]))
    s = Searcher(enc, trans, call_fixed, allowed, args.time)
    ok = s.solve()
    print("search", "SATPATH" if ok else "NONE", "nodes", s.nodes, "depth", len(s.path), "maxdepth", s.max_depth, "elapsed", f"{time.time()-s.start_time:.3f}")
    if ok:
        print("calls", "".join("B" if s.call_assign[i] else "P" for i in range(1, enc.n + 1)))
        print("path", " ".join(f"{node}:{state}" for node, state in s.path[:50]), "...")
    return 0 if ok else 2


if __name__ == "__main__":
    raise SystemExit(main())
