from __future__ import annotations

import itertools
import math
from collections import defaultdict
from typing import Any

import pandas as pd
import pytest
from pytest_mock import MockerFixture

from vivarium import Component
from vivarium.framework.engine import Builder, SimulationContext
from vivarium.framework.population import PopulationManager, SimulantData
from vivarium.framework.values import ValuesManager

# FIXME: Streamline with already-existing classes in tests/helpers.py
PIE_COL_NAMES = ["pie", "pi"]
PIES = ["apple", "chocolate", "pecan", "pumpkin", "sweet_potato"]
PIS = [math.pi**i for i in range(1, 11)]
PIE_RECORDS = [(pie, pi) for pie, pi in itertools.product(PIES, PIS)]
PIE_DF = pd.DataFrame(data=PIE_RECORDS, columns=PIE_COL_NAMES)
CUBE_COL_NAMES = ["cube", "cube_string"]
CUBE = [i**3 for i in range(len(PIE_RECORDS))]
CUBE_STRING = [str(i**3) for i in range(len(PIE_RECORDS))]
CUBE_DF = pd.DataFrame(
    zip(CUBE, CUBE_STRING),
    columns=CUBE_COL_NAMES,
    index=PIE_DF.index,
)


class PieComponent(Component):
    def setup(self, builder: Builder) -> None:
        builder.population.register_initializer(
            initializer=self.make_pie, columns=PIE_COL_NAMES
        )

    def make_pie(self, pop_data: SimulantData) -> None:
        self.population_view.initialize(self.get_initial_state(pop_data.index))

    def get_initial_state(self, index: pd.Index[int]) -> pd.DataFrame:
        return PIE_DF


class CubeComponent(Component):
    def setup(self, builder: Builder) -> None:
        builder.population.register_initializer(
            initializer=self.cubify, columns=CUBE_COL_NAMES
        )

    def cubify(self, pop_data: SimulantData) -> None:
        self.population_view.initialize(self.get_initial_state(pop_data.index))

    def get_initial_state(self, index: pd.Index[int]) -> pd.DataFrame:
        return CUBE_DF


@pytest.fixture(scope="function")
def pies_and_cubes_pop_mgr(mocker: MockerFixture) -> PopulationManager:
    """A mocked PopulationManager with some private columns set up.

    This fixture is tied directly to the PieComponent and CubeComponent helper classes.

    """

    class _PopulationManager(PopulationManager):
        def __init__(self) -> None:
            super().__init__()
            self._private_columns: pd.DataFrame = pd.concat([PIE_DF, CUBE_DF], axis=1)

        def _add_constraint(self, *args: Any, **kwargs: Any) -> None:
            pass

    mgr = _PopulationManager()

    # Use SimulationContext just for builder and mock as appropriate
    sim = SimulationContext()
    builder = sim._builder
    mocker.patch.object(ValuesManager, "logger", mocker.Mock(), create=True)
    mocker.patch.object(ValuesManager, "resources", mocker.Mock(), create=True)
    mocker.patch.object(ValuesManager, "add_constraint", mocker.Mock(), create=True)
    mocker.patch.object(ValuesManager, "_population_mgr", mgr, create=True)
    mocked_attribute_pipelines = {}
    sim._lifecycle.set_state("setup")
    mgr.setup(builder)
    sim._lifecycle.set_state("post_setup")
    sim._lifecycle.set_state("population_creation")

    for col in mgr._private_columns.columns:
        mocked_attribute_pipelines[col] = mocker.Mock()
    mgr._attribute_pipelines = mocked_attribute_pipelines
    mgr._private_column_metadata = defaultdict(
        list,
        {
            "pie_component": PIE_COL_NAMES,
            "cube_component": CUBE_COL_NAMES,
        },
    )
    # Change lifecycle phase to ensure tracked queries are applied appropriately
    mocker.patch.object(mgr, "get_current_state", lambda: "on_time_step")
    return mgr
