from numpy import * 
from Integrator1D import *
from ForgetfulIntegrator1D import *
# Plotting
from pylab import *


# define parameters
# using Yang et al.'s normalized variables
ThetaMax = pi	# tunneling occurs for |v_C| > V_T --> |theta| > pi
gamma = 1./3. # gamma = pi / (RC omega_p)
a = 1.49 #1.49 #1.95 #1.77#1.49	# a = V_b / V_T
b = 2.0		# b = V_p / V_T
#b = 0.		# b = 0 turns off ac pump --> simple relaxation oscillations
#Either make t_0 a parameter for the integrator or change it there
#t_0 = 0.0	# in normalized time, i.e. t --> omega_p * t
#Same with t_f
#t_f = 2.0	# in normalized time, i.e. t --> omega_p * t
#dt = .001	# IMPORTANT to have this small enough! see what works
dt = .01
#Ntransients = 50000 # number of transient timesteps to throw away
Ntransients = 15000 #50000 	# interesting to compare transient behavior, to see how v_C becomes entrained
timesteps = 30000	# is int type okay?
t_0 = 0.0
# Define IC
theta_0 = [0.,3.]#[-1.8,.3]	# must have |theta_0| < pi

# time vector created in integrator
#t = arange(t_0,(t_f+dt),dt)	# in normalized time, i.e. t --> omega_p * t;	need to add dt to reach tf

# define 1D ODE for normalized isolated SETJ 
def d_theta(a,b,gamma, theta,t):
	dTheta = -gamma/pi*theta + gamma*(a - b*cos(t))	# make sure this isn't returning a tuple!
	return dTheta

# pump already defined in integrator
#def pump_normalized(b,t):	# not ODE, just plain function
#	return b*cos(t)


#transientTheta, transientpump, transientTime, transienttau_i = RK1DIntegrator(a,b,gamma,theta_0[0],ThetaMax,d_theta,dt,Ntransients,t_0)	# evolve transient behavior
theta_start,theta_n,t_start,tau_i_transient = ForgetfulRK1DIntegrator(a,b,gamma,theta_0[0],ThetaMax,d_theta,dt,Ntransients,t_0)
TrajTheta, pump, Time, tau_i = RK1DIntegrator(a,b,gamma,theta_start,ThetaMax,d_theta,dt,timesteps,t_start)		# record steady state behavior

# for second IC... make this automatic for number_IC's in theta_0
#transientTheta, transientpump, transientTime, transienttau_i = RK1DIntegrator(a,b,gamma,theta_0[1],ThetaMax,d_theta,dt,Ntransients,t_0)	# evolve transient behavior
theta_start,theta_n,t_start,tau_i_transient = ForgetfulRK1DIntegrator(a,b,gamma,theta_0[1],ThetaMax,d_theta,dt,Ntransients,t_0)
TrajTheta2, pump, Time, tau_i2 = RK1DIntegrator(a,b,gamma,theta_start,ThetaMax,d_theta,dt,timesteps,t_start)		# record steady state behavior


# notice that (tau_i[n] - tau_i[n-1])/pi --> 4 in steady state for isolated SETJ

figure(1)
#figure(2)
clf()
#show()
#ion()
TitleString = 'time series for isolated SETJ voltages, using fourth-order Runge-Kutta integration and impulsive difference eqns\na = %g, b = %g, gamma = %g' % (a,b,gamma)
title(TitleString)
xlabel('tau = omega_p t')   # set x-axis label
ylabel('theta = 2 pi C v_C / q') # set y-axis label
labelname = ['pump voltage','junction voltage for theta_0 = %g' % theta_0[0],'junction voltage for theta_0 = %g' % theta_0[1]] 
plot(Time, TrajTheta, 'b-', antialiased=True,label=labelname[1])
plot(Time, TrajTheta2, 'r-', antialiased=True,label=labelname[2])
plot(Time, pump, 'g--', antialiased=True,label=labelname[0])
legend(loc='best')
show()

figure(2)
clf()
#show()
#ion()
TitleString = 'limit cycle for isolated SETJ\na = %g, b = %g, gamma = %g' % (a,b,gamma)
#TitleString = 'phase portrait for isolated SETJ\na = %g, b = %g, gamma = %g' % (a,b,gamma)
title(TitleString)
xlabel('junction voltage')   # set x-axis label
ylabel('pump voltage') # set y-axis label
labelname2 = ['theta_0 = %g' % theta_0[0],'theta_0 = %g' % theta_0[1]]
plot(TrajTheta, pump, 'b-', antialiased=True,label=labelname2[0]) # show limit cycle for entrainment to pump (limit cycle would actually be after the transients have died out though)
plot(TrajTheta2, pump, 'r-', antialiased=True,label=labelname2[1]) # show limit cycle for entrainment to pump (limit cycle would actually be after the transients have died out though)
legend(loc='best')
show()