"""Local search for a D-day cover. State = D partitions into 4 boats of 6.
Move = swap two people in different boats on the same day. Minimize #uncovered
pairs (cover redundancy is fine). Restart-friendly; D=6 has lots of slack."""
import random
from itertools import combinations
from lib import verify, all_pairs

N, BOATS, SIZE = 24, 4, 6


def boat_of(day):
    """map person -> boat index for a day given as list of 4 boats."""
    m = {}
    for bi, boat in enumerate(day):
        for p in boat:
            m[p] = bi
    return m


def random_day(rng):
    people = list(range(N))
    rng.shuffle(people)
    return [people[i*SIZE:(i+1)*SIZE] for i in range(BOATS)]


def coverage_count(schedule):
    cnt = {}
    for day in schedule:
        for boat in day:
            for a, b in combinations(sorted(boat), 2):
                cnt[(a, b)] = cnt.get((a, b), 0) + 1
    return cnt


def missing(schedule):
    cnt = coverage_count(schedule)
    return len(all_pairs() - set(cnt.keys()))


def search(D, seed, max_iter=200000):
    rng = random.Random(seed)
    day1 = [list(range(0, 6)), list(range(6, 12)), list(range(12, 18)), list(range(18, 24))]
    schedule = [day1] + [random_day(rng) for _ in range(D - 1)]
    cnt = coverage_count(schedule)
    cur_missing = len(all_pairs() - set(cnt.keys()))
    best = cur_missing
    T = 2.0
    for it in range(max_iter):
        if cur_missing == 0:
            break
        d = rng.randrange(1, D)              # don't disturb day 1
        b1, b2 = rng.sample(range(BOATS), 2)
        i = rng.randrange(SIZE)
        j = rng.randrange(SIZE)
        p = schedule[d][b1][i]
        q = schedule[d][b2][j]
        # delta on coverage counts: removing p from b1, q from b2; adding swap
        def pair_delta(remove_sets, add_sets):
            # returns change in #uncovered = (newly uncovered) - (newly covered)
            dmiss = 0
            for (a, b) in remove_sets:
                k = (min(a, b), max(a, b))
                if cnt.get(k, 0) == 1:
                    dmiss += 1
            for (a, b) in add_sets:
                k = (min(a, b), max(a, b))
                if cnt.get(k, 0) == 0:
                    dmiss -= 1
            return dmiss
        b1_rest = [schedule[d][b1][k] for k in range(SIZE) if k != i]
        b2_rest = [schedule[d][b2][k] for k in range(SIZE) if k != j]
        removed = [(p, x) for x in b1_rest] + [(q, x) for x in b2_rest]
        added = [(q, x) for x in b1_rest] + [(p, x) for x in b2_rest]
        dmiss = pair_delta(removed, added)
        if dmiss <= 0 or rng.random() < pow(2.718, -dmiss / T):
            # apply swap and update cnt
            for (a, b) in removed:
                k = (min(a, b), max(a, b)); cnt[k] -= 1
                if cnt[k] == 0: del cnt[k]
            for (a, b) in added:
                k = (min(a, b), max(a, b)); cnt[k] = cnt.get(k, 0) + 1
            schedule[d][b1][i], schedule[d][b2][j] = q, p
            cur_missing += dmiss
            best = min(best, cur_missing)
        T = max(0.05, T * 0.99997)
    return schedule, cur_missing


if __name__ == "__main__":
    import json, sys
    D = int(sys.argv[1]) if len(sys.argv) > 1 else 6
    for seed in range(500):
        sched, miss = search(D, seed)
        if miss == 0:
            ok, info = verify(sched)
            print(f"seed {seed}: SOLVED D={D}  full_cover={ok}  "
                  f"meetings={info['total_meetings']} excess={info['repeat_excess']}")
            out = [[sorted(b) for b in day] for day in sched]
            with open(f"solution{D}.json", "w") as f:
                json.dump(out, f)
            print(f"saved solution{D}.json")
            break
    else:
        print(f"no full cover found for D={D}")