# Complete PYOMO implementation

import pyomo.environ as pyo
from pyomo.opt import SolverFactory

def train_station_optimization():
    """Optimize train assignments to stations to maximize passengers served"""
    
    # 1. MODEL CREATION
    model = pyo.ConcreteModel()
    
    # 2. DATA SETUP
    # Data from the database schema
    stations = [1, 2, 3]
    trains = [101, 102, 103]
    
    # Passenger demand for each train at each station
    demand_data = {
        (101, 1): 250,
        (102, 2): 350,
        (103, 3): 150,
        (101, 2): 300,
        (102, 3): 400
    }
    
    # Number of platforms available at each station
    platforms_data = {
        1: 3,
        2: 4,
        3: 5
    }
    
    # Service compatibility of each train at each station
    compatibility_data = {
        (101, 1): True,
        (102, 2): True,
        (103, 3): False,
        (101, 2): True,
        (102, 3): True
    }
    
    # 3. SETS
    model.Trains = pyo.Set(initialize=trains)
    model.Stations = pyo.Set(initialize=stations)
    
    # 4. PARAMETERS
    model.Demand = pyo.Param(model.Trains, model.Stations, initialize=demand_data, default=0)
    model.Platforms = pyo.Param(model.Stations, initialize=platforms_data)
    model.Compatible = pyo.Param(model.Trains, model.Stations, initialize=compatibility_data, default=False)
    
    # 5. VARIABLES
    model.x = pyo.Var(model.Trains, model.Stations, within=pyo.Binary)
    
    # 6. OBJECTIVE FUNCTION
    def obj_rule(model):
        return sum(model.Demand[i, j] * model.x[i, j] for i in model.Trains for j in model.Stations)
    model.objective = pyo.Objective(rule=obj_rule, sense=pyo.maximize)
    
    # 7. CONSTRAINTS
    
    # Platform Capacity Constraint
    def platform_capacity_rule(model, j):
        return sum(model.x[i, j] for i in model.Trains) <= model.Platforms[j]
    model.platform_capacity_constraint = pyo.Constraint(model.Stations, rule=platform_capacity_rule)
    
    # Service Compatibility Constraint
    def service_compatibility_rule(model, i, j):
        return model.x[i, j] <= model.Compatible[i, j]
    model.service_compatibility_constraint = pyo.Constraint(model.Trains, model.Stations, rule=service_compatibility_rule)
    
    # Exclusive Assignment Constraint
    def exclusive_assignment_rule(model, i):
        return sum(model.x[i, j] for j in model.Stations) <= 1
    model.exclusive_assignment_constraint = pyo.Constraint(model.Trains, rule=exclusive_assignment_rule)
    
    # 8. SOLVING WITH GUROBI
    solver = SolverFactory('gurobi')
    
    # Solve the model
    results = solver.solve(model, tee=True)
    
    # 9. RESULT PROCESSING
    if results.solver.termination_condition == pyo.TerminationCondition.optimal:
        print("Optimal solution found!")
        print(f"Optimal value: {pyo.value(model.objective)}")
        
        # Extract variable values
        print("\nTrain assignments:")
        for i in model.Trains:
            for j in model.Stations:
                if pyo.value(model.x[i, j]) > 0.5:  # Binary variable, check if assigned
                    print(f"Train {i} assigned to Station {j}")
        
    elif results.solver.termination_condition == pyo.TerminationCondition.infeasible:
        print("Problem is infeasible")
    elif results.solver.termination_condition == pyo.TerminationCondition.unbounded:
        print("Problem is unbounded")
    else:
        print(f"Solver terminated with condition: {results.solver.termination_condition}")
    
    return model

# Run the optimization
train_station_optimization()