# Complete PYOMO implementation - Retry Attempt 4

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

def train_scheduling_optimization():
    """Optimize train scheduling considering weather-induced delays."""
    
    # 1. MODEL CREATION
    model = pyo.ConcreteModel()
    
    # 2. DATA SETUP
    # Define train and station sets
    train_ids = [1, 2, 3]
    station_ids = [101, 102, 103]
    
    # Adjusted travel times for each train at each station
    adjusted_travel_times = {
        (1, 101): 125.0,
        (2, 102): 135.0,
        (3, 103): 115.0
    }
    
    # Weather coefficients and base delay
    base_delay = 1.0
    precipitation_coefficient = 0.1
    wind_speed_coefficient = 0.05
    
    # Weather data for each station
    precipitation_levels = {101: 0.2, 102: 0.3, 103: 0.1}
    wind_speeds = {101: 5.0, 102: 3.0, 103: 4.0}
    
    # 3. SETS
    model.T = pyo.Set(initialize=train_ids)
    model.S = pyo.Set(initialize=station_ids)
    
    # 4. PARAMETERS
    model.adjusted_travel_time = pyo.Param(model.T * model.S, initialize=adjusted_travel_times, within=pyo.Reals)
    model.precipitation_level = pyo.Param(model.S, initialize=precipitation_levels, within=pyo.NonNegativeReals)
    model.wind_speed = pyo.Param(model.S, initialize=wind_speeds, within=pyo.NonNegativeReals)
    
    # 5. VARIABLES
    model.d = pyo.Var(model.S, within=pyo.NonNegativeReals)
    
    # 6. OBJECTIVE FUNCTION
    def obj_rule(model):
        return sum(model.adjusted_travel_time[t, s] * model.d[s] for t in model.T for s in model.S)
    model.objective = pyo.Objective(rule=obj_rule, sense=pyo.minimize)
    
    # 7. CONSTRAINTS
    def delay_factor_rule(model, s):
        return model.d[s] == base_delay + precipitation_coefficient * model.precipitation_level[s] + wind_speed_coefficient * model.wind_speed[s]
    model.delay_factor_constraint = pyo.Constraint(model.S, rule=delay_factor_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("\nDelay factors:")
        for s in model.S:
            d_val = pyo.value(model.d[s])
            print(f"d[{s}] = {d_val:.3f}")
        
    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}")

# Execute the optimization
train_scheduling_optimization()