from collections import defaultdict
from dataclasses import dataclass

import phylogenie._typings as pgt
from phylogenie.skyline import SkylineParameterLike, skyline_parameter
from phylogenie.tree_node import TreeNode
from phylogenie.treesimulator.core import STATE, Model
from phylogenie.treesimulator.parameterizations.core import (
    SingleReactantEventFunction,
    StochasticEvent,
)

RECOVERED_COUNTS = "recovered_counts"
NODE_IMMUNITIES = "node_immunities"

LAST_ANTIGENIC_ESCAPE_ID = "last_antigenic_escape_id"
LAST_FITNESS_GAIN_ID = "last_fitness_gain_id"

TRANSMISSION_RATES = "transmission_rates"
ESCAPE_FRACTIONS = "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:
    state: str
    escape_fraction: float

    def _get_susceptibles(self, model: Model) -> dict[frozenset[int], int]:
        antigenic_id, _ = _get_mutation_ids(self.state)
        return {
            recovered_set: int(self.escape_fraction * count) if recovered_set else count
            for recovered_set, count in model[RECOVERED_COUNTS].items()
            if antigenic_id not in recovered_set
        }

    def reactant_combinations(self, model: Model) -> int:
        susceptibles = self._get_susceptibles(model).values()
        return sum(susceptibles) * model.count_active_nodes(self.state)

    def apply(self, model: Model):
        susceptibles = self._get_susceptibles(model)
        susceptible_sets, susceptible_counts = zip(*susceptibles.items())

        (susceptible_set,) = model.rng.choices(susceptible_sets, susceptible_counts)

        parent_node = model.draw_active_node(self.state)
        stem_node, new_node = model.birth_from(parent_node, self.state)
        model[NODE_IMMUNITIES][stem_node] = model[NODE_IMMUNITIES][parent_node]
        del model[NODE_IMMUNITIES][parent_node]

        model[RECOVERED_COUNTS][susceptible_set] -= 1
        model[NODE_IMMUNITIES][new_node] = susceptible_set


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

    def apply_to_node(self, model: Model, node: TreeNode):
        antigenic_escape_id, fitness_gain_id = _get_mutation_ids(node[STATE])

        new_antigenic_escape_id = model[LAST_ANTIGENIC_ESCAPE_ID] + 1
        model[LAST_ANTIGENIC_ESCAPE_ID] = new_antigenic_escape_id

        new_state = _make_state(new_antigenic_escape_id, fitness_gain_id)
        new_node = model.migrate(node, new_state)
        model[NODE_IMMUNITIES][new_node] = {
            *model[NODE_IMMUNITIES][node],
            antigenic_escape_id,
        }
        del model[NODE_IMMUNITIES][node]

        new_escape_fraction = self.escape_fraction()
        model[ESCAPE_FRACTIONS][new_antigenic_escape_id] = new_escape_fraction
        model.add_run_event(
            StochasticEvent(
                rate=model[TRANSMISSION_RATES][fitness_gain_id],
                fn=AntigenicEscapeTransmission(
                    state=new_state, escape_fraction=new_escape_fraction
                ),
            )
        )


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

    def apply_to_node(self, model: Model, node: TreeNode):
        antigenic_escape_id, fitness_gain_id = _get_mutation_ids(node[STATE])

        new_fitness_gain_id = model[LAST_FITNESS_GAIN_ID] + 1
        model[LAST_FITNESS_GAIN_ID] = new_fitness_gain_id

        new_state = _make_state(antigenic_escape_id, new_fitness_gain_id)
        new_node = model.migrate(node, new_state)
        model[NODE_IMMUNITIES][new_node] = model[NODE_IMMUNITIES][node]
        del model[NODE_IMMUNITIES][node]

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

        model.add_run_event(
            StochasticEvent(
                rate=new_transmission_rate,
                fn=AntigenicEscapeTransmission(
                    state=new_state,
                    escape_fraction=model[ESCAPE_FRACTIONS][antigenic_escape_id],
                ),
            )
        )


def _finalize_infection(model: Model, node: TreeNode):
    infected_set = model[NODE_IMMUNITIES][node]
    del model[NODE_IMMUNITIES][node]
    antigenic_escape_id, _ = _get_mutation_ids(node[STATE])
    model[RECOVERED_COUNTS][frozenset(infected_set | {antigenic_escape_id})] += 1
    return node


class AntigenicEscapeSampling(SingleReactantEventFunction):
    def apply_to_node(self, model: Model, node: TreeNode):
        _finalize_infection(model, node)
        model.sample(node)


class AntigenicEscapeRecovery(SingleReactantEventFunction):
    def apply_to_node(self, model: Model, node: TreeNode):
        _finalize_infection(model, node)
        model.remove(node)


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_gains_mutations_fitness_scaler: pgt.Distribution,
    antigenic_escape_mutations_rate: SkylineParameterLike,
    antigenic_escape_mutations_escape_fraction: 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)

    recovered_counts: defaultdict[frozenset[int], int] = defaultdict(int)
    recovered_counts[frozenset()] = susceptibles

    model = Model(
        init_state=init_state,
        init_metadata={
            RECOVERED_COUNTS: recovered_counts,
            NODE_IMMUNITIES: defaultdict(set),
            LAST_ANTIGENIC_ESCAPE_ID: init_antigenic_escape_id,
            LAST_FITNESS_GAIN_ID: init_fitness_gain_id,
            TRANSMISSION_RATES: {init_fitness_gain_id: init_transmission_rate},
            ESCAPE_FRACTIONS: {init_antigenic_escape_id: init_escape_fraction},
        },
    )
    model.add_event(
        StochasticEvent(
            rate=init_transmission_rate,
            fn=AntigenicEscapeTransmission(
                state=init_state, escape_fraction=init_escape_fraction
            ),
        )
    )
    model.add_event(
        StochasticEvent(
            rate=skyline_parameter(recovery_rate), fn=AntigenicEscapeRecovery()
        )
    )
    model.add_event(
        StochasticEvent(
            rate=skyline_parameter(sampling_rate), fn=AntigenicEscapeSampling()
        )
    )
    model.add_event(
        StochasticEvent(
            rate=skyline_parameter(fitness_gains_mutations_rate),
            fn=FitnessGainMutation(
                fitness_scaler=fitness_gains_mutations_fitness_scaler
            ),
        )
    )
    model.add_event(
        StochasticEvent(
            rate=skyline_parameter(antigenic_escape_mutations_rate),
            fn=AntigenicEscapeMutation(
                escape_fraction=antigenic_escape_mutations_escape_fraction
            ),
        )
    )
    return model
