import itertools as it
import random


class GAOptimizedRandomGenerator:
    """Generate even random sequences according to a predefined TL ration (Ralph, 2014)"""

    def __init__(self, choices, trials=64, tl=1, pool_size=10):
        self.tl, self.trials, self.choices, self.pool_size = tl, trials, choices, pool_size
        self.pool = []
        self._init_pool(pool_size)

    def generate(self):
        generation = 0
        best_parent = self.find_best_parents(1)[0]
        while self.fitness(best_parent) > 0.1 and generation < 10000:
            generation += 1
            self.pool = self.mutate()
            self.pool = self.crossover_all()
            best_parent = self.find_best_parents(1)
            print(''.join(best_parent[0]), ' ',self.calculate_tl_ratio(best_parent))
        return list(best_parent[0])

    def _init_pool(self, pool_size):
        self.pool.clear()
        all_comb = it.combinations_with_replacement(self.choices, self.trials)
        sample = random.sample(list(all_comb), pool_size)
        self.pool.extend(sample)
        return self.pool

    def find_best_parents(self, count=1):
        sorted_pool = sorted(self.pool, key=lambda ss: self.fitness(ss))
        return sorted_pool[:count]

    def fitness(self, seq):
        # add fitness for uniform distribution of all stimuli
        return abs(self.calculate_tl_ratio(seq) - self.tl)

    def crossover_all(self):
        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):
        pos = random.randint(0, self.trials)
        return [seq1[:pos] + seq2[pos:], seq1[:pos] + seq2[pos:]]

    def mutate(self):
        # pass
        return self.pool

    @staticmethod
    def calculate_tl_ratio(seq):
        """Calculates the T:L ratio of a sequence. It only works for 2-back for now."""
        targets = 0.0
        lures = 0.0
        for index in range(2, len(seq)):
            item = seq[index]
            if item == seq[index-2]:
                targets += 1
            elif item == seq[index-1] or item == seq[index-3]:
                lures += 1
        # avoid division by zero
        if lures == 0:
            lures = 1
        return targets/lures


if __name__ == '__main__':
    generator = GAOptimizedRandomGenerator(['a', 'b', 'c', 'd', 'e', 'f'], trials=32)
    s = generator.generate()
    tl_ratio = generator.calculate_tl_ratio(s)

    print(s)
    print('GA-Optimized Sequence: %s' % ''.join(s), 'with tl_ratio=%f' % tl_ratio)