""" 
    Ball on a linear spring in 3-D, in vacuum, also subject to unit 
    downward gravity. Gets pretty chaotic. 
    
    Change the last digits in the initial conditions to see it!
    Red symbol marks the position of the ball at t=100.
    It'll shift unpredictably across the allowed region of
    motion, as a result of a tiny change in the input data or 
    in numerical treatment (such as the value of time step dt)
"""

import numpy as np
import matplotlib.pyplot as plt
# This import registers the 3D projection, but is otherwise unused.
from mpl_toolkits.mplot3d import Axes3D  # noqa: F401 unused import

T = 100       # finish time
dt = 0.000025 # time step
num_steps = int(T/dt)
skip = 50     # compute energy skipping this many steps
#  storage for later plots
xs = np.empty(num_steps+1)
ys = np.empty(num_steps+1)
zs = np.empty(num_steps+1)
errs = np.empty(num_steps//50+1)
# Set initial values
x, y, z = (0.000002 +0.000001*.1, -0.0000020, 0.3)
print("initial position",x,y,z,", mass at at rest.")
vx, vy, vz  = (0.,0.,0.) 

# Step through time, calculating the derivatives
# and updating variables. Leapfrog integration scheme
for i in range(num_steps+1):
    xs[i], ys[i], zs[i] = (x,y,z)
    r = np.sqrt(x*x + y*y + z*z)
    f_r = -4.*(r-1.)/r  # radial acceleration / r

    fx = x * f_r      # x/r etc. are the components of radial versor
    fy = y * f_r
    fz = z * f_r -1.  # -1 is downward gravity of unit strength
    
    vx = vx + dt*fx    
    vy = vy + dt*fy     
    vz = vz + dt*fz 
    
    x = x + dt*vx    
    y = y + dt*vy    
    z = z + dt*vz 
# error of energy (synchronized correctly)
    if (i%skip == 0):
        xsyn = x -dt*vx/2
        ysyn = y -dt*vy/2
        zsyn = z -dt*vz/2
        rsyn = (xsyn**2+ysyn**2+zsyn**2)**0.5
        Phi = 2.*(rsyn-1.)**2  +zsyn     # potential energy/mass
        E = Phi + (vx*vx+vy*vy+vz*vz)/2  # total energy/mass
        if (i==0): 
            E0 = E
        errs[i//skip] = E/E0 -1   # relative energy error
# Plot
fig = plt.figure(dpi=140)
ax = fig.gca(projection='3d')

ax.plot(xs, ys, zs, lw=0.5, alpha=0.66)
ax.scatter(xs[0],ys[0],zs[0],marker='s',c='b')
ax.scatter(x,y,z,marker='o',c='r')
ax.set_xlabel("X Axis")
ax.set_ylabel("Y Axis")
ax.set_zlabel("Z Axis")
ax.set_title("Ball on Hook's spring k=4, gravity g=1")


plt.show()
print("dt=",dt,".  ending time=",dt*num_steps,"   ",num_steps," steps")
print("final position",x,y,z)
plt.figure(dpi=130)
plt.cla()
plt.xlabel(' time')
plt.ylabel(' dE/E$_0$')
plt.grid(True)
plt.plot(np.linspace(0,100,num_steps//20),errs[0:-1])
plt.show()