#!/usr/bin/python

#Part of BeamSimulator.sh

import sys

pi = 3.141592654;
camau = 1822.89;
chaev = 27.2114;


#BASH outputs strings for decimals, convert everything to floats
ZStart = float(sys.argv[1])
ZEnd = float(sys.argv[2])
nZ = int(sys.argv[3])
Mass1 = float(sys.argv[4])
Mass2 = float(sys.argv[5])
EMinWP = float(sys.argv[6])
EMaxWP = float(sys.argv[7])
Z0 = float(sys.argv[8])
ZAna = float(sys.argv[9])
ZOptStart = float(sys.argv[10])
dZ = float(sys.argv[11])
nZscat = int(sys.argv[12])
ZScatEnd = float(sys.argv[13])
RStart = float(sys.argv[14])
REnd = float(sys.argv[15])
nR = float(sys.argv[16])

#print
#print "START python script WPcheck.py"
#print
#print "Center of wave packet at Z = ", Z0
#print "EminWP: ", EMinWP, "eV"
#print "EmaxWP: ", EMaxWP, "eV"
#print
#print "ZStart = ", ZStart
#print "ZEnd = ", ZEnd
#print "nZ = ", nZ
#print "Mass1 = ", Mass1
#print "Mass2 = ", Mass2
#print "EMinWP = ", EMinWP
#print "EMaxWP = ", EMaxWP
#print "Z0 = ", Z0
#print "ZAna = ", ZAna
#print "ZOptStart = ", ZOptStart
#print "dZ = ", dZ
# -------- convert to atomic units -------------------

EMinWP = EMinWP/chaev;
EMaxWP = EMaxWP/chaev;

Mass1 = camau * Mass1;
Mass2 = camau * Mass2;

# -------- compute some basic numbers ---------------
#Mark... Waarom moeten de imports hier?
#Als ik ze eerder doe, dan verlies ik precision.... Ik snap maar niet waarom.

from math import *
from cmath import *
from numpy import *
from numpy.fft import *



TotalMass = Mass1 + Mass2;
ReducedMass = Mass1 * Mass2 / TotalMass;

kmin = - sqrt( 2.0 * TotalMass * EMinWP );
kmax = - sqrt( 2.0 * TotalMass * EMaxWP );
k0 = (kmin + kmax) / 2.0;
alpha = ((k0 - kmin)**2) / 4.0;
EAverage = (k0**2)/(2.0* TotalMass);

# ---------- build initial wavepacket on Z grid -----------
ZGrid = linspace( ZStart, ZEnd, nZ );
ZScatGrid = linspace( ZStart, ZScatEnd, nZscat );
RGrid = linspace( RStart, REnd, nR );

anrm = (2.0*alpha/pi)**0.25;
arkz = 1j * sqrt( 2.0 * TotalMass * EAverage );

#print "anrm = ", anrm, "  arkz = ", arkz, " alpha = ", alpha

WP = []; AbsWP = [];
for z in ZGrid:
   expz = - alpha * (z - Z0)*(z - Z0) - arkz * z;
   psi = anrm * exp(expz);
   WP.append( psi );
   AbsWP.append( abs(psi) );

indexZofZ0 = abs(ZGrid - Z0).argmin();
indexZofZAna = abs(ZGrid - ZAna).argmin();
indexZofZOptStart = abs(ZGrid - ZOptStart).argmin();

NormInitialWP =  0.0;
for i in range( indexZofZAna, indexZofZOptStart ):
 NormInitialWP += dZ * AbsWP[ i ] * AbsWP[ i ];

#print
#print "Specular Grid in Z:"
#print ZGrid
#print
#print "Scattering Grid in Z:"
#print ZScatGrid
#print
#print "Grid in R:"
#print RGrid
#print
#print "ABS(Wavepacket) on ZSpec grid:"
#print AbsWP;
#print
#print "ABS(Wavepacket) at ZAna:", AbsWP[ indexZofZAna ];
#print "ABS(Wavepacket) at Z0:", AbsWP[ indexZofZ0 ];
#print "ABS(Wavepacket) at ZOptStart:", AbsWP[ indexZofZOptStart ];
#print
print "  Norm of WP between ZANA and ZSpecOptStart:", NormInitialWP
#print "  ZANA:          ", ZAna
#print "  ZSpecOptStart: ", ZOptStart
#print "  WPcenter:      ", Z0
if( abs( NormInitialWP - 1.0 ) > 1.0E-9 ): 
	print "WARNING the wavepacket ranges outside the ZAna to ZOptStart range!";
        print AbsWP[ indexZofZAna:indexZofZOptStart ]

	print
	print "Specular Grid in Z:"
	print ZGrid
print
print "ABS(Wavepacket) in Z between ZANA and ZSpecOptStart:"
print AbsWP[ indexZofZAna:indexZofZOptStart ]
print

# ---------- build the kZ grid -------------------------
kZGrid = ZGrid;

dkZ = 2.0*pi/(nZ*dZ);

kZGrid[ 0 ] = 0.0;
kZGrid[ nZ/2 ] = nZ * 0.5 * dkZ;

for i in range( 2, nZ/2 + 1 ):
   kZGrid[ i - 1 ] = (i-1) * dkZ;
   kZGrid[ nZ - i + 1 ] = -kZGrid[ i - 1 ];
#print
#print "dkZ in a.u.:", dkZ;
#print "The maximum kZ (dkZ*nZ/2) in a.u.:", kZGrid[ nZ/2 ], (dkZ*nZ/2);
#print "The minimum kZ (-dkZ*(nZ/2-1)) in a.u.:", kZGrid[ nZ/2 + 1 ], (-dkZ*(nZ/2-1));
#print
#print "kmin in a.u. was:", kmin;
#print "kmax in a.u. was:", kmax;
#print "k0 in a.u. was:", k0;
#print
#print "Grid in kZ: "
#print kZGrid;
#print
if( kmin < kZGrid[ nZ/2 + 1 ]): print "  WARNING kmin of WP smaller than lowest negative momentum on kZGrid!";
if( kmax > kZGrid[ nZ/2 ]): print "  WARNING kmax of WP bigger than highest momentum on kZGrid!";

# ------------------ fourier transform WP to k ----------------
kWP = dZ  / sqrt(2.0*pi) * fft( WP );
AbskWP = abs( kWP );
Maxindex = abs( kWP ).argmax();

kNormInitialWP = 0.0;
for i in range( nZ/2+1, nZ ):
 kNormInitialWP += dkZ * AbskWP[ i ] * AbskWP[ i ];

#print
#print "ABS(Wavepacket) on grid in kZ:" 
#print AbskWP;
#print
#print "ABS(Wavepacket) in kZ at zero momentum:", AbskWP[ 0 ];
#print "ABS(Wavepacket) in kZ at lowest momentum:", AbskWP[ nZ/2+1 ];
#print "kZ value for which ABS(Wavepacket) on kZ grid is biggest:", kZGrid[ Maxindex ];
#print "k0 in a.u. was:", k0;
#print "ABS(Wavepacket) in kZ at highest value:", AbskWP[ Maxindex ];
#print

print "  Norm of WP in kZ for negative momenta only:", kNormInitialWP
if( abs( kNormInitialWP - 1.0 ) > 1.0E-9 ): print "WARNING the wavepacket has too much positive momenta!";
#print
#print "END python script"
#print