#!/usr/bin/env python
# stackframemdocs - concatenate frame stack mdoc files and make a valid image mdoc
#
# Author: David Mastronarde
#
# $Id: stackframemdocs,v 937343107256 2023/02/19 22:44:16 mast $
#

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

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

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

options = ['frame:FrameSetMdoc:FNM:Name of one frame mdoc file',
           'imageroot:ImageRootOfOutputName:FN:Image file as root name of output ' + \
           '.mdoc file', 'mode:DataMode:I:Mode of file for DataMode entry',
           'size:SizeInXandY:IP:Size in X and Y for ImageSize entry',
           'pixel:PixelSpacing:F:Pixel size in Angstroms for PixelSpacing entry',
           'help:usage:B:']

PipExitOnError(False, prefix)
(numOpts, numNonOpts) = PipParseInput(sys.argv, options)

ifHelp = PipGetBoolean('help', 0)
if numOpts + numNonOpts < 2 or ifHelp:
   PipPrintHelp(progname, 0, 2, 1)
   sys.exit(0)

# Figure out how many names are coming and get output root name
outRoot = PipGetString('ImageRootOfOutputName', '')
numInNonOpt = numNonOpts
numInByOpt = PipNumberOfEntries('FrameSetMdoc')
if not outRoot:
   numInNonOpt -= 1
   if numInNonOpt + numInNonOpt == 0:
      exitError('No input file names entered')
   if not numNonOpts:
      exitError('No image root name for output entered')
   outRoot = PipGetNonOptionArg(numInNonOpt)

# Get the list of input files
fileList = []
for ind in range(numInByOpt):
   fileList.append(PipGetString('FrameSetMdoc', ''))
for ind in range(numInNonOpt):
   fileList.append(PipGetNonOptionArg(ind))

# Other options
dataMode = PipGetInteger('DataMode', 0)
modeEntered = 1 - PipGetErrNo()
(nxIn, nyIn) = PipGetTwoIntegers('SizeInXandY', 0, 0)
sizeEntered = 1 - PipGetErrNo()
pixelSize = PipGetFloat('PixelSpacing', 0)
pixelEntered = 1 - PipGetErrNo()

allLines = []

# Loop on input files
for ind in range(len(fileList)):
   sectLines = []
   frameLines = readTextFile(fileList[ind])
   gotSect = False
   if not ind:

      # For first file, get pixel size if not entered
      if not pixelEntered:
         pixelSize = optionValue(frameLines, 'PixelSpacing', FLOAT_VALUE, numVal = 1,
                                 otherSep = '=')
         if not pixelSize:
            exitError('There is no valid PixelSpacing entry in the first frame mdoc; ' +\
                      'you must enter the pixel size')

      # and if we need mode or size, look for frame stack and get its header
      if not modeEntered or not sizeEntered:
         subPath = optionValue(frameLines, 'SubFramePath', STRING_VALUE, otherSep = '=')
         if not subPath:
            exitError('There is no entry for SubFramePath in the first frame mdoc; ' +\
                      ' you must enter the mode and X/Y size')
         subPath = subPath.replace('\\', '/')
         (oldDir, stack) = os.path.split(subPath)
         if not os.path.exists(stack):
            exitError(stack + ' does not exist in current directory; you must enter ' + \
                      'mode and X/Y size')

         try:
            (nx, ny, nz, mode, px, py, pz) = getmrc(stack)
         except ImodpyError:
            exitFromImodError(progname)

         if not modeEntered:
            dataMode = mode
         if not sizeEntered:
            nxIn = nx
            nyIn = ny

      # Compose global section
      allLines.append('ImageFile = ' + outRoot)
      allLines.append(fmtstr('ImageSize = {} {}', nxIn, nyIn))
      allLines.append('DataMode = ' + str(dataMode))
      allLines.append(fmtstr('PixelSpacing = {}', pixelSize))
         
   # loop on lines and look for FrameSet, modify that, and append the rest   
   for line in frameLines:
      if not gotSect and line.startswith('[FrameSet ='):
         gotSect = True
         line = fmtstr('[ZValue = {}]', ind)
      if gotSect:
         sectLines.append(line)

   if not gotSect:
      exitError(fileList[ind] + ' does not contain a FrameSet section')


   # Add lines to output
   allLines.append('')
   allLines += sectLines

writeTextFile(outRoot + '.mdoc', allLines)
sys.exit(0)
   
