main.py

import argparse
import math
import warnings
from decimal import Decimal, ROUND_HALF_UP
from pathlib import Path

import pandas as pd
from ortools.sat.python import cp_model

WIDTH = 336  # mm
WIDTH_TOLERANCE_MIN = 3  # mm
WIDTH_TOLERANCE_MAX = 4  # mm
PATH_TO_COLLECTION = "mod_collection.csv"
CATEGORY_WIDTHS = {"half": 320}
LENGTH_SCALE = 10
DEFAULT_SOLVER_TIME_LIMIT = 60


def parse_args(args=None):
    parser = argparse.ArgumentParser(
        description="Sort a board game collection into organizer cubes."
    )
    parser.add_argument(
        "--path",
        "-p",
        default=PATH_TO_COLLECTION,
        help=f"Path to the modified collection CSV file (default: {PATH_TO_COLLECTION}).",
    )
    parser.add_argument(
        "--max-time",
        type=float,
        default=DEFAULT_SOLVER_TIME_LIMIT,
        help="Maximum CP-SAT solve time per category in seconds (default: 60).",
    )
    return parser.parse_args(args)


def csv2dict(path):
    bg_dict = pd.read_csv(path, sep=";")
    bg_dict["length"] = pd.to_numeric(bg_dict["length"], errors="coerce")
    bg_dict["avgweight"] = pd.to_numeric(bg_dict["avgweight"], errors="coerce")
    bg_dict = bg_dict.dropna(subset=["length"]).copy()
    zero_weight_games = bg_dict.loc[
        bg_dict["avgweight"] == 0, "objectname"
    ].tolist()
    if zero_weight_games:
        warnings.warn(
            "Games with weight 0: " + ", ".join(zero_weight_games),
            UserWarning,
            stacklevel=2,
        )
    bg_dict["category"] = (
        bg_dict["category"]
        .replace(r"^\s*$", "default", regex=True)
        .fillna("default")
    )
    # Group by 'category'
    grouped = bg_dict.groupby('category')
    result_dict = {}
    # Iterate over each group
    for category, group in grouped:
        # Initialize the nested dictionary for the current category
        category_dict = {}
        for index, row in group.iterrows():
            # Create the nested dictionary with 'objectname' as key and [avgweight, length] as value
            category_dict[row['objectname']] = [row['avgweight'], row['length']]
        # Add the nested dictionary to the result dictionary
        result_dict[category] = category_dict
    return result_dict


def width_for_category(category):
    return CATEGORY_WIDTHS.get(category, WIDTH)


def _to_length_units(value):
    return int(
        (Decimal(str(value)) * LENGTH_SCALE).to_integral_value(
            rounding=ROUND_HALF_UP
        )
    )


def solve_cubes(
    bg_dict,
    width=WIDTH,
    width_tolerance_min=WIDTH_TOLERANCE_MIN,
    width_tolerance_max=WIDTH_TOLERANCE_MAX,
    max_time_seconds=DEFAULT_SOLVER_TIME_LIMIT,
):
    if not bg_dict:
        return [], []
    if max_time_seconds <= 0:
        raise ValueError("Solver time limit must be positive")

    min_used_width = _to_length_units(width - width_tolerance_max)
    max_used_width = _to_length_units(width - width_tolerance_min)
    if min_used_width <= 0 or min_used_width > max_used_width:
        raise ValueError("Width tolerances produce an invalid usable-width range")

    games = sorted(bg_dict, key=lambda game: bg_dict[game][0], reverse=True)
    lengths = [_to_length_units(bg_dict[game][1]) for game in games]
    for game, length in zip(games, lengths):
        if length <= 0:
            raise ValueError(f"{game!r} must have a positive length")
        if length > max_used_width:
            raise ValueError(
                f"{game!r} is wider than the maximum usable cube width"
            )

    cube_count = math.ceil(sum(lengths) / max_used_width)
    cube_indexes = range(cube_count)
    model = cp_model.CpModel()
    assigned = {
        (game_index, cube_index): model.new_bool_var(
            f"game_{game_index}_cube_{cube_index}"
        )
        for game_index in range(len(games))
        for cube_index in cube_indexes
    }
    game_is_assigned = []

    for game_index in range(len(games)):
        is_assigned = model.new_bool_var(f"assigned_{game_index}")
        model.add(
            sum(
                assigned[game_index, cube_index]
                for cube_index in cube_indexes
            )
            == is_assigned
        )
        game_is_assigned.append(is_assigned)

    cube_is_used = []
    for cube_index in cube_indexes:
        is_used = model.new_bool_var(f"used_{cube_index}")
        games_in_cube = sum(
            assigned[game_index, cube_index]
            for game_index in range(len(games))
        )
        used_width = sum(
            lengths[game_index] * assigned[game_index, cube_index]
            for game_index in range(len(games))
        )
        model.add(games_in_cube >= is_used)
        model.add(games_in_cube <= len(games) * is_used)
        model.add(used_width <= max_used_width * is_used)
        cube_is_used.append(is_used)

    for cube_index in range(cube_count - 1):
        model.add(cube_is_used[cube_index] >= cube_is_used[cube_index + 1])
        used_width = sum(
            lengths[game_index] * assigned[game_index, cube_index]
            for game_index in range(len(games))
        )
        model.add(used_width >= min_used_width * cube_is_used[cube_index + 1])

    target_width = (min_used_width + max_used_width) // 2
    prefix_width = 0
    displacement_vars = []
    for game_index, length in enumerate(lengths):
        target_cube = min(
            (prefix_width + length // 2) // target_width,
            cube_count - 1,
        )
        displacement = model.new_int_var(
            0, cube_count - 1, f"displacement_{game_index}"
        )
        model.add(
            displacement
            == sum(
                abs(cube_index - target_cube)
                * assigned[game_index, cube_index]
                for cube_index in cube_indexes
            )
        )
        displacement_vars.append(displacement)
        prefix_width += length

    max_displacement = model.new_int_var(
        0, cube_count - 1, "max_displacement"
    )
    model.add_max_equality(max_displacement, displacement_vars)
    max_total_displacement = len(games) * (cube_count - 1)
    max_ordering_cost = (
        (cube_count - 1) * (max_total_displacement + 1)
        + max_total_displacement
    )
    assigned_count = sum(game_is_assigned)
    assigned_width = sum(
        length * game_is_assigned[game_index]
        for game_index, length in enumerate(lengths)
    )
    missing_count = len(games) - assigned_count
    missing_width = sum(lengths) - assigned_width
    packing_cost = missing_count * (sum(lengths) + 1) + missing_width
    ordering_cost = (
        max_displacement * (max_total_displacement + 1)
        + sum(displacement_vars)
    )
    model.minimize(packing_cost)

    packing_solver = cp_model.CpSolver()
    packing_solver.parameters.max_time_in_seconds = max_time_seconds * 0.65
    packing_solver.parameters.num_search_workers = 8
    packing_status = packing_solver.solve(model)
    if packing_status not in (cp_model.FEASIBLE, cp_model.OPTIMAL):
        raise RuntimeError(
            "CP-SAT did not find a packing before the solver time limit"
        )

    model.add(packing_cost <= packing_solver.value(packing_cost))
    for variable in assigned.values():
        model.add_hint(variable, packing_solver.value(variable))
    for variable in game_is_assigned + cube_is_used + displacement_vars:
        model.add_hint(variable, packing_solver.value(variable))
    model.add_hint(max_displacement, packing_solver.value(max_displacement))
    model.minimize(
        packing_cost * (max_ordering_cost + 1) + ordering_cost
    )

    solver = cp_model.CpSolver()
    solver.parameters.max_time_in_seconds = max_time_seconds * 0.35
    solver.parameters.num_search_workers = 8
    status = solver.solve(model)
    if status not in (cp_model.FEASIBLE, cp_model.OPTIMAL):
        solver = packing_solver

    cubes = []
    for cube_index in cube_indexes:
        if not solver.value(cube_is_used[cube_index]):
            break
        cube = [
            game
            for game_index, game in enumerate(games)
            if solver.value(assigned[game_index, cube_index])
        ]
        cubes.append(cube)
    missing_games = [
        game
        for game_index, game in enumerate(games)
        if not solver.value(game_is_assigned[game_index])
    ]
    return cubes, missing_games


def write_result(
    sorted_cubes, missing_games, bg_dict, nb_games, file_name, category
):
    with open(file_name, "a") as f:
        f.write("Category: " + category + "\n")
        f.write("Total number of games: " + str(nb_games))
        f.write("\nCP-SAT sorted cubes")
        for s in sorted_cubes:
            games = [f"{game} ({bg_dict[game][0]})" for game in s]
            f.write("\n" + str(games))
        f.write("\n\nMissing Games\n")
        missing = [
            f"{game} ({bg_dict[game][0]})" for game in missing_games
        ]
        f.write(str(missing))
        f.write("\n\n")


def main(argv=None):
    args = parse_args(argv)
    bg_dict = csv2dict(Path(args.path))
    with open("result.txt", "w") as f:
        f.write("")
    for k, v in bg_dict.items():
        sorted_cubes, missing_games = solve_cubes(
            v,
            width=width_for_category(k),
            max_time_seconds=args.max_time,
        )
        write_result(sorted_cubes, missing_games, v, len(v), "result.txt", k)


if __name__ == "__main__":
    main()