#!/usr/bin/env python3
"""
Gurobipy Implementation for Student Classroom Assignment Problem
"""

import gurobipy as gp
from gurobipy import GRB

def optimize_student_assignment():
    """Optimize student assignment to classrooms to minimize maximum class size."""
    
    # 1. MODEL & DATA SETUP
    model = gp.Model("student_assignment")
    
    # Example data (replace with actual data loading)
    students = [101, 102, 103, 104, 105]
    classrooms = [1, 2, 3]
    classroom_capacities = {1: 25, 2: 30, 3: 35}
    max_students_per_classroom = {1: 25, 2: 30, 3: 35}
    
    # CRITICAL: Validate array lengths before loops
    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): model.addVar(vtype=GRB.BINARY, name=f"x_{s}_{c}") 
         for s in students for c in classrooms}
    
    # Continuous variable for maximum number of students in any classroom
    z = model.addVar(vtype=GRB.CONTINUOUS, name="z", lb=0)
    
    # 3. OBJECTIVE FUNCTION
    model.setObjective(z, GRB.MINIMIZE)
    
    # 4. CONSTRAINTS
    
    # Student Assignment Constraint: Each student must be assigned to exactly one classroom
    for s in students:
        model.addConstr(gp.quicksum(x[s, c] for c in classrooms) == 1, name=f"student_assignment_{s}")
    
    # Classroom Capacity Constraint: Total students in each classroom must not exceed its capacity
    for c in classrooms:
        model.addConstr(gp.quicksum(x[s, c] for s in students) <= classroom_capacities[c], name=f"capacity_{c}")
    
    # Maximum Students Constraint: Total students in any classroom must be <= z
    for c in classrooms:
        model.addConstr(gp.quicksum(x[s, c] for s in students) <= z, name=f"max_students_{c}")
    
    # 5. SOLVING & RESULTS
    model.optimize()
    
    if model.status == GRB.OPTIMAL:
        print(f"Optimal value: {model.objVal}")
        for s in students:
            for c in classrooms:
                if x[s, c].x > 0.5:  # Check if assignment is True
                    print(f"Student {s} is assigned to Classroom {c}")
    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_student_assignment()