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

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

# number of frames to skip (i.e. from equilibration)
skip=0
# set folder name that contains the jobs
job_folder="4-jobs"
# set atom selection for concatenation
concat_sel="chain A"
# set atom selection to unwrap
unwrap_sel="chain A"
# set atom selection to wrap to
wrap_to_sel="chain A"
# set atom selection to wrap
wrap_sel="chain A"
# set atom selection for alignment
align_sel="chain A"
# set atom selection for EDA
eda_sel="chain A and name CA"

#######################################################################################
## SET VARS                                                                          ##
#######################################################################################

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

# create suffixes for output files
concat_sel_s="${concat_sel// /_}"
eda_sel_s="${eda_sel// /_}"

#######################################################################################
## 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) )
}

#######################################################################################
## CONCATENATE, UNWRAP, ALIGN INDVIDUAL REPLICATES                                   ##
#######################################################################################

function concat_individual() {

  get_systems

  for system in "${systems[@]}"; do
    
    dcds=()
    jobs=( $(ls -d ../${job_folder}/*/${system}/) )
    
    for job in ${jobs[@]}; do
      job=${job%%/}
      rel_job_path=${job#../${job_folder}/}
      mkdir -p ${out_dir}/individual/${rel_job_path}
      dcds=( $(ls ${job}/*sim_*.dcd | sort -V) )
      
      for dcd in ${dcds[@]}; do
        filename=${dcd##*/}
        filename=${filename%%.*}
        # UNWRAP AND ALIGN
        echo 'package require pbctools' > temp.tcl
        echo "pbc unwrap -all -sel \""${unwrap_sel}"\"" >> temp.tcl
        echo "pbc wrap -sel \""${wrap_sel}"\" -all -compound chain -center com -centersel \""${wrap_to_sel}"\"" >> temp.tcl
        echo "set out [atomselect top \""${concat_sel}"\"]" >> temp.tcl
        echo "animate write dcd ${out_dir}/individual/${rel_job_path}/temp-${filename}-${concat_sel_s}.dcd sel \$out waitfor all" >> temp.tcl
        for (( i=0; i<1; i++)); do
	        echo "\$out writepsf ${out_dir}/individual/${rel_job_path}/${system}-${concat_sel_s}.psf" >> temp.tcl
          echo "\$out writepdb ${out_dir}/individual/${rel_job_path}/${system}-${concat_sel_s}.pdb" >> temp.tcl
	      done

        echo "sync" >> temp.tcl
        echo 'exit' >> temp.tcl

        # run temp.tcl
        vmd -dispdev text -eofexit ${job}/${system}_solv_ion.psf ${dcd} < "temp.tcl"
        sync
        
        rm "temp.tcl"
      done
      
      dcds=( $(ls ${out_dir}/individual/${rel_job_path}/temp-*-${concat_sel_s}.dcd | sort -V) )
      
      echo "set ref [atomselect top \""${align_sel}"\" frame 0]" > temp.tcl
      echo "set sel [atomselect top \""${align_sel}"\"]" >> temp.tcl
      echo 'set all [atomselect top "all"]' >> temp.tcl
      echo 'set n [molinfo top get numframes]' >> temp.tcl
      echo 'for { set i 1 } { $i < $n } { incr i } {' >> temp.tcl
      echo '  $sel frame $i' >> temp.tcl
      echo '  $all frame $i' >> temp.tcl
      echo '  $all move [measure fit $sel $ref]' >> temp.tcl
      echo '}' >> temp.tcl
      echo "animate write dcd ${out_dir}/individual/${rel_job_path}/${system}-${concat_sel_s}.dcd beg ${skip} waitfor all" >> temp.tcl
      echo "sync" >> temp.tcl
      echo 'exit' >> temp.tcl

      # run temp.tcl
      vmd -dispdev text -eofexit ${out_dir}/individual/${rel_job_path}/${system}-${concat_sel_s}.psf ${dcds[@]} < "temp.tcl"
      sync
      rm "temp.tcl" 
      rm ${out_dir}/individual/${rel_job_path}/temp*
    done
  done
}

#######################################################################################
## CONCATENATE AND ALIGN ALL REPLICATES                                              ##
#######################################################################################

function concat_combined() {

  get_systems
  
  mkdir -p ${out_dir}/combined
  
  for system in ${systems[@]}; do
  

    # create dcd array and copy pdb/psf
    dcds=( $(ls ${out_dir}/individual/*/${system}/${system}-${concat_sel_s}.dcd | sort -V) )
    psfs=( $(ls ${out_dir}/individual/*/${system}/${system}-${concat_sel_s}.psf | sort -V) )
    pdbs=( $(ls ${out_dir}/individual/*/${system}/${system}-${concat_sel_s}.pdb | sort -V) )
    
    cp ${psfs[1]} ${out_dir}/combined/${system}-${concat_sel_s}.psf
    cp ${pdbs[1]} ${out_dir}/combined/${system}-${concat_sel_s}.pdb

    # ALIGN
    echo "set ref [atomselect top \""${align_sel}"\" frame 0]" >> temp.tcl
    echo "set sel [atomselect top \""${align_sel}"\"]" >> temp.tcl
    echo 'set all [atomselect top "all"]' >> temp.tcl
    echo 'set n [molinfo top get numframes]' >> temp.tcl
    echo 'for { set i 1 } { $i < $n } { incr i } {' >> temp.tcl
    echo '  $sel frame $i' >> temp.tcl
    echo '  $all frame $i' >> temp.tcl
    echo '  $all move [measure fit $sel $ref]' >> temp.tcl
    echo '}' >> temp.tcl
    # write aligned trr and dcd output files
    echo "animate write trr ${out_dir}/combined/${system}-${concat_sel_s}.trr waitfor all" >> temp.tcl
    echo "animate write dcd ${out_dir}/combined/${system}-${concat_sel_s}.dcd waitfor all" >> temp.tcl
    echo 'sync' >> temp.tcl
    echo 'exit' >> temp.tcl

    vmd -dispdev text -eofexit ${out_dir}/combined/${system}-${concat_sel_s}.psf ${dcds[@]} < "temp.tcl"
    rm "temp.tcl"
    sync

  done
}

#######################################################################################
## CONCATENATE AND ALIGN ALL SYSTEMS (EDA)                                           ##
#######################################################################################

function concat_eda() {

  get_systems
  
  mkdir -p ${out_dir}/eda/
  
  # concat systems to eda_sel
  for system in ${systems[@]}; do

    echo "set out [atomselect top \""${eda_sel}"\"]" > temp.tcl
    echo "set last_out [atomselect top \""${eda_sel}"\" frame 0]" >> temp.tcl
    
    # write aligned dcd
    echo "animate write dcd ${out_dir}/eda/${system}-${eda_sel_s}.dcd sel \$out waitfor all" >> temp.tcl
    echo "\$last_out writepsf ${out_dir}/eda/${system}-${eda_sel_s}.psf" >> temp.tcl
    echo "\$last_out writepdb ${out_dir}/eda/${system}-${eda_sel_s}.pdb" >> temp.tcl
    echo 'sync' >> temp.tcl
    echo 'exit' >> temp.tcl
    
    vmd -dispdev text -eofexit ${out_dir}/combined/${system}-${concat_sel_s}.psf ${out_dir}/combined/${system}-${concat_sel_s}.dcd < "temp.tcl"
    sync
    rm "temp.tcl"
  done
  
  dcds=( $(ls ${out_dir}/eda/*-${eda_sel_s}.dcd | sort -V) )
  psfs=( $(ls ${out_dir}/eda/*-${eda_sel_s}.psf | sort -V) )
  pdbs=( $(ls ${out_dir}/eda/*-${eda_sel_s}.pdb | sort -V) )
  
  cp ${psfs[1]} ${out_dir}/eda/combined-${eda_sel_s}.psf
  cp ${pdbs[1]} ${out_dir}/eda/combined-${eda_sel_s}.pdb
  
  # concat and align
  echo "set ref [atomselect top \""${eda_sel}"\" frame 0]" > temp.tcl
  echo "set sel [atomselect top \""${eda_sel}"\"]" >> temp.tcl
  echo 'set all [atomselect top "all"]' >> temp.tcl
  echo 'set n [molinfo top get numframes]' >> temp.tcl
  echo 'for { set i 1 } { $i < $n } { incr i } {' >> temp.tcl
  echo '  $sel frame $i' >> temp.tcl
  echo '  $all frame $i' >> temp.tcl
  echo '  $all move [measure fit $sel $ref]' >> temp.tcl
  echo '}' >> temp.tcl
  echo "animate write dcd ${out_dir}/eda/combined-${eda_sel_s}.dcd waitfor all" >> temp.tcl
  echo 'sync' >> temp.tcl
  echo 'exit' >> temp.tcl
    
  vmd -dispdev text -eofexit ${out_dir}/eda/combined-${eda_sel_s}.psf ${dcds[@]} < "temp.tcl"
  
  sync
  rm temp.tcl
  for system in ${systems[@]}; do
    rm "${out_dir}/eda/${system}-${eda_sel_s}.dcd"
    rm "${out_dir}/eda/${system}-${eda_sel_s}.pdb"
    rm "${out_dir}/eda/${system}-${eda_sel_s}.psf"
  done
}

concat_individual
#concat_combined
#concat_eda