#!/usr/bin/env bash

# Authors: Ilia Nikiforov (nikif002 |AT| umn DOT edu)
#          Daniel S. Karls (karl0100 |AT| umn DOT edu)
#          University of Minnesota
# Date: 2019-06-02

# This Test Driver computes the cohesive energy and equilibrium
# lattice constant for a 2D hexagonal lattice using Polak-Ribiere
# conjugate gradient static minimization in LAMMPS. A choice of
# the structure type, the atomic species, and an initial guess at the
# equilibrium lattice spacing are supplied by the user through
# pipeline.stdin.tpl. See the attached README for details.

# Function which outputs to stderr
echoerr() { echo "$@" 1>&2; }

# Function for calculated isolated atomic energies
function get_isolated_atom_energy {
    local model=$1
    local species=$2
    local thisdir=$3
    local isolated_atom_energy

    # Replace placeholder strings in the isolated_atom.lammps.tpl 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).
    sed -e "s/sed_model_string/"$model"/"                                       \
        -e "s/sed_species_string/$species/"                                     \
        ""$thisdir"/isolated_atom.lammps.tpl" > output/isolated_atom.lammps.in

    # Run LAMMPS using the isolated_atom.lammps.in input file and write to isolated_atom.lammps.log
    lammps -in output/isolated_atom.lammps.in > output/isolated_atom.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 isolated_atom_energy <<< `grep "Isolated atom energy: [0-9.e-]* eV" output/isolated_atom.lammps.log | cut -d' ' -f4 | sed ':a;N;$!ba;s/\n/ /g'`

    # Check to see that the energy parsed from LAMMPS is actually a number (in case there was a LAMMPS error of some sort)
    if ! [[ "${isolated_atom_energy}" =~ ^[0-9e.-]+ ]] ; then
        echo "Error: Isolated atom energy parsed from LAMMPS is not numerical.  Check the LAMMPS log for errors.  Exiting..."
        echoerr "Error: Isolated atom energy parsed from LAMMPS is not numerical.  Check the LAMMPS log for errors.  Exiting..."
        exit 1
    fi

    echo "${isolated_atom_energy}"
}

# Read the KIM Model name and initial lattice constant from pipeline.stdin.tpl
# (the former is passed using @< MODELNAME >@, which the
# pipeline will automatically fill in once a compatible Model is found).
# Also read the two species that will comprise the lattice
echo "Enter a model name:"
read modelname
echo "$modelname"
echo "Enter the choice of structure:"
echo "1: graphene-like"
echo "2: 2H TMD structure"
echo "3: 1T TMD structure"
read structure_type
echo "$structure_type"
echo "Please enter an initial lattice constant (Angstroms):"
read initial_lattice_constant
echo "$initial_lattice_constant"
echo "Please enter species 1:"
read species1
echo "$species1"
echo "Please enter species 2:"
read species2
echo "$species2"

# Determine extra info for input
if [ "$structure_type" == "1" ]; then      #graphene-like
    basis="basis 0.3333333333333333 0.6666666666666667 0.0 basis 0.6666666666666667 0.3333333333333333 0.0"
    tmd_uncomm="#"
    if [ "$species1" == "$species2" ]; then
        number_species=1
        species2=" "
        basis_atoms=" "
        number_atoms_species1=2
        number_atoms_species2=0
      else
        number_species=2
        basis_atoms="basis 1 1 basis 2 2"
        number_atoms_species1=1
        number_atoms_species2=1
    fi
elif [ "$structure_type" == "2" ]; then   #2H
    if [ "$species1" == "$species2" ]; then
        echo "Error: In TMD structures, the two species must be different.  Exiting..."
        echoerr "Error: In TMD structures, the two species must be different.  Exiting..."
        exit 1
    else
        tmd_uncomm=" "
        number_species=2
        basis_atoms="basis 1 1 basis 2 2 basis 3 2"
        basis="basis 0.0 0.0 0.005 basis 0.3333333333333333 0.6666666666666667 0.00625 basis 0.3333333333333333 0.6666666666666667 0.00375"
        number_atoms_species1=1
        number_atoms_species2=2
    fi
elif [ "$structure_type" == "3" ]; then  #1T
    if [ "$species1" == "$species2" ]; then
        echo "Error: In TMD structures, the two species must be different.  Exiting..."
        echoerr "Error: In TMD structures, the two species must be different.  Exiting..."
        exit 1
    else
        tmd_uncomm=" "
        number_species=2
        basis_atoms="basis 1 1 basis 2 2 basis 3 2"
        basis="basis 0.0 0.0 0.005 basis 0.3333333333333333 0.6666666666666667 0.00625 basis 0.6666666666666667 0.3333333333333333 0.00375"
        number_atoms_species1=1
        number_atoms_species2=2
    fi
else
    echo "Error: Invalid structure selection.  Exiting..."
    echoerr "Error: Invalid structure selection.  Exiting..."
    exit 1
fi

thisdir=`dirname "$0"` # Directory of this Test Driver executable

#
# Calculate the energy of an isolated atom of each species
#
if [ "$number_species" == 1 ]; then

  echo ""
  isolatedatomenergy="$(get_isolated_atom_energy $modelname $species1 $thisdir)"
  isolatedatomenergysum=`echo "$isolatedatomenergy" $number_atoms_species1 | awk '{print $1*$2}'`
  echo "Isolated atom energy: ${isolatedatomenergy} eV"
  echo ""

else
  # @@ TODO: Do species 1 and species 2 and then
  echo ""
  isolatedatomenergy1="$(get_isolated_atom_energy $modelname $species1 $thisdir)"
  isolatedatomenergy2="$(get_isolated_atom_energy $modelname $species2 $thisdir)"
  echo "Isolated atom energy ($species1): ${isolatedatomenergy1} eV"
  echo "Isolated atom energy ($species2): ${isolatedatomenergy2} eV"
  isolatedatomenergysum=`echo "$isolatedatomenergy1 $number_atoms_species1 $isolatedatomenergy2 $number_atoms_species2" | awk '{print $1*$2 + $3*$4}'`
  echo ""
fi

echo "Sum of isolated atom energies: ${isolatedatomenergysum} eV"
echo ""

# Replace the string 'sed_model_string' in the main.lammps.tpl input file
# script template with the name of the KIM Model being used.  Also replace
# the string 'sed_initial_lattice_constant_string' with the value supplied
# through stdin. Finally, replace the number of species and species names,
# and the LAMMPS basis command.
# The resulting  file will be stored in the Test Result folder (which may be
# referenced as the 'output' directory).
thisdir=`dirname "$0"` # The directory of this Test Driver executable
sed -e "s/sed_tmd_uncomm_string/$tmd_uncomm/"                              \
    -e "s/sed_model_string/"$modelname"/"                                  \
    -e "s/sed_initial_lattice_constant_string/$initial_lattice_constant/"  \
    -e "s/sed_number_species_string/$number_species/"                      \
    -e "s/sed_species1_string/$species1/"                                  \
    -e "s/sed_species2_string/$species2/"                                  \
    -e "s/sed_basis_atoms_string/$basis_atoms/"                            \
    -e "s/sed_basis_string/$basis/"                                        \
    -e "s/sed_isolated_atom_energy_sum/$isolatedatomenergysum/"            \
    ""$thisdir"/main.lammps.tpl" > output/lammps.in

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

# Parse the LAMMPS output log and extract the final surface tension (to
# indicate how converged it is to 0), cohesive energy, and equilibrium lattice
# constant.
finaltension=`grep "Final surface tension = [0-9.e+\-]\+ bar\*angstrom" output/lammps.log | cut -d' ' -f5`
ecohesive=`grep "Cohesive energy = [0-9.e+\-]\+ eV/atom" output/lammps.log | cut -d' ' -f4`
latticeconstant=`grep "Equilibrium lattice constant = [0-9.e+\-]\+ angstrom" output/lammps.log | cut -d' ' -f5`

if [ "$structure_type" == "2" -o "$structure_type" == "3" ]; then
    zposition=`awk "NR==$numberoflines-3" output/lammps.log | awk '{print $(NF-1)}'`
    if ! [[ $zposition =~ ^[0-9.e-]+ ]] ; then
        echo "Error: chalcogen z position parsed from LAMMPS log is not a numeric value.  Check the LAMMPS log for errors.  Exiting..."
        echoerr "Error: chalcogen z position parsed from LAMMPS log is not a numeric value.  Check the LAMMPS log for errors.  Exiting..."
        exit 1
    fi
fi


# Check that the results we obtained are actually numbers (in case there was a LAMMPS error of some sort)
if ! [[ $finaltension =~ ^[0-9.e-]+ ]] ; then
    echo "Error: Final surface tension parsed from LAMMPS log is not a numeric value.  Check the LAMMPS log for errors.  Exiting..."
    echoerr "Error: Final surface tension parsed from LAMMPS log is not a numeric value.  Check the LAMMPS log for errors.  Exiting..."
    exit 1
elif ! [[ $ecohesive =~ ^[0-9.e-]+ ]] ; then
    echo "Error: Cohesive energy parsed from LAMMPS log is not a numeric value.  Check the LAMMPS log for errors.  Exiting..."
    echoerr "Error: Cohesive energy parsed from LAMMPS log is not a numeric value.  Check the LAMMPS log for errors.  Exiting..."
    exit 1
elif ! [[ $latticeconstant =~ ^[0-9.e-]+ ]] ; then
    echo "Error: Equilibrium lattice constant parsed from LAMMPS log is not a numeric value.  Check the LAMMPS log for errors.  Exiting..."
    echoerr "Error: Equilibrium lattice constant parsed from LAMMPS log is not a numeric value.  Check the LAMMPS log for errors.  Exiting..."
    exit 1
fi

# Echo result to stdout
echo ""
echo "Optimization complete"
echo "    lattice constant: ${latticeconstant} Angstroms"
echo "     cohesive energy: ${ecohesive} eV/atom"
echo ""

# If there was only one species, we set species2 to a whitespace for the KIM
# potential specification. Set it back to what it actually is for writing
# to results.edn

if [ $number_species == 1 ]; then
    species2="$species1"
fi


# Determine extra info for output
if [ "$structure_type" == "1" ]; then
    shortname="graphene-like"
    basiscoord="      [   0.3333333333333333    0.6666666666666667    0 ]\n      [   0.6666666666666667  0.3333333333333333 0 ]"
    if [ "$species1" == "$species2" ]; then
        layergroup="p6/mmm"
        wycklett="\"2b\""
        wyckspec="\"$species1\""
        wyckcoord="      [   0.3333333333333333    0.6666666666666667    0 ]"
    else
        layergroup="p-6m2"
        wycklett="\"1b\"\n\"1c\""
        wyckspec="\"$species1\"\n\"$species2\""
        wyckcoord="      [   0.3333333333333333    0.6666666666666667    0 ]\n      [   0.6666666666666667  0.3333333333333333 0 ]"
    fi
elif [ "$structure_type" == "2" ]; then
    shortname="2H"
    basiscoord="      [   0    0    0 ]\n      [   0.3333333333333333    0.6666666666666667    $zposition ]\n      [   0.3333333333333333    0.666666666666666 -$zposition ]"
    layergroup="p-6m2"
    wycklett="\"1a\"\n\"2e\""
    wyckspec="\"$species1\"\n\"$species2\""
    wyckcoord="      [   0    0    0 ]\n      [   0.3333333333333333    0.6666666666666667    $zposition ]"
else
    shortname="1T"
    basiscoord="      [   0    0    0 ]\n      [   0.3333333333333333    0.6666666666666667    $zposition ]\n      [   0.666666666666666 0.3333333333333333    -$zposition ]"
    layergroup="p-3m1"
    wycklett="\"1a\"\n\"2c\""
    wyckspec="\"$species1\"\n\"$species2\""
    wyckcoord="      [   0    0    0 ]\n      [   0.3333333333333333    0.6666666666666667    $zposition ]"
fi

# Write results and crystallography info to results.edn

sed -e "s/_LATCONST_/${latticeconstant}/"  \
    -e "s/_ECOHESIVE_/${ecohesive}/"       \
    -e "s/_SPEC1_/${species1}/"            \
    -e "s/_SPEC2_/${species2}/"            \
    -e "s/_SHORTNAME_/${shortname}/"       \
    -e "s/_BASISCOORD_/${basiscoord}/"      \
    -e "s,_LAYERGROUP_,"${layergroup}","     \
    -e "s/_WYCKLETT_/${wycklett}/"         \
    -e "s/_WYCKSPEC_/${wyckspec}/"         \
    -e "s/_WYCKCOORD_/${wyckcoord}/"       \
    ""$thisdir"/results.edn.tpl" > output/results.edn

# Write final surface tension to file to indicate level of convergence
echo "The final surface tension corresponding to the cohesive energy and lattice constant given in results.edn is:" > output/finalpressure.txt
echo "${finaltension} bar*angstrom" >> output/finalsurfacetension.txt