import logging
import random

#trials = 100              # Number of total trials
#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
#D  = trials - L1 - L2 - T  # Number of distractors

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, 
               alphabet=['A','B','C','D','E','F'], 
               T = 20, 
               L1 = 10, 
               L2 = 10):
    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) -> str:
    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)]

    # distractor
    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))