# Complete DOCPLEX implementation

from docplex.mp.model import Model

def optimize_advisor_allocation():
    """Optimize the allocation of student advisors to students with pets."""
    
    # 1. MODEL & DATA SETUP
    mdl = Model(name="pets_1")
    
    # Data from the database
    students = [1, 2, 3]
    advisors = [201, 202, 203]
    max_students = {201: 5, 202: 6, 203: 7}
    max_weights = {201: 100.0, 202: 120.0, 203: 150.0}
    initial_weights = {201: 50.0, 202: 60.0, 203: 70.0}
    
    # Validate array lengths
    assert len(students) == 3, "Mismatch in number of students"
    assert len(advisors) == 3, "Mismatch in number of advisors"
    
    # Safe range for indexing
    safe_range_students = range(len(students))
    safe_range_advisors = range(len(advisors))
    
    # 2. VARIABLES
    x = {(i, j): mdl.binary_var(name=f"x_{i}_{j}") for i in safe_range_students for j in safe_range_advisors}
    w = {j: mdl.continuous_var(name=f"w_{j}", lb=0) for j in safe_range_advisors}
    M = mdl.continuous_var(name="M", lb=0)
    
    # 3. OBJECTIVE FUNCTION
    mdl.minimize(M)
    
    # 4. CONSTRAINTS
    
    # Maximum workload constraint
    for j in safe_range_advisors:
        mdl.add_constraint(M >= mdl.sum(x[i, j] for i in safe_range_students) + w[j], ctname=f"max_workload_{j}")
    
    # Student assignment constraint
    for i in safe_range_students:
        mdl.add_constraint(mdl.sum(x[i, j] for j in safe_range_advisors) == 1, ctname=f"student_assignment_{i}")
    
    # Advisor capacity constraints
    for j in safe_range_advisors:
        mdl.add_constraint(mdl.sum(x[i, j] for i in safe_range_students) <= max_students[advisors[j]], ctname=f"max_students_{j}")
        mdl.add_constraint(w[j] <= max_weights[advisors[j]], ctname=f"max_weight_{j}")
    
    # Initial pet weight assignment
    for j in safe_range_advisors:
        mdl.add_constraint(w[j] == initial_weights[advisors[j]], ctname=f"initial_weight_{j}")
    
    # 5. SOLVING & RESULTS
    solution = mdl.solve()
    
    if solution:
        print(f"Optimal value: {solution.objective_value}")
        for i in safe_range_students:
            for j in safe_range_advisors:
                if solution.get_value(x[i, j]) > 1e-6:
                    print(f"Student {students[i]} assigned to Advisor {advisors[j]}")
    else:
        print("No solution found")
        print(f"Status: {mdl.solve_details.status}")

    return mdl

# Run the optimization
optimize_advisor_allocation()