#!/usr/bin/env python3
"""
DOCplex implementation for student classroom assignment optimization
"""

from docplex.mp.model import Model

def optimize_student_assignments():
    """Optimize student assignments to classrooms to minimize maximum class size"""
    
    # 1. MODEL & DATA SETUP
    mdl = Model(name="student_classroom_assignment")
    
    # Data from the problem context
    classrooms = [1, 2, 3]
    students = [101, 102, 103, 104, 105]
    
    classroom_capacities = {1: 25, 2: 30, 3: 35}
    max_students_per_classroom = {1: 25, 2: 30, 3: 35}
    
    # CRITICAL: Validate array lengths to prevent IndexError
    assert len(classrooms) == len(classroom_capacities) == len(max_students_per_classroom), "Array length mismatch"
    
    # 2. VARIABLES
    # Binary decision variables for student assignments
    x = {(s, c): mdl.binary_var(name=f"x_{s}_{c}") for s in students for c in classrooms}
    
    # Continuous variable for maximum number of students in any classroom
    z = mdl.continuous_var(name="z", lb=0)
    
    # 3. OBJECTIVE FUNCTION
    # Minimize the maximum number of students in any classroom
    mdl.minimize(z)
    
    # 4. CONSTRAINTS
    
    # Student Assignment Constraint: Each student must be assigned to exactly one classroom
    for s in students:
        mdl.add_constraint(mdl.sum(x[s, c] for c in classrooms) == 1, ctname=f"student_assignment_{s}")
    
    # Classroom Capacity Constraint: Total students in each classroom must not exceed its capacity
    for c in classrooms:
        mdl.add_constraint(mdl.sum(x[s, c] for s in students) <= classroom_capacities[c], ctname=f"capacity_{c}")
    
    # Maximum Students Constraint: Total students in any classroom must be <= z
    for c in classrooms:
        mdl.add_constraint(mdl.sum(x[s, c] for s in students) <= z, ctname=f"max_students_{c}")
    
    # 5. SOLVING & RESULTS
    solution = mdl.solve()
    
    if solution:
        print(f"Optimal value: {solution.objective_value}")
        for s in students:
            for c in classrooms:
                if solution.get_value(x[s, c]) > 0.5:
                    print(f"Student {s} is assigned to Classroom {c}")
    else:
        print("No solution found")
        print(f"Status: {mdl.solve_details.status}")
    
    return mdl

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