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]

# initialize three blocks
for b in range(2):
    block = design.Block()
    block.set_factor("index", b)
    for tindex, item in enumerate(nback_sequence):
        trial = design.Trial()
        trial.set_factor("stimulus", item)
        trial.set_factor("index", tindex)
        for x in range(10):
            stim = stimuli.TextLine(item, text_size=200, background_colour=[0,0,0], text_colour=[255,0,0])
            c = stimuli.Canvas(stim.surface_size)
            stim.plot(c)
            c._get_surface().set_alpha(x*25)
            c.preload()
            trial.add_stimulus(c)
        block.add_trial(trial, random_position=False)
    exp.add_block(block)

expyriment.control.start(exp)

# present trials
for block in exp.blocks:
    for trial in block.trials:
        exp.clock.wait(1000 - fixation.present())
        for stim in trial.stimuli:
            stim.present()
            exp.clock.wait(5)
        pressed_key, rt = exp.keyboard.wait(response_keys)
        exp.data.add([block.get_factor("index"), 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)