CCSegLongPreprocess

#!/usr/bin/python

# command line processing

import os
import sys
import numpy
import getopt
import errno

import shutil
import subprocess

import CCSegUtils

import CCSegSubjectSelection

import CCSegLong

import distutils.spawn

import FLIRT

import tempfile
import nibabel

def criticalError(errorString):
	print "Error: " + str(errorString)
	quit()

def usage():
	print "Usage: " + sys.argv[0] + " [options] <input directory> <output directory>"
	print
	print "\tDESCRIPTION"
	print "\t\tLongitudinal preprocessing script. Registers all images to first timepoint for each subject then all images to the MNI152 template"
	print "\tExample: for each subject <foo>, there will be the following files in <input directory>"
	print "\t\t<foo>_tp1.nii.gz, <foo>_tp2.nii.gz, <foo>_tp3.nii.gz, ..."
	print "(note the suffixes dont have to be in the _tpXX format, this is just an example)"
	print "\t\tThere must be a line in the long-subject-timepoint-file file containing"
	print "\t\t<foo>_tp1 <foo>_tp2 <foo>_tp3 ..."
	print "\tIn the <output directory>, this program will produce:"
	print "\t\t<foo>_tp1_to_std.mat: FLIRT matrix transforming <foo>_tp1 to the MNI152 template"
	print "\t\tThen for each of the other timepoint files <foo>_tpX"
	print "\t\t<foo>_tpX_to_<foo>_tp1.mat: FLIRT matrix transforming <foo>_tpX to the first timepoint <foo>_tp1"
	print "\t\t<foo>_tpX_to_std.mat: FLIRT matrix transforming <foo>_tpX to the MNI152 template"

	print "\t\t"

	print
	print "\tMANDATORY OPTIONS"
	print CCSegLong.longOptionText()
	print "\tOTHER OPTIONS"
	print "\t\t--save-reg-to-first: save an image that is the result of registering all images for a subject to the first timepoint"
	print "\t\t\tAs per the example above, saves <foo>_tpX_to_<foo>_tp1.nii.gz"
	print "\t\t--save-reg-to-std: save an image that is the result of registering all images for a subject to the MNI152 template"
	print "\t\t\tAs per the example above, saves <foo>_tpX_to_std.nii.gz"
	print "\t\t--parallel-mode=: which mechanism to use to run the registrations in parallel, supported values are:"
	print "\t\t\tparallel (default if installed): Uses the GNU parallel program to use all host CPUs"
	print "\t\t\tnone: No parallel, runs all registrations in serial"
	print
	print CCSegLong.longOptionFileFormat()


	#CCSegSubjectSelection.printHelpText()

def criticalError(errorString):
	print "Error: " + str(errorString)
	quit()

def doFLIRT(inputDirectory, outputDirectory, mainAndTimePoints, saveRegToStd, saveRegToFirst, parallelMode, useSQForm = True):
	D = 15

	flirtCost = 'mutualinfo'
	flirtInterp = 'trilinear'

	flirtBase = [os.environ['FSLDIR'] + '/bin/flirt', '-searchrx', str(-D), str(D), '-searchry', str(-D), str(D), '-searchrz', str(-D), str(D), '-cost', flirtCost, '-interp', flirtInterp]
	#flirtBase = [os.environ['FSLDIR'] + '/bin/flirt', '-searchrx', str(-D), str(D), '-searchry', str(-D), str(D), '-searchrz', str(-D), str(D), '-cost', flirtCost, '-interp', flirtInterp]
#flirtBase = ['qsub', '-b', 'y', '-q', 'linux.q', '-V', '-cwd', '-N', 'IMAGEHD', os.environ['FSLDIR'] + '/bin/flirt', '-dof', '6', '-searchrx', str(-D), str(D), '-searchry', str(-D), str(D), '-searchrz', str(-D), str(D), '-cost', flirtCost, '-interp', flirtInterp]
	
	parallelSwitch = '-j+0'

	if parallelMode == 'parallel':
		parallelPipe = subprocess.Popen(['parallel', parallelSwitch], stdin = subprocess.PIPE, stdout = subprocess.PIPE)

	QSUBBase = ['qsub', '-b', 'y', '-N', 'ADNI', '-V', '-cwd', '-q', 'freesurfer.q']
	#fp = open('commands', 'w')

	for curSubject in sorted(mainAndTimePoints.iterkeys()):
		refFile = os.path.join(inputDirectory, curSubject + '.nii.gz')

		if mainAndTimePoints[curSubject] != None:
			for curTimePoint in range(len(mainAndTimePoints[curSubject])):
				inFile = os.path.join(inputDirectory, mainAndTimePoints[curSubject][curTimePoint])
				outPrefix = os.path.join(outputDirectory, mainAndTimePoints[curSubject][curTimePoint] + '_to_' + curSubject)
				
				if not os.path.isfile(outPrefix + '.mat'):
					#print outPrefix + '.mat' + " adding"
					#commandString = flirtApplyXFMBase[:]
					#commandString.extend(['-in', inFile, '-ref', refFile, '-out', outPrefix, '-init', outPrefix + '.mat'])
				#else:
					commandString = flirtBase[:]
					commandString.extend(['-in', inFile, '-ref', refFile, '-omat', outPrefix + '.mat', '-dof', '6'])
					if useSQForm == True:
						commandString.append('-usesqform')
					if saveRegToFirst == True:
						commandString.extend(['-out', outPrefix])
				
					if parallelMode == 'none':
						subprocess.call(commandString)
					if parallelMode == 'parallel':
						parallelPipe.stdin.write(" ".join(commandString) + "\n")
					if parallelMode == 'qsub':
						curCommand = QSUBBase[:]
						curCommand.extend(commandString)
						subprocess.call(curCommand)

				#print " ".join(commandString)
				#fp.write(" ".join(commandString) + "\n")
				#os.system(commandString)
				#subprocess.call(commandString)

	refFile = CCSegUtils.MNI152FLIRTTemplate()
# then register the first timepoint to MNI152

	for curSubject in sorted(mainAndTimePoints.iterkeys()):
		inFile = os.path.join(inputDirectory, curSubject)
		outPrefix = os.path.join(outputDirectory, curSubject + '_to_std')
		
		if not os.path.isfile(outPrefix + '.mat') or useSQForm == False:
			#commandString = flirtApplyXFMBase[:]
			#commandString.extend(['-in', inFile, '-ref', refFile, '-out', outPrefix, '-init', outPrefix + '.mat'])
		#else:
			commandString = flirtBase[:]
			commandString.extend(['-in', inFile, '-ref', refFile, '-omat', outPrefix + '.mat', '-dof', '7'])
			if useSQForm == True:
				commandString.append('-usesqform')
			if saveRegToStd == True:
				commandString.extend(['-out', outPrefix])

			if parallelMode == 'none':
				subprocess.call(commandString)
			if parallelMode == 'parallel':
				print " ".join(commandString)
				parallelPipe.stdin.write(" ".join(commandString) + "\n")
			if parallelMode == 'qsub':
				curCommand = QSUBBase[:]
				curCommand.extend(commandString)
				subprocess.call(curCommand)
		#fp.write(" ".join(commandString) + "\n")
#fp.close()
	if parallelMode == 'parallel':
		print "parallel, running registrations to first timepoints and template"
		parallelPipe.communicate()
		parallelPipe.stdin.close()
		print "parallel, finished"
		
	if parallelMode == 'qsub':
		quit()
	#if saveRegToStd == True:
	#	if parallelMode == 'parallel':
	#		parallelPipe = subprocess.Popen(['parallel', parallelSwitch], stdin = subprocess.PIPE, stdout = subprocess.PIPE)
	#		parallelPipe.communicate()
	#		parallelPipe.stdin.close()


def main():
	
	longOptStrings = ["long-subject-timepoint-file=", 'save-reg-to-first', 'save-reg-to-std', 'parallel-mode=']
	longOptStrings.extend(CCSegSubjectSelection.getOptions())	
	opts, args = getopt.getopt(sys.argv[1:], "h", longOptStrings)
	
	# initialise options

	# mandatory arguments
	inputDirectory = None
	outputDirectory = None

	longSubjectTimepointFile = None
	saveRegToFirst = False
	saveRegToStd = False
	# do the mandatory arguments first
	
	parallelMode = 'parallel'

	if len(args) != 2:
		print "The number of arguments must be 2"
		usage()
		exit()
	numpy.set_printoptions(precision = 3, formatter = {'all':lambda x: "%.3f" % x})

	inputDirectory = args[0]
	outputDirectory = args[1]
	
	#print "inputDirectory = " + inputDirectory
	#print "outputDirectory = " + outputDirectory

	for o, a in opts:
		if o == '-h':
			usage()
			quit()
		if o == '--long-subject-timepoint-file':
			longSubjectTimepointFile = a
		if o == '--parallel-mode':
			parallelMode = a
		if o == '--save-reg-to-first':
			saveRegToFirst = True
		if o == '--save-reg-to-std':
			saveRegToStd = True

	#for o, a in opts:
	
	if not os.path.isdir(inputDirectory):
		criticalError("The input directory specified was not found")
	
	if longSubjectTimepointFile == None:
		criticalError("The longitudinal file name was not given in the option --long-subject-timepoint-file")

	parallelModes = ['none', 'parallel', 'qsub']

	if not parallelMode in parallelModes:
		criticalError("Parallel mode not supported: " + parallelMode)
	
	if parallelMode == 'parallel':
		parallelExecutable = distutils.spawn.find_executable('parallel')
		if parallelExecutable == None:
			criticalError("Parallel mode was \"parallel\" but the parallel binary was not found")
	
	mainAndTimePoints, timePointsToReference = CCSegLong.loadLongSubjectFile(longSubjectTimepointFile)
	
	for curFile in timePointsToReference:
		if CCSegUtils.findNIFTIFromPrefix(os.path.join(inputDirectory, curFile)) == None:
			criticalError("One of the files in the long file was not in the input directory: " + curFile)

	try:
		os.makedirs(outputDirectory)
	except OSError as exc: # Python >2.5
		if exc.errno == errno.EEXIST and os.path.isdir(outputDirectory):
			pass
		else:
			raise Exception

	doFLIRT(inputDirectory, outputDirectory, mainAndTimePoints, saveRegToStd, saveRegToFirst, parallelMode, useSQForm = True)
	
	flirtInterp = 'trilinear'
	
	flirtApplyXFMBase = [os.environ['FSLDIR'] + '/bin/flirt', '-interp', flirtInterp, '-applyxfm']
	convertXFMBase = [os.environ['FSLDIR'] + '/bin/convert_xfm']
	
#	# set up initial commands for FLIRT
#	D = 15
#
#	flirtCost = 'mutualinfo'
#	flirtInterp = 'trilinear'
#
#	flirtBase = [os.environ['FSLDIR'] + '/bin/flirt', '-searchrx', str(-D), str(D), '-searchry', str(-D), str(D), '-searchrz', str(-D), str(D), '-cost', flirtCost, '-interp', flirtInterp]
#	#flirtBase = [os.environ['FSLDIR'] + '/bin/flirt', '-searchrx', str(-D), str(D), '-searchry', str(-D), str(D), '-searchrz', str(-D), str(D), '-cost', flirtCost, '-interp', flirtInterp]
##flirtBase = ['qsub', '-b', 'y', '-q', 'linux.q', '-V', '-cwd', '-N', 'IMAGEHD', os.environ['FSLDIR'] + '/bin/flirt', '-dof', '6', '-searchrx', str(-D), str(D), '-searchry', str(-D), str(D), '-searchrz', str(-D), str(D), '-cost', flirtCost, '-interp', flirtInterp]
#	
#	parallelSwitch = '-j+0'
#
#	if parallelMode == 'parallel':
#		parallelPipe = subprocess.Popen(['parallel', parallelSwitch], stdin = subprocess.PIPE, stdout = subprocess.PIPE)
#
#	#fp = open('commands', 'w')
#
#	for curSubject in sorted(mainAndTimePoints.iterkeys()):
#		refFile = os.path.join(inputDirectory, curSubject + '.nii.gz')
#
#		if mainAndTimePoints[curSubject] != None:
#			for curTimePoint in range(len(mainAndTimePoints[curSubject])):
#				inFile = os.path.join(inputDirectory, mainAndTimePoints[curSubject][curTimePoint])
#				outPrefix = os.path.join(outputDirectory, mainAndTimePoints[curSubject][curTimePoint] + '_to_' + curSubject)
#				
#				if not os.path.isfile(outPrefix + '.mat'):
#					#commandString = flirtApplyXFMBase[:]
#					#commandString.extend(['-in', inFile, '-ref', refFile, '-out', outPrefix, '-init', outPrefix + '.mat'])
#				#else:
#					commandString = flirtBase[:]
#					commandString.extend(['-in', inFile, '-ref', refFile, '-omat', outPrefix + '.mat', '-dof', '6', '-usesqform'])
#					if saveRegToFirst == True:
#						commandString.extend(['-out', outPrefix])
#				
#					if parallelMode == 'none':
#						subprocess.call(commandString)
#					if parallelMode == 'parallel':
#						parallelPipe.stdin.write(" ".join(commandString) + "\n")
#
#				#print " ".join(commandString)
#				#fp.write(" ".join(commandString) + "\n")
#				#os.system(commandString)
#				#subprocess.call(commandString)
#
#	refFile = CCSegUtils.MNI152FLIRTTemplate()
## then register the first timepoint to MNI152
#
#	for curSubject in sorted(mainAndTimePoints.iterkeys()):
#		inFile = os.path.join(inputDirectory, curSubject)
#		outPrefix = os.path.join(outputDirectory, curSubject + '_to_std')
#		
#		if not os.path.isfile(outPrefix + '.mat'):
#			#commandString = flirtApplyXFMBase[:]
#			#commandString.extend(['-in', inFile, '-ref', refFile, '-out', outPrefix, '-init', outPrefix + '.mat'])
#		#else:
#			commandString = flirtBase[:]
#			commandString.extend(['-in', inFile, '-ref', refFile, '-omat', outPrefix + '.mat', '-dof', '7', '-usesqform'])
#			if saveRegToStd == True:
#				commandString.extend(['-out', outPrefix])
#
#			if parallelMode == 'none':
#				subprocess.call(commandString)
#			if parallelMode == 'parallel':
#				print " ".join(commandString)
#				parallelPipe.stdin.write(" ".join(commandString) + "\n")
#		#fp.write(" ".join(commandString) + "\n")
##fp.close()
#	if parallelMode == 'parallel':
#		print "parallel, running registrations to first timepoints and template"
#		parallelPipe.communicate()
#		parallelPipe.stdin.close()
#		print "parallel, finished"
#	
#	if saveRegToStd == True:
#		if parallelMode == 'parallel':
#			parallelPipe = subprocess.Popen(['parallel', parallelSwitch], stdin = subprocess.PIPE, stdout = subprocess.PIPE)
#
# check for crazy registrations
# apply the transform and if most of the image is zero then do it without the sqform
	refFile = CCSegUtils.MNI152FLIRTTemplate()

	redoNoSQForm = dict()

#	for curSubject in sorted(mainAndTimePoints.iterkeys()):
#		
#		#NIITempDir = tempfile.mkdtemp()
##		outPrefix = os.path.join(outputDirectory, curSubject + '_to_std')
#		toStdNII = os.path.join(outputDirectory, curSubject + "_to_std.nii.gz")
#		
#		#myEnv = os.environ.copy()
#		#myEnv['FSLOUTPUTTYPE'] = 'NIFTI'
#		#commandString = [os.environ['FSLDIR'] + '/bin/flirt', '-interp', flirtInterp, '-applyxfm']
#		#commandString.extend(['-in', inFile, '-ref', refFile, '-out', toStdNII, '-init', outPrefix + '.mat'])
#		
#		#P = subprocess.Popen(commandString, env = myEnv)
#		#P.communicate()
#		#print " ".join(commandString)
#		#subprocess.call(commandString)
#		NII = nibabel.load(toStdNII)
#
#		NumNonZero = numpy.count_nonzero(NII.get_data() > 0)
#		#TotalSize = NII.get_data().size
#		TotalSize = numpy.prod(NII.shape)
#		#print NII
#		proportionCovered = float(NumNonZero) / float(TotalSize)
#		print curSubject + " " + str(NumNonZero) + " " + str(TotalSize) + " " + str(float(NumNonZero) / float(TotalSize))
#		if proportionCovered < 0.1:
#			redoNoSQForm[curSubject] = None
#		#shutil.rmtree(NIITempDir)
#		#A = numpy.loadtxt(outPrefix + ".mat")
#		#rotmat, skew, scales, transl, angles = FLIRT.fsl_decomp_aff(numpy.matrix(A))
#		
#		#scales = numpy.diag(scales)
#		#if numpy.any(numpy.abs(transl) > 50):
#		#	print curSubject + " " + str(transl.T)
#		
#	if len(redoNoSQForm) > 0:
#		print "redoing without usesqform:"
#		for curSubject in sorted(redoNoSQForm.iterkeys()):
#			print "\t" + curSubject
#		doFLIRT(inputDirectory, outputDirectory, redoNoSQForm, saveRegToStd, saveRegToFirst, parallelMode, useSQForm = False)
#
	parallelSwitch = '-j+0'

	if parallelMode == 'parallel':
		parallelPipe = subprocess.Popen(['parallel', parallelSwitch], stdin = subprocess.PIPE, stdout = subprocess.PIPE)
		
# concatenate the transforms and put the other files in MNI152 space 
	refFile = CCSegUtils.MNI152FLIRTTemplate()
	for curSubject in sorted(mainAndTimePoints.iterkeys()):
		firstMat = os.path.join(outputDirectory, curSubject + '_to_std.mat') 

		if mainAndTimePoints[curSubject] != None:

			for curTimePoint in range(len(mainAndTimePoints[curSubject])):
				outMat = os.path.join(outputDirectory, mainAndTimePoints[curSubject][curTimePoint] + '_to_std.mat')
				commandString = convertXFMBase[:]

				commandString.extend(['-omat', outMat, '-concat', firstMat, os.path.join(outputDirectory, mainAndTimePoints[curSubject][curTimePoint] + '_to_' + curSubject + '.mat')])
				subprocess.call(commandString)
				
				if saveRegToStd == True:
					commandString = flirtApplyXFMBase[:]
					inFile = os.path.join(inputDirectory, mainAndTimePoints[curSubject][curTimePoint])
					outPrefix = os.path.join(outputDirectory, mainAndTimePoints[curSubject][curTimePoint] + "_to_std")
					
					commandString.extend(['-in', inFile, '-ref', refFile, '-out', outPrefix, '-init', outMat])
					
					if parallelMode == 'none':
						subprocess.call(commandString)
					if parallelMode == 'parallel':
						parallelPipe.stdin.write(" ".join(commandString) + "\n")
				#print " ".join(commandString)
				#fp.write(" ".join(commandString) + "\n")
				#subprocess.call(commandString)
	#fp.close()
	#if saveRegToStd == True:
	if parallelMode == 'parallel':
		print "parallel, applying transforms"
		parallelPipe.communicate()
		parallelPipe.stdin.close()
		print "parallel, finished"
#def main():

if __name__ == "__main__":
	main()