from typing import Dict

import pandas as pd

from vivarium.framework.engine import Builder


class Observer:
    configuration_defaults = {
        "mortality": {
            "life_expectancy": 80,
        }
    }

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

    # noinspection PyAttributeOutsideInit
    def setup(self, builder: Builder):
        self.life_expectancy = builder.configuration.mortality.life_expectancy
        self.population_view = builder.population.get_view(["age", "alive"])

        builder.value.register_value_modifier("metrics", self.metrics)

    def metrics(self, index: pd.Index, metrics: Dict):
        pop = self.population_view.get(index)
        metrics["total_population_alive"] = len(pop[pop.alive == "alive"])
        metrics["total_population_dead"] = len(pop[pop.alive == "dead"])

        metrics["years_of_life_lost"] = (
            self.life_expectancy - pop.age[pop.alive == "dead"]
        ).sum()

        return metrics
