# three equations of lorenz system:
# 1. x'= sigma(y-x)
# 2. y'= x(rho-z)-y
# 3. z'= xy - beta*z
# we must choose proper sigma rho and beta so stimulate air
#  rho=28
#  sigma = 10
#  beta = 8/3


#import function
from scipy.integrate import odeint

import numpy as np
import matplotlib.pyplot as plt
import matplotlib.animation as animation
import mpl_toolkits.mplot3d.axes3d as p3

# lorenz system function
def lorenz(loc,t,rho,beta,sigma):
    (x,y,z) = loc
    return np.array([sigma*(y-x),x*(rho-z)-y,x*y-beta*z])


#set t range
t = np.arange(0, 30.0 , 0.01) #也可設定 0.1,比較快


#use ode to solve lorenz system
sol1 = odeint(lorenz,(0,0,0),t,args=(28,8/3,10))
sol2 = odeint(lorenz,(2.0,2.0,8.0),t,args=(28,8/3,10))
sol3 = odeint(lorenz,(2.0,2.0,8.1),t,args=(28,8/3,10))
sol4 = odeint(lorenz,(1.0,1.0,10.0),t,args=(28,8/3,10))


sol1 = sol1.T
sol2 = sol2.T
sol3 = sol3.T
sol4 = sol4.T


#draw pic

lorenzpic = plt.figure()
ax = p3.Axes3D(lorenzpic)


#set axes and range
ax.set_xlabel('X')
ax.set_ylabel('Y')
ax.set_zlabel('Z')
ax.set_xlim3d([-100,100])
ax.set_ylim3d([-100,100])
ax.set_zlim3d([0,100])

#set ball and line
ball1, = ax.plot([],[],linestyle='None',marker='o',markersize=8,markeredgecolor='r',color='black',markeredgewidth=2)
ball2, = ax.plot([],[],linestyle='None',marker='o',markersize=8,markeredgecolor='y',color='black',markeredgewidth=2)
ball3, = ax.plot([],[],linestyle='None',marker='o',markersize=8,markeredgecolor='b',color='black',markeredgewidth=2)
ball4, = ax.plot([],[],linestyle='None',marker='o',markersize=8,markeredgecolor='green',color='black',markeredgewidth=2)
line1, = ax.plot([],[],color='r')
line2, = ax.plot([],[],color='y')
line3, = ax.plot([],[],color='b')
line4, = ax.plot([],[],color='green')


def init():
    ball1.set_data(sol1[0:2,0])
    ball1.set_3d_properties(sol1[2,0])
    ball2.set_data(sol2[0:2,0])
    ball2.set_3d_properties(sol2[2,0])
    ball3.set_data(sol3[0:2,0])
    ball3.set_3d_properties(sol3[2,0])
    ball4.set_data(sol4[0:2,0])
    ball4.set_3d_properties(sol4[2,0])
    line1.set_data(sol1[0:2,0])
    line1.set_3d_properties(sol1[2,0])
    line2.set_data(sol2[0:2,0])
    line2.set_3d_properties(sol2[2,0])
    line3.set_data(sol3[0:2,0])
    line3.set_3d_properties(sol3[2,0])
    line4.set_data(sol4[0:2,0])
    line4.set_3d_properties(sol4[2,0])

    return ball1,ball2,ball3,ball4,line1,line2,line3,line4

def animate(i):
    ball1.set_data(sol1[0:2,i])
    ball1.set_3d_properties(sol1[2,i])
    ball2.set_data(sol2[0:2,i])
    ball2.set_3d_properties(sol2[2,i])
    ball3.set_data(sol3[0:2,i])
    ball3.set_3d_properties(sol3[2,i])
    ball4.set_data(sol4[0:2,i])
    ball4.set_3d_properties(sol4[2,i])
    ball1.set_data(sol1[0:2,i])
    line1.set_3d_properties(sol1[2,:i])
    line2.set_data(sol2[0:2,:i])
    line2.set_3d_properties(sol2[2,:i])
    line3.set_data(sol3[0:2,:i])
    line3.set_3d_properties(sol3[2,:i])
    line4.set_data(sol4[0:2,:i])
    line4.set_3d_properties(sol4[2,:i])

    return ball1,ball2,ball3,ball4,line1,line2,line3,line4

line_ani = animation.FuncAnimation(lorenzpic,animate,np.arange(1,t.size),init_func=init,interval=20,repeat=False,blit=False)
plt.show()