#!/usr/bin/env python3
"""
Gurobipy Implementation for School Bus Driver Assignment Problem
"""

import gurobipy as gp
from gurobipy import GRB

def optimize_driver_assignment():
    """Optimize driver to school assignments to minimize total travel distance."""
    
    # 1. MODEL & DATA SETUP
    model = gp.Model("school_bus_assignment")
    
    # Data from the problem
    drivers = [1, 2, 3]
    schools = [1, 2, 3]
    
    # Distance matrix: driver_id -> school_id -> distance
    distance_matrix = {
        1: {1: 10.5, 2: 15.3, 3: 20.1},
        2: {1: 18.2, 2: 8.7, 3: 22.4},
        3: {1: 19.8, 2: 21.3, 3: 7.5}
    }
    
    # CRITICAL: Validate array lengths before loops
    assert len(drivers) == len(distance_matrix), "Driver count mismatch in distance matrix"
    for driver in drivers:
        assert len(schools) == len(distance_matrix[driver]), "School count mismatch in distance matrix"
    
    # 2. VARIABLES
    # Binary decision variables: x[driver_id][school_id]
    x = model.addVars(drivers, schools, vtype=GRB.BINARY, name="x")
    
    # 3. OBJECTIVE FUNCTION
    # Minimize total travel distance
    model.setObjective(
        gp.quicksum(distance_matrix[i][j] * x[i, j] for i in drivers for j in schools),
        GRB.MINIMIZE
    )
    
    # 4. CONSTRAINTS
    
    # Each school must be assigned exactly one driver
    for j in schools:
        model.addConstr(
            gp.quicksum(x[i, j] for i in drivers) == 1,
            name=f"school_{j}_assignment"
        )
    
    # Each driver can be assigned to at most one school
    for i in drivers:
        model.addConstr(
            gp.quicksum(x[i, j] for j in schools) <= 1,
            name=f"driver_{i}_limit"
        )
    
    # 5. SOLVING & RESULTS
    model.optimize()
    
    if model.status == GRB.OPTIMAL:
        print(f"Optimal value: {model.objVal}")
        for i in drivers:
            for j in schools:
                if x[i, j].x > 0.5:  # Check if assignment is made
                    print(f"Driver {i} assigned to School {j} with distance {distance_matrix[i][j]}")
    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_driver_assignment()