#!/usr/bin/env python3
"""
DOCplex implementation for train scheduling optimization with weather delays
"""

from docplex.mp.model import Model

def optimize_train_scheduling():
    """Optimize train scheduling to minimize total travel time with weather delays"""
    
    # 1. MODEL & DATA SETUP
    mdl = Model(name="train_scheduling")
    
    # Data from the problem description
    train_ids = [1, 2, 3]
    base_travel_times = [120.0, 150.0, 180.0]
    max_travel_times = [200.0, 250.0, 300.0]
    precipitation = [0.5, 1.0, 1.5]
    wind_speed_mph = [10.0, 15.0, 20.0]
    
    # CRITICAL: Validate array lengths to prevent IndexError
    assert len(train_ids) == len(base_travel_times) == len(max_travel_times) == len(precipitation) == len(wind_speed_mph), "Array length mismatch"
    safe_range = range(min(len(train_ids), len(base_travel_times), len(max_travel_times), len(precipitation), len(wind_speed_mph)))  # Safe indexing
    
    # 2. VARIABLES
    # Decision variables: total travel time for each train
    t = {i: mdl.continuous_var(name=f"t_{train_ids[i]}", lb=0) for i in safe_range}
    
    # 3. OBJECTIVE FUNCTION
    # Minimize the sum of total travel times for all trains
    objective = mdl.sum(t[i] for i in safe_range)
    mdl.minimize(objective)
    
    # 4. CONSTRAINTS
    # Pattern B: Individual constraints in loop with safe indexing
    for i in safe_range:
        # Minimum Travel Time Constraint
        mdl.add_constraint(t[i] >= base_travel_times[i] + 0.1 * precipitation[i] + 0.05 * wind_speed_mph[i], ctname=f"min_travel_time_{train_ids[i]}")
        
        # Maximum Travel Time Constraint
        mdl.add_constraint(t[i] <= max_travel_times[i], ctname=f"max_travel_time_{train_ids[i]}")
    
    # 5. SOLVING & RESULTS
    solution = mdl.solve()
    
    if solution:
        print(f"Optimal value: {solution.objective_value}")
        for i in safe_range:
            value = solution.get_value(t[i])
            print(f"t[{train_ids[i]}] = {value:.3f}")
    else:
        print("No solution found")
        print(f"Status: {mdl.solve_details.status}")
    
    return mdl

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