# Complete DOCPLEX implementation - Retry Attempt 4

from docplex.mp.model import Model

def train_station_optimization():
    # 1. MODEL & DATA SETUP
    mdl = Model(name="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 array lengths
    assert len(passenger_waiting_time) == len(number_of_passengers), "Data length mismatch"
    
    # Extract unique trains and stations
    trains = set(t for t, s in passenger_waiting_time.keys())
    stations = set(s for t, s in passenger_waiting_time.keys())
    
    # 2. VARIABLES
    x = mdl.binary_var_dict(keys=passenger_waiting_time.keys(), name="x")
    
    # 3. OBJECTIVE FUNCTION
    objective = mdl.sum(
        number_of_passengers[t, s] * passenger_waiting_time[t, s] * x[t, s]
        for t, s in passenger_waiting_time.keys()
    )
    mdl.minimize(objective)
    
    # 4. CONSTRAINTS
    
    # Platform Capacity Constraint
    for s in stations:
        mdl.add_constraint(
            mdl.sum(x[t, s] for t in trains if (t, s) in x) <= station_platforms[s],
            ctname=f"platform_capacity_{s}"
        )
    
    # Maximum Waiting Time Constraint
    for t, s in passenger_waiting_time.keys():
        mdl.add_constraint(
            passenger_waiting_time[t, s] * x[t, s] <= 15,
            ctname=f"max_waiting_time_{t}_{s}"
        )
    
    # 5. SOLVING & RESULTS
    solution = mdl.solve()
    
    if solution:
        print(f"Optimal value: {solution.objective_value}")
        for t, s in passenger_waiting_time.keys():
            value = solution.get_value(x[t, s])
            if value > 1e-6:
                print(f"Train {t} assigned to Station {s}: {value:.0f}")
    else:
        print("No solution found")
        print(f"Status: {mdl.solve_details.status}")
    
    return mdl

# Execute the optimization
train_station_optimization()