#!/usr/bin/env python
#import numpy as np
cimport numpy as c_np

def euler_cython(f,z0,z,n_steps,h):
	#cython implementation of euler integration
	cdef int i
	z[0] = z0
	for i in range(n_steps-1):
		z[i+1] = z[i] + h*f(z[i])

def runge_kutta_4_cython(f,z0,z,n_steps,h):
	#cython implementation of runge kutta integration
	cdef c_np.ndarray[double, ndim=1] k1,k2,k3,k4
	cdef int i
	z[0] = z0
	for i in range(n_steps-1):
		k1 = h*f(z[i])
		k2 = h*f(z[i]+.5*k1)
		k3 = h*f(z[i]+.5*k2)
		k4 = h*f(z[i]+k3)
		z[i+1] = z[i] + k1/6.+k2/3.+k3/3.+k4/6.

def adaptive_runge_kutta_4_cython(f,dt,t,z,z0,h0=1.,error_threshold=1e-3,max_steps=1e4):
	#cython implementation of adaptive runge kutta
	#need to keep track of time values for z values we return
	cdef double beta,oneoverbeta #this is the adjustment parameter for the adaptive step choice
	beta = 1.001
	oneoverbeta = 1/beta
	cdef int i
	cdef double h, t_now
	z[0] = z0
	h = h0 					#starting step size
	i=0 						#counter, this will also keep track of number of steps used
	t_now=0. 						#current time value
	def step(h, zi): #should do cdef
		cdef c_np.ndarray[double, ndim=1] k1,k2,k3,k4
		k1 = h*f(zi)
		k2 = h*f(zi+.5*k1)
		k3 = h*f(zi+.5*k2)
		k4 = h*f(zi+k3)
		return zi + k1/6.+k2/3.+k3/3.+k4/6.
	mag = lambda x: x[0]**2 + x[1]**2
	def error(h,zi):
		return mag((step(h,zi)-step(.5*h,step(.5*h,zi))))/(h**5)
	while t_now<dt:
		if i >= max_steps:
			assert(0) #did not allocate enough space
		while error(h,z[i]) < error_threshold: #increase until we bust
			h = beta*h
			#print '\t',h
		while error(h,z[i]) > error_threshold: #then decrease until we just pass
			h = oneoverbeta*h
			#print '\t',h
		#print error(h,z[i]),error_threshold,h,error(h,z[i]) < error_threshold, error(beta*h,z[i]) > error_threshold
		#z[i+1] = step(h,z[i]) #stupid update
		z[i+1] = step(.5*h,step(.5*h,z[i])) + (step(.5*h,step(.5*h,z[i]))-step(h,z[i]))/15. #5-th order update
		
		t[i+1] = t_now
		t_now += h
		i += 1
	return i