from mip import Model, xsum, maximize, BINARY
from pprint import pprint

# Example parts
wall = [
    {"id": 1, "cost": 2000, "gain": 5, "type": "wall"},
    {"id": 2, "cost": 2300, "gain": 6, "type": "wall"}
]

floor = [
    {"id": 1, "cost": 1500, "gain": 3, "type": "floor"},
    {"id": 2, "cost": 1600, "gain": 3.1, "type": "floor"}
]

roof = [
    {"id": 1, "cost": 1000, "gain": 2, "type": "roof"},
    {"id": 2, "cost": 1100, "gain": 2.3, "type": "roof"}
]

# To solve this, we are solving a constrained Knapsack problem
# Maximize sum(gain_g . x_g) for g in groups
# subject to sum(cost_g . x_g) <= C
# subject to sum(x_g) <= 1 for g in groups
# x_g in {0, 1} for g in groups
#
# The first sum, which is the objective of the optimisation provlem, ensures that we are maximising the gain
# for the selected parts
# The second sum (and the first constraint) ensures that the cost of the selected parts is less than or equal to C
# The third sum (and the second constraint) ensures that at most one part from each group is selected
# The last constraint ensures that the decision variables are binary

# group all the parts
components = [wall, floor, roof]


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

    def __init__(self, components, max_cost):
        self.components = components
        self.max_cost = max_cost
        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
        ]

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

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

        # 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_gain = self.m.objective.x
        self.solution_cost = sum([component['cost'] for component in self.solution])


opt = GainOptimiser(components, max_cost=4000)

# Setup the knackpack problem
# This sets the objective & contraints
opt.setup()

# Solve the problem
opt.solve()

pprint(opt.solution)
print("total cost:", opt.solution_cost)
print("total gain:", opt.solution_gain)

# A bigger problem:
wall = [
    {"id": 1, "cost": 2000, "gain": 5, "type": "wall"},
    {"id": 2, "cost": 2300, "gain": 6, "type": "wall"},
    {"id": 3, "cost": 2200, "gain": 5.5, "type": "wall"},
    {"id": 4, "cost": 2500, "gain": 6.2, "type": "wall"},
    {"id": 5, "cost": 2100, "gain": 5.1, "type": "wall"},
    {"id": 6, "cost": 2400, "gain": 6.1, "type": "wall"},
    {"id": 7, "cost": 2000, "gain": 5.2, "type": "wall"}
]

floor = [
    {"id": 1, "cost": 1500, "gain": 3, "type": "floor"},
    {"id": 2, "cost": 1600, "gain": 3.1, "type": "floor"},
    {"id": 3, "cost": 1550, "gain": 3.2, "type": "floor"},
    {"id": 4, "cost": 1650, "gain": 3.3, "type": "floor"},
    {"id": 5, "cost": 1500, "gain": 3.4, "type": "floor"},
    {"id": 6, "cost": 1550, "gain": 3.5, "type": "floor"},
    {"id": 7, "cost": 1600, "gain": 3.6, "type": "floor"}
]

roof = [
    {"id": 1, "cost": 1000, "gain": 2, "type": "roof"},
    {"id": 2, "cost": 1100, "gain": 2.3, "type": "roof"},
    {"id": 3, "cost": 1200, "gain": 2.6, "type": "roof"},
    {"id": 4, "cost": 1300, "gain": 2.9, "type": "roof"},
    {"id": 5, "cost": 1100, "gain": 2.5, "type": "roof"},
    {"id": 6, "cost": 1200, "gain": 2.7, "type": "roof"},
    {"id": 7, "cost": 1300, "gain": 2.8, "type": "roof"}
]

heating = [
    {"id": 1, "cost": 3000, "gain": 7, "type": "heating"},
    {"id": 2, "cost": 3200, "gain": 7.2, "type": "heating"},
    {"id": 3, "cost": 3100, "gain": 7.1, "type": "heating"},
    {"id": 4, "cost": 3300, "gain": 7.3, "type": "heating"},
    {"id": 5, "cost": 3000, "gain": 7.4, "type": "heating"}
]

hot_water = [
    {"id": 1, "cost": 2500, "gain": 6.5, "type": "hot water"},
    {"id": 2, "cost": 2600, "gain": 6.6, "type": "hot water"},
    {"id": 3, "cost": 2500, "gain": 6.7, "type": "hot water"},
    {"id": 4, "cost": 2700, "gain": 6.8, "type": "hot water"},
    {"id": 5, "cost": 2500, "gain": 6.9, "type": "hot water"}
]

solar = [
    {"id": 1, "cost": 5000, "gain": 10, "type": "solar"},
    {"id": 2, "cost": 5500, "gain": 11, "type": "solar"},
    {"id": 3, "cost": 5300, "gain": 10.5, "type": "solar"},
    {"id": 4, "cost": 5200, "gain": 10.2, "type": "solar"},
    {"id": 5, "cost": 5400, "gain": 10.8, "type": "solar"}
]

heat_pumps = [
    {"id": 1, "cost": 4000, "gain": 9, "type": "heat pumps"},
    {"id": 2, "cost": 4200, "gain": 9.2, "type": "heat pumps"},
    {"id": 3, "cost": 4100, "gain": 9.1, "type": "heat pumps"},
    {"id": 4, "cost": 4300, "gain": 9.3, "type": "heat pumps"},
    {"id": 5, "cost": 4000, "gain": 9.4, "type": "heat pumps"}
]

components2 = [
    wall,
    floor,
    roof,
    heating,
    hot_water,
    solar,
    heat_pumps
]

opt2 = GainOptimiser(components2, max_cost=22000)

# Setup
opt2.setup()

# Solve the problem
opt2.solve()

pprint(opt2.solution)
print("total cost:", opt2.solution_cost)
print("total gain:", opt2.solution_gain)