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

import benchmarks.common as common

nback_sequence = nb_gm_002.SequenceGenerator(common.choices).generate(24, 8, 2, 2)

exp = design.Experiment("Skewed 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, 2):
    block = design.Block()
    block.set_factor("block", b)
    for item in nback_sequence:
        trial = design.Trial()
        trial.set_factor("stimulus", item)
        stim = stimuli.TextLine(item, text_size=200)
        stim.preload()
        trial.add_stimulus(stim)
        block.add_trial(trial, random_position=False)
    exp.add_block(block)

for block in exp.blocks:
    for trial in block.trials:
        exp.clock.wait(1000 - fixation.present())
        trial.stimuli[0].present()
        pressed_key, rt = exp.keyboard.wait(response_keys)
        exp.data.add([block.get_factor("block"), trial.get_factor("stimulus"), pressed_key, rt])
        exp.clock.wait(1000 - expyriment.stimuli.BlankScreen().present() - trial.stimuli[0].unload())

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