#!/usr/bin/env bash
export VMDNOCUDA=1
#export VMDNOOPTIX=1
export VMDNOOSPRAY=1

#######################################################################################
## ADJUSTABLE PARAMETERS                                                             ##
#######################################################################################

# set step size for dcd file incrementing
step_size=10000000
# set folder name that contains the jobs
job_folder="4-new_jobs"
# set sel1
sel1="chain B and not resname PHQ"
# set sel2
sel2="chain A"
# only run specific systems
my_systems=("ALL")

#######################################################################################
## FUNCTIONS                                                                         ##
#######################################################################################

function get_systems() {
  systems=()
  all_systems=()
  _jobs=( $(ls -d ../${job_folder}/*/*/) )
  for _job in ${_jobs[@]}; do
    all_systems+=( $(basename ${_job}) )
  done
  systems+=( $(for _system in ${all_systems[@]}; do echo ${_system}; done | sort -u) )
}

function get_ffs() {
  ffs=()
  for ff in ../0-forcefields/*; do
    if [[ -f "${ff}" ]]; then
      ffs+=(${ff})
    fi
  done
}

#######################################################################################
## ENERGY ANALYSIS                                                                   ##
#######################################################################################

# set folder name for output files
out_dir="$(basename ${0})"
out_dir="${out_dir%%.*}2"

function namd_energy() {

  get_systems
  get_ffs
  
  mkdir -p ${out_dir}
  
  ffs_energy=""
  for ff in ${ffs[@]}; do
    ffs_energy+="-par \"${ff}\" "
  done
    

  for system in "${systems[@]}"; do

    if echo ${my_systems[@]} | grep -w ${system} > /dev/null; then
      
      jobs=( $(ls -d ../${job_folder}/*/${system}/) )
      
      for job in ${jobs[@]}; do
        job=${job%%/}
        date=${job#*/*/}
        date=${date%/*}
        dcds=( $(ls ${job}/*sim_*.dcd | sort -V) )
        xscs=( $(ls ${job}/*sim_*.xsc | sort -V) )
      
        # Run NAMDEnergy
        echo 'package require namdenergy' > temp.tcl
        echo "set sel1 [atomselect top \"${sel1}\"]" >> temp.tcl
        echo "set sel2 [atomselect top \"${sel2}\"]" >> temp.tcl
        echo "namdenergy -vdw -elec -nonb -sel \$sel1 \$sel2 -ofile \""${out_dir}/${system}-${date}.csv"\" -tempname \""${system}-${date}"\" -switch 10 -cutoff 12 -diel 1.0 -T 310 -timemult 2 -stride 10000 -extsys "${xscs[0]}" -pme ${ffs_energy} -exe "/home/bryan/bin/namd2" -plot" >> temp.tcl
        echo 'exit' >> temp.tcl

        # run temp.tcl
        vmd -dispdev text -eofexit ${job}/${system}_solv_ion.psf ${dcds[@]:1} < "temp.tcl"
        rm "temp.tcl"
      done
  
    awk 'FNR == 1 { nfiles++; ncols = NF }
       { for (i = 1; i < NF; i++) sum[FNR,i] += $i
         if (FNR > maxnr) maxnr = FNR
       }
       END {
           for (line = 1; line <= maxnr; line++)
           {
               for (col = 1; col < ncols; col++)
                    printf "  %.2f", sum[line,col]/nfiles;
               printf "\n"
           }
       }' ${out_dir}/${system}-* > ${out_dir}/${system}-combined.csv
   fi
   
   rm FFTW_NAMD*
   
   plot_energy ${system} ${out_dir}/${system}-combined
     
  done
}

function plot_energy() {
gnuplot -persist <<-GNUPLOT_INPUT
  set terminal pngcairo enhanced size 1600,1200 font "arial,24" linewidth 2
  set output "${2}.png"
  set autoscale
  set title "SrtA-{/Arial-Italic Boc} interaction energy" font "arial,40"
  set xlabel "Time (ns)" font "arial,32"
  set ylabel "Energy (kcal/mol)" font "arial,32"
  set border 1+2
  #set xtics nomirror
  set ytics nomirror
  set grid
  set grid noxtics
  set key box off
  set yrange [-25:0]
  #set format y "%+-2.f"
  set style data lines
  set xzeroaxis ls 8
  
  # linetypes
  #set style line 1 pi 0 lw 1 ps 1 pt 0 lc rgb '#000000'

  plot "${2}.csv" using (\$1 / 50):5 with linespoints ls 8 pi 0 pt 0
  
  exit

GNUPLOT_INPUT
}

function plot_energy2() {
gnuplot -persist <<-GNUPLOT_INPUT
  set terminal pngcairo enhanced size 1600,1200 font "arial,24" linewidth 2
  set output "ALL-combined-both.png"
  set datafile separator ","
  set autoscale
  set title "SrtA-{/Arial-Italic Boc} interaction energy" font "arial,40"
  set xlabel "Time (ns)" font "arial,32"
  set ylabel "Energy (kcal/mol)" font "arial,32"
  set border 1+2
  #set xtics nomirror
  set ytics nomirror
  set grid
  set grid noxtics
  set key box right bottom
  set yrange [-80:0]
  #set format y "%+-2.f"
  set style data filledcurves
  #set xzeroaxis ls 8
  
  # linetypes
  set style line 1 pi 0 lw 1 ps 1 pt 0 lc rgb '#696969'
  set style line 2 pi 0 lw 1 ps 1 pt 0 lc rgb 'black'

  plot 'ALL-combined-both.csv' using (\$1 / 50):(\$3 + \$4) title 'SS peptide' with filledcurves y1=0 ls 1, \
  '' using (\$1 / 50):(\$3) title '{/Arial-Italic Boc}' with filledcurves y1=0 ls 2
  
  exit

GNUPLOT_INPUT
}

namd_energy