import vpython as vp

# https://www.glowscript.org/#/user/GlowScriptDemos/folder/Examples/program/DoublePendulum-VPython/edit

# Double pendulum

# The analysis is in terms of Lagrangian mechanics.
# The Lagrangian variables are angle of upper bar, angle of lower bar,
# measured from the vertical.

# Bruce Sherwood

# Corrections to the Lagrangian calculations by Owen Long, UC. Riverside

vp.scene.width = vp.scene.height = 600
vp.scene.range = 1.8
vp.scene.title = "A double pendulum"


def display_instructions():
    s = """In GlowScript programs:
To rotate "camera", drag with right button or Ctrl-drag.
To zoom, drag with middle button or Alt/Option depressed, or use scroll wheel.
  On a two-button mouse, middle is left + right.
To pan left/right and up/down, Shift-drag.
Touch screen: pinch/extend to zoom, swipe or two-finger rotate."""
    vp.scene.caption = s


# Display text below the 3D graphics:
display_instructions()

g = 9.8
M1 = 2.0
M2 = 1.0
d = 0.05  # thickness of each bar
gap = 2*d  # distance between two parts of upper, U-shaped assembly
L1 = 0.5  # physical length of upper assembly; distance between axles
L1display = L1+d  # show upper assembly a bit longer than physical, to overlap axle
L2 = 1  # physical length of lower bar
L2display = L2+d/2  # show lower bar a bit longer than physical, to overlap axle
# Coefficients used in Lagrangian calculation
A = (1/4)*M1*L1**2+(1/12)*M1*L1**2+M2*L1**2
B = (1/2)*M2*L1*L2
C = g*L1*(M1/2+M2)
D = M2*L1*L2/2
E = (1/12)*M2*L2**2+(1/4)*M2*L2**2
F = g*L2*M2/2

hpedestal = 1.3*(L1+L2)  # 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 pedestal to center of U-shaped upper assembly
# top of inner bar of U-shaped upper assembly
pedestal_top = vp.vec(0, hpedestal/2, 0)

theta1 = 1.3*vp.pi/2  # initial upper angle (from vertical)
theta1dot = 0  # initial rate of change of theta1
theta2 = 0  # initial lower angle (from vertical)
theta2dot = 0  # initial rate of change of theta2

pedestal = vp.box(pos=pedestal_top-vp.vec(0, hpedestal/2, offset),
                  size=vp.vec(wpedestal, 1.1*hpedestal, wpedestal),
                  color=vp.vec(0.4, 0.4, 0.5))
base = vp.box(pos=pedestal_top-vp.vec(0, hpedestal+tbase/2, offset),
              size=vp.vec(wbase, tbase, wbase),
              color=pedestal.color)
axle1 = vp.cylinder(pos=pedestal_top-vp.vec(0, 0, gap/2-d/4), axis=vp.vec(0, 0, -1),
                    size=vp.vec(offset, d/4, d/4), color=vp.color.yellow)

bar1 = vp.box(pos=pedestal_top+vp.vec(L1display/2-d/2, 0, -(gap+d)/2),
              size=vp.vec(L1display, d, d), color=vp.color.red)
bar1.rotate(angle=-vp.pi/2, axis=vp.vec(0, 0, 1),
            origin=vp.vec(axle1.pos.x, axle1.pos.y, bar1.pos.z))
bar1.rotate(angle=theta1, axis=vp.vec(0, 0, 1), origin=vp.vec(
    axle1.pos.x, axle1.pos.y, bar1.pos.z))

bar1b = vp.box(pos=pedestal_top+vp.vec(L1display/2-d/2, 0, (gap+d)/2),
               size=vp.vec(L1display, d, d), color=bar1.color)
bar1b.rotate(angle=-vp.pi/2, axis=vp.vec(0, 0, 1),
             origin=vp.vec(axle1.pos.x, axle1.pos.y, bar1b.pos.z))
bar1b.rotate(angle=theta1, axis=vp.vec(0, 0, 1), origin=vp.vec(
    axle1.pos.x, axle1.pos.y, bar1b.pos.z))

pivot1 = vp.vec(axle1.pos.x, axle1.pos.y, 0)

axle2 = vp.cylinder(pos=pedestal_top+vp.vec(L1, 0, -(gap+d)/2), axis=vp.vec(0, 0, 1),
                    size=vp.vec(gap+d, axle1.size.y/2, axle1.size.y/2), color=axle1.color)
axle2.rotate(angle=-vp.pi/2, axis=vp.vec(0, 0, 1),
             origin=vp.vec(axle1.pos.x, axle1.pos.y, axle2.pos.z))
axle2.rotate(angle=theta1, axis=vp.vec(0, 0, 1), origin=vp.vec(
    axle1.pos.x, axle1.pos.y, axle2.pos.z))

bar2 = vp.box(pos=axle2.pos+vp.vec(L2display/2-d/2, 0, (gap+d)/2),
              size=vp.vec(L2display, d, d), color=vp.color.green)

bar2.rotate(angle=-vp.pi/2,  axis=vp.vec(0, 0, 1),
            origin=vp.vec(axle2.pos.x, axle2.pos.y, bar2.pos.z))
bar2.rotate(angle=theta2,  axis=vp.vec(0, 0, 1), origin=vp.vec(
    axle2.pos.x, axle2.pos.y, bar2.pos.z))

dt = 0.0002
t = 0

while True:
    vp.rate(1/dt)
    # Calculate accelerations of the Lagrangian coordinates=
    atheta1 = ((E*C/B)*vp.sin(theta1)-F*vp.sin(theta2))/(D-E*A/B)
    atheta2 = -(A*atheta1+C*vp.sin(theta1))/B
    # Update velocities of the Lagrangian coordinates=
    theta1dot = theta1dot+atheta1*dt
    theta2dot = theta2dot+atheta2*dt
    # Update Lagrangian coordinates=
    dtheta1 = theta1dot*dt
    dtheta2 = theta2dot*dt
    theta1 = theta1+dtheta1
    theta2 = theta2+dtheta2

    bar1.rotate(angle=dtheta1, axis=vp.vec(0, 0, 1), origin=pivot1)
    bar1b.rotate(angle=dtheta1, axis=vp.vec(0, 0, 1), origin=pivot1)
    pivot2 = vp.vec(axle2.pos.x, axle2.pos.y, pivot1.z)
    axle2.rotate(angle=dtheta1, axis=vp.vec(0, 0, 1), origin=pivot1)
    bar2.rotate(angle=dtheta2, axis=vp.vec(0, 0, 1), origin=pivot2)
    pivot2 = vp.vec(axle2.pos.x, axle2.pos.y, pivot1.z)
    bar2.pos = pivot2 + bar2.axis/2

    t = t+dt
