import numpy as np
import matplotlib.pyplot as plt

n = 1000000
h = 0.0001
T = h*n
t = np.arange(0,T,h)

PX =[]
PY =[]
PZ =[]

sigma = 10
b = 8/3
r = 28

def dxdt(x,y,z): 
    return -sigma * x + sigma * y
def dydt(x,y,z): 
    return r * x - y -x * z
def dzdt(x,y,z): 
    return -b * z + x * y

Xst = (b*(r-1))**0.5
Yst = (b*(r-1))**0.5
Zst = r - 1

x = np.empty(n)
y = np.empty(n)
z = np.empty(n)
x[0] = 2 
y[0] = 1.1 
z[0] = 20 

for i in range(n-1):
    k_1 = dxdt(x[i] , y[i] , z[i])
    j_1 = dydt(x[i] , y[i] , z[i])
    m_1 = dzdt(x[i] , y[i] , z[i])
       
    k_2 = dxdt(x[i] + k_1 * h / 2 , y[i] + j_1 * h / 2, z[i] + m_1 * h / 2)
    j_2 = dydt(x[i] + k_1 * h / 2 , y[i] + j_1 * h / 2, z[i] + m_1 * h / 2)
    m_2 = dzdt(x[i] + k_1 * h / 2 , y[i] + j_1 * h / 2, z[i] + m_1 * h / 2)

    k_3 = dxdt(x[i] + k_2 *h / 2 , y[i] + j_2 * h / 2 , z[i] + m_2 * h / 2)
    j_3 = dydt(x[i] + k_2 *h / 2 , y[i] + j_2 * h / 2 , z[i] + m_2 * h / 2)
    m_3 = dzdt(x[i] + k_2 *h / 2 , y[i] + j_2 * h / 2 , z[i] + m_2 * h / 2)

    k_4 = dxdt(x[i] + k_3 *h , y[i] + j_3 * h , z[i] + m_3 * h)
    j_4 = dydt(x[i] + k_3 *h , y[i] + j_3 * h , z[i] + m_3 * h)
    m_4 = dzdt(x[i] + k_3 *h , y[i] + j_3 * h , z[i] + m_3 * h)

    x[i+1] = x[i] + h/6 * (k_1 + 2*k_2 + 2*k_3 + k_4 )
    y[i+1] = y[i] + h/6 * (j_1 + 2*j_2 + 2*j_3 + j_4 )
    z[i+1] = z[i] + h/6 * (m_1 + 2*m_2 + 2*m_3 + m_4 )    

    if (z[i+1] - Zst) > 0 and (z[i] - Zst) < 0:
        PX.append(x[i])
        PY.append(y[i])
        PZ.append(Zst)

fig = plt.figure()
ax = fig.add_subplot(projection='3d')
ax.plot(x, y, z, color='lime', linewidth=0.5)
ax.scatter(PX, PY, PZ, marker=".", color='black')
ax.scatter(Xst, Yst, Zst, marker="o", color='blue')
ax.scatter(-Xst, -Yst, Zst, marker="o", color='blue')
plt.show()