#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Created on Mon Jun  3 13:56:46 2019

@author: kerkmann
"""

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

# predictor corrector method of 3rd order (1 Punkt)
def pc(x,xm1,t,f,k):
    predictor = x + k/2*(3*f(x,t) - f(xm1,t))
    corrector = x + k/12*(8*f(x,t) - f(xm1,t) + 5*f(predictor,t))
    return corrector

# improved Eulers method (from exercise 27)
def Euler_improved(x,t,f,k):
    zwischen = x + k/2*f(x,t)
    return x + k*f(zwischen,t+k/2)

### MAIN PROGRAM ###

# parameters
P = 10
r = 28
b = 8/3
t_0 = 0
T = 250
k = 0.01
n = 0

f = lambda x,t: np.array([-P*x[0] + P*x[1], r*x[0] - x[1] - x[0]*x[2], x[0]*x[1] - b*x[2]])

# initial data
U = [np.array([0.4,1,9])]
t = t_0

# first step (use 2nd order improved Eulers method)
U.append(Euler_improved(U[0],t,f,k))
n = n+1
t = t+k

# loop until final time (1 Punkt)
while t<T:
    # match final time
    if (t+k) > T:
        k = T - t
        
    U.append(pc(U[n],U[n-1],t,f,k))
    n = n+1
    t = t+k
    
# 3D plot (1 Punkt)
U = np.array(U)
fig = plt.figure()
ax = fig.gca(projection='3d')
ax.plot(U[:,0], U[:,1], U[:,2], linewidth=.1)