#!/usr/bin/env python
# runraptor - Runs RAPTOR, taking care of several options and post-processing the model
#
# Author: David Mastronarde
#
# $Id: runraptor,v 937343107256 2023/02/19 22:44:16 mast $

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

# 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()
   raptorBin = os.path.join(os.environ['IMOD_DIR'], 'bin')
else:
   sys.stdout.write(prefix + " IMOD_DIR is not defined!\n")
   sys.exit(1)

#
# load IMOD Libraries
from pysed import *
from pip import exitError, setExitPrefix
setExitPrefix(prefix)

if os.getenv('RAPTOR_BIN'):
   raptorBin = os.environ['RAPTOR_BIN']
if not os.path.exists(raptorBin):
   exitError(raptorBin + 'does not exist.   You need to leave RAPTOR_BIN undefined to '+\
             'use the RAPTOR in IMOD, or define it with the path to a RAPTOR executable')

if len(sys.argv) < 5:
   prnstr("""Usage: runraptor [options] image_stack
   Runs RAPTOR on the given image_stack, which may be a raw stack or preali file
   Options other than -P[ID] for a PID are RAPTOR options and MUST include:
     -diam #  Bead diameter in pixels
     -mark #  Number of markers to track

   This script takes care of adding -exec, -path, -input, -track, and -output
       options.
   It must be run inside the data directory where files are located.
   Coarse alignment must be run first""")
   sys.exit(1)

# Set up a unique directory for the run
outdir = 'raptor1'
for i in range(1, 100):
   if not os.path.exists('raptor' + str(i)):
      outdir = 'raptor' + str(i)
      break

# Get options and build option string for raptor
options = []
gotmark = 0
gotdiam = 0
for i in range(1, len(sys.argv)):
   if sys.argv[i].startswith('-P') and sys.argv[i] in '-PID':
      printPID(True)
   else:
      options.append(sys.argv[i])
   if sys.argv[i].startswith('-ma'):
      gotmark = 1
   if sys.argv[i].startswith('-di'):
      gotdiam = 1

infile = sys.argv[len(sys.argv) - 1]

# The extension from splitext includes the . !
(rootname, ext) = os.path.splitext(infile)
useraw = ext == '.st' or ((ext in standardTypeExtensions() or ext == '.tif' or 
                           ext == '.tiff') and not rootname.endswith('_preali'))
if rootname.endswith('_preali'):
   rootname = rootname[:-7]
if os.path.basename(rootname) != rootname:
   exitError('You must run this script in the data directory')
if not os.path.exists(infile):
   exitError('The input file ' + infile + ' does not exist in the current directory')
if useraw and not os.path.exists(rootname + '.prexg'):
   exitError('Should be run on a raw stack only after running coarse alignment')

if not gotmark or not gotdiam:
   exitError('You must enter options with a diameter in pixels and a number of markers')

outname = rootname + '_raptor.fid'

# For a raw stack, figure out distortion parameters to send to xfmodel
distort = ""
binning = ""
gradient = ""
gotNewst = False
if useraw:
   for comf in ['prenewst.com', 'prenewsta.com']:
      if os.path.exists(comf):
         newstlines = readTextFile(comf)
         for l in newstlines:
            if gotNewst:
               lsplit = l.split()
               if lsplit[0] == 'DistortionField' and len(lsplit) > 1:
                  distort = '-distort ' + lsplit[1]
               if lsplit[0] == 'ImagesAreBinned' and len(lsplit) > 1:
                  binning = '-binning ' + lsplit[1]
               if lsplit[0] == 'GradientFile' and len(lsplit) > 1:
                  gradient = '-gradient ' + lsplit[1]
               
            elif l.find('newstack') >= 0:
               if l.find(infile) >= 0:
                  lsplit = l.split()
                  for ind in range(len(lsplit) - 1):
                     if lsplit[ind].startswith('-dist'):
                        distort = '-distort ' + lsplit[ind+1]
                     if lsplit[ind].startswith('-image'):
                        binning = '-binning ' + lsplit[ind+1]
                     if lsplit[ind].startswith('-grad'):
                        gradient = '-gradient ' + lsplit[ind+1]
                  break

               elif '-Standard' in l:
                  gotNewst = True
                  
# Build a command string with quoted arguments
comstr = fmtstr('"{0}/RAPTOR" -exec "{0}" -path . -inp {1} -track -out {2}', raptorBin,
                infile, outdir)
for opt in options:
   comstr += ' "' + opt + '"'

prnstr('RAPTOR results for this run will be in ' + outdir)
sys.stdout.flush()
try:
   runcmd(comstr, None)
except ImodpyError:
   logfile = outdir + '/align/' + rootname + '_RAPTOR.log'
   if os.path.exists(logfile):
      loglines = readTextFile(logfile)
      for l in loglines:
         if l.find('ERROR:') >= 0:
            prnstr(l)
   exitFromImodError(progname)

makeBackupFile(outname)

# Fix the pixel size and other features in the model
sedout = outname
if useraw:
   sedout = rootname + '.' + str(os.getpid())
try:
   (nx, ny, nz, mode, px, py, pz) = getmrc(infile)
except ImodpyError:
   exitFromImodError(progname)

sedcom = ["/symbol *circle/d",
          "/^size *7/d",
          fmtstr("/^units/a/refcurscale {} {} {}/", px, py, pz),
          "/^drawmode/a/symbol    0/",
          "/^drawmode/a/symsize   7/"]

pysed(sedcom, outdir + '/IMOD/' + rootname + '.fid.txt', sedout)

# If did raw stack, now align the model to the preali
if useraw:
   try:
      runcmd(fmtstr('xfmodel -xf {}.prexg {} {} {} {} {}',
                    rootname, distort, binning, gradient, sedout, outname), None)
   except ImodpyError:
      exitFromImodError(progname)

   cleanupFiles([sedout])

prnstr("The model that will load on the coarse aligned stack is:")
prnstr(outname)
sys.exit(0)
