#!/usr/bin/env python

"""     dtiDroid.py
       
        Copyright 2008 SBIA
        contact 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"
  sys.exit(2)

try:
  import convertDtNifti
except:
  print "dtiDroid needs convertDtNifti.py!"
  sys.exit(2)

#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)

  def isNiftiFile(fileName):
    """returns True if fileName is a nifti file"""
    try:
      return nifticlib.is_nifti_file(fileName) > 0
    except AttributeError,e:
      return clib.is_nifti_file(fileName) > 0

SVN_FILE_VERSION  ="$Id: dtiDroid_in.py 252 2009-06-16 00:14:19Z kanterae@UPHS.PENNHEALTH.PRV $"

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

def usage():
  print r"""
dtiDroid--

  Performs deformable registration on Nifti Diffusion Tensor Images

Usage: dtiDroid.py [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:
  [-M --registerMethod]  Choose which method is used to register the dti (not implemented)
                         Default is 'tensorgeometryregistration_v2.pl'
  [-i --geoFeatIters]    Iterations for geometric feature-driven registration. Default is 50
  [-l --rotFeatIters]    Iterations for rotational feature -driven registration. Default is 0
  [-c --smoothness]      Deformation field smoothness factor [0~1]. Default is 0.5
  [-E --smoothIters]     Deformation field smoothing iterations. Default is 1
  [-r --searchSize]      Neighborhood search size at each resolution low:med:hi. Default 12:10:8 (works best)
                         ** smaller neighborhoods do better locally, but are off globally **
  [-m --similarity]      The similarity degree. Default is 0.75 (works best)
  [-R --stepRatio]       Step ratio at each iteration. Default is 0.02
  [-f --useFSL]          Use fsl flirt for affine registration. Default is YES
  [-w --workingDir]      Specify a working directory. If not specified, a tmpDir is created and used
  [-o --outputDir]       The directory to output results. Defaults to the location of the input file

  [-u --usage, -h --help]  Display this message
  [-V --Version]           Display version information
  [-v --verbose]           Turns on verbose output

Examples:
  default parameters:
    dtiDroid.py -d tensor_nii.hdr -t JHU_nii.hdr -p tensor_to_JHU -l3
  
  recommendation for mouse brains:
    dtiDroid.py -d mouse_nii.hdr -t template_nii.hdr -p mouse_to_template -c0.8 -E3 -l3
"""

def main():
  verbose = 0
  simulate = 0

  #the defaults
  registerMethod = "tensorgeometryregistration_v2.pl"
  outDir         = None
  workingDir     = None
  rotFeatIters   = "0"
  geoFeatIters   = "50"
  smoothness     = "0.5"
  smoothIters    = "1"
  searchSize     = ["12","10","8"]
  similarity     = "0.75"
  stepRatio      = "0.02"
  useFsl         = True

  rOpts          = 0
  try:
    opts, files = getopt.gnu_getopt(sys.argv[1:], "hd:o:p:vVuM:w:t:i:l:f:c:E:r:m:R:",
      ["help", "dataFile=","outputDir=","prefix=","verbose","Version","usage",
      "registerMethod=","workingDir=","template=","geoFeatIters=","rotFeatIters=", "useFsl=","smoothness=","smoothIters=","searchSize=","similarity=","stepRatio="])

  except getopt.GetoptError, err:
    usage()
    print err.message # will print something like "option -a not recognized"
    sys.exit(2)

  for o, a in opts:
    if o in ("-v", "--verbose"):
      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 ("-M", "--registerMethod"):
      registerMethod = a #is ignored...
      print "Sorry registerMethod is not implemented yet"
      sys.exit(2)
    elif o in ("-l", "--rotFeatIters"):
      rotFeatIters = a
    elif o in ("-i", "--geoFeatIters"):
      geoFeatIters = a
    elif o in ("-c", "--smoothness"):
      smoothness = a
    elif o in ("-E", "--smoothIters"):
      smoothIters = a
    elif o in ("-r", "--searchSize"):
      searchSize = a.split(":")
    elif o in ("-m", "--similarity"):
      similarity = a
    elif o in ("-R", "--stepRatio"):
      stepRatio = a
    else:
      assert False, "unhandled option"

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

  return Run(dataFile, templateFile, prefix=prefix, registerMethod=registerMethod, geoFeatIters=geoFeatIters, rotFeatIters=rotFeatIters, smoothness=smoothness, smoothIters=smoothIters, searchSize=searchSize, similarity=similarity, stepRatio=stepRatio, useFsl=useFsl, workingDir=workingDir, outDir=outDir, verbose=verbose, simulate=simulate)
  
def Run(dataFile, templateFile, prefix=None, registerMethod="tensorgeometryregistration_v2.pl", geoFeatIters="50", rotFeatIters="0", smoothness="0.5", smoothIters="1", searchSize=["12","10","8"], similarity="0.75", stepRatio="0.02", useFsl=True, workingDir=None, outDir=None, verbose=0, simulate=0):
  """
  register a nifti tensor to a nifti template. Volumes must be the same dimensions
    dataFile       - the input nifti tensor file
    template       - the template nifti tensor file
    prefix         - the warped output volume prefix
    registerMethod - registration method (not available)
    geoFeatIters   - how many geometric feature iterations
    rotFeatIters   - how many rotational feature iterations
    smoothness     - Smoothness factor for deformation field [0-1]
    smoothIters    - Smoothing iterations for def field
    searchSize     - Neighborhood search size at each resolution ['low','med','hi']
    similarity     - Similarity degree
    stepRatio      - Step ratio for each iteration
    useFSL         - use FSL for affine registration
    workingDir     - keep intermediate files in a working dir
    outDir         - the directory for the output files
    verbose        - verbose output if >0
  """
  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.py")
  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):
    print "*** File does not exist: " + dataFile
    sys.exit(2)
  if not os.path.isfile(dataFile):
    print "*** File does not exist: " + templateFile
    sys.exit(2)
  if dataFile == templateFile:
    print "*** Input is the same as template: " + dataFile
    sys.exit(2)
  if not os.path.exists(outDir):
  	os.makedirs(outDir)
  
  if not isNiftiFile(dataFile) or not isNiftiFile(templateFile):
    print "*** Subject and template images must be Nifti files!"
    sys.exit(2)
    
  #check input images...
  data_im = NiftiImage(dataFile)
  mod_im = NiftiImage(templateFile)

  data_dims = data_im.getExtent()
  mod_dims  = mod_im.getExtent()
  
  data_type = data_im.raw_nimg.datatype
  mod_type  = mod_im.raw_nimg.datatype
  
  data_intent = data_im.raw_nimg.intent_code
  mod_intent  = mod_im.raw_nimg.intent_code
  
  if data_intent != 1005 or mod_intent != 1005:
    print "*** Subject and template images must tensor files!"
    sys.exit(2)
  
  if data_dims[0] != data_dims[1]:
    print "*** Image X and Y dimentions must be the same, sorry!"
    sys.exit(2)
    
  if not data_dims == mod_dims:
    print "*** Subject and template image dimensions must match, sorry!"
    sys.exit(2)

  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 = [
    registerMethod,
    "TensorStructRegistration_V2.0",
    "splitfile",
    "reversedeformationfield",
    "hammer2xdr",
    "dtiPD",
    "warpDT.V5",
    "combinetwodeformationfields",
    "rigidlyregisterdt.sh",
    "dtiFA",
    "FloatFA2Byte",
    "extractventricle",
    "linearlyregisterbyfsl.sh",
    "LnWarpDT",
    "makeavwheader",
    "fslmat2mymat",
    "cropslices",
    "transform3dimg",
    "convertDtNifti.py",
  ]
  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...
  ret = convertDtNifti.Run(templateFile, "template_dt", outDir="./", niftiToDt=True, verbose=verbose)

  #we need to convert the nii file to dt...
  dtConvCmd2 = ["convertDtNifti.py"]
  dtConvCmd2.extend(["-n"])
  dtConvCmd2.extend(["-o","./"])
  dtConvCmd2.extend(["-d",dataFile])
  dtConvCmd2.extend(["-p","input_dt"])
  if verbose > 0:
    dtConvCmd2.extend(["-v"])

  status = execCmd(dtConvCmd2,verbose=verbose)


  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(["-i", str(geoFeatIters)])
  regCmd.extend(["-l", str(rotFeatIters)])
  regCmd.extend(["-c", str(smoothness)])
  regCmd.extend(["-E", str(smoothIters)])
  regCmd.extend(["-r", ",".join([str(x) for x in searchSize])])
  regCmd.extend(["-m", str(similarity)])
  regCmd.extend(["-R", str(stepRatio)])
  if not useFsl:
    regCmd.append("-F")

  status = execCmd(regCmd,verbose=verbose,simulate=simulate,quiet=True)


  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"

  shutil.copy(fieldFile,outputBase+".Field")

  #make a dt file with the matching header
  shutil.copy(dtFile,"output_dt.img")
  shutil.copy("template_dt.hdr","output_dt.hdr")

  dtConvCmd2 = ["convertDtNifti.py"]
  dtConvCmd2.extend(["-o",outDir])
  dtConvCmd2.extend(["-d","output_dt.hdr"])
  dtConvCmd2.extend(["-p",prefix])
  
  status = execCmd(dtConvCmd2,verbose=verbose)

  #go back to where we started
  os.chdir(origCwd)

  #delete the temp
  if workingDir == None:
    shutil.rmtree(cwdDir)

  outputFile = os.path.join(outDir,prefix + ".hdr")
  if not os.path.isfile(outputFile):
    raise SbiaException("Registration output file not found: " + outputFile)
  
  return {'filename':outputFile, 'template': templateFile, 'description':" ".join(regCmd)}
  
if __name__ == '__main__': main()
