#!/usr/bin/env python3
import lzma
import os
import re
import sys
from functools import lru_cache

from extract_hypergraph import extract, detect_n, detect_m_and_ranges


def min_cover(bag, edges_by_v, edges, limit):
    bag = set(bag)
    if not bag:
        return []
    best = None

    def rec(uncovered, chosen):
        nonlocal best
        if not uncovered:
            if best is None or len(chosen) < len(best):
                best = chosen[:]
            return True
        if len(chosen) >= limit or (best is not None and len(chosen) >= len(best)):
            return False
        v = min(uncovered, key=lambda x: sum(1 for e in edges_by_v[x] if edges[e] & uncovered))
        cand = sorted(edges_by_v[v], key=lambda e: -len(edges[e] & uncovered))
        any_ok = False
        for e in cand:
            cov = edges[e] & uncovered
            if not cov:
                continue
            chosen.append(e)
            if rec(uncovered - cov, chosen):
                any_ok = True
            chosen.pop()
        return any_ok

    rec(frozenset(bag), [])
    return best


def greedy_order(n, edges, W):
    edges_by_v = [[] for _ in range(n)]
    primal = [0] * n
    for ei, e in enumerate(edges):
        vs = list(e)
        for v in vs:
            edges_by_v[v].append(ei)
        for i, u in enumerate(vs):
            for v in vs[i+1:]:
                primal[u] |= 1 << v
                primal[v] |= 1 << u
    rem = (1 << n) - 1
    order = []
    covers = [None] * n
    max_cover = 0
    while rem:
        best = None
        candidates = []
        # Try all remaining vertices. n is small enough for this profiler.
        x = rem
        while x:
            lb = x & -x
            v = lb.bit_length() - 1
            nbrs = primal[v] & rem
            bag = {v}
            y = nbrs
            while y:
                b = y & -y
                bag.add(b.bit_length() - 1)
                y -= b
            cov = min_cover(bag, edges_by_v, edges, W)
            if cov is not None:
                score = (len(cov), nbrs.bit_count(), sum(len(edges[e]) for e in cov))
                if best is None or score < best[0]:
                    best = (score, v, nbrs, cov)
            x -= lb
        if best is None:
            return None, None, None
        _, v, nbrs, cov = best
        covers[v] = cov
        max_cover = max(max_cover, len(cov))
        order.append(v)
        # Fill in clique on remaining neighbors.
        y = nbrs
        while y:
            b = y & -y
            u = b.bit_length() - 1
            primal[u] |= (nbrs & ~(1 << u))
            y -= b
        rem &= ~(1 << v)
    return order, covers, max_cover


def order_var(n, i, j):
    if i == j:
        raise ValueError
    if i < j:
        # row i has variables for i+1..n-1.
        return i * (2 * n - i - 1) // 2 + (j - i)
    return order_var(n, j, i)


def a_var(n, order_count, i, j):
    return order_count + 1 + i * n + j


def p_var(n, p_start, i, j):
    if i == j:
        return None
    c = j if j < i else j - 1
    return p_start + i * (n - 1) + c


def build_semantic(n, m, edges, W, order, covers, w_start, p_start):
    order_count = n * (n - 1) // 2
    a_start = order_count + 1
    w_end = w_start + n * m - 1
    p_end = p_start + n * (n - 1) - 1
    max_sem = p_end
    val = bytearray(max_sem + 1)  # 0 false, 1 true
    pos = [0] * n
    for t, v in enumerate(order):
        pos[v] = t
    # Order vars.
    for i in range(n):
        for j in range(i + 1, n):
            val[order_var(n, i, j)] = 1 if pos[i] < pos[j] else 0
    # Recompute chordal completion and covers' bags.
    primal = [0] * n
    for e in edges:
        vs = list(e)
        for a, u in enumerate(vs):
            for v in vs[a+1:]:
                primal[u] |= 1 << v
                primal[v] |= 1 << u
    rem = (1 << n) - 1
    for v in order:
        nbrs = primal[v] & rem & ~(1 << v)
        y = nbrs
        while y:
            b = y & -y
            u = b.bit_length() - 1
            val[a_var(n, order_count, v, u)] = 1
            y -= b
        y = nbrs
        while y:
            b = y & -y
            u = b.bit_length() - 1
            primal[u] |= (nbrs & ~(1 << u))
            y -= b
        rem &= ~(1 << v)
    # w covers
    for i, cov in enumerate(covers):
        for e in cov:
            val[w_start + i * m + e] = 1
    # p variables: exactly vertices that occur in selected hyperedges of each bag row.
    unions = [set() for _ in range(n)]
    for j, cov in enumerate(covers):
        for e in cov:
            unions[j].update(edges[e])
    for i in range(n):
        for j in range(n):
            pv = p_var(n, p_start, i, j)
            if pv is not None:
                val[pv] = 1 if i in unions[j] else 0
    return val, max_sem


def check_semantic(path, val, max_sem, max_report=10):
    bad = 0
    total = 0
    examples = []
    with lzma.open(path, "rt", encoding="ascii", errors="replace") as f:
        for line in f:
            if line.startswith("c") or line.startswith("p"):
                continue
            cl = [int(x) for x in line.split()]
            if cl and cl[-1] == 0:
                cl.pop()
            if any(abs(x) > max_sem for x in cl):
                continue
            total += 1
            ok = False
            for lit in cl:
                b = val[abs(lit)] != 0
                if (lit > 0 and b) or (lit < 0 and not b):
                    ok = True
                    break
            if not ok:
                bad += 1
                if len(examples) < max_report:
                    examples.append(cl)
    return total, bad, examples


def main():
    for path in sys.argv[1:]:
        n, m, w_start, p_start, edges_list, primal = extract(path)
        edges = [frozenset(e) for e in edges_list]
        mo = re.search(r"\.hg_(\d+)\.cnf", os.path.basename(path))
        W = int(mo.group(1)) if mo else 4
        order, covers, mx = greedy_order(n, edges, W)
        print(os.path.basename(path), "n", n, "m", m, "W", W, "order", order is not None, "maxcov", mx)
        if order:
            val, max_sem = build_semantic(n, m, edges, W, order, covers, w_start, p_start)
            total, bad, examples = check_semantic(path, val, max_sem)
            print(" semantic clauses", total, "bad", bad, "examples", examples[:3])


if __name__ == "__main__":
    main()
