import expyriment
from expyriment import design, control, stimuli, misc
from generators import progressive_random

num_of_blocks = 2
trials_per_block = 16
n = 2
choices = ['A', 'B', 'C', 'D', 'E', 'F', 'G', 'H']

control.defaults.window_mode = True
generators = [progressive_random.SequenceGenerator(choices, trials=trials_per_block, n=n) for b in range(0, num_of_blocks)]

exp = design.Experiment("Progressive N-Back Task")
control.initialize(exp)
exp.data_variable_names = ["block", "stimulus", "pressed_key", "rt"]

fixation = stimuli.FixCross()
fixation.preload()

# left and right arrow keys for responses
response_keys = [misc.constants.K_LEFT, misc.constants.K_RIGHT]

expyriment.control.start(exp)


# initialize three blocks
for b in range(0, num_of_blocks):
    block = design.Block()
    block.set_factor("block", b)
    block.set_factor("trials", trials_per_block)
    exp.add_block(block)

for block in exp.blocks:
    trials = block.get_factor("trials")
    block_index = block.get_factor("block")
    for i in range(0, trials):
        trial = design.Trial()
        item = generators[block_index].next_trial()
        trial.set_factor("stimulus", item)
        stim = stimuli.TextLine(item, text_size=200)
        exp.clock.wait(1000 - fixation.present() - stim.preload())
        stim.present()

        # Wait for response, remove duration param for forced-choice.
        pressed_key, rt = exp.keyboard.wait(response_keys, duration=2000)
        exp.data.add([block.get_factor("block"), trial.get_factor("stimulus"), pressed_key, rt])
        exp.clock.wait(1000 - expyriment.stimuli.BlankScreen().present() - stim.unload())

expyriment.control.end("Goodbye!", goodbye_delay=1000)