"""
Exxample from lecture 3 nov. 
Solve a system of four ODEs describing
the trajectory of a ball. 
"""

from ODESolver import ForwardEuler
import numpy as np
import matplotlib.pyplot as plt

def f(t, u):
    x, vx, y, vy = u
    g = 9.81
    return [vx, 0, vy, -g]

# initiell posisjon, start i origo
x = 0; y = 0 

# initiell hastighet og vinkel 
v0 = 5; theta = 80 * np.pi / 180
vx = v0*np.cos(theta); vy = v0*np.sin(theta) 

U0 = [x, vx, y, vy]

solver = ForwardEuler(f)
solver.set_initial_condition(U0)
T = 1.0
N = 100
t, u = solver.solve((0, T), N)

# u er en array av [x, vx, y, vy]-arrays, plott y mot x: 
x = u[:,0]; y = u[:,2]
plt.plot(x, y)
plt.show()