Tatami/tests/test_masterplan.py

"""Integration tests for the masterplan pipeline (group building + assignment)."""

import numpy as np
import pytest

from tatami.classes import Group
from tatami.traveltimes import reduce_distance_matrix
from tatami.tatami_masterplan import (
    assign_courses,
    fast_total_time,
    get_after_party_group,
    get_courses,
    get_masterplan,
    load_csv_to_participants,
)
from conftest import make_participants, make_timedelta_matrix


def full_matrix(participants, after_party_group, seed=0):
    """Synthetic uuid-indexed Timedelta matrix over participants + after party."""
    all_people = participants + [after_party_group.main_member]
    rng = np.random.default_rng(seed)
    n = len(all_people)
    seconds = rng.integers(60, 1200, size=(n, n)).astype(float)
    np.fill_diagonal(seconds, 0)
    return make_timedelta_matrix(all_people, seconds)


class TestGetAfterPartyGroup:
    def test_single_member_with_no_penalty(self):
        group = get_after_party_group("party street 1")
        assert len(group.members) == 1
        assert group.main_member.address == "party street 1"
        assert group.main_member.get_penalty().total_seconds() == 0


class TestLoadCsv:
    def test_parses_tab_separated_file(self, tmp_path):
        csv = tmp_path / "config.csv"
        csv.write_text(
            "name\taddress\tphone\tkitchen_size\tallergies\n"
            "Alice\tStreet 1\t111\t8\tnone\n"
            "Bob\tStreet 2\t222\t5\tpeanuts\n"
        )
        participants = load_csv_to_participants(str(csv))
        assert [p.name for p in participants] == ["Alice", "Bob"]
        assert participants[0].address == "Street 1"
        assert participants[1].kitchen_size == 5


class TestAssignCourses:
    def test_assigns_cycling_courses_and_hosts(self):
        groups = [Group(members=[p]) for p in make_participants(6)]
        courses = ["starter", "main", "dessert"] * 2
        assign_courses(groups, courses)
        assert [g.course for g in groups] == courses
        for i, group in enumerate(groups):
            expected_host_uuids = {groups[j].uuid for j in get_courses(i, len(groups))}
            assert {h.uuid for h in group.hosts} == expected_host_uuids
            assert group.uuid in {h.uuid for h in group.hosts}


class TestGetMasterplan:
    @pytest.mark.parametrize("n", [6, 12, 18])
    def test_group_count_and_coverage(self, n):
        participants = make_participants(n, kitchen_sizes=list(np.linspace(0, 10, n)))
        after_party = get_after_party_group("party street")
        matrix = full_matrix(participants, after_party)

        group_dicts, participant_dicts = get_masterplan(
            participants, matrix, after_party
        )

        expected_groups = 3 * (n // 6)
        assert len(group_dicts) == expected_groups
        assert len(participant_dicts) == n

        # Every participant ends up in exactly one group.
        assigned = [uuid for g in group_dicts for uuid in g["members"]]
        assert len(assigned) == n
        assert set(assigned) == {p.uuid for p in participants}

    def test_topology_of_output(self):
        participants = make_participants(18, kitchen_sizes=list(np.linspace(0, 10, 18)))
        after_party = get_after_party_group("party street")
        matrix = full_matrix(participants, after_party)

        group_dicts, _ = get_masterplan(participants, matrix, after_party)
        by_uuid = {g["uuid"]: g for g in group_dicts}

        # Courses are balanced across the three slots.
        course_counts = {}
        for g in group_dicts:
            course_counts[g["course"]] = course_counts.get(g["course"], 0) + 1
        assert course_counts == {"starter": 3, "main": 3, "dessert": 3}

        for g in group_dicts:
            host_courses = sorted(by_uuid[h]["course"] for h in g["hosts"])
            assert host_courses == ["dessert", "main", "starter"]
            assert g["uuid"] in g["hosts"]  # a group hosts its own course

    def test_does_not_mutate_input_participant_order(self):
        participants = make_participants(12)
        original_order = list(participants)
        after_party = get_after_party_group("party street")
        matrix = full_matrix(participants, after_party)
        get_masterplan(participants, matrix, after_party)
        assert participants == original_order


class TestCostConsistency:
    def test_fast_total_time_matches_group_get_total_time(self):
        # The optimizer cost (fast_total_time on the penalty-baked reduced matrix)
        # must equal the sum of per-group route times computed independently by
        # Group.get_total_time on the raw matrix.
        n = 6
        participants = make_participants(n, kitchen_sizes=[10, 8, 6, 4, 2, 0])
        after_party = get_after_party_group("party street")
        groups = [Group(members=[p]) for p in participants]
        courses = ["starter", "main", "dessert"] * (n // 3)
        assign_courses(groups, courses)

        full = full_matrix(participants, after_party, seed=5)

        reduced = reduce_distance_matrix(full, [*groups, after_party])
        reduced_seconds = np.array(
            [[pd_to_seconds(x) for x in row] for row in reduced.to_numpy()],
            dtype=float,
        )
        optimizer_cost = fast_total_time(reduced_seconds, list(range(n)))

        independent_cost = sum(
            g.get_total_time(full, after_party).total_seconds() for g in groups
        )
        assert optimizer_cost == pytest.approx(independent_cost)


def pd_to_seconds(value):
    import pandas as pd

    return pd.Timedelta(value).total_seconds()