/***************************************************************************************
 * 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 <limits.h>
#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 100000     // maximum number of focus pixels allowed
#define MAX_NUM_SETS 100 //maximum number of image sets to align

#define DEBUG 1

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 *);
void quickSort(short unsigned int*, int);


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, *BIP;
  int curr_y_offset, curr_x_offset;
  int curr_pixel;
  int num_test_pixels, test_pixel_index, test_pixel_thresh, brightfield_pixel_thresh;
  short unsigned int *short_pixel_list;


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

  int n;

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

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

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

  ROW_SIZE = yrows;
  COL_SIZE = xcols;

  max_total_offset = (MAX_OFF*2)-1;


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

  if((align_image = (char*) malloc(MAX_FILESIZE)) == NULL){
    fprintf(stderr, "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+6]);

    //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(stderr, "%s\n", align_fn);

    //open file
    fprintf(stderr, "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;

  if( (short_pixel_list=(short unsigned int*)malloc(10000*sizeof(short unsigned int)))==NULL){
    fprintf(stderr, "ERROR allocating memory for pixel list.  Exiting\n");
    exit(42);
  }
  

  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);
    BIP = (short unsigned int*) base_image;

    //determine the background threshold value
    for(i=0; i<xcols*yrows; i++){
      if((i%1000)==0){
	*(short_pixel_list + (i/1000)) = *(BIP+i);
      }
    }
    quickSort(short_pixel_list, 1000);

    brightfield_pixel_thresh = *(short_pixel_list + 500);
    fprintf(stdout, "Brightfield pixel threshold = %d\n", brightfield_pixel_thresh);

    //threshold brightfield image
    for(i=0; i<xcols*yrows; i++){
      if(*(BIP+i) < brightfield_pixel_thresh){
	*(BIP+i) = 0;
      }
    }
        
    
    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 min pixel value to consider a pixel a test pixel
      for(i=0; i<xcols*yrows; i++){
	if((i%1000)==0){//only take 1 out of every thousand pixels
	  *(short_pixel_list + (i/1000)) = *(IP+i);
	}
      }
      quickSort(short_pixel_list, 1000);
      test_pixel_thresh = *(short_pixel_list + 900); //take the 90th percentile val
      fprintf(stdout, "Test pixel threshold = %d, ", test_pixel_thresh);
      

      //fill test pixel array for current fluorescence image
      test_pixel_index = 0;
      for(i=0; i<xcols*yrows; i++){
	if(*(IP+i) > test_pixel_thresh){
	  FOCUS_PIXEL_X[test_pixel_index] = ((int) i % yrows) - MAX_OFF;
	  FOCUS_PIXEL_Y[test_pixel_index] = ((int) floorl(((float)i) / ((float)yrows))) - MAX_OFF;
	  test_pixel_index++;
	}
      }
      num_test_pixels = test_pixel_index;
      fprintf(stdout, " num_test_pixels==%d\n", num_test_pixels);
      

      // Determine scores for all possible offsets 
      BEST_SCORE = LONG_LONG_MAX-1;
  
      //iterate over all possible offsets in X and Y
      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;
	  
	  //at each X,Y offset, iterate over all 'focus pixels' and sum their vals in the
	  //'shifted' fluorescence image
	  for (curr_focus_pix = 0; curr_focus_pix < num_test_pixels; curr_focus_pix++) {
	    //count7 is the Y coordinate
	    count7 = FOCUS_PIXEL_Y[curr_focus_pix] + curr_y_offset;
	    if (count7 > 0) {
	      //INDEX is the position in the brightfield image of the current X, Y
	      //'test pixel' location plus current offset
	      INDEX = (PRE_MULT[count7]) + FOCUS_PIXEL_X[curr_focus_pix] + curr_x_offset;
	      //protect against offsets where corner 'focus pixels' fall outside the image
	      if (INDEX >= 0 && INDEX < TOT_SIZE) {
		count6+= (long long int) BIP[INDEX];
		count5++;
	      }
	    }
	  }
	  
	  //score is ratio of fluorescence image 'focus pixel' value sum to number of 'focus pixels'
	  //SCORE = (long long int) count6/count5;

#ifdef DEBUG
	  //	  fprintf(stdout, "%d\t%d\t%d\t%Lf\t%lld\n", current_frame, curr_y_offset, curr_x_offset, (long double)count6, BEST_SCORE);
#endif

	  if (count6 < BEST_SCORE) {
	    BEST_Y = curr_y_offset - MAX_OFF;
	    BEST_X = curr_x_offset - MAX_OFF;
	    BEST_SCORE = count6;
	  }
	}
      }
      
      // Report optimal offset
      if((BEST_X == MAX_OFF)||(BEST_X == -1*MAX_OFF)||(BEST_Y == MAX_OFF)||(BEST_Y == -1*MAX_OFF)){
	fprintf(stdout, "ERROR, frame %d cycle %d: offset(s) reported is the maximum specified.\n", current_frame, curr_set);
	BEST_X = 0;
	BEST_Y = 0;
      }
      fprintf(*(outfile + curr_set), "%s\t%d\t%d\t%lld\n", curr_fn_index, -1*BEST_X, -1*BEST_Y, BEST_SCORE);
    }//end for(current_set)
  }//end for(current_frame)
  for(i=0; i<num_sets; i++){
    fclose(*(outfile + i));
  }
  free(short_pixel_list);
  free(BIP);
  free(IP);
}//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(stderr, "Opening file %s to load image into memory.\n", FILENAME);
  
  // Open the TIFF image
  if((image = TIFFOpen(FILENAME, "r")) == NULL){
    fprintf(stderr, "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(stderr, "Read error on input strip number %d\n", stripCount);
      exit(42);
    }
    imageOffset += result;
  }
  // Close image
  TIFFClose(image);   
}



/* Quicksort implementation copied from:
http://www.personal.kent.edu/~rmuhamma/Algorithms/MyAlgorithms/Sorting/quickSort.htm
*/
void q_sort(short unsigned int numbers[], long long int left, long long int right){
  long long int pivot;
  long long int l_hold;
  long long int r_hold;

  l_hold = left;
  r_hold = right;
  pivot = numbers[left];
  while (left < right)
  {
    while ((numbers[right] >= pivot) && (left < right))
      right--;
    if (left != right)
    {
      numbers[left] = numbers[right];
      left++;
    }
    while ((numbers[left] <= pivot) && (left < right))
      left++;
    if (left != right)
    {
      numbers[right] = numbers[left];
      right--;
    }
  }
  numbers[left] = (short unsigned int) pivot;
  pivot = left;
  left = l_hold;
  right = r_hold;
  if (left < pivot)
    q_sort(numbers, left, pivot-1);
  if (right > pivot)
    q_sort(numbers, pivot+1, right);
}

void quickSort(short unsigned int numbers[], int array_size)
{
  q_sort(numbers, 0, (long long int) (array_size) - 1);
}