from numpy import arange, abs, pi, cos

def RK1D(a,b,gamma,theta,t,f,dt):		# should be easily modified to higher D, just by making theta and fhigher D, and making methods consistent for the k's
	
	k1 = dt * f(a,b,gamma, theta, t)
	t +=.5*dt
	k2 = dt * f(a,b,gamma, theta + k1 / 2.0, t)
	k3 = dt * f(a,b,gamma, theta + k2 / 2.0, t)
	t +=.5*dt
	k4 = dt * f(a,b,gamma, theta + k3, t)

	theta = theta + ( k1 + 2.0 * k2 + 2.0 * k3 + k4 ) / 6.0

	return theta 

def ForgetfulRK1DIntegrator(a,b,gamma,IC,ThetaMax,f,dt,timesteps,t_0):	
	"""IC was a list, but now just being used as a single input 
	r is a list of parameters
	ThetaMax is the normalized V_T... ThetaMax should be = pi if using Yang method
	f is the differential eqn describing the 1D ODE
	
	ForgetfulRK1DIntegrator only remembers the most recent theta, and the list of theta_n, and returns theta, theta_n, tau_i after the specified # of timesteps
	"""
	#theta = IC[0]
	theta = IC	# just do this until higher Dim
	#TrajTheta = [theta,]
	#Time  = [t_0,]
	t = t_0
	tau_i = [ ] 	# = [tau_1, tau_2, tau_3... ]
	theta_n = [ ] # leave out theta_0	#[theta,]	# will hold [theta_0, theta_1, theta_2, ..., theta_n = theta(t = pi2n), ... ]
	# for simple return map, we only need [theta_0, theta1]
	# for bifurcation map as fn of IC, keep ~ 1000 iterates?... this could take a while
	#pump = [b*cos(Time[0]),]	# will hold pump voltages
	
	p = 1.0
	for step in arange(1,timesteps+1):
		t = t_0 + dt * step
		#Time.append(t)
		theta = RK1D(a,b,gamma,theta,t,f,dt)
		#TrajTheta.append(theta)
		if abs(theta) > ThetaMax:
			if theta > ThetaMax:
				theta -= 2*pi 	
			if theta < -ThetaMax:	
				theta += 2*pi	
			tau_i.append(dt * step)		# each entry holds the time of a tunneling event
		# Turn this on/off to return pump voltage (would have to change return and accept values)
		if t > t_0 + 2*pi*p:	# theta_(n+p) = theta(t=2(n+p)pi) + t_0????
			p += 1.0
			theta_n.append(theta)
			#thetaN.append(theta)
		#pump.append(b*cos(t))
	return theta,theta_n,t,tau_i