Walkingsquaresinbox_walkingsquaresinbox.py

#squares walking on and in squares
__author__ = 'drandreasdr'
import numpy as np
import math
import pygame
from pygame import gfxdraw
math.tau = 2*math.pi

colorwhite = (255, 255, 255)
colorblack = (0, 0, 0)

class WalkingSquaresInBox:
    def __init__(self, animationrect, period):
        self.animationrect = animationrect
        self.period = period

        #Numbers of cycles (i.e. 90 degree turns):
        self.ncyclesA = 8
        self.ncyclesB = 8*3
        #self.ncyclesABcontact = 8*3/2

        #Periods:
        ##Time of one 90 degree rotation of A and B, resp:
        self.periodA = self.period/self.ncyclesA
        self.periodB = self.period/self.ncyclesB
        ##Time during which a pair of corners in A and B are in contact:
        self.periodABcontact = self.periodB*2

        #Rotation angles in one period:
        self.rotAngleA = math.tau/4

        #Global time for animation: (repeats after the time "self.period")
        self.elapsedtime = 0

        #Geometry:
        self.sidelength = self.animationrect.width/3

        #Predetermined stuff:
        ##List of A-ground contact points positions and indices, and initial angles, respectively:
        self.indA_cAg_list = np.cumsum([0, 2, 1, 2, 1, 2, 1, 2]) % 4
        posx_cAg_list = self.animationrect.left + self.sidelength*np.array([2, 3, 3, 2, 1, 0, 0, 1])
        posy_cAg_list = self.animationrect.top + self.sidelength*np.array([3, 2, 1, 0, 0, 1, 2, 3])
        self.pos_cAg_list = np.array([posx_cAg_list, posy_cAg_list]).T
        self.alpha0_list = math.tau/4*(np.array([1, 2, 2, 3, 3, 4, 4, 1]) % 4)

    def tick(self, dt):
        #Advance time:
        self.elapsedtime = (self.elapsedtime + dt) % self.period

    def draw(self, displaysurface):
        #Count number of periods elapsed and measure time current ones have lasted:
        nelapsedperiodsA, elapsedtimeA = divmod(self.elapsedtime, self.periodA)
        nelapsedperiodsAB, elapsedtimeABcontact = divmod(self.elapsedtime, self.periodABcontact)

        #Indices, angles, positions depending on the above:
        ##For square A:
        indA_cAg = self.indA_cAg_list[nelapsedperiodsA]
        pos_cAg = self.pos_cAg_list[nelapsedperiodsA, :]

        alpha0 = self.alpha0_list[nelapsedperiodsA]
        alpharel = self.rotAngleA*elapsedtimeA/self.periodA
        alpha = (alpha0 - alpharel) % math.tau
        ###Help objects:
        unitvectorA = np.array([math.cos(alpha), -math.sin(alpha)])
        unitvectorAperp = np.array([math.cos(alpha + math.tau/4), -math.sin(alpha + math.tau/4)])
        pos_cornersA_list = np.zeros([4, 2])
        pos_cornersA_list[indA_cAg] = pos_cAg
        pos_cornersA_list[(indA_cAg + 1) % 4] = pos_cAg + self.sidelength*unitvectorA
        pos_cornersA_list[(indA_cAg + 2) % 4] = pos_cAg + self.sidelength*(unitvectorA + unitvectorAperp)
        pos_cornersA_list[(indA_cAg + 3) % 4] = pos_cAg + self.sidelength*unitvectorAperp
        ##For square B:
        indA_cAB = (2 + nelapsedperiodsAB) % 4
        pos_cAB = pos_cornersA_list[indA_cAB, :]
        beta0 = (alpha + math.tau/2 + math.tau/4*((indA_cAB - indA_cAg - 1) % 4)) % math.tau
        betarel = math.tau/2*elapsedtimeABcontact/self.periodABcontact
        beta = (beta0 + betarel) % math.tau
        ###Help objects:
        unitvectorB = np.array([math.cos(beta), -math.sin(beta)])
        unitvectorBperp = np.array([math.cos(beta + math.tau/4), -math.sin(beta + math.tau/4)])
        pos_cornersB_list = np.array([pos_cAB,
                                      pos_cAB + self.sidelength*unitvectorB,
                                      pos_cAB + self.sidelength*(unitvectorB + unitvectorBperp),
                                      pos_cAB + self.sidelength*unitvectorBperp])
        ###(note: I don't care about consistency in numbering of square B:s corners between AB-periods)

        #Draw:
        pos_cornersA_int_list = [[int(y) for y in x] for x in pos_cornersA_list.tolist()]
        pos_cornersB_int_list = [[int(y) for y in x] for x in pos_cornersB_list.tolist()]
        pygame.draw.rect(displaysurface, colorwhite, self.animationrect, 0)
        gfxdraw.aapolygon(displaysurface, pos_cornersA_int_list, colorblack)
        gfxdraw.aapolygon(displaysurface, pos_cornersB_int_list, colorblack)
        gfxdraw.filled_polygon(displaysurface, pos_cornersA_int_list, colorblack)
        gfxdraw.filled_polygon(displaysurface, pos_cornersB_int_list, colorblack)