"""
A set partitioning model of a wedding seating problem
Adaptation where an initial solution is given to solvers: CPLEX_CMD, GUROBI_CMD, COIN_CMD

Authors: Stuart Mitchell 2009, Franco Peschiera 2019
"""

import pulp

max_tables = 5
max_table_size = 4
guests = "A B C D E F G I J K L M N O P Q R".split()


def happiness(
    table: tuple[str]
    | tuple[str, str, str, str]
    | tuple[str, str]
    | tuple[str, str, str],
) -> int:
    """
    Find the happiness of the table
    - by calculating the maximum distance between the letters
    """
    return abs(ord(table[0]) - ord(table[-1]))


# create list of all possible tables
possible_tables = [tuple(c) for c in pulp.allcombinations(guests, max_table_size)]

seating_model = pulp.LpProblem("Wedding Seating Model", pulp.LpMinimize)

# create a binary variable to state that a table setting is used
_table_keys = ["_".join(t) for t in possible_tables]
vars_by_key = seating_model.add_variable_dict(
    "table_%s", (_table_keys,), 0, 1, pulp.LpInteger
)
x = {t: vars_by_key["_".join(t)] for t in possible_tables}

seating_model += pulp.lpSum([happiness(table) * x[table] for table in possible_tables])

# specify the maximum number of tables
seating_model += (
    pulp.lpSum([x[table] for table in possible_tables]) <= max_tables,
    "Maximum_number_of_tables",
)

# A guest must seated at one and only one table
for guest in guests:
    seating_model += (
        pulp.lpSum([x[table] for table in possible_tables if guest in table]) == 1,
        f"Must_seat_{guest}",
    )

# I've taken the optimal solution from a previous solving. x is the variable dictionary.
solution = {
    ("M", "N"): 1.0,
    ("E", "F", "G"): 1.0,
    ("A", "B", "C", "D"): 1.0,
    ("I", "J", "K", "L"): 1.0,
    ("O", "P", "Q", "R"): 1.0,
}
for k, v in solution.items():
    x[k].setInitialValue(v)

solver = pulp.COIN_CMD(msg=True, warmStart=True)
# solver = pulp.CPLEX_CMD(msg=True, warmStart=True)
# solver = pulp.GUROBI_CMD(msg=True, warmStart=True)
# solver = pulp.CPLEX_PY(msg=True, warmStart=True)
# solver = pulp.GUROBI(msg=True, warmStart=True)
seating_model.solve(solver)


print(f"The chosen tables are out of a total of {len(possible_tables)}:")
for table in possible_tables:
    if x[table].value() == 1.0:
        print(table)