from dataclasses import dataclass

import phylogenie._typings as pgt
from phylogenie.skyline import SkylineParameterLike, skyline_parameter
from phylogenie.treesimulator.core import STATE_KEY, Model, StochasticEvent
from phylogenie.treesimulator.parameterizations.common import (
    Recovery,
    Sampling,
    StateDependentEvent,
)

INITIAL_SUSCEPTIBLES_KEY = "population_size"
SUSCEPTIBLES_KEY = "susceptibles"
RECOVERED_KEY = "recovered"

LAST_ANTIGENIC_ESCAPE_ID_KEY = "last_antigenic_escape_id"
LAST_FITNESS_GAIN_ID_KEY = "mutation_id"

TRANSMISSION_RATES_KEY = "transmission_rates"
ESCAPE_FRACTIONS_KEY = "escape_fractions"


def _make_state(antigenic_id: int, mutation_id: int) -> str:
    return f"{antigenic_id}.{mutation_id}"


def _get_mutation_ids(state: str) -> tuple[int, int]:
    ant, mut = state.split(".")
    return int(ant), int(mut)


@dataclass(kw_only=True)
class AntigenicEscapeTransmission(StochasticEvent):
    state: str
    escape_fraction: float

    def _get_other_recovered(self, model: Model) -> float:
        antigenic_escape_id, _ = _get_mutation_ids(self.state)
        return (
            model[INITIAL_SUSCEPTIBLES_KEY]
            - model[SUSCEPTIBLES_KEY]
            - model.count_active_nodes()
            - model[RECOVERED_KEY][antigenic_escape_id]
        )

    def _get_effective_susceptibles(self, model: Model) -> float:
        susceptibles = model[SUSCEPTIBLES_KEY]
        return susceptibles + self.escape_fraction * self._get_other_recovered(model)

    def get_propensity(self, model: Model) -> float:
        return (
            self.rate.get_value_at_time(model.current_time)
            * model.count_active_nodes(self.state)
            * self._get_effective_susceptibles(model)
        )

    def apply(self, model: Model):
        weight_S = model[SUSCEPTIBLES_KEY]
        weight_R = self.escape_fraction * self._get_other_recovered(model)
        if model.rng.random() > (weight_R / (weight_S + weight_R)):
            model[SUSCEPTIBLES_KEY] -= 1

        model.birth_from(model.draw_active_node(self.state), self.state)


@dataclass(kw_only=True)
class AntigenicEscapeMutation(StateDependentEvent):
    escape_fraction: pgt.Distribution

    def apply(self, model: Model):
        node = model.draw_active_node(self.state)
        _, fitness_gain_id = _get_mutation_ids(node[STATE_KEY])

        new_antigenic_escape_id = model[LAST_ANTIGENIC_ESCAPE_ID_KEY] + 1
        model[LAST_ANTIGENIC_ESCAPE_ID_KEY] = new_antigenic_escape_id

        new_state = _make_state(new_antigenic_escape_id, fitness_gain_id)
        model.migrate(node, new_state)

        new_escape_fraction = self.escape_fraction()
        model[ESCAPE_FRACTIONS_KEY][new_antigenic_escape_id] = new_escape_fraction
        model.add_run_stochastic_event(
            AntigenicEscapeTransmission(
                state=new_state,
                rate=model[TRANSMISSION_RATES_KEY][fitness_gain_id],
                escape_fraction=new_escape_fraction,
            )
        )


@dataclass(kw_only=True)
class FitnessGainMutation(StateDependentEvent):
    fitness_scaler: pgt.Distribution

    def apply(self, model: Model):
        node = model.draw_active_node(self.state)
        antigenic_escape_id, fitness_gain_id = _get_mutation_ids(node[STATE_KEY])

        new_fitness_gain_id = model[LAST_FITNESS_GAIN_ID_KEY] + 1
        model[LAST_FITNESS_GAIN_ID_KEY] = new_fitness_gain_id

        new_state = _make_state(antigenic_escape_id, new_fitness_gain_id)
        model.migrate(node, new_state)

        transmission_rate = model[TRANSMISSION_RATES_KEY][fitness_gain_id]
        new_transmission_rate = transmission_rate * self.fitness_scaler()
        model[TRANSMISSION_RATES_KEY][new_fitness_gain_id] = new_transmission_rate

        model.add_run_stochastic_event(
            AntigenicEscapeTransmission(
                state=new_state,
                rate=new_transmission_rate,
                escape_fraction=model[ESCAPE_FRACTIONS_KEY][antigenic_escape_id],
            )
        )


@dataclass(kw_only=True)
class AntigenicEscapeRecovery(Recovery):
    def apply(self, model: Model):
        node = model.draw_active_node(self.state)
        model.remove(node)
        antigenic_escape_id, _ = _get_mutation_ids(node[STATE_KEY])
        model[RECOVERED_KEY][antigenic_escape_id] += 1


def get_SIR_with_mutations_model(
    init_transmission_rate: SkylineParameterLike,
    init_escape_fraction: float,
    recovery_rate: SkylineParameterLike,
    sampling_rate: SkylineParameterLike,
    fitness_gains_mutations_rate: SkylineParameterLike,
    fitness_scalers_distribution: pgt.Distribution,
    antigenic_escape_mutations_rate: SkylineParameterLike,
    antigenic_escape_fractions_distribution: pgt.Distribution,
    susceptibles: int,
) -> Model:
    init_antigenic_escape_id = 1
    init_fitness_gain_id = 1
    init_state = _make_state(init_antigenic_escape_id, init_fitness_gain_id)
    init_transmission_rate = skyline_parameter(init_transmission_rate)
    model = Model(
        init_state=init_state,
        init_metadata={
            INITIAL_SUSCEPTIBLES_KEY: susceptibles,
            SUSCEPTIBLES_KEY: susceptibles,
            RECOVERED_KEY: {init_antigenic_escape_id: 0},
            LAST_ANTIGENIC_ESCAPE_ID_KEY: init_antigenic_escape_id,
            LAST_FITNESS_GAIN_ID_KEY: init_fitness_gain_id,
            TRANSMISSION_RATES_KEY: {init_fitness_gain_id: init_transmission_rate},
            ESCAPE_FRACTIONS_KEY: {init_antigenic_escape_id: init_escape_fraction},
        },
    )
    model.add_stochastic_event(
        AntigenicEscapeTransmission(
            rate=init_transmission_rate,
            state=init_state,
            escape_fraction=init_escape_fraction,
        )
    )
    model.add_stochastic_event(
        FitnessGainMutation(
            rate=skyline_parameter(fitness_gains_mutations_rate),
            fitness_scaler=fitness_scalers_distribution,
        )
    )
    model.add_stochastic_event(
        AntigenicEscapeMutation(
            rate=skyline_parameter(antigenic_escape_mutations_rate),
            escape_fraction=antigenic_escape_fractions_distribution,
        )
    )
    model.add_stochastic_event(
        AntigenicEscapeRecovery(rate=skyline_parameter(recovery_rate))
    )
    model.add_stochastic_event(
        Sampling(rate=skyline_parameter(sampling_rate), removal=True)
    )
    return model
