#!/usr/bin/env python3
"""
DOCplex implementation for school bus driver assignment optimization
"""

from docplex.mp.model import Model

def school_bus_optimization():
    """Optimize driver-school assignments to minimize total travel distance"""
    
    # 1. MODEL & DATA SETUP
    mdl = Model(name="school_bus_assignment")
    
    # Data from the distance matrix
    drivers = [1, 2, 3]
    schools = [1, 2, 3]
    distances = {
        (1, 1): 10.5,
        (1, 2): 15.3,
        (1, 3): 20.1,
        (2, 1): 18.2,
        (2, 2): 8.7,
        (2, 3): 22.4,
        (3, 1): 19.8,
        (3, 2): 21.3,
        (3, 3): 7.5
    }
    
    # CRITICAL: Validate data consistency
    assert len(drivers) == 3, "Expected 3 drivers"
    assert len(schools) == 3, "Expected 3 schools"
    assert len(distances) == 9, "Expected 9 distance entries"
    
    # 2. VARIABLES
    # Binary decision variables for assignments
    x = {(i, j): mdl.binary_var(name=f"x_{i}_{j}") for i in drivers for j in schools}
    
    # 3. OBJECTIVE FUNCTION
    # Minimize total travel distance
    objective = mdl.sum(distances[(i, j)] * x[(i, j)] for i in drivers for j in schools)
    mdl.minimize(objective)
    
    # 4. CONSTRAINTS
    
    # Each school must be assigned exactly one driver
    for j in schools:
        mdl.add_constraint(mdl.sum(x[(i, j)] for i in drivers) == 1, ctname=f"school_{j}_assignment")
    
    # Each driver can be assigned to at most one school
    for i in drivers:
        mdl.add_constraint(mdl.sum(x[(i, j)] for j in schools) <= 1, ctname=f"driver_{i}_limit")
    
    # 5. SOLVING & RESULTS
    solution = mdl.solve()
    
    if solution:
        print(f"Optimal value: {solution.objective_value}")
        for i in drivers:
            for j in schools:
                if solution.get_value(x[(i, j)]) > 0.5:
                    print(f"Driver {i} assigned to School {j}")
    else:
        print("No solution found")
        print(f"Status: {mdl.solve_details.status}")
    
    return mdl

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