#!/usr/bin/env python
#
#       dtiDroid.py
#       
#       Copyright 2008 Luke Bloy <lbloy@seas.upenn.edu>
##

import os, sys, getopt, tempfile, shutil, subprocess, string, time

try:
  from numpy import *
  from nifti import *
except:
  print "dtiDroid depends on numpy and pynifti sorry";

#Try to import some utility methods if not define them.
try:
  from SbiaUtils import *
except:  
  class SbiaException(Exception):
    def __init__(self, msg):
      self.message = msg 
    def __str__(self):
      return repr(self.message)

  def execCmd(cmdArray,verbose=False,simulate=False,quite=False,
                noFail=False):
    """
    Executes cmdarray
    """
    if (verbose or simulate):
      print string.join(cmdArray," ");
    
    if (simulate):
      ret = [0,"Simualted","Simulated"]

    else:
      output = ["",""]
      try:
        p = subprocess.Popen(cmdArray, stdout=subprocess.PIPE
          ,stderr=subprocess.PIPE);
        
        output = array(p.communicate());
        status = p.returncode;
        ret = [status,output[0],output[1]]  
      except OSError,e: #this might be dumb!
        msg = string.join(cmdArray," ") + "\tCaused an OSError\n"
        msg = msg + "\t"+e.message
        sys.stderr.write(msg)
        raise e
      
    if (verbose or simulate):
      print "return Value: %d" % ret[0];

    #pass stderr and stdout along.
    if not quite and not simulate:
      sys.stdout.write(output[0])
      sys.stderr.write(output[1])
    
    #if this failed then thorw an exception
    if (ret[0] != 0 and not noFail):
      msg = string.join(cmdArray," ") + "   Failed"
      raise SbiaException(msg)

    return ret

  def checkCmdExist(executable):
    if (execCmd(["which",executable],quite=True)[0] != 0):
      str = "Error can't find " + executable + " please check your path!"
      raise SbiaException(str)
  
verbose = 0
simulate = 0

SVN_FILE_VERSION  ="$Id: dtiDroid_in.py 191 2008-12-04 00:00:18Z kanterae $"

def version():
  str = """
dtiDroid-
    Release       : @RELEASE_ID@
    Svn Revision  : @SVN_REV@
    File Version  : """ + SVN_FILE_VERSION;
  print str;

def usage():
  print r"""
dtiDroid--
  
  dtiDroid performs deformable registration on Diffusion Tensor Images
  
Usage: dtiDroid [OPTIONS] 

Required Options:
  [-d --dataFile]        Specify the tensor data file      (required)
  [-t --template]        Specify the tensor template file  (required)
  [-p --prefix]          Specify the prefix of the output file   (required)

Options:
  [-r --registerMethod]  Choose which method is used to register the dti.
                         currently not implemented. default is tensorgeometryregistration_v2.pl
  [-l --ornFeatIters]    Number of iterations for tensor registration using orientation features.
                         Default is 0
  [-f --useFSL]          Should fsl flirt be used as affine registration.
                         Default is YES
  [-w --workingDir]      Specify a working directory. If none is specified a
                         tmpDir is created and used.
  [-o --outputDir]       The output directory to write the results
                         Defaults to the location of the input file
  
  [-u --usage]           Display this message
  [-h --help]            Display this message
  [-V --Version]         Display version information
"""


def main():
  global verbose;

  #the defaults
  registerMethod = "tensorgeometryregistration_v2.pl"
  outDir         = None
  workingDir     = None
  ornFeatIters   = "0"
  rOpts          = 0
  useFsl         = True 
  
  try:
    opts, files = getopt.gnu_getopt(sys.argv[1:], "hd:o:p:vVur:w:t:l:f:",
      ["help", "dataFile=","outputDir=","prefix=","verbose","Version","usage",
      "registerMethod=","workingDir=","template=","ornFeatIters=","useFsl="])
  
  except getopt.GetoptError, err:
    print str(err) # will print something like "option -a not recognized"
    usage()
    sys.exit(2)

  for o, a in opts:
    if o == "-v":
      verbose+=1
    elif o in ("-h", "--help","-u","--usage"):
      usage()
      sys.exit(0)
    elif o in ("-V", "--Version"):
      version()
      sys.exit(0)
    elif o in ("-d", "--dataFile"):
      dataFile = a
      rOpts+=1
    elif o in ("-t", "--template"):
      templateFile = a
      rOpts+=1
    elif o in ("-o", "--outputDir"):
      outDir = a
    elif o in ("-w", "--workingDir"):
      workingDir = a
    elif o in ("-p", "--prefix"):
      prefix = a
      rOpts+=1
    elif o in ("-f", "--useFsl"):
      optArg = a.strip().upper()
      if optArg == 'YES' or optArg == 'Y':
        useFsl = True
      else:
        useFsl = False
          
    elif o in ("-r", "--registerMethod"):
      registerMethod = a #is ignored...
      print "Sorry this is not implemented"
      usage()
      sys.exit(2)
    elif o in ("-l", "--ornFeatIters"):
      ornFeatIters = a
    else:
      assert False, "unhandled option"

  if (rOpts != 3):
    usage()
    print "Error: please specify all required options"
    sys.exit(2)

  if outDir == None:
    outDir = os.path.split(dataFile)[0]

  outDir      = os.path.realpath(outDir)
  outputBase  = outDir + os.path.sep + prefix
  outputBase  = os.path.realpath(outputBase)

  ###Check inputs!!!!
  checkCmdExist("convertDtNifti")
  checkCmdExist(registerMethod)

  #expand the files into absolute paths and check
  dataFile      = os.path.realpath(dataFile)
  templateFile  = os.path.realpath(templateFile)
  if not os.path.isfile(dataFile):
    cryandexit("*** File does not exist!", dataFile)
  if not os.path.isfile(dataFile):
    cryandexit("*** File does not exist!", templateFile)
  if dataFile == templateFile:
    cryandexit("*** Input is the same as template!", dataFile)
  if not os.path.exists(outDir):
  	os.makedirs(outDir)
  	
  #check dimensions...
  data_im = NiftiImage(dataFile)
  mod_im = NiftiImage(templateFile)

  data_dims = data_im.getExtent()
  mod_dims  = mod_im.getExtent()
  if ( not data_dims == mod_dims ):
    cryandexit("*** Subject and template image dimensions don't match!")
      
  data_voxDims  = data_im.getVoxDims()
  mod_voxDims   = mod_im.getVoxDims()

  #build the strings for these methods.
  dataDims_str = str(data_dims[0])+","+str(data_dims[1])+","+str(data_dims[2])
  modDims_str = str(mod_dims[0])+","+str(mod_dims[1])+","+str(mod_dims[2])

  dataVoxDims_str = str(data_voxDims[0])+","+str(data_voxDims[1])+","+str(data_voxDims[2])
  modVoxDims_str = str(mod_voxDims[0])+","+str(mod_voxDims[1])+","+str(mod_voxDims[2])
  
  #might want to check dependencies:
  depList = [
    "tensorgeometryregistration_v2.pl",
    "TensorStructRegistration_V2.0",
    "splitfile",
    "reversedeformationfield",
    "hammer2xdr",
    "dtiPD",
    "warpDT.V5",
    "combinetwodeformationfields",
    "rigidlyregisterdt.sh",
    "dtiFA",
    "FloatFA2Byte",
    "extractventricle",
    "linearlyregisterbyfsl.sh",
    "LnWarpDT",
#    "getFileSize",
    "makeavwheader",
    "fslmat2mymat",
    "cropslices",
    "transform3dimg",
    "convertDtNifti",
  ]
  if (useFsl):
    depList.append("flirt")
  
  for d in depList:
    checkCmdExist(d)

  #make a tempory directory.
  origCwd = os.path.realpath(os.curdir)

  if workingDir == None:
    cwdDir = tempfile.mkdtemp()
  else:
    cwdDir = workingDir

  if (not os.path.exists(cwdDir)):
    os.makedirs(cwdDir)
  elif ( not os.path.isdir(cwdDir) ):
    print "Sorry specified working directory is a FILE!"
    sys.exit(2)
    
  os.chdir(cwdDir)

  #we need to convert the nii file to dt...
  dtConvCmd1 = ["convertDtNifti"]
  dtConvCmd1.extend(["-n"])
  dtConvCmd1.extend(["-o","./"])
  dtConvCmd1.extend(["-d",templateFile])
  dtConvCmd1.extend(["-p","template_dt"])
  if (verbose):
    dtConvCmd1.extend(["-v"])

  status = execCmd(dtConvCmd1,verbose=verbose)
  if ( status[0] != 0 ):
    print string.join(dtConvCmd1," ")
    print "*** return Failed with message"
    print status[2]
    print
    sys.exit(2)

  #we need to convert the nii file to dt...
  dtConvCmd2 = ["convertDtNifti"]
  dtConvCmd2.extend(["-n"])
  dtConvCmd2.extend(["-o","./"])
  dtConvCmd2.extend(["-d",dataFile])
  dtConvCmd2.extend(["-p","input_dt"])
  if (verbose):
    dtConvCmd2.extend(["-v"])
  
  status = execCmd(dtConvCmd2,verbose=verbose)
  if ( status[0] != 0 ):
    print string.join(dtConvCmd2," ")
    print "*** return Failed with message"
    print status[2]
    print
    sys.exit(2)
    
#  regCmd = [registerMethod]
#  regCmd.extend(["./template_dt.img"])
#  regCmd.extend(["./input_dt.img"])
#  regCmd.append(ornFeat)
  
#  echo "tensorgeometryregistration_v2.pl -M $mdl_id -D $mdl_path -S $sub -X $dx,$dy,$dz -V $mvx,$mvy,$mvz -v $vx,$vy,$vz -l $orn_feat"
  #time tensorgeometryregistration_v2.pl -M $mdl_id -D $mdl_path -S $sub -X $dx,$dy,$dz -V $mvx,$mvy,$mvz -v $vx,$vy,$vz -l $orn_feat > ${sub_id}.registration.log 2>&1

  regCmd = ["tensorgeometryregistration_v2.pl"]
  regCmd.extend(["-M", "template_dt.img" ])
  regCmd.extend(["-D", "./" ])
  regCmd.extend(["-S", "./input_dt.img" ])
  regCmd.extend(["-X", dataDims_str ])
  regCmd.extend(["-V", modVoxDims_str ])
  regCmd.extend(["-v", dataVoxDims_str ])
  regCmd.extend(["-l",ornFeatIters])
  if (not useFsl):
    regCmd.append("-F")
  
  status = execCmd(regCmd,verbose=verbose)
  if ( status[0] != 0 ):
    print string.join(regCmd," ")
    print "*** return Failed with message"
    print status[2]
    print
    sys.exit(2)

  if (verbose):
    print status[1]
    
  logFile = open(outputBase+"_registration.log","w")
  logFile.write("STDOUT ******")
  logFile.write(status[1])
  logFile.write("STDERR ******")
  logFile.write(status[2])
  logFile.close()
  
  ##ok we are all done'
  #lets move the results to where we want them...
  fieldFile = "Field.1.img.Mdl2Subj"
  dtFile    = "Warped.Subj.dt"
#  logFile   = "input_dt.img.registration.log"

  shutil.copy(fieldFile,outputBase+".Field")
#  shutil.copy(logFile,outputBase+"_registration.log")

  #make a dt file with the matching header
  shutil.copy(dtFile,"output_dt.img")
  shutil.copy("template_dt.hdr","output_dt.hdr")
  
  dtConvCmd2 = ["convertDtNifti"]
  dtConvCmd2.extend(["-o",outDir])
  dtConvCmd2.extend(["-d","output_dt.hdr"])
  dtConvCmd2.extend(["-p",prefix]) 

  status = execCmd(dtConvCmd2,verbose=verbose)
  if ( status[0] != 0 ):
    print string.join(dtConvCmd2," ")
    print "*** return Failed with message"
    print status[2]
    print
    sys.exit(2)

  #go back to where we started.
  os.chdir(origCwd)
  
  #delete the temp
  if workingDir == None:
    shutil.rmtree(cwdDir)

  return(0)
  
if __name__ == '__main__': main()
