# Complete PYOMO implementation

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

def flight_optimization():
    """Flight scheduling optimization using Pyomo"""
    
    # 1. MODEL CREATION
    model = pyo.ConcreteModel()
    
    # 2. DATA SETUP
    flights = [101, 102, 103]
    aircrafts = [1, 2, 3]
    employees = [10, 20, 30]
    
    flight_prices = {101: 5000.0, 102: 7500.0, 103: 10000.0}
    flight_distances = {101: 300.0, 102: 450.0, 103: 600.0}
    
    aircraft_distances = {1: 500.0, 2: 700.0, 3: 1000.0}
    
    employee_salaries = {10: 3000.0, 20: 4000.0, 30: 5000.0}
    
    budget = 12000.0  # Example budget constraint
    
    # CRITICAL: Validate array lengths before indexing
    assert len(flights) == len(flight_prices) == len(flight_distances), "Flight data length mismatch"
    assert len(aircrafts) == len(aircraft_distances), "Aircraft data length mismatch"
    assert len(employees) == len(employee_salaries), "Employee data length mismatch"
    
    # 3. SETS
    model.F = pyo.Set(initialize=flights)
    model.A = pyo.Set(initialize=aircrafts)
    model.E = pyo.Set(initialize=employees)
    
    # 4. PARAMETERS
    model.price = pyo.Param(model.F, initialize=flight_prices)
    model.distance = pyo.Param(model.F, initialize=flight_distances)
    model.aircraft_distance = pyo.Param(model.A, initialize=aircraft_distances)
    model.salary = pyo.Param(model.E, initialize=employee_salaries)
    
    # 5. VARIABLES
    model.x = pyo.Var(model.F, model.A, model.E, within=pyo.Binary)
    
    # 6. OBJECTIVE FUNCTION
    def obj_rule(model):
        return sum((model.price[f] + model.salary[e]) * model.x[f, a, e] for f in model.F for a in model.A for e in model.E)
    model.objective = pyo.Objective(rule=obj_rule, sense=pyo.minimize)
    
    # 7. CONSTRAINTS
    
    # Flight Assignment Constraint
    def flight_assignment_rule(model, f):
        return sum(model.x[f, a, e] for a in model.A for e in model.E) == 1
    model.flight_assignment_constraint = pyo.Constraint(model.F, rule=flight_assignment_rule)
    
    # Aircraft Distance Capability Constraint
    def aircraft_distance_rule(model, a):
        return sum(model.distance[f] * model.x[f, a, e] for f in model.F for e in model.E) <= model.aircraft_distance[a]
    model.aircraft_distance_constraint = pyo.Constraint(model.A, rule=aircraft_distance_rule)
    
    # Employee Salary Budget Constraint
    def salary_budget_rule(model):
        return sum(model.salary[e] * model.x[f, a, e] for f in model.F for a in model.A for e in model.E) <= budget
    model.salary_budget_constraint = pyo.Constraint(rule=salary_budget_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):.2f}")
        
        # Extract variable values
        print("\nFlight Assignments:")
        for f in model.F:
            for a in model.A:
                for e in model.E:
                    if pyo.value(model.x[f, a, e]) > 0.5:  # Binary variable, check if assigned
                        print(f"Flight {f} assigned to Aircraft {a} and Employee {e}")
        
    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()