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

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

# set sel1 and sel1 name
sel1=("resid 171" "acceptor")
# set sel2 and sel2 name
#sel2=("chain B and resname PHQ" "{/Arial-Italic Boc}")
sel2=("resid 140 204" "donor")
# job folder
job_folder="4-jobs"
# cutoff
cutoff=3.0


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

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

# set output suffix
suffix="${sel1[0]}_vs_${sel2[0]}"
suffix="${suffix// /_}"

# create array of systems

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


#######################################################################################
## CALC CONTACTS                                                                     ##
#######################################################################################

function calc_contacts_num() {
  
  mkdir -p ${out_dir}
  echo ${systems[@]}
  
  for system in ${systems[@]}; do
    
    psf=$(ls 1-concat/combined/${system}-*.psf)
    dcd=$(ls 1-concat/combined/${system}-*.dcd)
  
    echo 'set n [molinfo top get numframes]' > temp.tcl
    echo "set sel1 [atomselect top \""${sel1[0]}"\"]" >> temp.tcl
    echo "set sel2 [atomselect top \""${sel2[0]}"\"]" >> temp.tcl
    echo "set fp [open "${out_dir}/${system}-${suffix}.csv" w]" >> temp.tcl
    
    # cycle over the trajectory
    echo 'for {set i 0} {$i <= $n} {incr i} {' >> temp.tcl
    echo '  $sel1 frame $i' >> temp.tcl
    echo '  $sel2 frame $i' >> temp.tcl
    echo "  set contacts [measure contacts ${cutoff} \$sel1 \$sel2]" >> temp.tcl
    echo '  set count [llength [lindex $contacts 0]]' >> temp.tcl
    echo '  puts $fp "${i} ${count}"' >> temp.tcl
    echo '}'  >> temp.tcl
    echo 'close $fp'  >> temp.tcl
    
    vmd -dispdev text -eofexit ${psf} ${dcd} < "temp.tcl"
    sync
    rm "temp.tcl"
    
  done
  
}
  
  
calc_contacts_num
  
function plot_contacts_num() {

  systems=()
  names=()
  
  for (( i = 0 ; i < ${#systems_names[@]} ; i+=2 )); do
    systems+=(${systems_names[${i}]})
  done
  
  for (( i = 1 ; i < ${#systems_names[@]} ; i+=2 )); do
    names+=("${systems_names[${i}]}")
  done

  echo 'set terminal pngcairo enhanced size 1600,1200 font "arial,24" linewidth 2' > temp
  echo "set output \"${out_dir}/combined-${suffix}.png\"" >> temp
  echo 'set autoscale' >> temp
  echo "set title \"${sel1[1]}-${sel2[1]} contacts\" font \"arial,40\"" >> temp
  echo 'set xlabel "ns" font "arial,32"' >> temp
  echo 'set ylabel "# Contacts" font "arial,32"' >> temp
  echo 'set border 1+2' >> temp
  echo 'set xtics nomirror out' >> temp
  echo 'set ytics nomirror' >> temp
  echo 'set grid' >> temp
  echo 'set grid noxtics' >> temp
  echo 'set key box top right' >> temp
  echo '#set yrange [-25:0]' >> temp
  echo '#set format y "%+-2.f"' >> temp
  echo 'set style data lines' >> temp
  echo '#set xzeroaxis ls 8' >> temp
  
  # linetypes
  echo '#set style line 1 pi 0 lw 1 ps 1 pt 0 lc rgb "black"' >> temp
  echo 'set style line 1 pi 0 lw 1 ps 1 pt 0 lc rgb "red"' >> temp
  echo 'set style line 2 pi 0 lw 1 ps 1 pt 0 lc rgb "green"' >> temp
  echo 'set style line 3 pi 0 lw 1 ps 1 pt 0 lc rgb "blue"' >> temp
  echo 'set style line 4 pi 0 lw 1 ps 1 pt 0 lc rgb "orange"' >> temp
  
  # plot
  echo 'plot \' >> temp
  
  for (( i = 0 ; i < ${#systems[@]} ; i++ )); do
    if [ $i -eq $((${#systems[@]} - 1)) ]; then
      echo "\"${out_dir}/${systems[$i]}-${suffix}.csv\" using 1:2 with linespoints ls $((${i} + 1)) title \""${names[$i]}"\"" >> temp
    else
      echo "\"${out_dir}/${systems[$i]}-${suffix}.csv\" using 1:2 with linespoints ls $((${i} + 1)) title \""${names[$i]}"\", \\" >> temp
    fi
  done
  
  echo 'exit' >> temp

  gnuplot temp
  #rm temp
}