from skga.hbrkga.brkga_mp_ipr.exploitation_method import ExploitationMethod
from skga.hbrkga.nn_decoder_PT import NNDecoder
from bayes_opt import BayesianOptimization
from bayes_opt import UtilityFunction
from typing import List
import random

class BayesianOptimizer(ExploitationMethod):
    def __init__(self, decoder: NNDecoder, e = 0.3, steps = 3, percentage = 1.0, verbose = 0):

        self._decoder = decoder
        self._e = e
        self._steps = steps
        self._percentage = percentage

        pbounds = {}
        for boundIdx in range(len(self._decoder._limits)):
            pbounds["g"+str(boundIdx)] = (0, 1)

        self._bayesian_optimizer = BayesianOptimization(
            f=self._score_for_bayesian_optimization,
            pbounds=pbounds,
            verbose=verbose
        )

    ###########################################################################

    def initialize(self, initial_populations: List[Population]) -> List[Population]:
        # Add the individuals to the list of explored solutions
        for population in initial_populations:
            for individual_idx in range(len(population.chromosomes)):
                individual = population.chromosomes[individual_idx]
                value = [x for (x,y) in population.fitness if y == individual_idx][0]
                self._bayesian_optimizer.register(individual, value)
        return initial_populations

    ###########################################################################

    def new_chromosomes(self, new_chromosomes: List[BaseChromosome]) -> List[BaseChromosome]:
        # Replace clones with suggestions from the optimizer
        for chromosome_idx in range(len(new_chromosomes)):
            if self._is_repeated_chromosome(new_chromosomes[chromosome_idx]):
                print(f"Repeated chromosome -> {new_chromosomes[chromosome_idx]}")
                new_chromosomes[chromosome_idx] = self.suggest_a_chromosome(new_chromosomes[chromosome_idx])
                print(f"Replaced by -> {new_chromosomes[chromosome_idx]}\n")
        return new_chromosomes

    ###########################################################################

    def suggest_a_chromosome(self, chromosome: BaseChromosome) -> BaseChromosome:
        # Suggest new chromosome for new_chromosomes()
        nextChromosome = [0]*len(chromosome)
        try:
            utility = UtilityFunction(kind="ucb", kappa=2.5, xi=0.0)
            pbounds = {}
            for geneIdx in range(len(nextChromosome)):
                pbounds["g"+str(geneIdx)] = (0,1)
            self._bayesian_optimizer.set_bounds(new_bounds=pbounds)
            temp = self._bayesian_optimizer.suggest(utility)
            for geneIdx in temp:
                nextChromosome[int(geneIdx[1:])] = temp[geneIdx]
            print("Replaced chromosome with suggestion")
        except Exception:
            for geneIdx in range(len(nextChromosome)):
                nextChromosome[geneIdx] = random.uniform(0, 1)
            print("Replaced chromosome randomly")
        return nextChromosome

    ###########################################################################

    def new_population(self, new_population: Population) -> Population:
        for individual_idx in range(len(new_population.chromosomes)):
            try:
                self._bayesian_optimizer.register(new_population.chromosomes[individual_idx], [x for (x,y) in new_population.fitness if y == individual_idx][0])
            except Exception:
                #print(f"Could not register {new_population.chromosomes[individual_idx]}")
                pass

            pbounds = {}
            currentChromosome = {}
            for geneIdx in range(len(new_population.chromosomes[individual_idx])):
                min = new_population.chromosomes[individual_idx][geneIdx] - self._e
                max = new_population.chromosomes[individual_idx][geneIdx] + self._e
                if min < 0:
                    min = 0
                if max > 1:
                    max = 1
                pbounds["g"+str(geneIdx)] = (min,max)
                currentChromosome["g"+str(geneIdx)] = new_population.chromosomes[individual_idx][geneIdx]

            self._bayesian_optimizer.set_bounds(new_bounds=pbounds)

            try:     
                self._bayesian_optimizer.maximize(
                    n_iter = self._steps
                )
            except Exception:
                # Sometimes an exception happens in sklearn at the beginning of Bayesian Optimization, usually just for the first chromosome.
                print(f"Could not maximize {new_population.chromosomes[individual_idx]}")

            testedChromosomesInsideCurrentBounds = []
            for testedChromosome in self._bayesian_optimizer.res:
                insideCurrentBounds = True
                for testedGene in testedChromosome["params"]:
                    if not (testedChromosome["params"][testedGene] >= pbounds[testedGene][0] and testedChromosome["params"][testedGene] <= pbounds[testedGene][1]):
                        insideCurrentBounds = False
                        break
                if insideCurrentBounds:
                    testedChromosomesInsideCurrentBounds.append(testedChromosome)

            bestChromosomeInsideCurrentBounds = testedChromosomesInsideCurrentBounds[0]
            for chromosome in testedChromosomesInsideCurrentBounds[1:]:
                if chromosome["target"] > bestChromosomeInsideCurrentBounds["target"]:
                    bestChromosomeInsideCurrentBounds = chromosome

            newChromosome = [0]*len(new_population.chromosomes[individual_idx])
            for geneIdx in bestChromosomeInsideCurrentBounds["params"]:
                newChromosome[int(geneIdx[1:])] = bestChromosomeInsideCurrentBounds["params"][geneIdx]

            if not newChromosome in new_population.chromosomes:
                new_population.chromosomes[individual_idx] = newChromosome
                for idx in range(len(new_population.fitness)):
                    if new_population.fitness[idx][1] == individual_idx:
                        new_population.fitness.pop(idx)
                        new_population.fitness.append((bestChromosomeInsideCurrentBounds["target"], individual_idx))
                        break

        return new_population

    ###########################################################################

    def _score_for_bayesian_optimization(self, **genes):
        chromosome = [0] * len(genes)
        for geneIdx in genes:
            chromosome[int(geneIdx[1:])] = genes[geneIdx]
        return (self._decoder.decode(chromosome = chromosome, rewrite = False, percentage = self._percentage))

    ###########################################################################

    def _is_repeated_chromosome(self, chromosome: BaseChromosome) -> bool:
        #Checks if the provided chromose has been explored before.
        repeated = False
        temp = {}
        for geneIdx in range(len(chromosome)):
            temp[f"g{geneIdx}"] = chromosome[geneIdx]
        for res in self._bayesian_optimizer.res:
            if res["params"] == temp:
                repeated = True
                break
        return repeated

    ###########################################################################

    def new_mutant(self, new_mutant: BaseChromosome) -> BaseChromosome:
        #Replace mutants with suggestions from BO
        #new_mutant = self.suggest_a_chromosome(new_mutant)
        return new_mutant