# - - - - - - - - - - - - - - - - - - - - - - - - - - - -
# Lennard-Jones potential for chemical bonds
# find amplitude and period of oscillations
# - - - - - - - - - - - - - - - - - - - - - - - - - - - -
#
import numpy as np
import matplotlib.pyplot as plt
import scipy.optimize as opt
import scipy.integrate as integrate
import sys

# global variables
r0 = 5.0       # parameters for potential
dd0 = 1.0
energy0 = 0.0

# - - - - - - - - - - - - - - - - - - - - - - - - - - - -
def LennardJones(rr):
    """ Lennard-Jones potential
    """
    aux = (r0/rr)**6
    lj  = (dd0*aux - 2.0)*aux    
    return lj

# - - - - - - - - - - - - - - - - - - - - - - - - - - - -
def LennardJones0(rr):
    """ Lennard-Jones potential minus an energy constant
    """
    lj = LennardJones(rr)    
    return energy0-lj


# - - - - - - - - - - - - - - - - - - - - - - - - - - - -
def periodIntegrand(rr):
    """ 
    The integrand to calculate period of oscillations 
    from energy conservation
    """
    lj = LennardJones(rr)
    aux = 2*(energy0-lj)
    aux = 1.0/np.sqrt(aux)    
    return aux
    
    
# - - - - - - - - - - - - - - - - - - - - - - - - - - - -
def graph_LJ():
    """plot the Lennard-Jones potential
    """
    plt.figure(figsize=(6,6))             # prepare graph
    plt.xlabel('$r$',fontsize=20)
    plt.ylabel('$V$',fontsize=20)
    plt.title("Lennard-Jones Potential")
    xmin = int(rmin/2.0); xmax = 2*rmin
    ymin = vmin - 1.0; ymax = 3.0
    plt.axis([xmin,xmax,ymin,ymax])
    
    xx = np.linspace(xmin,xmax,100)
    yy = LennardJones(xx)
    line = plt.plot(xx,yy,'k-')        #plot potential
    yy0 = 0.0*xx
    line = plt.plot(xx,yy0,'k:')       # plot dotted line at zero level
    dx = 2; dy = 1
    plt.xticks(np.arange(xmin,xmax+dx,dx),fontsize=14)
    plt.yticks(np.arange(ymin,ymax+dy,dy),fontsize=14)
    #plt.text(0.09,30,r'$P$',fontsize=20)

    plt.show()
    

# - - - - - - - - - - - - - - - - - - - - - - - - - - - -
# Main program

# Analytical results
rmin = dd0**(1./6.)*r0
print "Minimum of potential at ", rmin
vmin = -1.0/dd0
print "Minimum of potential is ", vmin
omega0 = np.sqrt(72.)/(r0*dd0**(2./3.))
print "Frequency of small oscillations ", omega0
period0 = 2*np.pi/omega0
print "Period of small oscillations ", period0

# make a graph of the potential
#graph_LJ()


# allowed region of motion
#energy = raw_input("Enter energy: ")
#energy0 = float(energy)
energy0 = -0.9
r1 = rmin/(1+np.sqrt(1+dd0*energy0))**(1./6.)
r2 = rmin/(1+np.sqrt(1-dd0*energy0))**(1./6.)
print "Interval of oscillating motion: (found analytically)",r1,r2

r1 = opt.newton(LennardJones0,0.9*rmin,tol=1.e-4,maxiter=200)
r2 = opt.newton(LennardJones0,1.1*rmin,tol=1.e-4,maxiter=200)
print "Interval of oscillating motion: (found numerically)",r1,r2

# period of oscillation (found by integration of energy function)
r1 = r1  #+0.0001
r2 = r2  #-0.0001
period = integrate.quadrature(periodIntegrand,r1,r2,tol=1.e-4)
print "Period of oscillation: ", 2*period[0]