import expyriment
from expyriment import design, control, stimuli, misc
from generators.nb_gm_004 import SequenceGenerator

import benchmarks.common as common
num_of_blocks = 2
trials_per_block = 8
n = 2
lures = 1
targets = 2

control.defaults.window_mode = True
generators = [SequenceGenerator(common.choices, 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")
    # TODO recalculate N based on performance
    generators[block_index].reset(trials_per_block, targets, lures, n=n)
    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)