# -*- coding: utf-8 -*-
"""
Created on Mon Mar  8 12:12:15 2021

@author: turinici : TP no. 1 :ex 2.17: stability of Explicit Euler

for z'(t) = L * z(t)  with L = i  = sqrt(-1)
With notation z = x+i*y : 
x,y=z
x' = Re(z') = Re( i * z) = Re(i*(x+iy))= - y     
y' = Im(z') = Im(i*z) = Im(i*(x+iy)) = x
ODE equation is : [x,y]' = [-y,x]

"""

import numpy as np
import matplotlib.pyplot as plt
%matplotlib inline
#%matplotlib auto

#parameters
T=100.*2*np.pi # final time = 100*2*pi
N=5000 # number of time steps
h = T/N # 
U0 = [1.,0.]

def ftxt(t,z):
    """
    define the function entering the ODE x'(t) = f(t,x(t))
    Parameters
    ----------
    t : time
    z : list of components of dimension d

    Returns
    -------
    a list, value of the function 

    """
    x,y=z
    return [-y,x]

def Explicit_Euler(h,N,ode_function,initial_value):
    """
    Parameters
    ----------
    h : time step
    N : number of time steps (integer)
    input_function : function to integrate
    initial_value : initial value of the ODE

    Returns
    -------
    List of approximate solution obtained by Explicit Euler scheme of step h.
    """

    U= [initial_value]*(N+1) # creates a list on N+1 elements, filled with value 0

    for ii in range(N):
        U[ii+1]=[uiik+h*xk for xk,uiik in zip( ode_function(ii*h,U[ii]), U[ii])]
#        U[ii+1]=list(np.array(U[ii])+h*np.array(ode_function(ii*h,U[ii])))
    
    return(U)
    
# define the time grid    
trange = np.linspace(0,T,num=N+1,endpoint=True)   
# solve by EE the ODE : use EE at any time step and put in a numpy array : solution
solution = Explicit_Euler(h,N,ftxt,U0)

solutionx = [z[0] for z in solution]
solutiony = [z[1] for z in solution]


# plot the results
plt.figure(1)
plt.subplot(2,1,1)
plt.plot(trange,solutionx,'-r')
plt.xlabel('time (t)')
plt.ylabel('z')
plt.legend(['Real(z(t))'])
plt.subplot(2,1,2)
#plt.plot(solution)
plt.plot(trange,solutiony,'-b')
plt.xlabel('time (t)')
plt.ylabel('z')
plt.legend(['Im(z(t))'])