#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Created on Mon May 27 11:21:22 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)

# (c) (1 Punkt)
def Euler_explicit(x,t,f,k):
    return x + k*f(x,t)

# (d) (1 Punkt)
def Euler_improved(x,t,f,k):
    zwischen = x + k/2*f(x,t)
    return x + k*f(zwischen,t+k/2)

# (e) (1 Punkt)
def RungeKutta4(x,t,f,k):
    Y1 = x
    Y2 = x + k/2*f(Y1,t)
    Y3 = x + k/2*f(Y2,t+k/2)
    Y4 = x + k*f(Y3,t+k/2)
    return x + k/6*(f(Y1,t) + 2*f(Y2,t+k/2) + 2*f(Y3,t+k/2) + f(Y4,t+k/2))

### MAIN PROGRAM ###
    
# (a) (system) (1 Punkt)
f = lambda x,t: np.array([x[2], x[3], -x[0]/np.sqrt(x[0]**2 + x[1]**2)**3, -x[1]/np.sqrt(x[0]**2 + x[1]**2)**3])
init = lambda eps: np.array([1-eps, 0, 0, np.sqrt((1+eps)/(1-eps))])

# (b) (exact solution) (1 Punkt)
t = np.linspace(0,20,201)
x = np.zeros_like(t)
y = np.zeros_like(t)
eps = .2
for n in range(len(t)):
    time = t[n]
    
    g = lambda u: u - eps*np.sin(u) - time
    dg = lambda u: 1- eps*np.cos(u)
    u = Newton(g,dg,0)[0]
    
    x[n] = np.cos(u) - eps
    y[n] = np.sqrt(1-eps**2)*np.sin(u)
    
plt.plot(x,y)

# (f) (1 Punkt)

for eps in [0.1,0.3]:
    mu = init(eps)
    
    # exact solution
    t = np.linspace(0,20,201)
    x = np.zeros_like(t)
    y = np.zeros_like(t)
    for n in range(len(t)):
        time = t[n]
        # exakt
        g = lambda u: u - eps*np.sin(u) - time
        dg = lambda u: 1- eps*np.cos(u)
        u = Newton(g,dg,0)[0]
        
        x[n] = np.cos(u) - eps
        y[n] = np.sqrt(1-eps**2)*np.sin(u)
    
    for spaces in [100,200,400]:
        t = np.linspace(0,2000,spaces+1)
        k = 20/spaces
        
        # explicit Euler method
        ee = np.zeros((4,spaces+1))
        ee[:,0] = mu

        for n in np.arange(1,len(t)):
            ee[:,n] = Euler_explicit(ee[:,n-1],t[n-1],f,k)
        
        # improved Euler method
        ve = np.zeros((4,spaces+1))
        ve[:,0] = mu
        
        for n in np.arange(1,len(t)):
            ve[:,n] = Euler_improved(ve[:,n-1],t[n-1],f,k)
            
        # classical Runge-Kutta method of fourth order
        rk4 = np.zeros((4,spaces+1))
        rk4[:,0] = mu
        
        for n in np.arange(1,len(t)):
            rk4[:,n] = RungeKutta4(rk4[:,n-1],t[n-1],f,k)
        
        plt.figure()
        plt.title('epsilon = {}, k = {}'.format(eps,k))
        plt.plot(ee[0,:],ee[1,:],label='explicit Euler method')
        plt.plot(ve[0,:],ve[1,:],label='improved Euler method')
        plt.plot(rk4[0,:],rk4[1,:],label='classical Runge-Kutta method')
        plt.plot(x,y,label='exact solution')
        plt.legend()