import logging
import random

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

  seq = []

  def __init__(self,
               N = 2, 
               trials = 100,                          # Number of total trials
               alphabet=['A','B','C','D','E','F'], 
               T = 20,                                # Number of targets
               L1 = 10,                               # Number of lures (foil) similar to the (N+1)-back
               L2 = 10                                # 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 = SkewedRandomGenerator()
  seq = generator.generate()
  print('Skewed Random Sequence: %s' % ''.join(seq))