"""
    @file  convertDtNifti.py
    @brief Convert between interleaved analyze DT tensors and Nifti tensor files.

    @note This file originally has been copied from the trunk of the
          sbiaUtilities project revision 297 on 4/24/2012.

    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, sys, getopt, tempfile, shutil
from numpy import *
from nifti import *
from sbia.dtidroid.SbiaUtils import *

def usage():
  """prints usage information"""
  print r"""
%(EXEC)s--
  Convert from interleaved Analyze tensors (6,x,y,z) where the 6 directions are in {xx,yy,zz,xy,xz,yz}
  order to Nifti tensors (x,y,z,1,6) where the directions are in lower-triangular order {xx,xy,yy,xz,yz,zz}
  Also works for vectors, from (3,x,y,z) to (x,y,z,1,3) order. Use -n for the inverse operation

Usage: %(EXEC)s [OPTIONS] 

Required Options:
  [-d --dataFile]  Specify the input data file             (required)
                   You need to specify a header (see makeNiftiHeader.py if your data didn't come with a header)
                   also, you need a corresponding .img file (you may need to rename .dat or .d file to .img)
  [-p --prefix]    Specify the prefix of the output file   (required)

Options:
  [-r --orient]    The orientation that should be written in the Nifti header (required if not present)
                   For example, LAS is +x=Left +y=Anterior +z=Superior
                   ** Make sure this represents the actual orientation of your file **
                   **      This program will NOT attempt to reorient your file      **
                        
  [-n --niftiInput]   The input data is a Nifti tensor and the output should be a dt file
  [-o --outputDir]    The output directory to write the results. Defaults to the location of the input file
                        
  [-u --usage -h --help]  Display this message
  [-V --Version]          Display version information
  [-v --verbose]          Verbose output
""" % {'EXEC':basis.get_executable_name()}

def main():
  """parses command-line input and calls Run"""
  #the defaults
  outDir = None
  orient = None
  niftiToDt = False # Defaults to DTtoNifti
  
  rOpts = 0
  try:
    opts, files = getopt.gnu_getopt(sys.argv[1:], "hd:o:p:vVunr:",
      ["help", "dataFile=", "outputDir=","prefix=","verbose","Version",
      "usage","niftiInput","orient="])
  
  except getopt.GetoptError, err:
    usage()
    print str(err) # will print something like "option -a not recognized"
    sys.exit(2)

  verbose = 0
  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"):
      basis.print_version('convertDtNifti',
            copyright="Copyright (c) 2009, 2012 University of Pennsylvania. All rights reserved.")
      sys.exit(0)
    elif o in ("-d", "--dataFile"):
      dataFile = a
      rOpts+=1
    elif o in ("-p", "--prefix"):
      prefix = a
      rOpts+=1
    elif o in ("-o", "--outputDir"):
      outDir = a
    elif o in ("-n", "--niftiInput"):
      niftiToDt = True
    elif o in ("-r", "--orient"):
      orient = a
    else:
      assert False, "unhandled option"

  if rOpts != 2:
    usage()
    print "*** Please specify both input as well as a prefix!"
    sys.exit(2)
  
  return Run(dataFile, prefix, outDir=outDir, orient=orient, niftiToDt=niftiToDt, verbose=verbose)

def Run(dataFile, prefix, outDir=None, orient=None, niftiToDt=False, verbose=0):
  """
  dataFile - the input file
  prefix   - the prefix of the output file
  outDir   - the output directory
  orient   - a three letter orientation code (ex. LAS), needed if input is not Nifti
  niftiToDt - go from nifti to dt format
  verbose  - verbose output if >0
  """
  
  if orient == None and not isNiftiFile(dataFile):
    usage()
    print "*** Input file contains no orientation information ***"
    print "***       Please use -r to supply orientation      ***"
    sys.exit(2)
  
  #expand the files into absolute paths.
  dataFile  = os.path.realpath(dataFile)
  idStr     = os.path.basename(dataFile).rsplit(os.path.extsep,1)[0]
  
  #check input
  if not fileExists(dataFile):
    print dataFile + " could not be found!"
    sys.exit(2)
  
  #TODO: Check data file extension for either .hdr or nii

  if idStr == prefix:
    usage()
    print "\nTo avoid confusion the prefix must be different from the base"
    print "of the input data file"
    sys.exit(2)
    
  if outDir==None:
    outDir=os.path.split(dataFile)[0]
  if not isWritable(outDir):
    cryandexit("Output directory is not writable!", outDir)

  outputFile  = outDir + os.path.sep + prefix + ".hdr"
  outputFile  = os.path.realpath(outputFile)

  if verbose > 0:
    print "Converting:  "+dataFile
    if niftiToDt:
      print "  from Nifti to DT"
    else:
      print "  from DT to Nifti"
    print "Output File: "+outputFile

  if niftiToDt:
    nii = NiftiImage(dataFile)
    try:
      nt,one,nz,ny,nx = nii.data.shape
      if nt == 6: # tensor case
        dtNii = NiftiImage(reshape(rollaxis(nii.data,0,5),(nz,ny,nx,6))[ :,:,:, [0,2,5,1,3,4] ],header=nii.header)
      elif nt == 3: # vector case
        dtNii = NiftiImage(reshape(rollaxis(nii.data,0,5),(nz,ny,nx,3)),header=nii.header)
      else:
        raise ValueError
    except ValueError,e:
      print "This may not be a tensor or vector file, or header is wrong!"
      sys.exit(2)
      
    # clean up the header just a little.
    dtNii.raw_nimg.intent_code = 0
    dtNii.raw_nimg.intent_p1 = 0.0
    dtNii.extensions.clear()
    dtNii.save(outputFile)
    if verbose:
      print "New shape:",dtNii.getExtent()
    
  else: # we are converting from DT to nifti...
    dtnii = NiftiImage(dataFile)
    try:
      nz,ny,nx,nt = dtnii.data.shape
      if nt == 6:
        nii = NiftiImage(reshape(rollaxis(dtnii.data,3,0),(6,1,nz,ny,nx))[[0,3,1,4,5,2],:,:,:,:],header=dtnii.header)
        nii.raw_nimg.intent_code = 1005
        nii.raw_nimg.intent_p1 = 3.0
      elif nt == 3:
        nii = NiftiImage(reshape(rollaxis(dtnii.data,3,0),(3,1,nz,ny,nx)),header=dtnii.header)
        nii.raw_nimg.intent_code = 1007
      else:
        raise ValueError
    except ValueError,e:
      print "This may not be a tensor or vector file, or header is wrong!"
      sys.exit(2)
    
    if orient != None:
      if orient == "RAS":
        nii.setQFac(1)
        nii.setQuaternion([0,0,0])
      elif orient == "RAI":
        nii.setQFac(-1)
        nii.setQuaternion([0,0,0])
      elif orient == "RPI":
        nii.setQFac(1)
        nii.setQuaternion([1,0,0])
      elif orient == "RPS":
        nii.setQFac(-1)
        nii.setQuaternion([1,0,0])
      elif orient == "LAI":
        nii.setQFac(1)
        nii.setQuaternion([0,1,0])
      elif orient == "LAS":
        nii.setQFac(-1)
        nii.setQuaternion([0,1,0])
      elif orient == "LPS":
        nii.setQFac(1)
        nii.setQuaternion([0,0,1])
      elif orient == "LPI":
        nii.setQFac(-1)
        nii.setQuaternion([0,0,1])
      else:
        print "*** Orient code " + orient + " not recognized!"
        sys.exit(2)
    
      #nii.setSForm(nii.getQForm())
      nii.setXFormCode('qform','scanner')
      #nii.setXFormCode('sform','scanner')
      if verbose:
        print "Set Nifti orientation to:",nii.getQOrientation(as_string=True)
      
    nii.save(outputFile)
  
  if niftiToDt:
    description="Converted from Nifti to dt"
  else:
    description="Converted from dt to Nifti"
  
  ret = {'filename':outputFile, 'orientation':nii.getQOrientation(as_string=True), 'description':description}
  
  return ret

if __name__ == '__main__': main()
