#!/usr/bin/env python3
"""Solve a multiple knapsack problem using the CP-SAT solver."""
from ortools.sat.python import cp_model

def main():
    data = {}
    data["weights"] = [48, 30, 42, 36, 36, 48, 42, 42, 36, 24, 30, 30, 42, 36, 36]
    data["values"] = [10, 30, 25, 50, 35, 30, 15, 40, 30, 35, 45, 10, 20, 30, 25]
    assert len(data["weights"]) == len(data["values"])
    data["num_items"] = len(data["weights"])
    data["all_items"] = range(data["num_items"])

    colors = [10, 11, 12]
    data["colors"] = [10, 11, 10, 10, 12, 10, 11, 10, 10, 11, 10, 10, 11, 10, 10]
    for c in data["colors"]:
        assert c in colors, f"Color {c} unknow !"
    assert len(data["colors"]) == len(data["values"])
    data["num_colors"] = len(data["colors"])
    data["all_colors"] = range(data["num_colors"])

    data["bin_capacities"] = [100, 100, 100, 100, 100]
    data["num_bins"] = len(data["bin_capacities"])
    data["all_bins"] = range(data["num_bins"])

    # Create the CP-SAT model.
    model = cp_model.CpModel()

    # Variables.
    # x[i, b] = 1 if item i is packed in bin b.
    x = {}
    for b in data["all_bins"]:
            for i in data["all_items"]:
                x[i, b] = model.new_bool_var(f"x_{i}_{b}")
    
    # y[c, b] = 1 if color c is packed in bin b.
    y = {}
    for b in data["all_bins"]:
            for c in colors:
                y[c, b] = model.new_bool_var(f"y_{c}_{b}")

    # Constraints.
    # Each item is assigned to at most one bin.
    for i in data["all_items"]:
        model.add(sum(x[i, b] for b in data["all_bins"]) <= 1)

    # The amount packed in each bin cannot exceed its capacity.
    for b in data["all_bins"]:
        model.add(
            sum(x[i, b] * data["weights"][i] for i in data["all_items"])
            <= data["bin_capacities"][b]
        )

    # a color is present in a bin if at least one item of this color is present in the bin
    for b in data["all_bins"]:
            for c in colors:
                model.add(sum(x[i, b] for i in data["all_items"] if data["colors"][i] == c) >= 1).only_enforce_if(y[c, b])
                model.add(sum(x[i, b] for i in data["all_items"] if data["colors"][i] == c) == 0).only_enforce_if(y[c, b].negated())

    # Each bin has at most one color
    for b in data["all_bins"]:
        model.add(sum(y[c, b] for c in colors) <= 1)

    # Objective.
    # Maximize total value of packed items.
    model.maximize(
       sum(x[i, b] * data['values'][i] for i in data["all_items"] for b in data["all_bins"]) 
    )

    # Solve 
    solver = cp_model.CpSolver()
    status = solver.solve(model)

    # Print solution
    if status == cp_model.OPTIMAL or status == cp_model.FEASIBLE:
        print(f"Total packed value: {solver.objective_value}")
        total_weight = 0
        total_value = 0
        for b in data["all_bins"]:
            print(f"Bin {b}")
            bin_weight = 0
            bin_value = 0
            bin_colors = set()
            for i in data["all_items"]:
                if solver.boolean_value(x[i, b]):
                    print(f"Item {i} weight: {data['weights'][i]} value: {data['values'][i]}  color: {data['colors'][i]}")
                    bin_weight += data["weights"][i]
                    bin_value += data["values"][i]
                    bin_colors.add(data["colors"][i])
            print(f"Packed bin weight: {bin_weight}")
            print(f"Packed bin value: {bin_value}")
            print(f"Color bin value: {bin_colors}\n")
            total_weight += bin_weight
            total_value += bin_value
        print(f"Total packed weight: {total_weight}")
        print(f"Total packed value: {total_value}")

    else:
        print("The problem does not have an optimal solution.")


if __name__ == "__main__":
    main()