"""
Try to construct a 5-day schedule for 24 participants, 4 boats of 6 each day,
using local search / simulated annealing.
"""
import itertools
import random
import math
from collections import defaultdict

N = 24
BOAT_SIZE = 6
BOATS_PER_DAY = 4
DAYS = 5
TOTAL_PAIRS = math.comb(N, 2)

def pairs_of_boat(boat):
    return set(itertools.combinations(boat, 2))

def pairs_of_day(day):
    pairs = set()
    for boat in day:
        pairs.update(pairs_of_boat(boat))
    return pairs

def random_partition():
    nums = list(range(N))
    random.shuffle(nums)
    return [tuple(sorted(nums[i*BOAT_SIZE:(i+1)*BOAT_SIZE])) for i in range(BOATS_PER_DAY)]

def random_schedule():
    return [random_partition() for _ in range(DAYS)]

def count_covered_pairs(schedule):
    """Count distinct pairs covered by the schedule."""
    all_pairs = set()
    for day in schedule:
        for boat in day:
            all_pairs.update(itertools.combinations(boat, 2))
    return len(all_pairs)

def missing_pairs(schedule):
    """Return set of pairs not yet covered."""
    all_pairs = set()
    for day in schedule:
        for boat in day:
            all_pairs.update(itertools.combinations(boat, 2))
    return set(itertools.combinations(range(N), 2)) - all_pairs

def perturb_schedule(schedule):
    """Randomly perturb the schedule by swapping two people across boats."""
    sched = [list(day) for day in schedule]
    day_idx = random.randrange(DAYS)
    day = [list(boat) for boat in sched[day_idx]]
    
    # Pick two different boats
    b1, b2 = random.sample(range(BOATS_PER_DAY), 2)
    # Pick people to swap
    p1 = random.randrange(BOAT_SIZE)
    p2 = random.randrange(BOAT_SIZE)
    
    day[b1][p1], day[b2][p2] = day[b2][p2], day[b1][p1]
    
    sched[day_idx] = [tuple(sorted(boat)) for boat in day]
    return sched

def local_search(max_iterations=200000, num_restarts=20):
    """Local search / hill climbing with restarts."""
    
    best_overall = TOTAL_PAIRS
    best_schedule_overall = None
    
    for restart in range(num_restarts):
        sched = random_schedule()
        covered = count_covered_pairs(sched)
        
        best_local = covered
        best_local_sched = [list(d) for d in sched]
        
        no_improve = 0
        
        for iteration in range(max_iterations):
            new_sched = perturb_schedule(sched)
            new_covered = count_covered_pairs(new_sched)
            
            if new_covered >= covered:
                sched = new_sched
                covered = new_covered
                no_improve = 0
                
                if covered > best_local:
                    best_local = covered
                    best_local_sched = [list(d) for d in sched]
                    print(f"  Restart {restart}, iter {iteration}: {covered}/{TOTAL_PAIRS}")
                    
                    if covered == TOTAL_PAIRS:
                        print(f"  FOUND! Schedule: {sched}")
                        return sched
            else:
                no_improve += 1
                # Random restart within this run if stuck
                if no_improve > 5000:
                    sched = random_schedule()
                    covered = count_covered_pairs(sched)
                    no_improve = 0
        
        print(f"Restart {restart} done: best = {best_local}/{TOTAL_PAIRS}")
        if best_local > best_overall:
            best_overall = best_local
            best_schedule_overall = best_local_sched
    
    return best_schedule_overall

def exhaustive_search_one_day(fixed_days, max_trials=100000):
    """Given 4 fixed days, try many random 5th days."""
    covered_so_far = set()
    for day in fixed_days:
        for boat in day:
            covered_so_far.update(itertools.combinations(boat, 2))
    
    unaccounted = set(itertools.combinations(range(N), 2)) - covered_so_far
    print(f"Need to cover {len(unaccounted)} more pairs with day 5")
    
    best = 0
    best_day = None
    
    for trial in range(max_trials):
        day5 = random_partition()
        new_pairs = 0
        for boat in day5:
            for p in itertools.combinations(boat, 2):
                if p in unaccounted:
                    new_pairs += 1
        
        if new_pairs > best:
            best = new_pairs
            best_day = day5
            print(f"  Trial {trial}: covers {new_pairs}/{len(unaccounted)} missing pairs")
            if new_pairs == len(unaccounted):
                return day5
    
    return best_day, best, unaccounted

if __name__ == "__main__":
    print(f"Total pairs to cover: {TOTAL_PAIRS}")
    print(f"Searching for 5-day schedule...\n")
    
    # First try local search
    result = local_search(max_iterations=50000, num_restarts=10)
    
    if result:
        if count_covered_pairs(result) == TOTAL_PAIRS:
            print(f"\nFOUND 5-DAY SOLUTION!")
            for i, day in enumerate(result):
                print(f"  Day {i+1}: {day}")
        else:
            print(f"\nBest found: {count_covered_pairs(result)}/{TOTAL_PAIRS}")
    else:
        print("No solution found in local search.")