#!/usr/bin/env python3
"""
DOCplex implementation for restaurant allocation optimization problem
"""

from docplex.mp.model import Model

def restaurant_allocation_optimization():
    """Optimize student allocation to restaurants to maximize satisfaction within budget"""
    
    # 1. MODEL & DATA SETUP
    mdl = Model(name="restaurant_allocation")
    
    # Sample data from the problem description
    students = [1, 2, 3]
    restaurants = [1, 2, 3]
    
    # Satisfaction scores (StuID, ResID, score)
    satisfaction_scores = {
        (1, 1): 4.5,
        (1, 2): 3.0,
        (2, 1): 3.8
    }
    
    # Restaurant capacities (ResID, capacity)
    restaurant_capacities = {
        1: 50,
        2: 75,
        3: 100
    }
    
    # Visits_Restaurant (StuID, ResID, visited)
    visits_restaurant = {
        (1, 1): True,
        (2, 2): True,
        (3, 3): True
    }
    
    # Budget
    budget = 15000
    
    # 2. VARIABLES
    # Binary decision variables: x[i,j] = 1 if student i is assigned to restaurant j
    x = {(i, j): mdl.binary_var(name=f"x_{i}_{j}") for i in students for j in restaurants}
    
    # 3. OBJECTIVE FUNCTION
    # Maximize total student satisfaction
    objective = mdl.sum(satisfaction_scores.get((i, j), 0) * x[i, j] for i in students for j in restaurants)
    mdl.maximize(objective)
    
    # 4. CONSTRAINTS
    
    # Total Spending Constraint
    total_spending = mdl.sum(satisfaction_scores.get((i, j), 0) * x[i, j] for i in students for j in restaurants)
    mdl.add_constraint(total_spending <= budget, ctname="total_spending")
    
    # Restaurant Capacity Constraint
    for j in restaurants:
        capacity = restaurant_capacities.get(j, 0)
        mdl.add_constraint(mdl.sum(x[i, j] for i in students) <= capacity, ctname=f"capacity_{j}")
    
    # Student Assignment Constraint
    for i in students:
        mdl.add_constraint(mdl.sum(x[i, j] for j in restaurants) == 1, ctname=f"assignment_{i}")
    
    # 5. SOLVING & RESULTS
    solution = mdl.solve()
    
    if solution:
        print(f"Optimal value: {solution.objective_value}")
        for i in students:
            for j in restaurants:
                if solution.get_value(x[i, j]) > 0.5:
                    print(f"Student {i} assigned to Restaurant {j}")
    else:
        print("No solution found")
        print(f"Status: {mdl.solve_details.status}")
    
    return mdl

# Execute the optimization
if __name__ == "__main__":
    restaurant_allocation_optimization()