import itertools as it
import random

import benchmarks.common as common
import scipy.stats

class ProgressiveGAGenerator:
    """Generate a sequence progressively according to a predefined TL ratio and an even distribution"""

    def __init__(self, choices, trials, tl=2.0, pool_size=100, n=3):
        """Initialize the genetic algorithm optimizer for n-back sequences.
        :param choices:
        :param trials:
        :param tl:
        :param pool_size:
        :param n:
        """
        self.tl, self.trials, self.choices, self.pool_size, self.n = tl, trials, choices, pool_size, n
        self.pool = []

        self.tl_norm = scipy.stats.norm(self.tl, 0.5)
        self.skewness_norm = scipy.stats.norm(0, 0.5)

        self.__init_pool(pool_size)

    def generate(self):
        """Generate a sequence of trials based on passed parameters. TL ratio and distribution are expected to be
        close to the desired ones but not exactly the same.
        :return: a sequence of items in "list" format.
        """
        generation = 0
        best_parent = self.__find_best_parents(1)[0]
        while self.cost(best_parent) > 0.1 and generation < 1000:
            generation += 1
            if random.random() > 0.5:
                self.pool = list(map(lambda s: self.mutate(s), self.pool))
            self.pool = self.crossover_all()
            best_parent = self.__find_best_parents(1)[0]
            print(best_parent, 'cost=%f' % self.cost(best_parent))
        return best_parent

    def __init_pool(self, pool_size) -> list:
        """
        Initialize solution pool.
        :param pool_size: Num of initial random solutions
        :return: initial pool of
        """
        print("Initializing the pool...")
        self.pool.clear()
        all_comb = it.combinations_with_replacement(self.choices, self.trials)
        sample = random.sample(list(all_comb), pool_size)
        self.pool.extend(map(lambda _: ''.join(_), sample))
        return self.pool

    def __find_best_parents(self, count=1):
        """
        Find best gene(s) or parent(s) from the current pool.
        :param count: Number of desired best parents to be returned. Default is 1.
        :return: A list of most fit sequences.
        """
        sorted_pool = sorted(self.pool, key=lambda ss: self.cost(ss))
        return sorted_pool[:count]


    def crossover_all(self):
        """
        Perform random crossover for all pairs.
        :return: new pool
        """
        new_pool = []
        for i in range(int(self.pool_size/2)):
            seq1 = self.pool[i*2]      # change to weighted random
            seq2 = self.pool[i*2 + 1]  # change to weighted random
            new_pool.extend(self.crossover(seq1, seq2))

        return new_pool

    def crossover(self, seq1, seq2):
        """
        Crossover two sequences.
        :param seq1:
        :param seq2:
        :return:
        """
        pos = random.randint(0, self.trials)
        return [seq1[:pos] + seq2[pos:], seq2[:pos] + seq1[pos:]]

    def mutate(self, seq):
        if random.random() > 0.5:
            pos = random.randint(0, len(seq)-1)
            seq_list = list(seq)
            seq_list[pos] = random.choice(self.choices)
            return ''.join(seq_list)
        return seq

    def cost(self, seq):
        targets, lures = common.count_targets_and_lures(seq, self.n)
        #print(seq, targets, lures)
        tl = targets / lures if lures >0 else targets
        c1, c2 = self.skewness_cost(seq), self.tl_cost(tl)
        return c1+c2

    def skewness_cost(self, seq):
        even_ratio = self.trials / len(self.choices)
        costs = {c: abs(seq.count(c) - even_ratio) for c in self.choices}
        max_deviation_from_even_dist = max(list(costs.values()))
        return 1.0 - self.skewness_norm.pdf(max_deviation_from_even_dist) / self.skewness_norm.pdf(0)

    def tl_cost(self, tl):
        return 1.0 - self.tl_norm.pdf(tl) / self.tl_norm.pdf(self.tl)


if __name__ == '__main__':

    generator = ProgressiveGAGenerator(['a', 'b', 'c', 'd', 'e', 'f', 'g', 'h'], trials=24, n=4)
    sq = generator.generate()

    print('Progressively-Optimized Sequence: %s' % sq)