#!/usr/bin/env python3
"""
Gurobipy Implementation for Restaurant Allocation Optimization
"""

import gurobipy as gp
from gurobipy import GRB

def optimize_restaurant_allocation():
    """Optimize student allocation to restaurants to maximize satisfaction within budget."""
    
    # 1. MODEL & DATA SETUP
    model = gp.Model("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
    
    # CRITICAL: Validate array lengths before loops
    assert len(students) > 0, "No students provided"
    assert len(restaurants) > 0, "No restaurants provided"
    assert len(satisfaction_scores) > 0, "No satisfaction scores provided"
    assert len(restaurant_capacities) > 0, "No restaurant capacities provided"
    assert len(visits_restaurant) > 0, "No visits data provided"
    
    # 2. VARIABLES
    # Binary decision variables: x[i,j] = 1 if student i is assigned to restaurant j
    x = model.addVars(students, restaurants, vtype=GRB.BINARY, name="x")
    
    # 3. OBJECTIVE FUNCTION
    # Maximize total student satisfaction
    model.setObjective(
        gp.quicksum(satisfaction_scores.get((i, j), 0) * x[i, j] 
                    for i in students for j in restaurants), 
        GRB.MAXIMIZE
    )
    
    # 4. CONSTRAINTS
    
    # Total Spending Constraint
    model.addConstr(
        gp.quicksum(visits_restaurant.get((i, j), 0) * satisfaction_scores.get((i, j), 0) * x[i, j] 
                    for i in students for j in restaurants) <= budget, 
        name="total_spending"
    )
    
    # Restaurant Capacity Constraint
    for j in restaurants:
        model.addConstr(
            gp.quicksum(x[i, j] for i in students) <= restaurant_capacities[j], 
            name=f"capacity_{j}"
        )
    
    # Student Assignment Constraint
    for i in students:
        model.addConstr(
            gp.quicksum(x[i, j] for j in restaurants) == 1, 
            name=f"assignment_{i}"
        )
    
    # 5. SOLVING & RESULTS
    model.optimize()
    
    if model.status == GRB.OPTIMAL:
        print(f"Optimal value: {model.objVal}")
        for i in students:
            for j in restaurants:
                if x[i, j].x > 1e-6:
                    print(f"Student {i} assigned to Restaurant {j}")
    elif model.status == GRB.INFEASIBLE:
        print("Problem is infeasible")
    elif model.status == GRB.UNBOUNDED:
        print("Problem is unbounded")
    
    return model

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