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

from extract_highlevel import extract
from greedy_probe import Propagator
from level_greedy import build_components, solve_level


class Search:
    def __init__(self, n, skills, conflicts, clauses, units_after, budget, seed=1):
        self.n = n
        self.skills = skills
        self.budget = budget
        self.rng = random.Random(seed)
        self.nvars = max(abs(l) for c in clauses for l in c)
        self.full_prop = Propagator(self.nvars, clauses)
        self.obj_prop = Propagator(self.nvars, clauses[units_after[0][0] + 1 :])
        self.comps = build_components(n, conflicts)
        self.skill_of = [set() for _ in range(n + 1)]
        for si, (agents, req, _, _) in enumerate(skills):
            for a in agents:
                self.skill_of[a].add(si)
        self.opts = []
        for comp in self.comps:
            levels = [0] + sorted(comp)
            covers = [set()]
            for a in levels[1:]:
                covers.append(set(self.skill_of[a]))
            self.opts.append((levels, covers))
        self.skill_comps = [[] for _ in skills]
        for ci, (_, covers) in enumerate(self.opts):
            for si in set().union(*covers) if len(covers) > 1 else []:
                self.skill_comps[si].append(ci)
        self.nodes = 0
        self.dead_obj = 0
        self.dead_avail = 0

    def selected_set(self, level):
        return {self.opts[ci][0][lv] for ci, lv in enumerate(level) if lv > 0}

    def objective_ok(self, level):
        return self.obj_prop.propagate(self.selected_set(level), self.n) is not None

    def dfs(self, level, deficits, approx_cost, deadline, depth=0):
        self.nodes += 1
        if self.nodes % 5000 == 0 and time.time() > deadline:
            return None
        if max(deficits) <= 0:
            sel = self.selected_set(level)
            if self.full_prop.propagate(sel, self.n) is not None:
                return sel
            return None

        # Basic availability check and branch-skill selection.
        best_si = None
        best_cands = None
        for si, d in enumerate(deficits):
            if d <= 0:
                continue
            cands = []
            avail = 0
            for ci in self.skill_comps[si]:
                cur = level[ci]
                if si in self.opts[ci][1][cur]:
                    avail += 1
                    continue
                # cheapest level that covers this skill
                nl = None
                for l in range(cur + 1, len(self.opts[ci][0])):
                    if si in self.opts[ci][1][l]:
                        nl = l
                        break
                if nl is not None:
                    avail += 1
                    cands.append((ci, nl))
            if avail < d:
                self.dead_avail += 1
                return None
            key = (len(cands), avail, -d)
            if best_si is None or key < best_key:
                best_si = si
                best_key = key
                best_cands = cands

        choices = []
        for ci, minl in best_cands:
            cur = level[ci]
            for nl in range(minl, len(self.opts[ci][0])):
                dc = nl - cur
                if self.budget is not None and approx_cost + dc > self.budget:
                    break
                new = self.opts[ci][1][nl] - self.opts[ci][1][cur]
                gain = sum(1 for si in new if deficits[si] > 0)
                pressure = sum(deficits[si] for si in new if deficits[si] > 0)
                if gain <= 0:
                    continue
                score = (gain * 1000 + pressure * 30) / dc
                choices.append((-score, dc, ci, nl))
        choices.sort()
        # Diversify equivalent high-score choices.
        head = choices[: min(12, len(choices))]
        self.rng.shuffle(head)
        choices[: len(head)] = head

        for _, dc, ci, nl in choices:
            old = level[ci]
            level[ci] = nl
            if self.obj_prop.propagate(self.selected_set(level), self.n) is None:
                self.dead_obj += 1
                level[ci] = old
                continue
            newdef = deficits[:]
            for si in self.opts[ci][1][nl] - self.opts[ci][1][old]:
                newdef[si] -= 1
            got = self.dfs(level, newdef, approx_cost + dc, deadline, depth + 1)
            if got is not None:
                return got
            level[ci] = old
            if time.time() > deadline:
                return None
        return None


def main():
    for arg in sys.argv[1:]:
        path = Path(arg)
        n, skills, conflicts, units_after, clauses = extract(path)
        m = re.search(r"_(\d+)_0\.\d+_(\d+)\.cnf", path.name)
        budget = int(m.group(2)) if m else None
        start = time.time()
        ans = None
        stats = None
        for seed in range(1, 20):
            s = Search(n, skills, conflicts, clauses, units_after, budget, seed)
            # Start from an objective-feasible greedy prefix if available.
            level = [0] * len(s.opts)
            deficits = [req for _, req, _, _ in skills]
            got = s.dfs(level, deficits, 0, start + 60)
            stats = (s.nodes, s.dead_obj, s.dead_avail)
            if got:
                ans = got
                break
            if time.time() - start > 60:
                break
        print(path.name, "SAT" if ans else "no", "time", round(time.time() - start, 2), "stats", stats)
        if ans:
            print(" size", len(ans), "agents", sorted(ans))


if __name__ == "__main__":
    main()
