# Complete DOCPLEX implementation - Retry Attempt 4

from docplex.mp.model import Model

def school_bus_optimization():
    # 1. MODEL & DATA SETUP
    mdl = Model(name="school_bus")
    
    # Data from the problem context
    drivers = [1, 2, 3]
    schools = [101, 102, 103]
    is_full_time = {1: True, 2: True, 3: True}
    required_drivers = {101: 3, 102: 4, 103: 5}
    years_mismatch = {
        (1, 101): 2, (1, 102): 0, (1, 103): 0,
        (2, 101): 0, (2, 102): 1, (2, 103): 0,
        (3, 101): 0, (3, 102): 0, (3, 103): 0
    }
    
    # Validate data lengths
    assert len(drivers) == len(is_full_time)
    assert len(schools) == len(required_drivers)
    
    # 2. VARIABLES
    x = {(i, j): mdl.binary_var(name=f"x_{i}_{j}") for i in drivers for j in schools}
    
    # 3. OBJECTIVE FUNCTION
    objective = mdl.sum(years_mismatch[i, j] * x[i, j] for i in drivers for j in schools)
    mdl.minimize(objective)
    
    # 4. CONSTRAINTS
    # Each driver can be assigned to at most one school
    for i in drivers:
        if is_full_time[i]:
            mdl.add_constraint(mdl.sum(x[i, j] for j in schools) <= 1, ctname=f"driver_{i}_assignment")
    
    # Each school must receive the exact number of drivers it requires
    for j in schools:
        mdl.add_constraint(mdl.sum(x[i, j] for i in drivers) == required_drivers[j], ctname=f"school_{j}_requirement")
    
    # 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:  # Binary variable threshold
                    print(f"Driver {i} assigned to School {j}")
    else:
        print("No solution found")
        print(f"Status: {mdl.solve_details.status}")

    return mdl

# Run the optimization
school_bus_optimization()