#!/usr/bin/python
# 2022-04: Python2 to Python3 conversion by Allison Hsiang
# 2022-07: Change from external ImageJ/FIJI installation to pyimagej API by Allison Hsiang

import os
import sys
import textwrap
import glob
import struct
import platform
import subprocess
import time
import shutil
import imagej
import jpype


def initImageJ(libPath):
    '''
    Initializes ImageJ instance with necessary plugins.
    '''
    # Specify ImageJ plugin directory
    plugins_dir = os.path.join(libPath, 'plugins')
    jpype.addClassPath(f'{plugins_dir}/*')

    # We copy the plugin directory to the working directory
    # Then we change the working directory to the one containing this script
    # This is a workaround; otherwise the JVM can't find the plugins
    # Pretty sure this is a pyimagej bug
    curr_dir = os.path.dirname(os.path.realpath(__file__))
    curr_plugins_dir = os.path.join(curr_dir, 'plugins')
    if not os.path.exists(curr_plugins_dir):
        shutil.copytree(plugins_dir, curr_plugins_dir)
    os.chdir(curr_dir)

    ij = imagej.init(mode='headless')

    return ij, curr_plugins_dir


def writeMacro(settings, obj):
    '''
    Writes macro for FIJI to generate heightmap from z-stack images.
    '''
    with open(obj.macro,'w') as f:
        if settings['macro']:
            stackReg = '\nrun("StackReg ", "transformation=[Scaled Rotation]");'
        else:
            stackReg = ''

        macroText = textwrap.dedent("""\
                    File.openSequence("{0}");{1}
                    run("Stack Focuser Batch ", "enter={2} generate");
                    selectImage("Focused_{3}");
                    saveAs("Tiff","{4}");
                    selectImage("HeightMap_{3}");
                    saveAs("Tiff","{5}");
                    close("*");
                    """)

        zStackDir = obj.stack
        kernelHM = settings['kernel_heightmap']
        tempWindow = obj.name.split('.')[0]
        focusedFileName = obj.fijiFocused
        heightmapFileName = obj.heightmap

        toWrite = macroText.format(zStackDir, stackReg, kernelHM, tempWindow, focusedFileName, heightmapFileName)
        f.write(toWrite)


def runFIJI(settings, obj, ij):
    '''
    Runs FIJI with written macro file.
    '''
    if not settings['docker']: # Use pyimagej
        with open(obj.macro, 'r') as f:
            macro = f.read()
        ij.py.run_macro(macro)

    else: # Use Docker's local FIJI installation
        command = '/tmp/Fiji.app/ImageJ-linux64'
        args=[command, '--headless', '-macro', obj.macro]
        # Suppress call output by redirecting output to DEVNULL
        FNULL = open(os.devnull,'w')
        subprocess.run(args, stdout=FNULL, stderr=subprocess.STDOUT)


def makeHeightMap(settings, obj, ij):
    '''
    Wrapper function for height map generation.
    '''
    # Start timer on object processing
    start = time.time()
    # Write FIJI macro, run FIJI on object, and save checkpoint
    writeMacro(settings, obj)
    runFIJI(settings, obj, ij)
    # End timer
    end = time.time()
    time_elapsed = end - start
    print ('\tINFO: Time elapsed: {:.3f} seconds\n'.format(time_elapsed))
