# Complete PYOMO implementation

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

def flight_optimization():
    """Flight schedule optimization using Pyomo"""
    
    # 1. MODEL CREATION
    model = pyo.ConcreteModel()
    
    # 2. DATA SETUP
    # Example data for demonstration purposes
    flight_numbers = ['FL123', 'FL456', 'FL789', 'FL101', 'FL202']
    costs = [1200.0, 2500.0, 1800.0, 1500.0, 1600.0]
    mandatory_flights = ['FL789', 'FL101', 'FL202']
    
    # Operational limits
    max_flights = 4
    max_departures = {'AirportA': 3, 'AirportB': 2}
    max_arrivals = {'AirportA': 3, 'AirportB': 2}
    
    # CRITICAL: Validate array lengths before indexing
    assert len(flight_numbers) == len(costs), "Array length mismatch"
    
    # 3. SETS
    model.Flights = pyo.Set(initialize=flight_numbers)
    model.MandatoryFlights = pyo.Set(initialize=mandatory_flights)
    
    # 4. PARAMETERS
    model.Cost = pyo.Param(model.Flights, initialize={flight_numbers[i]: costs[i] for i in range(len(flight_numbers))})
    
    # 5. VARIABLES
    model.x = pyo.Var(model.Flights, within=pyo.Binary)
    
    # 6. OBJECTIVE FUNCTION
    def obj_rule(model):
        return sum(model.Cost[flight] * model.x[flight] for flight in model.Flights)
    model.objective = pyo.Objective(rule=obj_rule, sense=pyo.minimize)
    
    # 7. CONSTRAINTS
    
    # Maximum number of flights operated
    def max_flights_rule(model):
        return sum(model.x[flight] for flight in model.Flights) <= max_flights
    model.max_flights_constraint = pyo.Constraint(rule=max_flights_rule)
    
    # Mandatory flights
    def mandatory_flights_rule(model, flight):
        return model.x[flight] == 1
    model.mandatory_flights_constraint = pyo.Constraint(model.MandatoryFlights, rule=mandatory_flights_rule)
    
    # Maximum departures from each airport
    def max_departures_rule(model, airport):
        # Example: flights departing from each airport
        departures_from_airport = {'AirportA': ['FL123', 'FL456'], 'AirportB': ['FL789', 'FL101']}
        return sum(model.x[flight] for flight in departures_from_airport[airport]) <= max_departures[airport]
    model.max_departures_constraint = pyo.Constraint(max_departures.keys(), rule=max_departures_rule)
    
    # Maximum arrivals at each airport
    def max_arrivals_rule(model, airport):
        # Example: flights arriving at each airport
        arrivals_at_airport = {'AirportA': ['FL789', 'FL202'], 'AirportB': ['FL123', 'FL456']}
        return sum(model.x[flight] for flight in arrivals_at_airport[airport]) <= max_arrivals[airport]
    model.max_arrivals_constraint = pyo.Constraint(max_arrivals.keys(), rule=max_arrivals_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)
    
    # 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("\nFlight operation decisions:")
        for flight in model.Flights:
            x_val = pyo.value(model.x[flight])
            if x_val > 1e-6:  # Only print flights that are operated
                print(f"Flight {flight} is operated.")
        
    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
flight_optimization()