diff --git a/py/ga_optimized_random_generator.py b/py/ga_optimized_random_generator.py
index 91a5e33..f1083ff 100644
--- a/py/ga_optimized_random_generator.py
+++ b/py/ga_optimized_random_generator.py
@@ -5,75 +5,120 @@
 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
+    def __init__(self, choices, trials=64, tl=1, pool_size=10, n=2):
+        """
+        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._init_pool(pool_size)
+        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]
+        best_parent = self.__find_best_parents(1)
         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])
+            best_parent = self.__find_best_parents(1)
+            print(''.join(best_parent[0]), ' ', self.calculate_tl_ratio(best_parent, self.n))
+        return best_parent
 
-    def _init_pool(self, pool_size):
+    def __init_pool(self, pool_size):
+        """
+
+        :param pool_size: DNA size or number of parents in the GA pool.
+        :return: A DNA or pool in list format.
+        """
         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):
+    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.fitness(ss))
         return sorted_pool[:count]
 
+    def __dist_fitness(self, seq, trials):
+        """
+        Calculate fitness according to the similarity to the desired uniform distribution.
+        :param seq:
+        :param trials:
+        :return:
+        """
+        pass
+
     def fitness(self, seq):
+        """
+        Calculate overall fitness of a sequence (block of trials).
+        :param seq:
+        :return:
+        """
         # add fitness for uniform distribution of all stimuli
-        return abs(self.calculate_tl_ratio(seq) - self.tl)
+        return abs(self.calculate_tl_ratio(seq, self.n) - self.tl)
 
     def crossover_all(self):
+        """
+        Do random crossover for all pairs.
+        :return:
+        """
         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
+            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:], seq1[:pos] + seq2[pos:]]
+        return [seq1[:pos] + seq2[pos:], seq2[:pos] + seq1[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."""
+    def calculate_tl_ratio(seq, N):
+        """Calculates the T/L ratio in a block of trials."""
         targets = 0.0
         lures = 0.0
-        for index in range(2, len(seq)):
+        for index in range(N, 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
+            if item == seq[index - N]:
+                targets += 1.0
+            elif item == seq[index - (N-1)] or item == seq[index - (N+1)]:
+                lures += 1.0
+        if lures - 0.0 < 0.001:  # avoid division by zero
+            lures = 0.001
         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)
+    tl_ratio = generator.calculate_tl_ratio(s, 2)
 
-    print(s)
-    print('GA-Optimized Sequence: %s' % ''.join(s), 'with tl_ratio=%f' % tl_ratio)
+    print('GA-Optimized Sequence: %s' % ''.join(list(s[0])), 'with tl_ratio=%f' % tl_ratio)