"""Sức chứa tích lũy trên thời gian nguyên, kiểm tra bằng vét cạn độc lập."""
from itertools import product
import clingo

PROGRAM = r"""
#const cap=3.
time(0..3). task(a,2,2). task(b,2,2).
startok(a,0). startok(a,1). startok(b,0). startok(b,2).
1 { start(J,S) : startok(J,S) } 1 :- task(J,D,Q).
busy(J,T) :- start(J,S), task(J,D,Q), time(T), S <= T, T < S+D.
usage(T,N) :- time(T), N = #sum { Q,J : busy(J,T), task(J,D,Q) }.
:- usage(T,N), N > cap.
#show start/2.
"""

def asp_schedules(capacity):
    ctl = clingo.Control(["0", "-c", f"cap={capacity}"])
    ctl.add("base", [], PROGRAM)
    ctl.ground([("base", [])])
    schedules = set()
    def collect(model):
        starts = {str(a.arguments[0]): a.arguments[1].number
                  for a in model.symbols(shown=True)}
        schedules.add((starts["a"], starts["b"]))
    result = ctl.solve(on_model=collect)
    assert result.exhausted and not result.interrupted
    return schedules

def oracle(capacity):
    result = set()
    for sa, sb in product((0, 1), (0, 2)):
        # Duyệt trực tiếp từng đơn vị thiết bị ở mỗi thời điểm.
        loads = [sum(2 for start in (sa, sb) if start <= t < start + 2)
                 for t in range(4)]
        if max(loads) <= capacity:
            result.add((sa, sb))
    return result

for capacity in (1, 3, 4):
    actual = asp_schedules(capacity)
    assert actual == oracle(capacity)
    print(f"Sức chứa {capacity}: {len(actual)} lịch, {sorted(actual)}")
assert asp_schedules(3) == {(0, 2)}
assert len(asp_schedules(4)) == 4