# BakersMap.py: Plot out iterates of the Dissipative Baker's Map
#
# The Dissipative Baker's Map is given by
#     x_n+1 = f( x_n , y_n )
#     y_n+1 = g( x_n , y_n )
# with
#     f( x_n , y_n ) = 2 * x_n (mod 1)
#     g( x_n , y_n ) = a * y_n,       if x_n <= 1/2
#                    = 1/2 + a * y_n, if x_n > 1/2
#
# The state space is the unit square:
#       (x,y) in [0,1] x [0,1] (unless defined by user)
# and the control parameter ranges in
#       k >= 0
# 

# Import plotting routines
from pylab import *

import numpy as np

# the 2 different (but equivalent) maps

def StandardMap(k,x,y):
	# We'll assume, for now, that (x,y) is in [0,1] x [0,1]
	# Calculate next y first, since it depends on current x
	# a '-' in the y equation shifts and is used
	# for the manifold plot for better visual
	y=(y+k/(2*pi)*sin(2*pi*x))
	x=(x+y)
	return  x%(1),y%(1)
	
def StandardMapShift(k,x,y):
	# We'll assume, for now, that (x,y) is in [0,1] x [0,1]
	# Calculate next y first, since it depends on current x
	# a '-' in the y equation shifts and is used
	# for the manifold plot for better visual
	y=(y-k/(2*pi)*sin(2*pi*x))
	x=(x+y)
	return  x%(1),y%(1)
	

def StandardMapShift_ynotmod1(k,x,y):
	# We'll assume, for now, that (x,y) is in [0,1] x [0,1]
	# Calculate next y first, since it depends on current x
	# a '-' in the y equation shifts and is used
	# for the manifold plot for better visual
	y=(y-k/(2*pi)*sin(2*pi*x))
	x=(x+y)
	return  x%(1),y



def RandInitialIterate(k, N, Map, n_ic=20, box = [0.0,1.0,0.0,1.0]):
	# finds random initial points in the box and iterates them
	# assumes user knows what they are doing
	
	num = raw_input('How many initial conditions? (leave blank for default)\n')
	if num != '': # if number is entered set it to n_ic
		n_ic = int(num)
	
	size = raw_input('Enter size of box as [xmin, xmax, ymin, ymax]:\nleave blank for default unit square\n ')
	if size != '': # if size is entered set to box
		box = eval(size)
	
	
	xlen = box[1]-box[0]
	ylen = box[3]-box[2]
	inits=[]
	for i in xrange(n_ic):
		xtemp = (random()*xlen+box[0])
		ytemp = (random()*ylen+box[2])
		inits.append([xtemp,ytemp])

	for nums in inits:
		# Set up arrays of iterates (x_n,y_n) and set the initital condition
		x=[nums[0]]
		y=[nums[1]]

		# The main loop that generates iterates and stores them
		for n in xrange(0,N):
		  # at each iteration calculate (x_n+1,y_n+1)
		  xtemp, ytemp = Map(k,x[n],y[n])
		  # and append to lists x and y
		  x.append( xtemp )
		  y.append( ytemp )
		plot(x,y,',')
		  
	# Setup the plot
	xlabel('theta(n)') # set horizontal axis label
	ylabel('p(n)') # set vertical axis label
	title("Standard Map with k =" + str(k)) # set plot title
	# Plot the time series
		
	# Make sure the iterates appear in a unit square
	#axis('equal')
	axis(box)

	# Display the plot in a window
	show()
	
	
def LineIterate(k, N, Map, n_ic=5000, endpts = ((0.3,0.2),(0.3,0.8))):
	# iterates a line segment
	
	num = raw_input('How many initial conditions? (leave blank for default)\n')
	if num != '': # if number is entered set it to n_ic
		n_ic = int(num)
	
	# assumes  user follows directions precisely
	print 'Enter end points of a line segment [xmin,xmax, ymin, ymax]:'
	segmt = raw_input('leave blank for default\n ')
		
	if segmt != '': # if size is entered set to box
		endpts = eval(segmt)
	
	# create line segment
	xmin = endpts[0][0]
	ymin = endpts[0][1]
	xmax = endpts[1][0]
	ymax = endpts[1][1]
	
	xinc = (xmax - xmin) / n_ic 
	yinc = (ymax - ymin) / n_ic
	# An array of initial conditions on a straight line
	x = []
	y = []
	
	for i in xrange(0,n_ic):
		x.append(xmin + xinc * i)
		y.append(ymin + yinc * i)
	# plot first line segment
	plot(x,y,',',label='0')
	
	# iterate the line segment
	for n in xrange(N):
	
		Tx=[]
		Ty=[]
		# now for the mapping
		for j in xrange(len(x)):
			xtemp,ytemp=Map(k,x[j],y[j])
			Tx.append(xtemp)
			Ty.append(ytemp)
		# plot iterated line segment (different color)	
		x = Tx
		y = Ty
		
		plot(x,y, ',', label=str(n+1))
	
	# Setup the plot
	xlabel('theta(n)') # set x-axis label
	ylabel('p(n)') # set y-axis label
	title("Standard Map with k =" + str(k)) # set plot title
	legend()

	# Display the plot in a window
	show()
	
def IterArray(k, N, Map, initialpts=np.array([(0,1./4.),(0,1./2.),(0,3./4.),(0,1./5.),(0,3./20.)])):
	# allows user to designate initial conditions for the standard map
	# i guess it doesn't require an array, a list should do just fine
	
	its = raw_input('Enter an assortment of initial points or just the name of the object. (leave blank for default):\n')
	if its != '': # if number is entered set it to n_ic
		initialpts = eval(its)
	
	
	for init in initialpts:
		# Set up arrays of iterates (x_n,y_n) and set the initital condition
		x=[init[0]]
		y=[init[1]]
		
		'''# throw away transients
		xtemp, ytemp = StandardMap(k,x[0],y[0])
		
		# uncomment to throw away transients
		for n in xrange(10):
		  # at each iteration calculate (x_n+1,y_n+1)
		  xtemp, ytemp = StandardMap(k,xtemp,ytemp)		
		
		x[0],y[0] = StandardMap(k,xtemp,ytemp)
		'''
		# The main loop that generates iterates and stores them
		for n in xrange(0,N):
		  # at each iteration calculate (x_n+1,y_n+1)
		  xtemp, ytemp = Map(k,x[n],y[n])
		  # and append to lists x and y
		  x.append( xtemp )
		  y.append( ytemp )
		  #plot each trjectory, top one includes label 
		  #plot(x,y,',',label = 'P_0: ' + str(y[0]))
		  plot(x,y,',')
		  
	# Setup the plot
	xlabel('theta(n)') # set horizontal axis label
	ylabel('p(n)') # set vertical axis label
	title("Standard Map with k =" + str(k)) # set plot title

	# Then run the standard map iterate to fill in some quasiperiodic orbits
	# this function will also plot
	#StandardMapIterate()
	show()


#StandardMapIteration()
	
def Manifold(k,N,Map, n_ic=10000, fxpt = (0.5,0)):
	# computes and plots manifold
	# currently only works for the (0.5,0) fixed point of the shifted Map
	# because the eigvalues are easier to compute, but it can be easily modified
	#
	# the code is primarily the same as LineIterate() except
	# that the initial line segment is very small and computed in the direction of 
	# eigen-direction
	
	# currently set up only for the shifted map but can be easily generalized
	Map = StandardMapShift_ynotmod1
	
	delta = 1e-4
	# initial fixed point
	x = fxpt[0]
	y = fxpt[1]
	# slope is along the unstable eigenvector  
	#eigvec = (1,-.5*k*cos(2*pi*x0)+.5*(k**2*cos(2*pi*x0)**2+4*k*cos(2*pi*x0))**.5)
	eigval = .5*(2+k-(k*(k+4))**.5)
	eigvec = (1, 1-eigval)
	# small displacement from fixed point along stable manifold
	p1 = x - delta*eigvec[0]
	q1 = y - delta*eigvec[1]
	# include both directions of manifold
	p2 = x + delta*eigvec[0]
	q2 = y + delta*eigvec[1]
	
	# 
	endpts = ((p1,q1),(p2,q2))
	
	# create line segment
	xmin = endpts[0][0]
	ymin = endpts[0][1]
	xmax = endpts[1][0]
	ymax = endpts[1][1]
	
	xinc = (xmax - xmin) / n_ic 
	yinc = (ymax - ymin) / n_ic
	# An array of initial conditions on a straight line
	x = []
	y = []
	
	for i in xrange(0,n_ic):
		x.append(xmin + xinc * i)
		y.append(ymin + yinc * i)
	# plot first line segment
	plot(x,y,',',label='0')
	
	# iterate the line segment (I set to 15 iterations because it looked nice)
	for n in xrange(15):
		Tx=[]
		Ty=[]
		# now for the mapping
		for j in xrange(len(x)):
			xtemp,ytemp=Map(k,x[j],y[j])
			Tx.append(xtemp)
			Ty.append(ytemp)
		# plot iterated line segment (different color)	
		x = Tx
		y = Ty
		
	plot(x,y, ',', label=str(n+1))
	
	# Setup the plot
	xlabel('theta(n)') # set x-axis label
	ylabel('p(n)') # set y-axis label
	title("Standard Map with k =" + str(k)) # set plot title

	# Display the plot in a window
	show()