"""SAT model: does a D-day schedule covering all pairs exist?
Variables x[p,d,b] = person p in boat b on day d.
Used to (a) corroborate D=5 is UNSAT, (b) find/confirm D=6 SAT."""
import sys
from itertools import combinations
from pysat.formula import IDPool
from pysat.card import CardEnc, EncType
from pysat.solvers import Cadical153

N, BOATS, SIZE = 24, 4, 6


def solve(DAYS, break_sym=True, time_budget=None):
    vp = IDPool()
    x = lambda p, d, b: vp.id(f"x_{p}_{d}_{b}")
    cnf = []

    # exactly one boat per person per day
    for p in range(N):
        for d in range(DAYS):
            lits = [x(p, d, b) for b in range(BOATS)]
            cnf.append(lits[:])                      # at least one
            for a, c in combinations(lits, 2):       # at most one
                cnf.append([-a, -c])

    # each boat exactly SIZE people per day
    for d in range(DAYS):
        for b in range(BOATS):
            lits = [x(p, d, b) for p in range(N)]
            enc = CardEnc.equals(lits=lits, bound=SIZE, vpool=vp, encoding=EncType.seqcounter)
            cnf.extend(enc.clauses)

    # coverage: each pair shares some boat on some day
    for p, q in combinations(range(N), 2):
        clause = []
        for d in range(DAYS):
            for b in range(BOATS):
                y = vp.id(f"y_{p}_{q}_{d}_{b}")
                cnf.append([-y, x(p, d, b)])
                cnf.append([-y, x(q, d, b)])
                clause.append(y)
        cnf.append(clause)

    if break_sym:
        # Day 0 fixed to bands: person p -> boat p//SIZE
        for p in range(N):
            cnf.append([x(p, 0, p // SIZE)])
        # Day 1: person 0..3 seeded into distinct boats (partial), and 0 in boat 0
        for b in range(BOATS):
            cnf.append([x(b, 1, b)])

    solver = Cadical153(bootstrap_with=cnf)
    sat = solver.solve()
    model = solver.get_model() if sat else None
    res = None
    if sat:
        ms = set(l for l in model if l > 0)
        schedule = []
        for d in range(DAYS):
            day = [[] for _ in range(BOATS)]
            for p in range(N):
                for b in range(BOATS):
                    if x(p, d, b) in ms:
                        day[b].append(p)
            schedule.append([sorted(bt) for bt in day])
        res = schedule
    solver.delete()
    return sat, res


if __name__ == "__main__":
    D = int(sys.argv[1]) if len(sys.argv) > 1 else 5
    print(f"Solving D={D} ...", flush=True)
    sat, sched = solve(D)
    print("SAT" if sat else "UNSAT")
    if sat:
        import json
        from lib import verify
        ok, info = verify(sched)
        print("verify full cover:", ok, info['missing'])
        with open(f"solution{D}_sat.json", "w") as f:
            json.dump(sched, f)
        print(f"saved solution{D}_sat.json")