from typing import Callable

import numpy as np

from projects.genetic_algorithms.base import (
    _Domain, _Mutator, _RandomSampler, _Scorer, _SimpleSorter)


class SimpleGeneticOptimizer:
    __slots__ = ['_domain',
                 '_objective',
                 '_population_size',
                 '_n_generations',
                 '_seed',
                 '_random_sampler',
                 '_scorer',
                 '_sorter',
                 '_mutator']

    def __init__(self,
                 domain: _Domain,
                 objective: Callable,
                 population_size: int = 30,
                 n_generations: int = 20,
                 mutation_probability: float = 0.2,
                 crossover_probability: float = 0.2,
                 seed: int = 111):
        assert population_size > 0
        assert n_generations > 0

        self._domain = domain
        self._objective = objective
        self._population_size = population_size
        self._n_generations = n_generations
        self._seed = seed

        rng = np.random.default_rng(self._seed)
        sampler_seed = rng.integers(low=1, high=1e10)
        mutator_seed = rng.integers(low=1, high=1e10)

        self._random_sampler = _RandomSampler(domain=self._domain,
                                              seed=sampler_seed)
        self._scorer = _Scorer(objectives=[self._objective],
                               domain=self._domain)
        self._sorter = _SimpleSorter()
        self._mutator = _Mutator(domain=self._domain,
                                 mutation_probability=mutation_probability,
                                 crossover_probability=crossover_probability,
                                 seed=mutator_seed)

    def perform_optimization(self):
        n = self._population_size
        population = self._random_sampler.get_random_population(n)
        self._scorer.score_population(population)

        for i in range(1, self._n_generations):
            population = self._mutator.mutate_population(population)
            self._scorer.score_population(population)
            population = self._sorter.sort_population(population)
            population = population[:n]

        return population
