"""
    @file  dtidroid.py
    @brief Perform deformable registration on NIfTI diffusion tensor images.

    Copyright (c) 2008, 2009, 2012 University of Pennsylvania.

    This file is part of DTI-DROID.

    DTI-DROID is free software: you can redistribute it and/or modify
    it under the terms of the GNU General Public License as published by
    the Free Software Foundation, either version 3 of the License, or
    (at your option) any later version.

    DTI-DROID is distributed in the hope that it will be useful,
    but WITHOUT ANY WARRANTY; without even the implied warranty of
    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
    GNU General Public License for more details.

    You should have received a copy of the GNU General Public License
    along with DTI-DROID.  If not, see <http://www.gnu.org/licenses/>.

    Contact: SBIA Group <sbia-software at uphs.upenn.edu>
"""

@BASIS_PYTHON_UTILITIES@

import os
import sys
import getopt
import tempfile
import shutil
from nifti import NiftiImage
try:    from nifti import clib
except: from nifti import nifticlib

# ============================================================================
# help
# ============================================================================

# ----------------------------------------------------------------------------
def print_help():
  print r"""
%(EXEC)s--
  Performs deformable registration of diffusion tensor images.

Usage:
  %(EXEC)s [options]

Required arguments:
  [-d --dataFile]        Specify the tensor data file.
  [-t --template]        Specify the tensor template file.
  [-p --prefix]          Specify the prefix of the output file.

Optional arguments:
  [-i --geoFeatIters]    Iterations for geometric feature-driven registration. (default: 50)
  [-l --rotFeatIters]    Iterations for rotational feature-driven registration. (default: 0)
  [-c --smoothness]      Deformation field smoothness factor [0~1]. (default: 0.5)
  [-E --smoothIters]     Deformation field smoothing iterations. (default: 1)
  [-r --searchSize]      Neighborhood search size at each resolution low:med:hi.
                         ** smaller neighborhoods do better locally, but are off globally **
                         (default: 12:10:8, works best in our experience).
  [-m --similarity]      The similarity degree. (default: 0.75, works best in our experience)
  [-R --stepRatio]       Step ratio at each iteration. (default: 0.02)
  [-A --skipAffine]      Skip affine registration (if your data is already affinely registered).
  [-f --useFSL]          Use fsl flirt for affine registration. (default: on)
  [-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.

  [-h --help]            Display help and exit.
  [--version]            Display version information and exit.
  [-v --verbose]         Increase verbosity of output messages.

Examples:
  default parameters:
    %(EXEC)s -d tensor_nii.hdr -t JHU_nii.hdr -p tensor_to_JHU -l3
  
  recommendation for mouse brains:
    %(EXEC)s -d mouse_nii.hdr -t template_nii.hdr -p mouse_to_template -c0.8 -E3 -l3

Contact:
  SBIA Group <sbia-softare at uphs.upenn.edu>
""" % {'EXEC': basis.get_executable_name()}


# ============================================================================
# auxiliary functions
# ============================================================================

# ----------------------------------------------------------------------------
def checkCmdExists(executable):
  """Check if command exists."""
  if not basis.is_known_target(executable):
      try:
        basis.which(executable)
      except basis.WhichError:
        raise Exception("Error: Cannot find command " + executable + "! Please check your PATH.")

# ----------------------------------------------------------------------------
def isNiftiFile(fileName):
  """Check if a given file is a NIfTI image."""
  try:
    return nifticlib.is_nifti_file(fileName) > 0
  except:
    return clib.is_nifti_file(fileName) > 0

# ============================================================================
# main
# ============================================================================

# ----------------------------------------------------------------------------
def main():
  verbose = 0
  simulate = 0

  #the defaults
  registerMethod = "tensorgeometryregistration"
  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
  skipAffine     = False

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

  except getopt.GetoptError, err:
    print_help()
    print str(err) # 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"):
      print_help()
      sys.exit(0)
    elif o == "--version":
      basis.print_version('dtiDroid.py')
      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
    elif o in ("-A", "--skipAffine"):
      skipAffine = True
    else:
      assert False, "unhandled option"

  if rOpts != 3:
    print_help()
    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, skipAffine=skipAffine, workingDir=workingDir, outDir=outDir, verbose=verbose, simulate=simulate)
  
def run(dataFile, templateFile, prefix, 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, skipAffine=False, 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
    skipAffine     - skip affine registration
    workingDir     - keep intermediate files in a working dir
    outDir         - the directory for the output files
    verbose        - verbose output if >0
    simulate       - don't execute anything if >0
  """
  if outDir == None:
    outDir = os.path.split(dataFile)[0]
  
  outDir      = os.path.realpath(outDir)
  outputBase  = os.path.join(outDir, prefix)
  outputBase  = os.path.realpath(outputBase)
  
  #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 dependencies
  depList = [
    # auxiliary executables
    registerMethod,
    "TensorStructRegistration",
    "splitfile",
    "reversedeformationfield",
    "hammer2xdr",
    "dtiPD",
    "nWarpDTV5",
    "combinetwodeformationfields",
    "rigidlyregisterdt",
    "dtiFA",
    "FloatFA2Byte",
    "extractventricle",
    "linearlyregisterbyfsl",
    "LnWarpDT",
    "makeavwheader",
    "fslmat2mymat",
    "cropslices",
    "transform3dimg",
    "convertDtNifti",
    "adjustSizeNifti"
  ]
  if useFsl:
    depList.append("flirt")
  for d in depList:
    checkCmdExists(d)
  
  #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
  
  data_orn = data_im.getQOrientation()
  mod_orn  = mod_im.getQOrientation()
  
  if data_intent != 1005 or mod_intent != 1005:
    print "*** Subject and template images must be tensor files!"
    sys.exit(2)
  
  if mod_dims[0] != mod_dims[1]:
    print "*** Model X and Y dimentions must be the same, sorry!"
    sys.exit(2)
  
  if data_orn != mod_orn:
    print "*** Subject and template orientations must match, sorry!"
    sys.exit(2)
  
  data_voxDims  = data_im.getVoxDims()
  mod_voxDims   = mod_im.getVoxDims()
  
  # build the strings for these methods
  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])  
  
  # make a tempory directory
  origCwd = os.path.realpath(os.curdir)
  if workingDir == None:
    exec_name = basis.get_executable_name()
    if os.environ.has_key( 'SBIA_TMPDIR' ):
      cwDir = tempfile.mkdtemp(prefix=exec_name,dir=os.environ['SBIA_TMPDIR'])
    else:
      cwDir = tempfile.mkdtemp(prefix=exec_name)
  else:
    cwDir = workingDir
  cwDir = os.path.realpath(cwDir)
  if not os.path.exists(cwDir):
    os.makedirs(cwDir)
  elif not os.path.isdir(cwDir):
    print("***Working dir is not a directory!")
    sys.exit(2)
  
  # cd int tmp
  os.chdir(cwDir)
  
  try:
    # match dimensions
    if data_dims != mod_dims:
      adjCmd = ["adjustSizeNifti"]
      adjCmd.extend(["-d",dataFile])
      adjCmd.extend(["-p","input_adj"])
      adjCmd.extend(["-o","./"])
      if data_dims[0] != mod_dims[0]:
        adjCmd.extend(["-x",str(mod_dims[0])])
      if data_dims[1] != mod_dims[1]:
        adjCmd.extend(["-y",str(mod_dims[1])])
      if data_dims[2] != mod_dims[2]:
        adjCmd.extend(["-z",str(mod_dims[2])])
      if verbose > 0:
        adjCmd.extend(["-v"])
      basis.execute_process(adjCmd, verbose=verbose)
      dataFile = "./input_adj.hdr"
      
    # we need to convert the nii template file to dt...
    dtConvCmd2 = ["convertDtNifti"]
    dtConvCmd2.extend(["-n"])
    dtConvCmd2.extend(["-o","./"])
    dtConvCmd2.extend(["-d",templateFile])
    dtConvCmd2.extend(["-p","template_dt"])
    if verbose > 0:
      dtConvCmd2.extend(["-v"])
    
    basis.execute_process(dtConvCmd2, verbose=verbose)
    
    # we need to convert the nii data file to dt...
    dtConvCmd2 = ["convertDtNifti"]
    dtConvCmd2.extend(["-n"])
    dtConvCmd2.extend(["-o","./"])
    dtConvCmd2.extend(["-d",dataFile])
    dtConvCmd2.extend(["-p","input_dt"])
    if verbose > 0:
      dtConvCmd2.extend(["-v"])
    
    basis.execute_process(dtConvCmd2, verbose=verbose)
    
    regCmd = ["tensorgeometryregistration"]
    regCmd.extend(["-M", "template_dt.img" ])
    regCmd.extend(["-D", "./" ])
    regCmd.extend(["-S", "./input_dt.img" ])
    regCmd.extend(["-X", modDims_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 skipAffine:
      regCmd.append("-A")
    elif not useFsl:
      regCmd.append("-F")
    
    status = basis.execute_process(regCmd, allow_fail=True, stdout=True, quiet=True, simulate=simulate, verbose=verbose)
    
    logFile = open(outputBase+"_registration.log","w")
    logFile.write("STDOUT ******")
    logFile.write(status[1])
    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.copy2(fieldFile,outputBase+".Field")
    
    #make a dt file with the matching header
    shutil.copy2(dtFile,"output_dt.img")
    shutil.copy2("template_dt.hdr","output_dt.hdr")
    
    dtConvCmd2 = ["convertDtNifti"]
    dtConvCmd2.extend(["-o",outDir])
    dtConvCmd2.extend(["-d","output_dt.hdr"])
    dtConvCmd2.extend(["-p",prefix])
    
    basis.execute_process(dtConvCmd2, verbose=verbose)
  except Exception,e:
    raise e
  finally:
    #go back to where we started
    os.chdir(origCwd)
    #delete the temp
    if workingDir == None:
      shutil.rmtree(cwDir)
  
  outputFile = os.path.join(outDir,prefix + ".hdr")
  if not os.path.isfile(outputFile):
    raise Exception("Registration output file not found: " + outputFile)
  elif NiftiImage(outputFile).max == 0:
    raise Exception("Output file is blank: " + outputFile)
  
  return {'filename':outputFile, 'template': templateFile, 'description':" ".join(regCmd)}
 
# ----------------------------------------------------------------------------
if __name__ == '__main__': main()
