# The henonheiles.py program 
# The equation of motion derived from H = 1/2*(px^2+py^2+A*x^2+B*y^2)+D*y*x^2-C*y^3/3
# In the regular Henon-Heiles Hamiltonian A=B=C=D=1.0
# All the variables are dimentionless

# Call module
import numpy as np
import string as st
import pylab as pl
from ProjectFunctions import *
import visual as vi

# Start collecting run variables
PlotType = pType()

# Plots Axis detiles
PlotAxis = pAxis(PlotType)
     

# Hamiltonian and Energy Data
# interesting energy levels 0<E<1/6
a,b,c,d = HamConst()
E,Estr = eLevel()

# Number of iteration
nIterate = nIter(PlotType)
dt = setDt(0.001)


# Start to generate plots
print '\n'
print 'Iteresting region for initial points in the x-y plane:'
print '(0,1),(sqrt(3)/2,-1/2), (-sqrt(3)/2,-1/2)'

# Default values
tRajectories = {'1':'ST' , '2':'AT'}
nTrajectories = {'1':1, '2':10}
ChooseTrajectory = '1'
tTrajectory = tRajectories[ChooseTrajectory]
nTrajectory = nTrajectories[ChooseTrajectory]
ATinitPoint = '1'	# Choosing initial position randomly along equipotential E


# Plot one point trajectory or array of points
if PlotType == 'PC':
    print '\n'
    print 'Choose:'
    print '(1) One point trajectory'
    print '(2) Array of points trajectories'
    inpStr = '[%s]: ' % ChooseTrajectory
    stemp = raw_input(inpStr)
    if (stemp in ['1','2']):
	if stemp != ChooseTrajectory:
	    tTrajectory = tRajectories[stemp]
	    nTrajectory = nTrajectories[stemp]
	    ChooseTrajectory = stemp
	if stemp=='2':   	# Array trajectory
	    inpStr = 'Number of points in the array [%g]: ' % nTrajectory
	    stemp = raw_input(inpStr)
	    if stemp !='':
		nTrajectory = int(stemp)          

if tTrajectory == 'ST':
    print '\n'
    print 'Initial position for single point trajectory'
    
    print 'Inetersting initial points for E = 1/12'
    print '[x ,    y  , py ]'
    print '[0 , 0.4999, 0]  -  Eliptic'
    print '[0 , -0.366, 0]  -  Eliptic'
    print '[0 , 0     , 0]  -  Orbital'
    
    print 'Inetersting initial points for E = 1/6'
    print '[x ,    y  , py ]'
    print '[0 , 0.9999, 0]  -  Eliptic'
    print '[0 , 0.1   , 0]  -  Orbital'
    print '[0 , -0.15, 0]  -  Chaotic'
    
    
    # Default Initial position
    x = [0.0]
    y = [0.1]
    py = [0.0]
    px = []
    inpStr = 'Initial value for x is [%.3f]: ' % x[0]
    stemp = raw_input(inpStr)
    if stemp != '':
	x[0] = float(stemp)
    inpStr = 'Initial value for y is [%.3f]: ' % y[0]
    stemp = raw_input(inpStr)
    if stemp != '':
	y[0] = float(stemp)
    inpStr = 'Initial value for py is [%.3f]: ' % py[0]
    stemp = raw_input(inpStr)
    if stemp != '':
	py[0] = float(stemp)
    x0 = x[0]
    y0 = y[0]
    py0 = py[0]
    try:
	Epx(a,b,c,d,E,x0,y0,py0)
    except ValueError:
	print 'Not good initial conditions'
    else:
	ppx = Epx(a,b,c,d,E,x0,y0,py0)
	px.append(ppx)
	px0 = px[0]
    
if tTrajectory == 'AT':  
    ymin,ymax = yRange(E)
    ystep = (ymax - ymin)/nTrajectory
    
    
# Start the calculation
# Single point plots
if tTrajectory == 'ST':
    if PlotType == '3D':
	# Ploting	
	vi.scene.title = "Henon Heiles"
	vi.scene.center = vi.vector(0,0,.4)
	vi.scene.height = 600
	vi.scene.width = 600
	vi.scene.background = (1,1,1)
	vi.scene.forward = (0,1,-0.5)
	vi.scene.lights = [ vi.vector(-0.25, -0.5, -1.0) ,vi.vector(0.25, 0.5, 1.0) ]
	# Draw axes
	vi.curve( pos = [ (-1,0,0), (1,0,0) ], color = (0,0,0), radius = 0.005 )
	vi.curve( pos = [ (0,-1,0), (0,1,0) ], color = (0,0,0), radius = 0.005 )
	vi.curve( pos = [ (0,0,0), (0,0,1) ], color = (0,0,0), radius = 0.005 )
	henon = vi.curve( color = vi.color.black, radius=0.005 )
	dit = 1./nIterate
	col = 0
	for i in range(nIterate):
	    xx,yy,ppx,ppy = RK4DIntegrator(a,b,c,d,x[-1],y[-1],px[-1],py[-1],dxdt,dydt,dpxdt,dpydt,dt)	
	    x.append(xx)
	    y.append(yy)
	    px.append(ppx)
	    py.append(ppy)
	    col += dit
	    if 'p_x' in PlotAxis:
		henon.append( pos=(xx,ppx,yy),color=(col,col, 1-col))
	    else:
		henon.append( pos=(yy,ppy,xx),color=(col,col, 1-col))
    if PlotType == 'PC':
	# Considering only x = 0 plane intersection
	xx = x[0]
	yy = y[0]
	ppx = px[0]
	ppy = py[0]
	xold = 0.0
	for i in range(nIterate):
	    xx,yy,ppx,ppy = RK4DIntegrator(a,b,c,d,xx,yy,ppx,ppy,dxdt,dydt,dpxdt,dpydt,dt)	
	    test = xold * xx 	# If test < 0 it means that we croosed the x = 0 plane
	    xold = xx
	    if test < 0:	# Collecting only the points that cross the x = 0 plane
		y.append(yy)
		py.append(ppy)	  
	    #x.append(xx)
	    #y.append(yy)
	    #px.append(ppx)
	    #py.append(ppy)
    
# Multi point plot
if tTrajectory == 'AT':
    titStr = 'Henon-Heiles with E = %s ' % (Estr)
    for nn in range(nTrajectory):
	x = []
	y = []
	px = []
	py = []
	# The variables that 
	xx = 0.0
	xold = 0.0
	yy = ymin + nn * ystep
	ppy = 0.0
	ppx = Epx(a,b,c,d,E,xx,yy,ppy)
	x.append(xx)
	y.append(yy)
	py.append(ppy)
	px.append(ppx)
	
	for i in range(nIterate):
	    err = abs(E-EE(xx,yy,ppx,ppy,a,b,c,d))
	    if err > 1e-6:
		print 'Check numerical error: ',err
	    xx,yy,ppx,ppy = RK4DIntegrator(a,b,c,d,xx,yy,ppx,ppy,dxdt,dydt,dpxdt,dpydt,dt)
	    test = xold * xx 	# If test < 0 it means that we croosed the x = 0 plane
	    xold = xx
	    if test < 0:	# Collecting only the points that cross the x = 0 plane
		y.append(yy)
		py.append(ppy)	    
	if 'p_x' in PlotAxis:
	    print 'this option is not programed'
	    #pl.plot(xpc,pxpc, 'r,')
	else:
	    col = nn*1.0/float(nTrajectory)
	    pl.plot(y,py, ',', color=(col,col, 1-col))
	
	    
	    
# Ploting 

if (PlotType == 'PC') and (tTrajectory == 'ST'):
    pl.figure(1,(8,6))
    if 'p_x' in PlotAxis:
	pl.plot(x,px, 'b,')
	pl.xlabel('x') # set x-axis label
	pl.ylabel('$p_x$') # set y-axis label
	if tTrajectory == 'AT':
	    titStr = 'Henon-Heiles with E = %s ' % (Estr)
	else:
	    pl.plot([x0],[px0],'ro')
	    titStr = 'Henon-Heiles with E = %s , ($x_0 , px_0$) = (%.3f,%.3f) ' % (Estr,x0,px0)
    else:
	pl.plot(y,py, 'b,')
	pl.xlabel('y') # set x-axis label
	pl.ylabel('$p_y$') # set y-axis label
	if tTrajectory == 'ST':
	    titStr = 'Henon-Heiles with E = %s ' % (Estr)
	else:
	    titStr = 'Henon-Heiles with E = %s , ($y_0 , py_0$) = (%.3f,%.3f) ' % (Estr,y0,py0)
	    pl.plot([y0],[py0],'ro')
    
if (PlotType == 'PC'):   
    pl.title(titStr) # set plot title
    pl.axis([-0.6, 1, -0.6, 0.6])
    pl.show()

raw_input('Press any key to continue...')