from mip import Model, xsum, minimize, BINARY, OptimizationStatus
from utils.logger import setup_logger

logger = setup_logger()


class CostOptimiser:
    """
    This class is used to minimise cost, given a constrained minimum gain
    """

    # We add an optional buffer to the minimum gain to allow for slack in the optimisation
    BUFFER = 0.2

    def __init__(self, components, min_gain, verbose=False):
        self.components = components
        self.min_gain = min_gain
        self.gain_constraint = None
        self.m = None
        self.variables = []
        self.solution = []

        self.solution_cost = None
        self.solution_gain = None
        self.verbose = verbose

    @classmethod
    def calculate_sap_gain_with_slack(cls, min_gain: int | float):
        """
        Adds a small amount of buffer to the minimum gain, to account for possible error in SAP predictions
        :param min_gain: Numerical value for the minimum gain
        :return:
        """
        if min_gain == 0:
            return min_gain
        elif min_gain <= 5:
            return min_gain + 0.25
        elif min_gain <= 20:
            return min_gain + 0.5
        else:
            return min_gain + 0.75

    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
        ]

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

        # Add constraints
        # This constrain ensures that sum of gain_ig * x_ig >= min_gain, where gain_ig represents the gain for the ith
        # component
        # in group g, and x_ig is the binary decision variable for the ith component in group g
        gain_expression = xsum(
            item['gain'] * var for group, group_vars in zip(self.components, self.variables) for item, var in
            zip(group, group_vars)
        ) >= self.min_gain

        self.gain_constraint = self.m.add_constr(gain_expression)

        # At most one item from each group
        # 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 gain constraint
        self.m.remove(self.gain_constraint)
        # Add slack variable
        s = self.m.add_var(lb=0)

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

        # Modify the objective to penalize the use of slack
        penalty = 10000  # you can adjust this based on how much you want to penalize the use of slack
        self.m.objective = minimize(
            xsum(
                component['cost'] * 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
        self.m.optimize()

        if self.m.status == OptimizationStatus.INFEASIBLE:
            # Turn off logging - too noisy
            # logger.info("We have an infeasible model, setting up slack model")
            self.setup_slack()
            self.m.optimize()

        self.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
        ]

        # Get the selected items
        self.solution_cost = self.m.objective.x
        self.solution_gain = sum([component['gain'] for component in self.solution])