"""
    @file  adjustSizeNifti.py
    @brief This program can be used to crop or pad a nifti image.

    @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
import numpy
from nifti import NiftiImage
from sbia.dtidroid.SbiaUtils import *

def usage():
  """prints usage information"""
  print """
%(EXEC)s--
  Performs cropping or zero-padding on a nifti image. If you are working on an axial
  orientation such as LPS, then your x-axis will correspond to Right:Left, the y-axis
  to Anterior:Posterior and the z-axis to Inferior-Superior
  
Usage: %(EXEC)s [OPTIONS] 

Required Options:
  [-d --dataFile    hdr]  Specify the input Nifti data file       (required)
  [-p --prefix     pref]  Specify the prefix of the output file   (required)
  
Options:
  [-x --xPad        x:X]  Specify the number of slices to be added or removed on the x-axis
                          Give a target number of slices to uniformly adjust your image
  [-y --yPad        y:Y]  Similarly for the y axis
  [-z --zPad        z:Z]  Similarly for the z axis
  [-c --copyFrom    hdr]  Pad to match this file's dimensions

  [-T --dt]               Specify that the input file is a dt file (old-school tensor)
                          !! Note: The output will also be a dt file !!
  [-o --outputDir   dir]  The output directory to write the results. Defaults
                          to the location of the input file
  [-O --origin    R:A:S]  Specify the Origin of the new coordinate system in RAS coordinates
                          By default, output origin is calculated to align with input origin

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

Example:
  %(EXEC)s -z -15:20 -x 128 -y 128 -d image.hdr -p cropped
    Will remove 15 slices from the beginning & add 23 at the end of the z direction
    and adjust the x and y directions uniformly to be 128 slices
""" % {'EXEC':basis.get_executable_name()}
  
def main():
  """parses command-line input and calls Run"""
  #the defaults
  outDir = None
  xPad   = None
  yPad   = None
  zPad   = None
  dtFlag = False
  origin = None
  copyFrom = None
  
  rOpts = 0
  try:
    opts, files = getopt.gnu_getopt(sys.argv[1:], "hd:o:p:vVux:y:z:TO:c:",
      ["help", "dataFile=", "outputDir=","prefix=","verbose","Version", "usage"
      ,"xPad=","yPad=","zPad=","dt","origin=","copyFrom="])
  
  except getopt.GetoptError, err:
    usage()
    print str(err) # option 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('adjustSizeNifti',
            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 ["-x", "--xPad"]:
      xPad = a
    elif o in ["-y", "--yPad"]:
      yPad = a
    elif o in ["-z", "--zPad"]:
      zPad = a
    elif o in ["-T", "--dt"]:
      dtFlag = True
    elif o in ["-O", "--origin"]:
      origin = a.split(":")
    elif o in ["-c", "--copyFrom"]:
      copyFrom = a
    else:
      assert False, "unhandled option"

  if rOpts != 2:
    usage()
    print "*** Please specify all required options!"
    sys.exit(2)
  
  if copyFrom:
    imc = NiftiImage(copyFrom)
    imc_dim = imc.getExtent()
    xPad = imc_dim[0]
    yPad = imc_dim[1]
    zPad = imc_dim[2]
  
  if xPad == None and yPad == None and zPad == None:
    usage()
    print "*** No adjustments specified! Nothing to do"
    sys.exit(2)

  #parse voxel strings...

  xNew = [0,0]
  if xPad != None:
    parts = str(xPad).split(":")
    if len(parts) == 1:
      xNew = [int(xPad)]
    else:
      if parts[0] != "":
        xNew[0] = int(parts[0])
      if parts[1] != "":
        xNew[1] = int(parts[1])
  yNew = [0,0]
  if yPad != None:
    parts = str(yPad).split(":")
    if len(parts) == 1:
      yNew = [int(yPad)]
    else:
      if parts[0] != "":
        yNew[0] = int(parts[0])
      if parts[1] != "":
        yNew[1] = int(parts[1])
  zNew = [0,0]
  if zPad != None:
    parts = str(zPad).split(":")
    if len(parts) == 1:
      zNew = [int(zPad)]
    else:
      if parts[0] != "":
        zNew[0] = int(parts[0])
      if parts[1] != "":
        zNew[1] = int(parts[1])
  
  return Run(dataFile, prefix, outDir=outDir, xPad=xNew, yPad=yNew, zPad=zNew, origin=origin, dtFlag=dtFlag, verbose=verbose)

def Run(dataFile, prefix, outDir=None, xPad=[0,0], yPad=[0,0], zPad=[0,0], origin=None, dtFlag=False, verbose=0):
  """
  pad/crop a Nifti image to specifications
    dataFile   - the input image
    prefix     - the output prefix
    outDir     - the directory where the adjusted image should be output
    xPad       - a list [A,B] where A is the slices to be padded (croped if negative)
                 at the beginning of the X direction and B the slices to be padded (cropped)
                 at the end of the X direction
    yPad       - same as xPad for the Y direction
    zPad       - same as xPad for the Z direction
    origin     - origin in RAS coords as [ox,oy,oz]
    dtFlag     - True if dataFile is in the dt format
    verbose    - verbose output if >0
  """
  # expand the files into absolute paths.
  idStr    = getFileBase(dataFile)
  dataFile = os.path.realpath(dataFile)

  # check input
  if not os.path.isfile(dataFile):
    cryandexit("File not found",dataFile)
  if not isNiftiFile(dataFile):
    print "!! Warning: " + dataFile + " does not have a valid nifti header.\n!! Please use convertAnalyze2Nifti.py to make a proper header with orientation information!"
    #sys.exit(2)
  
  if origin != None and (len(origin) != 3 or not isNumber(origin, strOk=True)):
    cryandexit("Expecting 3 numerical values for the origin")
  
  # check prefix
  if idStr == prefix:
    cryandexit("To avoid confusion the prefix must be different from the base of the input data file",prefix)
  
  # set output
  if outDir==None:
    outDir=getFilePath(dataFile)
  if not isWritable(outDir):
    cryandexit("Output directory is not writable!", outDir)

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

  if verbose > 0:
    print "Converting: ",dataFile
    print "Output File:",outputFile

  # start processing
  im = NiftiImage(dataFile)
  if not xPad:
    xPad = [0,0]
  else:
    xPad = [int(x) for x in xPad]
  if not yPad:
    yPad = [0,0]
  else:
    yPad = [int(x) for x in yPad]
  if not zPad:
    zPad = [0,0]
  else:
    zPad = [int(x) for x in zPad]
  
  # nifti image
  if not dtFlag: 
    dims = im.data.shape
    ndims = numpy.array(dims)
    if len(xPad) == 1:
      dn = (xPad[0] - ndims[-1])/2
      if dn != (xPad[0] - ndims[-1])/2.0:
        xPad = [dn,dn+1] if dn > 0 else [dn+1,dn]
      else:
        xPad = [dn,dn]
    ndims[-1] += numpy.sum(xPad)
    if len(yPad) == 1:
      dn = (yPad[0] - ndims[-2])/2
      if dn != (yPad[0] - ndims[-2])/2.0:
        yPad = [dn,dn+1] if dn > 0 else [dn+1,dn]
      else:
        yPad = [dn,dn]
    ndims[-2] += numpy.sum(yPad)
    if len(zPad) == 1:
      dn = (zPad[0] - ndims[-3])/2
      if dn != (zPad[0] - ndims[-3])/2.0:
        zPad = [dn,dn+1] if dn > 0 else [dn+1,dn]
      else:
        zPad = [dn,dn]
    ndims[-3] += numpy.sum(zPad)
    if verbose:
      print "X adjust: ", xPad
      print "Y adjust: ", yPad
      print "Z adjust: ", zPad

    # figure out the new origin
    if origin == None: 
      origin = numpy.dot(im.qform,numpy.array([-xPad[0],-yPad[0],-zPad[0],1]))[0:3]
      #origin[0] = str((im.qoffset[0]/im.pixdim[0]+xPad[0])*im.pixdim[0])
      #origin[1] = str((im.qoffset[1]/im.pixdim[1]-yPad[0])*im.pixdim[1])
      #origin[2] = str((im.qoffset[2]/im.pixdim[2]-zPad[0])*im.pixdim[2])
    if verbose:
      print "New origin: ", origin
    
    #compute the ranges...
    xRangeNew = [max([0,xPad[0]]), xPad[0]+dims[-1]+min([0,xPad[1]])]
    yRangeNew = [max([0,yPad[0]]), yPad[0]+dims[-2]+min([0,yPad[1]])]
    zRangeNew = [max([0,zPad[0]]), zPad[0]+dims[-3]+min([0,zPad[1]])]

    xRangeOld = [max([0,-xPad[0]]), max([0,-xPad[0]]) + xRangeNew[1]-xRangeNew[0]]
    yRangeOld = [max([0,-yPad[0]]), max([0,-yPad[0]]) + yRangeNew[1]-yRangeNew[0]]
    zRangeOld = [max([0,-zPad[0]]), max([0,-zPad[0]]) + zRangeNew[1]-zRangeNew[0]]
    
    origData = im.data
    newData = ndarray(ndims,im.data.dtype)

    #collapse data into 4 dimensions
    if len(ndims) > 4:
      nVols = ndims[0:len(ndims)-3].prod();
      newData  = newData.reshape(nVols,ndims[-3],ndims[-2],ndims[-1])
      origData = origData.reshape(nVols,dims[-3],dims[-2],dims[-1])

    #copy the data over...
    if len(ndims) == 3:
      newData[zRangeNew[0]:zRangeNew[1],yRangeNew[0]:yRangeNew[1],xRangeNew[0]:xRangeNew[1]] \
      = origData[zRangeOld[0]:zRangeOld[1],yRangeOld[0]:yRangeOld[1],xRangeOld[0]:xRangeOld[1]]
    else:
      newData[:,zRangeNew[0]:zRangeNew[1],yRangeNew[0]:yRangeNew[1],xRangeNew[0]:xRangeNew[1]] \
      = origData[:,zRangeOld[0]:zRangeOld[1],yRangeOld[0]:yRangeOld[1],xRangeOld[0]:xRangeOld[1]]

    #expand it again
    if len(ndims) > 4:
      nVols = ndims[0:len(ndims)-3].prod()
      newData = newData.reshape(ndims)
  # DT FILE
  else:
    if verbose:
      print "DT flag detected"
    dims = im.data.shape
    ndims = array(dims)
    if len(xPad) < 2:
      dn = (xPad[0] - ndims[2])/2
      if dn != (xPad[0] - ndims[2])/2.0:
        xPad = [dn,dn+1] if dn > 0 else [dn+1,dn]
      else:
        xPad = [dn,dn]
    ndims[2] += sum(xPad)
    if len(yPad) < 2:
      dn = (yPad[0] - ndims[1])/2
      if dn != (yPad[0] - ndims[1])/2.0:
        yPad = [dn,dn+1] if dn > 0 else [dn+1,dn]
      else:
        yPad = [dn,dn]
    ndims[1] += sum(yPad)
    if len(zPad) < 2:
      dn = (zPad[0] - ndims[0])/2
      if dn != (zPad[0] - ndims[0])/2.0:
        zPad = [dn,dn+1] if dn > 0 else [dn+1,dn]
      else:
        zPad = [dn,dn]
    ndims[0] += sum(zPad)
    if verbose:
      print "X adjust: ", xPad
      print "Y adjust: ", yPad
      print "Z adjust: ", zPad

    #compute the ranges...
    xRangeNew = [max([0,xPad[0]]), xPad[0]+dims[2]+min([0,xPad[1]])]
    yRangeNew = [max([0,yPad[0]]), yPad[0]+dims[1]+min([0,yPad[1]])]
    zRangeNew = [max([0,zPad[0]]), zPad[0]+dims[0]+min([0,zPad[1]])]

    xRangeOld = [max([0,-xPad[0]]), max([0,-xPad[0]]) + xRangeNew[1]-xRangeNew[0]]
    yRangeOld = [max([0,-yPad[0]]), max([0,-yPad[0]]) + yRangeNew[1]-yRangeNew[0]]
    zRangeOld = [max([0,-zPad[0]]), max([0,-zPad[0]]) + zRangeNew[1]-zRangeNew[0]]
    
    origData = im.data
    newData = ndarray(ndims,im.data.dtype)
      
    #copy the data over...
    newData[zRangeNew[0]:zRangeNew[1],yRangeNew[0]:yRangeNew[1],xRangeNew[0]:xRangeNew[1],:] \
      = origData[zRangeOld[0]:zRangeOld[1],yRangeOld[0]:yRangeOld[1],xRangeOld[0]:xRangeOld[1],:]
  
  newNii = NiftiImage(newData,header=im.header)
  if origin != None: # new to test for != None because I get "ValueError: The truth value of an array with more than one element is ambiguous"
    newNii.setQOffset(origin)
    newNii.setSFormCode(0) #newNii.setSForm(newNii.getQForm())
  
  newNii.save(outputFile)
  
  if verbose > 0:
    print ">> Saved adjusted file"
    
  description = "adjusted size to " + str(newNii.getVolumeExtent())
  ret = {'filename':outputFile, 'description':description}
  
  return ret

if __name__ == '__main__': main()
