#!/usr/bin/env python

import pygame
import ode

default_density = 2500
default_radius = 0.5
white = (255, 255, 255)
blue  = (0, 0, 255)

class Pendulum(object):
    def __init__(self, ball, joint):
        self.ball = ball
        self.joint = joint

class Block(object):
    def __init__(self, block, anchors, joints, i_joints):
        self.block = block
        self.anchors = anchors
        self.joints = joints
        self.i_joints = i_joints

def coord((x, y, z)):
    return int(100 + x*20), int(300 - y*20)

def init_scene():
    # Position and density information
    pendulum_density = default_density
    pendulum_balls  = [(0, 8, 0), (3, 5, 0), ( 7, 4, 0), (12, 3, 0)]
    pendulum_joints = [(5, 8, 0), (7, 8, 0), (10, 8, 0), (12, 8, 0)]

    base_density = default_density
    base_supports = [(5, 12, 0), (12, 12, 0)]
    base_joints = [(5, 8, 0), (12, 8, 0)]
    
    # Create world
    world = ode.World()
    world.setGravity((0, -9.81, 0))
    
    # Create base
    b_body = ode.Body(world)
    b_mass = ode.Mass()
    b_mass.setBox(base_density, 13, 2, 2)
    b_body.setMass(b_mass)
    b_body.setPosition((8.5, 8.5, 0))
    
    # Create supports
    b_anchors = []
    b_joints = []
    i_joints = []
    for s_pos, j_pos in zip(base_supports, base_joints):  
        print "anchor at", s_pos
        # create the support ball
        s_body = ode.Body(world)
        s_mass = ode.Mass()
        s_mass.setSphere(default_density, default_radius)
        s_body.setMass(s_mass)
        s_body.setPosition(s_pos)
        
        # Create immovable joint
        i_joint = ode.BallJoint(world)
        i_joint.attach(s_body, ode.environment)
        i_joint.setAnchor(s_pos)
        print "connected", s_pos, "to environment."
        
        # Create flexible joint
        m_joint = ode.BallJoint(world)
        m_joint.attach(s_body, b_body)
        m_joint.setAnchor(j_pos)
        print "connected", s_pos, "to", j_pos
        
        b_anchors.append(s_body)
        b_joints.append(m_joint)
        i_joints.append(i_joint)
    
    base = Block(b_body, b_anchors, b_joints, i_joints)
    
    # Create pendulums
    p_list = []
    for p_pos, j_pos in zip(pendulum_balls, pendulum_joints):
        p_body = ode.Body(world)
        p_mass = ode.Mass()
        p_mass.setSphere(default_density, default_radius)
        p_body.setMass(p_mass)
        p_body.setPosition(p_pos)
        
        # Create & connect joint
        joint = ode.BallJoint(world)
        joint.attach(p_body, b_body)
        joint.setAnchor(j_pos)
        
        print "connected {0} to {1}".format(p_pos, j_pos)
        
        # Instantiate container class
        p_list.append(Pendulum(p_body, joint))
        
    # return base, p_list
    return world, base, p_list

def init_graphics():
    pygame.init()
    screen = pygame.display.set_mode((600, 600))
    return screen

def main(screen, world, base, p_list):
    clock = pygame.time.Clock()
    fps = 60
    dt = 1.0/fps
    
    while not pygame.event.get(pygame.QUIT):
        pygame.event.get([pygame.KEYDOWN, pygame.MOUSEMOTION])
        # update world
        world.step(dt)
        
        # render world
        screen.fill(white)
        for p in p_list:
            pygame.draw.circle(screen, blue, coord(p.ball.getPosition()), 10)
            pygame.draw.line(screen, blue, coord(p.joint.getAnchor()), coord(p.ball.getPosition()), 1)
        
        for p in p_list:
            pygame.draw.circle(screen, blue, coord(p.joint.getAnchor()), 10)
        pygame.draw.line(screen, blue, coord(p_list[0].joint.getAnchor()), coord(p_list[-1].joint.getAnchor()), 20)
        
        for a, j in zip(base.anchors, base.joints):
            pygame.draw.circle(screen, blue, coord(a.getPosition()), 10)
            pygame.draw.line(screen, blue, coord(a.getPosition()), coord(j.getAnchor()), 1)
            
        pygame.display.flip()
        
        # tick pygame's clock
        clock.tick(fps)
        
if __name__=="__main__":
    screen = init_graphics()
    world, base, p_list = init_scene()
    main(screen, world, base, p_list)