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

from extract_highlevel import extract
from greedy_probe import Propagator


def build_components(n, conflicts):
    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)
    return [sorted(v) for v in comps.values()]


def solve_level(n, skills, conflicts, budget, rng, seconds=10, objective_prop=None):
    skill_of = [set() for _ in range(n + 1)]
    for si, (agents, req, _, _) in enumerate(skills):
        for a in agents:
            skill_of[a].add(si)
    comps = build_components(n, conflicts)
    comp_opts = []
    for comp in comps:
        opts = [0] + comp
        covers = [set()]
        ok_nested = True
        prev = set()
        for a in comp:
            cur = set(skill_of[a])
            if not prev.issubset(cur):
                ok_nested = False
            covers.append(cur)
            prev = cur
        # Keep variable order by id; it is the encoded cost order in this family.
        comp_opts.append((opts, covers, ok_nested))

    start = time.time()
    while time.time() - start < seconds:
        level = [0] * len(comp_opts)
        deficits = [req for _, req, _, _ in skills]
        total_cost = 0
        for _step in range((budget or 200) + 100):
            if max(deficits) <= 0:
                break
            choices = []
            for ci, (opts, covers, _) in enumerate(comp_opts):
                cur = level[ci]
                base = covers[cur]
                maxl = len(opts) - 1
                for nl in range(cur + 1, maxl + 1):
                    dc = nl - cur
                    if budget is not None and total_cost + dc > budget:
                        break
                    gain = 0
                    new = covers[nl] - base
                    for si in new:
                        if deficits[si] > 0:
                            gain += 1
                    if gain == 0:
                        continue
                    # Tight skills get extra pressure.
                    pressure = sum(deficits[si] for si in new if deficits[si] > 0)
                    score = (gain * 1000 + pressure * 25) / dc
                    score += rng.random() * 5
                    choices.append((score, ci, nl))
            if not choices:
                break
            choices.sort(reverse=True)
            pool = choices[: min(len(choices), 40)]
            picked = None
            if rng.random() < 0.45:
                rng.shuffle(pool)
            for _, ci0, nl0 in pool:
                old0 = level[ci0]
                trial = {comp_opts[j][0][lv] for j, lv in enumerate(level) if lv > 0 and j != ci0}
                trial.add(comp_opts[ci0][0][nl0])
                if objective_prop is None or objective_prop.propagate(trial, n) is not None:
                    picked = (ci0, nl0)
                    break
            if picked is None:
                break
            ci, nl = picked
            old = level[ci]
            oldcov = comp_opts[ci][1][old]
            newcov = comp_opts[ci][1][nl]
            level[ci] = nl
            total_cost += nl - old
            for si in newcov - oldcov:
                deficits[si] -= 1

        if max(deficits) > 0:
            continue

        # Downgrade expensive choices where possible.
        changed = True
        while changed:
            changed = False
            order = list(range(len(comp_opts)))
            rng.shuffle(order)
            order.sort(key=lambda ci: level[ci], reverse=True)
            for ci in order:
                cur = level[ci]
                if cur == 0:
                    continue
                for nl in range(cur - 1, -1, -1):
                    remove = comp_opts[ci][1][cur] - comp_opts[ci][1][nl]
                    if all(deficits[si] + 1 <= 0 for si in remove):
                        for si in remove:
                            deficits[si] += 1
                        total_cost -= cur - nl
                        level[ci] = nl
                        changed = True
                        break

        if budget is not None and total_cost > budget:
            continue
        selected = {comp_opts[ci][0][lv] for ci, lv in enumerate(level) if lv > 0}
        return selected, total_cost
    return None


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)
        m = re.search(r"_(\d+)_0\.\d+_(\d+)\.cnf", path.name)
        budget = int(m.group(2)) if m else None
        prop = Propagator(nvars, clauses)
        obj_prop = Propagator(nvars, clauses[units_after[0][0] + 1 :])
        start = time.time()
        ans = None
        attempts = 0
        while time.time() - start < 60 and ans is None:
            attempts += 1
            got = solve_level(
                n,
                skills,
                conflicts,
                budget,
                random.Random(10007 * attempts),
                seconds=1,
                objective_prop=obj_prop,
            )
            if got is None:
                continue
            selected, cost = got
            val = prop.propagate(selected, n)
            if val is not None:
                ans = (selected, cost)
                break
        print(path.name, "attempts", attempts, "time", round(time.time() - start, 2), "SAT" if ans else "no")
        if ans:
            selected, cost = ans
            print(" cost", cost, "size", len(selected), "agents", sorted(selected))


if __name__ == "__main__":
    main()
