#!/bin/bash -e
###########################################################################
# This file is part of meadTools, version 2.2.
# 
# Copyright (c) 2001-2019, Instituto de Tecnologia Quimica e Biologica,
# Universidade Nova de Lisboa, Portugal.
# 
# meadTools is free software: you can redistribute it and/or modify it
# under the terms of the GNU General Public License as published by the
# Free Software Foundation, either version 2 of the License, or (at your
# option) any later version.
# 
# meadTools is distributed in the hope that it will be useful, but
# WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
# General Public License for more details.
# 
# You should have received a copy of the GNU General Public License
# along with meadTools.  If not, see <http://www.gnu.org/licenses/>.
# 
# For further details and info check the README file.
# 
# You can get meadTools at www.itqb.unl.pt/simulation
###########################################################################


############################################################################
# meadT : a program to compute Poisson-Boltzmann free energy terms for
#         the tautomeric binding of protons and/or electrons.
#
# SUMMARY: meadT is a tool to compute the individual and pairwise
# Poisson-Boltzmann (PB) free energy terms required for simulating the
# tautomeric protonation equilibrium in a molecular system.  It
# follows the formalism described in [1], where each tautomeric site
# is split into several non-tautomeric pseudosites.  Note that meadT
# is not a PB solver but rather a front-end to the multiflex program,
# which is part of Don Bashford's MEAD package.  Redox sites can also
# be included, as described in [2].  meadT is part of the meadTools
# package.
# 
# INPUT: meadT accepts the same command-line options and arguments as
# multiflex, plus the additional options -n, -b, -s and -m.  The
# option -n specifies the number of CPUs (actually cores) used for
# single-host parallelization; note that no multi-host parallelization
# is supported and that no check is made of whether the requested
# number of CPUs exists.  The option -b specifies the number of
# pseudosites to include in each block when computing interactions
# (see below).  The options -s and -m indicate the directories where
# the executables statepqr and multiflex can respectively be found (by
# default they are looked for in the running directory, eventually as
# symbolic links).  The files expected to be found by multiflex must
# be in the running directory, as usual (see MEAD's documentation for
# details).  Those files comply with the multiflex rules except for
# the following two cases.  First, all pseudosites in the .pqr file
# should be in the fully charged state, the reference state used by
# the formalism in [1] (in any case, this is checked before starting
# any calculations).  Second, each line in the .sites file should not
# contain a single site_type entry (indicating its corresponding .st
# file; see MEAD's documentation for details), but rather a sequence
# of the site_type entries for all the pseudosites associated with
# that site (eg, instead of 'LYS' you should have 'LYStau1 LYStau2
# LYStau3').  If some .potat files from previous calculations already
# exist, meadT uses them instead of repeating the corresponding
# calculations, as usual.
# 
# ALGORITHM: meadT consists of two main parts, namely the calculation
# of individual (pKcrg) and pairwise (Wij) terms.  As noted above,
# meadT follows the formalism in [1], which uses a fully charged
# reference state and defines for each pseudosite a pKcrg value (pK*
# in the notation of [1]) corresponding to the the pKa of that
# pseudosite in an otherwise fully charged system.  In contrast,
# multiflex assumes a fully neutral reference state, defining an
# intrinsic pKa (pKint) for each site in an otherwise fully neutral
# system.  In order to compute the pKcrg values, meadT makes a single
# multiflex run for each pseudosite while keeping the remaining
# pseudosites fully charged.  These individual multiflex calculations
# are the most time-consuming part of meadT, being split over the
# requested CPUs and run in different directories (the 'CPU numbers'
# reported in output just label the different parallel jobs, having no
# hardware relation).  The calculation of the interactions Wij is
# straightforward, since they are identical in either the fully
# neutral or fully charged states (see [1]).  Since all the required
# .potat files are created by the pKcrg calculations, the subsequent
# calculation of the Wij terms is usually very fast using a single
# multiflex run done in the standard way.  Actually, the use of
# parallelization is usually detrimental at this stage, which is very
# heavy on disk access (reading the .potat files); thus, Wij
# calculations are run in a single processor.  Furthermore, since
# memory may also be a problem at this stage (all .potat files are
# simultaneously loaded in the memory), a piecewise method is used,
# where the set of pseudosites is split into several blocks, whose
# intra-block and block-block interactions are then computed.
# 
# OUTPUT: The main output of meadT is similar to that of multiflex:
# one .pkcrg file containing the pKcrg values for all pseudosites, one
# .g file containing the Wij interactions between all pseudosites
# pairs, and all the .potat files.  Two logfiles are created
# containing the appended output sent by all multiflex runs to
# standard output and to standard error (files .out and .err).
# Several temporary files and directories are also created, being
# removed at program termination.
# 
# REFERENCES:
# [1] AM Baptista, CM Soares (2001) J. Phys. Chem. B, 105:293
# [2] AM Baptista, PJ Martel, CM Soares (1999) Biophys. J. 76:2978
############################################################################


############################################################################
# Define some variables and functions:

# Program name:
prog=`basename $0`

# Option defaults:
ncpus=1
bsize=300
statepqr_dir=.
multiflex_dir=.

usage="\
Usage: $prog [OPTIONS] MEAD_OPTIONS RUN_NAME
Options:
  -n #CPUS           Number of requested same-host CPUs (cores).
  -b BLOCK_SIZE      Number of pseudosites per block in Wij calculations.
  -s STATEPQR_DIR    Directory where statepqr is located.
  -m MULTIFLEX_DIR   Directory where multiflex is located.
  -h                 Shows this message.
Defaults:  -n $ncpus  -b $bsize  -s $statepqr_dir  -m $multiflex_dir"

function message {
  case "$1" in
    E ) shift; echo -e "$prog: Error: $*" >&2; exit 1;;
    U ) shift; echo -e "$prog: Error: $*\n$usage" >&2; exit 1;;
    W ) shift; echo -e "$prog: Warning: $*" >&2;;
    * ) message E "Wrong use of 'message' function.";;
  esac
}

############################################################################
# Parse arguments and do some checking:

[ $# -lt 1 ] && message U "Wrong number of arguments."
# None of -n, -b, -s, -m and -h are currently used by multiflex
# (as of MEAD 2.2.5), but this should be checked for each new MEAD
# version.
until [ -z "$1" ]; do
  case "$1" in
    -n ) shift; ncpus=$1;;
    -b ) shift; bsize=$1;;
    -s ) shift; statepqr_dir=$1;;
    -m ) shift; multiflex_dir=$1;;
    -h ) echo "$usage"; exit 0;;
    *  ) mead_args="$mead_args $1"; run=$1;;
  esac
  shift
done

# Check for multiflex and statepqr:
for f in $multiflex_dir/multiflex $statepqr_dir/statepqr; do
  [ ! -x $f ] && message E "No executable $f was found. Aborting."
done
# Switch directories to absolute path:
statepqr_dir=`(cd $statepqr_dir; pwd -P)`
multiflex_dir=`(cd $multiflex_dir; pwd -P)`

# Check if .pqr file is in the fully charged state:
! $statepqr_dir/statepqr r=c $run.pqr $run.sites |
  diff -q $run.pqr - &> /dev/null &&
  message E "File $run.pqr is not in the fully charged state. Aborting."

############################################################################
# Start by doing some initializations:

echo "Starting meadT at `date +%T`"

# Backup original .sites file and make extended .psites file:
cp $run.sites $run.sites0
awk '{for(i=2;i<=NF;i++) print $1,$i}' $run.sites > $run.psites

# Get total number of pseudosites and change ncpus if needed:
n=$(($(wc -l < $run.psites)))
[ $ncpus -gt $n ] && ncpus=$n &&
  message W "Requested #CPUs > #psites. #CPUs changed to $n."


############################################################################
# Compute pKint values, splitting into background jobs over CPUs:

# Compute pseudosite range for each CPU:
for (( cpu=0 ; cpu<$ncpus ; cpu++ )); do
  first=$((${last:=0} + 1))
  last=$(($first + ($n-$last)/($ncpus-$cpu) - !(($n-$last)%($ncpus-$cpu))))
  min[$cpu]=$first
  max[$cpu]=$last
  echo "Pseudosite range for CPU $cpu :  ${min[$cpu]}-${max[$cpu]}"
done

listPIDs=()

# Launch background subshells for all CPUs:
for (( cpu=0 ; cpu<$ncpus ; cpu++ )); do
  (
      # Create directory for this CPU and make needed links
      # (including eventual .potat files from previous runs):
      mkdir ${run}_cpu${cpu}
      cd ${run}_cpu${cpu}
      ln -s ../{$run.{pqr,mgm,ogm},*.st} .
      ls ../$run.*.potat &> /dev/null && ln -s ../$run.*.potat .

      # Run one multiflex for each pseudosite (to get its pKint value):
      for (( i=${min[$cpu]} ; i<=${max[$cpu]} ; i++ )); do
	  awk "NR==$i" ../$run.psites > $run.sites
	  printf "Pseudosite %4d (%15s) starting on CPU %d at %8s\n" \
		 $i $(awk '{print $2"-"$1}' $run.sites) $cpu $(date +%T)
	  $multiflex_dir/multiflex $mead_args > $run.out  2> $run.err
	  cat $run.out >> $run.out-app
	  cat $run.pkint >> $run.pkint-app
      done

      # Go back to main directory and end subshell:
      cd ..
      echo "CPU $cpu finished pKint runs at $(date +%T)"
  ) &
  listPIDs=(${listPIDs[@]} $!)
done

echo "Waiting for all CPUs to finish the pKint runs..."
for (( cpu=0 ; cpu<$ncpus ; cpu++ )); do
    if wait ${listPIDs[$cpu]}; then
	echo "Job finished normally in CPU $cpu"
    else
	echo "Job died in CPU $cpu. Exiting, but others may still be running."
	exit 1
    fi
done
echo "All CPUs finished the pKint runs at $(date +%T)"

# Get needed data from ${run}_cpu*/ directories:
rm -f $run.out-app $run.pkcrg
for (( cpu=0 ; cpu<$ncpus ; cpu++ )); do
  cat ${run}_cpu${cpu}/$run.out-app >> $run.out-app
  cat ${run}_cpu${cpu}/$run.pkint-app >> $run.pkcrg
done
for f in ${run}_cpu*/*.potat; do [ ! -h $f ] && mv $f .; done

############################################################################
# Compute Wij values, splitting into different blocks (on a single CPU):

# Run one mead run for each block (to get Wij values):
rm -f $run.g-app
for (( i=1 ; i<=$n ; i+=$bsize )); do
  isz=$(($(awk "NR>=$i && NR<$i+$bsize" $run.psites | tee $run.sites | wc -l)))
  printf "Block %d-%d starting at %8s\n" $i $(($i+$isz-1)) $(date +%T)
  $multiflex_dir/multiflex $mead_args > $run.out  2> $run.err
  # Fix pseudosite numbers:
  awk -v i=$i '{print $1+i-1,$2+i-1"   "$3}' $run.g >> $run.g-app
done

# Run one mead run for each pair of blocks (to get Wij values):
for (( i=1 ; i<=$n ; i+=$bsize )); do
  for (( j=$(($i+$bsize)) ; j<=$n ; j+=$bsize )); do
    isz=$(($(awk "NR>=$i && NR<$i+$bsize" $run.psites | tee $run.sites | wc -l)))
    jsz=$(($(awk "NR>=$j && NR<$j+$bsize" $run.psites | tee -a $run.sites | wc -l)))
    printf "Blocks %d-%d,%d-%d starting at %8s\n" \
      $i $(($i+$isz-1)) $j $(($j+$jsz-1)) $(date +%T)
    $multiflex_dir/multiflex $mead_args > $run.out  2> $run.err
    # Fix pseudosite numbers:
    awk -v i=$i -v j=$j -v b=$bsize \
      '{print $1+($1>b?j-b:i)-1,$2+($2>b?j-b:i)-1"   "$3}' $run.g >> $run.g-app
  done
done

echo "All block runs finished at $(date +%T)"

# Sort interactions and eliminate redundant ones:
echo "Sorting of interactions starting at $(date +%T)"
sort -T . -k 1n -k 2n -u $run.g-app > $run.g-aux
rm $run.g-app

# Put high Wij values in tautomeric pairs:
echo "Setting of \"infinite\" interactions starting at $(date +%T)"
awk -v high="1.000000e+00" '
NR==FNR{for(i=1;i<NF-1;i++) for(j=i+1;j<=NF-1;j++) p[n+i,n+j]=p[n+j,n+i]=1; n+=NF-1}
NR!=FNR{print $1, $2 "   " (p[$1,$2]==1 ? high : $3)}
' $run.sites0 $run.g-aux > $run.g

############################################################################
# Clean and exit:

mv $run.out-app $run.out
mv $run.sites0 $run.sites
rm -r ${run}_cpu*/ $run.{psites,g-aux,err,pkint,summ}

echo "Ending meadT at $(date +%T)"

exit 0

