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

logger = setup_logger()


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

    def __init__(self, components, max_cost):
        self.components = components
        self.max_cost = max_cost
        self.cost_constraint = None
        self.m = None
        self.variables = []
        self.solution = []
        self.solution_gain = None
        self.solution_cost = None

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

        # 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)

        # 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)

        # 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
        self.m.optimize()

        if self.m.status == OptimizationStatus.INFEASIBLE:
            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
        ]

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