import numpy as np
import pandas as pd


class Population:
    configuration_defaults = {
        "population": {
            "colors": ["red", "blue"],
        }
    }

    def __init__(self):
        self.name = "population"

    def setup(self, builder):
        self.colors = builder.configuration.population.colors

        columns_created = ["color", "entrance_time"]
        builder.population.initializes_simulants(
            self.on_initialize_simulants, columns_created
        )
        self.population_view = builder.population.get_view(columns_created)

    def on_initialize_simulants(self, pop_data):
        new_population = pd.DataFrame(
            {
                "color": np.random.choice(self.colors, len(pop_data.index)),
                "entrance_time": pop_data.creation_time,
            },
            index=pop_data.index,
        )
        self.population_view.update(new_population)
