import itertools
import math
import re
from collections.abc import Callable
from datetime import timedelta
from typing import Any

import numpy as np
import pandas as pd
import pytest
from layered_config_tree.main import LayeredConfigTree
from loguru import logger
from pandas.core.groupby.generic import DataFrameGroupBy
from pytest_mock import MockerFixture

from tests.framework.results.helpers import (
    BASE_POPULATION,
    FAMILIARS,
    HOUSE_CATEGORIES,
    NAME,
    NAME_COLUMNS,
    sorting_hat_serial,
    sorting_hat_vectorized,
    verify_stratification_added,
)
from vivarium.framework.event import Event
from vivarium.framework.lifecycle import lifecycle_states
from vivarium.framework.results import VALUE_COLUMN
from vivarium.framework.results.context import ResultsContext
from vivarium.framework.results.interface import PopulationFilter
from vivarium.framework.results.observation import AddingObservation, ConcatenatingObservation
from vivarium.framework.results.stratification import Stratification, get_mapped_col_name
from vivarium.types import ScalarMapper, VectorMapper


def _aggregate_state_person_time(x: pd.DataFrame) -> float:
    """Helper aggregator function for observation testing"""
    return len(x) * (28 / 365.25)


@pytest.fixture
def event() -> Event:
    return Event(
        name=lifecycle_states.COLLECT_METRICS,
        index=pd.Index([0]),
        user_data={},
        time=0,
        step_size=1,
    )


@pytest.mark.parametrize(
    "mapper, is_vectorized",
    [
        (sorting_hat_vectorized, True),
        (sorting_hat_serial, False),
    ],
    ids=["vectorized_mapper", "non-vectorized_mapper"],
)
def test_add_stratification_mappers(
    mapper: VectorMapper | ScalarMapper, is_vectorized: bool, mocker: MockerFixture
) -> None:
    ctx = ResultsContext()
    mocker.patch.object(ctx, "excluded_categories", {})

    assert NAME not in ctx.stratifications

    ctx.add_stratification(
        name=NAME,
        requires_attributes=NAME_COLUMNS,
        categories=HOUSE_CATEGORIES,
        excluded_categories=None,
        mapper=mapper,
        is_vectorized=is_vectorized,
    )
    assert verify_stratification_added(
        stratifications=ctx.stratifications,
        name=NAME,
        requires_attributes=NAME_COLUMNS,
        categories=HOUSE_CATEGORIES,
        excluded_categories=[],
        mapper=mapper,
        is_vectorized=is_vectorized,
    )


@pytest.mark.parametrize(
    "excluded_categories",
    [
        [],
        HOUSE_CATEGORIES[:1],
        HOUSE_CATEGORIES[:2],
        HOUSE_CATEGORIES[:3],
    ],
    ids=[
        "no_excluded_categories",
        "one_excluded_category",
        "two_excluded_categories",
        "all_but_one_excluded_categories",
    ],
)
def test_add_stratification_excluded_categories(
    excluded_categories: list[str], mocker: MockerFixture
) -> None:
    ctx = ResultsContext()
    builder = mocker.Mock()
    builder.configuration.stratification = LayeredConfigTree(
        {"default": [], "excluded_categories": {NAME: excluded_categories}}
    )
    builder.logging.get_logger.return_value = logger
    ctx.setup(builder)

    assert NAME not in ctx.stratifications

    ctx.add_stratification(
        name=NAME,
        requires_attributes=NAME_COLUMNS,
        categories=HOUSE_CATEGORIES,
        excluded_categories=excluded_categories,
        mapper=sorting_hat_vectorized,
        is_vectorized=True,
    )

    assert verify_stratification_added(
        stratifications=ctx.stratifications,
        name=NAME,
        requires_attributes=NAME_COLUMNS,
        categories=HOUSE_CATEGORIES,
        excluded_categories=excluded_categories,
        mapper=sorting_hat_vectorized,
        is_vectorized=True,
    )


@pytest.mark.parametrize(
    "name, categories, excluded_categories, msg_match",
    [
        (
            "duplicate_name",
            HOUSE_CATEGORIES,
            [],
            "Stratification name 'duplicate_name' is already used",
        ),
        (
            NAME,
            HOUSE_CATEGORIES + ["slytherin"],
            [],
            f"Found duplicate categories in stratification '{NAME}': ['slytherin']",
        ),
        (
            NAME,
            HOUSE_CATEGORIES + ["gryffindor", "slytherin"],
            [],
            f"Found duplicate categories in stratification '{NAME}': ['gryffindor', 'slytherin']",
        ),
        (
            NAME,
            HOUSE_CATEGORIES,
            ["gryfflepuff"],
            "Excluded categories {'gryfflepuff'} not found in categories",
        ),
    ],
    ids=[
        "duplicate_name",
        "duplicate_category",
        "duplicate_categories",
        "unknown_excluded_category",
    ],
)
def test_add_stratification_raises(
    name: str,
    categories: list[str],
    excluded_categories: list[str],
    msg_match: str,
    mocker: MockerFixture,
) -> None:
    ctx = ResultsContext()
    mocker.patch.object(ctx, "excluded_categories", {name: excluded_categories})
    # Register a stratification to test against duplicate stratifications
    ctx.add_stratification(
        name="duplicate_name",
        requires_attributes=["foo"],
        categories=["bar"],
        excluded_categories=None,
        mapper=sorting_hat_serial,
        is_vectorized=False,
    )
    with pytest.raises(ValueError, match=re.escape(msg_match)):
        ctx.add_stratification(
            name=name,
            requires_attributes=NAME_COLUMNS,
            categories=categories,
            excluded_categories=excluded_categories,
            mapper=sorting_hat_vectorized,
            is_vectorized=True,
        )


@pytest.mark.parametrize(
    "kwargs",
    [
        {
            "name": "living_person_time",
            "population_filter": PopulationFilter("is_alive == True and undead == False"),
            "requires_attributes": ["is_alive", "undead"],
            "when": lifecycle_states.COLLECT_METRICS,
        },
        {
            "name": "undead_person_time",
            "population_filter": PopulationFilter("undead == True"),
            "requires_attributes": ["undead"],
            "when": lifecycle_states.TIME_STEP_PREPARE,
        },
    ],
    ids=["valid_on_collect_metrics", "valid_on_time_step__prepare"],
)
def test_register_observation(kwargs: dict[str, Any]) -> None:
    ctx = ResultsContext()
    assert len(ctx.grouped_observations) == 0
    kwargs["results_formatter"] = lambda: None
    kwargs["stratifications"] = tuple()
    kwargs["aggregator_sources"] = []
    kwargs["aggregator"] = len
    kwargs["requires_attributes"] = []
    ctx.register_observation(
        observation_type=AddingObservation,
        **kwargs,
    )
    assert len(ctx.grouped_observations) == 1


def test_register_observation_duplicate_name_raises() -> None:
    ctx = ResultsContext()
    ctx.register_observation(
        observation_type=AddingObservation,
        name="some-observation-name",
        population_filter=PopulationFilter("some-pop-filter"),
        when="some-when",
        requires_attributes=[],
        results_formatter=lambda df: df,
        stratifications=(),
        aggregator_sources=[],
        aggregator=len,
    )
    with pytest.raises(
        ValueError, match="Observation name 'some-observation-name' is already used: "
    ):
        # register a different observation but w/ the same name
        ctx.register_observation(
            observation_type=ConcatenatingObservation,
            name="some-observation-name",
            population_filter=PopulationFilter("some-other-pop-filter"),
            when="some-other-when",
            requires_attributes=[],
            stratifications=None,
        )


@pytest.mark.parametrize(
    "aggregator_sources, aggregator, stratifications",
    [
        ([], len, ["house", "familiar"]),
        (["power_level"], sum, ["house", "familiar"]),
        ([], _aggregate_state_person_time, ["house", "familiar"]),
        ([], len, ["house"]),
        (["power_level"], sum, ["house"]),
        ([], _aggregate_state_person_time, ["house"]),
    ],
    ids=[
        "len_aggregator_two_stratifications",
        "sum_aggregator_two_stratifications",
        "custom_aggregator_two_stratifications",
        "len_aggregator_one_stratification",
        "sum_aggregator_one_stratification",
        "custom_aggregator_one_stratification",
    ],
)
def test_adding_observation_gather_results(
    aggregator_sources: list[str],
    aggregator: Callable[..., int | float],
    stratifications: list[str],
    event: Event,
    mocker: MockerFixture,
) -> None:
    """Test cases where every stratification is in gather_results. Checks for
    existence and correctness of results"""
    ctx = ResultsContext()
    mocker.patch.object(ctx, "get_tracked_query", return_value="", create=True)

    # Generate population DataFrame
    population = BASE_POPULATION.copy()
    for stratification in stratifications:
        population[get_mapped_col_name(stratification)] = population[stratification]

    # Set up stratifications
    if "house" in stratifications:
        ctx.add_stratification(
            name="house",
            requires_attributes=["house"],
            categories=HOUSE_CATEGORIES,
            excluded_categories=None,
            mapper=None,
            is_vectorized=True,
        )
    if "familiar" in stratifications:
        ctx.add_stratification(
            name="familiar",
            requires_attributes=["familiar"],
            categories=FAMILIARS,
            excluded_categories=None,
            mapper=None,
            is_vectorized=True,
        )
    observation = ctx.register_observation(
        observation_type=AddingObservation,
        name="foo",
        population_filter=PopulationFilter(),
        requires_attributes=aggregator_sources,
        aggregator_sources=aggregator_sources,
        aggregator=aggregator,
        stratifications=tuple(stratifications),
        when=lifecycle_states.COLLECT_METRICS,
        results_formatter=lambda: None,
    )

    groups = population.groupby(stratifications)
    if aggregator == sum:
        power_level_sums = groups[aggregator_sources].sum().squeeze()
        assert len(power_level_sums.unique()) == 1
        expected_result = power_level_sums.iat[0]
    else:
        group_sizes = groups.size()
        assert len(group_sizes.unique()) == 1
        num_stratifications = group_sizes.iat[0]
        expected_result = (
            num_stratifications if aggregator == len else num_stratifications * 28 / 365.25
        )

    i = 0
    for result, _measure, _updater in ctx.gather_results(
        population, event.name, [observation]
    ):
        assert result is not None
        assert all(
            math.isclose(actual_result, expected_result, rel_tol=0.0001)
            for actual_result in result.values
        )
        i += 1
    assert i == 1


def test_concatenating_observation_gather_results(
    event: Event, mocker: MockerFixture
) -> None:

    ctx = ResultsContext()
    mocker.patch.object(ctx, "get_tracked_query", return_value="", create=True)

    # Generate population DataFrame
    population = BASE_POPULATION.copy()
    # Mock out some extra columns that would be produced by the manager's _prepare_population() method
    population["current_time"] = pd.Timestamp(year=2045, month=1, day=1, hour=12)
    population["event_step_size"] = timedelta(days=28)
    population["event_time"] = pd.Timestamp(year=2045, month=1, day=1, hour=12) + timedelta(
        days=28
    )

    lifecycle_state = lifecycle_states.COLLECT_METRICS
    population_filter = PopulationFilter(query="house=='hufflepuff'")
    included_cols = ["familiar", "house"]
    observation = ctx.register_observation(
        observation_type=ConcatenatingObservation,
        name="foo",
        population_filter=population_filter,
        when=lifecycle_state,
        requires_attributes=included_cols,
        results_formatter=lambda _, __: pd.DataFrame(),
        stratifications=None,
    )

    filtered_pop = population.query(population_filter.query)

    i = 0
    for result, _measure, _updater in ctx.gather_results(
        population, event.name, [observation]
    ):
        assert result is not None
        assert result.equals(filtered_pop[["event_time"] + included_cols])
        i += 1
    assert i == 1


@pytest.mark.parametrize(
    "name, aggregator_sources, aggregator, stratifications",
    [
        ("wizard_count", [], len, ["house", "familiar"]),
        ("power_level_total", ["power_level"], sum, ["house", "familiar"]),
        (
            "wizard_time",
            [],
            _aggregate_state_person_time,
            ["house", "familiar"],
        ),
        ("wizard_count", [], len, ["familiar"]),
        ("power_level_total", ["power_level"], sum, ["familiar"]),
        (
            "wizard_time",
            [],
            _aggregate_state_person_time,
            ["familiar"],
        ),
    ],
    ids=[
        "len_aggregator_two_stratifications",
        "sum_aggregator_two_stratifications",
        "custom_aggregator_two_stratifications",
        "len_aggregator_one_stratification",
        "sum_aggregator_one_stratification",
        "custom_aggregator_one_stratification",
    ],
)
def test_gather_results_partial_stratifications_in_results(
    name: str,
    aggregator_sources: list[str],
    aggregator: Callable[..., int | float],
    stratifications: list[str],
    event: Event,
    mocker: MockerFixture,
) -> None:
    """Test cases where not all stratifications are observed for gather_results. This looks for existence of
    unobserved stratifications and ensures their values are 0"""
    ctx = ResultsContext()
    mocker.patch.object(ctx, "get_tracked_query", return_value="", create=True)

    # Generate population DataFrame
    population = BASE_POPULATION.copy()
    # Remove an entire category from a stratification
    population = population[population["familiar"] != "unladen_swallow"].reset_index()

    # Set up stratifications
    if "house" in stratifications:
        ctx.add_stratification(
            name="house",
            requires_attributes=["house"],
            categories=HOUSE_CATEGORIES,
            excluded_categories=None,
            mapper=None,
            is_vectorized=True,
        )
        population[get_mapped_col_name("house")] = population["house"].astype(
            pd.CategoricalDtype(categories=HOUSE_CATEGORIES, ordered=True)
        )
    if "familiar" in stratifications:
        ctx.add_stratification(
            name="familiar",
            requires_attributes=["familiar"],
            categories=FAMILIARS,
            excluded_categories=None,
            mapper=None,
            is_vectorized=True,
        )
        population[get_mapped_col_name("familiar")] = population["familiar"].astype(
            pd.CategoricalDtype(categories=FAMILIARS, ordered=True)
        )

    observation = ctx.register_observation(
        observation_type=AddingObservation,
        name=name,
        population_filter=PopulationFilter(),
        requires_attributes=aggregator_sources,
        aggregator_sources=aggregator_sources,
        aggregator=aggregator,
        stratifications=tuple(stratifications),
        when=lifecycle_states.COLLECT_METRICS,
        results_formatter=lambda: None,
    )

    for results, _measure, _formatter in ctx.gather_results(
        population, event.name, [observation]
    ):
        assert results is not None
        unladen_results = results.reset_index().query('familiar=="unladen_swallow"')
        assert len(unladen_results) > 0
        assert (unladen_results[VALUE_COLUMN] == 0).all()


def test_gather_results_with_empty_pop_filter(event: Event, mocker: MockerFixture) -> None:
    """Test case where pop_filter filters to an empty population. gather_results
    should return None.
    """
    ctx = ResultsContext()
    mocker.patch.object(ctx, "get_tracked_query", return_value="", create=True)

    # Generate population DataFrame
    population = BASE_POPULATION.copy()

    lifecycle_state = lifecycle_states.COLLECT_METRICS
    observation = ctx.register_observation(
        observation_type=AddingObservation,
        name="wizard_count",
        population_filter=PopulationFilter("house == 'durmstrang'"),
        requires_attributes=["house"],
        aggregator_sources=[],
        aggregator=len,
        stratifications=tuple(),
        when=lifecycle_state,
        results_formatter=lambda: None,
    )

    for result, _measure, _updater in ctx.gather_results(
        population, event.name, [observation]
    ):
        assert not result


def test_gather_results_with_no_stratifications(event: Event, mocker: MockerFixture) -> None:
    """Test case where we have no stratifications. gather_results should return one value."""
    ctx = ResultsContext()
    mocker.patch.object(ctx, "get_tracked_query", return_value="", create=True)

    # Generate population DataFrame
    population = BASE_POPULATION.copy()

    lifecycle_state = lifecycle_states.COLLECT_METRICS
    observation = ctx.register_observation(
        observation_type=AddingObservation,
        name="wizard_count",
        population_filter=PopulationFilter(),
        requires_attributes=[],
        aggregator_sources=None,
        aggregator=len,
        stratifications=tuple(),
        when=lifecycle_state,
        results_formatter=lambda: None,
    )

    assert len(ctx.stratifications) == 0
    assert (
        len(
            list(
                result
                for result, _measure, _updater in ctx.gather_results(
                    population, event.name, [observation]
                )
            )
        )
        == 1
    )


def test_bad_aggregator_stratification(event: Event, mocker: MockerFixture) -> None:
    """Test if an exception gets raised when a stratification that doesn't
    exist is attempted to be used, as expected."""
    ctx = ResultsContext()
    mocker.patch.object(ctx, "get_tracked_query", return_value="", create=True)

    # Generate population DataFrame
    population = BASE_POPULATION.copy()
    lifecycle_state = lifecycle_states.COLLECT_METRICS

    # Set up stratifications
    ctx.add_stratification(
        name="house",
        requires_attributes=["house"],
        categories=HOUSE_CATEGORIES,
        excluded_categories=None,
        mapper=None,
        is_vectorized=True,
    )
    ctx.add_stratification(
        name="familiar",
        requires_attributes=["familiar"],
        categories=FAMILIARS,
        excluded_categories=None,
        mapper=None,
        is_vectorized=True,
    )
    observation = ctx.register_observation(
        observation_type=AddingObservation,
        name="this_shouldnt_work",
        population_filter=PopulationFilter(),
        requires_attributes=[],
        aggregator_sources=[],
        aggregator=sum,
        stratifications=("house", "height"),  # `height` is not a stratification
        when=lifecycle_state,
        results_formatter=lambda: None,
    )

    with pytest.raises(KeyError, match="height"):
        for result, _measure, _updater in ctx.gather_results(
            population, event.name, [observation]
        ):
            print(result)


@pytest.mark.parametrize(
    ["lifecycle_state", "time", "expected_observations"],
    [
        (lifecycle_states.COLLECT_METRICS, 0, ["obs1", "obs3"]),
        (lifecycle_states.TIME_STEP_PREPARE, 0, ["obs2"]),
        (lifecycle_states.COLLECT_METRICS, 1, ["obs1"]),
    ],
    ids=["collect_metrics_time_0", "time_step_prepare_time_0", "collect_metrics_time_1"],
)
def test_get_observations(
    lifecycle_state: str, time: int, expected_observations: list[str]
) -> None:
    ctx = ResultsContext()
    register_observation_kwargs = {
        "observation_type": AddingObservation,
        "population_filter": PopulationFilter(),
        "requires_attributes": [],
        "results_formatter": lambda: None,
        "stratifications": (),
        "aggregator_sources": None,
        "aggregator": len,
    }

    ctx.register_observation(
        name="obs1", when=lifecycle_states.COLLECT_METRICS, **register_observation_kwargs  # type: ignore[arg-type]
    )
    ctx.register_observation(
        name="obs2", when=lifecycle_states.TIME_STEP_PREPARE, **register_observation_kwargs  # type: ignore[arg-type]
    )
    ctx.register_observation(
        name="obs3",
        when=lifecycle_states.COLLECT_METRICS,
        to_observe=lambda event: event.time == 0,
        **register_observation_kwargs,  # type: ignore[arg-type]
    )

    event = Event(
        name=lifecycle_state, index=pd.Index([0]), user_data={}, time=time, step_size=1
    )

    assert [obs.name for obs in ctx.get_observations(event)] == expected_observations


@pytest.mark.parametrize(
    "observation_names, stratification_names, expected_resources",
    [
        (["obs1", "obs2"], ["strat1", "strat2"], {"x", "y", "z", "v"}),
        (["obs3"], ["strat1", "strat2"], {"x", "y", "w", "v"}),
        ([], ["strat1"], {"x", "y"}),
        (["obs2"], [], {"y", "z"}),
        ([], [], set()),
    ],
    ids=[
        "obs_and_strat_with_overlap",
        "obs_and_strat_without_overlap",
        "no_observations",
        "no_stratifications",
        "neither",
    ],
)
@pytest.mark.parametrize("include_untracked", [True, False])
def test_get_required_attributes(
    observation_names: list[str],
    stratification_names: list[str],
    expected_resources: set[str],
    include_untracked: bool,
    mocker: MockerFixture,
) -> None:
    ctx = ResultsContext()
    mocker.patch.object(ctx, "get_tracked_query", return_value='foo == "bar"', create=True)

    all_observations = {}
    register_observation_kwargs = {
        "observation_type": AddingObservation,
        "population_filter": PopulationFilter(include_untracked=include_untracked),
        "when": lifecycle_states.COLLECT_METRICS,
        "results_formatter": lambda: None,
        "stratifications": (),
        "aggregator_sources": None,
        "aggregator": len,
    }

    all_observations["obs1"] = ctx.register_observation(
        name="obs1",
        requires_attributes=["x", "y"],
        **register_observation_kwargs,  # type: ignore[arg-type]
    )
    all_observations["obs2"] = ctx.register_observation(
        name="obs2",
        requires_attributes=["y", "z"],
        **register_observation_kwargs,  # type: ignore[arg-type]
    )
    all_observations["obs3"] = ctx.register_observation(
        name="obs3",
        requires_attributes=["w"],
        **register_observation_kwargs,  # type: ignore[arg-type]
    )

    all_stratifications = {}
    stratification_kwargs = {
        "categories": ["cat1", "cat2"],
        "excluded_categories": ["cat3"],
        "mapper": lambda df: df["a"] + df["b"],
        "is_vectorized": True,
    }
    all_stratifications["strat1"] = Stratification(
        name="strat1",
        requires_attributes=["x", "y"],
        **stratification_kwargs,  # type: ignore[arg-type]
    )
    all_stratifications["strat2"] = Stratification(
        name="strat2",
        requires_attributes=["x", "v"],
        **stratification_kwargs,  # type: ignore[arg-type]
    )

    observations = [all_observations[name] for name in observation_names]
    stratifications = [all_stratifications[name] for name in stratification_names]

    actual_columns = ctx.get_required_attributes(observations, stratifications)
    if observations and not include_untracked:
        expected_resources = expected_resources.union({"foo"})
    assert set(actual_columns) == expected_resources


def test_get_required_attributes_columns_from_query(mocker: MockerFixture) -> None:
    """Tests that columns used in queries are regardless of requires_attributes."""

    ctx = ResultsContext()
    mocker.patch.object(
        ctx, "get_tracked_query", return_value="pet == 'cat' and lives < 9", create=True
    )

    observation = ctx.register_observation(
        observation_type=AddingObservation,
        name="obs_with_query",
        population_filter=PopulationFilter(
            "color in ['black', 'white'] or name == 'Garfield'"
        ),
        when=lifecycle_states.COLLECT_METRICS,
        requires_attributes=["foo", "bar"],  # does NOT include any of the query columns
        results_formatter=lambda: None,
        stratifications=(),
        aggregator_sources=None,
        aggregator=len,
    )

    assert set(ctx.get_required_attributes([observation], [])) == {
        "pet",
        "lives",
        "color",
        "name",
        "foo",
        "bar",
    }


@pytest.mark.parametrize(
    "pop_filter",
    ['familiar=="cat"', 'familiar=="spaghetti_yeti"', ""],
    ids=["pop_filter", "pop_filter_empties_dataframe", "no_pop_filter"],
)
@pytest.mark.parametrize("include_untracked", [True, False])
@pytest.mark.parametrize("tracked_query", ["house=='hufflepuff'", "house=='whitehouse'", ""])
def test__filter_population(
    pop_filter: str, include_untracked: bool, tracked_query: str, mocker: MockerFixture
) -> None:
    population = BASE_POPULATION.copy()
    ctx = ResultsContext()
    mocker.patch.object(ctx, "get_tracked_query", return_value=tracked_query, create=True)

    filtered_pop = ctx._filter_population(
        population=population,
        population_filter=PopulationFilter(pop_filter, include_untracked=include_untracked),
    )
    expected = population.copy()
    if pop_filter:
        familiar = pop_filter.split("==")[1].strip('"')
        expected = expected[expected["familiar"] == familiar]
    if not include_untracked and tracked_query:
        house = tracked_query.split("==")[1].strip("'")
        expected = expected[expected["house"] == house]
    assert filtered_pop.equals(expected)


@pytest.mark.parametrize(
    "stratifications",
    [tuple(), ("new_col1",), ("new_col1", "new_col2")],
    ids=[
        "no_stratifications",
        "single_excluded_stratification",
        "two_excluded_stratifications",
    ],
)
def test__drop_na_stratifications(stratifications: tuple[str, ...]) -> None:
    population = BASE_POPULATION.copy()
    population["new_col1"] = "new_value1"
    population["new_col2"] = "new_value2"
    population.loc[population["new_col1"].notna(), "new_col2"] = np.nan
    # Add on the post-stratified columns
    for stratification in stratifications:
        mapped_col = f"{stratification}_mapped_values"
        population[mapped_col] = population[stratification]

    filtered_pop = ResultsContext()._drop_na_stratifications(
        population=population, stratification_names=stratifications
    )
    expected = population.copy()
    for stratification in stratifications:
        expected = expected[expected[stratification].notna()]
    assert filtered_pop.equals(expected)


@pytest.mark.parametrize(
    "stratifications, values",
    [
        (("familiar",), [FAMILIARS]),
        (("familiar", "house"), [FAMILIARS, HOUSE_CATEGORIES]),
        ((), "foo"),
    ],
)
def test__get_groups(stratifications: tuple[str, ...], values: str | list[list[str]]) -> None:
    filtered_pop = BASE_POPULATION.copy()
    # Generate the post-stratified columns
    for stratification in stratifications:
        mapped_col = f"{stratification}_mapped_values"
        filtered_pop[mapped_col] = filtered_pop[stratification]
    groups = ResultsContext()._get_groups(
        stratifications=stratifications, filtered_pop=filtered_pop
    )
    assert isinstance(groups, DataFrameGroupBy)
    if stratifications:
        combinations = set(itertools.product(*values))
        if len(values) == 1:
            # convert from set of tuples to set of strings
            combinations = set(comb[0] for comb in combinations)
        # Check that all familiars exist
        assert set(groups.groups.keys()) == combinations
        # Check that the entire population is included
        assert sum([len(value) for value in groups.groups.values()]) == len(BASE_POPULATION)
    else:
        item = groups.groups.popitem()
        # Check that there are no other groups
        assert not groups.groups
        # Check that the group is 'all' and includes the entire population
        key, val = item
        assert key == "all"
        assert val.equals(BASE_POPULATION.index)
