one-file-projects/pendel.py

#!/usr/bin/env python

from game import *

from random import randint

class Vector(object):
    def __init__(self, x, y):
        self.x = float(x)
        self.y = float(y)

    def __add__(self, other):
        return Vector( self.x + other.x, self.y + other.y )

    def __neg__(self):
        return Vector( -self.x, -self.y )

    def __mul__(self, other):
        return Vector( self.x * other, self.y * other )

    def __rmul__(self, other):
        return self.__mul__(other)

    def __div__(self, other):
        return Vector( self.x / other, self.y / other )

    def __truediv__(self, other):
        return self.__div__( other )

    def __sub__(self, other):
        return self.__add__(other.__neg__())

    def len(self):
        return math.sqrt( self.x * self.x + self.y * self.y )

    def astuple(self):
        return ( int(self.x) , int(self.y) )

    def __repr__(self):
        return "( %.2f, %.2f )" % (self.x, self.y)

class PendelSim(Game):
    def __init__(self):
        super(PendelSim,self).__init__(1000,600, 60)
        self.init_data()
        self.mousedown = False
        self.mousepos = (0,0)
        self.pathsurface = pygame.Surface( (self._width, self._height) )
        self.resetpath()

    def resetpath(self):
        BLACK = (0,0,0)
        self.pathsurface.fill( BLACK )


    def init_data(self):
        self.staff_lengths = staff_lengths = [ randint(150,200), randint(100,200), randint(100,150) ]
        fixed_joint = Vector( self._width/2, self._height/5 )
        joint_a = fixed_joint - Vector(staff_lengths[0],0)
        joint_b = joint_a - Vector(0,-staff_lengths[1])
        joint_c = joint_b - Vector(-staff_lengths[2],0)
        self.joints = [fixed_joint, joint_a, joint_b, joint_c ]
        self.joint_speed = [ Vector(0,0), Vector(0,0), Vector(0,0), Vector(0,0) ]
        self.old_joints = self.joints

    def on_update(self, dtime):
        new_joints = self.joints[:]
        for i in range(1,4):
            new_joints[i] += (0.99 * self.joint_speed[i] + Vector(0,40)) * dtime

        if self.mousedown:
            new_joints[1] = Vector(self.mousepos[0], self.mousepos[1])

        for i in range(0,4-1):
            relvec = new_joints[i+1] - new_joints[i]
            relvec = (relvec * self.staff_lengths[i]) / relvec.len()
            new_joints[i+1] = new_joints[i] + relvec

        for i in range(0,4):
            self.joint_speed[i] = (new_joints[i] - self.joints[i]) / dtime

        self.old_joints = self.joints
        self.joints = new_joints

    def on_render(self, surface):
        BLACK = (0,0,0)
        RED = (255,0,0)
        GREEN = (0,255,0)
        WHITE = (255,255,255)
        BLUE = (0,0,255)

        staff_colors = [ GREEN, RED, BLUE ]

        for i in range(2,4):
            pygame.draw.line(self.pathsurface, staff_colors[(i-1)%3], self.old_joints[i].astuple(), self.joints[i].astuple(), 2)

        surface.blit(self.pathsurface, (0,0))


        for i in range(4-1):
            pygame.draw.line(surface, staff_colors[i%3], self.joints[i].astuple(), self.joints[i+1].astuple(), 3)

        for i in range(4):
            pygame.draw.circle(surface, WHITE, self.joints[i].astuple(), 4)

        pygame.display.flip()

    def on_event(self, event):
        if event.type == MOUSEBUTTONDOWN:
            self.mousedown = True
            self.mousepos = event.pos[:]
        elif event.type == MOUSEBUTTONUP:
            self.mousedown = False
            self.mousepos = event.pos[:]
            self.resetpath()
        elif event.type == MOUSEMOTION:
            self.mousepos = event.pos[:]
        else:
            super(PendelSim, self).on_event(event)


def main():
    PendelSim().on_execute()

if __name__ == '__main__':
    main()