#!/bin/bash

#    pft2p
#    A script to parallelize pft2
#    Author: Eduardo Sanz-García
#    Created: 20090324    Modified: 20090420

#    Feel free to modified, tune up, and redistribute it.
#    Suggestion or questions to: esanzgar@gmail.com

#    Considerations:
#           1- Log files (pft2p ... > log) doesn't make as much sense as normal pft2 log files.
#              This is because multiple processes are writing at the same time in the 
#              same file

warning(){
	echo "WARNING ${0##*/}: $@" >&2
}

usage(){
cat>&2<<-EOF

	Usage: ${0##*/} [options] input.star [input.star]
	---------------------------------------------------------
	A script to parallelize pft2 in multicore machines.
	pft2p accepts the usual pft2 options, except -select,
	since selection numbers are use to paralellize pft2,
	and -mode 1 or sphere search.
	
	Additionally pft2p is able to catch the environmental
	variable PROCESSORS.

	For more sophisticated usages, like clusters, pft2p 
	also catch the environmental variables NODES and WORKDIR.

	Actions:
	-processors 2            Number of processor available.
	-keep                    Keep intermediate files.
 
EOF
}

INSTAR=                      # Input STAR file(s)
OUTSTAR=                     # Output STAR file
LIST_STAR=                   # List of all the individual STAR files (one per process)
PFT2OPTS=                    # PFT2 options
VERBOSE=

# Set PROCESSORS and NODES variables
[[ $NODES ]] && NODES=( $NODES )
[[ $NODES ]] && PROCESSORS=${#NODES[@]}      # Number of NODES overwrites number of PROCESSORS
NUM_PROC=${PROCESSORS:-1}                    # Number of processor available (default 1)

# Process argument options
[[ $# -eq 0 ]] && usage && exit

while [[ $# -gt 0 ]]; do
	case "$1" in
		-pro*  ) NUM_PROC="${2:-1}" && shift 2 ;;
		-out*  ) OUTSTAR="$2" && shift 2 ;;
		-ke*   ) KEEP=1 && shift ;;
		-sel*  ) warning "Since the selection numbers are use to paralellize"
				 warning "pft2, -select option is not acepted!"
				 exit -1 ;;
		-mod*  ) case "$2" in
					1 | sphere_search ) warning -mode 1 or sphere_search not allowed at this moment! && exit -1 ;;
					* ) PFT2OPTS="$PFT2OPTS $1 $2" && shift 2 ;;
					esac ;;
		-ver* ) VERBOSE="$2" && PFT2OPTS="$PFT2OPTS $1 $2" && shift 2 ;;
		*.star ) INSTAR="$INSTAR $1" && shift ;;
		* ) PFT2OPTS="$PFT2OPTS $1" && shift
	esac
done

# Calculate the size set: By default use all the particles (as pft2).
NUM_PART=$(bmg -ver 2 $INSTAR | grep Totals | cut -f2 | sed 's/ ..*$//')
SIZE_SETS=$(echo "($NUM_PART / $NUM_PROC) + 1" | bc)
bpartsel -all -sets $SIZE_SETS -out sets.star $INSTAR

# Remove previous FOM, since those will be used by bpartmulti.
# Avoid that orientations from previous runs were part of the end STAR result.
bstar -floatscale particle.fom,0,-1 -out sets.star sets.star

# Spawn pft2 jobs. It is necessary to write a different file per pft2 process
for ((i=1, j=0; i<=NUM_PROC; i++, j++)); do
	LIST_STAR="$LIST_STAR ${OUTSTAR%.star}_$i.star"
	command="pft2 $PFT2OPTS -select $i -output ${LIST_STAR##* } sets.star"
	[[ $NODES ]] && command="ssh ${NODES[$j]} \"cd ${WORKDIR} && source ${RC} && $command \""
	[[ $VERBOSE -gt 0 ]] && echo $command
	eval $command &
done

# Wait for individual STAR files
wait

# Merge the individual STAR files
bpartmulti -merge -output $OUTSTAR $LIST_STAR
bpartsel -all -output $OUTSTAR $OUTSTAR

# Remove all extra files
if [[ ! $KEEP ]] ;then
	rm sets.star
	rm $LIST_STAR
fi

exit 0