#!/bin/bash

########################################################################
##						Version 2.1 (18.04.2018)					  ##
##																	  ##
##  Minimal pipeline for preprocessing.								  ##
## It includes the following steps:									  ##
##		Data conversion to nifti									  ##
##		Computation of realignment of sms (GRE_nosms as reference)	  ##
##		Computation of fieldmap for unwarping (from SE_nosms)		  ##
##		Realignment and unwarping (1st resampling)					  ##
##			If cutting is desired, it will be done in this step		  ##
##		Brain extraction (BET based on FEAT)						  ##
##		High-pass filtering (based on FEAT)							  ##
##		Optional: cutting (if no unwarping was performed)			  ##
########################################################################


#-----------------------------------------------------------------------
# Information
#-----------------------------------------------------------------------
Author="Jorge Manuel"
Name="JPreprocessing"
Version="2.1"

Usage() {
cat <<- EOF
	Minimal pipeline for preprocessing.
	
	Usage:
	   JPreprocessing -i <input> -o <output> [options]

	Input:
	   Folder with Dicom images. If the images have been already converted
	   and are in the raw directory, just use "bla" or similar as input.
	   The program will detect the images and skip the import.
	
	Output:
	   Destination directory.

	Options:
	   -bet [BB]
	      Define brain extraction. BB is the brain-background threshold.
	      Default: -bet 0.1
	   -filter [HP_cutoff] [LP_cutoff]
	      Define temporal filter. The cutoffs are given in seconds (FWHM).
	      Put -1 in the cutoff to avoid performing it.
	      Default: -filter 200 -1
	   -cut [t_start] [t_size]
	      Initial and final volumes can be removed. The numeration follows
	      fslroi. Volumes start at 0 and t_size=-1 gives the full image extent.
	      Default: -cut 0 -1
	   --nocleanup
	      If this option is selected, none of the intermediary functional
	      images will be deleted (i.e. mcflirt/sms, topup/sms and bet/sms).
EOF
	exit 1
}


#-----------------------------------------------------------------------
# Parse arguments
#-----------------------------------------------------------------------
if [ $# == 0 ] ; then
    Usage
    exit 1;
fi

#Default values
BB_thresh=0.1
HP_cutoff=200
LP_cutoff=-1
t_start=0
t_size=-1
CleanUp=1

#Parser
while [[ $# -gt 0 ]]; do
	key="$1"

	case $key in
	-i)
		Input="$2"
		shift	#Pass argument
		;;
	-o)
		Dir="$2"
		shift
		;;
	-bet)
		BB_thresh="$2"
		shift
		;;
	-filter)
		HP_cutoff="$2"
		LP_cutoff="$3"
		shift 2
		;;
	-cut)
		t_start="$2"
		t_size="$3"
		shift 2
		;;
	--nocleanup)
		CleanUp=0
		;;
	-h)
		Usage
		exit 
		;;
	*)
		echo "${1} unknown option"
		Usage
		exit
		;;
	esac
	shift #Pass argument or value
done

{
#-----------------------------------------------------------------------
# Create folder structure
#-----------------------------------------------------------------------
mkdir -p \
	${Dir}/anat \
	${Dir}/prepro/bet \
	${Dir}/prepro/filter \
	${Dir}/prepro/mcflirt \
	${Dir}/prepro/topup \
	${Dir}/raw


#-----------------------------------------------------------------------
# Create log-files
#-----------------------------------------------------------------------
printf \
	'%s\n\tVersion: %s\n\n' \
	${Name} ${Version} \
	> ${Dir}/Log.txt
printf \
	'Used command:\n\t%s \\\n' \
	${Name} \
	>> ${Dir}/Log.txt
printf \
	'\t\t-i %s \\\n' \
	${Input} \
	>> ${Dir}/Log.txt
printf \
	'\t\t-o %s \\\n' \
	${Dir} \
	>> ${Dir}/Log.txt
printf \
	'\t\t-bet %s \\\n' \
	${BB_thresh} \
	>> ${Dir}/Log.txt
printf \
	'\t\t-filter %s %s \\\n' \
	${HP_cutoff} ${LP_cutoff} \
	>> ${Dir}/Log.txt
printf \
	'\t\t-cut %s %s' \
	${t_start} ${t_size} \
	>> ${Dir}/Log.txt
if [ ${CleanUp} == 1 ]; then
	printf \
		'\n\n' \
		>> ${Dir}/Log.txt
elif [ ${CleanUp} == 0 ]; then
	printf \
		' \\\n\t\t-nocleanup\n\n' \
		>> ${Dir}/Log.txt
fi
printf \
	'Log:\n' \
	>> ${Dir}/Log.txt



#-----------------------------------------------------------------------
# Import data
#-----------------------------------------------------------------------
if [ $(ls ${Dir}/raw | wc -l) -gt 0 ]; then
	printf " Import data:\tomitted (folder not empty)\n" >> ${Dir}/Log.txt
else
	printf " Import data:\tstarted\n" >> ${Dir}/Log.txt
	mkdir -p ${Dir}/raw/temp

	#Conversion
	dcm2nii -o ${Dir}/raw/temp ${Input} > ${Dir}/raw/Log.txt

	#Structural
	fslswapdim \
		${Dir}/raw/temp/*T1* \
		z -x y \
		${Dir}/raw/t1 \
		1> /dev/null
	fslorient \
		-forceneurological \
		${Dir}/raw/t1
	fslswapdim \
		${Dir}/raw/temp/*T2* \
		z -x y \
		${Dir}/raw/t2 \
		1> /dev/null
	fslorient \
		-forceneurological \
		${Dir}/raw/t2

	#GRE_nosms
	if [ $(ls ${Dir}/raw/temp/*GREnosms* | wc -l) -gt 1 ]; then
		echo "There is more than one GRE_nosms sequence."
		read -p "Please rename manually, extract the first volume and then press enter."
	else
		mv ${Dir}/raw/temp/*GREnosms* ${Dir}/raw/GRE_nosms.nii.gz
		fslroi \
			${Dir}/raw/GRE_nosms.nii.gz \
			${Dir}/raw/GRE_nosms_vol1.nii.gz \
			0 1
	fi

	#SE_nosms
	if [ $(ls ${Dir}/raw/temp/*SEnosms*AP* | wc -l) -gt 1 ] || [ $(ls ${Dir}/raw/temp/*SEnosms*PA* | wc -l) -gt 1 ]; then
		echo "There is more than one SE_nosms_* sequence."
		read -p "Please rename manually, merge the files and then press enter."
	else
		fslmerge \
				-t \
				${Dir}/raw/SE_nosms_AP_PA \
				${Dir}/raw/temp/*SEnosms*AP* ${Dir}/raw/temp/*SEnosms*PA*
	fi

	#sms
	n=$(ls ${Dir}/raw/temp/*GREsms* | wc -l) #Number of functional scans
	for i in $(seq $n); do
		fslmaths \
			$(ls ${Dir}/raw/temp/*GREsms* -rt | sed "${i}q;d") \
			${Dir}/raw/sms_${i}
	done

	#Cleanup
	rm -r ${Dir}/raw/temp
	printf " Import data:\tcompleted\n" >> ${Dir}/Log.txt
fi


#-----------------------------------------------------------------------
# Sequence info
#-----------------------------------------------------------------------
T_R=$(fslinfo ${Dir}/raw/sms_1 | grep pixdim4 | rev | cut -f1 -d" " | rev)
y_size=$(fslinfo ${Dir}/raw/sms_1 | grep -w dim2 | rev | cut -f1 -d" " | rev)
z_size=$(fslinfo ${Dir}/raw/sms_1 | grep -w dim3 | rev | cut -f1 -d" " | rev)
n=$(ls ${Dir}/raw/*sms_?.nii.gz | wc -l) #Number of functional scans


#-----------------------------------------------------------------------
# Motion correction
#-----------------------------------------------------------------------
function Realignment {
	printf " Realignment_%s:\tstarted\n" ${i} >> ${Dir}/Log.txt
	mkdir -p ${Dir}/prepro/mcflirt/sms_${i}

	#Realign
	mcflirt \
		-in ${Dir}/raw/sms_${i} \
		-mats -plots -stats -rmsrel -rmsabs -spline_final \
		-reffile ${Dir}/raw/GRE_nosms_vol1 \
		-out ${Dir}/prepro/mcflirt/sms_${i}/sms

	#Plot movement
	fsl_tsplot \
		-i ${Dir}/prepro/mcflirt/sms_${i}/*.par \
		-t 'MCFLIRT estimated rotations (radians)' \
		-u 1 \
		--start=1 \
		--finish=3 \
		-a x,y,z \
		-w 640 \
		-h 144 \
		-o ${Dir}/prepro/mcflirt/sms_${i}/rot.png
	fsl_tsplot \
		-i ${Dir}/prepro/mcflirt/sms_${i}/*.par \
		-t 'MCFLIRT estimated translations (mm)' \
		-u 1 \
		--start=4 \
		--finish=6 \
		-a x,y,z \
		-w 640 \
		-h 144 \
		-o ${Dir}/prepro/mcflirt/sms_${i}/trans.png
	fsl_tsplot \
		-i ${Dir}/prepro/mcflirt/sms_${i}/sms_abs.rms,${Dir}/prepro/mcflirt/sms_${i}/sms_rel.rms \
		-t 'MCFLIRT estimated mean displacement (mm)' \
		-u 1 \
		-w 640 \
		-h 144 \
		-a absolute,relative \
		-o ${Dir}/prepro/mcflirt/sms_${i}/disp.png

	printf " Realignment_%s:\tcompleted\n" ${i} >> ${Dir}/Log.txt
}

if [ "$(ls -A ${Dir}/prepro/mcflirt)" ]; then
	printf " Realignment:\tomitted (folder not empty)\n" >> ${Dir}/Log.txt
else
	for i in $(seq $n); do
		Realignment &
	done
fi
wait


#-----------------------------------------------------------------------
# Unwarping
#-----------------------------------------------------------------------
function JUnwarp_Fieldmap {
	printf " Fieldmap:\tstarted\n" >> ${Dir}/Log.txt
	#Acquisition parameters
	Dwell_time=$(cat ${Dir}/raw/Log.txt | grep "echo spacing:" | tail -n1 | cut -f4 -d" " | cut -f1 -d"m")
	Time=$(echo "scale=8; ${Dwell_time}*(${y_size}-1)/1000" | bc) #Total readout time (defined as the time from the centre of the first echo to the centre of the last) in seconds
	printf "0 -1 0 %s\n0 -1 0 %s\n0 1 0 %s\n0 1 0 %s\n" \
		"0${Time}" "0${Time}" "0${Time}" "0${Time}" \
		> ${Dir}/prepro/topup/acqparams.txt
	
	#Generate fieldmap
	topup \
		--imain=${Dir}/raw/SE_nosms_AP_PA.nii.gz \
		--datain=${Dir}/prepro/topup/acqparams.txt \
		--config=b02b0.cnf \
		--fout=${Dir}/prepro/topup/fieldmap \
		--out=${Dir}/prepro/topup/topup_results
	mv ${Dir}/raw/SE_nosms_AP_PA.topup_log ${Dir}/prepro/topup/
	
	#Unwarp GRE_nosms_vol1
	applytopup \
		-i ${Dir}/raw/GRE_nosms_vol1.nii.gz \
		-a ${Dir}/prepro/topup/acqparams.txt \
		-x 1 \
		-t ${Dir}/prepro/topup/topup_results \
		-o ${Dir}/prepro/topup/GRE_nosms_vol1_unwarped \
		-m jac
	
	printf " Fieldmap:\tcompleted\n" >> ${Dir}/Log.txt
}

function JUnwarp_Apply {
	printf " Unwarp_%s:\tstarted\n" ${i} >> ${Dir}/Log.txt
	mkdir -p \
		${Dir}/prepro/topup/sms_${i}/sms \
		${Dir}/prepro/topup/sms_${i}/sms_trafo
	
	#Prepare sms
	fslsplit \
		${Dir}/raw/sms_${i} \
		${Dir}/prepro/topup/sms_${i}/sms/vol
	
	#Calculate t_end for cutting purposes
	if [ ${t_size} == -1 ]; then
		t_size=$(fslinfo ${Dir}/raw/sms_${i} | grep -w dim4 | rev | cut -f1 -d" " | rev)
	fi
	t_end=$(echo "${t_start}+${t_size}-1" | bc)
	
	#Apply unwarping and realingment to sms
	for k in $(seq ${t_start} ${t_end}); do
		j=$(printf "%04d" ${k})
		mkdir ${Dir}/prepro/topup/sms_${i}/sms_trafo/vol${j}
		Line=$(echo "${j}+1" | bc)
		arr=($(sed ${Line}'!d' ${Dir}/prepro/mcflirt/sms_${i}/sms.par))
		echo \
			${arr[3]} ${arr[4]} ${arr[5]} ${arr[0]} ${arr[1]} ${arr[2]} \
			> ${Dir}/prepro/topup/sms_${i}/sms_trafo/vol${j}/topup_results_movpar.txt &
		cp \
			${Dir}/prepro/topup/topup_results_fieldcoef.nii.gz \
			${Dir}/prepro/topup/sms_${i}/sms_trafo/vol${j}/topup_results_fieldcoef.nii.gz &
	done

	wait
	for k in $(seq ${t_start} ${t_end}); do
		j=$(printf "%04d" ${k})
		applytopup \
		-i ${Dir}/prepro/topup/sms_${i}/sms/vol${j}.nii.gz \
		-a ${Dir}/prepro/topup/acqparams.txt \
		-x 1 \
		-t ${Dir}/prepro/topup/sms_${i}/sms_trafo/vol${j}/topup_results \
		-o ${Dir}/prepro/topup/sms_${i}/sms/vol${j}_unwarped.nii.gz \
		-m jac &
	done
	wait
	
	fslmerge \
		-t \
		${Dir}/prepro/topup/sms_${i}/sms_unwarped \
		${Dir}/prepro/topup/sms_${i}/sms/vol*_unwarped.nii.gz
	rm -r \
		${Dir}/prepro/topup/sms_${i}/sms \
		${Dir}/prepro/topup/sms_${i}/sms_trafo
	
	printf " Unwarp_%s:\tcompleted\n" ${i} >> ${Dir}/Log.txt
}

if [ "$(ls -A ${Dir}/prepro/topup)" ]; then
	printf " Unwarping:\tomitted (folder not empty)\n" >> ${Dir}/Log.txt
elif [ ! -f ${Dir}/raw/SE_nosms_AP_PA.nii.gz ]; then
	printf " Unwarping:\tomitted (no reference image)\n" >> ${Dir}/Log.txt
else
	JUnwarp_Fieldmap
	for i in $(seq $n); do
		JUnwarp_Apply
	done
fi


#-----------------------------------------------------------------------
# Brain extraction
#-----------------------------------------------------------------------
function JBET_Mask {
	printf " Mask:\tstarted\n" >> ${Dir}/Log.txt
	#Create and apply first mask
	for i in $(seq $n); do
		fslmaths \
			${Dir}/${input/'$i'/$i} \
			-Tmean \
			${Dir}/prepro/bet/mean_func_${i} &
	done
	wait
	
	fslmerge \
		-t \
		${Dir}/prepro/bet/mean_func \
		${Dir}/prepro/bet/mean_func_*
	rm ${Dir}/prepro/bet/mean_func_*
	fslmaths \
		${Dir}/prepro/bet/mean_func \
		-Tmean \
		${Dir}/prepro/bet/mean_func
	
	bet2 \
		${Dir}/prepro/bet/mean_func \
		${Dir}/prepro/bet/mask \
		-f 0.3 \
		-n \
		-m
	mv ${Dir}/prepro/bet/mask_mask.nii.gz ${Dir}/prepro/bet/mask.nii.gz
	rm ${Dir}/prepro/bet/mean_func.nii.gz
	
	for i in $(seq $n); do
		fslmaths \
			${Dir}/${input/'$i'/$i} \
			-mas \
			${Dir}/prepro/bet/mask \
			${Dir}/prepro/bet/sms_bet_${i} &
	done
	wait
	
	#Calculate intensity within the mask
	fslmerge \
		-t \
		${Dir}/prepro/bet/sms_bet \
		${Dir}/prepro/bet/sms_bet_*
	rm ${Dir}/prepro/bet/sms_bet_*
	I_masked=$(fslstats \
		${Dir}/prepro/bet/sms_bet \
		-p 2 -p 98)
	I_masked=$(echo ${I_masked} | cut -f2 -d" ")
	
	#Improve the mask using the brain/background threshold
	fslmaths \
		${Dir}/prepro/bet/sms_bet \
		-thr $(echo "scale=7; ${I_masked}*${BB_thresh}" | bc) \
		-Tmin \
		-bin ${Dir}/prepro/bet/mask \
		-odt char
	
	#Dilate the mask to remove holes
	fslmaths \
		${Dir}/prepro/bet/mask \
		-dilF \
		${Dir}/anat/wb_mask
	
	#Rename the mask, so that it is clear that it is undilated
	mv ${Dir}/prepro/bet/mask.nii.gz ${Dir}/prepro/bet/undilated_mask.nii.gz
	
	rm ${Dir}/prepro/bet/sms_bet.nii.gz
	
	printf " Mask:\tcompleted\n" >> ${Dir}/Log.txt
}

function JBET_Apply {
	printf " BET_%s:\tstarted\n" ${i} >> ${Dir}/Log.txt
	#Calculate the mean intensity of the masked data (for normalisation later)
	#Note how it is done with the undilated mask, like in FEAT
	I_mean=$(fslstats \
		${Dir}/${input/'$i'/$i} \
		-k ${Dir}/prepro/bet/undilated_mask \
		-p 50)
	
	
	#Apply the mask and set mean intensity (of the 4D image) to 10000.
	fslmaths \
		${Dir}/${input/'$i'/$i} \
		-mas ${Dir}/anat/wb_mask \
		-mul $(echo "scale=15; 10000/${I_mean}" | bc) \
		${Dir}/prepro/bet/sms_bet_${i}
	
	printf " BET_%s:\tcompleted\n" ${i} >> ${Dir}/Log.txt
}

#If the images were unwarped, then use those images, otherwise the realigned.
if [ "$(ls -A ${Dir}/prepro/topup)" ]; then
	input='prepro/topup/sms_$i/sms_unwarped.nii.gz'
else
	input='prepro/mcflirt/sms_$i/sms.nii.gz'
fi

if [ "$(ls -A ${Dir}/prepro/bet)" ]; then
	printf " BET:\tomitted (folder not empty)\n" >> ${Dir}/Log.txt
else
	JBET_Mask
	for i in $(seq $n); do
		JBET_Apply &
	done
	wait
fi


#-----------------------------------------------------------------------
# Temporal filtering
#-----------------------------------------------------------------------
function JFiltering {
	printf " Filter_%s:\tstarted\n" ${i} >> ${Dir}/Log.txt
	#Calculate mean to restore it after the time filtering.
	fslmaths \
		${Dir}/prepro/bet/sms_bet_${i} \
		-Tmean \
		${Dir}/prepro/filter/tempMean_${i}

	#Filtering
	if [ ${HP_cutoff} == -1 ]; then
		HP_thresh=-1
	else
		HP_thresh=$(echo "scale=15; ${HP_cutoff}/2/${T_R}" | bc)
	fi
	if [ ${LP_cutoff} == -1 ]; then
		LP_thresh=-1
	else
		LP_thresh=$(echo "scale=15; ${LP_cutoff}/2/${T_R}" | bc)
	fi
	fslmaths \
		${Dir}/prepro/bet/sms_bet_${i} \
		-bptf ${HP_thresh} ${LP_thresh} \
		-add ${Dir}/prepro/filter/tempMean_${i} \
		${Dir}/prepro/filter/filtered_func_data_${i}

	rm ${Dir}/prepro/filter/tempMean_${i}.nii.gz
	printf " Filter_%s:\tcompleted\n" ${i} >> ${Dir}/Log.txt
}

if [ "$(ls -A ${Dir}/prepro/filter)" ]; then
	printf " Filter:\tomitted (folder not empty)\n" >> ${Dir}/Log.txt
else
	for i in $(seq $n); do
		JFiltering &
	done
	wait
fi


#-----------------------------------------------------------------------
# Cutting, if not performed during unwarping
#-----------------------------------------------------------------------
function JCutting {
	printf " Cutting_%s:\tstarted\n" ${i} >> ${Dir}/Log.txt

	fslroi \
		${Dir}/prepro/filter/filtered_func_data_${i} \
		${Dir}/prepro/filter/filtered_func_data_${i} \
		${t_start} ${t_size}

	printf " Cutting_%s:\tcompleted\n" ${i} >> ${Dir}/Log.txt
}

for i in $(seq $n); do
	Volumes=$(fslinfo ${Dir}/prepro/filter/filtered_func_data_${i} | grep -w dim4 | rev | cut -f1 -d" " | rev)
	if [ ${Volumes} -gt ${t_size} ]; then
		JCutting &
	fi
	wait
done


#-----------------------------------------------------------------------
# Cleanup
#-----------------------------------------------------------------------
find ${Dir} -type d -empty -delete
if [ ${CleanUp} == 1 ]; then
	for i in $(seq $n); do
		rm ${Dir}/prepro/mcflirt/sms_${i}/sms.nii.gz
		rm -r ${Dir}/prepro/topup/sms_${i}
		rm ${Dir}/prepro/bet/sms_bet_${i}.nii.gz
	done
fi

} 2> ${Dir}/Log_err.txt