#!/usr/bin/env python
# Chunksetup - A script to set up command files to process volumes in chunks
#
# Author: David Mastronarde
#
# $Id: chunksetup,v bcbe6dbc0924 2023/10/21 04:07:29 mast $

progname = 'chunksetup'
prefix = 'ERROR: ' + progname + ' - '

# load System Libraries
import os, sys, glob

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

# Initializations - defaults for filename identifiers and tomopieces call
masterin = 'INPUTFILE'
masterout = 'OUTPUTFILE'
padstring = ''
overlap = ''
mega = ''
nofft = ''

# Fallbacks from ../manpages/autodoc2man 3 1 chunksetup
options = ["master:MasterComFile:FN:", "command:OneLineCommand:CH:",
           "input:InputImageFile:FN:", "output:OutputImageFile:FN:",
           "suffix:SuffixForOutputName:CH:", "format:FormatOfOutputFile:CH:",
           "p:PaddingPixels:I:", "o:OverlapPixels:I:", "m:MegavoxelMaximum:I:",
           "xm:XMaximumPieces:I:", "ym:YMaximumPieces:I:", "no:NoFFTSizes:B:",
           "param:ParameterFile:PF:", "help:usage:B:"]

(opts, nonopts) = PipReadOrParseOptions(sys.argv, options, progname, 3, 2, 1)

# Get file names
inInd = 0

mastercom = PipGetString('MasterComFile', '')
command = PipGetString('OneLineCommand', '')
outfile = PipGetString('OutputImageFile', '')
suffix = PipGetString('SuffixForOutputName', '')
infile = PipGetString('InputImageFile', '')
if mastercom and command:
   exitError('You cannot enter both a command file and a one-line command')
if outfile and suffix:
   exitError('You cannot enter both an output file name and a suffix for making the' +\
             ' name')
numAllowed = 3
if mastercom or command:
   numAllowed -= 1
if outfile or suffix:
   numAllowed -= 1
if infile:
   numAllowed -= 1
if numAllowed != nonopts:
   exitError(fmtstr('There are {} non-option arguments and there must be {} with the' +\
                    ' options entered', nonopts, numAllowed))

if not mastercom and not command:
   mastercom = PipGetNonOptionArg(0)
   inInd = 1

outInd = inInd
if not infile:
   infile = PipGetNonOptionArg(inInd)
   outInd = inInd + 1
if not infile:
   exitError('Image input file must be entered')

if not outfile and not suffix:
   outfile = PipGetNonOptionArg(outInd)

# Figure out final output format
nameStyle = 1
outFormat = os.getenv('IMOD_OUTPUT_FORMAT')
oformat = PipGetString('FormatOfOutputFile', '')
allowed = ('HDF', 'MRC', 'TIFF', 'TIF')
if oformat:
   if oformat.upper() not in allowed:
      exitError('Output format ' + oformat + ' is not a valid format')
   if oformat.upper() == 'TIFF':
      oformat = 'TIF'
   outFormat = oformat

if outFormat:
   outFormat = outFormat.lower()
   if outFormat == 'tiff':
      outFormat = 'tif'
   nameStyle = mapTypeExtensionToStyle(outFormat, True)
   if not nameStyle:
      nameStyle = 1
else:
   outFormat = 'mrc'

# Compose output name   
if suffix:
   (outRoot, outExt) = os.path.splitext(infile)
   if outExt.startswith('.'):
      outExt = outExt[1:]
   if outExt.upper() not in allowed:
      outRoot += '_' + outExt
   outfile = outRoot + '_' + suffix + '.' + outFormat

# Any names that are going to be split for directory need to have forward slashes
if mastercom:
   mastercom = mastercom.replace('\\', '/')
   (fullroot, comExt) = os.path.splitext(mastercom)
   if comExt not in ('.com', '.pcm'):
      comExt = '.' + defaultComExtension()
   rootname = os.path.basename(fullroot)
else:
   comSplit = command.split()
   fullroot = comSplit[0] + '-cs'
   rootname = fullroot
   comExt = '.' + defaultComExtension()
   mastercom = fullroot + comExt
   
allname = fullroot + '-all' + comExt
finishname = fullroot + '-finish' + comExt

# The input and output files may be relative to the location of the command
# file which need not be in current dir
masthead = os.path.dirname(mastercom)
infromhere = infile
cleandir = '.'
if masthead:
   infromhere = os.path.join(masthead, infile)
   cleandir = masthead

if not command and not os.path.exists(mastercom):
   exitError('Command file ' + mastercom + ' does not exist')
if not os.path.exists(infromhere):
   exitError('Image file file ' + infile + ' does not exist')

# Get rest of arguments
pad = PipGetInteger('PaddingPixels', 8)
overin = PipGetInteger('OverlapPixels', 8)
overlap = '-min ' + str(overin)
megain = PipGetInteger('MegavoxelMaximum', 0)
if not PipGetErrNo():
   mega = '-mega ' + str(megain)
maxxp = PipGetInteger('XMaximumPieces', -1)   
maxyp = PipGetInteger('YMaximumPieces', -1)
nof = PipGetBoolean('NoFFTSizes', 0)
if nof:
   nofft = '-nofft'

if not command:
   masterLines = readTextFile(mastercom, 'master command file')
else:
   masterLines = ['$' + command + ' "' + masterin + '"  "' + masterout + '"']
   
testin = False
testout = False
for l in masterLines:
   if l.find(masterin) >= 0:
      testin = True
   if l.find(masterout) >= 0:
      testout = True
   if testout and testin:
      break
else:
   exitError(fmtstr('The master command file does not contain both {} and {}', masterin,
                    masterout))

# Any filenames passed in a command line that could have a path must be quoted
tomopcom = fmtstr('tomopieces -tomo "{}" {} {} -xp {} -yp {} -zp {} {} -xmax {} -ymax {}',
                  infromhere, mega, nofft, pad, pad, pad, overlap, maxxp, maxyp)
try:
   ranlist = runcmd(tomopcom)
except ImodpyError:
   exitFromImodError(progname)

# Remove any previous files now in case the number has changed
cleanChunkFiles(fullroot)

# Get the number of pieces in X/Y/Z
if len(ranlist) < 3:
   exitError('First line of output from tomopieces does not have correct form')
npl = ranlist[0].split()
if len(npl) < 3:
   exitError('First line of output from tomopieces does not have correct form')
try:
   npiecex = int(npl[0])
   npiecey = int(npl[1])
   npiecez = int(npl[2])
except Exception:
   exitError('First line of output from tomopieces does not have correct form')
npiecetot = npiecex * npiecey * npiecez
if len(ranlist) != 1 + npiecetot + npiecex + npiecey + npiecez:
   exitError('Output from tomopieces does not have correct number of lines')

# Make the chunk coms
allines = ['# THIS IS A COMMAND FILE TO RUN ALL THE PIECES AND PUT THEM TOGETHER','#']
for ipc in range(npiecetot):
   numtext = fmtstr('-{:03d}', ipc + 1)
   minmaxes = ranlist[ipc + 1].strip().split(',')
   comfile = fullroot + numtext + comExt
   imagein = rootname + numtext + '.in'
   imageout = rootname + numtext + '.out'
   comlines = ['$taperoutvol -StandardInput',
               'InputFile ' + infile,
               'OutputFile ' + imagein,
               fmtstr('XMinAndMax {},{}', minmaxes[0], minmaxes[1]),
               fmtstr('YMinAndMax {},{}', minmaxes[2], minmaxes[3]),
               fmtstr('ZMinAndMax {},{}', minmaxes[4], minmaxes[5]),
               fmtstr('TaperPadsInXYZ {0},{0},{0}', pad)]
   if nofft:
      comlines.append('NoFFTSizes')

   comlines += pysed([fmtstr('/{}/s//{}/g', masterin, imagein),
                      fmtstr('/{}/s//{}/g', masterout, imageout)], masterLines)
   comlines += ["", fmtstr('$b3dremove {0}', imagein)]
   
   writeTextFile(comfile, comlines)

   allines += [fmtstr('$echo Working on piece  {}  of {}', ipc + 1, npiecetot),
               fmtstr('$vmstopy -x {} {}{}.log', comfile, rootname,
                      numtext)]
   
# Make the reassembly file
finishlines= ['# THIS COMMAND FILE REASSEMBLES THE PIECES', '#',
              '$assemblevol -StandardInput',
              'OutputFile ' + outfile]
addOutputFormatVarToLines(finishlines, nameStyle, outFormat, True)
ranind = npiecetot + 1
for ipc in range(npiecex):
   finishlines.append('StartEndToExtractInX ' + ranlist[ranind].strip())
   ranind += 1
for ipc in range(npiecey):
   finishlines.append('StartEndToExtractInY ' + ranlist[ranind].strip())
   ranind += 1
for ipc in range(npiecez):
   finishlines.append('StartEndToExtractInZ ' + ranlist[ranind].strip())
   ranind += 1

for ipc in range(npiecetot):
   numtext = fmtstr('-{:03d}', ipc + 1)
   finishlines.append('InputFile ' + rootname + numtext + '.out')

finishlines.append('$b3dremove -g ' + rootname + '-[0-9][0-9][0-9]*' + comExt +'* ' + \
                   rootname + \
                   '-[0-9][0-9][0-9]*.log* ' + rootname + '-[0-9][0-9][0-9]*.out*')
allines += ['$echo Reassembling pieces',
            '$vmstopy -x ' + finishname + ' ' + rootname + '-finish.log']

writeTextFile(finishname, finishlines)
writeTextFile(allname, allines)

if suffix:
   prnstr('Output file: ' + outfile + '    [CHS1]')
   prnstr('')

prnstr(fmtstr("""{} command files were generated and are ready to run.

You can run {} (with subm) to run all command files in 
sequence, reassemble the volume, and clean up intermediate image and 
command files.

If you have multiple processors available, you can use 
  processchunks machine_list {}
to do these operations in parallel, where machine_list is a comma-separated 
list of available machines or just the number of local processors""", npiecetot,
              os.path.basename(allname), rootname))

sys.exit(0)
