Tatami/src/tatami/tatami_masterplan.py

from tatami.classes import Participant, Group
from tatami.traveltimes import reduce_distance_matrix, get_participant_distance_matrix
import pandas as pd
import numpy as np
import random
from tqdm import tqdm


def get_after_party_group(address: str) -> Group:
    """
    Create a group for the after party with a single participant.
    """
    participant = Participant(
        name="After Party", address=address, phone="", kitchen_size=10, allergies=""
    )
    return Group(members=[participant], main_member=0)


def get_masterplan(
    participants: list[Participant],
    distance_matrix: pd.DataFrame,
    after_party_group: Group,
) -> tuple[list[dict], list[dict]]:
    groups_per_course = np.floor(len(participants) / 6).astype(int)
    participants = sorted(
        participants,
        key=lambda x: x.get_after_party_time(distance_matrix, after_party_group),
    )
    hosts = participants[: 3 * groups_per_course]
    semi_hosts = participants[3 * groups_per_course :]
    courses = ["starter", "main", "dessert"] * groups_per_course

    groups = []
    for host in hosts:
        group = Group(members=[host])
        groups.append(group)

    random.shuffle(semi_hosts)
    for i, member in enumerate(semi_hosts):
        groups[i % len(groups)].add_member(member)

    distance_matrix = reduce_distance_matrix(
        distance_matrix, [*groups, after_party_group]
    )
    best_order = run_simulated_annealing(
        groups,
        distance_matrix,
        initial_temperature=1000,
        cooling_rate=0.99,
        max_iterations=10000,
        multiprocessing=1,
    )
    assign_courses(best_order, courses)
    group_dicts = [group.dict() for group in best_order]
    participant_dicts = [participant.dict() for participant in participants]
    return group_dicts, participant_dicts


def assign_courses(groups: list[Group], courses: list[str]) -> None:
    """
    Assign courses to groups.

    Courses are assigned to every group first, then hosts are wired up: a group's
    hosts are sorted by course (``sort_hosts``), so all course assignments must be
    in place before any ``set_hosts`` call, otherwise hosts whose course is still
    unset get mis-ordered.
    """
    for group, course in zip(groups, courses):
        group.set_course(course)

    for i, group in enumerate(groups):
        hosts = [groups[j] for j in get_courses(i, len(groups))]
        group.set_hosts(hosts)


def run_simulated_annealing(
    groups: list[Group],
    reduced_distance_matrix: pd.DataFrame,
    initial_temperature: float,
    cooling_rate: float,
    max_iterations: int,
    multiprocessing: int = 1,
) -> list[Group]:
    """
    Run simulated annealing to find the optimal order of groups.
    """
    group_indices = [int(i) for i in range(len(groups))]

    # The reduced matrix holds pd.Timedelta values; convert to a float matrix of
    # seconds so it can be used numerically in the acceptance criterion
    # (np.exp cannot operate on Timedelta objects).
    distance_matrix = np.array(
        [
            [pd.Timedelta(x).total_seconds() for x in row]
            for row in reduced_distance_matrix.to_numpy()
        ],
        dtype=float,
    )

    if multiprocessing > 1:
        raise NotImplementedError("Multiprocessing is not implemented yet.")
    else:
        # Run simulated annealing without multiprocessing
        best_order = simulated_annealing(
            distance_matrix,
            group_indices,
            initial_temperature,
            cooling_rate,
            max_iterations,
        )
        best_ordererd_groups = [groups[i] for i in best_order]
        return best_ordererd_groups


def simulated_annealing(
    distance_matrix: np.ndarray,
    group_indices: list[int],
    initial_temperature: float,
    cooling_rate: float,
    max_iterations: int,
) -> list[int]:
    """
    Simulated annealing over assignments of groups to rotation slots.

    ``solution[slot]`` is the group index placed in that slot. Each iteration
    proposes a random two-slot swap and accepts it with the Boltzmann
    probability; the best solution seen is tracked and returned.
    """
    current = group_indices.copy()
    np.random.shuffle(current)
    current_cost = fast_total_time(distance_matrix, current)

    best = current.copy()
    best_cost = current_cost

    # Nothing to optimize with fewer than two slots.
    if len(current) < 2:
        return best

    current_temperature = initial_temperature

    for iteration in tqdm(range(max_iterations)):
        try:
            candidate = current.copy()
            i, j = np.random.choice(len(candidate), size=2, replace=False)
            candidate[i], candidate[j] = candidate[j], candidate[i]
            candidate_cost = fast_total_time(distance_matrix, candidate)

            delta = candidate_cost - current_cost
            if delta <= 0 or np.random.random() < np.exp(-delta / current_temperature):
                current, current_cost = candidate, candidate_cost
                if current_cost < best_cost:
                    best, best_cost = current.copy(), current_cost

            if iteration % 100 == 0:
                print(
                    f"Iteration {iteration}: current = {current_cost:.0f}s, "
                    f"best = {best_cost:.0f}s"
                )
            current_temperature *= cooling_rate
        except KeyboardInterrupt:
            print("Simulation interrupted. Returning best solution so far.")
            break
    return best


def get_courses(
    group_index: int,
    total_meetings: int,
):
    a = group_index
    b = (group_index + 1) % total_meetings
    c = (group_index - 4) % total_meetings
    order = sorted((a, b, c), key=lambda x: x % 3)
    return order


def fast_total_time(distance_matrix: np.ndarray, solution: list[int]) -> float:
    """Calculates the total travel time for a given slot-to-group assignment.

    The rotation topology is defined over *slots* ``0..n-1`` via ``get_courses``;
    ``solution[slot]`` gives the group placed in that slot. For each slot the
    occupying group travels starter-host -> main-host -> dessert-host -> after
    party, and the durations of those legs are summed across all slots.

    Args:
        distance_matrix (np.ndarray): 2D array where ``distance_matrix[i][j]`` is
            the travel time (seconds) from group ``i`` to group ``j``. Shape
            ``(n+1, n+1)``; the last row/column is the after-party location.
        solution (list[int]): Permutation mapping each slot to a group index.
    """
    n = len(solution)
    after_party_idx = n  # Since distance_matrix is (n+1)x(n+1)

    total_time = 0.0
    total_meetings = n

    for slot in range(n):
        a, b, c = get_courses(slot, total_meetings)
        ga, gb, gc = solution[a], solution[b], solution[c]
        total_time += distance_matrix[ga][gb]
        total_time += distance_matrix[gb][gc]
        total_time += distance_matrix[gc][after_party_idx]

    return total_time


def load_csv_to_participants(file_path: str) -> list[Participant]:
    """
    Load participants from a CSV file.
    """
    df = pd.read_csv(file_path, sep="\t")
    participants = []
    for _, row in df.iterrows():
        participant = Participant(
            name=row["name"],
            address=row["address"],
            phone=row["phone"],
            kitchen_size=row["kitchen_size"],
            allergies=row["allergies"],
        )
        participants.append(participant)
    return participants


if __name__ == "__main__":
    # Example usage
    participants = load_csv_to_participants("test-config.csv")
    after_party_group = get_after_party_group(
        "Sebastian-Kneipp-Straße 6, 76131 Karlsruhe"
    )
    distance_matrix = get_participant_distance_matrix(
        [*participants, after_party_group.main_member], mode="BICYCLE"
    )
    masterplan = get_masterplan(participants, distance_matrix, after_party_group)
    print(masterplan)
    print("Masterplan generated successfully.")