# Complete DOCPLEX implementation

from docplex.mp.model import Model

def railway_optimization():
    """Optimize manager to railway assignments to minimize total cost."""
    
    # 1. MODEL & DATA SETUP
    mdl = Model(name="railway_management")
    
    # Data from the database
    managers = [1, 2, 3]
    railways = [1, 2, 3]
    
    # Manager costs and capacities
    manager_costs = {1: 120, 2: 150, 3: 200}
    manager_capacities = {1: 4, 2: 6, 3: 8}
    
    # CRITICAL: Validate array lengths to prevent IndexError
    assert len(managers) == len(manager_costs) == len(manager_capacities), "Array length mismatch"
    
    # 2. VARIABLES
    # Binary decision variables: x[m][r] = 1 if manager m is assigned to railway r
    x = {(m, r): mdl.binary_var(name=f"x_{m}_{r}") for m in managers for r in railways}
    
    # 3. OBJECTIVE FUNCTION
    # Minimize the total cost of assigning managers to railways
    objective = mdl.sum(manager_costs[m] * x[(m, r)] for m in managers for r in railways)
    mdl.minimize(objective)
    
    # 4. CONSTRAINTS
    
    # Railway Management Constraint: Each railway must be managed by at least one manager
    for r in railways:
        mdl.add_constraint(mdl.sum(x[(m, r)] for m in managers) >= 1, ctname=f"railway_{r}_managed")
    
    # Manager Capacity Constraint: No manager can be assigned to more railways than their capacity
    for m in managers:
        mdl.add_constraint(mdl.sum(x[(m, r)] for r in railways) <= manager_capacities[m], ctname=f"manager_{m}_capacity")
    
    # 5. SOLVING & RESULTS
    solution = mdl.solve()
    
    if solution:
        print(f"Optimal value: {solution.objective_value}")
        for m in managers:
            for r in railways:
                if solution.get_value(x[(m, r)]) > 0:
                    print(f"Manager {m} is assigned to Railway {r}")
    else:
        print("No solution found")
        print(f"Status: {mdl.solve_details.status}")
    
    return mdl

# Run the optimization
railway_optimization()