from numpy import arange, abs, pi, cos, array

def RK1D(a,b,gamma,theta,t,f,dt):		# should be easily modified to higher D, just by making theta and f higher 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)# should change t to t+.5*dt
	k3 = dt * f(a,b,gamma, theta + k2 / 2.0, t)# should change t to t +.5*dt
	t+=.5*dt
	k4 = dt * f(a,b,gamma, theta + k3, t)# should change t to t + dt ******

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

	return theta 


def RK1D2(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
	#theta = theta[0]
	#print theta
	#theta = array(theta)
	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 = array(theta)
	theta = theta + ( k1 + 2.0 * k2 + 2.0 * k3 + k4 ) / 6.0

	return theta.tolist() #??******************[0] 
	

def RK1DIntegrator(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
	"""
	#theta = IC[0]
	theta = IC	# just do this until higher Dim
	TrajTheta = [theta,]
	Time  = [t_0,]
	tau_i = [ ] 	# = [tau_1, tau_2, tau_3... ]
	#theta_n = [theta,]	# will hold [theta_0, theta_1, theta_2, ..., theta_n = theta(t = pi2n), ... ]
	pump = [b*cos(Time[0]),]	# will hold pump voltages
	
	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:	# Yang mentions this condition early on but does not use it for eq. 6 and on... is it necessary? appropriate? --perhaps exp(-c*theta/pi) always beats out theta*c(a-b*cos(t)) for neg. theta?
				theta += 2*pi	# probably this condition should never be needed
			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)
		pump.append(b*cos(t))
	return TrajTheta,pump,Time,tau_i



def RK1DIntegrator2(a,b,gamma,IC,ThetaMax,f,dt,timesteps,t_0):	# use for 2-coupled SETJ
	"""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
	"""
	#theta = IC[0]
	theta = IC	# just do this until higher Dim
	TrajTheta = [theta,] # [[theta1_0,theta2_0,theta_in_0], ]
	Time  = [t_0,]
	tau_i = [ ] 	# = [tau_1, tau_2, tau_3... ]
	#theta_n = [theta,]	# will hold [theta_0, theta_1, theta_2, ..., theta_n = theta(t = pi2n), ... ]
	pump = [b*cos(Time[0]),]	# will hold pump voltages
	
	for step in arange(1,timesteps+1):
		t = t_0 + dt * step
		Time.append(t)
		theta = RK1D2(a,b,gamma,theta,t,f,dt)
		TrajTheta.append(theta)
		if abs(theta[0]) > ThetaMax:
			if theta[0] > ThetaMax:
				theta[0] -= 2*pi 	
			if theta[0] < -ThetaMax:	# Yang mentions this condition early on but does not use it for eq. 6 and on... is it necessary? appropriate? --perhaps exp(-c*theta/pi) always beats out theta*c(a-b*cos(t)) for neg. theta?
				theta[0] += 2*pi	# probably this condition should never be needed
			tau_i.append(dt * step)		# each entry holds the time of a tunneling event
		if abs(theta[1]) > ThetaMax:
			if theta[1] > ThetaMax:
				theta[1] -= 2*pi 	
			if theta[1] < -ThetaMax:	# Yang mentions this condition early on but does not use it for eq. 6 and on... is it necessary? appropriate? --perhaps exp(-c*theta/pi) always beats out theta*c(a-b*cos(t)) for neg. theta?
				theta[1] += 2*pi	# probably this condition should never be needed
			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)
		pump.append(b*cos(t))
	#return [TrajTheta[:][0],TrajTheta[:][1],TrajTheta[:][2]],pump,Time,tau_i
	return TrajTheta,pump,Time,tau_i