from numpy import * 
from Integrator1D import *
from ForgetfulIntegrator1D import *
# Plotting
from pylab import *
#from enthought.mayavi.mlab import axes,plot3d,title

# 2 Coupled SETJs

# define parameters
# using Yang's normalized variables
ThetaMax = pi	# tunneling occurs for |v_C| > V_T --> |theta| > pi
gamma = 1./3. # gamma = pi / (RC omega_p)
a = [1.77,1.6] #[1.77,-1.77] #1.49 #1.95 #1.77#1.49	# a = V_b / V_T
b = 2.0		# b = V_p / V_T
epsilon = .3 # what's a reasonable value?? # C_in / C
k = .1#sqrt(10.) # what's a reasonable value??# R_in / R
#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 = .01	# IMPORTANT to have this small enough! see what works
#dt = .01
#Ntransients = 50000 # number of transient timesteps to throw away
Ntransients = 5000 #50000 	# interesting to compare transient behavior, to see how v_C becomes entrained
timesteps = 50000	# is int type okay?
t_0 = 0.0
# Define IC
theta_0 = [-1.8,.3,.1]	# must have |theta_0| < pi, this is now [theta1_0, theta2_0, theta_in_0]

# 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):
	#print theta
	dTheta1 = gamma*(a[0] - b*cos(t) + (-(k+1)*theta[0] + theta[1] + theta[2])/(k*pi))	# make sure this isn't returning a tuple!
	dTheta2 = gamma*(a[1] - b*cos(t) + (-(k+1)*theta[1] + theta[0] - theta[2])/(k*pi))
	dTheta_in = gamma*(theta[0] - theta[1] - theta[2])/(k*epsilon*pi)
	return array([dTheta1,dTheta2,dTheta_in])

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

# To do:
#	get rid of transients, look at OneDimMap ex.
#	show the two limit cycles
#	include theta_n = theta(t = 2 n pi) in calculations (started in Integrator1D)
#		plot return map for theta_n+p vs. theta_n = theta(t = 2 n pi); include identity for stable points.
#		plot bifurcation diag. for theta_n vs. theta_0, many iterates after transients die out
#			--> will show where to find periodic solutions (P1, P2, chaotic, etc.)
#	plot bifurcation diag. for a, b, gamma... any better way to show it all?
#		-- !!! How about an interactive phase space set in 2/3 dimensions (most important parameters) and user can press ctrl-a + move cursor, ctrl-b + move cursor to move through those extra dimensions.  Or use sliders with a GUI.  Also use color where possible.
#		-- could 'Find' the parameters you want this way by traveling through the N-dimensional space... would be really cool in KeckCaves, could walk through 5 or more dimensions!--Pretty cool just to better understand higher dim.
#		-- could also find these parameters computationally by testing 'surrounding' areas and testing them for desired characteristics, then moving in the most promising direction... somewhat of a combination of Monte Carlo and finding a generalized minimization of action / potential
#	basins of attraction portrait (color regions differently depending on which basin they fall in: a vs. b? vs. gamma?

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

#[TrajTheta, TrajTheta2,TrajTheta_in], pump, Time, tau_i = RK1DIntegrator2(a,b,gamma,theta_start,ThetaMax,d_theta,dt,timesteps,t_start[-1])		# 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
TrajTheta1 = [ ]
TrajTheta2 = [ ] 
for line in TrajTheta:
	TrajTheta1.append(line[0])
	TrajTheta2.append(line[1])

figure(1)
#figure(2)
clf()
#show()
#ion()
TitleString = 'time series for two-coupled SETJ voltages, using fourth-order Runge-Kutta integration and impulsive difference eqns\na = %s, b = %g, gamma = %g, epsilon = %g, k = %g' % (`a`,b,gamma,epsilon,k)
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','junction1 voltage for theta1_0 = %g' % theta_0[0],'junction2 voltage for theta2_0 = %g' % theta_0[1]] 
plot(Time, TrajTheta1, '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 = 'two-coupled SETJ\na = %s, b = %g, gamma = %g, epsilon = %g, k = %g' % (`a`,b,gamma,epsilon,k)
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(TrajTheta1, 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()

figure(3)
clf()
#show()
#ion()
TitleString = 'phase portrait for two-coupled SETJ\na = %s, b = %g, gamma = %g,  epsilon = %g, k = %g' % (`a`,b,gamma,epsilon,k)
title(TitleString)
xlabel('junction1 voltage')   # set x-axis label
ylabel('junction2 voltage') # set y-axis label
labelname3 = ['r_0 = theta1_0 - theta2_0 = %g' % (theta_0[0] - theta_0[1]),]
plot(TrajTheta1, TrajTheta2, 'b-', antialiased=True,label=labelname3[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=labelname3[1]) # show limit cycle for entrainment to pump (limit cycle would actually be after the transients have died out though)
legend(loc='best')
show()

from enthought.mayavi.mlab import axes,plot3d,title
figure(4)
clf()
plot3d(TrajTheta1,TrajTheta2,pump,Time,color=(0.,0.,1.),tube_radius=0.03)#1)
title('%s attractor with a = %s, b = %g, gamma = %g,  epsilon = %g, k = %g' % ('2-Coupled SETJ', `a`,b,gamma,epsilon,k))
axes(color=(0.,0.,0.),extent=[-5.0,5.0,-5.0,5.0,0.0,50.0],ranges=[-5.0,5.0,-5.0,5.0,0.0,50.0])
# for presentation could compare b = 0.0 w/ b = 2.0; show how b = 2. has 2 different basins of attraction