#!/usr/bin/env python
# Squeezevol - run matchvol to squeeze or expand a volume
#
# Authors: Tor Mohling and David Mastronarde
#
# $Id: squeezevol,v 31c2edd1c296 2023/04/03 18:24:23 mast $
# Log at end

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

# load System Libraries
import os, sys, math

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

# Initializations
factor = 1.6
ix,iy,iz,ox,oy,oz,xs,ys,zs = 0,0,0,0,0,0,0,0,0
squeezex,squeezey,squeezez = 0,0,0
pixelx,pixely,pixelz = 0,0,0
linear = ""
pixelxyz = [1.0, 1.0, 1.0]
tempdir = ""
InputFile = ""
OutputFile = ""

# Fallbacks from ../manpages/autodoc2man 3 1 squeezevol
options = ["f:factor:F:", "e:expand:F:", "x:xFactor:F:", "y:yFactor:F:", "z:zFactor:F:",
           "ix:ixSize:I:", "iy:iySize:I:", "iz:izSize:I:", "t:tempdir:FN:",
           "l:linear:B:", ":PID:B:"]

#
# Process command-line using PIP
(opts, nonopts) = PipReadOrParseOptions(sys.argv, options, progname, 2, 1, 1)
if nonopts != 2:
   prnstr(prefix + "wrong number of arguments")
   PipPrintHelp(progname, 0, 1, 1)
   sys.exit(1)

doPID = PipGetBoolean('PID', 0)
printPID(doPID)

InputFile = PipGetNonOptionArg(0)
OutputFile = PipGetNonOptionArg(1)
if not os.path.exists(InputFile):
   exitError("input file " + InputFile + " does not exist")

t = PipGetString('tempdir', '')
if t != None and t != '':
   tempdir = "TemporaryDirectory " + t
if PipGetBoolean('linear', 0) == 1:
   linear = "InterpolationOrder 1"
factor = PipGetFloat('factor', factor)
ifFactor = 1 - PipGetErrNo()
expand = PipGetFloat('expand', 0.)
if expand and ifFactor:
   exitError("You can enter either -factor or -expand but not both")
if expand:
   factor = 1. / expand
xs = PipGetFloat('xFactor', factor)
ys = PipGetFloat('yFactor', factor)
zs = PipGetFloat('zFactor', factor)
ix = PipGetInteger('ixSize', 0)
ixEntered = 1 - PipGetErrNo()
iy = PipGetInteger('iySize', 0)
iyEntered = 1 - PipGetErrNo()
iz = PipGetInteger('izSize', 0)
izEntered = 1 - PipGetErrNo()
legacy = PipGetBoolean('LegacyInterpolation', 0)
PipDone()

if (math.fabs(xs - round(xs)) >= 0.001 or math.fabs(ys - round(ys)) >= 0.001) and \
   math.fabs(xs - ys) >= 0.001 and not legacy:
   legacy = 1
   prnstr('Using matchvol because reductions in X and Y are unequal and non-integer')

if not legacy and (xs < 1. or ys < 1. or zs < 1.):
   legacy = 1
   prnstr('Using matchvol because the volume is being expanded')

if not legacy and ixEntered + iyEntered + izEntered > 0:
   legacy = 1
   prnstr('Using matchvol because an output size was entered')
#
try:
   if legacy:
      (tx,ty,tz,mode,pixelxyz[0],pixelxyz[1],pixelxyz[2],origx,origy,origz,dmin,
       dmax,dmean) = getmrc(InputFile, True)
      #
      if not ixEntered:
         ix = tx
      if not iyEntered:
         iy = ty
      if not izEntered:
         iz = tz
      ox = int(ix / xs)
      oy = int(iy / ys)
      oz = int(iz / zs)
      squeezex = 1 / xs
      squeezey = 1 / ys
      squeezez = 1 / zs
      pixelx = round((xs * pixelxyz[0]), 3)
      pixely = round((ys * pixelxyz[1]), 3)
      pixelz = round((zs * pixelxyz[2]), 3)
      origx += 0.5 * (ox * xs - tx) * pixelxyz[0]
      origy += 0.5 * (oy * ys - ty) * pixelxyz[1]
      origz += 0.5 * (oz * zs - tz) * pixelxyz[2]
      #
      prnstr("Squeezing the volume with Matchvol...")
      matchin =  ["InputFile " + InputFile,
                  "OutputFile " + OutputFile,
                  tempdir,
                  linear,
                  fmtstr("OutputSizeXYZ {} {} {}", ox,oy,oz),
                  fmtstr("3DTransform {} 0 0 0 0 {} 0 0 0 0 {} 0", 
                         squeezex,squeezey,squeezez)]
      matchout = runcmd("matchvol -StandardInput", matchin, 'stdout')
      #
      prnstr(fmtstr("Adjusting pixel spacing in header to {} {} {}", pixelx,pixely,
                    pixelz))
      alterin = [OutputFile, 
                 "del",
                 fmtstr("{} {} {}", pixelx,pixely,pixelz),
                 "org",
                 fmtstr("{} {} {}", origx,origy,origz),
                 "done"]
      runcmd("alterheader", alterin)

   else:

      prnstr("Squeezing the volume with Binvol...")
      runcmd(fmtstr('binvol -xbin {} -ybin {} -zbin {} -anti -1 "{}" "{}"',
                    xs, ys, zs, InputFile, OutputFile), None, 'stdout')

except ImodpyError:
   exitFromImodError(progname)

sys.exit(0)
