#!/bin/bash

#################################################################
#   Copyright 2016-2024, Johns Hopkins University               #
#                                                               #
#   Permission is hereby granted, free of charge,               #
#   to any person obtaining a copy of this software             #
#   and associated documentation files (the "Software"),        #
#   to deal in the Software without restriction, including      #
#   without limitation the rights to use, copy, modify,         #
#   merge, publish, distribute, sublicense, and/or sell         #
#   copies of the Software, and to permit persons to            #
#   whom the Software is furnished to do so,                    #
#   subject to the following conditions:                        #
#                                                               #
#   The above copyright notice and this permission              #
#   notice shall be included in all copies or                   #
#   substantial portions of the Software.                       #
#                                                               #
#   THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY          #
#   OF ANY KIND, EXPRESS OR IMPLIED, INCLUDING 	                #
#   BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,       #
#   FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT.       #
#   IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT 	                #
#   HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES                    #
#   OR OTHER LIABILITY, WHETHER IN AN ACTION                    #
#   OF CONTRACT, TORT OR OTHERWISE, ARISING                     #
#   FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE             #
#   OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.               #
#################################################################

#################################################################
#   Generate a k-point grid by calling the servlet              #
#   This script is assumes that the user is calling             #
#   it from the directory containing the necessary files.       #
#   (although the script can be placed some place else)	        #
#                                                               #
#   Questions, comments, feedback, saw some bad coding?         #
#   Email us at kpoints@jhu.edu                                 #
#   by Pandu Wisesa	                                            #
#################################################################

version="C2023.04.28" 
echo "Running getKPoints script version" $version".";

# This switch controls the information regarding grid generation output file.
# The default is "TRUE", replacing it with "FALSE" or any other values will turn it off.
list_message="TRUE"

# These format the messages accordingly
warning_message()
{
  echo "*** WARNING: $1 ***"
}

error_message()
{
  echo "*** ERROR: $1 ***" 
}

# Check for early kill signals.
killed()
{
  error_message "Command was terminated by user."
}
trap 'killed' 1 2 3

failed_to_connect()
{
  error_message "There is a problem connecting to the database."
}
trap 'failed_to_connect' 6 7 10 27


# Check for format of the software package, right now only VASP is supported.
found_input_files="FALSE"
determine_input_format()
{
  input="$1"
  if [ "$input" == "POSCAR" ]; then
    echo "-vasp"
  fi 
}

# Generate a KPOINTS file for VASP.
# Due to how VASP filenames are, the files can be fetched directly by their names.
generate_grid_for_vasp()
{
  found_input_files="TRUE"
  if [ -f INCAR ]; then
    KPTS=$(curl -s https://muellergroup.jhu.edu/kpts/PreCalcServer/PreCalcServlet?format=vasp\&messagelist="$list_message"\&clientversion=$version --form "fileupload=@PRECALC" --form "fileupload=@POSCAR" --form "fileupload=@INCAR")
  else
    warning_message "No INCAR file detected. Using only symmetry information from structure file."
    KPTS=$(curl -s https://muellergroup.jhu.edu/kpts/PreCalcServer/PreCalcServlet?format=vasp\&messagelist="$list_message"\&clientversion=$version --form "fileupload=@PRECALC" --form "fileupload=@POSCAR")
  fi
  STAT=$(echo $?)
  if [ "$STAT" -ne 0 ]; then
    error_message "Can't connect to the database."
    exit 7
  else
    if [ "$list_message" == "TRUE" ]; then
      HEAD=$(echo "$KPTS" | sed '/+++SEPARATE_HERE+++/,$d')
      GRID=$(echo "$KPTS" | sed '/+++SEPARATE_HERE+++/,$!d' | sed 1d)
      OUT_OK="TRUE"
      if echo "$HEAD" | grep --quiet "ERROR"; then
        OUT_OK="FALSE"
      fi
      echo "$HEAD"
      if [ "$OUT_OK" == "TRUE" ]; then
        echo "$GRID" > KPOINTS
      fi
    else
      echo "$KPTS" > KPOINTS
    fi
  fi
  # Do a quick check of the integrity of the KPOINTS file.
  if [ -f KPOINTS ]; then
    if grep --quiet "HTTP Status" KPOINTS; then
      rm KPOINTS
      error_message "Invalid KPOINTS file was generated. Please check your input files and re-submit the request."
    fi
    if grep "ERROR" KPOINTS; then # Useful when "list_message=FALSE"
      rm KPOINTS 
    fi
  fi
}

# Check if PRECALC file exists
precalc_default="FALSE"
check_precalc()
{
  if [ -f PRECALC ]; then 
    echo "Generating grids..."
    precalc_default="FALSE"
  else # In the case PRECALC does not exists, create an empty file to feed in
    warning_message "No PRECALC file detected. Using default values."
    touch PRECALC
    # If you want to edit your own default PRECALC settings, you can uncomment and edit these lines.
    #echo "MINDISTANCE=28.1" >> PRECALC  
    #echo "INCLUDEGAMMA=FALSE" >> PRECALC
    echo "Generating grids..."
    precalc_default="TRUE"
  fi   
}

# Write output to screen
output_to_stdout() {
  if [ -e KPOINTS ]; then
    echo "=============== CONTENT OF KPOINTS ==============="
    cat KPOINTS
  fi
}

################################## main() #####################################
# These lines below are the "main method" of this script.
check_precalc
# Check if the first argument is the format or filename.
if [ "$1" == -* ]; then 
  format="$1"
  file1="$2"
else
  format=""
  file1="$1"
fi

# If no format is specified, scour through the files in this directories for input structure(s).
# In addition, generate the grids for every structures found.
if [ "$format" == "" ]; then 
  for files in *; do
    format=$(determine_input_format $files)
    if [ "$format" == "-vasp" ]; then
      generate_grid_for_vasp 
    fi
  done
# Otherwise generate the grid for that input file.
elif [ "$format" == "-vasp" ]; then
  generate_grid_for_vasp
fi

# Final cleanup
if [ "$precalc_default" == "TRUE" ]; then
  rm PRECALC
fi
if [ "$OUT_OK" == "FALSE" ] && [ -e KPOINTS ]; then
  rm KPOINTS
fi
if [ "$found_input_files" == "FALSE" ]; then
  error_message "Found no input files for grid generation! Exit."
  echo "Finished."
  exit 1
fi
echo "Finished"

# If write generating vectors, write output to screen.
if [ "$precalc_default" == "FALSE" ] && [ -e PRECALC ]; then
  write_vectors=$( grep -iE "^[ ]*WRITE_LATTICE_VECTORS" ./PRECALC 2> /dev/null | \
                   awk -F "[=#]" '{ print tolower($2) }' | tr -d [:space:] 2> /dev/null )
  if [ ! -z $write_vectors ] && [ $write_vectors == "true" ]; then
    output_to_stdout
  fi
fi
