# Complete GUROBIPY implementation - Retry Attempt 4

import gurobipy as gp
from gurobipy import GRB

def train_station_optimization():
    # 1. MODEL & DATA SETUP
    model = gp.Model("train_station_optimization")
    
    # Data from the database
    passenger_waiting_time = {
        (1, 1): 5,
        (2, 2): 10,
        (3, 3): 15
    }
    
    number_of_passengers = {
        (1, 1): 150,
        (2, 2): 100,
        (3, 3): 50
    }
    
    station_platforms = {
        1: 4,
        2: 3,
        3: 2
    }
    
    # Validate data consistency
    trains = set(t for t, s in passenger_waiting_time.keys())
    stations = set(s for t, s in passenger_waiting_time.keys())
    
    assert all((t, s) in number_of_passengers for t, s in passenger_waiting_time), "Missing passenger_count for some train-station pairs"
    assert all(s in station_platforms for s in stations), "Missing platform count for some stations"
    
    # 2. VARIABLES
    x = model.addVars(trains, stations, vtype=GRB.BINARY, name="x")
    
    # 3. OBJECTIVE FUNCTION
    model.setObjective(
        gp.quicksum(
            number_of_passengers[t, s] * passenger_waiting_time[t, s] * x[t, s]
            for t, s in passenger_waiting_time
        ),
        GRB.MINIMIZE
    )
    
    # 4. CONSTRAINTS
    
    # Platform Capacity Constraint
    for s in stations:
        model.addConstr(
            gp.quicksum(x[t, s] for t in trains) <= station_platforms[s],
            name=f"platform_capacity_{s}"
        )
    
    # Maximum Waiting Time Constraint
    for t, s in passenger_waiting_time:
        model.addConstr(
            passenger_waiting_time[t, s] * x[t, s] <= 15,
            name=f"max_waiting_time_{t}_{s}"
        )
    
    # 5. SOLVING & RESULTS
    model.optimize()
    
    if model.status == GRB.OPTIMAL:
        print(f"Optimal value: {model.objVal}")
        for t, s in x:
            if x[t, s].x > 1e-6:
                print(f"Train {t} assigned to Station {s}: {x[t, s].x:.0f}")
    elif model.status == GRB.INFEASIBLE:
        print("Problem is infeasible")
    elif model.status == GRB.UNBOUNDED:
        print("Problem is unbounded")
    
    return model

# Execute the optimization
train_station_optimization()