#!/usr/bin/env python3
import collections
import lzma
import random
import sys
import time
from pathlib import Path

from extract_highlevel import extract


def clause_satisfied(c, val):
    for lit in c:
        v = abs(lit)
        if val[v] == 0:
            continue
        if (lit > 0 and val[v] > 0) or (lit < 0 and val[v] < 0):
            return True
    return False


class Propagator:
    def __init__(self, nvars, clauses):
        self.nvars = nvars
        self.clauses = clauses
        self.clauses_by_lit = collections.defaultdict(list)
        for idx, c in enumerate(clauses):
            for lit in c:
                self.clauses_by_lit[lit].append(idx)

    def propagate(self, selected, n_agents):
        nvars = self.nvars
        clauses = self.clauses
        clauses_by_lit = self.clauses_by_lit
        val = [0] * (nvars + 1)
        q = collections.deque()

        def assign_lit(lit):
            v = abs(lit)
            b = 1 if lit > 0 else -1
            if val[v] == b:
                return True
            if val[v] == -b:
                return False
            val[v] = b
            q.append(lit)
            return True

        for i in range(1, n_agents + 1):
            if not assign_lit(i if i in selected else -i):
                return None
        for c in clauses:
            if len(c) == 1:
                if not assign_lit(c[0]):
                    return None

        while q:
            lit = q.popleft()
            false_lit = -lit
            for ci in clauses_by_lit.get(false_lit, []):
                c = clauses[ci]
                sat = False
                un = 0
                last = 0
                for x in c:
                    v = abs(x)
                    xv = val[v]
                    if xv == 0:
                        un += 1
                        last = x
                    elif (x > 0 and xv > 0) or (x < 0 and xv < 0):
                        sat = True
                        break
                if sat:
                    continue
                if un == 0:
                    return None
                if un == 1:
                    if not assign_lit(last):
                        return None
        for v in range(1, nvars + 1):
            if val[v] == 0:
                val[v] = -1
        for c in clauses:
            if not clause_satisfied(c, val):
                return None
        return val


def highlevel_ok(skills, selected):
    for agents, req, _, _ in skills:
        if sum(1 for a in agents if a in selected) < req:
            return False
    return True


def solve_once(n, skills, conflicts, propagator, rng, cost, budget):
    skill_of = [[] for _ in range(n + 1)]
    for si, (agents, req, _, _) in enumerate(skills):
        for a in agents:
            skill_of[a].append(si)
    forbidden = [False] * (n + 1)
    selected = set()
    deficits = [req for _, req, _, _ in skills]
    adj = [[] for _ in range(n + 1)]
    for a, b in conflicts:
        adj[a].append(b)
        adj[b].append(a)

    for _step in range(sum(deficits) + 200):
        if max(deficits) <= 0:
            break
        best_skill = None
        best_avail = 10**9
        # Bias toward tight, still-deficient skills.
        order = [i for i, d in enumerate(deficits) if d > 0]
        rng.shuffle(order)
        for si in order[: min(len(order), 40)]:
            avail = sum(1 for a in skills[si][0] if not forbidden[a] and a not in selected)
            key = avail * 10 - deficits[si]
            if key < best_avail:
                best_avail = key
                best_skill = si
        if best_skill is None:
            return None
        cand = []
        for a in skills[best_skill][0]:
            if forbidden[a] or a in selected:
                continue
            gain = 0
            for si in skill_of[a]:
                if deficits[si] > 0:
                    gain += 1
            if gain <= 0:
                continue
            # Lower cost is better; add noise to diversify.
            score = gain * 10000.0 / (1.0 + cost[a]) + rng.random() * 0.5
            cand.append((score, a))
        if not cand:
            return None
        cand.sort(reverse=True)
        pick_pool = cand[: min(len(cand), 5)]
        a = rng.choice(pick_pool)[1] if rng.random() < 0.35 else pick_pool[0][1]
        selected.add(a)
        for b in adj[a]:
            forbidden[b] = True
        for si in skill_of[a]:
            deficits[si] -= 1

    if max(deficits) > 0:
        return None

    # Prune redundant agents, highest estimated cost first.
    changed = True
    while changed:
        changed = False
        order = list(selected)
        order.sort(key=lambda a: (cost[a], len(skill_of[a])), reverse=True)
        for a in order:
            selected.remove(a)
            if highlevel_ok(skills, selected):
                changed = True
                for si in skill_of[a]:
                    deficits[si] += 1
            else:
                selected.add(a)

    if budget is not None and sum(cost[a] for a in selected) > budget:
        return None

    val = propagator.propagate(selected, n)
    if val is None:
        return None
    return selected, val


def main():
    for arg in sys.argv[1:]:
        path = Path(arg)
        n, skills, conflicts, units_after, clauses = extract(path)
        nvars = max(abs(l) for c in clauses for l in c)
        parent = list(range(n + 1))
        def find(x):
            while parent[x] != x:
                parent[x] = parent[parent[x]]
                x = parent[x]
            return x
        def union(a, b):
            ra, rb = find(a), find(b)
            if ra != rb:
                parent[rb] = ra
        for a, b in conflicts:
            union(a, b)
        comps = collections.defaultdict(list)
        for i in range(1, n + 1):
            comps[find(i)].append(i)
        cost = [1] * (n + 1)
        for comp in comps.values():
            for rank, a in enumerate(sorted(comp), 1):
                cost[a] = rank
        import re
        m = re.search(r"_(\d+)_0\.\d+_(\d+)\.cnf", path.name)
        budget = int(m.group(2)) if m else None
        propagator = Propagator(nvars, clauses)
        start = time.time()
        ans = None
        attempts = 0
        while time.time() - start < 20 and ans is None:
            attempts += 1
            rng = random.Random(12345 + attempts * 1009)
            ans = solve_once(n, skills, conflicts, propagator, rng, cost, budget)
        print(path.name, "attempts", attempts, "time", round(time.time()-start, 2), "SAT" if ans else "no")
        if ans:
            selected, _ = ans
            print(" selected", len(selected), "costproxy", sum(cost[a] for a in selected), "agents", sorted(selected)[:80])


if __name__ == "__main__":
    main()
