from mip import Model, xsum, maximize, BINARY, OptimizationStatus
from utils.logger import setup_logger
from typing import Mapping

logger = setup_logger()


class GainOptimiser:
    """
    This class is used to maximise gain, given a constrained cost
    """

    def __init__(
        self,
        components: list[list[Mapping[str, int | float | str]]],
        max_cost: float | int,
        max_gain: float | int | None,
        allow_slack: bool = True,
        verbose: bool = False
    ):
        """
        This function will try and maximise the gain, given a constrained cost. If we specific a max_gain, then the
        optimisation routine is constained to try not to exceed a maximum increase


        If the maximum gain (`max_gain`) is explicitly set to 0, the optimization routine interprets this as an
        instruction not to perform any optimization.

        :param components: List of components, where each component is a dictionary with keys "id", "cost" and "gain"
        :param max_cost: Maximum cost constraint
        :param max_gain: Maximum gain constraint
        :param allow_slack: If True, and the solution is infeasible, allows the model to use slack variables to relax
                            the cost constraint if the model. Defaults to True.
        :param verbose: If True, enables verbose logging
        """
        self.components = components
        self.max_cost = max_cost
        self.max_gain = max_gain
        self.cost_constraint = None
        self.max_gain_constraint = None
        self.m = None
        self.variables = []
        self.solution = []
        self.solution_gain = None
        self.solution_cost = None
        self.allow_slack = allow_slack
        self.verbose = verbose

    def setup(self):
        # Initialize Model
        self.m = Model("knapsack")

        # Set the verbosity level
        self.m.verbose = 1 if self.verbose else 0

        # Create variables
        self.variables = [
            [self.m.add_var(var_type=BINARY, name=str(component["id"])) for component in group] for group in
            self.components
        ]

        # This objective is the sum
        # gain_ig * x_ig, where gain_ig represents the gain for ith part in group g
        # and x_ig is the binary decision variable for the ith part in group g
        self.m.objective = maximize(
            xsum(
                component['gain'] * var for group, group_vars in zip(self.components, self.variables) for component, var
                in
                zip(group, group_vars)
            )
        )

        # This constrain ensures that sum of cost_ig * x_ig <= C, where cost_ig represents the cost for the ith
        # component
        # in group g, and x_ig is the binary decision variable for the ith component in group g
        cost_expression = xsum(
            item['cost'] * var for group, group_vars in zip(self.components, self.variables) for item, var in
            zip(group, group_vars)
        ) <= self.max_cost

        self.cost_constraint = self.m.add_constr(cost_expression)

        # Add an optional max gain constraint if max_gain is not None
        if self.max_gain is not None:
            max_gain_expression = xsum(
                component['gain'] * var for group, group_vars in zip(self.components, self.variables) for component, var
                in zip(group, group_vars)
            ) <= self.max_gain

            self.max_gain_constraint = self.m.add_constr(max_gain_expression)

        # This constraint ensures that at most one item from each group is selected
        # This is expressed by summing up the decision variables for each group and ensuring that the sum is <= 1
        for group_vars in self.variables:
            self.m += xsum(var for var in group_vars) <= 1

    def setup_slack(self):
        # Remove the original cost constraint
        self.m.remove(self.cost_constraint)

        if self.max_gain is not None:
            # Remove the original max gain constraint
            self.m.remove(self.max_gain_constraint)

        # Add slack variable
        s = self.m.add_var(lb=0)

        # Modify the constraint
        self.m += xsum(
            item['cost'] * var for group, group_vars in zip(self.components, self.variables) for item, var in
            zip(group, group_vars)
        ) + s <= self.max_cost

        # Modify the objective to penalize the use of slack
        penalty = -10000  # Negative penalty because we are maximizing
        self.m.objective = maximize(
            xsum(
                component['gain'] * var for group, group_vars in zip(self.components, self.variables) for component, var
                in
                zip(group, group_vars)
            ) + penalty * s
        )

    def solve(self):
        # Solve the problem

        if self.max_gain == 0:
            logger.info("Max gain is set to 0, no optimisation will be performed")
            # Nothing to do
            return

        self.m.optimize()

        solution = [
            item for group, group_vars in zip(self.components, self.variables) for item, var in zip(group, group_vars)
            if
            var.x >= 0.99
        ]

        if (self.m.status == OptimizationStatus.INFEASIBLE) or (
            (self.m.status == OptimizationStatus.OPTIMAL) and not len(solution)
        ):
            if self.allow_slack:
                # Turn off logging - too noisy
                # logger.info("We have an infeasible model, setting up slack model")
                self.setup_slack()
                self.m.optimize()
                solution = [
                    item for group, group_vars in zip(self.components, self.variables) for item, var in
                    zip(group, group_vars)
                    if
                    var.x >= 0.99
                ]
            else:
                logger.info("Infeasible but slack disabled - returning empty solution")

        self.solution = solution

        self.solution_gain = sum(component['gain'] for component in self.solution)
        self.solution_cost = sum([component['cost'] for component in self.solution])