import logging
import random
import csv

import heapq

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,
        targets=8,                              # Number of targets
        lures1=2,                               # Number of lures (foil) similar to the (N+1)-back
        lures2=2                                # Number of lures (foil) similar to the (N-1)-back
    ):
        self.n, self.choices, self.targets, self.lures1, self.lures2 = n, choices, targets, lures1, lures2
        self.seq = []
        self.trials = None

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

        self.trials = trials
        targets, lures1, lures2 = self.targets, self.lures1, self.lures2
        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"POOL: {''.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)
        return self.seq


if __name__ == '__main__':
    gen = SequenceGenerator(['1','2','3','4','5','6'], 3, 8, 0, 0)
    s =gen.generate(24, debug=True)
    print(f"RESU: {''.join(s)}")