# Complete GUROBIPY implementation

import gurobipy as gp
from gurobipy import GRB

def optimize_train_station_allocation():
    """Optimize train allocation to stations to maximize passengers served."""
    
    # 1. MODEL & DATA SETUP
    model = gp.Model("train_station")
    
    # Data from the problem context
    stations = [1, 2, 3]
    trains = [101, 102, 103]
    
    # Passenger demand for each train at each station
    demand = {
        (101, 1): 250, (101, 2): 300, (101, 3): 0,
        (102, 1): 0, (102, 2): 350, (102, 3): 400,
        (103, 1): 0, (103, 2): 0, (103, 3): 150
    }
    
    # Number of platforms available at each station
    platforms = {1: 3, 2: 4, 3: 5}
    
    # Compatibility of train services with stations
    compatibility = {
        (101, 1): True, (101, 2): True, (101, 3): False,
        (102, 1): False, (102, 2): True, (102, 3): True,
        (103, 1): False, (103, 2): False, (103, 3): False
    }
    
    # Validate data lengths
    assert len(demand) == len(trains) * len(stations), "Demand data length mismatch"
    assert len(platforms) == len(stations), "Platform data length mismatch"
    assert len(compatibility) == len(trains) * len(stations), "Compatibility data length mismatch"
    
    # 2. VARIABLES
    # Decision variables: x[i, j] = 1 if train i is assigned to station j
    x = model.addVars(trains, stations, vtype=GRB.BINARY, name="x")
    
    # 3. OBJECTIVE FUNCTION
    # Maximize total number of passengers served
    model.setObjective(gp.quicksum(demand[i, j] * x[i, j] for i in trains for j in stations), GRB.MAXIMIZE)
    
    # 4. CONSTRAINTS
    
    # Platform Capacity Constraint
    model.addConstrs((gp.quicksum(x[i, j] for i in trains) <= platforms[j] for j in stations), name="platform_capacity")
    
    # Service Compatibility Constraint
    model.addConstrs((x[i, j] <= compatibility[i, j] for i in trains for j in stations), name="service_compatibility")
    
    # Exclusive Assignment Constraint
    model.addConstrs((gp.quicksum(x[i, j] for j in stations) <= 1 for i in trains), name="exclusive_assignment")
    
    # 5. SOLVING & RESULTS
    model.optimize()
    
    if model.status == GRB.OPTIMAL:
        print(f"Optimal value: {model.objVal}")
        for i in trains:
            for j in stations:
                if x[i, j].x > 1e-6:
                    print(f"Train {i} assigned to Station {j} with demand {demand[i, j]}")
    elif model.status == GRB.INFEASIBLE:
        print("Problem is infeasible")
    elif model.status == GRB.UNBOUNDED:
        print("Problem is unbounded")
    
    return model

# Run the optimization
optimize_train_station_allocation()