#!/usr/bin/env python
# subtomosetup - make command files for reconstructing subvolumes around points
#
# Author: David Mastronarde
#
# $Id$
#

progname = 'subtomosetup'
prefix = 'ERROR: ' + progname + ' - '
yzRatioCrit = 2.
xzShiftDecimals = 2
maxChunks = 99990
minChunks = 8
maxChunkPerProc = 10
minChunkPerProc = 5
warnXtiltCrit = 0.3
debug = 0

# Compose the next command file name and increment the number. making a sync if indicated
def makeNextComName(doSync = False):
   global comNum
   if doSync:
      comName = rootWithDir + fmtstr('-{:03d}', comNum) + '-sync.com'
   else:
      comName = rootWithDir + fmtstr('-{:03d}', comNum) + '.com'
   comNum += 1
   return comName


#### MAIN PROGRAM  ####
#
# load System Libraries
import os, sys, copy

#
# Setup runtime environment
if os.getenv('IMOD_DIR') != None:
   IMOD_DIR = os.environ['IMOD_DIR']
   if sys.platform == 'cygwin' and sys.version_info[0] > 2:
      IMOD_DIR = IMOD_DIR.replace('\\', '/')
      if IMOD_DIR[1] == ':' and IMOD_DIR[2] == '/':
         IMOD_DIR = '/cygdrive/' + IMOD_DIR[0].lower() + IMOD_DIR[2:]
   sys.path.insert(0, os.path.join(IMOD_DIR, 'pylib'))
   from imodpy import *
   addIMODbinIgnoreSIGHUP()
else:
   sys.stdout.write(prefix + " IMOD_DIR is not defined!\n")
   sys.exit(1)

#
# load IMOD Libraries
from pip import *
from pysed import *
from tomocoords import *

# Fallbacks from ../manpages/autodoc2man 3 1 subtomosetup
options = ["root:RootName:CH:", "center:CenterPositionFile:FN:",
           "volume:VolumeModeled:FN:", "raw:RawStackFile:FN:", "axis:AxisAngle:F:",
           "objects:ObjectsToUse:LI:", "size:SizeInXYZ:IT:",
           "dir:DirectoryForOutput:FN:", "chunk:DirectoryForChunkFiles:FN:",
           "skip:SkipSubVolNumbers:B:", "stackvols:MakeVolumeStacks:I:",
           "com:CommandFile:FN:", "binali:NewAlignedBinning:I:",
           "newstcom:NewstackComFile:FN:", "unaligned:UseUnalignedImages:B:",
           "reduce:FourierReduceByFactor:I:", "zlevels:NumberOfZLevels:I:",
           "extent:ExtentOfZLevelsInNm:I:", "invert:InvertZLevelOffsets:I:",
           "adjust:AdjustForAlignZShift:B:", "ctfcom:CorrectionComFile:FN:",
           "erase:EraseFiducials:B:", "goldcom:GoldEraserComFile:FN:",
           "filter:FilterIn2D:B:", "2dcom:2DFilterComFile:FN:", "gpu:WhenToUseGPU:I:",
           "pixel:RawPixelSize:F:", "xform:AlignTransformFile:FN:",
           "reorient:ReorientionType:I:", "proc:ProcessorNumber:I:",
           "runs:RunsPerChunk:I:", "help:usage:B:"]

(rootName, volName, centerFile, pointFile, modelFile, objList, comFile, \
       reorientIn, enteredOrient) = getCommonOptions(options, progname)

checkList = []
boundPixels = parallelBoundarySize()
outDir = PipGetString('DirectoryForOutput', '')
skipVolNumbers = PipGetBoolean('SkipSubVolNumbers', 0)
comDir = PipGetString('DirectoryForChunkFiles', 0)
(xSize, ySize, zSize) = PipGetThreeIntegers('SizeInXYZ', 0, 0, 0)
if xSize < 1 or ySize < 1 or zSize < 1:
   exitError('Positive sizes must be entered for the reconstructions')

makeVolStacks = PipGetInteger('MakeVolumeStacks', 0)
if not PipGetErrNo() and makeVolStacks < 5:
   exitError('The volume stack entry must be at least 5')
numZlevels = PipGetInteger('NumberOfZLevels', 0)
maxZextent = PipGetFloat('ExtentOfZLevelsInNm', 0.)
if numZlevels and maxZextent:
   exitError('You cannot enter both a number of Z levels and their extent')
doCTF = numZlevels > 0 or maxZextent > 0
invertZoffsets = PipGetInteger('InvertZLevelOffsets', -1)
adjustForZShift = PipGetBoolean('AdjustForAlignZShift', 0)

if doCTF:
   ctfComFile = getOrDeriveComFile('CorrectionComFile', 'ctfcorrection', comFile,
                                   'CTF correction')
   ctfLines = readTextFile(ctfComFile, 'CTF correction command file')
   checkList.append((ctfComFile, 'CTF correction command file'))

doErase = PipGetBoolean('EraseFiducials', 0)
if doErase:
   eraseComFile =  getOrDeriveComFile('GoldEraserComFile', 'golderaser', comFile,
                                      'erasing gold')
   checkList.append((eraseComFile, 'gold erasing command file'))

# Determine filtering
doFilter = PipGetBoolean('FilterIn2D', 0)
if doFilter:
   mtfComFile = getOrDeriveComFile('2DFilterComFile', 'mtffilter', comFile, 
                                   '2D filtering')
   mtfLines = readTextFile(mtfComFile, '2D filtering command file')
   
(comExt, ifDual, rawRoot, typeExt, rawExt) = findRootAxisAndExtensions()
axisLet = setAxisLetter(rootName, ifDual)

# Get raw stack name 
stackName = PipGetString('RawStackFile', '')
if not stackName:
   (rawExt, numCandid) = getRawExtensionFallback(rawExt, rootName)

   if not numCandid or numCandid > 1:
      exitError('The raw stack name must be entered, it cannot be deduced')
   
   stackName = rootName + '.' + rawExt

# Get the aligned stack name, or give up and assume it
if typeExt:
   aliName = datasetFilename('.ali', rootName, typeExt)
else:
   tiltLines = readTextFile(comFile)
   aliName = optionValue(tiltLines, 'inputproj', 0, 1)
   if not aliName:
      aliName = rootName + '.ali'

# Get the axis angle and whether X/Y are transposed
(axisAngle, transposeXY) = getAxisAngleAndTranspose(axisLet)

# See if using raw input
useRawInput = PipGetBoolean('UseUnalignedImages', 0)
splitCorrSize = ''
if useRawInput:
   (aliXformFile, rawBinning, rawPixSize) = getEssentialRawOptions(rawRoot)
   try:
      (nxCorr, nyCorr, nzCorr) = getmrcsize(stackName)
   except ImodpyError:
      exitFromImodError(progname)

   splitCorrSize = fmtstr('-size {},{},{}', nxCorr // rawBinning, nyCorr // rawBinning,
                           nzCorr)

# See if remaking aligned stack
newAliBinning = PipGetInteger('NewAlignedBinning', 0)
if newAliBinning > 0:
   if useRawInput:
      exitError('You cannot use unaligned images and specify a new aligned stack binning')
   newstComFile = getOrDeriveComFile('NewstackComFile', 'newst', comFile, 
                                     'making new aligned stack')
   newstLines = readTextFile(newstComFile)

   # Make modifications now to binning and to size if it is present
   sedcom = sedDelAndAdd('BinByFactor', newAliBinning, 'TransformFile')
   aliOutSize = optionValue(newstLines, 'SizeToOutputInXandY', INT_VALUE)
   oldAliBin = optionValue(newstLines, 'BinByFactor', INT_VALUE, numVal = 1)
   if not oldAliBin:
      oldAliBin = 1
   if aliOutSize:
      for ind in (0, 1):
         aliOutSize[ind] = (aliOutSize[ind] * oldAliBin + newAliBinning - 1) // \
            newAliBinning
      sedcom.append(sedModify('SizeToOutputInXandY', 
                    fmtstr('{},{}', aliOutSize[0], aliOutSize[1])))
   newstLines = pysed(sedcom, newstLines)

   runLines = []
   gotRun = False

   # Get the lines for running it
   for line in newstLines:
      if gotRun and line.startswith('$'):
         break
      if not gotRun and line.startswith('$') and 'newstack' in line:
         gotRun = True
      elif gotRun:
         runLines.append(line)

   if not gotRun:
      exitError('Could not find newstack line in command file')

   # Make a stack with one section in order to get the size and header correct
   sedcom = sedDelAndAdd('SectionsToRead ', 0, 'TransformFile')
   runLines = pysed(sedcom, runLines)
   try:
      newstOut = runcmd('newstack -StandardInput', runLines)
   except ImodpyError:
      exitFromImodError(progname)

# Check that all the files exist
checkList += [(stackName, 'raw stack'),
              (volName, 'modeled volume'), (comFile, 'command file'),
              (centerFile, 'center position file')]
if not useRawInput and newAliBinning <= 0:
   checkList.append((aliName, 'aligned stack'))
if doCTF and adjustForZShift:
   if axisLet == None:
      exitError('Cannot determine axis for finding align command file')
   alignCom = 'align' + axisLet + '.com'
   checkList.append((alignCom, 'align command file'))

for (name, descrip) in checkList:
   if not os.path.exists(name):
      exitError('The ' + descrip + ', ' + name + ', does not exist')

# Check and create output directory
if outDir:
   if os.path.exists(outDir):
      if not os.path.isdir(outDir):
         exitError('The specified name for output directory already exists and is ' +\
                      'not a directory')
   else:
      try:
         os.mkdir(outDir)
      except OSError:
         exitError('Making directory for output, ' + outDir)

      try:
         outDir = os.path.relpath(outDir)
      except Exception:
         pass
      if os.path.isabs(outDir):
         prnstr('WARNING: ' + progname + ' - The output directory cannot be converted' + \
                ' to a relative path and may not work on other machines')

if comDir:
   if os.path.exists(comDir):
      if not os.path.isdir(comDir):
         exitError('The specified name for command file directory already exists and ' +\
                   'is not a directory')
   else:
      try:
         os.mkdir(comDir)
      except OSError:
         exitError('Making directory for command files, ' + comDir)


# Get possible entries for runs per chunk and # of processors
numProc = PipGetInteger('ProcessorNumber', 0)
zeroProcEntered = not PipGetErrNo() and numProc == 0
numRunsPerChunk = PipGetInteger('RunsPerChunk', 10)
if not PipGetErrNo() and numProc > 0:
   exitError('You cannot enter both -proc and -runs')

fullRec = None
(pid, cleanList, pointList, aliBinning, reorient, tiltLines, \
    nxRaw, nyRaw, nzRaw, pixXraw, pixYraw, pixZraw, \
    nxAli, nyAli, nzAli, pixXali, pixYali, pixZali, origXali, origYali, \
    origZali, nxVol, nyVol, nzVol, pixXvol, pixYvol, pixZvol, origXvol, \
    origYvol, origZvol, nxFull, nyFull, nzFull, pixXfull, pixYfull, pixZfull, \
    origXfull, origYfull, origZfull) = \
    getPointsAndHeaders(modelFile, objList, pointFile, progname, stackName, aliName, \
                        volName, fullRec, enteredOrient, reorientIn, comFile, \
                        useRawInput)

# Fix nzAli to be raw size for new ali binning or raw input, for raw it is used
# only for getting the maxSlices
if newAliBinning > 0 or useRawInput:
   nzAli = nzRaw

if optionValue(tiltLines, 'ExpandedByFactor', FLOAT_VALUE, numVal = 1):
   exitError('Reconstructions with an expansion factor applied are not (yet) supported')

if doCTF:
   rawArg = ''
   if useRawInput:
      rawArg = stackName
   (nx, ny, nz, xaxisTilt, tiltUseGPU, ctfXtilt, ctfUseGPU, ctfPixSize, ctfInput,
    ctfOutput) = getCTFoptionsCheckIfCorrected(progname, tiltLines, ctfLines, rawArg, \
                                               doFilter)
   (aliRoot, aliExt) = os.path.splitext(ctfOutput)
   if not useRawInput:
      splitCorrSize = fmtstr('-size {},{},{}', nxAli, nyAli, nzAli)
   if ctfUseGPU == None or ctfUseGPU < 0:
      ctfUseGPU = tiltUseGPU

   zShiftAdjustment = 0.
   if adjustForZShift:
      aliLines = readTextFile(alignCom, 'align command file')
      alignZshift = optionValue(aliLines, 'AxisZShift', FLOAT_VALUE, numVal = 1)
      shiftArr = optionValue(tiltLines, 'shift', 2, 1)
      if shiftArr and len(shiftArr) > 1:
         startingZshift = shiftArr[1]

      zShiftAdjustment = alignZshift + startingZshift

else:
   ctfOutput = aliName
   ctfInput = aliName
   ctfPixSize = pixXraw
   tiltUseGPU = optionValue(tiltLines, 'UseGPU', 1, 1, numVal = 1)
   if tiltUseGPU == None:
      tiltUseGPU = -1

# Restore the previous aligned stack
if newAliBinning > 0:
   cleanupFiles([aliName])
   if os.path.exists(aliName + '~'):
      try:
         os.rename(aliName + '~', aliName)
      except Exception:
         pass

(aliRoot, aliExt) = os.path.splitext(ctfOutput)
overrideGPU = PipGetInteger('WhenToUseGPU', -1)
if not PipGetErrNo():
   if overrideGPU == 0:
      ctfUseGPU = tiltUseGPU = -1
   elif overrideGPU == 1:
      ctfUseGPU = tiltUseGPU = 0
   elif overrideGPU > 1:
      ctfUseGPU = 0
      tiltUseGPU = -1

(comRoot, ext) = os.path.splitext(comFile)
if not comExt:
   comExt = ext
comRoot += '-sub'
rootWithDir = comRoot
optComDir = ''
if comDir:
   rootWithDir = os.path.join(comDir, comRoot)
   optComDir = '-dir "' + comDir + '"'
cleanChunkFiles(rootWithDir)
comNum = 1
tempStacks = ''

# Handle setting all the items for "aligned stack" when using raw input
# The ...ali values from the above calls are actually raw stack values
if useRawInput:
   checkForDistortion(axisLet, 1, progname)
   aliBinning = rawBinning
   if not rawPixSize:
      rawPixSize = getFallbackRawPixel(ctfPixSize, axisLet, comExt)
   headerPixSize = ctfPixSize * rawBinning
   if transposeXY:
      (nxAli, nyAli) = (nyAli, nxAli)
      origXali += pixXraw * (nxAli - nxRaw) / 2.
      origYali += pixYraw * (nyAli - nyRaw) / 2.
   nxAli = nxAli // rawBinning
   nyAli = nyAli // rawBinning
   pixXali *= rawBinning
   pixYali *= rawBinning
   ctfPixSize = rawPixSize * rawBinning

# Make command file for making new aligned stack
# adjust pixel size from com file for CTF correction by change in binning
if newAliBinning > 0:
   if debug:
      prnstr(fmtstr('newst {} -> {}', 
                    optionValue(newstLines, 'InputFile', STRING_VALUE), 
                    optionValue(newstLines, 'OutputFile', STRING_VALUE)))
   writeTextFile(makeNextComName(True), newstLines)
   ctfPixSize = (ctfPixSize * newAliBinning) / oldAliBin

# If using raw input, set up input name and optional binning com file
if useRawInput:
   ctfInput = stackName
   if rawBinning > 1:
      ctfInput = fmtstr('{}_red{}_tmp{}', aliRoot, rawBinning, aliExt)
      tempStacks += ' ' + ctfInput
      newstLines = [fmtstr('$newstack -ftreduce {} {} {}', rawBinning, stackName,
                           ctfInput)]
      if debug:
         prnstr(fmtstr('newstack reduce {} -> {}', stackName, ctfInput))
      writeTextFile(makeNextComName(True), newstLines)

# If filtering, do it in an initial command file
if doFilter:
   filtInput = ctfInput
   ctfInput = fmtstr('{}_filt_tmp{}', aliRoot, aliExt)
   tempStacks += ' ' + ctfInput
   mtfsed = [sedModify('InputFile', filtInput, delim = '|'),
             sedModify('OutputFile', ctfInput, delim = '|'),
             sedModify('PixelSize', ctfPixSize, delim = '|')]
   mtfMod = pysed(mtfsed, mtfLines, delim = '|')
   if debug:
      prnstr(fmtstr('mtffilter {} -> {}', filtInput, ctfInput))
   writeTextFile(makeNextComName(True), mtfMod)

if not doCTF:
   ctfOutput = ctfInput

if doErase:
   eraseLines = readTextFile(eraseComFile, 'Gold erasing command file')
   eraseName = fmtstr('{}_erase_tmp{}', aliRoot, aliExt)
   eraseSed = [sedModify('InputFile', ctfOutput, delim = '|'),
            sedModify('OutputFile', eraseName, delim = '|')]
   if useRawInput:
      rawEraseFid = backTransformEraseModel(eraseLines, aliXformFile, 10. * rawPixSize)
      eraseSed.append(sedModify('ModelFile', rawEraseFid, delim = '|'))
      tempStacks += ' ' + rawEraseFid

   eraseMod = pysed(eraseSed, eraseLines, delim = '|')
   eraseMod.append(fmtstr('$b3drename "{}" "{}"', eraseName, ctfOutput))
   
thickness = ySize
numSlices = zSize // aliBinning
if reorient:
   thickness = zSize
   numSlices = ySize // aliBinning

# Get values needed for making volume stack and model
xCenter = (xSize // aliBinning) / 2.
yCenter = (ySize // aliBinning) / 2.
zFinal = zSize // aliBinning
zCenter = zFinal / 2. - 0.5

# ASSUMING CENTER ALIGNED STACK, could use origins to overcome this
if transposeXY:
   (nxRaw, nyRaw) = (nyRaw, nxRaw)
sssx = (nxRaw - nxAli * aliBinning) // 2
sssy = (nyRaw - nyAli * aliBinning) // 2

sedcomBase = [sedModify('IMAGEBINNED', aliBinning, delim = '|'),
              sedModify('FULLIMAGE', fmtstr('{} {}', nxRaw, nyRaw), delim = '|'),
              sedModify('SUBSETSTART', fmtstr('{} {}', sssx, sssy), delim = '|'),
              sedModify('THICKNESS', thickness, delim = '|'),
              '|savework|d'] + \
   sedDelAndAdd('WIDTH', xSize, 'THICKNESS', delim = '|') + \
   sedDelAndAdd('XSubsetLoadRatio', 1.2, 'THICKNESS', delim = '|')

if ctfOutput != aliName:
   sedcomBase.append(sedModify('InputProjections', ctfOutput, delim = '|'))
if useRawInput:
   sedcomBase += sedDelAndAdd('UseUnalignedImages', 1, 'THICKNESS', delim = '|') + \
                 sedDelAndAdd('AlignTransformFile', aliXformFile, 'THICKNESS', 
                              delim = '|') + \
                 [sedModify('IMAGEBINNED', str(rawBinning), delim = '|')]

if typeExt and typeExt != 'mrc':
   sedcomBase.append('|IMOD_OUTPUT_FORMAT|d')

if doCTF:
   # Work out GPU for CTF which can be independent: if it is on it stays on
   # Do CTF in parallel if multiple procs explicitly entered and tilt is using GPUs,
   # or if not using a GPU for CTF and no number was entered - assume 8 in thta case
   maxSlices = 0
   procForCTF = numProc
   if numProc == 0:
      procForCTF = 8
   if numProc > 1 or (numProc == 0 and ctfUseGPU < 0):
      numChunks = 3 * procForCTF
      maxSlices = max(1, (nzAli + numChunks - 1) // numChunks)
                     
# Get format string for particle names to get equal digits on all
numDec = 1
numPts = len(pointList)
delMatch = '-[0-9]'
while numPts > 9:
   numDec += 1
   numPts = numPts // 10
   delMatch += '[0-9]'
numFormat = '-{:0' + str(numDec) + 'd}'
numPts = len(pointList)

# Loop through all the points, getting their command file text and z shifts
numPtTot = 0
minZshift = 1.e10
maxZshift = -1.e10
shiftList = []
sedList = []
volNames = []
for ptNum in range(numPts):
   numUse = numPtTot + 1
   if skipVolNumbers:
      numUse = ptNum + 1

   chunkBase = rootName + fmtstr(numFormat, numUse)

   # Use a forward slash so output is stable and tests work with Windows python
   if outDir:
      chunkBase = outDir + '/' + chunkBase
   chunkName = chunkBase + '.mrc'
   if reorient:
      chunkName = chunkBase + '.tmp'
   sedcom = copy.deepcopy(sedcomBase)
   sedcom.append(sedModify('OutputFile', chunkName, delim = '|'))
   point = pointList[ptNum]

   # Convert points from a point list by the header transformation to match scaled
   # values that came in from model conversion
   if not modelFile:
      point[0] = point[0] * pixXvol - origXvol;
      point[1] = point[1] * pixYvol - origYvol;
      point[2] = point[2] * pixZvol - origZvol;
   
   # Now need to get slice range and X/Z shifts.  X is easy and invariant
   xInAli = (point[0] + origXali) / pixXali
   xShift = aliBinning * (nxAli / 2. - xInAli)

   # For no reorientation, Y comes from Z, Y from Y; z shift is negative of coordinate
   if reorient == 0:
      yInAli = (point[2] + origYali) / pixYali
      zShift = -aliBinning * (point[1]) / pixXali

   # For rotation, Y comes from Y, Z from inversion of Y
   elif reorient < 0:
      yInAli = (point[1] + origYali) / pixYali
      zShift = aliBinning * (point[2]) / pixXali

   # For flip, the origins were not swapped in the header, so undo the origin that was
   # applied and adjust by origin that should have been applied
   else:
      yInAli = (point[1] + origYvol - origZvol - origYali) / pixYali
      zShift = -aliBinning * (point[2] + origZvol - origYvol) / pixXali


   # Get the slice range, skip if too far out of range
   sliceStart = int(round(yInAli - numSlices / 2.)) * aliBinning
   sliceEnd = sliceStart + numSlices * aliBinning - 1
   if yInAli < numSlices / 6. or yInAli > nyAli - numSlices / 6.:
      prnstr(fmtstr('WARNING: {} - Point # {} is skipped; it requires too many Y ' +\
                  'slices outside the reconstructable range for this aligned stack',
                    progname, ptNum + 1))
      continue

   # And set up for blank slices if partly out of the range
   # Get the binned slice range that Tilt will use.  It uses slices numbered from 1
   # and rounds up when binning, but (sl0 + 1 + bin - 1) / bin = sl0 / bin + 1
   # numbered from 1, which is oddly just sl0 / bin numbered from 0.
   newstRange = ''
   binSliceStart = sliceStart // aliBinning
   binSliceEnd = sliceEnd // aliBinning
   if sliceStart < 0 or binSliceEnd >= nyAli:
      if sliceStart < 0:
         numBlank = numSlices - (binSliceEnd + 1)
         newstRange = fmtstr('{}-{}', -numBlank, binSliceEnd)
         sliceStart = 0
      else:
         numBlank = numSlices - (nyAli - binSliceStart)
         newstRange = '0-' + str(numSlices - 1)
         sliceEnd = nyAli * aliBinning - 1
      prnstr(fmtstr('WARNING: {} - Point # {} is near the edge of the aligned stack' +\
                    ' in Y and requires {} blank slices', progname, ptNum + 1,
                    numBlank))

   # Finish the sed com, process the lines
   sedcom += sedDelAndAdd('SHIFT', fmtstr('{} {}', round(xShift, xzShiftDecimals),
                                          round(zShift, xzShiftDecimals)), 'THICKNESS',
                          delim = '|')
   sedcom += sedDelAndAdd('SLICE', fmtstr('{} {}', sliceStart, sliceEnd), 'THICKNESS',
                          delim = '|')
   if overrideGPU >= 0:
      sedcom += sedDelAndAdd('UseGPU', tiltUseGPU, 'THICKNESS', delim = '|')
   sedLines = pysed(sedcom, tiltLines, delim = '|')

   # Add blank slices if needed
   if newstRange:
      sedLines += [fmtstr('$newstack -blank -sec {} "{}" "{}"', newstRange, chunkName,
                          chunkBase + '.tmp2'), 
                   '$b3dremove "' + chunkName + '"']
      chunkName = chunkBase + '.tmp2'
   
   # Add final reorientation
   if reorient:
      oper = 'rotx'
      if reorient > 0:
         oper = 'flipyz'
      sedLines += [fmtstr('$clip {} "{}" "{}"', oper, chunkName, chunkBase + '.mrc'), 
                   '$b3dremove "' + chunkName + '"']

   # Maintain min/max, save the lines and the z shift
   minZshift = min(minZshift, zShift)
   maxZshift = max(maxZshift, zShift + 0.01)
   sedList.append(sedLines)
   shiftList.append(zShift)
   numPtTot += 1
   if makeVolStacks:
      volNames.append(chunkBase + '.mrc')

if not numPtTot:
   exitError('There are no points to do because all points are being skipped')
      
# if doing CTF, figure out the range of Z etc
if doCTF:
   ubPixSize = ctfPixSize / aliBinning
   fullZpixels = maxZshift - minZshift
   fullZnm = fullZpixels * ubPixSize
   if maxZextent > 0:
      numZlevels = int(math.ceil(fullZnm / maxZextent))
   if numZlevels < 2 and maxZextent > 0:
      prnstr(fmtstr('WARNING: {} - The Z extent of {:.0f} nm will result in only one Z' +\
                    ' level because the range of center positions is {:.0f} nm', progname,
                    maxZextent, fullZnm))

   zExtentNm = fullZnm / numZlevels
   prnstr(fmtstr('CTF corrections will be computed at {} levels that are {:.0f} nm thick',
                 numZlevels, zExtentNm))
   delZpixels = zExtentNm / ubPixSize
   numInLevel = []
   ptIndList = []
   invertOpt = optionValue(ctfLines, 'InvertTiltAngles', BOOL_VALUE)
   if invertZoffsets < 0 and invertOpt:
      invertZoffsets = 1
   defocusTol = optionValue(ctfLines, 'DefocusTol', INT_VALUE, numVal = 1) 
   zNmInt = int(round(zExtentNm))
   if defocusTol:
      defocusTol = min(defocusTol, zNmInt)
   else:
      defocusTol = zNmInt

   # Test for X tilt consistency and if it matters
   checkXtiltCtfVsRec(xaxisTilt, ctfXtilt, warnXtiltCrit, 0.1 * nyAli * pixXali,
                       zExtentNm, 'level', progname)

   # Find the number in each level
   for level in range(numZlevels):
      shiftLow = minZshift + level * delZpixels
      shiftHigh = shiftLow + delZpixels
      numTmp = 0
      levelInds = []
      for ind in range(numPtTot):
         if shiftList[ind] >= shiftLow and shiftList[ind] < shiftHigh:
            numTmp += 1
            levelInds.append(ind)
      numInLevel.append(numTmp)
      ptIndList.append(copy.deepcopy(levelInds))

else:
   numZlevels = 1
   delZpixels = maxZshift - minZshift
   numInLevel = [numPtTot]
   ptIndList = [list(range(numPtTot))]

# No processors entered and using GPU for tilt, assume 1
if numProc == 0 and tiltUseGPU != None and tiltUseGPU >= 0:
   numProc = 1

# Loop on levels if any
for level in range(numZlevels):
   numPtsLvl = numInLevel[level]
   if not numPtsLvl:
      continue
   maxChunkLvl = maxChunks
   numOptimal = 1000
   if doCTF:
      maxChunkLvl = int(round((0.9 * maxChunks * numPtsLvl) / numPtTot))
      numOptimal = (numOptimal * numPtsLvl) // numPtTot
      
   if numProc > 1 or zeroProcEntered:

      # If # of processors entered, try for a large # of chunks per processor but lower it
      # to give fewer than 1000 chunks; in any case limit chunks to maximum and to # pts
      numChunkPerProc = maxChunkPerProc
      while numChunkPerProc >= minChunkPerProc:
         numChunks = min(numPtsLvl, numChunkPerProc * numProc, maxChunkLvl)
         if numChunks < numOptimal:
            break;
         numChunkPerProc -= 1

   else:

      # Otherwise base it on default or entered # of runs per chunks; but it must be
      # raised if that gives too many chunks, or lower it if it is too few
      minRuns =  numPtsLvl // maxChunkLvl + 1
      minChunksLvl = min(numPtsLvl, minChunks)
      runsPerChunkLvl = max(numRunsPerChunk, minRuns)
      numChunks = max(minChunksLvl, (numPtsLvl + runsPerChunkLvl - 1) // runsPerChunkLvl)

   runsPerChunkLvl = numPtsLvl // numChunks

   if doCTF:
      offsetPix = ((level + 0.5) * delZpixels + minZshift + zShiftAdjustment) / aliBinning
      if invertZoffsets > 0:
         offsetPix = -offsetPix
      ctfsed = [sedModify('InputStack', ctfInput, delim = '|')] +\
               sedDelAndAdd('OffsetInZ', round(offsetPix, 2), 'DefocusFile',
                            delim = '|')+\
         sedDelAndAdd('UseGPU', ctfUseGPU, 'DefocusFile', delim = '|') + \
         sedDelAndAdd('DefocusTol', fmtstr('{}', defocusTol), 'DefocusFile',
                      delim = '|')

      if useRawInput or newAliBinning > 0:
         ctfsed.append(sedModify('PixelSize', ctfPixSize, delim = '|'))

      if useRawInput:
         ctfsed += ['|TransformFile|d'] + \
                   sedDelAndAdd('XAxisTilt', xaxisTilt, 'DefocusFile', delim = '|') + \
                   sedDelAndAdd('AxisAngle', axisAngle, 'DefocusFile', delim = '|')

      ctfMod = pysed(ctfsed, ctfLines, delim = '|')
      if debug:
         prnstr(fmtstr('ctfcorrection  {} -> {}', ctfInput, 
                       optionValue(ctfMod, 'OutputFileName', STRING_VALUE)))
   
      if maxSlices:
         ctfComOut = 'ctfcorrection.tmp.com'
         writeTextFile(ctfComOut, ctfMod)
         try:
            splitLines = runcmd(fmtstr('splitcorrection -i {} -o -m {} -b {} -uni ' +\
                                       '{} {} -r {} {}', comNum, maxSlices, boundPixels, 
                                       splitCorrSize, optComDir, comRoot, ctfComOut))
            numAdded = findSplitComNumber(splitLines, 'output of splitcorrection')
            comNum += numAdded

         except ImodpyError:
            exitFromImodError(progname)

      else:
         ctfComOut = makeNextComName(True)
         writeTextFile(ctfComOut, ctfMod)

   if doErase:
      if debug:
         prnstr(fmtstr('eraser  {} <-> {}', ctfOutput, eraseName))
      eraseName = makeNextComName(True)
      writeTextFile(eraseName, eraseMod)
      
   # Loop on chunks
   indInLevel = 0
   for chunk in range(numChunks):
      numInChunk = runsPerChunkLvl
      if chunk < numPtsLvl % numChunks:
         numInChunk += 1

      comName = makeNextComName(False)
      chunkLines = []
      if typeExt and typeExt != 'mrc':
         chunkLines = ['$setenv IMOD_OUTPUT_FORMAT MRC']
      for indInChunk in range(numInChunk):
         ptInd = ptIndList[level][indInLevel]
         if debug:
            prnstr(fmtstr('tilt  {} -> {}', 
                          optionValue(sedList[ptInd], 'InputProjections', STRING_VALUE),
                          optionValue(sedList[ptInd], 'OutputFile', STRING_VALUE)))
         indInLevel += 1

         # Add lines to chunk
         chunkLines += sedList[ptInd]

      # Write the file
      writeTextFile(comName, chunkLines)

   if doCTF:
      comName = makeNextComName(True)
      writeTextFile(comName, [fmtstr('$b3dremove {}', ctfOutput)])

if makeVolStacks:

   # Get number of stacks and format for numbering
   numStacks = (numPtTot + makeVolStacks - 1) // makeVolStacks
   numDec = 1
   numPts = numStacks
   while numPts > 9:
      numDec += 1
      numPts = numPts // 10
   numFormat = '-vol{:0' + str(numDec) + 'd}'

   # Make one sync file; loop on stacks
   comName = makeNextComName(True)
   comLines = []
   for group in range(numStacks):

      # Get range of subvols and name for output
      (start, end) = balancedGroupLimits(numPtTot, numStacks, group)
      stackRoot = rootName + fmtstr(numFormat, group + 1)
      if outDir:
         stackRoot = outDir + '/' + stackRoot
      comLines += ['$setenv IMOD_OUTPUT_FORMAT MRC',
                  '$newstack -StandardInput',
                  'OutputFile ' + stackRoot + '.mrc']

      # Add the input file names and accumulate point lines
      ptLines = []
      for ind in range(start, end + 1):
         comLines.append('InputFile ' + volNames[ind])
         contTime = ind + 1 - start
         ptLines.append(fmtstr('1 {} {} {} {} {}', contTime, xCenter, yCenter, zCenter,
                               contTime))

      # Write point file, add commands to convert to model, and remove point file
      writeTextFile(stackRoot + '.pt', ptLines)
      comLines.append(fmtstr('$alterheader -volstack {} "{}.mrc"', zFinal, stackRoot))
      comLines.append(fmtstr('$point2model -scat -time -sphere 3 -image "{0}.mrc"' +\
                             ' "{0}.pt" "{0}.mod"', stackRoot))
      comLines.append('$b3dremove "' + stackRoot + '.pt"')
      
   # After all runs, remove the subvols
   cleanRoot = rootName + delMatch
   if outDir:
      cleanRoot = outDir + '/' + cleanRoot
   comLines.append('$b3dremove -g "' + cleanRoot + '*.mrc"')
   writeTextFile(comName, comLines)

# Write finish file
finlines = [fmtstr('$b3dremove -g "{0}-[0-9][0-9][0-9]*.com*" ' + \
                     '"{0}-[0-9][0-9][0-9]*.log*" "{0}-finish*.com*"', rootWithDir)]
if tempStacks:
   finlines.append('$b3dremove ' + tempStacks)
if doCTF and maxSlices > 0:
   finlines.append('$b3dremove ' + ctfComOut)
writeTextFile(rootWithDir + '-finish.com', finlines)
prnstr(fmtstr('Created {} command files for {} subtomograms; run them with:', comNum,
              numPtTot))
prnstr(fmtstr('    "subm {0}*.com"   or   "processchunks ... {0}"', rootWithDir))
if doCTF and maxSlices:
   cleanupFiles([ctfComOut])
sys.exit(0)