import numpy as np
from vpython import *
from scipy.integrate import solve_ivp

scene = canvas(width=480, height=480)
floor = box(pos=vector(0.,-1.1,0.), length=2.2, height=0.01, width=1.2, opacity=0.2)
ball = sphere(radius=0.05, color=color.red)
rope = curve(color = vector(1.,1.,0.8), opacity=0.8)

speed = 0.2
nseg = 80
m, dm = 1., 0.1/nseg
g, R0, k = 9.8, 1.0, 1000.*nseg

t = 0.
y = np.zeros(nseg*4)
y[nseg*0:nseg*1] = np.linspace(0.,0.8,nseg,endpoint=False)+0.8/nseg
y[nseg*1:nseg*2] = np.linspace(0.,0.6,nseg,endpoint=False)+0.6/nseg

def f(t,y):
    rx = y[nseg*0:nseg*1]
    ry = y[nseg*1:nseg*2]
    vx = y[nseg*2:nseg*3]
    vy = y[nseg*3:nseg*4]
    
    dx = rx.copy()
    dy = ry.copy()
    dx[1:] -= rx[:-1]
    dy[1:] -= ry[:-1]
    R = (dx**2+dy**2)**0.5
    fs = -k*(R-R0/nseg)
    
    dydt = np.zeros(nseg*4)
    dydt[nseg*0:nseg*1] = vx
    dydt[nseg*1:nseg*2] = vy
    ax = dydt[nseg*2:nseg*3]
    ay = dydt[nseg*3:nseg*4]
    
    ax[:-1] = fs[:-1]*dx[:-1]/R[:-1]/dm \
            - fs[1: ]*dx[1: ]/R[1: ]/dm
    ay[:-1] = fs[:-1]*dy[:-1]/R[:-1]/dm \
            - fs[1: ]*dy[1: ]/R[1: ]/dm - g
    ax[-1 ] = fs[-1]*dx[-1]/R[-1]/m
    ay[-1 ] = fs[-1]*dy[-1]/R[-1]/m - g
    
    return dydt

while True:
    sol = solve_ivp(f, [t, t+0.040*speed], y)
    y   = sol.y[:,-1]
    t   = sol.t[-1]

    rx = y[nseg*0:nseg*1]
    ry = y[nseg*1:nseg*2]
    
    rope.clear()
    for i in range(nseg):
        rope.append(vector(rx[i],ry[i],0))
    ball.pos.x = rx[nseg-1]
    ball.pos.y = ry[nseg-1]
    
    rate(1./0.040)