#!/usr/bin/env python
"""
Created on Sun 19 Feb 2019/Magne Guttormsen
"""
from numpy import *
from matplotlib.pyplot import *

def psikw_r(x,t):
    return fac*(10.*sqrt(2/a)*sin(1*pi*x/a)*cos(w1*t)
               - 5.*sqrt(2/a)*sin(3*pi*x/a)*cos(w3*t)
               + 1.*sqrt(2/a)*sin(5*pi*x/a)*cos(w5*t)
                )

def psikw_i(x,t):
    return -fac*(10.*sqrt(2/a)*sin(1*pi*x/a)*sin(w1*t)
                - 5.*sqrt(2/a)*sin(3*pi*x/a)*sin(w3*t)
                + 1.*sqrt(2/a)*sin(5*pi*x/a)*sin(w5*t)
                )

mc2 =   0.511E+06       # eV
hbarc = 197.3           # nm eV
c=      299.792458      # nm/fs
a =     1.              # nm
fac=    1./sqrt(126.)

w1= 1 *pi*pi*c*hbarc/(2*mc2*a*a)
w3= 9 *pi*pi*c*hbarc/(2*mc2*a*a)
w5= 25*pi*pi*c*hbarc/(2*mc2*a*a)
#print w1

n = 100
t = 0
T = 1.41
dt = 0.01
x = linspace(0, a, n ) # in nm

ion()

figure()
# Plotter initialtilstand naar t = 0
lineR, = plot( x, psikw_r(x, t ),'b',alpha=0.7,label='Real' )
lineI, = plot( x, psikw_i(x, t ),'g',alpha=0.7,label='Imaginary' )
lineP, = plot(x,psikw_r(x,0)**2 + psikw_i(x,0)**2,'k',label='Probability density')
time_text = text(0.1, 2.5, r't = {:.2f} fs'.format(t))
legend(loc='upper right')

xlabel('$x$ [nm]')  # Tekst langs x-aksen
ylabel('$\Psi (x, t)$ [1/nm$^{1/2}$] or $|\Psi (x, t)|^2$ [1/nm]') # Tekst langs y-aksen
xlim([0, 1])        # Grenser x-akse
ylim([-2, 5])       # Grenser y-akse

draw()

while t < T:
    lineR.set_ydata( psikw_r( x, t ) )
    lineI.set_ydata( psikw_i( x, t ) )
    lineP.set_ydata(psikw_r(x,t)**2+psikw_i(x,t)**2)  # update the data
    time_text.set_text(r't = {:.2f} fs'.format(t))
    draw()
    pause(.1)
    t += dt

ioff()
show()