Disney schedule solver

solve.py

from ortools.sat.python import cp_model

# --- Model setup (same as before) ---
model = cp_model.CpModel()
num_slots = 12
slots = list(range(num_slots))
park_names = ['MK', 'EPCOT', 'HS', 'AK', 'None']
crowds = [
    [5,3,4,4,11], # Sunday
    [4,3,4,6,11], # Monday
    [2,2,4,3,11], # Tuesday
    [3,2,3,3,11], # Wednesday
    [6,3,4,3,11], # Thursday
    [2,3,4,4,11], # Friday
]
slot_vars = [model.NewIntVar(0,4,f'slot_{i}') for i in slots]
HOP_WEIGHT = 2


# --- Flexible slot counts ---
def add_slot_range(model, slot_vars, park_index, min_count, max_count, park_name):
    indicators = []
    for i, var in enumerate(slot_vars):
        ind = model.NewBoolVar(f'{park_name}_slot_{i}')
        model.Add(var == park_index).OnlyEnforceIf(ind)
        model.Add(var != park_index).OnlyEnforceIf(ind.Not())
        indicators.append(ind)
    model.Add(sum(indicators) >= min_count)
    model.Add(sum(indicators) <= max_count)

add_slot_range(model, slot_vars, 0, 3, 4, 'MK')
add_slot_range(model, slot_vars, 1, 3, 4, 'EPCOT')
add_slot_range(model, slot_vars, 2, 2, 2, 'HS')
add_slot_range(model, slot_vars, 3, 1, 1, 'AK')

# --- AM constraints ---
am_slots = [i for i in range(num_slots) if i % 2 == 0]
for park_index, park_name in enumerate(park_names[:-1]):
    am_indicators = []
    for slot in am_slots:
        ind = model.NewBoolVar(f'{park_name}_am_{slot}')
        model.Add(slot_vars[slot] == park_index).OnlyEnforceIf(ind)
        model.Add(slot_vars[slot] != park_index).OnlyEnforceIf(ind.Not())
        am_indicators.append(ind)
    model.Add(sum(am_indicators) >= 1)

# --- PM constraints ---
pm_slots = [i for i in range(num_slots) if i % 2 == 1]
evening_min = {'MK': 2, 'EPCOT': 1, 'HS': 1}
for park_name, min_count in evening_min.items():
    park_index = park_names.index(park_name)
    pm_indicators = []
    for slot in pm_slots:
        ind = model.NewBoolVar(f'{park_name}_pm_{slot}')
        model.Add(slot_vars[slot] == park_index).OnlyEnforceIf(ind)
        model.Add(slot_vars[slot] != park_index).OnlyEnforceIf(ind.Not())
        pm_indicators.append(ind)
    model.Add(sum(pm_indicators) >= min_count)

# --- Special events and restrictions ---
model.Add(slot_vars[2] == 4)  # Monday AM no park
model.Add(slot_vars[3] == 2) 
model.Add(slot_vars[4] == 4)  # Tuesday AM no park
model.Add(slot_vars[5] == 0) 
model.Add(slot_vars[6] == 0)
model.Add(slot_vars[11] != 0) # Friday PM cannot be MK

# EPCOT must happen by Wednesday
epcot_index = park_names.index('EPCOT')
epcot_slots = [slot_vars[i] for i in range(0, 8)]
epcot_indicators = []
for i, var in enumerate(epcot_slots):
    ind = model.NewBoolVar(f'epcot_by_wed_slot_{i}')
    model.Add(var == epcot_index).OnlyEnforceIf(ind)
    model.Add(var != epcot_index).OnlyEnforceIf(ind.Not())
    epcot_indicators.append(ind)
model.Add(sum(epcot_indicators) >= 1)

# --- Objective ---
crowd_terms = []
hop_terms = []
for day in range(6):
    am_slot = day*2
    pm_slot = am_slot + 1
    for slot in [am_slot, pm_slot]:
        crowd_expr = model.NewIntVar(0, 12, f'crowd_{slot}')
        model.AddElement(slot_vars[slot], crowds[day], crowd_expr)
        crowd_terms.append(crowd_expr)
    hop = model.NewBoolVar(f'hop_day_{day}')
    model.Add(slot_vars[am_slot] != slot_vars[pm_slot]).OnlyEnforceIf(hop)
    model.Add(slot_vars[am_slot] == slot_vars[pm_slot]).OnlyEnforceIf(hop.Not())
    hop_terms.append(hop * HOP_WEIGHT)
model.Minimize(sum(crowd_terms) + sum(hop_terms))

# --- Pretty print ---
def print_schedule(solution, park_names):
    for day in range(6):
        am_slot = day*2
        pm_slot = am_slot + 1
        crowd_am = crowds[day][solution[am_slot]]
        crowd_pm = crowds[day][solution[pm_slot]]
        day_name = ["Sunday","Monday","Tuesday","Wednesday","Thursday","Friday"][day]
        print(f"{day_name}: Morning: {park_names[solution[am_slot]]} crowd=({crowd_am if crowd_am < 10 else 0}/10), Evening: {park_names[solution[pm_slot]]} crowd=({crowd_pm}/10)")

# --- Solution collector ---
class TopSolutionCollector(cp_model.CpSolverSolutionCallback):
    def __init__(self, slot_vars, park_names, max_solutions=3):
        super().__init__()
        self.slot_vars = slot_vars
        self.park_names = park_names
        self.solutions = []
        self.max_solutions = max_solutions

    def on_solution_callback(self):
        if len(self.solutions) >= self.max_solutions:
            self.StopSearch()
            return
        assign = [self.Value(var) for var in self.slot_vars]
        self.solutions.append({
            "objective": self.ObjectiveValue(),
            "assign": assign
        })

# --- Solve ---
solver = cp_model.CpSolver()
collector = TopSolutionCollector(slot_vars, park_names, max_solutions=3)
solver.parameters.enumerate_all_solutions = True
solver.Solve(model, collector)

# Pretty print the first, second, and third solutions
for i, sol in enumerate(collector.solutions):
    print(f"\nSolution {i+1} - Objective: {sol['objective']}")
    print_schedule(sol['assign'], park_names)