Model/recommendations/optimiser/StrategicOptimiser.py

from enum import Enum
from mip import OptimizationStatus
from typing import Mapping, Optional, TypedDict, List
from recommendations.optimiser.CostOptimiser import CostOptimiser
from recommendations.optimiser.GainOptimiser import GainOptimiser


class Measure(TypedDict):
    id: str
    cost: float
    gain: float


class Strategies(Enum):
    CASE_1_TRY_MIN_COST_WITH_CONSTRAINTS = "case_1_try_min_cost_with_constraints"
    CASE_1_SOLVE_MAX_GAIN_UNDER_BUDGET = "case_1_solve_max_gain_under_budget"
    CASE_2_SOLVE_MAX_GAIN_UNDER_BUDGET = "case_2_solve_max_gain_under_budget"
    CASE_3_SOLVE_MIN_COST_FOR_TARGET = "case_3_solve_min_cost_for_target"


class StrategicOptimiser:
    """
    Domain-level optimiser implementing logical optimisation logic.

    Behaviour:

    1) If both budget and target_gain are provided:
        - Minimise cost subject to:
            gain >= target_gain
            cost <= budget
        - If infeasible:
            maximise gain subject to cost <= budget

    2) If only budget is provided:
        - Maximise gain under budget

    3) If only target_gain is provided:
        - Minimise cost to achieve gain

    """

    def __init__(
        self,
        components: list[list[Mapping[str, int | float | str]]],
        budget: Optional[float] = None,
        target_gain: Optional[float] = None,
        allow_slack: bool = False,
        verbose: bool = False,
    ) -> None:

        if not components:
            raise ValueError("Components cannot be empty.")

        if budget is None and target_gain is None:
            raise ValueError("At least one of budget or target_gain must be provided.")

        self.components = components
        self.budget = budget
        self.target_gain = target_gain
        self.verbose = verbose
        self.allow_slack = allow_slack

        self.solution: List[Measure] = []
        self.solution_cost: float = 0.0
        self.solution_gain: float = 0.0

        # For debugging purposes, we keep a record of which option was selected
        self.strategy_used: Optional[Strategies] = None

    def solve(self) -> None:
        """
        Primary entry point for solving the optimisation problem based on the provided budget and target gain.
        :return:
        """

        # Case 1: budget + target
        if self.budget is not None and self.target_gain is not None:
            # Given:
            # Budget B
            # Target gain G
            #
            # We want the solution to:
            #
            # Primary problem (P1)
            # min cost
            # subject to
            #
            # gain >= G
            # cost <= B
            # multiple-choice constraints
            #
            # If (P1) is feasible → that solution is exactly what you want.
            # If (P1) is infeasible → solve the following problem (P2):
            #
            # max gain
            # subject to
            #
            # cost <= B
            if self._try_min_cost_with_constraints():
                # Keep a record of the strategy used to solve the problem, for debugging purposes
                self.strategy_used = Strategies.CASE_1_TRY_MIN_COST_WITH_CONSTRAINTS
                return
            self._solve_max_gain_under_budget()
            self.strategy_used = Strategies.CASE_1_SOLVE_MAX_GAIN_UNDER_BUDGET
            return

        # Case 2: budget only
        if self.budget is not None:
            self._solve_max_gain_under_budget()
            self.strategy_used = Strategies.CASE_2_SOLVE_MAX_GAIN_UNDER_BUDGET
            return

        # Case 3: target only
        self._solve_min_cost_for_target()
        self.strategy_used = Strategies.CASE_3_SOLVE_MIN_COST_FOR_TARGET
        return

    # ---------------------------------------------------------
    # Internal Functions
    # ---------------------------------------------------------

    def _try_min_cost_with_constraints(self) -> bool:
        """
        Try to minimise cost while satisfying:
            gain >= target_gain
            cost <= budget
        """

        opt = CostOptimiser(
            self.components,
            min_gain=self.target_gain,
            verbose=self.verbose,
            allow_slack=self.allow_slack
        )

        opt.setup()
        opt.add_budget_constraint(self.budget)
        opt.solve()

        if opt.m.status == OptimizationStatus.INFEASIBLE:
            return False

        self._store_solution(opt.solution)
        return True

    def _solve_max_gain_under_budget(self) -> None:

        opt = GainOptimiser(
            self.components,
            max_cost=self.budget,
            max_gain=None,
            allow_slack=self.allow_slack,
            verbose=self.verbose
        )

        opt.setup()
        opt.solve()

        self._store_solution(opt.solution)

    def _solve_min_cost_for_target(self) -> None:

        opt = CostOptimiser(
            self.components,
            min_gain=self.target_gain,
            verbose=self.verbose
        )

        opt.setup()
        opt.solve()

        self._store_solution(opt.solution)

    def _store_solution(self, solution: List[Measure]) -> None:
        self.solution = solution
        self.solution_cost = sum(m["cost"] for m in solution)
        self.solution_gain = sum(m["gain"] for m in solution)