#!/usr/bin/env python3
import collections
import os
import re
import sys

sys.path.append(os.path.dirname(__file__))
from extract_hypergraph import extract
from try_order import a_var, check_semantic, order_var, p_var


def minfill_td(adj0):
    n = len(adj0)
    adj = [set(s) for s in adj0]
    rem = set(range(n))
    elim = []
    bags = []
    later = []
    while rem:
        best = None
        for v in rem:
            nb = adj[v] & rem
            l = list(nb)
            fill = 0
            for i, u in enumerate(l):
                au = adj[u]
                for w in l[i + 1 :]:
                    if w not in au:
                        fill += 1
            score = (fill, len(nb), v)
            if best is None or score < best[0]:
                best = (score, v, nb)
        _, v, nb = best
        elim.append(v)
        later.append(set(nb))
        bags.append(set(nb) | {v})
        l = list(nb)
        for i, u in enumerate(l):
            for w in l[i + 1 :]:
                adj[u].add(w)
                adj[w].add(u)
        rem.remove(v)

    pos = {v: i for i, v in enumerate(elim)}
    bag_of = {v: i for i, v in enumerate(elim)}
    parent = [-1] * n
    for idx, nb in enumerate(later):
        if nb:
            u = min(nb, key=lambda x: pos[x])
            parent[idx] = bag_of[u]
    children = [[] for _ in range(n)]
    roots = []
    for i, p in enumerate(parent):
        if p >= 0:
            children[p].append(i)
        else:
            roots.append(i)
    return elim, bags, parent, children, roots


def compute_arcs(n, edges, order):
    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
    arcs = [[False] * n for _ in range(n)]
    for v in order:
        nbrs = primal[v] & rem & ~(1 << v)
        y = nbrs
        while y:
            b = y & -y
            u = b.bit_length() - 1
            arcs[v][u] = True
            y -= b
        y = nbrs
        while y:
            b = y & -y
            u = b.bit_length() - 1
            primal[u] |= nbrs & ~(1 << u)
            y -= b
        rem &= ~(1 << v)
    return arcs


def p_closure(n, edges, covers, arcs, pos):
    p = [[False] * n for _ in range(n)]
    for i in range(n):
        for j in range(n):
            if i != j and pos[i] < pos[j]:
                p[j][i] = True
    for row, cov in enumerate(covers):
        for e in cov:
            for v in edges[e]:
                if v != row:
                    p[v][row] = True
    changed = True
    while changed:
        changed = False
        for j in range(n):
            for k in range(n):
                if not arcs[j][k]:
                    continue
                for i in range(n):
                    if i != j and i != k and p[i][k] and not p[i][j]:
                        p[i][j] = True
                        changed = True
    return p


def build_val(n, m, order, covers, ws, ps, arcs, p, pos):
    oc = n * (n - 1) // 2
    max_sem = ps + n * (n - 1) - 1
    val = bytearray(max_sem + 1)
    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
    for i in range(n):
        for j in range(n):
            if arcs[i][j]:
                val[a_var(n, oc, i, j)] = 1
    for i, cov in enumerate(covers):
        for e in cov:
            val[ws + i * m + e] = 1
    for i in range(n):
        for j in range(n):
            if i != j and p[i][j]:
                val[p_var(n, ps, i, j)] = 1
    return val, max_sem


def order_from_td(nh, m, inc, bags, children, roots, root_choice=None):
    incident = [set() for _ in range(m)]
    for e, (a, b) in enumerate(inc):
        incident[a].add(e)
        incident[b].add(e)
    chi = []
    for bag in bags:
        s = set()
        for x in bag:
            s |= incident[x]
        chi.append(s)

    root = roots[0] if root_choice is None else root_choice
    depth = [-1] * len(bags)
    depth[root] = 0
    st = [root]
    while st:
        x = st.pop()
        for c in children[x]:
            depth[c] = depth[x] + 1
            st.append(c)
    if any(d < 0 for d in depth):
        return None

    f = [None] * nh
    for e in range(nh):
        cand = [i for i, c in enumerate(chi) if e in c]
        f[e] = min(cand, key=lambda i: depth[i])
    order = sorted(range(nh), key=lambda e: (-depth[f[e]], f[e], e))
    covers = [set(bags[f[e]]) for e in range(nh)]
    return order, covers, depth, f, chi


def main():
    for path in sys.argv[1:]:
        n, m, ws, ps, edges_list, _ = extract(path)
        edges = [frozenset(x) for x in edges_list]
        inc = [[] for _ in range(n)]
        for h, vs in enumerate(edges):
            for v in vs:
                inc[v].append(h)
        if any(len(x) != 2 for x in inc):
            print(os.path.basename(path), "not-2-regular")
            continue
        adj = [set() for _ in range(m)]
        for a, b in inc:
            adj[a].add(b)
            adj[b].add(a)
        _, bags, parent, children, roots = minfill_td(adj)
        data = order_from_td(n, m, inc, bags, children, roots)
        if data is None:
            print(os.path.basename(path), "disconnected-td")
            continue
        order, covers, depth, f, chi = data
        pos = [0] * n
        for t, v in enumerate(order):
            pos[v] = t
        arcs = compute_arcs(n, edges, order)
        p = p_closure(n, edges, covers, arcs, pos)
        val, max_sem = build_val(n, m, order, covers, ws, ps, arcs, p, pos)
        total, bad, examples = check_semantic(path, val, max_sem)
        W = re.search(r"hg_(\d+)", os.path.basename(path))
        print(
            os.path.basename(path),
            "W",
            W.group(1) if W else "?",
            "td_maxbag",
            max(map(len, bags)),
            "cover_max",
            max(map(len, covers)),
            "bad",
            bad,
            "examples",
            examples[:3],
            "f_hist",
            collections.Counter(f).most_common(3),
        )


if __name__ == "__main__":
    main()
