#!/usr/bin/env python
#fixed step euler and runge-kutta for harmonic oscillator
from __future__ import division
from numpy import *
from matplotlib import pyplot as plt
import pyximport
pyximport.install(setup_args={"include_dirs":get_include()})
from integration import *

'''
#moved the integration to cython
def euler_python(z0,n_steps,h):
	#python implementation of euler step
	z=array([z0]) #this will accumulate function values over time
	for i in range(n_steps-1):
		z_new = z[-1] + h*f(z[-1])
		z = vstack((z,z_new))
	return z
'''

def f(z):
	#we are simplifying a second order equation in y to a first order equation in z
	#we know that ddy = -y, so let z = (y,dy)
	#then dz = (dy,ddy) = (dy, -y) = (z[1],-z[0])
	return array([z[1],-z[0]])

def compare_fixed_h(ti,tf,h,z0):
	from time import time
	n_steps = int((tf-ti)/h)
	t = linspace(ti,tf,n_steps)
	es = empty((n_steps,2))
	rks = empty((n_steps,2))
	t0 = time()
	euler_cython(f,z0,es,n_steps,h)
	print "Euler took: %.3f s"%(time()-t0)
	t0 = time()
	runge_kutta_4_cython(f,z0,rks,n_steps,h)
	print "Runge Kutta took: %.3f s"%(time()-t0)

	plt.figure()
	plt.plot(t,es[...,0],color='b', label='euler')
	plt.plot(t,rks[...,0],color='g',label='runge kutta 4')
	plt.xlim([ti,tf])
	plt.ylim([-2,2])
	plt.xlabel('t')
	plt.ylable('y')
	plt.title('Euler vs. Runge Kutta for $\ddot{y}+y=0$, h=%.4f'%(h))
	plt.legend(loc='lower left')


def plot_over_h(ti,tf,z0,hs):
	#for this ode, we know the solution is y=cos(t)
	ale_euler = []
	ale_rks = []
	end_euler = []
	end_rks = []

	for h in hs:
		n_steps = int((tf-ti)/h)
		t = linspace(ti,tf,n_steps)
		es = empty((n_steps,2))
		rks = empty((n_steps,2))
		euler_cython(f,z0,es,n_steps,h)
		runge_kutta_4_cython(f,z0,rks,n_steps,h)
		
		ale_euler.append( sum(absolute(es[...,0]-cos(t)))/n_steps )
		ale_rks.append( sum(absolute(rks[...,0]-cos(t)))/n_steps )
		end_euler.append(absolute(es[-1]-array([cos(tf),-sin(tf)])) )
		end_rks.append(absolute(rks[-1]-array([cos(tf),-sin(tf)])) )

	#average local error
	plt.figure()
	plt.loglog(hs,ale_euler,color='b', label='euler')
	plt.loglog(hs,ale_rks,color='g',label='runge kutta 4')
	plt.xlim([0,hs[-1]])
	plt.xlabel('h')
	plt.ylabel('average local error')
	plt.grid(True)
	plt.title('Average local error of Euler vs. Runge Kutta for $\ddot{y}+y=0$')
	plt.legend(loc='upper left')

	#ending value
	plt.figure()
	plt.loglog(hs,[x[0] for x in end_euler],color='b', label='euler')
	plt.loglog(hs,[x[0] for x in end_rks],color='g',label='runge kutta 4')
	plt.xlim([0,hs[-1]])
	plt.xlabel('h')
	plt.ylabel('ending local error')
	plt.grid(True)
	plt.title('Ending error of Euler vs. Runge Kutta for $\ddot{y}+y=0$')
	plt.legend(loc='upper left')

	#ending slope
	plt.figure()
	plt.loglog(hs,[x[1] for x in end_euler],color='b', label='euler')
	plt.loglog(hs,[x[1] for x in end_rks],color='g',label='runge kutta 4')
	plt.xlim([0,hs[-1]])
	plt.xlabel('h')
	plt.ylabel('ending local derivative error')
	plt.grid(True)
	plt.title('Ending derivative error of Euler vs. Runge Kutta for $\ddot{y}+y=0$')
	plt.legend(loc='upper left')

if __name__ == '__main__':
	ti = 0.
	tf = 100*pi
	h = .005
	z0 = array([1,0]) #z = (y,dy)
	
	#compare_fixed_h(ti,tf,h,z0)
	
	hs = logspace(-4, -1, num=12)
	plot_over_h(ti,tf,z0,hs)

	plt.show()