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

sys.path.append(os.path.dirname(__file__))
from extract_hypergraph import extract
from test_td_order import build_val, compute_arcs, p_closure
from try_order import check_semantic


class Node:
    __slots__ = ("lam", "chi", "children")

    def __init__(self, lam, chi, children):
        self.lam = lam
        self.chi = chi
        self.children = children


def bits(mask):
    while mask:
        b = mask & -mask
        yield b.bit_length() - 1
        mask -= b


def popcount(mask):
    return mask.bit_count()


def mask_from(xs):
    m = 0
    for x in xs:
        m |= 1 << x
    return m


class DetK:
    def __init__(self, n_h, m, inc, W):
        self.n_h = n_h
        self.m = m
        self.inc = inc
        self.W = W
        self.adj = [0] * m
        self.inc_edges = [0] * m
        for e, (a, b) in enumerate(inc):
            self.adj[a] |= 1 << b
            self.adj[b] |= 1 << a
            self.inc_edges[a] |= 1 << e
            self.inc_edges[b] |= 1 << e
        self.fail = set()
        self.calls = 0

    def union_edges(self, vmask):
        out = 0
        for v in bits(vmask):
            out |= self.inc_edges[v]
        return out

    def components(self, edges_mask, sep_mask):
        rem = edges_mask & ~sep_mask
        comps = []
        while rem:
            start_bit = rem & -rem
            q = [start_bit.bit_length() - 1]
            seen = 0
            rem &= ~start_bit
            while q:
                v = q.pop()
                seen |= 1 << v
                nb = self.adj[v] & rem
                while nb:
                    b = nb & -nb
                    u = b.bit_length() - 1
                    rem &= ~b
                    q.append(u)
                    nb -= b
            comps.append(seen)
        return comps

    def conn_between(self, comp_mask, sep_mask):
        out = 0
        for e, (a, b) in enumerate(self.inc):
            ina = (comp_mask >> a) & 1
            inb = (comp_mask >> b) & 1
            if ina ^ inb:
                if ((sep_mask >> a) & 1) or ((sep_mask >> b) & 1):
                    out |= 1 << e
        return out

    def cover_candidates(self, conn_edges, current_edges, limit=200):
        if conn_edges == 0:
            return [0]
        cand_vertices = 0
        for e in bits(conn_edges):
            a, b = self.inc[e]
            cand_vertices |= (1 << a) | (1 << b)
        cover_by = {}
        for v in bits(cand_vertices):
            cover_by[v] = self.inc_edges[v] & conn_edges
        # Forced choices for degree-one boundary edges reduce branching.
        results = []

        def rec(uncovered, chosen, cand_mask):
            if len(results) >= limit:
                return
            if uncovered == 0:
                results.append(chosen)
                return
            if popcount(chosen) >= self.W:
                return
            # Lower bound: greedy maximal matching on uncovered boundary graph.
            used_v = 0
            lb = 0
            for e in bits(uncovered):
                a, b = self.inc[e]
                if ((used_v >> a) & 1) or ((used_v >> b) & 1):
                    continue
                used_v |= (1 << a) | (1 << b)
                lb += 1
            if popcount(chosen) + lb > self.W:
                return
            # Branch on the least-covered remaining edge.
            best_e = None
            best_opts = None
            for e in bits(uncovered):
                a, b = self.inc[e]
                opts = [v for v in (a, b) if (cand_mask >> v) & 1]
                if best_opts is None or len(opts) < len(best_opts):
                    best_e = e
                    best_opts = opts
            opts = sorted(
                best_opts,
                key=lambda v: (
                    -popcount(cover_by[v] & uncovered),
                    0 if ((current_edges >> v) & 1) else 1,
                    v,
                ),
            )
            for v in opts:
                rec(uncovered & ~cover_by[v], chosen | (1 << v), cand_mask & ~(1 << v))

        rec(conn_edges, 0, cand_vertices)
        # Prefer covers that use current component vertices and split well.
        results = sorted(
            set(results),
            key=lambda s: (
                popcount(s),
                -popcount(s & current_edges),
                s,
            ),
        )
        return results[:limit]

    def decomp(self, edges_mask, conn_edges, depth=0):
        self.calls += 1
        key = (edges_mask, conn_edges)
        if key in self.fail:
            return None
        if popcount(edges_mask) <= self.W:
            lam = edges_mask
            chi = self.union_edges(edges_mask)
            return Node(lam, chi, [])

        covs = self.cover_candidates(conn_edges, edges_mask)
        add_options_cache = {}
        for cov in covs:
            in_cov = cov & edges_mask
            add_options = [0]
            if in_cov == 0:
                if popcount(cov) >= self.W:
                    continue
                opts = []
                for a in bits(edges_mask):
                    sep = cov | (1 << a)
                    comps = self.components(edges_mask, sep)
                    maxc = max((popcount(c) for c in comps), default=0)
                    opts.append((maxc, -len(comps), -popcount(self.adj[a] & edges_mask), a))
                opts.sort()
                add_options = [1 << a for _, _, _, a in opts[:80]]
            for add in add_options:
                sep = cov | add
                in_sep = sep & edges_mask
                if in_sep == 0:
                    continue
                comps = self.components(edges_mask, sep)
                if len(comps) == 1 and comps[0] == edges_mask:
                    continue
                comps.sort(key=lambda c: popcount(c), reverse=True)
                child_nodes = []
                ok = True
                for comp in comps:
                    child_conn = self.conn_between(comp, sep)
                    child = self.decomp(comp, child_conn, depth + 1)
                    if child is None:
                        ok = False
                        break
                    child_nodes.append(child)
                if ok:
                    chi = conn_edges | self.union_edges(in_sep)
                    return Node(sep, chi, child_nodes)
        self.fail.add(key)
        return None


def collect_nodes(root):
    nodes = []
    parent = []

    def dfs(node, p):
        idx = len(nodes)
        nodes.append(node)
        parent.append(p)
        for ch in node.children:
            dfs(ch, idx)

    dfs(root, -1)
    return nodes, parent


def leo_from_htd(nh, nodes, parent):
    children = [[] for _ in nodes]
    root = 0
    for i, p in enumerate(parent):
        if p >= 0:
            children[p].append(i)
        else:
            root = i
    depth = [0] * len(nodes)
    st = [root]
    for x in st:
        for c in children[x]:
            depth[c] = depth[x] + 1
            st.append(c)
    f = [None] * nh
    for v in range(nh):
        cand = [i for i, nd in enumerate(nodes) if (nd.chi >> v) & 1]
        if not cand:
            return None
        f[v] = min(cand, key=lambda i: depth[i])
    order = sorted(range(nh), key=lambda v: (-depth[f[v]], f[v], v))
    covers = [set(bits(nodes[f[v]].lam)) for v in range(nh)]
    return order, covers, depth, f


def main():
    for path in sys.argv[1:]:
        n, m, ws, ps, edges_list, _ = extract(path)
        W = int(re.search(r"hg_(\d+)", os.path.basename(path)).group(1))
        inc = [[] for _ in range(n)]
        for h, vs in enumerate(edges_list):
            for v in vs:
                inc[v].append(h)
        if any(len(x) != 2 for x in inc):
            print(os.path.basename(path), "not-2-reg")
            continue
        det = DetK(n, m, [tuple(x) for x in inc], W)
        root = det.decomp((1 << m) - 1, 0)
        print(os.path.basename(path), "W", W, "calls", det.calls, "fail", len(det.fail), "found", root is not None)
        if root is None:
            continue
        nodes, parent = collect_nodes(root)
        leo = leo_from_htd(n, nodes, parent)
        if leo is None:
            print(" no-leo")
            continue
        order, covers, depth, f = leo
        print(" nodes", len(nodes), "cover_max", max(map(len, covers)))
        pos = [0] * n
        for t, v in enumerate(order):
            pos[v] = t
        edges = [frozenset(x) for x in edges_list]
        arcs = compute_arcs(n, edges, order)
        for _ in range(n):
            p = p_closure(n, edges, covers, arcs, pos)
            added = 0
            for i in range(n):
                for j in range(n):
                    if not arcs[i][j]:
                        continue
                    for k in range(n):
                        if i != k and j != k and arcs[j][k] and p[i][k] and not arcs[i][k]:
                            if pos[i] < pos[k]:
                                arcs[i][k] = True
                                added += 1
            if added == 0:
                break
        val, max_sem = build_val(n, m, order, covers, ws, ps, arcs, p, pos)
        total, bad, ex = check_semantic(path, val, max_sem)
        print(" semantic_bad", bad, ex[:3])


if __name__ == "__main__":
    main()
