#!/bin/ksh

# Script to find the partial occupancies for STM image
# computations. This script is to be executed in the directory
# that contains the VASP ground state output files.
# It requires a command line argument that specifies the energy
# range (below the Fermi level or top of VBM) to integrate the LDOS
# over.
# JMS NRL 1/9/03


if [ $# -ne 1 ]; then
	echo "Usage: stm_occup [V]"
	exit 1
else
	bias_voltage=$1
	shift 1
fi

outfile=OUTCAR

# Fermi-level and position in file
tmp=$(grep -n "E-fermi" $outfile | tail -1)
ef_line=$(echo $tmp | sed 's/://g' | awk '{print $1}')
ef=$(echo $tmp | sed 's/://g' | awk '{print $3}')

# Upper bound is the Fermi level, lower bound is Ef-V
ebottom=$(echo $ef $bias_voltage | awk '{printf "%12.8f \n", $1+(-1)*$2}')
echo $ef $ebottom
#exit 0

# Image empty states
nl=$(echo $bias_voltage | awk '$1<0 {print "E"}'|grep "E" | wc -l | awk '{print $1}')
if [ $nl -eq 1 ]; then
    etmp=$ebottom
    ebottom=$ef
    ef=$etmp
fi


# Number of bands and k-points
tmp=$(grep "NBANDS" $outfile)
nk=$(echo $tmp | awk '{print $4}')
nb=$(echo $tmp | awk '{print $9	}')
echo "Number of K-points:" $nk "Number of bands:" $nb

# Number of spins
tmp=$(grep "ISPIN" $outfile)
nspin=$(echo $tmp | awk '{print $3}')
echo "Number of spins:" $nspin

##################################################################################
# Loop over the k-points and bands
start1=$((ef_line+2))
if [ $nspin -eq 2 ]; then
    start1=$((start1+2))
fi
end1=$((start1+nb+2))
ik=1
start=$start1
end=$end1
#echo $start $end
while [ $ik -le $nk ]; do
	sed -n ''$((start+3))','$end' p' $outfile >| tmp.dump.$ik # write all the band energies to file
	ik=$((ik+1))
	start=$((end+1))
	end=$((start+nb+2))
done

# Do the second spin in the case of polarized calculations
if [ $nspin -eq 2 ]; then
    start=$((end1+3))
    end=$((start+nb+2))
    ik=1
    while [ $ik -le $nk ]; do
	sed -n ''$((start+3))','$end' p' $outfile >| tmp.dump.$ik.2
	ik=$((ik+1))
	start=$((end+1))
	end=$((start+nb+2))
    done
fi
##################################################################################

# Now we loop over the k-points and spin and search through the eigenvalues in the 
# requested range. Those define the lower and upper bounds of bands which should
# be occupied to compute the integrated LDOS in that range of bias voltage.
set -A nbbottom 0
set -A nbtop 0
ik=1
while [ $ik -le $nk ]; do
    awk '$2>='"$ebottom"' && $2<= '"$ef"' {print $1, $2, $3}' tmp.dump.$ik >| tmp.dump
    nbbottom[$ik]=$(head -1 tmp.dump | awk '{print $1}')
    nbtop[$ik]=$(tail -1 tmp.dump | awk '{print $1}')
    ik=$((ik+1))
done
if [ $nspin -eq 2 ]; then
    set -A nbbottom2 0
    set -A nbtop2 0 
    ik=1
    while [ $ik -le $nk ]; do
	awk '$2>='"$ebottom"' && $2<= '"$ef"' {print $1, $2, $3}' tmp.dump.$ik.2 >| tmp.dump
	nbbottom2[$ik]=$(head -1 tmp.dump | awk '{print $1}')
	nbtop2[$ik]=$(tail -1 tmp.dump | awk '{print $1}')
	ik=$((ik+1))
    done
fi


# Now we construct the strings for the FERWE and FERDO tokens in the vasp input file.
# The loop over bands is the fastest, followed by k-points
ik=1
string=''
while [ $ik -le $nk ]; do
	first=${nbbottom[$ik]}
	last=${nbtop[$ik]}
	cnt1=$((first-1))
	cnt3=$((nb-$last))
	cnt2=$(echo $last $first | awk '{print $1-$2+1}')
	string=$(echo $string $cnt1*0.0 $cnt2*1.0 $cnt3*0.0)
	ik=$((ik+1))
done


if [ $nspin -eq 2 ]; then
    string2=''
    ik=1
    while [ $ik -le $nk ]; do
	first=${nbbottom2[$ik]}
	last=${nbtop2[$ik]}
	cnt1=$((first-1))
	cnt3=$((nb-$last))
	cnt2=$(echo $last $first | awk '{print $1-$2+1}')
	string2=$(echo $string2 $cnt1*0.0 $cnt2*1.0 $cnt3*0.0)
	ik=$((ik+1))
    done
fi



# Parameters for the INCAR file
echo
echo "You should comment out ISMEAR parameter in the INCAR file and append the following:"
echo 
echo
echo "#################################################################################"
echo "# Compute the integrated LDOS(E) b/n " $ebottom "and" $evbm "for STM Simulations:"
echo "ISMEAR = -2"
echo "NELM = 0"
echo "LWAVE = .FALSE."
echo "FERWE =" $string
echo "FERDO =" $string2
echo "##################################################################################"

# Clean up
rm -f tmp.dump tmp.dump.*

exit 0