import qaekwy as qw


def solve_aircraft_fleet(
    aircraft_weights: list[int],
    aircraft_values: list[int],
    fuel_efficiencies: list[float],
    ranges: list[int],
    maintenance_costs: list[int],
    fleet_size: int,
    max_distance: int,
):
    # Instantiate the model
    m = qw.Model()

    num_aircraft = len(aircraft_weights)

    # Define variables: indicates if an aircraft type is selected (binary 0 or 1)
    # Note: If the logic allowed multiple of the same type, domain high could be 10
    selected_aircraft = [
        m.integer_variable(name=f"aircraft_{i}", domain=(0, 1))
        for i in range(num_aircraft)
    ]

    # Constraint: Total weight must be within fleet capacity
    m.constraint(
        sum(aircraft_weights[i] * selected_aircraft[i] for i in range(num_aircraft))
        <= fleet_size
    )

    # Constraint: Combined range must meet the maximum distance requirement
    m.constraint(
        sum(ranges[i] * selected_aircraft[i] for i in range(num_aircraft))
        >= max_distance
    )

    # Constraint: Total fuel consumption limit
    # (efficiency * selected * range)
    m.constraint(
        sum(
            fuel_efficiencies[i] * selected_aircraft[i] * ranges[i]
            for i in range(num_aircraft)
        )
        <= max_distance / 2
    )

    # Constraint: Total maintenance cost limit
    m.constraint(
        sum(maintenance_costs[i] * selected_aircraft[i] for i in range(num_aircraft))
        <= fleet_size * 1000
    )

    # Constraint: Limit on number of aircraft per type (max 10)
    for i in range(num_aircraft):
        m.constraint(selected_aircraft[i] <= 10)

    # Constraint: At least one long-range aircraft (> 5000) must be included
    long_range_indices = [i for i, r in enumerate(ranges) if r > 5000]
    m.constraint(
        sum(selected_aircraft[i] for i in long_range_indices) >= 1
    )

    # Objective: Maximize the total value of the fleet
    total_value = m.integer_variable(
        name="total_value",
        expression=sum(
            aircraft_values[i] * selected_aircraft[i] for i in range(num_aircraft)
        ),
    )

    m.maximize(total_value)

    # Solve using Branch-and-bound searcher
    return m.solve_one(searcher="bab")


if __name__ == "__main__":
    # Problem parameters
    weights = [200, 300, 400, 500]
    values = [3000, 4000, 5000, 6000]
    efficiencies = [0.05, 0.06, 0.07, 0.08]
    ranges_list = [6000, 7000, 8000, 9000]
    maint_costs = [100, 120, 150, 180]
    limit_fleet = 1000
    limit_dist = 10000

    # Solve the problem
    solution = solve_aircraft_fleet(
        weights, values, efficiencies, ranges_list, maint_costs, limit_fleet, limit_dist
    )

    # Results output
    if solution:
        print("--- Selected Fleet Composition ---")
        for var_name, value in solution.items():
            if "aircraft" in str(var_name):
                print(f"{var_name}: {value}")
        print(f"\nTotal Fleet Value: {solution.total_value}")
    else:
        print("No feasible fleet configuration found.")