# This demo was written by Peter Scarfe and Max Sargent

# Import the various libraries that we will be using
import pyglet
from psychopy import visual, event, core
import numpy as np

# Define our colours
black = 0
white = 1
grey = white /  2
red = [white, black, black]
green = [black, white, black]

# Get a list of the screens avaliable
display = pyglet.canvas.get_display()
screens = display.get_screens()

# Get the length of the list of screens
numScreens = len(screens)

# We want to present on the secondary monitor (assuming two monitors connected). Note screens are counted from 0.
screenNum = numScreens - 1;

# Get the size of the screen in pixels
myScreen = screens[screenNum]
widthPix = myScreen.width
heightPix = myScreen.height

# Create a full screen window on our secondary monitor. We use pyglet as our screen backend. Note the rbg1 colour profile.
# We color our screen grey and will use pixel units.
mywin = visual.Window(fullscr = True, units = 'pix', color = [black, black, black], winType = 'pyglet', colorSpace = 'rgb1', screen = screenNum)

# Dimensions of the back and front squares
squareDim = 200

# Square positions: we will draw the left squares inefficiently, and the right 
# squares efficiently in a single draw call
leftBackPos = [-widthPix * 0.25, 0]
leftFrontPos = [-(widthPix * 0.25) - 100, -100]
rightPos = np.array([[widthPix * 0.25, 0], [(widthPix * 0.25) - 100, -100]])

# Number of squares we will be drawing on each side
numSquares = rightPos.shape[0]

# Square colors for the efficent drawing
squareColors = np.array([red, green])

# Create rectangle stimuli for the left side: back and front
leftBackRect = visual.Rect(
    win = mywin, width = squareDim, height = squareDim,
    fillColor = red, colorSpace = 'rgb1', pos = leftBackPos
)
leftForwardRect = visual.Rect(
    win = mywin, width = squareDim, height = squareDim,
    fillColor = green, colorSpace = 'rgb1', pos = leftFrontPos
)

# For the right hand squares we use and ElementArrayStim to draw them both in one go rather 
# then sequentially, as with the left squares. This is simpler and more effiecient especially
# for many squares
rightSquares = visual.ElementArrayStim(
    win = mywin,
    units = 'pix',
    nElements = numSquares,
    elementTex = None,
    elementMask = None,
    xys = rightPos,
    sizes = squareDim,
    colors = squareColors,
    colorSpace = 'rgb1'
)

# Draw all the squares. We have to issue two drawing commands for the left hand squares as we
# draw them seperately, the oder of drawing being determined by the two drawing call. For
# the right hand squares we issue a single drawing call and the order is determined by the order
# of the position coordinates and colors we set up and passed to ElementArrayStim
leftBackRect.draw()
leftForwardRect.draw()
rightSquares.draw()

# Flip to the screen
mywin.flip()

# Wait for key press before continuing. When pressed close the screen and exit PsychoPy.
event.waitKeys()
mywin.close()
core.quit()