"""
Fixed search - normalize all pairs as (min, max) to avoid double counting.
"""
import itertools
import math
import random as rnd

N = 24
TOTAL_PAIRS = math.comb(N, 2)

def normalize_pair(p, q):
    return (min(p, q), max(p, q))

def count_pairs(sched):
    """Count distinct pairs covered by the schedule."""
    pairs = set()
    for day in sched:
        for boat in day:
            for p, q in itertools.combinations(boat, 2):
                pairs.add(normalize_pair(p, q))
    return len(pairs)

def missing_pairs_set(sched):
    pairs = set()
    for day in sched:
        for boat in day:
            for p, q in itertools.combinations(boat, 2):
                pairs.add(normalize_pair(p, q))
    all_pairs = set(normalize_pair(i, j) for i in range(N) for j in range(i+1, N))
    return all_pairs - pairs

# Check that the simple affine schedule actually covers 240
GF5 = [0, 1, 2, 3, 4]
all_points = [(x, y) for x in GF5 for y in GF5]
removed = (4, 4)
points = [p for p in all_points if p != removed]
pt_to_idx = {pt: i for i, pt in enumerate(points)}

def get_parallel_classes():
    classes = []
    for slope in range(6):
        klass = []
        if slope <= 4:
            for b in GF5:
                klass.append([(x, (slope*x + b) % 5) for x in GF5])
        else:
            for c in GF5:
                klass.append([(c, y) for y in GF5])
        classes.append(klass)
    return classes

classes = get_parallel_classes()

adapted_classes = []
for klass in classes:
    filtered_lines = []
    for line in klass:
        filtered = [p for p in line if p != removed]
        filtered_lines.append(filtered)
    size4_line = [l for l in filtered_lines if len(l) == 4][0]
    size5_lines = [l for l in filtered_lines if len(l) == 5]
    adapted_classes.append((size4_line, size5_lines))

def verify_partition(boats):
    all_elems = set()
    for boat in boats:
        if len(boat) != 6:
            return False
        for x in boat:
            if x < 0 or x >= N:
                return False
            all_elems.add(x)
    return len(all_elems) == N

def simple_partition(s4, s5):
    a, b, c, d = [pt_to_idx[p] for p in s4]
    s5_i = [[pt_to_idx[p] for p in line] for line in s5]
    return [
        tuple(sorted(s5_i[0] + [a])),
        tuple(sorted(s5_i[1] + [b])),
        tuple(sorted(s5_i[2] + [c])),
        tuple(sorted(s5_i[3] + [d])),
    ]

# Build simple schedule (sorted boats)
simple_sched = [simple_partition(s4, s5) for s4, s5 in adapted_classes]
print(f"Simple schedule: {count_pairs(simple_sched)}/{TOTAL_PAIRS}")
print(f"Missing: {len(missing_pairs_set(simple_sched))}")

# Strategy 1: (a,b) together, (c,d) together
def strat1_partition(s4, s5, p_idx=0, q_idx=0):
    a, b, c, d = [pt_to_idx[p] for p in s4]
    s5_i = [[pt_to_idx[p] for p in line] for line in s5]
    p = s5_i[0][p_idx]
    q = s5_i[1][q_idx]
    return [
        tuple(sorted([x for x in s5_i[0] if x != p] + [a, b])),
        tuple(sorted([x for x in s5_i[1] if x != q] + [c, d])),
        tuple(sorted(s5_i[2] + [p])),
        tuple(sorted(s5_i[3] + [q])),
    ]

# Strategy 3: all 4 together
def strat3_partition(s4, s5, si=0, pi=0, qi=1):
    a, b, c, d = [pt_to_idx[p] for p in s4]
    s5_i = [[pt_to_idx[p] for p in line] for line in s5]
    p = s5_i[si][pi]
    q = s5_i[si][qi]
    remaining = [x for x in s5_i[si] if x != p and x != q]
    other = [j for j in range(4) if j != si]
    return [
        tuple(sorted([a, b, c, d, p, q])),
        tuple(sorted(s5_i[other[0]] + [remaining[0]])),
        tuple(sorted(s5_i[other[1]] + [remaining[1]])),
        tuple(sorted(s5_i[other[2]] + [remaining[2]])),
    ]

# Try random combinations of strategies
print("\n=== Random search ===")
best_cov = 0
best_sched = None

for trial in range(100000):
    schedule = []
    for ci, (s4, s5) in enumerate(adapted_classes):
        strategy = rnd.choice(['simple', 's1', 's3'])
        if strategy == 'simple':
            day = simple_partition(s4, s5)
        elif strategy == 's1':
            pi = rnd.randrange(5)
            qi = rnd.randrange(5)
            day = strat1_partition(s4, s5, pi, qi)
        else:
            si = rnd.randrange(4)
            pi, qi = rnd.sample(range(5), 2)
            day = strat3_partition(s4, s5, si, pi, qi)
        schedule.append(day)
    
    if not all(verify_partition(d) for d in schedule):
        continue
    
    cov = count_pairs(schedule)
    if cov > best_cov:
        best_cov = cov
        best_sched = [list(d) for d in schedule]
        print(f"  Trial {trial}: {best_cov}/{TOTAL_PAIRS}")
        if cov == TOTAL_PAIRS:
            print("FOUND!")
            break

print(f"\nBest: {best_cov}/{TOTAL_PAIRS}")

# Local search
if best_sched:
    print("\n=== Local search ===")
    current = [list(d) for d in best_sched]
    current_cov = count_pairs(current)
    best = current_cov
    
    for it in range(100000):
        d = rnd.randrange(6)
        day = [list(b) for b in current[d]]
        b1, b2 = rnd.sample(range(4), 2)
        p1 = rnd.randrange(6)
        p2 = rnd.randrange(6)
        day[b1][p1], day[b2][p2] = day[b2][p2], day[b1][p1]
        new_day = [tuple(sorted(b)) for b in day]
        
        if not verify_partition(new_day):
            continue
        
        new_sched = current[:d] + [new_day] + current[d+1:]
        cov = count_pairs(new_sched)
        
        if cov >= current_cov:
            current = new_sched
            current_cov = cov
            if cov > best:
                best = cov
                print(f"  Iter {it}: {best}/{TOTAL_PAIRS}")
                if best == TOTAL_PAIRS:
                    print(f"\nFOUND!")
                    best_sched = current
                    import json
                    print(f"\nJSON:\n{json.dumps(best_sched)}")
                    raise SystemExit(0)
        
        # Restart
        if it > 0 and it % 20000 == 0:
            current = [list(d) for d in best_sched]
            current_cov = count_pairs(current)
            print(f"  Restart at {it}: {current_cov}/{TOTAL_PAIRS}")

print(f"Best: {best}/{TOTAL_PAIRS}")
if best == TOTAL_PAIRS:
    import json
    print(json.dumps(best_sched))