from mip import Model, xsum, minimize, BINARY


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

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

        self.solution_cost = None
        self.solution_gain = 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
        ]

        # 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
        self.m += 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

        # 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 solve(self):
        # Solve the problem
        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])