#!/bin/bash

#    breconstructp
#    A script to parallelize breconstruct
#    Author: Eduardo Sanz-García
#    Created: 20090417    Modified: 20090420

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

#    Considerations:
#           1- Log files (breconstructp ... > log) doesn't make as much sense as normal breconstruct log files.
#              This is because multiple processes are writing at the same time in the 
#              same file
#           2- The splitted STAR file can be not totally equilibrated (same number of particle in each), hence
#              some weighting in badd should be done.

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

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

	Usage: ${0##*/} [options] input.star [input.star]
	------------------------------------------------------------
	A script to parallelize breconstruct in multicore machines.
	breconstructp accepts the usual breconstruct options, except
	-fom, -bootstrap and -number, these two last, since selection
	numbers are use to paralellize breconstruct.
	
	Additionally breconstructp is able to catch the environmental
	variable PROCESSORS.

	For more sophisticated usages, like clusters, breconstructp 
	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
OUTRECONS=                   # Output reconstruction
LIST_STAR=                   # List of all the individual STAR files (one per process)
BRECONSOPTS=                 # BRECONS 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 ;;
		-reco* ) OUTRECONS="$2" && EXT=${2##*.} && shift 2 ;;
		-al*   ) ALL=1 && BRECONSOPTS="$BRECONSOPTS $1" && shift ;;
		-ke*   ) KEEP=1 && shift ;;
		-mul*  ) MUL=$2 && BRECONSOPTS="$BRECONSOPTS $1 $2" && shift 2 ;;
		-num*  ) warning "Since the selection numbers are use to paralellize"
				 warning "breconstruct, -number option is not acepted!"
				 exit -1 ;;
		-bo*   ) warning "-bootstrap option not implemented jet!"
				 exit -1 ;;
		-fo*   ) warning "-fom option not implemented jet!"
				 exit -1 ;;
		-ver*  ) VERBOSE="$2" && BRECONSOPTS="$BRECONSOPTS $1 $2" && shift 2 ;;
		*.star ) INSTAR="$INSTAR $1" && shift ;;
		* ) BRECONSOPTS="$BRECONSOPTS $1" && shift
	esac
done

# Calculate the size set (only use selected particles, except when -all option is used).
if [[ $ALL ]]; then
	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
else
	NUM_PART=$(bmg -ver 2 $INSTAR | grep Totals | cut -f3 | sed 's/ ..*$//')
	SIZE_SETS=$(echo "($NUM_PART / $NUM_PROC) + 1" | bc)
	bpartsel -sets $SIZE_SETS -out sets.star $INSTAR
fi

# Spawn breconstruct jobs. It is necessary to write a different file per breconstruct process
for ((i=1, j=0; i<=NUM_PROC; i++, j++)); do
	LIST_STAR="$LIST_STAR ${OUTSTAR%.star}_$i.star"
	LIST_RECONS="$LIST_RECONS ${OUTRECONS%.$EXT}_$i.$EXT"
	if [[ $OUTSTAR ]]; then
		command="breconstruct $BRECONSOPTS -number $i -reconstruction ${LIST_RECONS##* } -output ${LIST_STAR##* } sets.star"
	else
		command="breconstruct $BRECONSOPTS -number $i -reconstruction ${LIST_RECONS##* } sets.star"
	fi
	[[ $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 indiviual reconstructions
if [[ $MUL ]]; then
	tmp=$MUL
	while (( MUL > 0 ));do
		badd -output ${OUTRECONS%.$EXT}_0$MUL.$EXT ${LIST_RECONS//.$EXT/_0$MUL.$EXT}
		(( MUL -= 1 ))
	done
	MUL=$tmp
else
	badd -output $OUTRECONS $LIST_RECONS
fi

# Merge the individual STAR files
if [[ $OUTSTAR ]]; then
	bpartmulti -merge -output $OUTSTAR $LIST_STAR
	bpartsel -consolidate -output $OUTSTAR $OUTSTAR
fi

# Remove all extra files
if [[ ! $KEEP ]] ;then
	rm sets.star
	if [[ $MUL ]]; then
		while (( $MUL > 0 ));do
			rm ${LIST_RECONS//.$EXT/_0$MUL.$EXT}
			(( MUL -= 1))
		done
	else
		rm $LIST_RECONS
	fi
	[[ $OUTSTAR ]] && rm $LIST_STAR
fi

exit 0