#!/usr/bin/env bash

function get_next_step() {

  # get system name from path
  system=$(basename "${job}")

  # get prev_step and prev_step_num from parsing coor filenames
  last_step_coorname=$(find ${job} -name "*.coor" -type f | sort -V | tail -n 1)
  if [[ "${last_step_coorname}" == "" ]]; then
    prev_step=""
  else
    filename=$(basename ${last_step_coorname})
    filename=${filename%%.*}
    prev_step=${filename##${system}_}
    prev_step=${prev_step%%_*}
  fi
  case "${prev_step}" in
    "4-sim")
      this_step="4-sim"
      prev_step_num=${filename##*[^0-9]}
      final_step_num=$((${prev_step_num}+${step_size})) 
      ;;
    "3-heat")
      this_step="4-sim"
      prev_step_num=0
      final_step_num=${step_size} 
      ;;
    "2-min")
      this_step="3-heat"
      prev_step_num=0
      final_step_num=0 
      ;;
    "1-min")
      this_step="2-min"
      prev_step_num=0
      final_step_num=0 
      ;;
    "")
      if [[ -f "${job}/${system}.pdb" && -f "${job}/${system}.psf" ]]; then
        this_step="1-min"
        prev_step_num=0
        final_step_num=0
      else
        echo "No requisite pdb or psf files, exiting..."
        exit 1
      fi
      ;;
  esac
}


function create_conf() {

  step_size=$((${final_step_num} - ${prev_step_num}))
  
  echo "Creating conf file using job:"${job}" system:"${system}" this_step:"$this_step" prev_step:$prev_step_num final_step:$final_step_num"

  #######################################################################################
  ## CREATE NAMD CONFIGURATION FILES                                                   ##
  #######################################################################################

  case "${this_step}" in
  
    ##########################
    ## MINIMIZATION 1       ##
    ##########################
    "1-min")
    
      conf="${job}/${system}_${this_step}.conf"
      out="${job}/${system}_${this_step}.out"
      
      echo "structure ${system}.psf" > ${conf}
      echo "coordinates ${system}.pdb" >> ${conf}
      echo "paraTypeCharmm on" >> ${conf}
      for ff in "${ffs[@]}"; do
        echo "parameters "${ff}"" >> ${conf}
      done
      echo "temperature 0" >> ${conf}
      echo "" >> ${conf}
      echo "# Force-Field Parameters" >> ${conf}
      echo "exclude scaled1-4" >> ${conf}
      echo "1-4scaling 1.0" >> ${conf}
      echo "cutoff 12." >> ${conf}
      echo "switching on" >> ${conf}
      echo "switchdist 10." >> ${conf}
      echo "pairlistdist 14" >> ${conf}
      echo "" >> ${conf}
      echo "# Integrator Parameters" >> ${conf}
      echo "timestep 1.0  ;# 1fs/step" >> ${conf}
      echo "nonbondedFreq 1" >> ${conf}
      echo "fullElectFrequency 2" >> ${conf}
      echo "stepspercycle 10" >> ${conf}
      echo "" >> ${conf}
      echo "# Output" >> ${conf}
      echo "outputName ${system}_${this_step}" >> ${conf}
      echo "outputEnergies 100" >> ${conf}
      echo "outputPressure 100" >> ${conf}
      # Minimize
      echo "minimize 15000" >> ${conf}
      ;;
    ##########################
    ## MINIMIZATION 2       ##
    ##########################
    "2-min")
     
      conf="${job}/${system}_${this_step}.conf"
      out="${job}/${system}_${this_step}.out"
     
      # Solvate and ionize
      solv_ion
      
      # create pcell.txt
      calc_pcell_min2

      # Source pcell.txt for PBC values
      source "${job}/${system}_${this_step}_pcell.txt"

      # Begin writing NAMD conf file
      echo "structure ${system}_solv_ion.psf" > ${conf}
      echo "coordinates ${system}_solv_ion.pdb" >> ${conf}
      echo "paraTypeCharmm on" >> ${conf}
      for ff in "${ffs[@]}"; do
        echo "parameters "${ff}"" >> ${conf}
      done
      echo "temperature 0" >> ${conf}
      echo "" >> ${conf}
      echo "# Periodic boundary conditions" >> ${conf}
      echo "wrapWater on" >> ${conf}
      echo "wrapAll on" >> ${conf}
      echo "cellOrigin	$CORX	$CORY	$CORZ" >> ${conf}
      echo "cellBasisVector1	$CBVX	0.0	0.0" >> ${conf}
      echo "cellBasisVector2	0.0	$CBVY	0.0" >> ${conf}
      echo "cellBasisVector3	0.0	0.0	$CBVZ" >> ${conf}
      echo "# Force-Field Parameters" >> ${conf}
      echo "exclude scaled1-4" >> ${conf}
      echo "1-4scaling 1.0" >> ${conf}
      echo "cutoff 12." >> ${conf}
      echo "switching on" >> ${conf}
      echo "switchdist 10." >> ${conf}
      echo "pairlistdist 14" >> ${conf}
      echo "" >> ${conf}
      echo "# Integrator Parameters" >> ${conf}
      echo "timestep 2.0  ;# 2fs/step" >> ${conf}
      echo "rigidBonds all  ;# needed for 2fs steps" >> ${conf}
      echo "nonbondedFreq 1" >> ${conf}
      echo "fullElectFrequency 2" >> ${conf}
      echo "stepspercycle 10" >> ${conf}
      echo "" >> ${conf}
      echo "# Output" >> ${conf}
      echo "outputName ${system}_${this_step}" >> ${conf}
      echo "outputEnergies 10000" >> ${conf}
      echo "outputPressure 10000" >> ${conf}
      # Minimize
      echo "minimize 15000" >> ${conf}
      ;;
    ##########################
    ## EQUILIBRATION        ##
    ##########################
    "3-heat")
    
      conf="${job}/${system}_${this_step}.conf"
      out="${job}/${system}_${this_step}.out"
      
      # create pcell.txt
      calc_pcell_heat

      # Source pcell.txt for PBC values
      source "${job}/${system}_${this_step}_pcell.txt"

      echo "structure ${system}_solv_ion.psf" > ${conf}
      echo "coordinates ${system}_solv_ion.pdb" >> ${conf}
      echo "bincoordinates ${system}_2-min.coor" >> ${conf}
      echo "#extendedSystem ${system}_2-min.xsc" >> ${conf}
      echo "cellOrigin	$CORX	$CORY	$CORZ" >> ${conf}
      echo "cellBasisVector1	$CBVX	0.0	0.0" >> ${conf}
      echo "cellBasisVector2	0.0	$CBVY	0.0" >> ${conf}
      echo "cellBasisVector3	0.0	0.0	$CBVZ" >> ${conf}
      echo "" >> ${conf}
      echo "firsttimestep ${prev_step_num}" >> ${conf}
      echo "" >> ${conf}
      echo "paraTypeCharmm on" >> ${conf}
      for ff in "${ffs[@]}"; do
        echo "parameters "${ff}"" >> ${conf}
      done
      echo "temperature 0" >> ${conf}
      echo "#margin 2" >> ${conf}
      echo "" >> ${conf}
      echo "# Periodic boundary conditions" >> ${conf}
      echo "wrapWater on" >> ${conf}
      echo "wrapAll on" >> ${conf}
      echo "" >> ${conf}
      echo "# Full Electrostatics" >> ${conf}
      echo "PME on" >> ${conf}
      echo "PMEGridSpacing 1.0" >> ${conf}
      echo "# Force-Field Parameters" >> ${conf}
      echo "exclude scaled1-4" >> ${conf}
      echo "1-4scaling 1.0" >> ${conf}
      echo "cutoff 12." >> ${conf}
      echo "switching on" >> ${conf}
      echo "switchdist 10." >> ${conf}
      echo "pairlistdist 14" >> ${conf}
      echo "" >> ${conf}
      echo "# Integrator Parameters" >> ${conf}
      echo "timestep 2.0  ;# 2fs/step" >> ${conf}
      echo "rigidBonds all  ;# needed for 2fs steps" >> ${conf}
      echo "nonbondedFreq 1" >> ${conf}
      echo "fullElectFrequency 2  " >> ${conf}
      echo "stepspercycle 10" >> ${conf}
      echo "" >> ${conf}
      echo "# Output" >> ${conf}
      echo "outputName ${system}_${this_step}" >> ${conf}
      echo "outputEnergies 1000" >> ${conf}
      echo "outputPressure 1000" >> ${conf}
      echo "dcdfreq 1000" >> ${conf}
      echo "" >> ${conf}
      echo "# Constant Temperature Control" >> ${conf}
      echo "langevin on    ;# do langevin dynamics" >> ${conf}
      echo "langevinDamping 0.5   ;# damping coefficient (gamma) of 0.5/ps" >> ${conf}
      echo "langevinTemp 310" >> ${conf}
      echo "langevinHydrogen no    ;# don't couple langevin bath to hydrogens" >> ${conf}
      echo "" >> ${conf}
      echo "# Constant Pressure Control (variable volume)" >> ${conf}
      echo "useGroupPressure yes ;# needed for 2fs steps" >> ${conf}
      echo "useFlexibleCell yes ;# no for water box, yes for membrane" >> ${conf}
      echo "useConstantRatio yes  ;# no for water box, yes for membrane" >> ${conf}
      echo "langevinPiston on" >> ${conf}
      echo "langevinPistonTarget 1.01325 ;#  in bar -> 1 atm" >> ${conf}
      echo "langevinPistonPeriod 100." >> ${conf}
      echo "langevinPistonDecay 50." >> ${conf}
      echo "langevinPistonTemp 310" >> ${conf}
      echo "" >> ${conf}
      # Equilibration loop
      echo "set freq 100" >> ${conf}
      echo 'for {set i 10} {$i <= 310} {incr i 10} {' >> ${conf}
      echo 'reinitvels $i' >> ${conf}
      echo 'langevinTemp $i' >> ${conf}
      echo 'run $freq' >> ${conf}
      echo '}' >> ${conf}
      # Stabilize
      echo 'run 10000' >> ${conf}
      ;;
    ##########################
    ## SIMULATION           ##
    ##########################
    "4-sim")
    
      conf="${job}/${system}_${this_step}_${final_step_num}.conf"
      out="${job}/${system}_${this_step}_${final_step_num}.out"
      
      # Begin writing NAMD conf file
      echo "structure ${system}_solv_ion.psf" > ${conf}
      echo "coordinates ${system}_solv_ion.pdb" >> ${conf}

      # if this is the first timestep, begin with heat
      if [[ ${prev_step_num} == 0 ]]; then
        echo "temperature 310" >> ${conf}
        echo "bincoordinates ${system}_3-heat.coor" >> ${conf}
        # create pcell.txt
        calc_pcell_heat
        source "${job}/${system}_${this_step}_pcell.txt"
        echo "cellOrigin	$CORX	$CORY	$CORZ" >> ${conf}
        echo "cellBasisVector1	$CBVX	0.0	0.0" >> ${conf}
        echo "cellBasisVector2	0.0	$CBVY	0.0" >> ${conf}
        echo "cellBasisVector3	0.0	0.0	$CBVZ" >> ${conf}
      # else resume
      else
        echo "binvelocities ${system}_${this_step}_${prev_step_num}.vel" >> ${conf}
        echo "bincoordinates ${system}_${this_step}_${prev_step_num}.coor" >> ${conf}
        echo "extendedSystem ${system}_${this_step}_${prev_step_num}.xsc" >> ${conf}
      fi

      echo "" >> ${conf}
      echo "firsttimestep ${prev_step_num}" >> ${conf}
      echo "" >> ${conf}
      echo "#margin 2" >> ${conf}
      echo "paraTypeCharmm on" >> ${conf}
      for ff in "${ffs[@]}"; do
        echo "parameters "${ff}"" >> ${conf}
      done
      echo "" >> ${conf}
      echo "# Periodic boundary conditions" >> ${conf}
      echo "wrapWater on" >> ${conf}
      echo "wrapAll on" >> ${conf}
      echo "" >> ${conf}
      echo "# Full Electrostatics" >> ${conf}
      echo "PME on" >> ${conf}
      echo "PMEGridSpacing 1.0" >> ${conf}
      echo "" >> ${conf}
      echo "# Force-Field Parameters" >> ${conf}
      echo "exclude scaled1-4" >> ${conf}
      echo "1-4scaling 1.0" >> ${conf}
      echo "cutoff 12." >> ${conf}
      echo "switching on" >> ${conf}
      echo "switchdist 10." >> ${conf}
      echo "pairlistdist 14" >> ${conf}
      echo "" >> ${conf}
      echo "# Integrator Parameters" >> ${conf}
      echo "timestep 2.0  ;# 2fs/step" >> ${conf}
      echo "rigidBonds all  ;# needed for 2fs steps" >> ${conf}
      echo "nonbondedFreq 1" >> ${conf}
      echo "fullElectFrequency 2  " >> ${conf}
      echo "stepspercycle 10" >> ${conf}
      echo "" >> ${conf}
      echo "# Output" >> ${conf}
      echo "outputName ${system}_${this_step}_${final_step_num}" >> ${conf}
      echo "outputEnergies 10000" >> ${conf}
      echo "outputPressure 10000" >> ${conf}
      echo "dcdfreq 10000" >> ${conf}
      echo "" >> ${conf}
      echo "# Constant Temperature Control" >> ${conf}
      echo "langevin on    ;# do langevin dynamics" >> ${conf}
      echo "langevinDamping 0.5   ;# damping coefficient (gamma) of 0.5/ps" >> ${conf}
      echo "langevinTemp 310" >> ${conf}
      echo "langevinHydrogen no    ;# don't couple langevin bath to hydrogens" >> ${conf}
      echo "" >> ${conf}
      echo "# Constant Pressure Control (variable volume)" >> ${conf}
      echo "useGroupPressure yes ;# needed for 2fs steps" >> ${conf}
      echo "useFlexibleCell yes ;# no for water box, yes for membrane" >> ${conf}
      echo "useConstantRatio yes  ;# no for water box, yes for membrane" >> ${conf}
      echo "langevinPiston on" >> ${conf}
      echo "langevinPistonTarget 1.01325 ;#  in bar -> 1 atm" >> ${conf}
      echo "langevinPistonPeriod 100." >> ${conf}
      echo "langevinPistonDecay 50." >> ${conf}
      echo "langevinPistonTemp 310" >> ${conf}
      echo "" >> ${conf}
      echo "run ${step_size}" >> ${conf}
      ;;
  esac
}

function solv_ion() {
  # Create temporary VMD script to solvate, ionize, and generate PBC measurements
  # Load molecule
  echo "mol load psf [glob ${job}/${system}.psf] namdbin [glob ${job}/${system}_${prev_step}.coor]" > ${job}/temp_solv_ion.tcl
  echo "set all [atomselect top all]" >> ${job}/temp_solv_ion.tcl
  echo "\$all writepdb ${job}/${system}_${prev_step}.pdb" >> ${job}/temp_solv_ion.tcl
  # Solvate
  echo "package require solvate" >> ${job}/temp_solv_ion.tcl
  echo "solvate ${job}/${system}.psf ${job}/${system}_${prev_step}.pdb -o ${job}/${system}_solv -s WT -x 13 -y 13 -z 13 +x 13 +y 13 +z 13 -b 2.4" >> ${job}/temp_solv_ion.tcl
  # Ionize
  echo "package require autoionize" >> ${job}/temp_solv_ion.tcl
  echo "autoionize -psf ${job}/${system}_solv.psf -pdb ${job}/${system}_solv.pdb -o ${job}/${system}_solv_ion -sc 0.15" >> ${job}/temp_solv_ion.tcl

  # Run VMD
  vmd -dispdev text -eofexit < "${job}/temp_solv_ion.tcl"
  rm ${job}/temp_solv_ion.tcl
}

function calc_pcell_min2() {
  # Load molecule
  echo "mol load psf [glob ${job}/${system}_solv_ion.psf] pdb [glob ${job}/${system}_solv_ion.pdb]" > ${job}/temp_build_box.tcl

  # Get PBC box measurements from VMD and output to pcell.txt
  echo "set fileId [open ${job}/${system}_${this_step}_pcell.txt "w"]" >> ${job}/temp_build_box.tcl
  echo "set all [atomselect top "not lipid"]" >> ${job}/temp_build_box.tcl
  echo 'set minmax [measure minmax $all]' >> ${job}/temp_build_box.tcl
  echo 'set vec [vecsub [lindex $minmax 1] [lindex $minmax 0]]' >> ${job}/temp_build_box.tcl
  echo 'puts $fileId "#!/usr/bin/env bash"' >> ${job}/temp_build_box.tcl
  echo 'puts $fileId "CBVX=[lindex $vec 0]"' >> ${job}/temp_build_box.tcl
  echo 'puts $fileId "CBVY=[lindex $vec 1]"' >> ${job}/temp_build_box.tcl
  echo 'puts $fileId "CBVZ=[lindex $vec 2]"' >> ${job}/temp_build_box.tcl
  echo 'set center [measure center $all]' >> ${job}/temp_build_box.tcl
  echo 'puts $fileId "CORX=[lindex $center 0]"' >> ${job}/temp_build_box.tcl
  echo 'puts $fileId "CORY=[lindex $center 1]"' >> ${job}/temp_build_box.tcl
  echo 'puts $fileId "CORZ=[lindex $center 2]"' >> ${job}/temp_build_box.tcl
  echo 'close $fileId' >> ${job}/temp_build_box.tcl

  # Run VMD
  vmd -dispdev text -eofexit < "${job}/temp_build_box.tcl"
  rm "${job}/temp_build_box.tcl"
}

function calc_pcell_heat() {
  # Load molecule
  echo "mol load psf [glob ${job}/${system}_solv_ion.psf] namdbin [glob ${job}/${system}_${prev_step}.coor]" > ${job}/temp_build_box.tcl

  # Get PBC box measurements from VMD and output to pcell.txt
  echo "set fileId [open ${job}/${system}_${this_step}_pcell.txt "w"]" >> ${job}/temp_build_box.tcl
  echo "set all [atomselect top all]" >> ${job}/temp_build_box.tcl
  echo 'set minmax [measure minmax $all]' >> ${job}/temp_build_box.tcl
  echo 'set vec [vecsub [lindex $minmax 1] [lindex $minmax 0]]' >> ${job}/temp_build_box.tcl
  echo 'puts $fileId "#!/usr/bin/env bash"' >> ${job}/temp_build_box.tcl
  echo 'puts $fileId "CBVX=[lindex $vec 0]"' >> ${job}/temp_build_box.tcl
  echo 'puts $fileId "CBVY=[lindex $vec 1]"' >> ${job}/temp_build_box.tcl
  echo 'puts $fileId "CBVZ=[lindex $vec 2]"' >> ${job}/temp_build_box.tcl
  echo 'set center [measure center $all]' >> ${job}/temp_build_box.tcl
  echo 'puts $fileId "CORX=[lindex $center 0]"' >> ${job}/temp_build_box.tcl
  echo 'puts $fileId "CORY=[lindex $center 1]"' >> ${job}/temp_build_box.tcl
  echo 'puts $fileId "CORZ=[lindex $center 2]"' >> ${job}/temp_build_box.tcl
  echo 'close $fileId' >> ${job}/temp_build_box.tcl

  # Run VMD
  vmd -dispdev text -eofexit < "${job}/temp_build_box.tcl"
  rm "${job}/temp_build_box.tcl"
}

function calc_pcell_sim() {
  # Load molecule
  echo "mol load psf [glob ${job}/${system}_solv_ion.psf] namdbin [glob ${job}/${system}_${prev_step}_${prev_step_num}.coor]" > ${job}/temp_build_box.tcl

  # Get PBC box measurements from VMD and output to pcell.txt
  echo "set fileId [open ${job}/${system}_${this_step}_${prev_step_num}_pcell.txt "w"]" >> ${job}/temp_build_box.tcl
  echo "set all [atomselect top all]" >> ${job}/temp_build_box.tcl
  echo 'set minmax [measure minmax $all]' >> ${job}/temp_build_box.tcl
  echo 'set vec [vecsub [lindex $minmax 1] [lindex $minmax 0]]' >> ${job}/temp_build_box.tcl
  echo 'puts $fileId "#!/usr/bin/env bash"' >> ${job}/temp_build_box.tcl
  echo 'puts $fileId "CBVX=[lindex $vec 0]"' >> ${job}/temp_build_box.tcl
  echo 'puts $fileId "CBVY=[lindex $vec 1]"' >> ${job}/temp_build_box.tcl
  echo 'puts $fileId "CBVZ=[lindex $vec 2]"' >> ${job}/temp_build_box.tcl
  echo 'set center [measure center $all]' >> ${job}/temp_build_box.tcl
  echo 'puts $fileId "CORX=[lindex $center 0]"' >> ${job}/temp_build_box.tcl
  echo 'puts $fileId "CORY=[lindex $center 1]"' >> ${job}/temp_build_box.tcl
  echo 'puts $fileId "CORZ=[lindex $center 2]"' >> ${job}/temp_build_box.tcl
  echo 'close $fileId' >> ${job}/temp_build_box.tcl

  # Run VMD
  vmd -dispdev text -eofexit < "${job}/temp_build_box.tcl"
  rm "${job}/temp_build_box.tcl"
}

function copy_ffs() {
  ffs=()
  for ff in ../0-forcefields/*; do
    if [[ -f "${ff}" ]]; then
      # strip path
      ff_name=${ff##*/}
      # copy ff and append to ffs array
      cp -f "${ff}" "${job}/${ff_name}"
      ffs+=(${ff##*/})
    fi
  done
}

function backup() {
  if [[ -f "${out}" ]]; then
   mv "${out}" "${out}.bak"
  fi
}