"""
Jankovic et al Experiment 2b: TO RUN, REMOVE COMMENT SYMBOLS "#" FROM LINES 590-592
Experiment 1 found no alerting when Flash occurred in Maljkovic & Nakayama (1994) compound task 
Experiment 2b: to examine effect of alerting Flash in second component task (target discrimination/corner task)
Conditions: ColourRepeat x Flash
Each trial coded as a member of a pair as in E2a (this was not necessary and does not have to factor into analysis; was done simply because I had to code, run and analyze the experiment in 24 hours before adding the data to a poster to be printed for a conference the following day)
Alerting stimulus: brightening of the screen background for 48 ms; 100 ms flash --> display SOA

@author: nadja

"""

from psychopy import visual, event, core, misc, gui
import random 
from datetime import datetime
from win32api import GetSystemMetrics

##define an error that we can raise (from Bert)
class ExpError(Exception):
    def __init__(self, value):
        self.value=value
    def __str__(self):
        return self.value

# Class to define the Maljkovic & Nakayama diamond shaped items with either L or R corner cut off
class DiamondItem(object):
    def __init__(self, Display, Trial, CutPos, Colour, ItemCenter = [0,0], isTarget = False):
        super(DiamondItem,self).__init__()
        self.Trial = Trial
        self.Display = Display
        self.isTarget = isTarget
        self.CutPos = CutPos
        self.ItemCenter = ItemCenter
        self.Colour = Colour
        self.Shape = visual.ShapeStim(
            win = self.Display.Trial.Block.Experiment.win,
            lineColor = [0,0,0],
            fillColor = [0,0,0],
            fillColorSpace = 'rgb',
            lineColorSpace = 'rgb',
            ##N&M: 1.0 x 1.0 deg of vis angle, with .14 deg cutoff (at 57 cm away, will be 1 cm by 1 cm, with .14 cm cutoff)
            ## On Lab screens: 1.0 cm x 1.0 cm, with .15 cut off (at 57 cm away)
            vertices = [(20,0),(0,20),(-10,10),(-10,-10),(0,-20)],
            pos = self.ItemCenter,
            closeShape = True
            )
        ## Change shape colour based on Colour property.
        try:
            if self.Colour == "Red":
                self.Shape.lineColor = [1,-1,-1]
                self.Shape.fillColor = [1,-1,-1]
            elif self.Colour == "Green":
                self.Shape.fillColor = [-1,-.09,-1]
                self.Shape.lineColor = [-1,-.09,-1]
            elif self.Colour == "None":
                pass
        except:
            raise ExpError("Error in setting DiamondItem's self.Colour. Acceptable arguments: 'Red', 'Green', None.")
        ## Change orientation based on Left or Right cut position. In degrees. 
        try:
            if self.CutPos == "Left":
                self.Shape.ori = 0
            elif self.CutPos == "Right":
                self.Shape.ori = 180
        except:
            raise ExpError("Error in setting DiamondItem's self.CutPos. Acceptable Arguments: 'Left' or 'Right'")

    def drawItem(self):
        self.Shape.draw()

    def targetPosition(self):
        if self.isTarget:
            return self.ItemCenter
        else:
            return None 

# Class to define display. Unlike the other experiments, here there is always one item at screen center.
class Display(object):
    def __init__(self, Trial, TargetColour, SetSize, TargetShapeID, isOddballPresent):
        super(Display,self).__init__()
        self.isOddballPresent = isOddballPresent
        self.Trial = Trial
        self.TargetColour = TargetColour
        self.SetSize = SetSize   
        self.TargetShapeID = TargetShapeID
        self.Display_Items_List = []

    def createDisplay(self):
            target = DiamondItem(Display = self, Trial = self.Trial, CutPos = self.TargetShapeID, Colour = self.TargetColour, ItemCenter = [0,0], isTarget = True)
            self.Display_Items_List.append(target)

    def draw(self):
        for i in range(len(self.Display_Items_List)):
            self.Display_Items_List[i].drawItem()

# Instead of a fixation cross, have a fixation "box" comprised of rotated Ls, forming the vertices. 
class Lshape(object):
    def __init__(self, FixBox, pos, ori):
        self.FixBox = FixBox
        self.pos = pos
        self.ori = ori
        self.shape = visual.ShapeStim(
            win = self.FixBox.Trial.Block.Experiment.win,
            lineColor = [-1,-1,-1],
            lineColorSpace = 'rgb',
            pos = self.pos,
            ori = self.ori,
            autoDraw = self.FixBox.autoDraw,
            vertices = [(0,5),(0,0),(5,0)],
            closeShape = False)
    def draw(self):
        self.shape.draw()

# Because item always at screen center, have fixation box rather than a fixation cross in the middle. 
class FixBox(object):
    def __init__(self,Trial,autoDraw = False):
        self.Trial = Trial
        self.autoDraw = autoDraw
    def draw(self):
        pos_oris = [[(-40,-40), 0],[(-40,40), 90],[(40,40),180],[(40,-40),270]]
        Shapes = []
        for po in pos_oris:
            shape = Lshape(FixBox = self, pos = po[0], ori = po[1])
            Shapes.append(shape)
        for sh in Shapes:
            sh.draw()

# Lots of object attributes not used in this experiment; holdover from E2a
class Trial(object):
    TrialNo = 0
    def __init__(self, Experiment, DoubletOrder, isOddballPresent, CounterBalanceOpt, Block, Repeat, TargetColour, TargetShapeID, Flash, SetSize, isCatchTrial= False):
        super(Trial, self).__init__()
        self.Experiment = Experiment
        self.DoubletOrder = DoubletOrder
        self.isOddballPresent = isOddballPresent
        self.CounterBalanceOpt = CounterBalanceOpt
        self.Block = Block
        self.Repeat = Repeat
        self.TargetColour = TargetColour
        self.TargetShapeID = TargetShapeID
        self.Flash = Flash
        self.SetSize = SetSize
        self.isCatchTrial = isCatchTrial
        self.FixBox = FixBox(Trial = self)

    def createTrialDisplay(self):
        self.TrialDisplay = Display(Trial = self, isOddballPresent = self.isOddballPresent, TargetColour = self.TargetColour, SetSize = self.SetSize, TargetShapeID = self.TargetShapeID)
        self.TrialDisplay.createDisplay()
        self.TargetPos = [0,0]
            
    def runTrial(self):        
        Flash_Duration = .044   ## want to approximate flash lasting 50 msec. Best we can do is 48 msec. To be safe from adding an extra frame, subtract from 48 half of the refresh rate
        Flash_Jitter = random.randint(8,12)/10.0  ## delivers a random floating point number between .5 and 1.2 (to jitter the period before the start of the trial and the Flash. 
        ISI_Duration = .044   ## to achieve total ISI duration of "50" msec - see Experiment 1 for more details
        self.FixBox.autoDraw = True
        self.FixBox.draw()
        timer = core.Clock()
        self.Block.Experiment.win.flip()
        timer.reset()
        core.wait(Flash_Jitter)
        self.Block.Experiment.win.flip()
        ## screen background changes to white for stim duration
        if self.Flash == "Flash":
            timer.reset()
            while timer.getTime() < Flash_Duration:
                self.Block.Experiment.win.color = [1,1,1]
                self.Block.Experiment.win.flip()
            ##Set screen colour back to grey:
            self.Block.Experiment.win.color = [0,0,0]
            self.Block.Experiment.win.flip()
        ## screen remains gray (with fixbox) for same duration
        else:
            timer.reset()
            while timer.getTime() < Flash_Duration:
                pass
        self.Block.Experiment.win.color = [0,0,0] 
        self.Block.Experiment.win.flip()

        ## ISI between the Flash (or absence of the flash) and the display onset
        timer.reset()
        while timer.getTime() < ISI_Duration:
            pass
        ## Present the stimulus display
        self.TrialDisplay.draw()
        self.Block.Experiment.win.flip()
        timer.reset()
        # Collect response to display and record time
        self.keys = event.waitKeys(
            maxWait = 2,
            keyList = ['z','m'],
            timeStamped = timer
            )
        self.Block.Experiment.win.flip()
        self.FixBox.autoDraw = False
        self.Block.Experiment.win.flip()

    def Response(self):
        if self.keys:
            KeyPress = self.keys[0][0]   ## to get the actual key pressed: index 0 of self.keys tuple
        else:
            KeyPress = "None"
        return KeyPress

    def Acc(self):
        if self.TargetShapeID == "Left":
            if self.keys and (self.keys[0][0]== 'z'):
                return 1
            else:
                return 0
        elif self.TargetShapeID == "Right":
            if self.keys and (self.keys[0][0]== 'm'):
                return 1
            else:
                return 0

    def RT(self):
        if self.keys:
            ReacTime = self.keys[0][1]
        else:
            ReacTime = "None"
        return ReacTime

    ## str function to write to the datafile
    def __str__(self):
        ## Add 1 to both blockNumber and TrialNo, because these start at 0
        return str(self.Block.Experiment.Age)  +"," + str(self.Block.Experiment.Gender)+ ","+ str(self.Block.Experiment.Hand)+","+str(self.Block.Experiment.ColourBlindness)+"," + \
            str(self.Block.blockNumber()+1)+ ","+ str(self.TrialNo+1)+","+ str(self.DoubletOrder)+","+str(self.isOddballPresent)+","+ str(self.Repeat)+","+ \
            str(self.Flash)+","+str(self.TargetColour)+","+str(self.TargetShapeID)+","+str(self.Response())+ ","+str(self.Acc())+"," +str(self.RT())

    ## for the first line of the datafile. 
    def DataHeader(self):
        return "Age,Gender,Hand,R/G_Colourblindness,BlockNo,TrialNo,DoubletOrder,isOddballPresent,Repeat,Flash,TargetColour,TargetShapeID,Response,Acc,RT"

    def saveTrial(self):
        ## saves the trial's DataHeader() and __str__()
        if self.Block.blockNumber() == 0 and self.TrialNo == 0:
            ## this is the first trial of the first block: write the dataheader to the outputfile
            self.Block.Experiment.DataFile.write(self.DataHeader()+'\n')
        ## write the trial's data to the file:
        self.Block.Experiment.DataFile.write(self.__str__()+'\n')

# again, several Block attributes (e.g. CounterBalanceOpt) held over from E2a
class Block(object):
    def __init__(self, Experiment, CounterBalanceOpt, TrialsPerBlock, TrialsPerCondition, isPracBlock = False):
        super(Block, self).__init__()
        self.Experiment = Experiment
        self.CounterBalanceOpt = CounterBalanceOpt
        self.TrialsPerCondition = TrialsPerCondition
        self.TrialsPerBlock = TrialsPerBlock
        self.isPracBlock = isPracBlock
        self.Trials = []
        self.Trial_Types = []
        self.Test_Trials = []

    ##Trial Type Codes: 2x2x2 = 8 trial types
    ##   1) NoRepeat/Repeat = self.Trial_Types[0]
    ##   2) NoFlash/Flash = self.Trial_Types[1]
    ##   3) CutLeft/CutRight = self.Trial_Types[3]

    def buildDoubletTypes(self):
        # Type Two Trials: isOddballPresent x Repeat x Flash x Colour
        self.TypeTwos= [["Right","Repeat","Flash","Red"],["Right","Repeat","Flash","Green"],["Right","Repeat","NoFlash","Red"],["Right","Repeat","NoFlash","Green"],\
            ["Right","NoRepeat","Flash","Red"],["Right","NoRepeat","Flash","Green"],["Right","NoRepeat","NoFlash","Red"],["Right","NoRepeat","NoFlash","Green"],\
            ["Left","Repeat","Flash","Red"],["Left","Repeat","Flash","Green"],["Left","Repeat","NoFlash","Red"],["Left","Repeat","NoFlash","Green"],\
            ["Left","NoRepeat","Flash","Red"],["Left","NoRepeat","Flash","Green"],["Left","NoRepeat","NoFlash","Red"],["Left","NoRepeat","NoFlash","Green"]]
        # Type One Trials: isOddballPresent x Flash
        self.TypeOnes= [["Right","Flash",],["Right","NoFlash"],["Left","Flash",],["Left","NoFlash"]]
        NumProgramConditions = len(self.TypeOnes) * len(self.TypeTwos)   # = 128 combinations 
        DoubletsPerBlock = self.TrialsPerBlock / 2
        NumDoubletTypesPerBlock = int(DoubletsPerBlock/NumProgramConditions)
        #Populate self.DoubletTypes with all possible combinations from 
        for i in range(NumDoubletTypesPerBlock):
            for two in range(len(self.TypeTwos)):
                for one in range(len(self.TypeOnes)):
                    self.Trial_Types.append([self.TypeOnes[one], self.TypeTwos[two]])
    
    def buildTrials(self):
        self.buildDoubletTypes()
        random.shuffle(self.Trial_Types)
        for trial in range(len(self.Trial_Types)):
            TrialOne = Trial(
                Experiment = self.Experiment,
                Block = self,
                CounterBalanceOpt = self.CounterBalanceOpt,
                DoubletOrder = 1,
                isOddballPresent = True,
                Repeat = "None",
                Flash = self.Trial_Types[trial][0][1],
                TargetColour = None,
                TargetShapeID = self.Trial_Types[trial][0][0],
                SetSize = 1)
            if self.Trial_Types[trial][1][1] == "Repeat":
                TrialOne.TargetColour = self.Trial_Types[trial][1][3]
            elif self.Trial_Types[trial][1][1] =="NoRepeat":
                if self.Trial_Types[trial][1][3] == "Red":
                    TrialOne.TargetColour = "Green"
                elif self.Trial_Types[trial][1][3] == "Green":
                    TrialOne.TargetColour = "Red"
            TrialTwo = Trial(
                Experiment = self.Experiment,
                Block = self,
                CounterBalanceOpt = self.CounterBalanceOpt,
                DoubletOrder = 2,
                isOddballPresent = True,
                Repeat = self.Trial_Types[trial][1][1],
                Flash = self.Trial_Types[trial][1][2],
                TargetColour = self.Trial_Types[trial][1][3],
                TargetShapeID = self.Trial_Types[trial][1][0],
                SetSize = 1)
            self.Trials.append(TrialOne)
            self.Trials.append(TrialTwo)
        for t in self.Trials:
            t.createTrialDisplay()

    def buildPracTrials(self):
        self.Prac_Types=[["Right","Repeat","Flash","Red"],["Right","Repeat","Flash","Green"],["Right","Repeat","NoFlash","Red"],["Right","Repeat","NoFlash","Green"],\
            ["Right","NoRepeat","Flash","Red"],["Right","NoRepeat","Flash","Green"],["Right","NoRepeat","NoFlash","Red"],["Right","NoRepeat","NoFlash","Green"],\
            ["Left","Repeat","Flash","Red"],["Left","Repeat","Flash","Green"],["Left","Repeat","NoFlash","Red"],["Left","Repeat","NoFlash","Green"],\
            ["Left","NoRepeat","Flash","Red"],["Left","NoRepeat","Flash","Green"],["Left","NoRepeat","NoFlash","Red"],["Left","NoRepeat","NoFlash","Green"]]
        for pt in range(len(self.Prac_Types)):
            self.Trials.append(Trial(Experiment = 
                    self.Experiment,
                    Block = self,
                    CounterBalanceOpt = self.CounterBalanceOpt,
                    DoubletOrder = 1,
                    isOddballPresent = None,
                    Repeat = "None",
                    Flash = self.Prac_Types[pt][2],
                    TargetColour = self.Prac_Types[pt][3],
                    TargetShapeID = self.Prac_Types[pt][0],
                    SetSize = 1))
        random.shuffle(self.Trials)
        for t in self.Trials:
            t.createTrialDisplay()

    def runTrials(self):
        trialcounter = 0
        if self.Experiment.TestProgram == 0:
            for t in self.Trials:
                t.TrialNo = trialcounter
                t.runTrial()
                trialcounter+=1
        elif self.Experiment.TestProgram == 1:
            for c in range(10):
                self.Trials[c].TrialNo = trialcounter
                self.Trials[c].runTrial()
                self.Test_Trials.append(self.Trials[c])
                trialcounter+=1
        if self.isPracBlock == False:
            self.saveTrials()
    
    def saveTrials(self):
        if self.Experiment.TestProgram == 0:
            for t in self.Trials:
                t.saveTrial()
        elif self.Experiment.TestProgram == 1:
            for test in self.Test_Trials:
                test.saveTrial()
    
    def blockNumber(self):
        ## find this block in parent experiment's collection:
        try:
            idx=self.Experiment.Blocks.index(self)   ## built in python method: list.index()
            return idx
        except:
            raise ExpError('Block blockNumber() method: unable to index block object')

    def initBlock(self):
        if self.isPracBlock == False:
            self.buildTrials()
        elif self.isPracBlock == True:
            self.buildPracTrials()

class Experiment(object):
    def __init__(self, NumBlocks, TrialsPerBlock, TrialsPerCondition, screen = 0, datafile = ""):
        super(Experiment, self).__init__()
        self.NumBlocks = NumBlocks
        self.TrialsPerBlock = TrialsPerBlock
        self.TrialsPerCondition = TrialsPerCondition
        self.Blocks = []

        expGUI = gui.Dlg(title="Experiment", screen = 0)
        expGUI.addField("ParticipantID:" , "0")
        expGUI.addField("Session:", "1")
        expGUI.addField("Test? (Programmer only)", choices = [0,1])
        configinfo = expGUI.show()
        if expGUI.OK:
            self.ParticipantID = configinfo[0]
            self.SessionNo = configinfo[1] 
            self.TestProgram = configinfo[2]
        else:
            print ("Experiment.init(): Cancelled by user")
            core.quit()
        demoGUI = gui.Dlg(title = "Demographic Information", screen = 0)
        demoGUI.addField("Age:", "0")
        demoGUI.addField("Gender:", choices = ["Female", "Male", "Other"])
        demoGUI.addField("Dominant Hand:", choices = ["Right", "Left"])
        demoGUI.addField("Red/Green Colourblindness?:", choices = ["No", "Yes","Not sure"])
        demoinfo = demoGUI.show()
        if demoGUI.OK:
            self.Age = demoinfo[0]
            self.Gender = demoinfo[1]
            self.Hand = demoinfo[2]
            self.ColourBlindness = demoinfo[3]

        ## If nothing passed as datafile argument, then a new .csv file will be opened to write to
        if datafile == "":
            d = datetime.now()
            ## fn = 20181112188_s999_1.csv
            filename = str(d.year)+str(d.month)+str(d.day)+str(d.hour)+str(d.minute)+"_"+"s"+str(self.ParticipantID)+"_"+str(self.SessionNo)+".csv"
            self.DataFile = open(filename, 'w')

        ## If filename is passed to datafile argument, then the file will be opened to append to
        else:
            self.DataFile = open(datafile, 'a')

        self.win = visual.Window(
            size = [GetSystemMetrics(0), GetSystemMetrics(1)],   ## How to fit to all screens. from win32api import GetSystemMetrics ; width = GetSystemMetics(0), height = GetSystemMetrics(1)
            units = 'pix',
            fullscr = True,
            color = [0,0,0],   ## grey
            allowGUI = False,  ## hide mouse
            waitBlanking = True
            )

    def drawInstructions(self):
        ## Create and draw Instructions screen from file. FILE MUST BE IN SAME FOLDER AS THE PROGRAM!
        try:
            self.InstructionsSlide = visual.ImageStim(
                win = self.win,
                image = "Instructions.bmp"
                )

        ## Display error text instead of 'Instructions.bmp' if instructions file not found
        except IOError:
            self.InstructionsSlide = visual.TextStim(
            win = self.win,
            color = (1,1,1),
            text = "Oops! Unable to display Instructions. Press space to continue",
            font = "Calibri",
            height = 40)
            
        self.InstructionsSlide.draw()

        ## Create BlockScreen ("Block # / press space to continue") in between each block. 
    def drawBlockScreen(self, BlockNo):
        self.BlockScreen = visual.TextStim(
            win = self.win,
            color = (1,1,1),
            pos = [0, 30],
            text = "Block " + str(BlockNo + 1),
            font = "Calibri",
            height = 40)
        self.pressSpace = visual.TextStim(
            win = self.win,
            color = (1,1,1),
            pos = [0, -30],
            text = "(press the spacebar to continue)",
            font = "Calibri",
            height = 30)
        self.BlockScreen.draw()
        self.pressSpace.draw()

    def drawPracStartScreen(self):
        self.PracStartScreen = visual.TextStim(
            win = self.win,
            color = (1,1,1),
            pos = [0, 30],
            text = "About to start the practice trials.",
            font = "Calibri",
            height = 30)
        self.pressSpaceBegin = visual.TextStim(
            win = self.win,
            color = (1,1,1),
            pos = [0, -30],
            text = "(press the spacebar to begin)",
            font = "Calibri",
            height = 25)
        self.PracStartScreen.draw()
        self.pressSpaceBegin.draw()

    def drawPracEndScreen(self):
        self.PracEndScreen = visual.TextStim(
            win = self.win,
            color = (1,1,1),
            pos = [0, 30],
            text = "That's the end of the practice trials. Ready to begin the experiment?",
            font = "Calibri",
            height = 30)
        self.pressSpaceExp = visual.TextStim(
            win = self.win,
            color = (1,1,1),
            pos = [0, -40],
            text = "(press the spacebar to start the experiment)",
            font = "Calibri",
            height = 25)
        self.PracEndScreen.draw()
        self.pressSpaceExp.draw()

    def drawDebrief(self):
        ## Create debrief slide imagestim object
        try:
            self.DebriefSlide = visual.ImageStim(
                win = self.win,
                image = "Debrief.bmp"   # FILE MUST BE IN SAME FOLDER AS PROGRAM
                )
        ## Display textstim error message instead of 'Debrief.bmp' if file not found
        except IOError:
            self.DebriefSlide = visual.TextStim(
                win = self.win,
                color = (1,1,1),
                text = "Oops! Unable to display Debrief. Press spacebar to exit",
                font = "Calibri",
                height = 40)
        self.DebriefSlide.draw()

    def createPracBlock(self):
        self.PracBlock = Block(Experiment = self, CounterBalanceOpt = None, TrialsPerBlock = self.TrialsPerBlock, TrialsPerCondition = self.TrialsPerCondition, isPracBlock = True)
        self.PracBlock.initBlock()

    def createBlocks(self):
        for i in range(self.NumBlocks):
            self.Blocks.append(Block(Experiment = self, CounterBalanceOpt = None, TrialsPerBlock = self.TrialsPerBlock, TrialsPerCondition = self.TrialsPerCondition, isPracBlock = False))
        random.shuffle(self.Blocks)

    def run(self):
        self.drawInstructions()
        self.win.flip()
        self.createPracBlock()
        self.createBlocks()
        event.waitKeys(keyList =['space'])
        self.win.flip()
        self.drawPracStartScreen()
        self.win.flip()
        event.waitKeys(keyList =['space'])
        self.PracBlock.runTrials()
        self.win.flip()
        self.drawPracEndScreen()
        self.win.flip()
        event.waitKeys(keyList=['space'])
        self.win.flip()
        self.runAllBlocks()
        self.drawDebrief()
        self.win.flip()
        event.waitKeys()

    def runAllBlocks(self):
        for b in range(len(self.Blocks)):
            self.drawBlockScreen(BlockNo = b)   ## draws screen that displays block number & asks for keypress (spacebar) to continue
            self.win.flip()
            self.Blocks[b].initBlock()   ## initialize the block while BlockScreen is being displayed
            keypress = event.waitKeys(keyList = ['space','escape'])
            ## If need to exit out of experiment, press esc key during block break. Will take you to debrief screen. 
            if keypress == ["escape"]:
                self.drawDebrief()
                self.win.flip()
                event.waitKeys()
                self.cleanup()
            else:
                self.win.flip()
                self.runBlock(BlockNo=b) #run a single block of trials
    
    ## Run one block of trials
    def runBlock(self, BlockNo=0):
        self.Blocks[BlockNo].runTrials()

    ## closes the datafile and the window; DO NOT FORGET TO CALL AT END OF EXPERIMENT 
    def cleanup(self):
        self.DataFile.close()
        self.win.close()


## ~~~~~~~~~~~~~~~~~~~~ RUN EXPERIMENT ~~~~~~~~~~~~~~~~~~~~ ##

## -------------test-------------##

# To test program, select '1' for 'Test?' on GUI

#test = Experiment(NumBlocks = 1, TrialsPerBlock = 128, TrialsPerCondition = 48)
#test.run()
#test.cleanup()

## ------------EXP--------------##

## 512 total trials = 128 trials per condition(both doublets, 1 & 2, included)

#Exp = Experiment(NumBlocks = 4, TrialsPerBlock = 128, TrialsPerCondition = 48)
#Exp.run()
#Exp.cleanup()