# -*- coding: utf-8 -*-
"""
Spyder Editor

Exercice 2.17

1/ Test that sqrt(2)*sqrt(2)-2 is not zero


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

T=100*2*np.pi
N=10_000
dt=T/N

def f(z):
    """
    input z of shape (2,)
    output f(z) of shape (2,)
    """
    return np.array([-z[1],z[0]])

solEE=np.zeros((N+1,2))
#start from (1,0) as initial value
solEE[0,:]=np.array([1.,0.])

trange=np.linspace(0,T,N+1,endpoint=True)

for ii in range(N):    
    solEE[ii+1,:]=solEE[ii,:]+ dt*f(solEE[ii,:]) 
    
plt.figure('x(t)')
plt.plot(trange,solEE[:,0])
plt.plot(trange,solEE[:,1])
plt.legend(['Re(z)','Im(z)'])


def sol_IE_f(z):
    """
    input z of shape (2,)
    output f(z) = (....) of shape (2,)
    """
    x,y= z
    return np.array([ x-dt*y, y+dt*x])/(1+dt**2)

solIE=np.zeros((N+1,2))
#start from (1,0) as initial value
solIE[0,:]=np.array([1.,0.])
for ii in range(N):    
    solIE[ii+1,:]=sol_IE_f(solIE[ii,:]) 
    
plt.figure('IE x(t)')
plt.subplot(1,2,1)
plt.plot(trange,solIE[:,0])
plt.legend(['Re(z)'])
plt.subplot(1,2,2)
plt.plot(trange,solIE[:,1])
plt.legend(['Im(z)'])