#!/usr/bin/env python
"""
Created on Sun 27 Oct 2013.
Modified to solve the home-exam in FYS2140 March 2014.

Uses the Crank Nicolson scheme to solve the time dependent Schrodinger equation
for a harmonic oscillator.

Animation is done using the matplotlib.pyplot library.

@author Kristoffer Braekken
@author Are Raklev
"""
# Tools for sparse matrices
import scipy.sparse as sparse
import scipy.sparse.linalg

# Numerical tools
from numpy import *

# Plotting library
from matplotlib.pyplot import *

"""Physical constants"""
_mpcsq = 938.27e6          # Rest energy for a proton [eV]
_mClcsq = 34.9689*931.48e6 # Rest energy for a ^36Cl isotope  [eV]
_mucsq = _mpcsq*_mClcsq / (_mpcsq + _mClcsq) # Reduced mass rest energy [eV]
print 'Rest energy for reduced mass: ', _mucsq/1000./1000., ' MeV'
_hbarc = 197.3             # [eV nm]
_c = 3.0e5                 # Speed of light [nm / ps]


# Function definitions
def Psi0( x ):
    '''
    Initial state for a stationairy gaussian wave packet.
    '''
    x1 = 0.240   # [nm] Start here
    a =  0.004 # [nm] Width of packet

    A = ( 1. / ( 2 * pi * a**2 ) )**0.25
    K1 = exp( - ( x - x1 )**2 / ( 4. * a**2 ) )

    return A * K1

def morse( x, V0 = 10.2, x0 = 0.127, b = 12.6 ):
    """
    The Morse potential. Calculated in eV units.
    
    @param V0 potential minimum value [eV]
    @param x0 potential minimum [nm]
    @param b potential shape [nm^-1]
    """
    
    potential = V0 * ( exp( -2.*b * (x-x0) ) - 2.*exp( -b * (x-x0)) )
    
    return potential

# Main programme
if __name__ == '__main__':
    
    # Define some numbers for x-axis
    nx = 2000   # Number of points in x direction
    dx = 0.0002 # Distance between x points [nm]

    # Interval [a,b], use only points on the side of x=0 where we are
    a = 0.0
    b = nx * dx
    x = linspace( a, b, nx )   # x is now an array containing all x-values used
    
    # Time parameters for animation
    T = 0.613                  # How long to run simulation [ps]
    dt = 1e-5                  # Size of time step [ps]
    t = 0                      # Start-time
    time_steps = int( T / dt ) # Number of time steps

    # Constants - save time by calculating outside of loop
    k1 = - ( 1j * _hbarc * _c) / (2. * _mucsq )
    k2 = ( 1j * _c ) / _hbarc
    # print k1
    # print k2
    
    # Create the matrix containing central differences. It it used to
    # approximate the second derivative.
    data = ones((3, nx))
    data[1] = -2*data[1]
    diags = [-1,0,1]
    D2 = k1 / dx**2 * sparse.spdiags(data,diags,nx,nx)
    # print D2

    # Identity Matrix
    I = sparse.identity(nx)

    # Create the diagonal matrix containing the potential.
    V_data = morse(x)
    V_diags = [0]
    V = k2 * sparse.spdiags(V_data, V_diags, nx, nx)
    # print V

    # Put mmatplotlib in interactive mode for animation
    ion()

    # Setup the figure before starting animation
    fig = figure() # Create window
    ax = fig.add_subplot(111) # Add axes
    Psi = Psi0(x)   # Create the initial state as an array Psi from the array x
    line, = ax.plot( x, abs(Psi)**2, label='$|\Psi(x,t)|^2$' ) # Fetch the line object
    ax.axis([-0.2, 0.5, -11, 100])

    # Also draw a green line illustrating the potential
    ax.plot( x, morse( x ), label='$V(x)$' )

    # Add other properties to the plot to make it elegant
    fig.suptitle("Solution of Schrodinger's equation with Morse potential") # Title of plot
    ax.grid('on') # Square grid lines in plot
    ax.set_xlabel('$x$ [nm]') # X label of axes
    ax.set_ylabel('$|\Psi(x, t)|^2$ [1/nm] og $V(x)$ [eV]') # Y label of axes
    ax.legend(loc='best') # Adds labels of the lines to the window
    draw() # Draws first window

    # Find expectation value and standard devation for position x by integrating
    expx = sum(x*abs(Psi)**2*dx)                      # Expectation value
    sigmax = sqrt(sum(x*x*abs(Psi)**2*dx) - expx**2)  # Standard deviation
    
    # print expx
    expx_list = []               # List to store expectation value
    sigmax_list = []             # List to store standard deviation
    time_list = []               # List to store time
    expx_list.append(expx)
    sigmax_list.append(sigmax)
    time_list.append(0)
 
    # Dont plot all frames.
    # Creates better "framerate" when requiring low time step.
    PLOT_EVERY = 10
    counter = 0 # Counter for drawing frames

    # Time loop
    while t < T:
        """
            For each iteration: Solve the system of linear equations:
            (I - k/2*D2) u_new = (I + k/2*D2)*u_old
        """
        # Set the elements of the equation
    	A = (I - dt/2*(D2 + V))
    	b = (I + dt/2. * (D2 + V)) * Psi

        # Calculate the new Psi
    	Psi = sparse.linalg.spsolve(A,b)

        # Update time
    	t += dt
        
        # Update counter for frame
        counter += 1
        
        # Choose which frames to plot and when to calculate expectation value
        if counter == PLOT_EVERY:
            # Plot this new state
            line.set_ydata( abs(Psi)**2 ) # Update the y values of the Psi line
            draw() # Update the plot
            counter = 0

            # Calculate expectation value for position by integrating and store in list
            expx = sum(x*abs(Psi)**2*dx)
            sigmax = sqrt(sum(x*x*abs(Psi)**2*dx) - expx**2)
            expx_list.append(expx)
            sigmax_list.append(sigmax)
            time_list.append(t)
            # print expx

    # New figure for expectation value and standard deviation versus time
    figure()
    plot(time_list,expx_list, label= r'$\langle x \rangle$')
    plot(time_list,sigmax_list, label= r'$\sigma_x$')
    axis([0, T, 0.0, 0.3])         # Set axis range
    xlabel('$t$ [ps]')                 # Label for x-axis
    ylabel(r'$\langle x \rangle$ [nm] og $\sigma_x$ [nm]') # Label for y-axis
    legend(loc='best') # Adds labels of the lines to the window
    savefig('expx_morse.eps')          # Save as .eps figure
    
    # Turn off interactive mode
    ioff()

    # Add show so that windows do not automatically close
    show()