from mip import Model, xsum, maximize, BINARY

# 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

C = 4000

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

# Initialize Model
m = Model("knapsack")

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

# 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
m.objective = maximize(
    xsum(
        component['gain'] * var for group, group_vars in zip(groups, vars) 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
m += xsum(item['cost'] * var for group, group_vars in zip(groups, vars) for item, var in zip(group, group_vars)) <= C

# 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 vars:
    m += xsum(var for var in group_vars) <= 1

# Solve the problem
m.optimize()

# Get the selected items
selected_items = [
    item for group, group_vars in zip(groups, vars) for item, var in zip(group, group_vars) if var.x >= 0.99
]

total_gain = m.objective.x
actual_cost = sum([component['cost'] for component in selected_items])
print("Selected items:", selected_items)