#import all libraries
import numpy as np
import matplotlib.pyplot as plt
import matplotlib.animation as animation
from mpl_toolkits.mplot3d import Axes3D

#RK4
def rungeKutta4(f, t_i, y_i, dt):
  f_1 = f(t_i, y_i)
  f_2 = f(t_i + dt/2, y_i + (1/2)*dt*f_1)
  f_3 = f(t_i + dt/2, y_i + (1/2)*dt*f_2)
  f_4 = f(t_i + dt, y_i + dt*f_3)

  y_next = y_i + (1/6)*dt*(f_1 + 2*f_2 + 2*f_3 + f_4)

  return y_next


'''
Attributes
----------
- ts (list)      : saves the time steps
- sigma (float)  : value for sigma
- b (float)      : value for b
- r (float)      : value for r
- dt (float)     : value for the step size (in seconds)
- xs (list)      : list of xs
- ys (list)      : list of ys
- zs (list)      : list of zs


Methods
-------
- update        : uses RK4 to update the variables
- derivatives   : returns the derivatives of the system variables
- display       : display 3D trajectory graph
'''

class ConvectionRoll:
    '''
    This function initializes the class. The attributes are described above.
    '''
    def __init__(self, x0, y0, z0, sigma, b, r, dt):
        self.xs  = [x0]
        self.ys = [y0]
        self.zs  = [z0]
        self.step = dt

        self.sigma = sigma
        self.b = b
        self.r = r
        
        self.ts = [0]
        self.state = (self.xs[-1], self.ys[-1], self.zs[-1])

    def derivative(self, time, state):
        '''
        return the derivatives of our state variables
        '''
        x,y,z = state
        sigma, b, r = self.sigma, self.b, self.r

        x_prime = sigma*(y-x)
        y_prime = r*x - y - x*z
        z_prime = x*y - b*z

        return np.array((x_prime, y_prime, z_prime))

    def update(self):
        '''
        Updates the variables using RK4. 
        '''
        state = self.state
        next_state = rungeKutta4(self.derivative, self.ts[-1], state, self.step)
		        
        x, y, z = next_state

        self.xs.append(x)
        self.ys.append(y)
        self.zs.append(z)

        self.state = next_state

        t_next = self.ts[-1] + self.step
        self.ts.append(t_next)

    def display(self):
        fig = plt.figure(figsize=(14,10))
        points = np.array([self.xs, self.ys, self.zs])
        
        ax = fig.add_subplot(111, projection = '3d')
        ax.plot(points[0], points[1], points[2], marker = '.', linewidth = 0.5)
        ax.scatter(*points.T[0], color = 'red') # initial point
        
        plt.title(f'$\sigma$ = {self.sigma}, r ={self.r:0.2f}, \
        b ={self.b:0.2f}, Time = {self.ts[-1]:0.2f}')
        ax.set_xlabel('x (t)')
        ax.set_ylabel('y (t)')
        ax.set_zlabel('z (t)')
        plt.show()

    

#initialize the system with given values
sys = ConvectionRoll(0,1,0, sigma=10, b=8/3, r=20, dt= 0.01)

#let the system run for a while before diplaying the system
for i in range(5000):
    sys.update()
sys.display()

plt.figure(figsize=(22,12))
rs = np.linspace(20, 30, 12)

for i in range(len(rs)//2):
    plt.subplot(3,2,i+1)

    r = rs[2*i]
    sys = ConvectionRoll(0,1,0, sigma=10, b=8/3, r=r, dt= 0.01)
    for _ in range(5000):
        sys.update()
    plt.plot(sys.ts, sys.xs, label = 'r = '+str(np.round(r,2)))

    r = rs[2*i+1]
    sys = ConvectionRoll(0,1,0, sigma=10, b=8/3, r=r, dt= 0.01)
    for _ in range(5000):
        sys.update()
    plt.plot(sys.ts, sys.xs, label = 'r = '+str(np.round(r,2)))

    plt.xlabel('Time, t')
    plt.ylabel('x(t)')
    plt.legend()
plt.show()

plt.figure(figsize=(22,12))
rs = np.linspace(20, 30, 12)

for i in range(len(rs)//2):
    plt.subplot(3,2,i+1)

    r = rs[2*i]
    sys = ConvectionRoll(0,1,0, sigma=10, b=8/3, r=r, dt= 0.01)
    for _ in range(5000):
        sys.update()
    plt.plot(sys.ts, sys.ys, label = 'r = '+str(np.round(r,2)))

    r = rs[2*i+1]
    sys = ConvectionRoll(0,1,0, sigma=10, b=8/3, r=r, dt= 0.01)
    for _ in range(5000):
        sys.update()
    plt.plot(sys.ts, sys.ys, label = 'r = '+str(np.round(r,2)))

    plt.xlabel('Time, t')
    plt.ylabel('y(t)')
    plt.legend()
plt.show()

plt.figure(figsize=(22,12))
rs = np.linspace(20, 30, 12)

for i in range(len(rs)//2):
    plt.subplot(3,2,i+1)

    r = rs[2*i]
    sys = ConvectionRoll(0,1,0, sigma=10, b=8/3, r=r, dt= 0.01)
    for _ in range(5000):
        sys.update()
    plt.plot(sys.ts, sys.zs, label = 'r = '+str(np.round(r,2)))

    r = rs[2*i+1]
    sys = ConvectionRoll(0,1,0, sigma=10, b=8/3, r=r, dt= 0.01)
    for _ in range(5000):
        sys.update()
    plt.plot(sys.ts, sys.zs, label = 'r = '+str(np.round(r,2)))

    plt.xlabel('Time, t')
    plt.ylabel('z(t)')
    plt.legend()
plt.show()

fig = plt.figure(figsize=(16,16))
rs = np.linspace(20, 30, 6)

for i in range(len(rs)):
    ax = fig.add_subplot(3,2,i+1, projection = '3d')

    r = rs[i]
    sys = ConvectionRoll(0,1,0, sigma=10, b=8/3, r=r, dt= 0.01)
    for _ in range(10000):
        sys.update()

    points = np.array([sys.xs, sys.ys, sys.zs])
    ax.plot(points[0], points[1], points[2], marker = '.', linewidth = 0.5)
    ax.scatter(*points.T[0], color = 'red')
    
    plt.title(f'$\sigma$ = {sys.sigma}, r ={sys.r:0.2f}, \
    b ={sys.b:0.2f}, Time = {sys.ts[-1]:0.2f}')
    ax.set_xlabel('x (t)')
    ax.set_ylabel('y (t)')
    ax.set_zlabel('z (t)')

plt.show()



fig = plt.figure(figsize=(16,16))
rs = np.linspace(24, 25, 6)

for i in range(len(rs)):
    ax = fig.add_subplot(3,2,i+1, projection = '3d')

    r = rs[i]
    sys = ConvectionRoll(0,1,0, sigma=10, b=8/3, r=r, dt= 0.01)
    for _ in range(10000):
        sys.update()

    points = np.array([sys.xs, sys.ys, sys.zs])
    ax.plot(points[0], points[1], points[2], marker = '.', linewidth = 0.5)
    ax.scatter(*points.T[0], color = 'red')
    
    plt.title(f'$\sigma$ = {sys.sigma}, r ={sys.r:0.2f}, \
    b ={sys.b:0.2f}, Time = {sys.ts[-1]:0.2f}')
    ax.set_xlabel('x (t)')
    ax.set_ylabel('y (t)')
    ax.set_zlabel('z (t)')

plt.show()



y0_1 = 1
y0_2 = 1.0000001
sys1 = ConvectionRoll(0, y0_1, 0, sigma=10, b = 8/3, r = 28, dt = 0.01)
sys2 = ConvectionRoll(0, y0_2, 0, sigma=10, b = 8/3, r = 28, dt = 0.01)

for i in range(10000):
    sys1.update()
    sys2.update()

sys1.display()
sys2.display()

del_x = np.array(sys1.xs) - np.array(sys2.xs)
del_y = np.array(sys1.ys) - np.array(sys2.ys)
del_z = np.array(sys1.zs) - np.array(sys2.zs)
difference = np.sqrt(del_x**2 + del_y**2 + del_z**2)

plt.figure(figsize = (14,6))

plt.subplot(121)
plt.plot(sys1.ts, difference)
plt.xlabel('Time, t')
plt.ylabel('Difference')
plt.title('Regular Plot')

plt.subplot(122)
plt.loglog(sys1.ts, difference)
plt.xlabel('Time, t')
plt.ylabel('Difference')
plt.title('Log-Log Plot')

plt.show()

from IPython.display import HTML

def make_animation(systems, n_update = 30, n_animate = 50, interval = 200):
    n = len(systems)

    fig = plt.figure(figsize=(14,10))
    ax = fig.add_subplot(111, projection = '3d')
    lines = [None]*n

    for k in range(n):
        sys = systems[k]
        points = np.array([sys.xs, sys.ys, sys.zs])
        lines[k], = ax.plot(points[0], points[1], points[2], marker = '.', 
                        linewidth = 0.5, color = 'C'+str(k))

    def animate(i):
        for _ in range(n_update):
            for sys in systems:
                sys.update()
        
        for k in range(n):
            sys = systems[k]
            points = np.array([sys.xs, sys.ys, sys.zs])
            lines[k], = ax.plot(points[0], points[1], points[2], marker = '.', 
                            linewidth = 0.5, color = 'C'+str(k))
        ax.set_title(f'$\sigma$ = {sys.sigma}, r = {sys.r:0.2f}, b = {sys.b:0.2f}, Time = {sys.ts[-1]:0.2f}', loc = 'center')
        return lines

    # Setting the axes properties
    ax.set_xlim3d([-20, 20])
    ax.set_xlabel('X')

    ax.set_ylim3d([-20, 20])
    ax.set_ylabel('Y')

    ax.set_zlim3d([0, 40])
    ax.set_zlabel('Z')

    

    # starts animation
    ani = animation.FuncAnimation(fig, animate, n_animate, interval = interval, blit = True)
    
    return ani



y0_1 = 1
sys1 = ConvectionRoll(0, y0_1, 0, sigma=10, b = 8/3, r = 20, dt = 0.01)

ani = make_animation([sys1], n_update=20, n_animate=120)
HTML(ani.to_html5_video())

y0_1 = 1
sys1 = ConvectionRoll(0, y0_1, 0, sigma=10, b = 8/3, r = 28, dt = 0.01)

ani = make_animation([sys1], n_update=20, n_animate=120)
HTML(ani.to_html5_video())


sys = ConvectionRoll(0,1,0, sigma=10, b=8/3, r=20, dt= 0.01)

n_animate = 44 
interval = 800

fig = plt.figure(figsize = (10,6))
ax = plt.subplot(ylim=(-20, 20))

lines = [None]


rs = np.linspace(20,30,n_animate)

def animate(i):
    ax.clear()
    r = rs[i]
    sys = ConvectionRoll(0,1,0, sigma=10, b=8/3, r=r, dt= 0.01)
    for _ in range(6000):
        sys.update()

    lines[0], = ax.plot(sys.ts, sys.xs)

    # Setting the axes properties
    ax.set_xlabel('Time, t')

    ax.set_ylim([-20, 20])
    ax.set_ylabel('x (t)')
    
    ax.set_title(f'$\sigma$ = {sys.sigma}, b = {sys.b:0.2f}, r = {sys.r:0.2f}, Time = {sys.ts[-1]:0.2f}', 
                 loc = 'center', fontweight = 3, fontsize = 15)
    plt.grid()
    return lines

# starts animation
ani = animation.FuncAnimation(fig, animate, n_animate, interval = interval, blit = True)
    
HTML(ani.to_html5_video())



sys = ConvectionRoll(0,1,0, sigma=10, b=8/3, r=20, dt= 0.01)

n_animate = 44 
interval = 800

fig = plt.figure(figsize = (10,6))
ax = plt.subplot(ylim=(-20, 20))

lines = [None]


rs = np.linspace(20,30,n_animate)

def animate(i):
    ax.clear()
    r = rs[i]
    sys = ConvectionRoll(0,1,0, sigma=10, b=8/3, r=r, dt= 0.01)
    for _ in range(6000):
        sys.update()

    lines[0], = ax.plot(sys.ts, sys.ys)

    # Setting the axes properties
    ax.set_xlabel('Time, t')

    ax.set_ylim([-20, 20])
    ax.set_ylabel('y (t)')
    
    ax.set_title(f'$\sigma$ = {sys.sigma}, b = {sys.b:0.2f}, r = {sys.r:0.2f}, Time = {sys.ts[-1]:0.2f}', 
                 loc = 'center', fontweight = 3, fontsize = 15)
    plt.grid()
    return lines

# starts animation
ani = animation.FuncAnimation(fig, animate, n_animate, interval = interval, blit = True)
    
HTML(ani.to_html5_video())



sys = ConvectionRoll(0,1,0, sigma=10, b=8/3, r=20, dt= 0.01)

n_animate = 44 
interval = 800

fig = plt.figure(figsize = (10,6))
ax = plt.subplot(ylim=(0, 50))

lines = [None]


rs = np.linspace(20,30,n_animate)

def animate(i):
    ax.clear()
    r = rs[i]
    sys = ConvectionRoll(0,1,0, sigma=10, b=8/3, r=r, dt= 0.01)
    for _ in range(6000):
        sys.update()

    lines[0], = ax.plot(sys.ts, sys.zs)

    # Setting the axes properties
    ax.set_xlabel('Time, t')

    ax.set_ylim([0, 50])
    ax.set_ylabel('z (t)')
    
    ax.set_title(f'$\sigma$ = {sys.sigma}, b = {sys.b:0.2f}, r = {sys.r:0.2f}, Time = {sys.ts[-1]:0.2f}', 
                 loc = 'center', fontweight = 3, fontsize = 15)
    plt.grid()
    return lines

# starts animation
ani = animation.FuncAnimation(fig, animate, n_animate, interval = interval, blit = True)
    
HTML(ani.to_html5_video())



sys = ConvectionRoll(0,1,0, sigma=10, b=8/3, r=20, dt= 0.01)

n_animate = 41 
interval = 600

fig = plt.figure(figsize=(12,9))
ax = fig.add_subplot(111, projection = '3d')
lines = [None]

points = np.array([sys.xs, sys.ys, sys.zs])
lines[0], = ax.plot(points[0], points[1], points[2], marker = '.', 
                    linewidth = 0.5, color = 'blue')

# Setting the axes properties
ax.set_xlim3d([-20, 20])
ax.set_xlabel('X')

ax.set_ylim3d([-20, 20])
ax.set_ylabel('Y')

ax.set_zlim3d([0, 40])
ax.set_zlabel('Z')

rs = np.linspace(20,30,n_animate)

def animate(i):
    ax.clear()
    r = rs[i]
    sys = ConvectionRoll(0,1,0, sigma=10, b=8/3, r=r, dt= 0.01)
    for _ in range(4000):
        sys.update()

    points = np.array([sys.xs, sys.ys, sys.zs])
    lines[0], = ax.plot(points[0], points[1], points[2], marker = '.', 
                    linewidth = 0.5, color = 'blue')
    
    ax.set_title(f'$\sigma$ = {sys.sigma}, b = {sys.b:0.2f}, r = {sys.r:0.2f}, Time = {sys.ts[-1]:0.2f}', 
                 loc = 'center', fontweight = 3, fontsize = 15)
    
    # Setting the axes properties
    ax.set_xlim3d([-20, 20])
    ax.set_xlabel('X')

    ax.set_ylim3d([-20, 20])
    ax.set_ylabel('Y')

    ax.set_zlim3d([0, 40])
    ax.set_zlabel('Z')

    return lines

# starts animation
ani = animation.FuncAnimation(fig, animate, n_animate, interval = interval, blit = True)
    
HTML(ani.to_html5_video())


y0_1 = 1
y0_2 = 1.0000001
sys1 = ConvectionRoll(0, y0_1, 0, sigma=10, b = 8/3, r = 28, dt = 0.01)
sys2 = ConvectionRoll(0, y0_2, 0, sigma=10, b = 8/3, r = 28, dt = 0.01)

n_update = 30

initial_steps = 2000
for _ in range(initial_steps):
    sys1.update()
    sys2.update()

fig = plt.figure(figsize=(14,10))
ax = fig.add_subplot(111, projection = '3d')
lines = [None, None]

points = np.array([sys1.xs, sys1.ys, sys1.zs])
lines[0], = ax.plot(points[0][-n_update:], points[1][-n_update:], 
                    points[2][-n_update:], marker = '.', 
                    linewidth = 0.5, color = 'blue')
points = np.array([sys2.xs, sys2.ys, sys2.zs])
lines[1], = ax.plot(points[0][-n_update:], points[1][-n_update:], 
                    points[2][-n_update:], marker = '.', 
                    linewidth = 0.5, color = 'red')

# Setting the axes properties
ax.set_xlim3d([-20, 20])
ax.set_xlabel('x(t)')

ax.set_ylim3d([-20, 20])
ax.set_ylabel('y(t)')

ax.set_zlim3d([0, 40])
ax.set_zlabel('z(t)')

def animate(i):
    ax.clear()
    for _ in range(n_update):
        sys1.update()
        sys2.update()
    
    points = np.array([sys1.xs, sys1.ys, sys1.zs])
    lines[0], = ax.plot(points[0][-n_update:], points[1][-n_update:], 
                        points[2][-n_update:], marker = '.', 
                        linewidth = 0.5, color = 'blue')
    points = np.array([sys2.xs, sys2.ys, sys2.zs])
    lines[1], = ax.plot(points[0][-n_update:], points[1][-n_update:], 
                        points[2][-n_update:], marker = '.', 
                        linewidth = 0.5, color = 'red')

    # Setting the axes properties
    ax.set_xlim3d([-20, 20])
    ax.set_xlabel('x(t)')

    ax.set_ylim3d([-20, 20])
    ax.set_ylabel('y(t)')

    ax.set_zlim3d([0, 40])
    ax.set_zlabel('z(t)')

    ax.set_title(f'$\sigma$ = {sys1.sigma}, b = {sys1.b:0.2f}, r = {sys1.r:0.2f}, Time = {sys1.ts[-1]:0.2f}', 
                loc = 'center', fontweight = 3, fontsize = 15)

    return lines


# starts animation
ani = animation.FuncAnimation(fig, animate, 150, interval = 200, blit = True)
HTML(ani.to_html5_video())