Aircraft Fleet Problem

from qaekwy.engine import DirectEngine
from qaekwy.model.constraint.relational import RelationalExpression
from qaekwy.model.specific import SpecificMaximum
from qaekwy.model.variable.integer import IntegerExpressionVariable, IntegerVariable
from qaekwy.model.modeller import Modeller
from qaekwy.model.searcher import SearcherType


class AircraftFleetProblem(Modeller):
    def __init__(
            self,
            aircraft_weights,
            aircraft_values,
            fuel_efficiencies,
            ranges,
            maintenance_costs,
            fleet_size, max_distance):
        super().__init__()

        num_aircraft = len(aircraft_weights)
        selected_aircraft = [
                IntegerVariable(var_name=f"aircraft_{i}", domain_low=0, domain_high=1)
                for i in range(num_aircraft)
            ]

        # Constraint: Total weight of selected aircraft should be less than or equal to the fleet size
        weight_constraint = RelationalExpression(
            sum(aircraft_weights[i] * selected_aircraft[i]
            for i in range(num_aircraft)) <= fleet_size
        )

        # Constraint: The maximum distance traveled by the fleet should be greater than or equal to the maximum distance
        distance_constraint = RelationalExpression(
            sum(ranges[i] * selected_aircraft[i]
            for i in range(num_aircraft)) >= max_distance
        )

        # Constraint: Total fuel consumption should be within limits
        fuel_consumption_constraint = RelationalExpression(
            sum(fuel_efficiencies[i] * selected_aircraft[i] * ranges[i]
            for i in range(num_aircraft)) <= max_distance / 2
        )

        # Constraint: Total maintenance cost should be within limits
        maintenance_cost_constraint = RelationalExpression(
            sum(maintenance_costs[i] * selected_aircraft[i]
            for i in range(num_aircraft)) <= fleet_size * 1000
        )

        # Constraint: Maximum number of aircraft of a particular type should not be exceeded
        for aircraft_type in range(num_aircraft):
            number_of_aircraft_type_constraint = RelationalExpression(
                sum(selected_aircraft[i] for i in range(num_aircraft) if aircraft_type == i) <= 10
            )
            self.add_constraint(number_of_aircraft_type_constraint)

        # Constraint: The fleet should have at least one long-range aircraft
        long_range_aircraft = [i for i in range(num_aircraft) if ranges[i] > 5000]
        number_of_long_range_aircraft_constraint = RelationalExpression(
            sum(selected_aircraft[i] for i in long_range_aircraft) >= 1
        )
        self.add_constraint(number_of_long_range_aircraft_constraint)

        # Maximize the total value of the selected aircraft
        total_value = IntegerExpressionVariable(
            var_name="total_value",
            expression=sum(aircraft_values[i] * selected_aircraft[i] for i in range(num_aircraft))
        )

        self.add_variable(total_value)
        self.add_objective(SpecificMaximum(total_value))

        # Add variables and constraints to the problem
        for aircraft_var in selected_aircraft:
            self.add_variable(aircraft_var)

        self.add_constraint(weight_constraint)
        self.add_constraint(distance_constraint)
        self.add_constraint(fuel_consumption_constraint)
        self.add_constraint(maintenance_cost_constraint)

        # Set the search strategy to Depth-First Search (DFS)
        self.set_searcher(SearcherType.DFS)


if __name__ == "__main__":
    # Create a Qaekwy engine for direct interaction
    qaekwy_engine = DirectEngine()

    # Define the aircraft weights, values, and fleet size
    aircraft_weights = [200, 300, 400, 500]
    aircraft_values = [3000, 4000, 5000, 6000]
    fuel_efficiencies = [0.05, 0.06, 0.07, 0.08]
    ranges = [6000, 7000, 8000, 9000]
    maintenance_costs = [100, 120, 150, 180]
    fleet_size = 1000
    max_distance = 10000

    # Create the aircraft fleet problem instance
    aircraft_fleet_problem = AircraftFleetProblem(
        aircraft_weights,
        aircraft_values,
        fuel_efficiencies,
        ranges,
        maintenance_costs,
        fleet_size,
        max_distance
    )

    # Request the Qaekwy engine to solve the aircraft fleet problem
    response = qaekwy_engine.model(model=aircraft_fleet_problem)

    # Retrieve the list of solutions from the response
    list_of_solutions = response.get_solutions()
    solution = list_of_solutions[0]

    # Print the solution(s) obtained for the aircraft fleet problem
    print("Selected aircraft:")
    for solution_item, solution_value in solution.items():
        print(f"{solution_item}: {solution_value}")