import numpy as np
from numpy import *
from matplotlib import pyplot as plt

# ODE:
# ddy + y = 0
# 
# x1 = y
# x2 = dy
# 
# dx1 = x2 		( = dy)
# dx2 = -x1 	( = ddy)
# 
# f(x1,x2) = [dx1, dx2]
def f(x1,x2):
	return [x2, -x1]

# RK4:
# k1 = h*f(x,y(x))
# k2 = h*f(x+h/2,y(x)+k1/2)
# k3 = h*f(x+h/2,y(x)+k2/2)
# k4 = h*f(x+h,y(x)+k3)
# y(x + h) = y(x) + k1/6 + k2/3 + k3/3 + k4/6

steps = linspace(0.001,0.1,50)
g_errors = []
l_errors = []
# h = 0.1
for h in steps:
	[x1, x2] = [1, 0]
	t = arange(0,10*pi,h)
	xx1 = []
	xx2 = []
	x1_exacts = []
	errors = []
	for i in t:
		k1 = h*array(f(x1,x2))
		k2 = h*array(f(x1+k1[0]/2,x2+k1[1]/2))
		k3 = h*array(f(x1+k2[0]/2,x2+k2[1]/2))
		k4 = h*array(f(x1+k3[0],x2+k3[1]))

		[x1, x2] = [x1, x2] + k1/6 + k2/3 + k3/3 + k4/6
		xx1.append(x1)
		xx2.append(x2)
		x1_exact = cos(i)
		x1_exacts.append(x1_exact)
		error = abs(x1-x1_exact)
		errors.append(error)

	# plt.figure()
	# plt.plot(t,xx1)
	# plt.plot(t,x1_exacts)
	# plt.figure()
	# plt.plot(t,errors)
	# plt.show()

	print("Step size = %.3f" %(h))
	print("Global error = %.2f" % (error))
	print("Average local error = %.2f" % (mean(errors)))
	g_errors.append(error)
	l_errors.append(mean(errors))
	

plt.figure()
gerr, = plt.semilogy(steps,g_errors)
lerr, = plt.semilogy(steps,l_errors)
plt.xlabel('Step Size')
plt.ylabel('Error')
plt.legend([gerr, lerr],['Global Error', 'Local Error'],'best')
plt.title('Runge-Kutta (4th order) Error vs Step-Size')
plt.show()