# Complete GUROBIPY implementation

import gurobipy as gp
from gurobipy import GRB

def optimize_dormitory_assignment():
    """Optimize student dormitory assignments to minimize allergy penalties."""
    
    # 1. MODEL & DATA SETUP
    model = gp.Model("allergy_1")
    
    # Data from the database schema
    dormitory_capacity = {1: 50, 2: 100, 3: 150}
    dormitory_allergy_friendly = {1: True, 2: False, 3: True}
    allergy_penalty = {
        (101, 2): 20.0,
        (102, 2): 25.0,
        (103, 2): 15.0
    }
    
    students = [101, 102, 103]
    dormitories = [1, 2, 3]
    
    # CRITICAL: Validate array lengths before loops
    assert len(dormitory_capacity) == len(dormitory_allergy_friendly) == len(dormitories), "Array length mismatch"
    
    # 2. VARIABLES
    # Binary decision variables for student-dormitory assignments
    x = model.addVars(students, dormitories, vtype=GRB.BINARY, name="x")
    
    # 3. OBJECTIVE FUNCTION
    # Minimize the total penalty for assigning students with allergies to non-allergy-friendly dormitories
    model.setObjective(
        gp.quicksum(allergy_penalty.get((i, j), 0) * x[i, j] for i in students for j in dormitories),
        GRB.MINIMIZE
    )
    
    # 4. CONSTRAINTS - CORRECT SYNTAX PATTERNS
    
    # Each student must be assigned to exactly one dormitory
    model.addConstrs(
        (gp.quicksum(x[i, j] for j in dormitories) == 1 for i in students),
        name="student_assignment"
    )
    
    # The number of students assigned to a dormitory cannot exceed its capacity
    model.addConstrs(
        (gp.quicksum(x[i, j] for i in students) <= dormitory_capacity[j] for j in dormitories),
        name="dormitory_capacity"
    )
    
    # Students with allergies cannot be assigned to non-allergy-friendly dormitories
    model.addConstrs(
        (x[i, j] == 0 for i in students for j in dormitories if not dormitory_allergy_friendly[j] and (i, j) in allergy_penalty),
        name="allergy_constraint"
    )
    
    # 5. SOLVING & RESULTS
    model.optimize()
    
    if model.status == GRB.OPTIMAL:
        print(f"Optimal value: {model.objVal}")
        for i in students:
            for j in dormitories:
                if x[i, j].x > 1e-6:
                    print(f"Student {i} assigned to Dormitory {j}")
    elif model.status == GRB.INFEASIBLE:
        print("Problem is infeasible")
    elif model.status == GRB.UNBOUNDED:
        print("Problem is unbounded")
    
    return model

# Run the optimization
optimize_dormitory_assignment()