# Complete DOCPLEX implementation

from docplex.mp.model import Model

def optimize_student_dormitory_assignment():
    # 1. MODEL & DATA SETUP
    mdl = Model(name="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 to prevent IndexError
    assert len(dormitory_capacity) == len(dormitory_allergy_friendly) == len(dormitories), "Array length mismatch"
    
    # 2. VARIABLES
    x = {(i, j): mdl.binary_var(name=f"x_{i}_{j}") for i in students for j in dormitories}
    
    # 3. OBJECTIVE FUNCTION
    objective = mdl.sum(allergy_penalty.get((i, j), 0) * x[i, j] for i in students for j in dormitories)
    mdl.minimize(objective)
    
    # 4. CONSTRAINTS
    
    # Each student must be assigned to exactly one dormitory
    for i in students:
        mdl.add_constraint(mdl.sum(x[i, j] for j in dormitories) == 1, ctname=f"assign_{i}")
    
    # The number of students assigned to a dormitory cannot exceed its capacity
    for j in dormitories:
        mdl.add_constraint(mdl.sum(x[i, j] for i in students) <= dormitory_capacity[j], ctname=f"capacity_{j}")
    
    # Students with allergies cannot be assigned to non-allergy-friendly dormitories
    for i in students:
        for j in dormitories:
            if not dormitory_allergy_friendly[j]:
                mdl.add_constraint(x[i, j] == 0, ctname=f"allergy_{i}_{j}")
    
    # 5. SOLVING & RESULTS
    solution = mdl.solve()
    
    if solution:
        print(f"Optimal value: {solution.objective_value}")
        for i in students:
            for j in dormitories:
                if solution.get_value(x[i, j]) > 1e-6:
                    print(f"Student {i} assigned to Dormitory {j}")
    else:
        print("No solution found")
        print(f"Status: {mdl.solve_details.status}")

    return mdl

# Run the optimization
optimize_student_dormitory_assignment()