import logging
import random


class SequenceGenerator:
    """Generates random sequence of stimuli for the n-back task. Implementation is based on Ralph (2014)."""

    seq = []

    def __init__(self,
                N=2,
                trials=10,                          # Number of total trials
                alphabet=['A', 'B', 'C', 'D', 'E', 'F'],
                T=2,                                # Number of targets
                L1=1,                               # Number of lures (foil) similar to the (N+1)-back
                L2=1                                # Number of lures (foil) similar to the (N-1)-back
                ):
        self.N, self.alphabet, self.trials, self.T, self.L1, self.L2 = N, alphabet, trials, T, L1, L2
        self.D = trials - T - L1 - L2

    def generate(self) -> list:
        trial = 1
        self.seq = []
        while trial <= self.trials:
            self.seq += self.random_stimulus(trial)
            trial += 1
        return self.seq

    def random_stimulus(self, trial):
        rnd = random.randint(1, self.trials - trial + 1)
        T, L1, L2 = self.T, self.L1, self.L2
        if rnd <= T and len(self.seq) >= self.N:
            self.T -= 1
            return self.seq[-self.N]
        elif T < rnd <= T + L1  and len(self.seq) >= self.N+1:
            self.L1 -= 1
            return self.seq[-(self.N+1)]
        elif T + L1 < rnd <= T + L1 + L2  and len(self.seq) >= self.N-1:
            self.L2 -= 1
            return self.seq[-(self.N-1)]

        # distract
        self.D -= 1
        alphabet = [item for item in self.alphabet if item not in self.seq[-self.N-1:-self.N+1]]
        return random.choice(alphabet)


if __name__ == '__main__':
    generator = SequenceGenerator()
    seq = generator.generate()
    print('Skewed Random Sequence: %s' % ''.join(seq))