#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Created on Fri May 31 13:58:54 2019

@author: kerkmann
"""

import numpy as np
import matplotlib.pyplot as plt

# explicit Euler method (from exercise 27)
def Euler_explicit(x,t,f,k):
    return x + k*f(x,t)

# Heun's method (1 Punkt)
def Heun(x,t,f,k):
    F1 = f(x,t)
    F2 = f(x + k*F1,t+k)
    return x + k/2*(F1 + F2)

### MAIN PROGRAM ###

# parameters
a_1 = 5
a_2 = 0.1
b_1 = 3.5
b_2 = 2
d_1 = 0.4
d_2 = 0.01
T = 3000
k = 1
ks = []
t = 0
ts = [t]
n = 0
again = False

# onestep error control
eps = 1e-6

# right hand side of IVP
f_1 = lambda z: a_1*z/(1+b_1*z)
f_2 = lambda z: a_2*z/(1+b_2*z)
f = lambda x,t: np.array([x[0]*(1-x[0]) - f_1(x[0])*x[1], f_1(x[0])*x[1] - f_2(x[1])*x[2] - d_1*x[1], f_2(x[1])*x[2] - d_2*x[2]])

# initial data
U = [np.array([0.4,1,9])]
while t<T:
    # match final time
    if (t+k) > T:
        k = T - t

    # approximation with Euler
    U1 = Euler_explicit(U[n],t,f,k)
    
    # approximation with Heun
    U2 = Heun(U[n],t,f,k)
    
    # error estimate (1 Punkt)
    err = np.max(abs(U2-U1))
    
    # new stepsize (1 Punkt)
    k = k*min(2,max(0.2,0.9*np.sqrt(eps/err)))
    
    # was the step successful? (1 Punkt)
    if err <= eps:
        # step successful
        ks.append(k)
        n = n+1
        t = t+k
        ts.append(t)
        U.append(U2)
    else:
        # step unsuccessful
        # repeat step, nothing to do
        pass

# convert list to array
U = np.array(U)
ts = np.array(ts)
ks = np.array(ks)

# plot solution for t=2000 to t=3000 (1 Punkt)
ind = ts>2000
labels = ['c(t)','u(t)','v(t)']

plt.figure(0)
for sol,label in zip(U.transpose(),labels):
    plt.plot(ts[ind],sol[ind],label=label)
plt.legend()

# (1 Punkt)
plt.figure(1)
plt.plot(range(n),ks)