/***************************************************************************************
 * register.c                                                                          *
 *                                                                                     *
 * Takes a raw transmitted light image, and N fluorescence images, and returns the     *
 * offsets (x,y) to translate the fluorescent images by to bring them into register    *
 * with the transmitted light image.  Offsets are reported in N files, 1 per image.    *
 *                                                                                     *
 * TO COMPILE: g++ register.c -O3 -ffloat-store -foptimize-sibling-calls               * 
 *             -fno-branch-count-reg -o register -ltiff -lm -I../../INCLUDE            * 
 *             -L../../LIB                                                             *
 *                                                                                     *
 * USAGE:                                                                              *
 *       ./register xcols yrows base_dir/ baseimg_id focus_pixels cutoff min_max       *
 *       max_off set_1 set_2 ... set_n                                                 *
 *                                                                                     *
 *       WHERE:                                                                        *
 *             xcols = number of colums in the images (e.g. 1000)                      *
 *             yrows = number of rows in the images (e.g. 1000)                        *
 *             base_dir/ = the base dir (e.g. /mnt/data_seq1/B6/)                      *
 *             baseimg = the cycle number of the base images (usually 999)             *
 *             focus_pixels = 2000 max number of focus pixels to use for registration  *
 *             max_off = 40                                                            *
 *             min_max 1                                                               *
 *             set_1                                                                   *
 *             set_2                                                                   *
 *             .                                                                       *
 *             .                                                                       *
 *             .                                                                       *
 *             set_N                                                                   *
 *                                                                                     *
 * WRITTEN BY: Jay Shendure, Greg Porreca (Church Lab) 01-30-2005                      *
 *                                                                                     *
 *                                                                                     *
 ***************************************************************************************/

#include <stdio.h>
#include "tiffio.h"
#include <string.h>
#include <stdlib.h>
#include <math.h>
#include <dirent.h>

#include <map>
#include <iostream>
#include <string>
#include <algorithm>
#include <iterator>
#include <fstream>
#include <fnmatch.h>

using namespace std;

#define MAX_ROOTPATH_LEN 100
#define MAX_FULLPATH_LEN 150
#define PATH_FULL_RAW "IMAGES/FULL/RAW/"
#define PATH_FULL_OBJECT "IMAGES/FULL/OBJECT/"
#define PATH_RAW "IMAGES/RAW/"
#define PATH_OBJECT "IMAGES/OBJECT/"
#define PATH_ALIGN "PROC/ALIGN/"
#define MAX_FILESIZE 2252800 //2200KB
#define MAX_OBJECTS 64000 //max objects per frame
#define MAX_FOCUS_PIXELS 20000     // maximum number of focus pixels allowed
#define MAX_NUM_SETS 100 //maximum number of image sets to align

unsigned short int yrows, xcols;
char ROOT_PATH[MAX_ROOTPATH_LEN], CYCLEID_BASE[5];
int STRIPSIZE = 8000;
int STRIPMAX = 25;  


//function prototypes
void load_image(char*, char*);
int scandir_sel(struct dirent *);


int main(int argc, char *argv[]){
 unsigned short int current_frame;
  char raw_image_dir[MAX_FULLPATH_LEN];
  char base_fn[MAX_FULLPATH_LEN];
  char curr_fn_index[5];
  int num_frames = 0;
  int num_sets = 0;
  struct dirent **namelist;
  int i;
  char* base_image;
  char* align_image;
  char align_fn[MAX_FULLPATH_LEN];;
  int set_index[MAX_NUM_SETS];
  int curr_set;
  char curr_set_index[4];
  FILE *outfile[MAX_NUM_SETS];
  short unsigned int* IP, *IP2;
  int curr_y_offset, curr_x_offset;
  int curr_pixel;


  int FOCUS_N, CUTOFF, MIN_MAX, MAX_OFF, FRAME_N;
  int INDEX;  
  int count, total_focus_pix, count3, curr_focus_pix, count5, count6, count7, max_total_offset;
  int FOCUS_PIXEL_X[MAX_FOCUS_PIXELS];
  int FOCUS_PIXEL_Y[MAX_FOCUS_PIXELS];
  int SCORE, BEST_X, BEST_Y, BEST_SCORE;
  int PRE_MULT[20000];
  int TOT_SIZE;
  unsigned short int ROW_SIZE, COL_SIZE;

  int n;

  map<string, int>:: iterator p;

  // Set warning handler for libtiff to void
  TIFFSetWarningHandler(0);

  if(argc < 10){
    fprintf(stdout, "\n");
    fprintf(stdout, "To use this program, execute as follows:\n");
    fprintf(stdout, " ./register xcols yrows base/dir/ 999 2000 threshold bead_color max_offset cyclenums\n");
    fprintf(stdout, "  WHERE:\n");
    fprintf(stdout, "        xcols = number of columns (in pixels) in image; 1000 with our camera\n");
    fprintf(stdout, "        yrows = number of rows (in pixels) in image; 1000 with our camera\n");
    fprintf(stdout, "        base/dir/ = full path to data directory; if the data directory is \"E2\", and the directory\n");
    fprintf(stdout, "                    path is /mnt/data1/, /base/dir/ would be \"/mnt/data1/E2/\"\n");
    fprintf(stdout, "        threshold = brightfield threshold used for find_object\n");
    fprintf(stdout, "        bead_color = 1 or -1 (as specified for find_object\n");
    fprintf(stdout, "        max_offset = maximum offset (in pixels) to translate image by to put it in register with the\n");
    fprintf(stdout, "                     base image; the larger this value, the longer it will take; we use 40\n");
    fprintf(stdout, "        cyclenums = list of the number of each cycle to process; for example, if you want to process\n");
    fprintf(stdout, "                    the images from directories 001, 002, 003, 004, and 005, specify \"1 2 3 4 5\"\n");
    fprintf(stdout, "\n");
    exit(0);
  }

  // Parse and report on command-line arguments
  yrows = atoi(argv[1]);
  xcols = atoi(argv[2]);
  strcpy(ROOT_PATH, argv[3]);
  strcpy(CYCLEID_BASE, argv[4]);
  FOCUS_N  = atoi(argv[5]); //2000
  CUTOFF   = atoi(argv[6]); //7000
  MIN_MAX  = atoi(argv[7]); // -1
  MAX_OFF  = atoi(argv[8]); // 40 
  num_sets = argc - 9;

  ROW_SIZE = yrows;
  COL_SIZE = xcols;

  max_total_offset = (MAX_OFF*2)-1;


  if((base_image = (char*) malloc(MAX_FILESIZE)) == NULL){
    fprintf(stdout, "Could not allocate enough memory for the image.\n");
    exit(42);
  }

  if((align_image = (char*) malloc(MAX_FILESIZE)) == NULL){
    fprintf(stdout, "Could not allocate enough memory for the image.\n");
    exit(42);
  }


  for(curr_set=0; curr_set < num_sets; curr_set++){

    set_index[curr_set] = atoi(argv[curr_set+9]);

    //generate complete filenames for align output files
    sprintf(curr_set_index, "%03d", *(set_index + curr_set));
    strcpy(align_fn, "");
    strcat(align_fn, ROOT_PATH);
    strcat(align_fn, PATH_ALIGN);
    strcat(align_fn, curr_set_index);
    strcat(align_fn, ".align");
    fprintf(stdout, "%s\n", align_fn);

    //open file
    fprintf(stdout, "Opening file %s for output.\n", align_fn);
    *(outfile + curr_set) = fopen(align_fn, "w");
    setbuf(*(outfile + curr_set), NULL);
  }

  if (MIN_MAX >= 0) {MIN_MAX = 1;} else {MIN_MAX = -1;}
  
  //determine number of RAW files to process
  sprintf(raw_image_dir, "%s%s%s/", ROOT_PATH, PATH_RAW, CYCLEID_BASE);
  num_frames = scandir(raw_image_dir, &namelist, (int (*)(const struct dirent *))(scandir_sel), alphasort);
  free(namelist);
  
  // Precalculate multiples of the number of columns
  for (count = 0; count < ROW_SIZE; count++) {
    PRE_MULT[count] = COL_SIZE * count;
  }
  TOT_SIZE = ROW_SIZE * COL_SIZE;
  

  for(current_frame = 1; current_frame < num_frames + 1; current_frame++){
    
    //current image index, as a 4-digit string
    sprintf(curr_fn_index, "%04d", current_frame);
    
    //generate complete filename for current raw image
    strcpy(base_fn, "");
    strcat(base_fn, ROOT_PATH);
    strcat(base_fn, PATH_RAW);
    strcat(base_fn, CYCLEID_BASE);
    strcat(base_fn, "/WL_");
    strcat(base_fn, curr_fn_index);
    strcat(base_fn, ".tif");
    
    //load base image for current frame
    load_image(base_fn, base_image);
    IP = (short unsigned int*) base_image;
    
    
    // Select focus pixels
    total_focus_pix = 0;
    //iterate over all pixels in the raw image
    for(curr_pixel = 0; curr_pixel < TOT_SIZE; curr_pixel++){
      if ((MIN_MAX*IP[curr_pixel]) > (MIN_MAX*CUTOFF)) {
	total_focus_pix++;
      }
    }
    count3 = (int) floorl(((float) total_focus_pix) / ((float) FOCUS_N));
    if(count3==0){
      fprintf(stdout, "count3 (floorl(total_focus_pix / FOCUS_N))==0; setting count3=1 frame %d\n", current_frame);
      count3 = 1;
    }
    total_focus_pix = 0;
    curr_focus_pix = 0;
    
    for(curr_pixel = 0; curr_pixel < TOT_SIZE; curr_pixel++){
      if ((MIN_MAX*IP[curr_pixel]) > (MIN_MAX*CUTOFF)) {
	total_focus_pix++;
	if ((total_focus_pix % count3) == 0 && curr_focus_pix < FOCUS_N) {
	  FOCUS_PIXEL_X[curr_focus_pix] = ((int) curr_pixel % COL_SIZE) - MAX_OFF;
	  FOCUS_PIXEL_Y[curr_focus_pix] = ((int) floorl(((float)curr_pixel) / ((float)COL_SIZE))) - MAX_OFF;
	  curr_focus_pix++;	    
	}
      }
    }
    
    for(curr_set = 0; curr_set < num_sets; curr_set++){
      
      //generate complete filename for current align image file
      sprintf(curr_set_index, "%03d", *(set_index + curr_set));
      strcpy(align_fn, "");
      strcat(align_fn, ROOT_PATH);
      strcat(align_fn, PATH_RAW);
      strcat(align_fn, curr_set_index);
      strcat(align_fn, "/SC_");
      strcat(align_fn, curr_fn_index);
      strcat(align_fn, ".tif");
      
      //load align image for current frame
      load_image(align_fn, align_image);
      IP = (short unsigned int*) align_image;
      
      // Determine scores for all possible offsets
      BEST_SCORE = 0;
  
      for (curr_y_offset = 0; curr_y_offset < max_total_offset; curr_y_offset++) {	
	for (curr_x_offset = 0; curr_x_offset < max_total_offset; curr_x_offset++) {
	  
	  count5 = 0;
	  count6 = 0;
	  
	  for (curr_focus_pix = 0; curr_focus_pix < FOCUS_N; curr_focus_pix++) {
	    count7 = FOCUS_PIXEL_Y[curr_focus_pix] + curr_y_offset;
	    if (count7 > 0) {
	      INDEX = (PRE_MULT[count7]) + FOCUS_PIXEL_X[curr_focus_pix] + curr_x_offset;
	      
	      if (INDEX >= 0 && INDEX < TOT_SIZE) {
		count6+= IP[INDEX];
		count5++;
	      }
	    }
	  }
	  
	  SCORE = (int) count6/count5;
	  if (SCORE > BEST_SCORE) {
	    BEST_Y = curr_y_offset - MAX_OFF;
	    BEST_X = curr_x_offset - MAX_OFF;
	    BEST_SCORE = SCORE;
	  }
	}
      }
      
      // Report optimal offset
      fprintf(*(outfile + curr_set), "%s\t%d\t%d\t%d\n", curr_fn_index, BEST_X, BEST_Y, BEST_SCORE);
    }//end for(current_set)
  }//end for(current_frame)
  for(i=0; i<num_sets; i++){
    fclose(*(outfile + i));
  }
}//end int(main)

//make scandir only return files matching pattern FILEMASK
int scandir_sel(struct dirent * ent){
  char* FILEMASK = "*.tif";
  if(1)
    return(!fnmatch(FILEMASK, ent->d_name, 0));
  return(0);
}


void load_image(char FILENAME[MAX_FULLPATH_LEN], char* img){
  TIFF *image;
  tsize_t stripSize;
  unsigned long imageOffset, result;
  int stripMax, stripCount;
  
  fprintf(stdout, "Opening file %s to load image into memory.\n", FILENAME);
  
  // Open the TIFF image
  if((image = TIFFOpen(FILENAME, "r")) == NULL){
    fprintf(stdout, "Could not open image %s\n", FILENAME);
    exit(42);
  }
  
  stripSize = TIFFStripSize (image);
  stripMax = TIFFNumberOfStrips (image);
  STRIPSIZE = stripSize;
  STRIPMAX = stripMax;
  imageOffset = 0;
  
  // Load data to buffer
  for (stripCount = 0; stripCount < stripMax; stripCount++){
    if((result = TIFFReadEncodedStrip (image, stripCount,
				       img + imageOffset,
				       stripSize)) == -1){
      fprintf(stdout, "Read error on input strip number %d\n", stripCount);
      exit(42);
    }
    imageOffset += result;
  }
  // Close image
  TIFFClose(image);   
}