from __future__ import print_function, division
from visual import *
#from visual.graph import *

#Double pendulum
#Adapted from the work of Bruce Sherwood and Akshay Murthy

scene.title = 'Double Pendulum'
scene.height = scene.width = 800
scene.autoscale = true
g = 9.80665
m1 = 1.0 #upper mass
m2 = 0.5 #lower mass
mrat = m2/m1 #ratio of lower mass to upper mass
l = 1 #the length of each rigid bar (same length and equal to 1 for simplicity)
d = 0.05 # thickness of each bar
gap = 1 #distance between two parts of upper assembly
L1display = l+d # show upper assembly a bit longer than physical, to overlap axle
L2display = l+d/2 # show lower bar a bit longer than physical, to overlap axle

hpedestal = 2.6*l #height of pedestal
wpedestal = 0.1 #width of pedestal
tbase = 0.05 #thickness of base
wbase = 8*gap #width of base
offset = 2*gap #from center of pedastal to center of assembly
top = vector(0,0,0) #top of inner bar of assembly
scene.center = top-vector(0,l,0)

phi1 = 2.3 #initial upper angle from vertical
phi2 = 2.1 #initial lower angle from vertical
phi1dot = 0 #initial time derivative of phi1
phi2dot = 0 #initial time derivative of phi2

pedestal = box(pos=top-vector(0,hpedestal/2.,offset),
               height=1.1*hpedestal, length=wpedestal, width=wpedestal,
               color=(0.4,0.4,0.5))
base = box(pos=top-vector(0,hpedestal+tbase/2.,offset),
                 height=tbase, length=wbase, width=wbase,
                 color=pedestal.color)
axle = cylinder(pos=top-vector(0,0,gap/2.-d/4.), axis=(0,0,-offset), radius=d/4., color=color.yellow)

frame1 = frame(pos=top)
bar1 = box(frame=frame1, pos=(L1display/2.-d/2.,0,-(gap+d)/2.), size=(L1display,d,d), color=color.white)
bar1b = box(frame=frame1, pos=(L1display/2.-d/2.,0,(gap+d)/2.), size=(L1display,d,d), color=color.white)
axle1 = cylinder(frame=frame1, pos=(l,0,-(gap+d)/2.), axis=(0,0,gap+d),
                 radius=axle.radius, color=axle.color)
frame1.axis = (0,-1,0)
frame2 = frame(pos=frame1.axis*l)
bar2 = box(frame=frame2, pos=(L2display/2.-d/2.,0,0), size=(L2display,d,d), color=color.white)
frame2.axis = (0,-1,0)
frame1.rotate(axis=(0,0,1), angle=phi1)
frame2.rotate(axis=(0,0,1), angle=phi2)
frame2.pos = top+frame1.axis*l
    
dt = 0.00005
t=0

while True:
    rate(1/dt)
    p1 = (m1+m2)*l**2*phi1dot+m2*l**2*phi2dot*cos(phi1-phi2) #Hamiltonian generalized momentum
    p2 = m2*l**2*phi2dot+m2*l**2*phi1dot*cos(phi1-phi2) #Hamiltonian generalized momentum
    A = p1*p2*sin(phi1-phi2)/(m1*l**2*(1+mrat*(sin(phi1-phi2))**2))
    B = (p1**2*mrat-2*p1*p2*mrat*cos(phi1-phi2)+p2**2*(1+mrat))*sin(2*(phi1-phi2))/(2*m1*l**2*(1+mrat*(sin(phi1-phi2))**2)**2)
    p1dot = -m1*(1+mrat)*g*l*sin(phi1)-A+B
    p2dot = -m1*mrat*g*l*sin(phi2)+A-B
    Y11 = dt*((p1-p2*cos(phi1-phi2))/(m1*l**2*(1+mrat*(sin(phi1-phi2))**2)))
    Y12 = dt*((p2*(1+mrat)-p1*mrat*cos(phi1-phi2))/(m1*l**2*(1+mrat*(sin(phi1-phi2))**2)))
    Y13 = dt*(-m1*(1+mrat)*g*l*sin(phi1)-A+B)
    Y14 = dt*(-m1*mrat*g*l*sin(phi2)+A-B)
    phi1a = phi1+0.5*Y11
    phi2a = phi2+0.5*Y12
    p1a = p1+0.5*Y13
    p2a = p2+0.5*Y14
    Aa = p1a*p2a*sin(phi1a-phi2a)/(m1*l**2*(1+mrat*(sin(phi1a-phi2a))**2))
    Ba = (p1a**2*mrat-2*p1a*p2a*mrat*cos(phi1a-phi2a)+p2**2*(1+mrat))*sin(2*(phi1a-phi2a))/(2*m1*l**2*(1+mrat*(sin(phi1a-phi2a))**2)**2)
    Y21 = dt*((p1a-p2a*cos(phi1a-phi2a))/(m1*l**2*(1+mrat*(sin(phi1a-phi2a))**2)))
    Y22 = dt*((p2a*(1+mrat)-p1a*mrat*cos(phi1a-phi2a))/(m1*l**2*(1+mrat*(sin(phi1a-phi2a))**2)))
    Y23 = dt*(-m1*(1+mrat)*g*l*sin(phi1a)-Aa+Ba)
    Y24 = dt*(-m1*mrat*g*l*sin(phi2a)+Aa-Ba)
    phi1b = phi1+0.5*Y21
    phi2b = phi2+0.5*Y22
    p1b = p1+0.5*Y23
    p2b = p2+0.5*Y24
    Ab = p1b*p2b*sin(phi1b-phi2b)/(m1*l**2*(1+mrat*(sin(phi1b-phi2b))**2))
    Bb = (p1b**2*mrat-2*p1b*p2b*mrat*cos(phi1b-phi2b)+p2b**2*(1+mrat))*sin(2*(phi1b-phi2b))/(2*m1*l**2*(1+mrat*(sin(phi1b-phi2b))**2)**2)
    Y31 = dt*((p1b-p2b*cos(phi1b-phi2b))/(m1*l**2*(1+mrat*(sin(phi1b-phi2b))**2)))
    Y32 = dt*((p2b*(1+mrat)-p1b*mrat*cos(phi1b-phi2b))/(m1*l**2*(1+mrat*(sin(phi1b-phi2b))**2)))
    Y33 = dt*(-m1*(1+mrat)*g*l*sin(phi1b)-Ab+Bb)
    Y34 = dt*(-m1*mrat*g*l*sin(phi2b)+Ab-Bb)
    phi1c = phi1+Y31
    phi2c = phi2+Y32
    p1c = p1+Y33
    p2c = p2+Y34
    Ac = p1c*p2c*sin(phi1c-phi2c)/(m1*l**2*(1+mrat*(sin(phi1c-phi2c))**2))
    Bc = (p1c**2*mrat-2*p1c*p2c*mrat*cos(phi1c-phi2c)+p2c**2*(1+mrat))*sin(2*(phi1c-phi2c))/(2*m1*l**2*(1+mrat*(sin(phi1c-phi2c))**2)**2)
    Y41 = dt*((p1c-p2c*cos(phi1c-phi2c))/(m1*l**2*(1+mrat*(sin(phi1c-phi2c))**2)))
    Y42 = dt*((p2c*(1+mrat)-p1c*mrat*cos(phi1c-phi2c))/(m1*l**2*(1+mrat*(sin(phi1c-phi2c))**2)))
    Y43 = dt*(-m1*(1+mrat)*g*l*sin(phi1c)-Ac+Bc)
    Y44 = dt*(-m1*mrat*g*l*sin(phi2c)+Ac-Bc)
    phi1 += (1/6)*(Y11+2*Y21+2*Y31+Y41)
    phi2 += (1/6)*(Y12+2*Y22+2*Y32+Y42)
    p1 += (1/6)*(Y13+2*Y23+2*Y33+Y43)
    p2 += (1/6)*(Y14+2*Y24+2*Y34+Y44)
    phi1dot = (p1-p2*cos(phi1-phi2))/(m1*l**2*(1+mrat*(sin(phi1-phi2))**2))
    phi2dot = (p2*(1+mrat)-p1*mrat*cos(phi1-phi2))/(m1*l**2*(1+mrat*(sin(phi1-phi2))**2))
    frame1.rotate(axis=(0,0,1), angle=phi1dot*dt)
    frame2.pos = top+frame1.axis*l
    frame2.rotate(axis=(0,0,1), angle=phi2dot*dt)
    t = t+dt
    print(phi2dot)