import numpy as np
import matplotlib.pyplot as plt


def lorenz(x, y, z, s=10, r=28, b=2.667):
    """
    Given:
       x, y, z: a point of interest in three dimensional space
       s, r, b: parameters defining the lorenz attractor
    Returns:
       x_dot, y_dot, z_dot: values of the lorenz attractor's partial
           derivatives at the point x, y, z
    """
    x_dot = s*(y - x)
    y_dot = r*x - y - x*z
    z_dot = x*y - b*z
    return x_dot, y_dot, z_dot


dt = 0.01
num_steps = 10000


class LorenzEquation:

    def __init__(self, a = 0., b = 1., c = 1.05):
        self.a = a
        self.b = b
        self.c = c
    
    def numerical_iteration_simulation(self):
        xs = np.empty(num_steps + 1)
        ys = np.empty(num_steps + 1)
        zs = np.empty(num_steps + 1)

        #Set initial Values
        xs[0], ys[0], zs[0] = (self.a, self.b, self.c)
        
        # Step through time, calculating the partial derivatives at the current point
        # and using them to estimate the next point 
        for i in range(num_steps): 
            x_dot, y_dot, z_dot = lorenz(xs[i], ys[i], zs[i])
            xs[i + 1] = xs[i] + (x_dot * dt)
            ys[i + 1] = ys[i] + (y_dot * dt)
            zs[i + 1] = zs[i] + (z_dot * dt)
        ax = plt.figure().add_subplot(projection='3d')
        ax.plot(xs, ys, zs, lw=0.5)
        ax.set_xlabel("X Axis")
        ax.set_ylabel("Y Axis")
        ax.set_zlabel("Z Axis")
        ax.set_title("Lorenz Attractor")
        plt.show()

LorenzEquation(0., 1., 1.05).numerical_iteration_simulation()
