from numpy import *
from pylab import *
from ForgetfulIntegrator1D import *
from OneDimMap import *
import random
import time

time.clock()	# let's see how long this takes
a = input('Please enter a value for a:  ')
print '\nPlease wait while I compute the Poincare map for this a-value...\n'
dt = .1	# more precise?
nTransients = 200.0#50.0
#nTransients = 0.0
nIterates = 50.0 #80.0
transienttimesteps = nTransients*2*pi/dt	# does this mess up
timesteps = nIterates*2*pi/dt # + nIterates	# add nIterates for now in case we lose a few timesteps...only length(theta_n) will count	# is int type okay? necessary?
t_0 = 0.0

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

class SETJmap_theta(OneDimMap):

	def __init__(self, ind_var = [-pi,pi], param = [a,2.0,1./3.,-pi]):	# ind_var: a, use random IC for theta_0
		#(self,theta_0): #,neighbors):
		#param = [1.77,2.0,1./3.]	# params = [a,b,gamma], could be ranges? maybe for specific applications?
		OneDimMap.__init__(self, ind_var = [-pi,pi], param = [a,2.0,1./3.,-pi])
		self.nTransients = nTransients 
		self.nIterates = nIterates
		#self.nIterates = 1.0
		self.independent_var = 'theta_0'
		self.ylabel = 'theta(t = 2 pi n), n = 1,2,3,... '
		self.name = 'Poincare map for a = '+`a`
		#self.neighbors = neighbors	# ...

	def iterate2(self,r):	# map from OneDMaps
		self.t = 0.0
		self.current,transienttheta,self.t,transient_tau = self.maptransients(r)
		self.theta_n = [ ]	# see if this helps
		#self.current,self.theta_n,self.t,tau_i_map = self.map(r)	#ignores transients!
		self.current,thetaN,self.t,tau_i_map = self.map(r)
		#for i in xrange(self.nTransients):	## fix this!***!***!***!
		#	self.current = self.map(r)		## Don't pass self!!	
		rsweep = [r,]*int(self.nIterates)
		#rsweep = [ ]   # The parameter value
		#theta = [ ]        # The iterates
		#theta = 
		#for i in xrange(self.nIterates):
		#	self.current = self.map(r)		## Don't pass self!!
		#	rsweep.append(r)			# --> [r,r,r,r,... ]
		#	theta.append( self.current )
		return thetaN, rsweep


	def maptransients(self,r):
		self.current = r #random.uniform(-1.,1.)*pi	# this will carry over to the non-transient map, should draw out any variability across theta_0
		theta_last,theta_n_transient,t_last,tau_i_map = ForgetfulRK1DIntegrator(a,self.b,self.gamma,self.current,self.ThetaMax,d_theta,dt,transienttimesteps,self.t) # only differs by transienttimesteps
		# see that r replaces 'a'
		
		return theta_last,theta_n_transient,t_last,tau_i_map


	def map(self,r):
		theta_last,theta_n,t_last,tau_i_map = ForgetfulRK1DIntegrator(a,self.b,self.gamma,r,self.ThetaMax,d_theta,dt,timesteps,self.t)
		
		return theta_last,theta_n,t_last,tau_i_map


MyMap = SETJmap_theta()


MyMap.nSteps = 1000
figure(1,(12,9))	# Setup the plot; It's numbered 1 and is 8 x 6 inches.
# Stuff parameter range into a string via the string formating commands.
TitleString  = MyMap.name
#TitleString += ' from [%g,%g]' % (self.ind[0],self.ind[1])
title(TitleString)
# Label axes
xlabel(MyMap.independent_var)
ylabel(MyMap.ylabel)	#{Theta(n)}')

rInc = (MyMap.ind[1]-MyMap.ind[0])/float(MyMap.nSteps)
plot([-pi,pi],[-pi,pi],'r-')
#color_run = ['ko','bo','go','mo','ro']
for r in arange(MyMap.ind[0],MyMap.ind[1],rInc):
	# Set the initial condition to the reference value
	MyMap.current = r	# since theta = theta_0 initially	#self.IC
	# Throw away the transient iterations
	x, rsweep = MyMap.iterate2(r)				## Don't pass self!!
	plot(rsweep, x, 'bo',markersize=.5) # Plot the list of (r,x) pairs as pixels, ms = .5
#savefig(self.name[0:8]+'.ps')
# Display plot in window
#axis([self.param[0],self.param[1],-10.,0.])
show()


#for initial_theta in arange(0,2*pi,.01):
#	for p in xrange(6)
#		MyMap.map(a)
#MyMap.show()
time.clock()