#!/usr/bin/env python3
"""
Pyomo 6.9.2 Implementation for Train Scheduling Optimization
"""

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

def train_scheduling_optimization():
    """Optimize train schedules to minimize total travel time considering weather delays."""
    
    # 1. MODEL CREATION
    model = pyo.ConcreteModel()
    
    # 2. DATA SETUP
    # Train data
    train_ids = [1, 2, 3]
    base_travel_times = {1: 120.0, 2: 150.0, 3: 180.0}
    max_travel_times = {1: 200.0, 2: 250.0, 3: 300.0}
    
    # Weather data
    precipitation = {1: 0.5, 2: 1.0, 3: 1.5}
    wind_speed_mph = {1: 10.0, 2: 15.0, 3: 20.0}
    
    # Validate array lengths
    assert len(train_ids) == len(base_travel_times) == len(max_travel_times) == len(precipitation) == len(wind_speed_mph), "Array length mismatch"
    
    # 3. SETS
    model.I = pyo.Set(initialize=train_ids)  # Set of trains
    
    # 4. PARAMETERS
    model.base_travel_time = pyo.Param(model.I, initialize=base_travel_times)
    model.max_travel_time = pyo.Param(model.I, initialize=max_travel_times)
    model.precipitation = pyo.Param(model.I, initialize=precipitation)
    model.wind_speed_mph = pyo.Param(model.I, initialize=wind_speed_mph)
    
    # 5. VARIABLES
    model.t = pyo.Var(model.I, within=pyo.NonNegativeReals)  # Total travel time for each train
    
    # 6. OBJECTIVE FUNCTION
    def obj_rule(model):
        return sum(model.t[i] for i in model.I)
    model.objective = pyo.Objective(rule=obj_rule, sense=pyo.minimize)
    
    # 7. CONSTRAINTS
    # Minimum Travel Time Constraint
    def min_travel_time_rule(model, i):
        return model.t[i] >= model.base_travel_time[i] + 0.1 * model.precipitation[i] + 0.05 * model.wind_speed_mph[i]
    model.min_travel_time_constraint = pyo.Constraint(model.I, rule=min_travel_time_rule)
    
    # Maximum Travel Time Constraint
    def max_travel_time_rule(model, i):
        return model.t[i] <= model.max_travel_time[i]
    model.max_travel_time_constraint = pyo.Constraint(model.I, rule=max_travel_time_rule)
    
    # 8. SOLVING WITH GUROBI
    solver = SolverFactory('gurobi')
    
    # Optional: Set solver options
    solver.options['TimeLimit'] = 300  # 5 minutes
    solver.options['MIPGap'] = 0.01    # 1% gap
    
    # Solve the model
    results = solver.solve(model, tee=True)  # tee=True shows solver output
    
    # 9. RESULT PROCESSING
    # Check solver status
    if results.solver.termination_condition == pyo.TerminationCondition.optimal:
        print("Optimal solution found!")
        print(f"Optimal value: {pyo.value(model.objective)}")
        
        # Extract variable values
        print("\nVariable values:")
        for i in model.I:
            t_val = pyo.value(model.t[i])
            print(f"t[{i}] = {t_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}")
    
    return model

# Execute the optimization
if __name__ == "__main__":
    train_scheduling_optimization()