# Complete DOCPLEX implementation

from docplex.mp.model import Model

def train_station_optimization():
    """Optimize train assignments to stations to maximize passengers served."""
    
    # 1. MODEL & DATA SETUP
    mdl = Model(name="train_station")
    
    # Data from the database
    stations = [
        {"Station_ID": 1, "Total_Passengers": 1200, "Number_of_Platforms": 3},
        {"Station_ID": 2, "Total_Passengers": 1800, "Number_of_Platforms": 4},
        {"Station_ID": 3, "Total_Passengers": 2200, "Number_of_Platforms": 5}
    ]
    
    train_station_data = [
        {"Train_ID": 101, "Station_ID": 1, "Service_Compatible": True},
        {"Train_ID": 102, "Station_ID": 2, "Service_Compatible": True},
        {"Train_ID": 103, "Station_ID": 3, "Service_Compatible": False},
        {"Train_ID": 101, "Station_ID": 2, "Service_Compatible": True},
        {"Train_ID": 102, "Station_ID": 3, "Service_Compatible": True}
    ]
    
    service_compatibility = [
        {"Train_ID": 101, "Station_ID": 1, "Compatible": True},
        {"Train_ID": 102, "Station_ID": 2, "Compatible": True},
        {"Train_ID": 103, "Station_ID": 3, "Compatible": False},
        {"Train_ID": 101, "Station_ID": 2, "Compatible": True},
        {"Train_ID": 102, "Station_ID": 3, "Compatible": True}
    ]
    
    passenger_demand = [
        {"Train_ID": 101, "Station_ID": 1, "Demand": 250},
        {"Train_ID": 102, "Station_ID": 2, "Demand": 350},
        {"Train_ID": 103, "Station_ID": 3, "Demand": 150},
        {"Train_ID": 101, "Station_ID": 2, "Demand": 300},
        {"Train_ID": 102, "Station_ID": 3, "Demand": 400}
    ]
    
    # Extract unique train and station IDs
    train_ids = set(d["Train_ID"] for d in passenger_demand)
    station_ids = set(d["Station_ID"] for d in stations)
    
    # 2. VARIABLES
    x = {(i, j): mdl.binary_var(name=f"x_{i}_{j}") for i in train_ids for j in station_ids}
    
    # 3. OBJECTIVE FUNCTION
    objective = mdl.sum(d["Demand"] * x[d["Train_ID"], d["Station_ID"]] for d in passenger_demand)
    mdl.maximize(objective)
    
    # 4. CONSTRAINTS
    
    # Platform Capacity Constraint
    for station in stations:
        station_id = station["Station_ID"]
        mdl.add_constraint(
            mdl.sum(x[i, station_id] for i in train_ids) <= station["Number_of_Platforms"],
            ctname=f"platform_capacity_{station_id}"
        )
    
    # Service Compatibility Constraint
    for sc in service_compatibility:
        if not sc["Compatible"]:
            mdl.add_constraint(
                x[sc["Train_ID"], sc["Station_ID"]] == 0,
                ctname=f"service_compatibility_{sc['Train_ID']}_{sc['Station_ID']}"
            )
    
    # Exclusive Assignment Constraint
    for train_id in train_ids:
        mdl.add_constraint(
            mdl.sum(x[train_id, j] for j in station_ids) <= 1,
            ctname=f"exclusive_assignment_{train_id}"
        )
    
    # 5. SOLVING & RESULTS
    solution = mdl.solve()
    
    if solution:
        print(f"Optimal value: {solution.objective_value}")
        for i in train_ids:
            for j in station_ids:
                value = solution.get_value(x[i, j])
                if value > 1e-6:
                    print(f"x[{i},{j}] = {value:.3f}")
    else:
        print("No solution found")
        print(f"Status: {mdl.solve_details.status}")
    
    return mdl

# Run the optimization
train_station_optimization()