from kartezio.core.components import Genotype, UpdatableComponent
from kartezio.mutation.base import Mutation
from kartezio.mutation.behavioral import MutationBehavior
from kartezio.mutation.decay import MutationDecay
from kartezio.mutation.edges import MutationEdges, MutationEdgesUniform
from kartezio.mutation.effect import MutationEffect, MutationUniform


class MutationHandler:
    def __init__(self, mutation: Mutation):
        self.mutation = mutation
        self.behavior = None
        self.decay = None
        self.effect = MutationUniform()
        self.edges_weights = MutationEdgesUniform()
        self.node_rate = None
        self.out_rate = None

    def set_behavior(self, behavior: MutationBehavior):
        self.behavior = behavior

    def set_decay(self, decay: MutationDecay):
        self.decay = decay

    def set_effect(self, effect: MutationEffect):
        self.effect = effect

    def set_edges(self, edges: MutationEdges):
        self.edges_weights = edges

    def set_mutation_rates(self, node_rate, out_rate):
        self.node_rate = node_rate
        self.out_rate = out_rate

    def compile(self, n_iterations: int):
        assert self.node_rate is not None, (
            "Node rate must be set before compiling the mutation handler."
        )
        assert self.out_rate is not None, (
            "Output rate must be set before compiling the mutation handler."
        )
        self.mutation.node_rate = self.node_rate
        self.mutation.out_rate = self.out_rate
        self.mutation.edges_weights = self.edges_weights
        self.mutation.parameters = self.effect
        self.mutation.parameters.compile(n_iterations)
        if self.behavior:
            self.behavior.set_mutation(self.mutation)
        if self.decay:
            self.decay.set_mutation(self.mutation)
            self.decay.compile(n_iterations)

    def collect_updatables(self):
        updatables = []
        if isinstance(self.mutation.parameters, UpdatableComponent):
            updatables.append(self.mutation.parameters)
        if self.behavior:
            if isinstance(self.behavior, UpdatableComponent):
                updatables.append(self.behavior)
        if self.decay:
            if isinstance(self.decay, UpdatableComponent):
                updatables.append(self.decay)
        return updatables

    def mutate(self, genotype: Genotype):
        if self.behavior:
            return self.behavior.mutate(genotype)
        else:
            return self.mutation.mutate(genotype)
