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

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

# set sel1 and sel1 name
sel1=("chain A and name CA" "SrtA")
# set sel2 and sel2 name
#sel2=("chain B and resname PHQ" "{/Arial-Italic Boc}")
sel2=("chain B" "SS")
# only run specific systems
#my_systems=("ALL")
systems_names=('ALL' '+{/Arial-Italic Boc} +Ca^{2+}' 'no_CA' '+{/Arial-Italic Boc} −Ca^{2+}' 'no_PHQ' '−{/Arial-Italic Boc} +Ca^{2+}' 'no_PHQ_no_CA' '−{/Arial-Italic Boc} −Ca^{2+}')
# cutoff
cutoff=3.0
# combine?
combine=1

#######################################################################################
## 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// /_}"


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

function calc_contacts() {

  systems=()
  
  for (( i = 0 ; i < ${#systems_names[@]} ; i+=2 )); do
    systems+=(${systems_names[i]})
  done
  
  mkdir -p ${out_dir}
  
  for system in ${systems[@]}; do
    
    psf=$(ls 1-concat/combined/${system}-*.psf)
    dcd=$(ls 1-concat/combined/${system}-*.dcd)

    echo 'set numframes [expr {[molinfo top get numframes] - 1}]' > temp.tcl
    echo 'set totframes [expr $numframes + 1]' >> temp.tcl
    echo "set A [atomselect top \""${sel1[0]}"\"]" >> temp.tcl
    echo "set B [atomselect top \""${sel2[0]}"\"]" >> temp.tcl
    
    # cycle over the trajectory
    echo 'for {set i 0} {$i <= $numframes} {incr i} {' >> temp.tcl
    echo '  foreach resID [$A get resid] {' >> temp.tcl
    echo '    set A_temp [atomselect top "resid $resID"]' >> temp.tcl
    echo '    $A_temp frame $i' >> temp.tcl
    echo '    $B frame $i' >> temp.tcl
    echo '    $A_temp update' >> temp.tcl
    echo '    $B update' >> temp.tcl
    echo "    foreach {listA tmp} [measure contacts ${cutoff} \$A_temp \$B] break" >> temp.tcl
    echo '    if {[llength $listA] > 0} {' >> temp.tcl
    echo '      lappend contactTable($resID) $i' >> temp.tcl
    echo '    } else {'  >> temp.tcl
    echo '      lappend contactTable($resID)' >> temp.tcl
    echo '    }' >> temp.tcl
    echo '    $A_temp delete' >> temp.tcl
    echo '  }' >> temp.tcl
    echo '  puts $i' >> temp.tcl
    echo '}' >> temp.tcl
    
    # open output file for writing
    echo "set fContAB \""${out_dir}/${system}-${suffix}-contacts.csv"\"" >> temp.tcl
    echo 'set fcAB [open $fContAB w]' >> temp.tcl
    # calculate percentage
    echo 'foreach resID [lsort -integer [array names contactTable]] {' >> temp.tcl
    echo '  set frames [llength $contactTable($resID)]' >> temp.tcl
    echo '  set percentage [expr 100.0 * $frames / $totframes]' >> temp.tcl
    # write to file
    echo '  puts $fcAB "$resID $percentage"' >> temp.tcl
    echo '  flush $fcAB' >> temp.tcl
    # set user values
    echo '  set sel1 [atomselect top "resid $resID"]' >> temp.tcl
    echo '  for {set i 0} {$i <= $numframes} {incr i} {' >> temp.tcl
    echo '    $sel1 frame $i' >> temp.tcl
    echo '    $sel1 set user $percentage' >> temp.tcl
    echo '  }' >> temp.tcl
    echo '  $sel1 delete' >> temp.tcl
    echo '}' >> temp.tcl
    # close output files
    echo 'close $fcAB' >> temp.tcl
    echo 'exit' >> temp.tcl
    
    vmd -dispdev text -eofexit ${psf} ${dcd} < "temp.tcl"
    sync
    rm "temp.tcl"
  
    plot_contacts "${system}"
    
  done
    
  if (( ${combine} != 0 )); then
    plot_combined_contacts "${systems_names[@]}"
  fi
  
}

function plot_contacts() {

  echo 'set terminal pngcairo enhanced size 1600,1200 font "arial,24" linewidth 2' > temp
  echo "set output \"${out_dir}/${1}-${suffix}-contacts.png\"" >> temp
  echo 'set autoscale' >> temp
  echo "set title \"${sel1[1]}-${sel2[1]} contacts\" font \"arial,40\"" >> temp
  echo 'set xlabel "Residue" 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 off' >> 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
  
  echo "plot "${out_dir}/${1}-${suffix}-contacts.csv" using 1:2 with impulses ls 8" >> temp
  
  echo 'exit' >> temp

  gnuplot temp
  rm temp
}


# probably want to handle this one manually
function plot_combined_contacts() {

  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 2400,1200 font "arial,24" linewidth 2' > temp
  echo "set output \"${out_dir}/combined-${suffix}-contacts.png\"" >> temp
  echo 'set autoscale' >> temp
  echo "set title \"${sel1[1]}-${sel2[1]} contacts\" font \"arial,40\"" >> temp
  echo 'set xlabel "Residue" font "arial,32"' >> temp
  echo 'set ylabel "% Contacts" font "arial,32"' >> temp
  echo 'set border 1+2' >> temp
  echo 'set xtics nomirror out rotate' >> temp
  echo 'set ytics nomirror' >> temp
  echo 'set grid' >> temp
  echo 'set grid noxtics' >> temp
  echo 'set key box top center' >> temp
  echo '#set style data histogram' >> temp
  echo 'set style fill solid border' >> 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

  # thin out xtic labels
  echo 'every(col) = (int(column(col))%10 ==0)?stringcolumn(1):""' >> temp

  # plot
  echo 'plot \' >> temp
  
  for (( i = 0 ; i < ${#systems[@]} ; i++ )); do
    if [ $i -eq $((${#systems[@]} - 1)) ]; then
      echo "\"${out_dir}/${systems[$i]}-${suffix}-contacts.csv\" using 2:xticlabels(every(1)) with histogram ls $((${i} + 1)) title \""${names[$i]}"\"" >> temp
    else
      echo "\"${out_dir}/${systems[$i]}-${suffix}-contacts.csv\" using 2:xticlabels(every(1)) with histogram ls $((${i} + 1)) title \""${names[$i]}"\", \\" >> temp
    fi
  done
  
  echo 'exit' >> temp

  gnuplot temp
  #rm temp

}


calc_contacts