from visual import*
from visual.graph import *
scene = display(title="Pendulum",width=500,height=500,range=2)
scene.autoscale = 0
L = 1
theta1 = -pi/4
theta2 = pi/3
s1 = sphere(pos = (L * sin(theta1), L * -cos(theta1), 0), radius = .08, color = color.red, m = .0025, p = vector(0,0,0))
s2 = sphere(pos = (L * sin(theta2), L * -cos(theta2), 0), radius = .08, color = color.red, m = .0025, p = vector(0,0,0))
c12 = vector(0,0,0)
g1 = s1.m * 9.8
g2 = s2.m * 9.8
pend1 = cylinder(pos = (0,0,0), axis = s1.pos, radius = .01, color = color.green)
pend2 = cylinder(pos = (0,0,0), axis = s2.pos, radius = .01, color = color.green)
spr12 = helix(pos = s1.pos, axis = s2.pos - s1.pos, radius = .05, thickness = .02, coils = 10, color = color.green)
k12 = 0.1

def proj(v, s):
    return dot(v,s)/dot(s,s)*s




t = 0
dt = .005
while(1):
    rate(600)
    s1.pos = (L * sin(theta1), L * -cos(theta1), 0)
    s2.pos = (L * sin(theta2), L * -cos(theta2), 0)
    c12 = (s1.m*s1.pos+s2.m*s2.pos)/(s1.m+s2.m)
    s1.rest = (s1.pos - c12)*0.7 + c12
    s2.rest = (s2.pos - c12)*0.7 + c12
    fpar1 = (sin(theta1) * g1) * norm(vector(L*sin(theta1-pi/2), L*-cos(theta1-pi/2),0)) + proj((s1.pos-s1.rest)*k12, norm(vector(L*sin(theta1-pi/2), L*-cos(theta1-pi/2),0)))
    fpar2 = sin(theta2) * g2 * norm(vector(L*sin(theta2-pi/2), L*-cos(theta2-pi/2),0)) + proj((s2.pos-s2.rest)*k12, norm(vector(L*sin(theta2-pi/2), L*-cos(theta2-pi/2),0)))
    s1.p += (fpar1 * dt)
    s2.p += (fpar2 * dt)
    s1.pos += (s1.p / s1.m) * dt
    s2.pos += (s2.p / s2.m) * dt
    spr12.axis = s2.pos - s1.pos
    spr12.pos = s1.pos
    theta1 = arcsin(s1.pos.x / L)
    theta2 = arcsin(s2.pos.x / L)
    pend1.axis = s1.pos
    pend2.axis = s2.pos
    t += dt
    #print mag(s1.pos), mag(s2.pos)