#!/usr/bin/env bash

# Author: Daniel S. Karls, University of Minnesota (karl0100 |AT| umn DOT edu)
# Date: 11/8/2018

# This example Test Driver uses LAMMPS to compute a cohesive energy versus lattice constant curve for a given cubic
# lattice (fcc, bcc, sc, diamond) of a single given species.  The curve is computed for lattice constants ranging from
# a_min_frac*a_0 to a_max_frac*a_0, where a_0, a_min_frac, and a_max_frac are specified via stdin.
# The parameter a_0 is typically approximately equal to the equilibrium lattice constant for the Model/species/lattice type being paired.
# A logarithmic scale is used such that most lattice spacings are about a_0. The precise scaling of and number of sample points going
# from a_min to a_0 and from a_0 to a_max is specified by two separate parameters passed from stdin.
# Please see README.txt for more details.

# Define function which prints to stderr
echoerr() { echo "$@" 1>&2; }

# Read the KIM Model name from stdin (this is passed through pipeline.stdin.tpl using @< MODELNAME >@, which the pipeline
# will automatically fill in once a compatible Model is found).
# Also pass the species, atomic mass (in g/mol), type of cubic lattice (bcc, fcc, sc, or diamond), a_0, a_min_frac, a_max_frac,
# number of sample spacings between a_min (= a_min_frac*a_0) and a_0, number of sample spacings between a_0 and a_max
# (= a_max_frac*a_0), and the two parameters governing the distribution of sample spacings around a_0 compared to a_min/a_max
# respectively.  Please see README.txt for more details on these parameters and how they are used.
echo "Please enter a valid KIM Model extended-ID:"
read modelname
echo "Please enter the species symbol (e.g. Si, Au, Al, etc.):"
read element
echo "Please enter the atomic mass of the species (g/mol):"
read mass
echo "Please enter the lattice type (bcc, fcc, sc, or diamond):"
read latticetypeinput
echo "Please specify a lattice constant (referred to as a_0 below) in meters about which the energy will be computed (This will usually be the equilibrium lattice constant.\
  Most of the volumes sampled will be about this lattice constant.):"
read a_0
echo "Please specify the smallest lattice spacing (referred to as a_min below) at which to compute the energy, expressed as a fraction of a_0 (for example, if you wish for\
 a_min to be equal to 0.8*a_0, please specify 0.8 for this value):"
read a_min_frac
echo "Please specify the largest lattice spacing (referred to as a_max below) at which to compute the energy, expressed as a multiple of a_0 (for example, if you wish for\
 a_max to be equal to 1.5*a_0, please specify 1.5 for this value):"
read a_max_frac
echo "Please enter the number of sample lattice spacings to compute which are >= a_min and < a_0 (one of these sample lattice spacings will be equal to a_min):"
read N_lower
echo "Please enter the number of sample lattice spacings to compute which are > a_0 and <= a_max (one of these sample lattice spacings will be equal to a_max):"
read N_upper
echo "Please enter a value of the lower sample spacing parameter (see README.txt for more details):"
read samplespacing_lower
echo "Please enter a value of the upper sample spacing parameter (see README.txt for more details):"
read samplespacing_upper

# Check that element string read in contains no spaces
if [[ "$element" =~ \  ]] ; then
    echo "Error: a space was detected in the element inputted. Please note that this Test supports only a single species. Exiting..."
    echoerr "Error: a space was detected in the element inputted. Please note that this Test supports only a single species. Exiting..."
    exit 1
fi

# Check that a_0 is numerical and strictly positive
if ! [[ "$a_0" =~ ^[0-9e.-]+ ]] ; then
    if [[ "${a_0}" == "[]" ]] ; then
        echo "Error: a_0 read in is empty. If using a query, check that it returns a non-empty value. Exiting..."
        echoerr "Error: a_0 read in is empty. If using a query, check that it returns a non-empty value. Exiting..."
        exit 1
    else
        echo "Error: a_0 read in is not numerical. Check pipeline.stdin for errors. Exiting..."
        echoerr "Error: a_0 read in is not numerical. Check pipeline.stdin for errors. Exiting..."
        exit 1
    fi
fi

a_0check=`echo $a_0 | awk '{if($1 <= 0.0) print "Not positive"}'`
if [ "$a_0check" == "Not positive" ]; then
    echo "Error: a_0 read in must be a positive number.  Exiting..."
    echoerr "Error: a_0 read in must be a positive number.  Exiting..."
    exit 1
fi

# Convert a_0 from meters to angstroms
a_0=`echo ${a_0} | awk '{print $1*1e10}'`

# Check that a_min_frac entered is positive and strictly less than 1
a_min_fraccheck=`echo $a_min_frac | awk '{if($1 > 0.0 && $1 < 1.0) print "a_min_frac OK"}'`
if [ "$a_min_fraccheck" != "a_min_frac OK" ]; then
    echo "Error: a_min_frac must be in the range (0,1)."
    echoerr "Error: a_min_frac must be in the range (0,1)."
    exit 1
else
    a_min=`echo $a_min_frac $a_0 | awk '{print $1*$2}'`
fi

# Check that a_min_frac entered is greater than 1
a_max_fraccheck=`echo $a_max_frac | awk '{if($1 > 1.0) print "a_max_frac OK"}'`
if [ "$a_max_fraccheck" != "a_max_frac OK" ]; then
    echo "Error: a_max_frac must be strictly greater than 1."
    echoerr "Error: a_max_frac must be strictly greater than 1."
    exit 1
else
    a_max=`echo $a_max_frac $a_0 | awk '{print $1*$2}'`
fi

# Check that the number of spacings are positive
N_lowercheck=`echo $N_lower | awk '{if($1 <= 0) print "Not positive"}'`
if [ "$N_lowercheck" == "Not positive" ]; then
    echo "Error: N_lower read in must be a positive number.  Exiting..."
    echoerr "Error: N_lower read in must be a positive number.  Exiting..."
    exit 1
fi

N_uppercheck=`echo $N_upper | awk '{if($1 <= 0) print "Not positive"}'`
if [ "$N_uppercheck" == "Not positive" ]; then
    echo "Error: N_upper read in must be a positive number.  Exiting..."
    echoerr "Error: N_upper read in must be a positive number.  Exiting..."
    exit 1
fi

# Check that samplespacing parameters are > 1
spacingparamcheck=`echo $samplespacing_lower $samplespacing_upper | awk '{if($1 <= 1.0 && $2 <=1.0) print 1; else if($1 <= 1.0 && $2 > 1.0) print 2; else if($1 > 1.0 && $2 <= 1.0) print 3; else print 4}'`
if [ "$spacingparamcheck" == 1 ]; then
    echo "Error: lower and upper sample spacing parameters must both be strictly greater than 1."
    echoerr "Error: lower and upper sample spacing parameters must both be strictly greater than 1."
    exit 1
elif [ "$spacingparamcheck" == 2 ]; then
    echo "Error: lower sample spacing parameter must be strictly greater than 1.  Exiting."
    echoerr "Error: lower sample spacing parameter must be strictly greater than 1.  Exiting."
    exit 1
elif [ "$spacingparamcheck" == 3 ]; then
    echo "Error: upper sample spacing parameter must be strictly greater than 1.  Exiting."
    echoerr "Error: upper sample spacing parameter must be strictly greater than 1.  Exiting."
    exit 1
fi

# Identify which of the cubic lattice types (bcc,fcc,sc,diamond) the user entered (case-insensitive).
if [ `echo $latticetypeinput | tr [:upper:] [:lower:]` = `echo bcc | tr [:upper:] [:lower:]`  ]; then
    latticetype="bcc"
    space_group="Im-3m"
    wyckoffcode="2a"
    basisatomcoords="[   0    0    0 ]\n      [ 0.5  0.5  0.5 ]"
    specieslist="\"${element}\"\n      \"${element}\""
elif [ `echo $latticetypeinput | tr [:upper:] [:lower:]` = `echo fcc | tr [:upper:] [:lower:]` ]; then
    latticetype="fcc"
    space_group="Fm-3m"
    wyckoffcode="4a"
    basisatomcoords="[   0    0    0 ]\n      [   0  0.5  0.5 ]\n      [ 0.5    0  0.5 ]\n      [ 0.5  0.5    0 ]"
    specieslist="\"${element}\"\n      \"${element}\"\n      \"${element}\"\n      \"${element}\""
elif [ `echo $latticetypeinput | tr [:upper:] [:lower:]` = `echo sc | tr [:upper:] [:lower:]` ]; then
    latticetype="sc"
    space_group="Pm-3m"
    wyckoffcode="1a"
    basisatomcoords="[ 0 0 0 ]"
    specieslist="\"${element}\""
elif [ `echo $latticetypeinput | tr [:upper:] [:lower:]` = `echo diamond | tr [:upper:] [:lower:]` ]; then
    latticetype="diamond"
    space_group="Fd-3m"
    wyckoffcode="8a"
    basisatomcoords="[    0     0     0 ]\n      [    0   0.5   0.5 ]\n      [  0.5   0.5     0 ]\n      [  0.5     0   0.5 ]\n      [ 0.75  0.25  0.75 ]\n      [ 0.25  0.25  0.25 ]\n      [ 0.25  0.75  0.75 ]\n      [ 0.75  0.75  0.25 ]"
    specieslist="\"${element}\"\n      \"${element}\"\n      \"${element}\"\n      \"${element}\"\n      \"${element}\"\n      \"${element}\"\n      \"${element}\"\n      \"${element}\""
else
    echo "Error: This Test supports only cubic lattices (specified by 'bcc', 'fcc', 'sc', or 'diamond'). Exiting..."
    echoerr "Error: This Test supports only cubic lattices (specified by 'bcc', 'fcc', 'sc', or 'diamond'). Exiting..."
    exit 1
fi

# Define the lattice spacings at which the energy will be computed.  See README.txt for more details.
latticeconst=`echo $a_0 $a_min $a_max $N_lower $N_upper $samplespacing_lower $samplespacing_upper |  awk '{for (i=0;i<=$5-1;++i){printf "%f ",$1+($3-$1)*(1-log(1+i*($7-1)/$5)/log($7))}}{for (i=$4;i>=0;--i){printf "%f ",$2+($1-$2)*log(1+i*($6-1)/$4)/log($6)}}'`
read -a lattice_const <<< "$latticeconst"

numberofattemptedspacings=`expr $N_lower + $N_upper + 1`

# Replace placeholder strings in the lammp.in.template input file script template.  The resulting
# lammps input file (lammps.in) will be stored in the Test Result folder (which may be referenced
# as the 'output' directory).
thisdir=`dirname "$0"` # Directory of this Test Driver executable
sed -e "s/sed_model_string/"$modelname"/"                                   \
    -e "s/sed_species_string/$element/"                                     \
    -e "s/sed_mass_string/$mass/"                                           \
    -e "s/sed_latticetype_string/$latticetype/"                             \
    -e "s/sed_numberofattemptedspacings_string/$numberofattemptedspacings/" \
    -e "s/sed_latticeconst_string/$latticeconst/"                           \
    ""$thisdir"/lammps.in.template" > output/lammps.in

# Run LAMMPS using the lammps.in input file and write to lammps.log
lammps -in output/lammps.in > output/lammps.log

# Parse LAMMPS output log and extract the cohesive energies corresponding to each lattice spacing into an array.  Lattice spacings
# which produced an error because they were too small will be ignored by the `grep` process.
read -a cohesive_energy <<< `grep "Cohesive energy = [0-9.e-]* eV/atom" output/lammps.log | cut -d' ' -f4 | sed ':a;N;$!ba;s/\n/ /g'`
numberofreturnedspacings=${#cohesive_energy[@]}

numberofreturnedspacings_check=`echo $numberofreturnedspacings | awk '{if($1 == 0) print "No cohesive energies returned"} {}'`
if [ "$numberofreturnedspacings_check" == "No cohesive energies returned" ]; then
    echo "Error: No cohesive energies were able to be parsed from the LAMMPS output log.  Check the LAMMPS output log for errors. Exiting..."
    echoerr "Error: No cohesive energies were able to be parsed from the LAMMPS output log.  Check the LAMMPS output log for errors. Exiting..."
    exit 1
fi

for ((i=1; i<=$numberofreturnedspacings;++i)); do
    j=`expr $i - 1`
    latconstarray="$latconstarray ${lattice_const[$j]} "
done

for ((i=1; i<=$numberofreturnedspacings;++i)); do
    j=`expr $i - 1`
    # Check to see that the cohesive energies parsed from LAMMPS are actually numbers (in case there was a LAMMPS error of some sort)
    if ! [[ "${cohesive_energy[$j]}" =~ ^[0-9e.-]+ ]] ; then
        echo "Error: Cohesive energies parsed from LAMMPS are not numerical.  Check the LAMMPS log for errors.  Exiting..."
        echoerr "Error: Cohesive energies parsed from LAMMPS are not numerical.  Check the LAMMPS log for errors.  Exiting..."
        exit 1
    fi
    ecoh=`echo ${cohesive_energy[$j]} | awk '{print $1*(-1)}'`
    ecoharray="$ecoharray ${ecoh} "
done

# Replace the placeholders in the EDN results template file (results.edn.tpl) with results
sed -e "s/_LATTICETYPE_/${latticetype}/"         \
    -e "s/_SPECIES_/${specieslist}/"             \
    -e "s/_LATCONSTARRAY_/${latconstarray}/"     \
    -e "s/_BASISATOMCOORDS_/${basisatomcoords}/" \
    -e "s/_ECOHARRAY_/${ecoharray}/"             \
    ""$thisdir"/results.edn.tpl" >  output/results.edn