#!/bin/bash
# tool for interfacing between pydefect and sxdefectalign

shopt -s extglob
########################################
# help
########################################
Help()
{
   # Display Help
   echo "Automated sxdefectalign analysis of rad-induced defects using pydefect directory structure."
   echo
   echo "Syntax: bash ~/work/bin/radDefects/radDefects/auto_sxda [-d|c|p|h]"
   echo "options:"
   echo "d     Defect categories to analyze: defect complexes and/or single intrinsic"
   echo "c     Charge offset between defect pairs"
   echo "p     Plotting program choice or no plotting: gnuplot or none"
   echo "h     Print this Help."
   echo
}

usage() { echo "Usage: $0 [-h] [-d <s|f|sf>] [-c <int>] [-p <gnuplot|none>]" 1>&2; exit 1; }

analyze_defects="sf"
chg_offset=0
plot_prog="gnuplot"

# get the options
while getopts ":h:d:c:p:" o; do
  case "${o}" in
    h) # display Help
      Help
      exit;;
    d) # Enter defect types to analyze
      analyze_defects=${OPTARG}
      ;;
    c) # Enter charge offset between defects in pair
      chg_offset=${OPTARG}
      ;;
    p) # Enter plotting program
      plot_prog=${OPTARG}
      ((plot_prog == "gnuplot" || plot_prog == "none")) || usage
      ;;
    \?) # Invalid option
      echo "Error: Invalid option"
      usage
      exit;;
  esac
done
shift $((OPTIND-1))

########################################
# set variables for directories
########################################
base_path="${PWD}"
SCRIPT_DIR="$(dirname "$(readlink -f "$0")")"

# pydefect directories
unitcell_path="${base_path}/unitcell/"
uc_opt_path="${unitcell_path}structure_opt/"
uc_band_path="${unitcell_path}band/"
uc_dos_path="${unitcell_path}dos/"
uc_dielec_path="${unitcell_path}dielectric/"

cpd_path="${base_path}/cpd/"

defect_path="${base_path}/defect/"
perfect_path="${defect_path}perfect/"

########################################
# unit conversion and utility funcs
########################################
# convert eV to Rydberg
eV2Ry () {
  echo "scale=10; ${1}*0.0734985857" | bc -l
}

# convert Angstroms to Bohr
Ang2Bohr () {
  echo "scale=10; ${1}*1.8897259886" | bc -l
}

# join array by delimiter
# https://stackoverflow.com/a/17841619/24613079
function join_by { local IFS="$1"; shift; echo "$*"; }

########################################
# sxdefectalign vasp examples
########################################
# sxdefectalign --ecut 30 --charge -2 --eps 8.9 --center 0.5,0.5,0.5 --relative --vdef vacancy/LOCPOT --vref bulk/LOCPOT --vasp
# xmgrace-nxy vline-eV-a0.dat
# sxdefectalign --ecut 30 --charge -2 --eps 8.9 --center 0.5,0.5,0.5 --relative --vdef vacancy/LOCPOT --vref bulk/LOCPOT --vasp -C -0.12

# potential terms sxdefectalign = pydefect
# V(long-range) = V_{PC,q}
# V(defect)-V(ref) = V_{q/b}
# V(defect)-V(ref)-V(long-range) = dV_{PC,q/b}

########################################
# perform sxdefectalign analysis
########################################
# get lattice parameters from optimized unitcell CONTCAR to determine averaging
declare -a x1_arr=($(sed -n "3p" "${uc_opt_path}CONTCAR"))
declare -a x2_arr=($(sed -n "4p" "${uc_opt_path}CONTCAR"))
declare -a x3_arr=($(sed -n "5p" "${uc_opt_path}CONTCAR"))

# calculate lattice parameters from vectors
a=$(bc -l <<< "e( (1/2) * l((${x1_arr[0]}^2) + (${x1_arr[1]}^2) + (${x1_arr[2]}^2)) )")
b=$(bc -l <<< "e( (1/2) * l((${x2_arr[0]}^2) + (${x2_arr[1]}^2) + (${x2_arr[2]}^2)) )")
c=$(bc -l <<< "e( (1/2) * l((${x3_arr[0]}^2) + (${x3_arr[1]}^2) + (${x3_arr[2]}^2)) )")

# convert lattice parameters to Bohr
a_bohr=$(Ang2Bohr ${a})
b_bohr=$(Ang2Bohr ${b})
c_bohr=$(Ang2Bohr ${c})

echo "Lattice parameters (Ang.): ${a}, ${b}, ${c}"
echo "Lattice parameters (Bohr): ${a_bohr}, ${b_bohr}, ${c_bohr}"

# dielectric tensor notation: eps_xx,eps_xy,eps_xz,eps_xy,eps_yy,eps_yz,eps_xz,eps_yz,eps_zz
# gather dielectric tensor from unitcell.yaml
declare -a eps_ele=($(niet ".ele_dielectric_const" "${unitcell_path}unitcell.yaml" -f newline | grep -Eo '[+-]?[0-9]+([.][0-9]+)?')) || exit
declare -a eps_ion=($(niet ".ion_dielectric_const" "${unitcell_path}unitcell.yaml" -f newline | grep -Eo '[+-]?[0-9]+([.][0-9]+)?')) || exit
declare -a eps_tensor=()

for i in "${!eps_ele[@]}"; do
  eps_tensor+=($(bc -l <<< "${eps_ele[i]}+${eps_ion[i]}"))
done

eps_tensor_str=$(join_by , ${eps_tensor[@]})
echo "Dielectric tensor: ${eps_tensor_str}"

for def_dir in ${defect_path}*_*/; do
  # species from POSCAR
  declare -a element_arr=($(sed '6q;d' "${def_dir}POSCAR"))
  
  # defect_site_charge format (e.g., Ga_i1_1)
  defect=$(basename ${def_dir})
  
  # split overall defect string into type (e.g., Ga_i1) and charge (e.g., 1)
  defect_type=$(echo "${defect}" | rev | cut -d"_" -f2-  | rev)
  defect_chg=$(echo "${defect}" | rev | cut -d"_" -f1  | rev)
  
  # split defect type into atom (e.g., Ga) and site (e.g., i1)
  defect_atom=$(echo "${defect_type}" | rev | cut -d"_" -f2-  | rev)
  defect_site=$(echo "${defect_type}" | rev | cut -d"_" -f1  | rev)
  
  # convert defect charge into difference in electron number
  ele_diff=$(bc -l <<< "-1*${defect_chg}")
  
  # strings for plot formatting
  defect_file_str="${defect}"
  defect_title_str="${defect_atom}_{${defect_site}, q=${defect_chg}}"
  
  echo "${defect_type} ${defect_chg}"
  
  # remove out_file if it already exists
  [ -f "${def_dir}/${defect_file_str}.sxda" ] && rm "${def_dir}/${defect_file_str}.sxda"
  echo "${defect_type}, q = ${defect_chg}" >> "${def_dir}/${defect_file_str}.sxda"
  
  # get and convert ENCUT from INCAR file
  encut_ev=$(grep ENCUT "${def_dir}INCAR" | grep -Eo '[0-9]+([.][0-9]+)?')
  encut_ry=$(eV2Ry ${encut_ev})
  
  echo "ENCUT = ${encut_ry} Ry"
  echo "ENCUT = ${encut_ry} Ry" >> "${def_dir}/${defect_file_str}.sxda"
  
  # get alignment term from defect_energy_info.yaml and scale by charge
  alignment_C=$(niet ".energy_corrections" "${def_dir}defect_energy_info.yaml" -f newline)
  alignment_C=${alignment_C//@(*: |\}*)/}
  if [[ ${defect_chg} -ne 0 ]]; then
    alignment_C=$(echo "scale=18; ${alignment_C} / ${defect_chg#-}" | bc -l)
  fi
  
  echo "C = ${alignment_C} eV"
  echo "C = ${alignment_C} eV" >> "${def_dir}/${defect_file_str}.sxda"
  
  # defect center notation: x,y,z
  # gather defect center position from defect_entry.json
  declare -a defect_center=($(niet ".defect_center" "${def_dir}defect_entry.json" -f newline | grep -Eo '[+-]?[0-9]+([.][0-9]+)?')) || exit
  defect_center_str=$(join_by , ${defect_center[@]})
  
  # should determine vacancy and interstitial positions for each charge and add to defect_entry.json
  # then can pull positions similar to defect center for use in sxda
  # use of EIGENVAL for default charge offset?
  
  # Ga_FP1 pos
  # declare -a vac_pos=(0.4666667362384516 0.3333333110428143 0.3330453429404585)
  # VV_GaN pos
  declare -a vac_pos=(0.1666666667000030  0.0833333332999970  0.2495566444000019)
  vac_pos_str=$(join_by , ${vac_pos[@]})
  
  # Ga_FP1 pos
  # declare -a int_pos=(0.2171215306571241 0.5828821174423098 0.7366903974866901)
  # VV_GaN pos
  declare -a int_pos=(0.1666666667000030  0.0833333332999970  0.4379433555999981)
  int_pos_str=$(join_by , ${int_pos[@]})
  
  echo "Defect center = (${defect_center_str})"
  echo "Vacancy position = (${vac_pos_str})"
  echo "Interstitial position = (${int_pos_str})"
  
  ########################################
  # perform sxdefectalign analysis
  ########################################
  cd ${def_dir}
  
  ########################################
  # Defect complexes
  # [[ ${def_dir} == *"FP"* || ${def_dir} == *"VV"* ]]; then
  ########################################
  if [[ ${def_dir} == *"VV"* ]]; then
    # check if complexes are selected for analysis
    if [[ "${analyze_defects}" != *"f"* ]]; then
      continue
    fi
    
    # positively charged complexes
    if [[ ${defect_chg} -gt 0 ]]; then
      echo "positive charge at interstitial"
      # create charge offset between interstitial and vacancy sites
      int_ele_diff=$(bc -l <<< "${ele_diff} - ${chg_offset}")
      vac_ele_diff=$(bc -l <<< "${chg_offset}")
      
      # convert defect charge into difference in electron number
      int_chg=$(printf "%.2f\n" $(bc -l <<< "-1*${int_ele_diff}"))
      vac_chg=$(printf "%.2f\n" $(bc -l <<< "-1*${vac_ele_diff}"))
    
      # reset plot name and title variables for FPs
      defect_file_str="${defect}_qv${vac_chg}qi${int_chg}"
      defect_title_str="${defect_atom}_{${defect_site}, q=${defect_chg}, qv=${vac_chg}, qi=${int_chg}}"
      mv -f "${defect}.sxda" "${defect_file_str}.sxda"
      
      sxdefectalign --ecut ${encut_ry} --pos ${int_pos_str} --charge ${int_ele_diff} --pos ${vac_pos_str} --charge ${vac_ele_diff} --tensor ${eps_tensor_str} --relative --vdef "${def_dir}LOCPOT" --vref "${perfect_path}LOCPOT" --vasp --format=matrix -C ${alignment_C} 2>&1 | tee -a "${defect_file_str}.sxda" || exit

    # negatively charged complexes
    elif [[ ${defect_chg} -lt 0 ]]; then
      echo "negative charge at vacancy"
      # create charge offset between interstitial and vacancy sites
      int_ele_diff=$(bc -l <<< "-1 * ${chg_offset}")
      vac_ele_diff=$(bc -l <<< "${ele_diff} + ${chg_offset}")
      
      # convert defect charge into difference in electron number
      int_chg=$(printf "%.2f\n" $(bc -l <<< "-1*${int_ele_diff}"))
      vac_chg=$(printf "%.2f\n" $(bc -l <<< "-1*${vac_ele_diff}"))
    
      # reset plot name and title variables for FPs
      defect_file_str="${defect}_qv${vac_chg}qi${int_chg}"
      defect_title_str="${defect_atom}_{${defect_site}, q=${defect_chg}, qv=${vac_chg}, qi=${int_chg}}"
      mv -f "${defect}.sxda" "${defect_file_str}.sxda"
      
      sxdefectalign --ecut ${encut_ry} --pos ${vac_pos_str} --charge ${vac_ele_diff} --pos ${int_pos_str} --charge ${int_ele_diff} --tensor ${eps_tensor_str} --relative --vdef "${def_dir}LOCPOT" --vref "${perfect_path}LOCPOT" --vasp --format=matrix -C ${alignment_C} 2>&1 | tee -a "${defect_file_str}.sxda" || exit

    # neutral complexes
    else
      echo "neutral charge"
      # create charge offset between interstitial and vacancy sites
      int_ele_diff=$(bc -l <<< "-1 * ( ${chg_offset} / 2 )")
      vac_ele_diff=$(bc -l <<< "${chg_offset} / 2")
      
      # convert defect charge into difference in electron number
      int_chg=$(printf "%.2f\n" $(bc -l <<< "-1*${int_ele_diff}"))
      vac_chg=$(printf "%.2f\n" $(bc -l <<< "-1*${vac_ele_diff}"))
    
      # reset plot name and title variables for FPs
      defect_file_str="${defect}_qv${vac_chg}qi${int_chg}"
      defect_title_str="${defect_atom}_{${defect_site}, q=${defect_chg}, qv=${vac_chg}, qi=${int_chg}}"
      mv -f "${defect}.sxda" "${defect_file_str}.sxda"
      
      sxdefectalign --ecut ${encut_ry} --pos ${int_pos_str} --charge ${int_ele_diff} --pos ${vac_pos_str} --charge ${vac_ele_diff} --tensor ${eps_tensor_str} --relative --vdef "${def_dir}LOCPOT" --vref "${perfect_path}LOCPOT" --vasp --format=matrix -C ${alignment_C} 2>&1 | tee -a "${defect_file_str}.sxda" || exit
    fi
    
    # previous treatment based on defect center
    # sxdefectalign --ecut ${encut_ry} --center ${defect_center_str} --charge ${ele_diff} --tensor ${eps_tensor_str} --average ${c_bohr} --relative --vdef "${def_dir}LOCPOT" --vref "${perfect_path}LOCPOT" --vasp --format=matrix -C ${alignment_C} 2>&1 | tee -a "${defect}.sxda" || exit

  ########################################
  # intrinsic point defects
  ########################################
  else
    # check if single point defects are selected for analysis
    if [[ "${analyze_defects}" != *"s"* ]]; then
      continue
    fi
    
    sxdefectalign --ecut ${encut_ry} --center ${defect_center_str} --charge ${ele_diff} --tensor ${eps_tensor_str} --average ${c_bohr} --relative --vdef "${def_dir}LOCPOT" --vref "${perfect_path}LOCPOT" --vasp --format=matrix -C ${alignment_C} 2>&1 | tee -a "${defect}.sxda" || exit
  fi
  
  ########################################
  # perform plotting using gnuplot
  ########################################
  if [ "${plot_prog}" = "gnuplot" ]; then
    # vline-eV-a{0,1,2}.dat matrix format
    # z V(long-range) V(defect)-V(ref) V(defect)-V(ref)-V(long-range)
    gnuplot> call -c "${SCRIPT_DIR}/plnr_avgs.p" ${element_arr[@]} "${defect_file_str}" "${defect_title_str}"
  
    # vAtoms.dat format
    # r V(long-range) V(defect)-V(ref) V(defect)-V(ref)-V(long-range) x y z ...
    # empty lines between different species (need to include in plotting now)
    # replace empty lines with double empty lines to enable gnuplot index use
    sed -i -e 's/^$/\n/' vAtoms.dat
    gnuplot> call -c "${SCRIPT_DIR}/atomic_sph_avg.p" ${element_arr[@]} "${defect_file_str}" "${defect_title_str}"
  fi
  
  cd ../../
done
