# Complete DOCPLEX implementation

from docplex.mp.model import Model

def flight_optimization():
    """Flight scheduling optimization using DOCPLEX"""

    # 1. MODEL & DATA SETUP
    mdl = Model(name="flight_optimization")

    # Data from the problem context
    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 to prevent IndexError
    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"

    # Safe indexing
    safe_flights = range(len(flights))
    safe_aircrafts = range(len(aircrafts))
    safe_employees = range(len(employees))

    # 2. VARIABLES
    # Decision variables: x[f, a, e] is 1 if flight f is assigned to aircraft a and employee e
    x = {(f, a, e): mdl.binary_var(name=f"x_{f}_{a}_{e}") for f in flights for a in aircrafts for e in employees}

    # 3. OBJECTIVE FUNCTION
    # Minimize total operational cost
    objective = mdl.sum((flight_prices[f] + employee_salaries[e]) * x[f, a, e] for f in flights for a in aircrafts for e in employees)
    mdl.minimize(objective)

    # 4. CONSTRAINTS

    # Flight Assignment Constraint: Each flight must be assigned to exactly one aircraft and one employee
    for f in flights:
        mdl.add_constraint(mdl.sum(x[f, a, e] for a in aircrafts for e in employees) == 1, ctname=f"assign_flight_{f}")

    # Aircraft Distance Capability Constraint
    for a in aircrafts:
        for e in employees:
            mdl.add_constraint(mdl.sum(flight_distances[f] * x[f, a, e] for f in flights) <= aircraft_distances[a], ctname=f"distance_cap_{a}_{e}")

    # Employee Salary Budget Constraint
    mdl.add_constraint(mdl.sum(employee_salaries[e] * x[f, a, e] for f in flights for a in aircrafts for e in employees) <= budget, ctname="salary_budget")

    # 5. SOLVING & RESULTS
    solution = mdl.solve()

    if solution:
        print(f"Optimal value: {solution.objective_value}")
        for f in flights:
            for a in aircrafts:
                for e in employees:
                    if solution.get_value(x[f, a, e]) > 0.5:  # Check if the binary variable is effectively 1
                        print(f"Flight {f} assigned to Aircraft {a} and Employee {e}")
    else:
        print("No solution found")
        print(f"Status: {mdl.solve_details.status}")

    return mdl

# Run the optimization
flight_optimization()