'''
Lorenz equation using RK4
'''

import numpy as np
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D


def oppgave_6_4_12(y_0,a,b,h):
    n = round((b-a)/h)
    t = a
    # start the integration loop
    y = np.zeros((n+1,3))
    y[0,:] = y_0
    for i in range(n):
        y[i+1,:] = RK4step(t,y[i,:],h)
        t = t+h
    # end for
    return y
    


def RK4step(t,y,h):
    # one step of explicit RK4
    # Input: current time t, current value y,  stepsize h
    # Output: approximate solution value at time t+h
    k1=ydot(t,y)
    k2=ydot(t+h/2,y+h/2*k1)
    k3=ydot(t+h/2,y+h/2*k2)
    k4=ydot(t+h,y+h*k3)
    y1=y + h/6*(k1+2*k2+2*k3+k4)
    return y1

def ydot(t,y):
    # right-hand side of the ODE
    s,r,b = 10., 28., 8./3.
    z = np.zeros(3)
    z[0] = -s*y[0] + s*y[1]
    z[1] = -y[0]*y[2]+r*y[0]-y[1]
    z[2] = y[0]*y[1]-b*y[2]
    return z

if __name__=='__main__':
    y_0=np.array([5.0,5.0,5.0])
    a=0.
    b=10.
    h=1.0E-03
    y = oppgave_6_4_12(y_0,a,b,h)

    # plot the approximation
    fig = plt.figure()
    ax = Axes3D(fig)
    ax.plot(xs=y[:,0], ys=y[:,1], zs=y[:,2],lw=2)
    ax.set_title('Approximate solution to Lorenz equations')
    ax.set_xlabel('x(t)')
    ax.set_ylabel('y(t)')
    ax.set_zlabel('z(t)')    
    plt.show()

    print('Final point\nx=%e, y=%e, z=%e\n' % (y[-1,0],y[-1,1],y[-1,2]))