from abc import abstractmethod

from kartezio.callback import Callback, Event, EventType
from kartezio.core.components import Endpoint, Fitness, Library, Preprocessing
from kartezio.core.initialization import RandomInit
from kartezio.evolution.decoder import DecoderCGP
from kartezio.evolution.population import Population, PopulationHistory
from kartezio.evolution.strategy import OnePlusLambda, Strategy
from kartezio.export import PythonClassWriter
from kartezio.helpers import Observable
from kartezio.mutation.base import PointMutation
from kartezio.mutation.behavioral import MutationBehavior
from kartezio.mutation.decay import MutationDecay
from kartezio.mutation.handler import MutationHandler
from kartezio.types import DataBatch, DataPopulation


class ObservableModel(Observable):
    def send_event(self, name: EventType, state: PopulationHistory):
        self.notify(Event(self.get_current_iteration(), name, state))

    def force_event(self, name: EventType, state: PopulationHistory):
        self.notify(Event(self.get_current_iteration(), name, state, force=True))

    @abstractmethod
    def get_current_iteration(self) -> int:
        pass


class GeneticAlgorithm:
    def __init__(self, initializer, mutation_handler, fitness: Fitness):
        """Initialize the genetic algorithm."""
        self.strategy: Strategy = OnePlusLambda(initializer, mutation_handler)
        self.population: Population | None = None
        self.fitness: Fitness = fitness
        self.current_iteration = 0
        self.n_iterations = 0

    def initialize(self, n_iterations: int) -> None:
        """Compile the genetic algorithm with given iterations and children."""
        self.population = self.strategy.compile(n_iterations)
        self.current_iteration = 0
        self.n_iterations = n_iterations

    def is_satisfying(self):
        return (
            self.current_iteration >= self.n_iterations
            or self.population.get_fitness()[0] == 0.0
        )

    def selection(self) -> PopulationHistory:
        return self.strategy.selection(self.population)

    def reproduction(self):
        self.strategy.reproduction(self.population)

    def evaluation(self, y_true: DataBatch, y_pred: DataPopulation):
        fitness = self.fitness.batch(y_true, y_pred, reduction="raw")
        self.population.set_raw_fitness(fitness)

    def next(self):
        self.current_iteration += 1


class KartezioCGP(ObservableModel):
    def __init__(
        self,
        n_inputs,
        n_nodes,
        n_chromosomes,
        libraries,
        endpoint,
        fitness,
        preprocessing=None,
    ):
        super().__init__()
        self.preprocessing = preprocessing
        self.decoder = DecoderCGP(n_inputs, n_nodes, n_chromosomes, libraries, endpoint)
        mutation = PointMutation(self.decoder.adapter)
        initializer = RandomInit(mutation)
        self.evolver = GeneticAlgorithm(initializer, MutationHandler(mutation), fitness)

    def collect_updatables(self):
        updatables = []
        updatables.extend(self.evolver.strategy.mutation_handler.collect_updatables())
        return updatables

    def initialize(self, n_iterations):
        self.evolver.initialize(n_iterations)
        self.clear()
        for updatable in self.collect_updatables():
            self.attach(updatable)

    def get_current_iteration(self):
        return self.evolver.current_iteration

    @property
    def population(self) -> Population:
        return self.evolver.population

    @property
    def elite(self):
        return self.population.get_elite()

    def evaluation(self, x: DataBatch, y: DataBatch):
        y_pred = self.decoder.decode_population(self.population, x)
        self.evolver.evaluation(y, y_pred)

    def evolve(self, x: DataBatch, y: DataBatch):
        history = []
        self.evaluation(x, y)
        state = self.evolver.selection()
        history.append(state)
        if state.changed:
            self.force_event(EventType.NEW_PARENT, state)
        self.force_event(EventType.START_LOOP, state)
        while not self.evolver.is_satisfying():
            self.send_event(EventType.START_STEP, state)
            self.evolver.reproduction()
            self.evaluation(x, y)
            state = self.evolver.selection()
            if state.changed:
                self.force_event(EventType.NEW_PARENT, state)
            history.append(state)
            self.send_event(EventType.END_STEP, state)
            self.evolver.next()
        self.force_event(EventType.END_LOOP, state)
        return self.elite, history

    def preprocess(self, x):
        if self.preprocessing:
            return self.preprocessing.call(x)
        return x

    def predict(self, x):
        """
        Predict the output of the model given the input.
        Apply the preprocessing if it exists.
        """
        x = self.preprocess(x)
        return self.decoder.decode(self.elite, x)

    def evaluate(self, x, y):
        y_pred, _ = self.predict(x)
        return self.evolver.fitness.batch(y, [y_pred])


class KartezioTrainer:
    def __init__(
        self,
        n_inputs: int,
        n_nodes: int,
        libraries: list[Library] | Library,
        endpoint: Endpoint,
        fitness: Fitness,
        preprocessing: Preprocessing | None = None,
        n_chromosomes=1,
    ):
        super().__init__()
        if not isinstance(libraries, list):
            libraries = [libraries]
        self.model = KartezioCGP(
            n_inputs,
            n_nodes,
            n_chromosomes,
            libraries,
            endpoint,
            fitness,
            preprocessing,
        )
        self.updatables = []

    def fit(
        self,
        n_iterations: int,
        x: DataBatch,
        y: DataBatch,
        callbacks: list[Callback] = [],
    ):
        # compile the Kartezio model
        self.model.initialize(n_iterations)

        # attach callbacks
        for callback in callbacks:
            callback.set_decoder(self.model.decoder)
            self.model.attach(callback)

        # preprocess data
        x = self.model.preprocess(x)

        # evolve the model
        return self.model.evolve(x, y)

    @property
    def decoder(self) -> DecoderCGP:
        return self.model.decoder

    def set_endpoint(self, endpoint: Endpoint):
        self.decoder.endpoint = endpoint

    @property
    def strategy(self) -> Strategy:
        return self.model.evolver.strategy

    def set_n_children(self, n_children):
        self.strategy.n_children = n_children

    def set_required_fps(self, required_fps):
        self.strategy.set_required_fps(required_fps)

    def set_gamma(self, gamma):
        self.strategy.set_gamma(gamma)

    @property
    def mutation(self) -> MutationHandler:
        return self.strategy.mutation_handler

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

    def set_behavior(self, behavior: MutationBehavior):
        self.mutation.set_behavior(behavior)
        self.mutation.behavior.set_decoder(self.decoder)

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

    def set_mutation_effect(self, effect):
        self.mutation.set_effect(effect)

    def set_mutation_edges(self, edges):
        self.mutation.set_edges(edges)

    def evaluate(self, x, y):
        return self.model.evaluate(x, y)

    def print_python_class(self, class_name):
        python_writer = PythonClassWriter(self.decoder)
        python_writer.to_python_class(class_name, self.model.population.get_elite())

    def display_elite(self):
        elite = self.model.elite
        print(f"Elite chromosome: {elite[0]}")
        print(f"Elite outputs: {elite.outputs}")

    def summary(self):
        # TODO: implement summary
        pass
