from scipy import arange
import matplotlib.pyplot as plt
import numpy as np
import random

# These are functions which calculate and plot transient length and cycle length for a discretized map

# This function calculates the transient length and cycle length of a discretized map, for some initial condition using the 'tortoise and the hare' method
def TransientCycle(map,r,N,ic):
	# hare iterates the map twice as fast as tortoise. this loop is run until tortoise=hare.
	# the value of tortoise will then be a multiple of the cycle length
	tortoise = map.iterate(ic,r,N,1,'reset')
	hare = map.iterate(ic,r,N,2,'reset')
	while tortoise != hare:
		tortoise = map.iterate(tortoise,r,N,1,'reset')
		hare = map.iterate(hare,r,N,2,'reset')
 
	# here we find the first state which is repeated with cycle length equal to the tortoise value found above
	# since this is a multiple of the smallest cycle length, this state also marks the start of the smallest cycle
	# the index is therefore the transient length "mu" of the map
	mu = 0
	tortoise, hare = ic, tortoise
	while tortoise != hare:
		tortoise = map.iterate(tortoise,r,N,1,'reset')
		hare = map.iterate(hare,r,N,1,'reset')
		mu += 1
 
	# here the smallest cycle length "lam"is calculated in a straightforward manner, starting at the transient length
	lam = 1
	hare = map.iterate(tortoise,r,N,1,'reset')
	while tortoise != hare:
		hare = map.iterate(hare,r,N,1,'reset')
		lam += 1
	return mu, lam

# this function calculates an average cycle/transient length of a map, by testing many initial conditions,
# in addition, noise is added to r and N with variance rvar and Nvar respectively, to allow for a 'binning' effect.
# this will allow us to smooth a plot of cycle length versus N for example.
def TCMean(map,r,rvar,N,Nvar,samples):
	mu, lam = 0.,0.
	# the average is calculated by iteratively adding up all of the cycle lengths, scaled by the total number of samples
	for i in xrange(samples):
		ic=random.random()
		r=r+ (random.random()-.5)*rvar
		N=N+ (random.random()-.5)*Nvar
		# the change in mu and lam are calculated
		deltamu, deltalam = TransientCycle(map,r,N,ic)
		# the changes are scaled and added to the totals
		mu,lam=mu+deltamu/float(samples) , lam+deltalam/float(samples)
	return mu, lam

# this function plots the transient and cycle lengths of a discretized map, with either a linear or log-log scale
def TCPlot(map,r,N,samples,scale='linear'):
	
	# if r and N are both lists, this is interpreted as attempting to vary both parameters in the bifurcation diagram
	if type(r)==type([]) and type(N)==type([]):
		raise NotImplementedError, '3D plots not yet supported.'
		
	# if r  and N are both numbers, the lengths are simply printed
	elif type(r)==type(1.) and type(N)==(1.):
		print "Average transient length: " + `TCMean(map,r,N,samples)[0]` +"\nAverage cycle length: " + `TCMean(map,r,N,samples)[1]`
	
	# if r is a list and N is a number, we vary r in the plot
	elif type(r)==type([]) and type(N)==type(1.):
		if len(r)<=2: rinc=float(r[1]-r[0])/128.
		else: rinc=r[2]
		# This is the main loop. It constructs the plot iteratively, to be shown upon completion
		rs=arange(r[0],r[1],rinc)
		mus=[] # a list of mu values
		lams=[] # a list of lam values
		
		#here the iterates are calculated
		for rval in rs:
			mu, lam = TCMean(map,rval,rinc,N,0.,samples)
			mus.append(mu)
			lams.append(lam)
		
		plt.subplot(211)
		# if the scale is log-log, modify the lists  and labels accordingly
		if scale == 'log':
			Ns=np.log(Ns)
			mus=np.log(mus)
			lams=np.log(lams)
			plt.ylabel('Log(Transient Length)')
			plt.xlabel('Log(Control Parameter r)')
		else:
			plt.ylabel('Transient Length')
			plt.xlabel('Control Parameter r')
		plt.plot(rs,mus,'b')
		plt.title(map.titlestring + ' N=' + `N`)
		
		plt.subplot(212)
		# if the scale is log-log, modify the lists  and labels accordingly
		if scale == 'log':
			plt.ylabel('Log(Cycle Length)')
			plt.xlabel('Log(Control Parameter r)')
		else:
			plt.ylabel('Cycle Length')
			plt.xlabel('Control Parameter r')

		plt.plot(rs,lams,'b')
	
	# if r is a number and N a list, we vary N in the plot
	elif type(r)==type(1.) and type(N)==type([]):
		if len(N)<=2: Ninc=float(N[1]-N[0])/128.
		else: Ninc=N[2]
		# This is the main loop. It constructs the plot iteratively, to be shown upon completion
		Ns=arange(N[0],N[1],Ninc)
		mus=[] # a list of mu values
		lams=[] # a list of lam values
		for Nval in Ns:
			mu, lam = TCMean(map,r,0.,Nval,Ninc,samples)
			mus.append(mu)
			lams.append(lam)
		
		plt.subplot(211)
		# if the scale is log-log, modify the lists  and labels accordingly
		if scale == 'log':
			Ns=np.log(Ns)
			mus=np.log(mus)
			lams=np.log(lams)
			plt.xlabel('Log(Grid Points N)')
			plt.ylabel('Log(Transient Length)')
		else:
			plt.xlabel('Grid Points N')
			plt.ylabel('Transient Length')
		plt.plot(Ns,mus,'b')
		plt.title(map.titlestring + ' r=' + `r`)
		
		plt.subplot(212)
		# if the scale is log-log, modify the lists  and labels accordingly
		if scale == 'log':
			plt.ylabel('Log(Cycle Length)')
			plt.xlabel('Log(Grid Points N)')
		else:
			plt.ylabel('Cycle Length')
			plt.xlabel('Grid Points N')
		#print lams[0:100]
		plt.plot(Ns,lams,'b')