#!/bin/ksh

# Location of binaries
BIN=$HOME/bin
a0=0.529177

# VASP outputfiles
outfile=OUTCAR
posfile=POSCAR

# Generate the LDOS(x,y,z) data from the VASP output
length=31.747 
width=22.449


# Number of real space points in the data
nx=69
ny=69
nz=29

# Central point of plot
center=1  # 1=> centered on defect, 2=> defect at corner



# First parameter on the command line is the LDOS desired if given
if [ $# -eq 2 ]; then
    ldos=$1
    shift 1
    defect_atom=$1
    shift 1
    echo $ldos $defect_atom	
elif [ $# -eq 0 ]; then
# Value of the LDOS to plot
    ldos=0.000005
# Set the position of the defect atom to -100 to not specify a point
    defect_atom=-100
else
    echo "Usage: stm_atom_positions <ldos> <numatom>"
    exit 2
fi

# Distance (a.u.) below top atomic plane to plot atomic positions
dist_below=10.0

# Height (Angstroms) above highest atomic plane to begin the 
# LDOS calculation for
start_height=1.0
start_height=$(echo $start_height $a0 | awk '{printf "%12.8f \n", $1/$2}')

# Height (in a.u.) above starting plane to consider
height=10.0


# Everything below should be automatic
###########################################################################
# Dump the atomic coordinates to a file
natom=$($BIN/vp | grep "Total number of atoms:"|awk '{print $5}')

echo $natom >| pos.dat
echo >> pos.dat
$BIN/vp | tail -$natom | awk '{printf "%s %12.9f %12.9f %12.9f %s \n", $1, $2, $3, $4, NR}' >> pos.dat


# Make a jmol file with the defect atom labeled.
rm -f pos2.dat
i=-2
while read label x y z tmp
do
	i=$((i+1))
	if [ $i -eq $defect_atom ]; then
		echo "ZR" $x $y $z $tmp >> pos2.dat
	else
		echo $label $x $y $z $tmp >> pos2.dat
	fi

done < pos.dat

# Make a file to use in processing atom positions
echo $natom >| pos.dat
echo >> pos.dat
$BIN/vp | tail -$natom | awk '{printf "%s %12.9f %12.9f %12.9f %s \n", $1, $2/'$a0', $3/'$a0', $4/'$a0', NR}' >> pos.dat


#_jms 4/2/03
# Make sure the POSCAR file is compatible with the input coordinates, otherwise warn the user and stop
a3=$(sed -n '5 p' $posfile | awk '{printf "%12.8f \n", $3}')
coordtype=$(sed -n '7 p' $posfile)
if [ $coordtype = 'Cartesian' ]; then
    a3=1.0
fi
echo $natom >| pos_check.dat
echo >> pos_check.dat
start=8
end=$(echo $start $natom | awk '{print $1+$2-1}')
sed -n ''$start','$end' p' $posfile | awk '{printf "%12.8f \n", $3*'"$a3"'/'$a0'}' >> pos_check.dat
tail -$natom pos.dat | awk '{printf "%6.3f \n", $4}' >| tmp.dump1
tail -$natom pos_check.dat | awk '{printf "%6.3f \n", $1}' >| tmp.dump2
paste tmp.dump1 tmp.dump2 | awk '{printf "%s %6.3f %6.3f %6.3f \n", NR, $1, $2, sqrt(($1-$2)^2)}' | sort -n -k 4 >| tmp.dump3
userpos3='f'
nl=$(tail -1 tmp.dump3 | awk '$4>=0.1 {print "F", $1, $2, $3, $4}'  | tee tmp.dump1 | wc -l | awk '{print $1}')
if [ $nl -eq 1 ]; then
    read check trash < tmp.dump1
    if [ $check = 'F' ]; then
	echo $check $trash
	echo "Your input coordinates differ signinificantly from the output vasp coordinates."
	echo "This may make your stm simulation incorrect."
	echo "You should provide your own pos3.dat file, give it the name my_pos3.dat"
	echo
	userpos3='t'
    fi
fi
rm -f tmp.dump1 tmp.dump2 tmp.dump3
#_jms 4/2/03

#_jms 4/2/03
rm -f pos3.dat
set -A def_position 0.0 0.0 0.0
if [ $userpos3 = 'f' ]; then
i=-2
while read label x y z tmp
do
	i=$((i+1))
	if [ $i -eq $defect_atom ]; then
		echo "ZR" $x $y $z $tmp >> pos3.dat
		set -A def_position $x $y $z
	else
		echo $label $x $y $z $tmp >> pos3.dat
	fi

done < pos.dat

else
    if [[ -s my_pos3.dat ]]; then
	echo
	echo "Thank you for providing a pos3.dat file."
	echo
	cp my_pos3.dat pos3.dat
    else
	echo "Looking for my_pos3.dat"
	echo "You should provide a pos3.dat file since your coordinates are difficult."
	exit 0
    fi
fi
#_jms 4/2/03


# Get the highest atomic plane of atoms
zmax=$(sort -n -k 4 pos3.dat | tail -1 | awk '{printf "%12.8f \n", $4}')
zstart=$(echo $zmax $start_height | awk '{printf "%12.8f \n", $1+$2}')


echo "Highest plane of atoms at z=" $zmax
echo "Defect atom at (xyz)=" ${def_position[*]}

# Determine the number of spin states
nspin=$(grep "ISPIN" $outfile | awk '{print $3}')

# Get the lattice vectors
#set -A a1 $(grep "A1" $outfile | sed 's/\,//g' | sed 's/(//g' | sed 's/)//g' | sed 's/-//g' | awk '{print $3, $4, $5}' )
#set -A a2 $(grep "A2" $outfile | sed 's/\,//g' | sed 's/(//g' | sed 's/)//g' | sed 's/-//g' | awk '{print $3, $4, $5}' )
set -A a1 $(sed -n '3 p' CHGCAR | awk '{printf "%12.8f %12.8f %12.8f \n", $1/'"$a0"', $2/'"$a0"', $3/'"$a0"'}' )
set -A a2 $(sed -n '4 p' CHGCAR | awk '{printf "%12.8f %12.8f %12.8f \n", $1/'"$a0"', $2/'"$a0"', $3/'"$a0"'}' )
#exit 0



# Get the atom number types
set -A ntype $(sed -n '6 p' $posfile)
echo "Number of atom types:" ${#ntype[*]}
echo "Multiplicities of these types:" ${ntype[*]}
ns=${#ntype[*]}
i=0
natoms=0
while [ $i -le $((ns-1)) ]; do
    nt=${ntype[$i]}
    let natoms=$natoms+$nt
    i=$((i+1))
done
echo "Total number of atoms:" $natoms
set -A type
set -A label $(grep "POTCAR:" $outfile | head -$ns | sed 's/POTCAR://g' | sed 's/US//g'| sed 's/PAW//g')

# Create a vasp2ab file
cat >| vasp2ab.in <<EOF
CHGCAR
${#ntype[*]}      # Number of atom types
EOF
i=0
while [ $i -le $((ns-1)) ]; do
echo ${label[$i]} 0 "# Atom Labels" >> vasp2ab.in
    i=$((i+1))
done
cat >> vasp2ab.in <<EOF2
$nspin     # Number of spin states
0          # Density (0) or Potential (1)
EOF2
$BIN/vasp2ab.exe < vasp2ab.in


if [ $nspin -eq 2 ]; then

# Total density
spin_component=1

cat >| cut3d.in <<EOF
txo_VASP_DEN
1                                             # unformatted
$spin_component                               # which spin component to print
4                                             # Volume
2                                             # Cartesian points
0.0 0.0  0.0                                  # X1
1.0 0.0  0.0                                  # X2
0.0 1.0  0.0                                  # X3
1                                             # Central point in Cartesian coordinates
${def_position[0]} ${def_position[1]} $zstart # Central point
$width                                        # Plane width
$length                                       # Plane length
$height                                       # Height 
3                                             # Up
$center                                             # Corner (2) or center (1)
$nx                                           # W resolution
$ny                                           # L resolution
$nz                                           # H resolution
data.in
1       # ASCII
1       # keep parameters
2       # don't run again
EOF

else

cat >| cut3d.in <<EOF
txo_VASP_DEN
1                                             # unformatted
4                                             # Volume
2                                             # Cartesian points
0.0 0.0  0.0                                  # X1
1.0 0.0  0.0                                  # X2
0.0 1.0  0.0                                  # X3
1                                             # Central point in Cartesian coordinates
${def_position[0]} ${def_position[1]} $zstart # Central point
$width                                        # Plane width
$length                                       # Plane length
$height                                       # Height 
3                                             # Up
$center                                             # Corner (2) or center (1)
$nx                                           # W resolution
$ny                                           # L resolution
$nz                                           # H resolution
data.in
1       # ASCII
1       # keep parameters
2       # don't run again
EOF

fi


# Run the cut3d program and make a file for simulating STM image
$BIN/cut3d_n < cut3d.in
echo  $((nx+1)) $((ny+1)) $((nz+1)) >| data.in2
cat data.in >> data.in2
mv -f data.in2 data.in

# Now we generate the atomic positions to include in the figure
echo $center >| atoms.in
echo $length $width >> atoms.in
zb=$(echo $zmax $dist_below | awk '{printf "%12.8f ", $1-$2}')
echo $zb
nk=$(sort -n -k 4 pos3.dat | awk '$4 >= '"$zb"' {printf "%s %12.8f %12.8f %12.8f \n", $1, $2, $3, $4-'"$zb"'-'"$dist_below"'}' | wc -l)
echo $nk >> atoms.in
echo ${a1[*]} >> atoms.in
echo ${a2[*]} >> atoms.in
def_position[2]=$(echo ${def_position[2]} $zb $dist_below  | awk '{printf "%12.8f \n", $1-$2-$3}')
echo ${def_position[0]} ${def_position[1]} ${def_position[2]} >> atoms.in
sort -n -k 4 pos3.dat | awk '$4 >= '"$zb"' {printf "%s %12.8f %12.8f %12.8f \n", $1, $2, $3, $4-'"$zb"'-'"$dist_below"'}' >> atoms.in




# Run the stm fitting code
cat >| stm.in <<EOF
data.in
stm.out
2         # 0 => simple search, 1 = cubic spline fit, 2 = exponential fit
$ldos
EOF

$BIN/stm_fit.exe < stm.in

# Produce some atomic coordinates
$BIN/2d.exe < atoms.in

# Check the range of the STM and the atoms
vischeck=1
if [ $vischeck -eq 0 ]; then

cat >| plot_atoms <<EOF
plot "atoms.out" using 2:3 with points
pause -1
EOF
	gnuplot "plot_atoms"

cat >| plot_stm <<EOF
plot "stm.out" using 4:5 with points
pause -1
EOF
	gnuplot "plot_stm"
fi


exit 0