import logging
import random
import csv

import heapq

import benchmarks.common as utils

class SequenceGenerator:
    """nb_gm_002 generator
    Generates skewed random sequence of stimuli for the n-back task, based on Ralph (2014).
    Each sequence contains specific fraction of matched trials (targets).
    """

    def __init__(
        self,
        choices: list,
        n=2
    ):
        self.n, self.choices = n, choices
        self.seq = []
        self.trials = None

    def generate(self, trials, targets, lures1, lures2, debug=False) -> list:

        self.trials = trials
        distractors = trials - targets - lures1 - lures2 - self.n
        pool = ['T'] * targets + ['1'] * lures1 + ['2'] * lures2 + ['-'] * distractors
        random.shuffle(pool)
        pool = ['-'] * self.n + pool
        if debug:
            print(f"MASK  : {''.join(pool)}")
        self.seq.clear()
        for i in range(len(pool)):
            t = pool[i]
            if t == 'T':
                self.seq += self.seq[i-self.n]
            elif t == '1':
                self.seq += self.seq[i-self.n - 1]
            elif t == '2':
                self.seq += self.seq[i-self.n + 1]
            else:  # distractors
                d = [c for c in self.choices if c not in self.seq[-self.n - 1:-self.n + 1]] if i > self.n else self.choices
                self.seq += random.choices(d)
        if debug:
            print(f"RESULT: {''.join(self.seq)}", '(t,l)=', utils.count_targets_and_lures(self.seq, self.n))
        return self.seq


if __name__ == '__main__':
    gen = SequenceGenerator(['A', 'B', 'C', 'D', 'E', 'F'], 2)
    s =gen.generate(24, 8, 4, 0, debug=True)

    with open("seqs.txt") as fh:
        for line in fh:
            print(line)