#!/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 *

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 fixed_error(ti,tf,z0):
	from time import time
	dt = tf-ti
	t0 = time()
	max_steps = 1e4
	z = empty((max_steps,2)) #allocate
	t = empty(max_steps) #allocate
	n_steps = adaptive_runge_kutta_4_cython(f,dt,t,z,z0,h0=.01,error_threshold=1e-9,max_steps=max_steps)
	t = t[:n_steps+1]
	z = z[:n_steps+1]
	print "adaptive Runge Kutta took: %.3f s"%(time()-t0)

	plt.figure()
	plt.plot(t+ti,z[...,0],color='g',label='runge kutta 4',marker='.')
	plt.xlim([ti,tf])
	plt.ylim([-2,2])
	plt.xlabel('t')
	plt.ylabel('y')
	plt.title('Adaptive Runge Kutta for $\ddot{y}+y=0$')
	plt.legend(loc='lower left')

	plt.figure()
	plt.plot(t[1:]+ti, t[1:]-t[:-1], color='g',label='runge kutta 4')
	plt.xlim([ti,tf])
	plt.ylim([0,1])
	plt.xlabel('t')
	plt.ylabel('step size')
	plt.title('Adaptive Runge Kutta step size for $\ddot{y}+y=0$')
	plt.legend(loc='lower left')


def plot_over_desired_error(ti,tf,z0,errors):
	from time import time
	avg_steps = []
	dt = tf-ti
	max_steps = 1e4
	for err in errors:
		t0 = time()
		z = empty((max_steps,2)) #allocate
		t = empty(max_steps) #allocate
		n_steps = adaptive_runge_kutta_4_cython(f,dt,t,z,z0,h0=.01,error_threshold=err,max_steps=max_steps)
		t = t[:n_steps+1]
		z = z[:n_steps+1]
		print "adaptive Runge Kutta took: %.3f s"%(time()-t0)
		avg_steps.append( average(t[1:]-t[:-1]) )
	print avg_steps
	plt.figure()
	ax = plt.gca()
	ax.set_xscale('log')
	ax.set_yscale('log')
	plt.plot(errors,avg_steps,color='b',marker='.')
	plt.xlim([0,errors[-1]])
	plt.xlabel('desired error')
	plt.ylabel('average step size')
	plt.grid(True)
	plt.title('Average step size of adaptive Runge Kutta for $\ddot{y}+y=0$')



if __name__ == '__main__':
	ti = 0.
	tf = 100*pi
	z0 = array([1,0]) #z = (y,dy)
	fixed_error(ti,tf,z0)
	#errors = logspace(-11, -3, num=10)
	#plot_over_desired_error(ti,tf,z0,errors)
	plt.show()