# ============================================================
# Dual Channel Fluorescence Voxel Counting Workflow - Jython
# Requires: FIJI / ImageJ
# Opens paired channel files (_C0, _C1, _C2), merges to greyscale,
# then processes and counts voxels at eleven threshold sensitivities.
# Displays maximum intensity Z projections for validation before saving.
# Records chosen threshold to a summary CSV.
# Saves all Voxel Counter fields to a single CSV per sample.
# Includes Isolated Cell Check with exact cell count for ROI segmentation.
# Includes Needs Cleanup option with exclusion ROI drawing on chosen threshold.
# ROI-corrected values appended to per-sample CSV alongside original values.
# Summary CSV records only ROI-corrected chosen threshold volume.
# Saves drawn ROIs to cell subfolders for future reference.
# Uses virtual stacks for faster file opening of large z-stacks.
#
# NOTE: This version does NOT call the Voxel Counter plugin UI.
#       Voxel counting is performed directly in Jython by iterating over
#       the binary stack pixel by pixel, computing all the same fields
#       as the Voxel Counter plugin without triggering any save prompts.
#
# NOTE: Jython uses Python 2.7 syntax inside FIJI.
#       Key differences from Python 3:
#       - print is a statement not a function: print "text"
#       - division of integers is floor division by default
#       - strings are bytes by default, not unicode
# ============================================================

from ij import IJ, WindowManager
from ij.gui import GenericDialog, WaitForUserDialog
from ij.plugin import ZProjector, ImageCalculator
from ij.plugin.frame import RoiManager
from ij.process import AutoThresholder
from java.io import File, FileWriter, BufferedWriter
from loci.plugins import BF
from loci.plugins.in import ImporterOptions
import java.lang.System as System
import os

# Voxel calibration constants -- adjust if acquisition settings change
PIXEL_WIDTH  = 0.104      # microns per pixel (X)
PIXEL_HEIGHT = 0.104      # microns per pixel (Y)
VOXEL_DEPTH  = 0.216683   # microns per pixel (Z)
VOXEL_UNIT   = "micron"

# ============================================================
# HELPER FUNCTIONS
# ============================================================

def get_directory(title):
    """
    Opens a FIJI directory chooser dialog and returns the selected path.
    Appends a trailing slash for safe path concatenation.
    """
    dc = IJ.getDirectory(title)
    if dc is None:
        raise RuntimeError("No directory selected for: " + title)
    return dc


def make_directory(path):
    """
    Creates a directory at the given path if it does not already exist.
    Uses Java's File class for compatibility inside FIJI's JVM.
    """
    d = File(path)
    if not d.exists():
        d.mkdirs()


def save_string(content, path):
    """
    Writes a string to a file at the given path, overwriting existing content.
    Uses Java's BufferedWriter for reliable file writing inside FIJI's JVM.
    """
    bw = BufferedWriter(FileWriter(path, False))
    bw.write(content)
    bw.close()


def append_string(content, path):
    """
    Appends a string to an existing file at the given path.
    Uses Java's BufferedWriter in append mode.
    """
    bw = BufferedWriter(FileWriter(path, True))
    bw.write(content)
    bw.close()


def open_virtual(path):
    """
    Opens an image as a virtual stack using Bio-Formats.
    Virtual stacks load slices from disk on demand rather than reading
    the entire file into RAM at once, significantly reducing initial
    load time for large z-stacks (~100MB files with 120 slices).
    """
    options = ImporterOptions()
    options.setId(path)
    options.setAutoscale(False)
    options.setVirtual(True)
    imps = BF.openImagePlus(options)
    return imps[0]


def z_project_max(imp):
    """
    Generates a Maximum Intensity Z Projection of the given ImagePlus stack.
    Returns the resulting projected ImagePlus.
    Uses ZProjector with MAX_METHOD for brightest-point projection across slices.
    """
    zp = ZProjector(imp)
    zp.setMethod(ZProjector.MAX_METHOD)
    zp.doProjection()
    return zp.getProjection()


def close_image(imp):
    """
    Safely closes an ImagePlus without triggering a save prompt.
    Sets imp.changes = False first to suppress FIJI's unsaved changes dialog.
    """
    if imp is not None:
        imp.changes = False
        imp.close()


def close_all_image_windows():
    """
    Closes all currently open image windows in FIJI.
    Called between samples to prevent window accumulation across the batch.
    Sets changes = False on each image to suppress save prompts.
    """
    ids = WindowManager.getIDList()
    if ids is None:
        return
    for image_id in ids:
        imp = WindowManager.getImage(image_id)
        close_image(imp)


def apply_voxel_calibration(imp):
    """
    Sets voxel dimensions on the given ImagePlus using the calibration constants
    defined at the top of the script. Preserves existing channel, slice and
    frame counts to avoid overwriting stack structure.
    """
    n_ch = imp.getNChannels()
    n_sl = imp.getNSlices()
    n_fr = imp.getNFrames()
    IJ.run(imp, "Properties...",
           "channels=" + str(n_ch) +
           " slices=" + str(n_sl) +
           " frames=" + str(n_fr) +
           " pixel_width=" + str(PIXEL_WIDTH) +
           " pixel_height=" + str(PIXEL_HEIGHT) +
           " voxel_depth=" + str(VOXEL_DEPTH) +
           " unit=" + VOXEL_UNIT)


def force_black_white_display(imp):
    """
    Forces a binary image to display correctly as black and white.
    Resets any threshold colouring, applies a greyscale LUT, inverts it
    so thresholded regions appear white, and sets the display range to 0-255.
    """
    IJ.run(imp, "Select None", "")
    imp.getProcessor().resetThreshold()
    IJ.run(imp, "Grays", "")
    IJ.run(imp, "Invert LUT", "")
    imp.setDisplayRange(0, 255)
    imp.updateAndDraw()


def format_number(value, digits):
    """
    Formats a float to a given number of decimal places using FIJI's IJ.d2s().
    Used for consistent numeric formatting in CSV output.
    """
    return IJ.d2s(float(value), digits)


def csv_field_name(field_name):
    """
    Appends units to field names where relevant for CSV column headers.
    Keeps field names consistent with what the Voxel Counter plugin outputs.
    """
    if field_name == "Voxel size":
        return field_name + " (um)"
    if field_name == "Thresholded volume":
        return field_name + " (um^3)"
    if field_name == "Average volume per slice":
        return field_name + " (um^3)"
    if field_name == "Total ROI volume":
        return field_name + " (um^3)"
    if field_name == "Volume of stack":
        return field_name + " (um^3)"
    return field_name


def calculate_voxel_counter_fields(imp):
    """
    Replicates all Voxel Counter plugin output fields by directly iterating
    over the pixels of the binary stack in Jython.
    This avoids calling the Voxel Counter plugin entirely, which eliminates
    the Results table save prompt that appears when calling the plugin from
    outside the IJM macro interpreter context.

    White pixels (value 255) are counted as thresholded (foreground) voxels.
    If a ROI is set on the image, counting is restricted to pixels within that ROI.

    Returns a dictionary mapping field name to string value, matching the
    field names used throughout the rest of the script.
    """
    apply_voxel_calibration(imp)

    roi = imp.getRoi()
    width = imp.getWidth()
    height = imp.getHeight()
    n_slices = imp.getStackSize()
    stack = imp.getStack()

    thresholded_voxels = 0
    roi_pixels_per_slice = 0

    if roi is None:
        # No ROI -- count all pixels in the full frame
        roi_pixels_per_slice = width * height
        for z in range(1, n_slices + 1):
            ip = stack.getProcessor(z)
            for y in range(height):
                for x in range(width):
                    if ip.getPixel(x, y) == 255:
                        thresholded_voxels += 1
    else:
        # ROI present -- count only pixels within the ROI boundary and mask
        bounds = roi.getBounds()
        mask = roi.getMask()

        # Count ROI pixels per slice for Total ROI Voxels calculation
        for y in range(bounds.y, bounds.y + bounds.height):
            for x in range(bounds.x, bounds.x + bounds.width):
                inside = True
                if mask is not None:
                    inside = mask.getPixel(x - bounds.x, y - bounds.y) != 0
                if inside:
                    roi_pixels_per_slice += 1

        # Count thresholded voxels within ROI across all slices
        for z in range(1, n_slices + 1):
            ip = stack.getProcessor(z)
            for y in range(bounds.y, bounds.y + bounds.height):
                for x in range(bounds.x, bounds.x + bounds.width):
                    inside = True
                    if mask is not None:
                        inside = mask.getPixel(x - bounds.x, y - bounds.y) != 0
                    if inside and ip.getPixel(x, y) == 255:
                        thresholded_voxels += 1

    total_roi_voxels = roi_pixels_per_slice * n_slices
    voxels_in_stack = width * height * n_slices
    avg_voxels_per_slice = float(thresholded_voxels) / float(n_slices)

    # Volume fraction as percentage of thresholded voxels within ROI
    volume_fraction = 0.0
    if total_roi_voxels > 0:
        volume_fraction = 100.0 * float(thresholded_voxels) / float(total_roi_voxels)

    # Get calibration from image for volumetric calculations
    cal = imp.getCalibration()
    voxel_volume = cal.pixelWidth * cal.pixelHeight * cal.pixelDepth
    unit = cal.getUnit()
    if unit is None or unit == "":
        unit = "pixel"

    values = {}
    values["Thresholded voxels"]       = str(thresholded_voxels)
    values["Average voxels per slice"] = format_number(avg_voxels_per_slice, 3)
    values["Total ROI Voxels"]         = str(total_roi_voxels)
    values["Volume fraction"]          = format_number(volume_fraction, 3)
    values["Voxels in stack"]          = str(voxels_in_stack)
    values["Voxel size"]               = (
        format_number(cal.pixelWidth, 6) + " x " +
        format_number(cal.pixelHeight, 6) + " x " +
        format_number(cal.pixelDepth, 6) + " " + unit
    )
    values["Thresholded volume"]       = format_number(thresholded_voxels * voxel_volume, 6)
    values["Average volume per slice"] = format_number(avg_voxels_per_slice * voxel_volume, 6)
    values["Total ROI volume"]         = format_number(total_roi_voxels * voxel_volume, 6)
    values["Volume of stack"]          = format_number(voxels_in_stack * voxel_volume, 6)

    return values


def run_voxel_counter(imp):
    """
    Runs the native Jython voxel counting implementation on the given ImagePlus.
    Wraps calculate_voxel_counter_fields() for consistent call syntax
    throughout the script.
    Returns a dictionary of field name to string value.
    """
    return calculate_voxel_counter_fields(imp)


def run_voxel_counter_with_roi(thresh_imp, roi, field_names):
    """
    Applies a given ROI to a thresholded ImagePlus and runs the native
    voxel counting implementation restricted to that ROI.
    Returns a dictionary of field name to string value.
    """
    thresh_imp.setRoi(roi)
    return run_voxel_counter(thresh_imp)


def clear_outside_roi_stack(imp, roi):
    """
    Clears (sets to 0) all pixels outside the given ROI across every slice
    of the stack. Used for the cleanup workflow to remove unwanted regions
    from the binary mask before counting.
    After clearing, forces correct black/white display on the image.
    """
    width = imp.getWidth()
    height = imp.getHeight()
    stack = imp.getStack()
    n_slices = imp.getStackSize()
    bounds = roi.getBounds()
    mask = roi.getMask()

    for z in range(1, n_slices + 1):
        ip = stack.getProcessor(z)
        for y in range(height):
            for x in range(width):
                inside = False
                if (x >= bounds.x and x < bounds.x + bounds.width and
                        y >= bounds.y and y < bounds.y + bounds.height):
                    if mask is None:
                        inside = True
                    else:
                        inside = mask.getPixel(x - bounds.x, y - bounds.y) != 0
                if not inside:
                    ip.putPixel(x, y, 0)

    IJ.run(imp, "Select None", "")
    force_black_white_display(imp)


def save_roi(roi, path, image_name, cell_num):
    """
    Saves a single ROI to a .zip file in the given directory using the ROI Manager.
    The ROI is added to a fresh ROI Manager instance, saved, then cleared.
    This allows the ROI to be reopened in FIJI later for reference.
    """
    rm = RoiManager.getInstance()
    if rm is None:
        rm = RoiManager()
    rm.reset()
    rm.addRoi(roi)
    rm.runCommand("Save", path + image_name + "_cell" + str(cell_num) + "_ROI.zip")
    rm.reset()


def draw_roi_on_image(imp, dialog_title, dialog_message):
    """
    Activates the freehand selection tool, brings the image to the front,
    and prompts the user to draw an ROI via a WaitForUserDialog.
    Uses WaitForUserDialog instead of GenericDialog so the user can interact
    with the image window to draw the ROI before clicking OK.
    Returns the drawn ROI, or None if no ROI was drawn.
    """
    imp.show()
    IJ.setTool("freehand")
    IJ.selectWindow(imp.getTitle())
    WaitForUserDialog(dialog_title, dialog_message).show()
    roi = imp.getRoi()
    return roi


def process_cell_with_roi(roi, cell_num, cell_count, image_name, sample_dir,
                           multipliers, multiplier_labels, field_names,
                           chosen, needs_cleanup, summary_file, all_values,
                           thresholded_volumes):
    """
    Processes a single cell ROI on the chosen threshold TIFF only.
    All eleven threshold columns are preserved in the per-sample CSV using
    the original uncorrected values. The chosen threshold column additionally
    has ROI-corrected rows appended at the bottom, clearly labelled.
    The summary CSV records only the ROI-corrected chosen threshold volume.
    Original uncorrected values are retained but noted as potentially inaccurate
    due to ROI overlap with other cells or excluded regions.
    """
    cell_dir = sample_dir + "cell_" + str(cell_num) + "/"
    make_directory(cell_dir)
    save_roi(roi, cell_dir, image_name, cell_num)

    # Open chosen threshold TIFF and run native voxel count with ROI applied
    chosen_imp = IJ.openImage(sample_dir + image_name + "_" + chosen + "pct_processed.tif")
    parsed_corrected = run_voxel_counter_with_roi(chosen_imp, roi, field_names)
    close_image(chosen_imp)

    chosen_index = multiplier_labels.index(chosen)

    # Build per-cell CSV with all original uncorrected values intact
    csv_content = "Field," + ",".join([l + "%" for l in multiplier_labels]) + "\n"
    for f in field_names:
        row = csv_field_name(f)
        for m in range(len(multipliers)):
            row += "," + all_values.get((m, f), "N/A")
        csv_content += row + "\n"

    # Append ROI-corrected section for chosen threshold only
    csv_content += "\n"
    csv_content += "--- CHOSEN THRESHOLD ROI CORRECTIONS (" + chosen + "%) ---\n"
    csv_content += "NOTE: ROI corrections apply to chosen threshold only. "
    csv_content += "Other columns are original uncorrected values and may be "
    csv_content += "inaccurate due to ROI overlap.\n"

    for f in field_names:
        row = csv_field_name(f) + " (ROI corrected)"
        for m in range(len(multipliers)):
            if m == chosen_index:
                row += "," + parsed_corrected.get(f, "N/A")
            else:
                row += ",N/A"
        csv_content += row + "\n"

    save_string(csv_content, cell_dir + image_name + "_cell" + str(cell_num) + "_voxels.csv")

    # Record ROI-corrected Thresholded Volume in summary CSV
    corrected_volume = parsed_corrected.get("Thresholded volume", "N/A")
    cleanup_note = " - NEEDS CLEANUP" if needs_cleanup else ""
    append_string(
        image_name + "," +
        str(cell_num) + "," +
        chosen + "," +
        corrected_volume + "," +
        "ROI segmented from multi-cell image" + cleanup_note + "\n",
        summary_file
    )
    IJ.log("Recorded: " + image_name + " cell " + str(cell_num) + " -- " + chosen + "% (ROI corrected)")


def process_cleanup_single_cell(image_name, sample_dir, multipliers, multiplier_labels,
                                field_names, chosen, summary_file, all_values,
                                thresholded_volumes):
    """
    Handles the cleanup workflow for a single cell image.
    Opens the pre-threshold and chosen threshold images side by side as max
    projections, prompts the user to draw a cleanup boundary ROI on the chosen
    threshold image to define the region to KEEP (everything outside is cleared),
    then prompts for a measurement ROI to define the exact cell to count.
    Saves both original and corrected values to the per-sample CSV.
    The summary CSV records only the corrected chosen threshold volume.
    """
    IJ.log("Cleanup requested for: " + image_name + " -- prompting cleanup boundary drawing.")

    pre_imp = IJ.openImage(sample_dir + image_name + "_pre_threshold.tif")
    chosen_imp = IJ.openImage(sample_dir + image_name + "_" + chosen + "pct_processed.tif")
    pre_proj = z_project_max(pre_imp)
    chosen_proj = z_project_max(chosen_imp)
    pre_proj.setTitle("Pre-Threshold Reference - " + image_name)
    chosen_proj.setTitle("Chosen Threshold (" + chosen + "%) - " + image_name)
    pre_proj.show()
    chosen_proj.show()
    IJ.run("Tile")

    # Draw cleanup boundary -- defines region to KEEP, everything outside is cleared
    excl_roi = draw_roi_on_image(
        chosen_proj,
        "Draw Cleanup Boundary - " + image_name,
        "Draw a freehand ROI around the analysis area to KEEP.\nEverything outside this boundary will be cleared.\nClick OK when done."
    )

    close_image(pre_proj)
    close_image(chosen_proj)
    close_image(pre_imp)
    close_image(chosen_imp)

    if excl_roi is None:
        IJ.log("WARNING: No cleanup boundary drawn for " + image_name + " -- processing without cleanup boundary.")

    if excl_roi is not None:
        rm = RoiManager.getInstance()
        if rm is None:
            rm = RoiManager()
        rm.reset()
        rm.addRoi(excl_roi)
        rm.runCommand("Save", sample_dir + image_name + "_cleanup_boundary_ROI.zip")
        rm.reset()
        IJ.log("Cleanup boundary ROI saved for: " + image_name)

    # Open chosen threshold TIFF and apply cleanup boundary
    chosen_thresh_imp = IJ.openImage(sample_dir + image_name + "_" + chosen + "pct_processed.tif")
    if excl_roi is not None:
        clear_outside_roi_stack(chosen_thresh_imp, excl_roi)

    # Show cleaned result alongside pre-threshold for measurement ROI drawing
    pre_imp = IJ.openImage(sample_dir + image_name + "_pre_threshold.tif")
    pre_proj = z_project_max(pre_imp)
    cleaned_proj = z_project_max(chosen_thresh_imp)
    pre_proj.setTitle("Pre-Threshold Reference - " + image_name)
    cleaned_proj.setTitle("Cleaned Chosen Threshold (" + chosen + "%) - " + image_name)
    pre_proj.show()
    cleaned_proj.show()
    IJ.run("Tile")

    # Draw measurement ROI -- defines the exact cell region to count
    measure_roi = draw_roi_on_image(
        cleaned_proj,
        "Draw Measurement ROI - " + image_name,
        "Draw a freehand ROI around the cell to MEASURE.\nThis is separate from the cleanup boundary.\nClick OK when done."
    )

    close_image(pre_proj)
    close_image(cleaned_proj)
    close_image(pre_imp)

    if measure_roi is not None:
        save_roi(measure_roi, sample_dir, image_name, 1)
        parsed_corrected = run_voxel_counter_with_roi(chosen_thresh_imp, measure_roi, field_names)
        IJ.run(chosen_thresh_imp, "Select None", "")
    else:
        IJ.log("WARNING: No measurement ROI drawn for " + image_name + " -- counting cleaned full image.")
        parsed_corrected = run_voxel_counter(chosen_thresh_imp)

    IJ.saveAs(chosen_thresh_imp, "Tiff", sample_dir + image_name + "_" + chosen + "pct_cleaned.tif")
    close_image(chosen_thresh_imp)

    chosen_index = multiplier_labels.index(chosen)

    # Build per-sample CSV with all original uncorrected values intact
    csv_content = "Field," + ",".join([l + "%" for l in multiplier_labels]) + "\n"
    for f in field_names:
        row = csv_field_name(f)
        for m in range(len(multipliers)):
            row += "," + all_values.get((m, f), "N/A")
        csv_content += row + "\n"

    # Append corrected section for chosen threshold only
    csv_content += "\n"
    csv_content += "--- CHOSEN THRESHOLD CLEANUP + MEASUREMENT ROI CORRECTIONS (" + chosen + "%) ---\n"
    csv_content += "NOTE: Cleanup and measurement ROI corrections apply to chosen threshold only. "
    csv_content += "Other columns are original uncorrected values and may be "
    csv_content += "inaccurate due to excluded region overlap.\n"

    for f in field_names:
        row = csv_field_name(f) + " (cleanup + measurement ROI corrected)"
        for m in range(len(multipliers)):
            if m == chosen_index:
                row += "," + parsed_corrected.get(f, "N/A")
            else:
                row += ",N/A"
        csv_content += row + "\n"

    save_string(csv_content, sample_dir + image_name + "_voxels.csv")

    corrected_volume = parsed_corrected.get("Thresholded volume", "N/A")
    append_string(
        image_name + ",1," + chosen + "," + corrected_volume +
        ",Single cell confirmed - cleanup boundary and measurement ROI applied\n",
        summary_file
    )
    IJ.log("Recorded: " + image_name + " -- " + chosen + "% with cleanup and measurement ROI correction")


# ============================================================
# MAIN WORKFLOW
# ============================================================

input_dir = get_directory("Input_folder containing image stacks")
output_dir = get_directory("Output_folder to save results")

# Ask for number of channels once at the start of the batch
gd = GenericDialog("Channel Settings")
gd.addMessage("How many channels per sample?")
gd.addChoice("Number of channels:", ["2", "3"], "2")
gd.showDialog()
if gd.wasCanceled():
    raise RuntimeError("Macro cancelled by user.")
n_channels = int(gd.getNextChoice())

# Eleven threshold multipliers to compare
# lower * multiplier adjusts sensitivity: lower value = more signal captured
multipliers       = [0.7, 0.75, 0.8, 0.85, 0.9, 0.95, 1.0, 1.05, 1.10, 1.15, 1.20]
multiplier_labels = ["70", "75", "80", "85", "90", "95", "100", "105", "110", "115", "120"]

# Channel suffixes -- extend if more channels are ever needed
channel_suffixes = ["_C0", "_C1", "_C2"]

# Voxel Counter field names to calculate and export
field_names = [
    "Thresholded voxels",
    "Average voxels per slice",
    "Total ROI Voxels",
    "Volume fraction",
    "Voxels in stack",
    "Voxel size",
    "Thresholded volume",
    "Average volume per slice",
    "Total ROI volume",
    "Volume of stack"
]

# Initialise summary CSV with header row
# Cell column added to distinguish multiple cells from the same image
summary_file = output_dir + "threshold_selections.csv"
save_string("Sample,Cell,Chosen Threshold (%),Thresholded Volume (um^3),Notes\n", summary_file)

input_files = os.listdir(input_dir)

ic = ImageCalculator()

for filename in input_files:

    # Suggest Java garbage collection between samples to prevent GC pauses
    # mid-processing caused by memory from previous sample lingering in heap
    System.gc()

    # Only process _C0 files to avoid processing each pair twice
    if not (filename.endswith("_C0.tif") or filename.endswith("_C0.tiff")):
        continue

    # Get base image name and file extension
    if filename.endswith("_C0.tif"):
        image_name = filename.replace("_C0.tif", "")
        ext = ".tif"
    else:
        image_name = filename.replace("_C0.tiff", "")
        ext = ".tiff"

    # Check all required channel files exist before proceeding
    all_found = True
    for c in range(n_channels):
        channel_file = image_name + channel_suffixes[c] + ext
        if not os.path.exists(input_dir + channel_file):
            IJ.log("WARNING: Missing channel file " + channel_file + " for " + image_name + " -- skipping.")
            all_found = False
    if not all_found:
        continue

    IJ.log("Processing: " + image_name)

    # Create a subfolder for this sample in the output directory
    sample_dir = output_dir + image_name + "/"
    make_directory(sample_dir)

    # Open first channel as virtual stack to reduce initial RAM load
    merged_imp = open_virtual(input_dir + image_name + "_C0" + ext)

    # Open and sum remaining channels one by one into a 32-bit greyscale stack
    # Summing amplifies genuine signal where both channels are bright relative
    # to background noise where only one channel contributes
    for c in range(1, n_channels):
        next_imp = open_virtual(input_dir + image_name + channel_suffixes[c] + ext)
        result_imp = ic.run("Add create 32-bit stack", merged_imp, next_imp)
        close_image(merged_imp)
        close_image(next_imp)
        merged_imp = result_imp

    # STEP 1: Despeckle -- reduces salt-and-pepper noise across all slices
    IJ.run(merged_imp, "Despeckle", "stack")

    # STEP 2: Auto Brightness/Contrast -- equivalent to pressing Auto once in B&C window
    IJ.run(merged_imp, "Enhance Contrast", "saturated=0.35")

    # Save pre-threshold image for reference
    IJ.saveAs(merged_imp, "Tiff", sample_dir + image_name + "_pre_threshold.tif")

    # Generate Maximum Intensity Z Projection of pre-threshold image for validation
    pre_thresh_proj = z_project_max(merged_imp)
    pre_thresh_proj.setTitle("Pre-Threshold - " + image_name)
    pre_thresh_proj.show()

    # STEP 3: Calculate threshold from middle slice
    # Uses Default auto threshold method on middle slice as a representative sample
    # The middle slice is most likely to contain the brightest in-focus signal
    stack = merged_imp.getStack()
    n_slices = merged_imp.getNSlices()
    middle_slice = int(round(n_slices / 2.0))
    middle_proc = stack.getProcessor(middle_slice)

    # AutoThresholder computes the lower threshold value using the Default method
    thresholder = AutoThresholder()
    histogram = middle_proc.getHistogram()
    lower = thresholder.getThreshold("Default", histogram)

    proj_list = []
    all_values = {}
    thresholded_volumes = []

    # Loop through each threshold multiplier
    for m, (multiplier, label) in enumerate(zip(multipliers, multiplier_labels)):

        # Duplicate merged image for this threshold run
        dup_imp = merged_imp.duplicate()

        # Apply threshold with multiplier to whole stack
        IJ.setThreshold(dup_imp, lower * multiplier, 65535)
        IJ.run(dup_imp, "Convert to Mask", "method=Default background=Dark black")

        # STEP 4: Binary Open -- cleans up cell edges and removes stray pixels
        IJ.run(dup_imp, "Open", "stack")

        # Force correct black/white display on binary image
        force_black_white_display(dup_imp)

        # STEP 5: Apply voxel calibration for correct volumetric calculations
        apply_voxel_calibration(dup_imp)

        # Generate Maximum Intensity Z Projection for validation
        proj = z_project_max(dup_imp)
        force_black_white_display(proj)
        proj.setTitle("Threshold " + label + "% - " + image_name)
        proj.show()
        proj_list.append(proj)

        # Save processed TIFF for quality control
        IJ.saveAs(dup_imp, "Tiff", sample_dir + image_name + "_" + label + "pct_processed.tif")

        # Run native Jython voxel counting -- no plugin call, no save prompt
        parsed = run_voxel_counter(dup_imp)

        for f in field_names:
            all_values[(m, f)] = parsed[f]

        thresholded_volumes.append(parsed.get("Thresholded volume", "N/A"))

        close_image(dup_imp)

    # Tile all projections including pre-threshold for side by side validation
    IJ.run("Tile")

    # ---- ISOLATED CELL CHECK ----
    gd2 = GenericDialog("Isolated Cell Check: " + image_name)
    gd2.addMessage("Review the projections -- leftmost is pre-threshold, followed by 70% through 120%.\nHow many cells are visible?")
    gd2.addChoice("Number of cells:", ["1", "2", "3"], "1")
    gd2.addCheckbox("Needs Cleanup?", False)
    gd2.addMessage("Which threshold do you prefer?")
    gd2.addChoice("Chosen threshold:", multiplier_labels, "80")
    gd2.showDialog()
    if gd2.wasCanceled():
        raise RuntimeError("Macro cancelled by user at: " + image_name)
    cell_count    = int(gd2.getNextChoice())
    needs_cleanup = gd2.getNextBoolean()
    chosen        = gd2.getNextChoice()

    # Close all projections including pre-threshold before further processing
    for proj in proj_list:
        close_image(proj)
    close_image(pre_thresh_proj)

    if cell_count == 1 and not needs_cleanup:
        # --- SINGLE CELL, NO CLEANUP NEEDED --- normal save
        csv_content = "Field," + ",".join([l + "%" for l in multiplier_labels]) + "\n"
        for f in field_names:
            row = csv_field_name(f)
            for m in range(len(multipliers)):
                row += "," + all_values.get((m, f), "N/A")
            csv_content += row + "\n"

        save_string(csv_content, sample_dir + image_name + "_voxels.csv")

        chosen_index  = multiplier_labels.index(chosen)
        chosen_volume = thresholded_volumes[chosen_index]
        append_string(
            image_name + ",1," + chosen + "," + chosen_volume + ",Single cell confirmed\n",
            summary_file
        )
        IJ.log("Recorded: " + image_name + " -- " + chosen + "%")

    elif cell_count == 1 and needs_cleanup:
        # --- SINGLE CELL, NEEDS CLEANUP ---
        # Delegate to helper function which handles cleanup boundary drawing,
        # measurement ROI drawing, voxel counting and CSV writing
        process_cleanup_single_cell(
            image_name, sample_dir, multipliers, multiplier_labels,
            field_names, chosen, summary_file, all_values, thresholded_volumes
        )

    else:
        # --- MULTIPLE CELLS --- ROI segmentation per cell
        IJ.log("*** MULTIPLE CELLS DETECTED: " + image_name + " -- Prompting ROI segmentation ***")

        excl_roi = None
        if needs_cleanup:
            # Draw a single cleanup boundary before cell segmentation
            # Applied to chosen threshold TIFF only
            IJ.log("Cleanup requested for: " + image_name + " -- prompting cleanup boundary before cell segmentation.")

            pre_imp    = IJ.openImage(sample_dir + image_name + "_pre_threshold.tif")
            chosen_imp = IJ.openImage(sample_dir + image_name + "_" + chosen + "pct_processed.tif")
            pre_proj    = z_project_max(pre_imp)
            chosen_proj = z_project_max(chosen_imp)
            pre_proj.setTitle("Pre-Threshold Reference - " + image_name)
            chosen_proj.setTitle("Chosen Threshold (" + chosen + "%) - " + image_name)
            pre_proj.show()
            chosen_proj.show()
            IJ.run("Tile")

            excl_roi = draw_roi_on_image(
                chosen_proj,
                "Draw Cleanup Boundary - " + image_name,
                "Draw a freehand ROI around the analysis area to KEEP for all cells.\nEverything outside this boundary will be cleared.\nClick OK when done."
            )

            close_image(pre_proj)
            close_image(chosen_proj)
            close_image(pre_imp)
            close_image(chosen_imp)

            if excl_roi is None:
                IJ.log("WARNING: No cleanup boundary drawn for " + image_name + " -- continuing without cleanup boundary.")
            else:
                rm = RoiManager.getInstance()
                if rm is None:
                    rm = RoiManager()
                rm.reset()
                rm.addRoi(excl_roi)
                rm.runCommand("Save", sample_dir + image_name + "_cleanup_boundary_ROI.zip")
                rm.reset()
                IJ.log("Cleanup boundary ROI saved for: " + image_name)

                # Apply cleanup boundary to chosen threshold TIFF only
                chosen_thresh_imp = IJ.openImage(sample_dir + image_name + "_" + chosen + "pct_processed.tif")
                clear_outside_roi_stack(chosen_thresh_imp, excl_roi)
                IJ.saveAs(chosen_thresh_imp, "Tiff", sample_dir + image_name + "_" + chosen + "pct_processed.tif")
                close_image(chosen_thresh_imp)
                IJ.log("Cleanup boundary applied to chosen threshold TIFF for: " + image_name)

        # Prompt for individual cell ROIs on chosen threshold TIFF only
        for cell_num in range(1, cell_count + 1):

            pre_imp    = IJ.openImage(sample_dir + image_name + "_pre_threshold.tif")
            chosen_imp = IJ.openImage(sample_dir + image_name + "_" + chosen + "pct_processed.tif")
            pre_proj    = z_project_max(pre_imp)
            chosen_proj = z_project_max(chosen_imp)
            pre_proj.setTitle("Pre-Threshold Reference - " + image_name)
            chosen_proj.setTitle("Chosen Threshold (" + chosen + "%) - " + image_name)
            pre_proj.show()
            chosen_proj.show()
            IJ.run("Tile")

            roi = draw_roi_on_image(
                chosen_proj,
                "Draw ROI - Cell " + str(cell_num) + " of " + str(cell_count),
                "Draw a freehand ROI around cell " + str(cell_num) + " of " + str(cell_count) + ".\nClick OK when done."
            )

            close_image(pre_proj)
            close_image(chosen_proj)
            close_image(pre_imp)
            close_image(chosen_imp)

            if roi is None:
                IJ.log("ROI drawing cancelled for cell " + str(cell_num) + " of " + image_name + " -- skipping.")
                continue

            process_cell_with_roi(
                roi, cell_num, cell_count, image_name, sample_dir,
                multipliers, multiplier_labels, field_names,
                chosen, needs_cleanup, summary_file, all_values,
                thresholded_volumes
            )

    close_image(merged_imp)
    close_all_image_windows()

IJ.log("=== Batch complete. Results saved to: " + output_dir + " ===")
IJ.log("=== Threshold selections saved to: " + summary_file + " ===")