#!/usr/bin/env python3
"""
Gurobipy Implementation for Railway Management Optimization
"""

import gurobipy as gp
from gurobipy import GRB

def railway_optimization():
    """Optimize manager assignments to railways to minimize total cost."""
    
    # 1. MODEL & DATA SETUP
    model = gp.Model("railway_management")
    
    # Data from the database
    managers = [1, 2, 3]
    railways = [1, 2, 3]
    
    # Manager costs and capacities
    cost_per_level = {1: 120, 2: 150, 3: 200}
    manager_capacity = {1: 4, 2: 6, 3: 8}
    
    # CRITICAL: Validate array lengths before loops
    assert len(managers) == len(cost_per_level) == len(manager_capacity), "Array length mismatch"
    
    # 2. VARIABLES
    # Binary decision variables: x[m, r] = 1 if manager m is assigned to railway r
    x = model.addVars(managers, railways, vtype=GRB.BINARY, name="x")
    
    # 3. OBJECTIVE FUNCTION
    # Minimize the total cost of assigning managers to railways
    model.setObjective(gp.quicksum(cost_per_level[m] * x[m, r] for m in managers for r in railways), GRB.MINIMIZE)
    
    # 4. CONSTRAINTS
    
    # Railway Management Constraint: Each railway must be managed by at least one manager
    for r in railways:
        model.addConstr(gp.quicksum(x[m, r] for m in managers) >= 1, name=f"railway_{r}_managed")
    
    # Manager Capacity Constraint: No manager can be assigned to more railways than their capacity
    for m in managers:
        model.addConstr(gp.quicksum(x[m, r] for r in railways) <= manager_capacity[m], name=f"manager_{m}_capacity")
    
    # 5. SOLVING & RESULTS
    model.optimize()
    
    if model.status == GRB.OPTIMAL:
        print(f"Optimal value: {model.objVal}")
        for m in managers:
            for r in railways:
                if x[m, r].x > 0.5:
                    print(f"Manager {m} is assigned to Railway {r}")
    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__":
    railway_optimization()