# Complete PYOMO implementation

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

def flight_optimization():
    """Optimize flight operations to minimize costs while ensuring operational efficiency."""
    
    # 1. MODEL CREATION
    model = pyo.ConcreteModel()
    
    # 2. DATA SETUP
    # Flight operational costs (example data)
    flight_costs = {101: 5000, 102: 4500, 103: 4800}
    
    # Employee salaries (example data)
    employee_salaries = {1: 3000, 2: 3200, 3: 3100}
    
    # Aircraft capacity data (example data)
    aircraft_capacity = {('JFK', 'LAX'): 8, ('LAX', 'JFK'): 8, ('ORD', 'DFW'): 5}
    
    # Flight origin-destination pairs (example data)
    flight_od_pairs = {101: ('JFK', 'LAX'), 102: ('LAX', 'JFK'), 103: ('ORD', 'DFW')}
    
    # Employee assignments (example data)
    employee_assignments = [(1, 101), (2, 102), (3, 103)]
    
    # Validate array lengths before indexing
    assert len(flight_costs) == len(flight_od_pairs), "Flight data length mismatch"
    assert len(employee_salaries) == len(set(eid for eid, _ in employee_assignments)), "Employee data length mismatch"
    
    # 3. SETS (Pyomo way to define indices)
    model.F = pyo.Set(initialize=flight_costs.keys())  # Set of flights
    model.E = pyo.Set(initialize=employee_salaries.keys())  # Set of employees
    model.OD = pyo.Set(initialize=aircraft_capacity.keys())  # Set of origin-destination pairs
    
    # 4. PARAMETERS (data containers)
    model.cost_f = pyo.Param(model.F, initialize=flight_costs)  # Flight operational costs
    model.salary_e = pyo.Param(model.E, initialize=employee_salaries)  # Employee salaries
    model.max_flights_od = pyo.Param(model.OD, initialize=aircraft_capacity)  # Maximum flights per OD pair
    model.flight_od = pyo.Param(model.F, initialize=flight_od_pairs)  # Flight origin-destination pairs
    
    # 5. VARIABLES
    # Binary decision variable for operating flights
    model.x_f = pyo.Var(model.F, within=pyo.Binary)
    
    # Binary decision variable for assigning employees to flights
    model.y_ef = pyo.Var(model.E, model.F, within=pyo.Binary)
    
    # 6. OBJECTIVE FUNCTION
    def obj_rule(model):
        return sum(model.cost_f[f] * model.x_f[f] for f in model.F) + \
               sum(model.salary_e[e] * sum(model.y_ef[e, f] for f in model.F) for e in model.E)
    model.objective = pyo.Objective(rule=obj_rule, sense=pyo.minimize)
    
    # 7. CONSTRAINTS
    
    # Flight Capacity Constraint
    def flight_capacity_rule(model, o, d):
        return sum(model.x_f[f] for f in model.F if model.flight_od[f] == (o, d)) <= model.max_flights_od[o, d]
    model.flight_capacity_constraint = pyo.Constraint(model.OD, rule=flight_capacity_rule)
    
    # Employee Assignment Constraint
    def employee_assignment_rule(model, f):
        return sum(model.y_ef[e, f] for e in model.E) >= model.x_f[f]
    model.employee_assignment_constraint = pyo.Constraint(model.F, rule=employee_assignment_rule)
    
    # 8. SOLVING WITH GUROBI (your available solver)
    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("\nFlight operations:")
        for f in model.F:
            if pyo.value(model.x_f[f]) > 0.5:  # Only print operated flights
                print(f"Flight {f} is operated.")
        
        print("\nEmployee assignments:")
        for e in model.E:
            for f in model.F:
                if pyo.value(model.y_ef[e, f]) > 0.5:  # Only print assigned employees
                    print(f"Employee {e} is assigned to flight {f}.")
        
    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
if __name__ == "__main__":
    flight_optimization()