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

h = 10
hlast = h
desired_error = 1e-10
[x1, x2] = [1, 0]
# t = arange(0,10*pi,h)
t = 0
tt = []
xx1 = []
xx2 = []
x1_exacts = []
errors = []
hh = []

while (t < 100.*pi):
	dError = 1
	while (dError > 0):
		# full step:
		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]))
		[x1f, x2f] = [x1, x2] + k1/6. + k2/3. + k3/3. + k4/6.

		# two half steps:
		k1 = h/2.*array(f(x1,x2))
		k2 = h/2.*array(f(x1+k1[0]/2.,x2+k1[1]/2.))
		k3 = h/2.*array(f(x1+k2[0]/2.,x2+k2[1]/2.))
		k4 = h/2.*array(f(x1+k3[0],x2+k3[1]))
		[x1h, x2h] = [x1, x2] + k1/6. + k2/3. + k3/3. + k4/6.

		k1 = h/2.*array(f(x1h,x2h))
		k2 = h/2.*array(f(x1h+k1[0]/2.,x2h+k1[1]/2.))
		k3 = h/2.*array(f(x1h+k2[0]/2.,x2h+k2[1]/2.))
		k4 = h/2.*array(f(x1h+k3[0],x2h+k3[1]))
		[x1h, x2h] = [x1h, x2h] + k1/6. + k2/3. + k3/3. + k4/6.

		# update step size
		error_estimate = x1f-x1h
		dError = (abs(error_estimate)-desired_error)
		
		if (dError > 0):
			h = h*0.76
			hlast = h
			print("Error too large... Redo'ing step")
		elif (dError < 0):
			hlast = h
			h = h*1.02
			print("Expanding step size")

	print("t %.3f \t eE %.3e \t dE %.3e \t h %.3e" %(t,error_estimate,dError,h))
	t += hlast
	tt.append(t)
	[x1, x2] = [x1f, x2f]

	hh.append(hlast)
	xx1.append(x1)
	xx2.append(x2)
	x1_exact = cos(t)
	x1_exacts.append(x1_exact)
	error = abs(x1-x1_exact)
	# print(error)
	errors.append(error)
	

# plt.figure()
# plt.plot(tt,xx1)
# plt.plot(tt,x1_exacts)
plt.figure()
plt.plot(tt,errors)
plt.xlabel('Time')
plt.ylabel('Global Error')
plt.title('Adaptive Stepping RK Error')
plt.figure()
plt.plot(tt,hh)
plt.xlabel('Time')
plt.ylabel('Step Size')
plt.title('Adaptive Stepping RK Step Size')
plt.show()