// Global variables to be configured for each use.
INDIR = "C:\\experiment_name\\Data_folder\\";
 // Input directory, or the location of your raw images. 
OUTDIR = "C:\\experiment_name\\";
 // Output directory, or the location that your processed images will be saved. 
TPNUM = 2;
 // Number of time points.
WELL_NUM = 96
 // Number of wells. 
CHANNELS = newArray("GFP", "RFP");
 // The channels in which the images were taken. 
PRIMARY = "GFP";
 // Primary channel used to generate ROIs.
GRID_TYPE = "snake by columns"; // The image acquisition order used by Grid/Collection Stitching within Fiji.
STITCH_ORDER = "[Up & Right]";
 // The order string to be used as per Grid/Collection Stitching within Fiji.
DIM = 2;
 // Dimensions of the montage. If one image was taken enter "1."
T_O = 12;
 // Tile overlap, or how much each image in a montage overlaps with its neighbors. 

// Background subtraction variables. 
BGSUB = true; // Set to true to background subtract after registration.
RB_BALL_RADIUS = 150; // Parameter for rolling-ball background subtraction.

//Set to hide active images to expedite computation.
setBatchMode("hide");

// Utility function to ensure output directories exist.
function ensurePathExists(path) {
  if (!File.exists(path)) {
    File.makeDirectory(path);
  }
}

function acquirePaths() {
  // Acquire timepoint directories from PRIMARY directory. 
  tp_to_img_paths = newArray(TPNUM*WELL_NUM);
  timepoints = getFileList(INDIR + "\\" + PRIMARY + "\\");
  pathsFound = 0;
  // For each timepoint.
  for(tp=0; tp < timepoints.length; tp++) {
    // Specify a timepoint directory within PRIMARY channel.
    tp_dir = INDIR + "\\" + PRIMARY + "\\" + timepoints[tp];
    // Acquire list of all column directories within timepoint.
    cols = getFileList(tp_dir);
    // For each column.
    for(col=0; col < cols.length; col++) {
      // Specify a column directory.
      col_dir = tp_dir + cols[col];
      // Acquire list of images within the column directory.
      imgs = getFileList(col_dir);
      // For each image.
      for(n=0; n < imgs.length; n++) {
        //Expected string example: A1_01.tif
        s = split(imgs[n], ".");
        // If image has a 'tif' extension, then remove underscore to acquire prefix.
        if(s[1] == 'tif') {
          t = split(s[0], '_');
          // As only require well number, select well value only once for first image and ignore rest.
          if(t[1] == '01') {
            tp_to_img_paths[pathsFound] = timepoints[tp] + cols[col] + t[0];
            pathsFound++;
          }
        }
      }
    }
  }
  return Array.trim(tp_to_img_paths, pathsFound);
}

// Stitching related functions.
function mkStitchedOutputDirectories(tp_to_img_paths) {
  // Ensure all output directories exist. 
  ensurePathExists(OUTDIR + "\\stitched");
  for(p=0; p < tp_to_img_paths.length; p++) {
    // Convert a path like: 'T1/col_01/A01' to 'T1'
    path = File.getParent(File.getParent(tp_to_img_paths[p]));
    for(ch=0; ch < CHANNELS.length; ch++) {
      tppath = OUTDIR + "\\stitched\\" + CHANNELS[ch] + "\\" + path;
      chpath = File.getParent(tppath);
      ensurePathExists(tppath);
      ensurePathExists(chpath);
    }
  }
}


// Stitch primary channel first. Primary channel is stitched first and output coordinates are used to stitch other channels.
function stitchPrimaryChannel(tp_to_img_paths) {
  for(i=0; i < tp_to_img_paths.length; i++)  {
    well = File.getName(tp_to_img_paths[i]);
    tp_to_col = File.getParent(tp_to_img_paths[i]);

    // Run stitching command and ensure output coordinates are saved.
    run("Grid/Collection stitching", "type=[Grid: " + GRID_TYPE + "] order=" + STITCH_ORDER + " grid_size_x=DIM grid_size_y=DIM tile_overlap=T_O first_file_index_i=1 directory=[" + INDIR + "\\" + PRIMARY + "\\" + tp_to_col + "] file_names=" + well + "_{ii}.tif output_textfile_name=" + well + "_coords.txt fusion_method=[Linear Blending] regression_threshold=0.001 max/avg_displacement_threshold=1.5 absolute_displacement_threshold=100 compute_overlap subpixel_accuracy computation_parameters=[Save computation time (but use more RAM)] image_output=[Fuse and display]");

    imgpath = OUTDIR + "\\stitched\\" + PRIMARY + "\\" + File.getParent(tp_to_col) + "\\" + well + ".tif";
    save(imgpath);
  }
}
  
// Stitch all channels that are not the primary one.
function stitchOtherChannels(tp_to_img_paths) {
  prim_dir = INDIR + "\\" + PRIMARY;
  for(i=0; i < tp_to_img_paths.length; i++)  {
    well = File.getName(tp_to_img_paths[i]);
    tp_to_col = File.getParent(tp_to_img_paths[i]);
    for(ch=0; ch < CHANNELS.length; ch++) { 
      // Ensure that channel about to be stitched is not the primary channel, which has already been stitched.
      if (CHANNELS[ch] != PRIMARY) {
        raw_img_dir = INDIR + "\\" + CHANNELS[ch];
        // Check that this channel has tiles to be stitched.
        if(File.exists(raw_img_dir + "\\" + tp_to_img_paths[i] + "_01.tif")) {
          // Copy coordinates used to stitch primary channel tiles for this channel.
          tp = File.getParent(File.getParent(tp_to_img_paths[i]));
          // Acquire coordinate file. Plugin requires they be in the same directory as input tile images.
          File.copy(prim_dir + "\\" + tp_to_img_paths[i] + "_coords.registered.txt", raw_img_dir + "\\" + tp_to_img_paths[i] + "_coords.registered.txt");
  
          run("Grid/Collection stitching", "type=[Positions from file] order=[Defined by TileConfiguration] directory=[" + raw_img_dir + "\\" + tp_to_col + "] layout_file=" + well + "_coords.registered.txt fusion_method=[Linear Blending] regression_threshold=0.01 max/avg_displacement_threshold=1.5 absolute_displacement_threshold=100 subpixel_accuracy computation_parameters=[Save computation time (but use more RAM)] image_output=[Fuse and display]");
  
          save(OUTDIR + "\\stitched\\" + CHANNELS[ch] + "\\" + File.getParent(tp_to_col) + "\\" + well + ".tif");
          run("Close All");
        }
      }
    }
  }
}

function stitch(tp_to_img_paths) {
  stitchPrimaryChannel(tp_to_img_paths);
  stitchOtherChannels(tp_to_img_paths);
}

function mkStackedOutputDirectories(tp_to_img_paths) {
  // Ensure all output directories exist.
  ensurePathExists(OUTDIR + "\\stacked");
  for(ch=0; ch < CHANNELS.length; ch++) {
    tppath = OUTDIR + "\\stacked\\" + CHANNELS[ch];
    ensurePathExists(tppath);
  }
}

// Register the primary channel, then register the other channels with the coordinates output for the primary channel.
function stackPrimaryChannel(tp_to_img_paths) {
  for(i=0; i < tp_to_img_paths.length; i++) {
    well = File.getName(tp_to_img_paths[i]);
    if(!File.exists(OUTDIR + "\\stacked\\" + PRIMARY + "\\" + well + ".tif")) {
      for(tp=1; tp < TPNUM+1; tp++) {
        well_path = OUTDIR + "\\stitched\\" + PRIMARY + "\\T" + tp + "\\" + well + ".tif";
        if(File.exists(well_path)) {
            open(well_path);
        }
      }

      // First stack the images.
      run("Images to Stack", "method=[Copy (center)] name=" + PRIMARY + "_stack title=[] use");
      setSlice(1);

      // Then register them.
      run("MultiStackReg", "stack_1=" + PRIMARY + "_stack action_1=Align file_1=[" + OUTDIR + "\\stacked\\" + PRIMARY + "\\" + well + "_stack_matrix.txt] stack_2=None action_2=Ignore file_2=[] transformation=[Translation] save");

      // Optionally, background subtract the stack before saving.
      if (BGSUB) { 
        run("Subtract Background...", "rolling=" + RB_BALL_RADIUS + " stack");
      }
      save(OUTDIR + "\\stacked\\" + PRIMARY + "\\" + well + ".tif");
    }
  }
}

function stackOtherChannels(tp_to_img_paths) {
  for(i=0; i < tp_to_img_paths.length; i++) {
    well = File.getName(tp_to_img_paths[i]);
    for(ch=0; ch < CHANNELS.length; ch++) {
      if(CHANNELS[ch] != PRIMARY) {
        if(!File.exists(OUTDIR + "\\stacked\\" + CHANNELS[ch] + "\\" + well + ".tif")) {
          // Copy coordinate information from primary channel to this directory for use.
          File.copy(OUTDIR + "\\stacked\\" + PRIMARY + "\\" + well + "_stack_matrix.txt", OUTDIR + "\\stacked\\" + CHANNELS[ch] + "\\" + well + "_stack_matrix.txt");

          // Keep count of how many timepoints are found. This will hep ensure that at least one image is found.
          tps_found = 0;
          //Open files to stack from each timepoint
          for(tp=1; tp < TPNUM+1; tp++) {
            well_path = OUTDIR + "\\stitched\\" + CHANNELS[ch] + "\\T" + tp + "\\" + well + ".tif";
            if(File.exists(well_path)) {
              open(well_path);
              tps_found += 1;
            }
          }

          // If at least one image is open, then stack.
          if(tps_found > 0) {
            run("Images to Stack", "method=[Copy (center)] name=" + CHANNELS[ch] + "_stack title=[] use");
            setSlice(1);
                
            // Register the stack.
            run("MultiStackReg", "stack_1=" + CHANNELS[ch] + "_stack action_1=[Load Transformation File] file_1=[" + OUTDIR + "\\stacked\\" + CHANNELS[ch] + "\\" + well + "_stack_matrix.txt] stack_2=None action_2=Ignore file_2=[] transformation=[Translation]");

            // Optionally, background subtract the stack before saving.
            if (BGSUB) { 
              run("Subtract Background...", "rolling=" + RB_BALL_RADIUS + " stack");
            }
            save(OUTDIR + "\\stacked\\" + CHANNELS[ch] + "\\" + well + ".tif");
            }
          }
        }
      }
      run("Close All");
  }
}

// Stacking related functions. 
function stack(tp_to_img_paths) {
  stackPrimaryChannel(tp_to_img_paths);
  stackOtherChannels(tp_to_img_paths);
}


function process() {
  // Acquire needed paths.
  tp_to_img_paths = acquirePaths();
  // Ensure all output paths exist.
  mkStitchedOutputDirectories(tp_to_img_paths);
  // Stitch images and output.
  stitch(tp_to_img_paths);

  if (TPNUM > 1) { 
    // Ensure all output paths exist.
    mkStackedOutputDirectories(tp_to_img_paths);
    // Stack images and output.
    stack(tp_to_img_paths);
  }
}

process();