import pandas as pd

from vivarium.framework.state_machine import State, Machine, Transition
from vivarium.framework.utilities import rate_to_probability
from vivarium.framework.values import list_combiner, union_post_processor


class DiseaseTransition(Transition):

    def __init__(self, input_state, output_state, cause_key, measure, rate_name, **kwargs):
        super().__init__(input_state, output_state, probability_func=self._probability, **kwargs)
        self.cause_key = cause_key
        self.measure = measure
        self.rate_name = rate_name

    # noinspection PyAttributeOutsideInit
    def setup(self, builder):
        rate = builder.configuration[self.cause_key][self.measure]

        self.base_rate = lambda index: pd.Series(rate, index=index)
        self.transition_rate = builder.value.register_rate_producer(self.rate_name,
                                                                    source=self._risk_deleted_rate)
        self.joint_population_attributable_fraction = builder.value.register_value_producer(
            f'{self.rate_name}.population_attributable_fraction',
            source=lambda index: [pd.Series(0, index=index)],
            preferred_combiner=list_combiner,
            preferred_post_processor=union_post_processor)

    def _probability(self, index):
        effective_rate = self.transition_rate(index)
        return rate_to_probability(effective_rate)

    def _risk_deleted_rate(self, index):
        return self.base_rate(index) * (1 - self.joint_population_attributable_fraction(index))


class DiseaseState(State):

    def __init__(self, state_name, cause_key, with_excess_mortality=False):
        super().__init__(state_name)
        self.cause_key = cause_key
        self.with_excess_mortality = with_excess_mortality

    # noinspection PyAttributeOutsideInit
    def setup(self, builder):
        """Performs this component's simulation setup.

        Parameters
        ----------
        builder : `engine.Builder`
            Interface to several simulation tools.
        """
        super().setup(builder)
        if self.with_excess_mortality:
            self._excess_mortality_rate = builder.configuration[self.cause_key].excess_mortality_rate
        else:
            self._excess_mortality_rate = 0

        self.clock = builder.time.clock()

        self.excess_mortality_rate = builder.value.register_rate_producer(
            f'{self.state_id}.excess_mortality_rate',
            source=self.risk_deleted_excess_mortality_rate
        )
        self.excess_mortality_rate_paf = builder.value.register_value_producer(
            f'{self.state_id}.excess_mortality_rate.population_attributable_fraction',
            source=lambda index: [pd.Series(0, index=index)],
            preferred_combiner=list_combiner,
            preferred_post_processor=union_post_processor
        )

        builder.value.register_value_modifier('mortality_rate', self.add_in_excess_mortality)
        self.population_view = builder.population.get_view(
            [self._model], query=f"alive == 'alive' and {self._model} == '{self.state_id}'")

    def add_transition(self, output, measure, rate_name, **kwargs):
        t = DiseaseTransition(self, output, self.cause_key, measure, rate_name, **kwargs)
        self.transition_set.append(t)
        return t

    def risk_deleted_excess_mortality_rate(self, index):
        return pd.Series(self._excess_mortality_rate, index=index) * (1 - self.excess_mortality_rate_paf(index))

    def add_in_excess_mortality(self, index, mortality_rates):
        affected = self.population_view.get(index)
        mortality_rates.loc[affected.index] += self.excess_mortality_rate(affected.index)

        return mortality_rates


class DiseaseModel(Machine):

    def __init__(self, disease, initial_state, **kwargs):
        super().__init__(disease, **kwargs)
        self.initial_state = initial_state.state_id

    # noinspection PyAttributeOutsideInit
    def setup(self, builder):
        super().setup(builder)
        config = builder.configuration[self.state_column]
        # Reasonable approximation for short duration diseases.
        case_fatality_rate = config.excess_mortality_rate / (config.excess_mortality_rate + config.remission_rate)
        cause_specific_mortality_rate = config.incidence_rate * case_fatality_rate

        self.cause_specific_mortality_rate = builder.value.register_rate_producer(
            f'{self.state_column}.cause_specific_mortality_rate',
            source=lambda index: pd.Series(cause_specific_mortality_rate, index=index)
        )
        builder.value.register_value_modifier('mortality_rate', modifier=self.delete_cause_specific_mortality)
        builder.value.register_value_modifier('metrics', modifier=self.metrics)

        creates_columns = [self.state_column]
        builder.population.initializes_simulants(self.on_initialize_simulants,
                                                 creates_columns=creates_columns)
        self.population_view = builder.population.get_view(['age', 'sex', self.state_column])

        builder.event.register_listener('time_step', self.on_time_step)

    def on_initialize_simulants(self, pop_data):
        condition_column = pd.Series(self.initial_state, index=pop_data.index, name=self.state_column)
        self.population_view.update(condition_column)

    def on_time_step(self, event):
        self.transition(event.index, event.time)

    def delete_cause_specific_mortality(self, index, rates):
        return rates - self.cause_specific_mortality_rate(index)

    def metrics(self, index, metrics):
        pop = self.population_view.get(index, query="alive == 'alive'")
        metrics[self.state_column + '_prevalent_cases'] = len(pop[pop[self.state_column] != self.initial_state])
        return metrics


class SISDiseaseModel:

    configuration_defaults = {
        'disease': {
            'incidence_rate': 0.005,
            'remission_rate': 0.05,
            'excess_mortality_rate': 0.01,
        }
    }

    def __init__(self, disease_name):
        self._name = disease_name
        self.configuration_defaults = {disease_name: SISDiseaseModel.configuration_defaults['disease']}

        susceptible_state = DiseaseState(f'susceptible_to_{self._name}', self._name)
        infected_state = DiseaseState(f'infected_with_{self._name}', self._name, with_excess_mortality=True)

        susceptible_state.allow_self_transitions()
        susceptible_state.add_transition(infected_state, measure='incidence_rate',
                                         rate_name=f'{infected_state.state_id}.incidence_rate')
        infected_state.allow_self_transitions()
        infected_state.add_transition(susceptible_state, measure='remission_rate',
                                      rate_name=f'{infected_state.state_id}.remission_rate')

        model = DiseaseModel(self._name,
                             initial_state=susceptible_state,
                             states=[susceptible_state, infected_state])
        self._sub_components = [model]

    @property
    def name(self):
        return self._name

    @property
    def sub_components(self):
        return self._sub_components
