#!/usr/bin/env python3
import collections
import itertools
import lzma
import sys
import time

MOVES = [(1, 2), (2, 1), (-1, 2), (-2, 1), (1, -2), (2, -1), (-1, -2), (-2, -1)]


def load_tours(path="/workspace/tmp/tours.txt"):
    tours = {}
    cur = None
    with open(path) as f:
        for line in f:
            if line.startswith("n "):
                cur = int(line.split()[1])
            elif line.strip():
                tours[cur] = [int(x) for x in line.split()]
    return tours


def read_cnf(path):
    clauses = []
    nvar = 0
    with lzma.open(path, "rt") if path.endswith(".xz") else open(path) as f:
        for line in f:
            if line.startswith("p "):
                nvar = int(line.split()[2])
            elif line.startswith(("c", "%")):
                continue
            else:
                lits = tuple(int(x) for x in line.split() if x != "0")
                if lits:
                    clauses.append(lits)
    return nvar, clauses


def extract_sparse_support(clauses):
    binset = {tuple(sorted(c)) for c in clauses if len(c) == 2}
    groups, seen = [], {}
    for c in clauses:
        if 2 <= len(c) <= 8:
            ok = True
            for a, b in itertools.combinations(c, 2):
                if tuple(sorted((-a, -b))) not in binset:
                    ok = False
                    break
            if ok:
                key = tuple(sorted(abs(x) for x in c))
                if key not in seen:
                    seen[key] = len(groups)
                    groups.append(list(c))

    occ = collections.defaultdict(list)
    sign = {}
    for gi, g in enumerate(groups):
        for lit in g:
            occ[abs(lit)].append(gi)
            sign[(gi, abs(lit))] = 1 if lit > 0 else -1

    adj = [set() for _ in groups]
    var_between = {}
    for v, gs in occ.items():
        if len(gs) == 2:
            a, b = gs
            adj[a].add(b)
            adj[b].add(a)
            var_between[tuple(sorted((a, b)))] = v

    color = [-1] * len(groups)
    for s in range(len(groups)):
        if color[s] >= 0:
            continue
        color[s] = 0
        stack = [s]
        while stack:
            u = stack.pop()
            for w in adj[u]:
                if color[w] < 0:
                    color[w] = color[u] ^ 1
                    stack.append(w)
                elif color[w] == color[u]:
                    raise RuntimeError("support graph is not bipartite")
    return groups, [list(a) for a in adj], color, var_between, sign


def canonical_biadj(n):
    v = n * n
    adj = [[] for _ in range(2 * v)]
    for r in range(n):
        for c in range(n):
            u = r * n + c
            for dr, dc in MOVES:
                rr, cc = r + dr, c + dc
                if 0 <= rr < n and 0 <= cc < n:
                    w = rr * n + cc
                    adj[u].append(v + w)
                    adj[v + w].append(u)
    return adj


def wl_union(cadj, tadj, tside):
    n1, n2 = len(cadj), len(tadj)
    adj, side = [], []
    for i, a in enumerate(cadj):
        adj.append(a)
        side.append(0 if i < n1 // 2 else 1)
    for i, a in enumerate(tadj):
        adj.append([n1 + x for x in a])
        side.append(tside[i])
    color = [(side[i], len(adj[i])) for i in range(n1 + n2)]
    mp = {}
    color = [mp.setdefault(x, len(mp)) for x in color]
    for _ in range(20):
        keys = [(color[i], tuple(sorted(color[j] for j in adj[i]))) for i in range(n1 + n2)]
        mp = {}
        new = [mp.setdefault(x, len(mp)) for x in keys]
        if new == color:
            break
        color = new
    return color[:n1], color[n1:]


def iso_biadj(n, tadj, tcolor):
    v = n * n
    cadj = canonical_biadj(n)
    ccol, tcol = wl_union(cadj, tadj, tcolor)
    tadj_set = [set(a) for a in tadj]
    cadj_set = [set(x - v for x in cadj[u]) for u in range(v)]
    tside = [[i for i, c in enumerate(tcolor) if c == s] for s in (0, 1)]
    dom = [
        [[t for t in tside[s] if tcol[t] == ccol[(0 if s == 0 else v) + u]] for u in range(v)]
        for s in (0, 1)
    ]

    maps = [[-1] * v, [-1] * v]
    used = [set(), set()]
    sys.setrecursionlimit(10000)

    def feasible(s, u, t):
        if s == 0:
            for w, tw in enumerate(maps[1]):
                if tw != -1 and ((tw in tadj_set[t]) != (w in cadj_set[u])):
                    return False
        else:
            for w, tw in enumerate(maps[0]):
                if tw != -1 and ((tw in tadj_set[t]) != (u in cadj_set[w])):
                    return False
        return True

    def rec():
        best = None
        bc = None
        for s in (0, 1):
            for u in range(v):
                if maps[s][u] != -1:
                    continue
                cand = [t for t in dom[s][u] if t not in used[s] and feasible(s, u, t)]
                if best is None or len(cand) < len(bc):
                    best = (s, u)
                    bc = cand
                if len(cand) <= 1:
                    break
            if best is not None and len(bc) <= 1:
                break
        if best is None:
            return True
        if not bc:
            return False
        s, u = best
        for t in bc:
            maps[s][u] = t
            used[s].add(t)
            if rec():
                return True
            used[s].remove(t)
            maps[s][u] = -1
        return False

    if not rec():
        raise RuntimeError("isomorphism failed")
    return maps


def orient_groups(groups, adj, sign):
    flip = [0] * len(groups)
    for root in range(len(groups)):
        if flip[root] != 0:
            continue
        flip[root] = 1
        stack = [root]
        while stack:
            g = stack.pop()
            gvars = {abs(x) for x in groups[g]}
            for h in adj[g]:
                shared = gvars & {abs(x) for x in groups[h]}
                v = next(iter(shared))
                need = flip[g] * sign[(g, v)] * sign[(h, v)]
                if flip[h] == 0:
                    flip[h] = need
                    stack.append(h)
                elif flip[h] != need:
                    raise RuntimeError("orientation conflict")
    return flip


def check(clauses, assign):
    bad = undec = 0
    residual = collections.Counter()
    for c in clauses:
        sat = False
        un = 0
        last_signs = []
        for lit in c:
            val = assign[abs(lit)]
            if val == 0:
                un += 1
                last_signs.append(1 if lit > 0 else -1)
            elif val == (1 if lit > 0 else -1):
                sat = True
                break
        if not sat:
            if un:
                undec += 1
                residual[(un, sum(1 for s in last_signs if s > 0), sum(1 for s in last_signs if s < 0))] += 1
            else:
                bad += 1
    return bad, undec, residual


def unit_prop(clauses, assign):
    changed = True
    while changed:
        changed = False
        for c in clauses:
            sat = False
            un = 0
            last = 0
            for lit in c:
                val = assign[abs(lit)]
                if val == 0:
                    un += 1
                    last = lit
                elif val == (1 if lit > 0 else -1):
                    sat = True
                    break
            if sat:
                continue
            if un == 0:
                return False
            if un == 1:
                val = 1 if last > 0 else -1
                v = abs(last)
                if assign[v] == 0:
                    assign[v] = val
                    changed = True
                elif assign[v] != val:
                    return False
    return True


def build_sparse_base(path, n):
    nvar, clauses = read_cnf(path)
    groups, adj, color, var_between, sign = extract_sparse_support(clauses)
    maps = iso_biadj(n, adj, color)
    flip = orient_groups(groups, adj, sign)
    assign = [0] * (nvar + 1)
    for (a, b), var in var_between.items():
        sem = flip[a] * sign[(a, var)]
        assign[var] = -sem
    tour = load_tours()[n]
    v = n * n
    for i, u in enumerate(tour):
        w = tour[(i + 1) % v]
        ga, gb = maps[0][u], maps[1][w]
        var = var_between[tuple(sorted((ga, gb)))]
        sem = flip[ga] * sign[(ga, var)]
        assign[var] = sem
    return nvar, clauses, assign


if __name__ == "__main__":
    path = sys.argv[1]
    n = int(sys.argv[2])
    t = time.time()
    nvar, clauses, assign = build_sparse_base(path, n)
    print("base", time.time() - t, check(clauses, assign)[:2])
    print("up", unit_prop(clauses, assign), check(clauses, assign)[:2])
    bad, undec, hist = check(clauses, assign)
    print("residual hist", hist.most_common(30))
