#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Created on Fri Jun 28 15:55:03 2019

@author: kerkmann
"""

import numpy as np
import matplotlib.pyplot as plt

# Newton method (from exercise 13)
def Newton(f,df,x0,tol=1e-8,max_it=100):
    '''Newton algorithm to find root of function f'''
    # Input parameters:
    # f        function to find root of
    # df       derivative of f
    # x0       initial guess / starting value
    # tol      tolerance to finish iteration
    # max_it   maximum number of iterations
            
    # Output parameters:
    # x        solution
    # it       number of iterations
    # n        size of next step
    
    
    it = 0
    d = np.linalg.solve(np.atleast_2d(df(x0)),np.atleast_1d(f(x0)))
    x = x0
    
    while np.linalg.norm(d,np.inf) > tol and it < max_it:
        x = x - d
        d = np.linalg.solve(np.atleast_2d(df(x)),np.atleast_1d(f(x)))
        it += 1
    
    # Warning if maximimum number of iterations has been reached
    if it == max_it:
        print('Warning: Maximum number of iterations reached. Newton\'s method might not have converged.')
        
    return x, it, np.linalg.norm(d)

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

# Runge Kutta method of fourth oder (from exercise 27)
def RungeKutta4(x,t,f,k):
    F1 = f(x,t)
    F2 = f(x + k/2*F1,t+k/2)
    F3 = f(x + k/2*F2,t+k/2)
    F4 = f(x + k*F3,t+k)
    return x + k/6*(F1 + 2*F2 + 2*F3 + F4)

# BDF1 method
def BDF1(x,t,f,df,k):
    g = lambda u: u - k*f(u,t+k) - x
    dg = lambda u: np.eye(len(x)) - k*df(u,t+k)
    # explicit Euler as initial guess
    x0 = Euler_explicit(x,t,f,k)
    return Newton(g,dg,x0)[0]

# BDF2 method
def BDF2(x,xm1,t,f,df,k):
    g = lambda u: 3*u-4*x+xm1 - 2*k*f(u,t+k)
    dg = lambda u: 3*np.eye(len(x)) - 2*k*df(u,t+k)
    # explicit Euler as initial guess
    x0 = Euler_explicit(x,t,f,k)
    return Newton(g,dg,x0)[0]

### MAIN PROGRAM ###    

# first case (3 Punkte)
eps = 1
d = 4

f = lambda u,t: np.array([u[1],1/eps*(d*(1-u[0]**2)*u[1] - u[0])])
df = lambda u,t: np.array([[0, 1],[- 1/eps*d*2*u[0]*u[1] - 1/eps,1/eps*(d*(1-u[0]**2))]])

t_0 = 0
T = 16

u0 = np.array([1,1])

k = .03

# initial data
U = [u0]
t = [t_0]
n = 0

while t[n]<T-1e-14:
    # match final time
    if (t[-1]+k) > T:
        k = T - t[-1]

    # approximation with Runge Kutta 4
    U.append(RungeKutta4(U[n],t[-1],f,k))
    n = n+1
    t.append(t[-1] + k)

U = np.array(U)
t = np.array(t)
plt.figure(1)
plt.plot(U[:,0],U[:,1],label='RungeKutta4')
plt.axis('equal')

# second case (3 Punkte)
eps = .001
d = 1

f = lambda u,t: np.array([u[1],1/eps*(d*(1-u[0]**2)*u[1] - u[0])])
df = lambda u,t: np.array([[0,1],[- 1/eps*d*2*u[0]*u[1] - 1/eps,1/eps*(d*(1-u[0]**2))]])

t_0 = 0
T = 4

u0 = np.array([2,0])

k = .03

# initial data
U = [u0]
t = [t_0]
n = 0

# one step with BDF1
U.append(BDF1(U[n],t[-1],f,df,k))
n = n+1
t.append(t[-1] + k)

while t[n]<T-1e-14:
    # match final time
    if (t[-1]+k) > T:
        k = T - t[-1]

    # approximation with BDF 2 method
    U.append(BDF2(U[n],U[n-1],t[-1],f,df,k))
    #U.append(RungeKutta4(U[n],t[-1],f,k))
    n = n+1
    t.append(t[-1] + k)

U = np.array(U)
t = np.array(t)
plt.figure(2)
#plt.plot(t,U[:,0])
plt.plot(U[:,0],U[:,1],label='BDF2')