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


def load_clauses(path):
    clauses = []
    nvars = nclauses = 0
    opener = lzma.open if path.endswith(".xz") else open
    with opener(path, "rt") as f:
        for line in f:
            if line.startswith("p"):
                _, _, nvars, nclauses = line.split()[:4]
                nvars, nclauses = int(nvars), int(nclauses)
                continue
            if not line or line[0] in "c%0":
                continue
            vals = tuple(int(x) for x in line.split() if x != "0")
            if vals:
                clauses.append(vals)
    return nvars, nclauses, clauses


def detect_ands(clauses):
    binaries = {frozenset(c) for c in clauses if len(c) == 2}
    gates = []
    for cl in clauses:
        if len(cl) != 4:
            continue
        candidates = []
        for out_lit in cl:
            inputs = tuple(-lit for lit in cl if lit != out_lit)
            if all(frozenset((inp, -out_lit)) in binaries for inp in inputs):
                candidates.append((out_lit, inputs))
        if len(candidates) != 1:
            raise RuntimeError(f"ambiguous AND gate {cl}: {candidates}")
        gates.append(candidates[0])
    return gates


def xor_const(clauses, vars_):
    good = []
    index = {v: i for i, v in enumerate(vars_)}
    for mask in range(8):
        ok = True
        for cl in clauses:
            sat = False
            for lit in cl:
                bit = (mask >> index[abs(lit)]) & 1
                if lit < 0:
                    bit ^= 1
                if bit:
                    sat = True
                    break
            if not sat:
                ok = False
                break
        if ok:
            good.append(mask)
    if len(good) != 4:
        return None
    parities = {mask.bit_count() & 1 for mask in good}
    if len(parities) != 1:
        return None
    return parities.pop()


def detect_xors(clauses):
    groups = collections.defaultdict(list)
    for cl in clauses:
        if len(cl) == 3:
            groups[tuple(sorted(abs(lit) for lit in cl))].append(cl)
    xors = []
    for vars_, cls in groups.items():
        if len(cls) != 4:
            continue
        const = xor_const(cls, list(vars_))
        if const is not None:
            xors.append((vars_, const))
    return xors


def partition_rank(comp, adj):
    remaining = set(comp)
    parts = []
    while remaining:
        first = next(iter(remaining))
        part = {x for x in remaining if x == first or x not in adj[first]}
        parts.append(sorted(part, key=lambda z: abs(z)))
        remaining -= part
    parts.sort(key=lambda p: (len(p), min(abs(x) for x in p)))
    return parts


def reconstruct(path):
    nvars, nclauses, clauses = load_clauses(path)
    ands = detect_ands(clauses)
    xors = detect_xors(clauses)
    out_to_gate = {abs(out): (out, inputs) for out, inputs in ands}

    base_lits = sorted({lit for _, inputs in ands for lit in inputs}, key=lambda x: abs(x))
    adj = {lit: set() for lit in base_lits}
    for _, inputs in ands:
        for a, b in itertools.combinations(inputs, 2):
            adj[a].add(b)
            adj[b].add(a)

    seen = set()
    comps = []
    for lit in base_lits:
        if lit in seen:
            continue
        stack = [lit]
        seen.add(lit)
        comp = []
        while stack:
            cur = stack.pop()
            comp.append(cur)
            for nxt in adj[cur]:
                if nxt not in seen:
                    seen.add(nxt)
                    stack.append(nxt)
        comps.append(sorted(comp, key=lambda x: abs(x)))
    comps.sort(key=lambda c: min(abs(x) for x in c))

    rank_of_lit = {lit: r for r, comp in enumerate(comps) for lit in comp}
    parts = [partition_rank(comp, adj) for comp in comps]
    part_of = {}
    for r, rank_parts in enumerate(parts):
        for p, part in enumerate(rank_parts):
            for q, lit in enumerate(part):
                part_of[lit] = (r, p, q)

    and_rank = {}
    for out_var, (_, inputs) in out_to_gate.items():
        ranks = {rank_of_lit[x] for x in inputs}
        if len(ranks) != 1:
            raise RuntimeError("AND gate crosses rank components")
        and_rank[out_var] = next(iter(ranks))

    xadj = collections.defaultdict(list)
    for gid, (vars_, _) in enumerate(xors):
        for v in vars_:
            xadj[v].append(gid)

    units = [cl[0] for cl in clauses if len(cl) == 1]
    cells = []
    for unit in units:
        const = 1 if unit > 0 else 0
        seen_gates = set()
        seen_vars = {abs(unit)}
        stack = [abs(unit)]
        leaves = []
        while stack:
            var = stack.pop()
            for gid in xadj[var]:
                if gid in seen_gates:
                    continue
                seen_gates.add(gid)
                vars_, xconst = xors[gid]
                const ^= xconst
                for nxt in vars_:
                    if nxt in seen_vars:
                        continue
                    seen_vars.add(nxt)
                    if nxt in out_to_gate:
                        leaves.append(nxt)
                    else:
                        stack.append(nxt)
        if len(leaves) != len(comps):
            raise RuntimeError(f"unit {unit} reaches {len(leaves)} leaves, expected {len(comps)}")
        leaf_by_rank = {and_rank[v]: v for v in leaves}
        if len(leaf_by_rank) != len(comps):
            raise RuntimeError("XOR component does not contain one leaf per rank")

        target = const
        for leaf in leaves:
            out_lit, _ = out_to_gate[leaf]
            if out_lit < 0:
                target ^= 1

        coord = [None, None, None]
        ref_leaf = leaf_by_rank[0]
        for lit in out_to_gate[ref_leaf][1]:
            _, part, pos = part_of[lit]
            coord[part] = pos
        cells.append((tuple(coord), target, leaf_by_rank))

    dims = tuple(len(part) for part in parts[0])
    if len({coord for coord, _, _ in cells}) != len(cells):
        raise RuntimeError("duplicate reconstructed tensor cell coordinates")

    all_coords = {coord for coord, _, _ in cells}
    rank_lit_map = []
    for r in range(len(comps)):
        occur = collections.defaultdict(set)
        for coord, _, leaf_by_rank in cells:
            _, inputs = out_to_gate[leaf_by_rank[r]]
            for lit in inputs:
                occur[lit].add(coord)

        mapping = {}
        for lit, coords in occur.items():
            matches = []
            for dim, size in enumerate(dims):
                for pos in range(size):
                    ideal = {coord for coord in all_coords if coord[dim] == pos}
                    if coords == ideal:
                        matches.append((dim, pos))
            if len(matches) != 1:
                raise RuntimeError(f"could not align rank {r} literal {lit}")
            mapping[matches[0]] = lit
        if len(mapping) != sum(dims):
            raise RuntimeError(f"rank {r} map has {len(mapping)} literals")
        rank_lit_map.append(mapping)

    target = 0
    for coord, bit, _ in cells:
        idx = (coord[0] * dims[1] + coord[1]) * dims[2] + coord[2]
        if bit:
            target |= 1 << idx

    return {
        "path": path,
        "nvars": nvars,
        "nclauses": nclauses,
        "clauses": clauses,
        "dims": dims,
        "rank": len(comps),
        "target": target,
        "rank_lit_map": rank_lit_map,
    }


def make_items(a, b, c):
    items = [0]
    factors = [(0, 0, 0)]
    matrices = []
    for v in range(1, 1 << b):
        for w in range(1, 1 << c):
            mask = 0
            for j in range(b):
                if (v >> j) & 1:
                    for k in range(c):
                        if (w >> k) & 1:
                            mask |= 1 << (j * c + k)
            matrices.append((v, w, mask))
    plane = b * c
    for u in range(1, 1 << a):
        for v, w, mat in matrices:
            bits = 0
            for i in range(a):
                if (u >> i) & 1:
                    bits |= mat << (i * plane)
            items.append(bits)
            factors.append((u, v, w))
    return items, factors


def item_search(a, b, c, rank, target, seconds, seed):
    rng = random.Random(seed)
    items, factors = make_items(a, b, c)
    nitems = len(items)
    deadline = time.time() + seconds
    best = 10**9
    best_sel = None

    while time.time() < deadline:
        selected = [0] * rank
        current = 0
        residual = target

        # Greedy-random start against the current residual.
        for t in range(rank):
            if rng.random() < 0.75:
                best_i = 0
                best_s = residual.bit_count()
                for _ in range(2500):
                    i = rng.randrange(1, nitems)
                    s = (residual ^ items[i]).bit_count()
                    if s < best_s:
                        best_s = s
                        best_i = i
                i = best_i
            else:
                i = rng.randrange(nitems)
            selected[t] = i
            current ^= items[i]
            residual ^= items[i]

        no_improve = 0
        temp = 2.5
        while time.time() < deadline and no_improve < 120:
            residual = current ^ target
            score = residual.bit_count()
            if score == 0:
                return [factors[i] for i in selected], 0
            if score < best:
                best = score
                best_sel = selected[:]

            order = list(range(rank))
            rng.shuffle(order)
            best_move = None
            best_score = 10**9
            for t in order:
                wanted = residual ^ items[selected[t]]
                local_i = selected[t]
                local_s = (wanted ^ items[local_i]).bit_count()
                for i, item in enumerate(items):
                    s = (wanted ^ item).bit_count()
                    if s < local_s:
                        local_s = s
                        local_i = i
                        if s == 0:
                            break
                if local_s < best_score:
                    best_score = local_s
                    best_move = (t, local_i)
                if local_s < score:
                    break

            t, new_i = best_move
            if best_score <= score or rng.random() < math.exp((score - best_score) / max(0.1, temp)):
                current ^= items[selected[t]] ^ items[new_i]
                selected[t] = new_i
                no_improve = 0 if best_score < score else no_improve + 1
            else:
                t = rng.randrange(rank)
                new_i = rng.randrange(nitems)
                current ^= items[selected[t]] ^ items[new_i]
                selected[t] = new_i
                no_improve += 1
            temp *= 0.992

    if best_sel is None:
        return None, best
    return None, best


def main(argv):
    seconds = float(os.environ.get("SECONDS_PER_INSTANCE", "5"))
    for path in argv[1:]:
        rec = reconstruct(path)
        a, b, c = rec["dims"]
        start = time.time()
        sol, best = item_search(a, b, c, rec["rank"], rec["target"], seconds, 12345)
        print(
            os.path.basename(path),
            "dims",
            rec["dims"],
            "rank",
            rec["rank"],
            "ok",
            sol is not None,
            "best",
            best,
            "time",
            f"{time.time() - start:.2f}",
            flush=True,
        )


if __name__ == "__main__":
    main(sys.argv)
