/***************************************************************************************
 * find_objects.c                                                                      *
 *                                                                                     *
 * Takes a raw transmitted light image, finds all objects, and outputs a new image     *
 * file where every object from the original image is numbered sequentially (same      *
 * as bwlabel in Matlab) in the 'object' file.                                         *
 *                                                                                     *
 * TO COMPILE: g++ find_objects.c -O3 -ffloat-store -foptimize-sibling-calls           * 
 *             -fno-branch-count-reg -o find_objects -ltiff -lm -I../../INCLUDE        * 
 *             -L../../LIB                                                             *
 *                                                                                     *
 * USAGE:                                                                              *
 *       ./find_objects xcols yrows base_dir/ cutoff {-1,1}                            *
 *       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/)                      *
 *             cutoff = the threshold value (e.g. 9000)                                *
 *             {-1, 1} = whether bead val is 0 or 65535 (black beads = -1, white = 1)  *
 *                                                                                     *
 * OUTPUT goes to base_dir/IMAGES/OBJECT/obj_counts.dat and to base_dir/IMAGES/OBJECT/ *
 *                                                                                     *
 * 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 <iostream>
#include <multimap.h>
#include <map>
#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/RAW/999/"
#define PATH_FULL_OBJECT "IMAGES/OBJECT/"
#define PATH_RAW "IMAGES/RAW/"
#define PATH_OBJECT "IMAGES/OBJECT/"
#define MAX_FILESIZE 2252800 //2200KB
#define MAX_OBJECTS 64000 //max objects per frame
#define OBJ_COUNT_OUTFILENAME "IMAGES/OBJECT/obj_counts.dat"

//function prototypes
char* load_image(char*);
int save_image(char*, char*, char*);
int assign(int, int);
int scandir_sel(struct dirent * ent);

//global variables
unsigned short int yrows, xcols;
char ROOT_PATH[MAX_ROOTPATH_LEN];
multimap <int, int> m1;
map <int, int> m2;
map <int, int> m3;
map <int, int> m4;
int STRIPSIZE = 8000;
int STRIPMAX = 250;  
//tsize_t STRIPSIZE;
//int STRIPMAX;

//
int main(int argc, char *argv[]){
  
  char *raw_image, *object_image;
  
  //--------------pathname generation
  char raw_image_dir[MAX_FULLPATH_LEN];
  unsigned short int current_frame;
  char curr_raw_fn[MAX_FULLPATH_LEN];
  char curr_obj_fn[MAX_FULLPATH_LEN];
  char obj_count_fn[MAX_FULLPATH_LEN];
  char curr_fn_index[5];
  int num_frames = 0;
  struct dirent **namelist;
  int i;
  
  FILE *obj_count_outfile;
  
  int CUTOFF, MIN_MAX;
  
  int count, count2, count3, count4, count5;  
  
  int flag1, flag2, N1, N2;
  
  short unsigned int* IP;
  short unsigned int* IP2;

  if(argc != 6){
    fprintf(stderr, "\n");
    fprintf(stderr, "To use find_objects, execute as follows:\n");
    fprintf(stderr, " ./find_objects xcols yrows /base/dir/ threshold bead_color\n");
    fprintf(stderr, "  WHERE:\n");
    fprintf(stderr, "        xcols = number of colums (in pixels) in the image; will be 1000 for Hamamatsu 1k EM\n");
    fprintf(stderr, "        yrows = number of rows (in pixels) in the image; will be 1000 for the Hamamatsu 1k EM\n");
    fprintf(stderr, "        /base/dir = full path to data dir; if data directory is called \"E2\", and it is located\n");
    fprintf(stderr, "                    in /mnt/data_seq1, /base/dir would be \"/mnt/data_seq1/E2/\"\n");
    fprintf(stderr, "        threshold = the value to use when thresholding the brightfield images such that beads\n");
    fprintf(stderr, "                    pass the threshold but the background does not\n");
    fprintf(stderr, "        bead_color = 1 if the beads appear white in the brightfield images (as they will with\n");
    fprintf(stderr, "                     the phase plan apo we recommend).  Set this to -1 if the beads appear black\n");
    fprintf(stderr, "                     (as they will if you use a non-phase objective and Abbe condenser)\n");
    fprintf(stderr, "\n");
    exit(0);
  }

  yrows = atoi(argv[1]);
  xcols = atoi(argv[2]);
  strcpy(ROOT_PATH, argv[3]);
  CUTOFF   = atoi(argv[4]);
  MIN_MAX  = atoi(argv[5]);
  if (MIN_MAX >= 0) {MIN_MAX = 1;} else {MIN_MAX = -1;}
  
  
  //turn annoying TIFF warnings off
  TIFFSetWarningHandler(0);
  
  //determine number of RAW files to process
  sprintf(raw_image_dir, "%s%s", ROOT_PATH, PATH_FULL_RAW);
  num_frames = scandir(raw_image_dir, &namelist, (int (*)(const struct dirent *))(scandir_sel), alphasort);
  free(namelist);
  
  //open object count file
  strcpy(obj_count_fn, "");
  strcat(obj_count_fn, ROOT_PATH);
  strcat(obj_count_fn, OBJ_COUNT_OUTFILENAME);
  obj_count_outfile = fopen(obj_count_fn, "w");
  
  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(curr_raw_fn, "");
    strcat(curr_raw_fn, ROOT_PATH);
    strcat(curr_raw_fn, PATH_FULL_RAW);
    strcat(curr_raw_fn, "WL_");
    strcat(curr_raw_fn, curr_fn_index);
    strcat(curr_raw_fn, ".tif");
    
    //generate complete filename for current object image
    strcpy(curr_obj_fn, "");
    strcat(curr_obj_fn, ROOT_PATH);
    strcat(curr_obj_fn, PATH_FULL_OBJECT);
    strcat(curr_obj_fn, curr_fn_index);
    strcat(curr_obj_fn, ".tif");
    
    //allocate memory and load tiff
    raw_image = load_image(curr_raw_fn);
    
    //allocate memory for object image
    if((object_image = (char *) malloc(MAX_FILESIZE)) == NULL){
      fprintf(stderr, "Could not allocate enough memory for the image.\n");
      exit(42);
    }
    
    //------------------- PERFORM REGISTRATION ---------------------------
    //
    
    // Clear maps
    
    m1.clear();
    m2.clear();
    m3.clear();
    m4.clear();
    
    IP = (short unsigned int*) raw_image;
    
    if((IP2 = (short unsigned int*) malloc(MAX_FILESIZE)) == NULL){
      fprintf(stderr, "Could not allocate enough memory for IP2.");
      exit(42);
    }
    
    count3=0;
    count2=1;  // store label count 
    
    for(count = 0; count < (yrows * xcols); count++){
      
      IP2[count] = 0;
      
      if ((MIN_MAX*IP[count]) > (MIN_MAX*CUTOFF)) {	
	
	count3++;
	
	flag1 = 0; flag2 = 0;
	N1 = count - xcols;
	N2 = count - 1;
	if ((N1 >= 0) && (IP2[N1] > 0) && ((MIN_MAX*IP[N1]) > (MIN_MAX*CUTOFF))) {flag1 = 1;} 
	if ((count % xcols != 0) && (IP2[N2]) > 0 && ((MIN_MAX*IP[N2]) > (MIN_MAX*CUTOFF))) {flag2 = 1;}
	if (flag1 == 0 && flag2 == 0) {IP2[count] = count2; count2++;}  // untouched by labeled already-seen neighbors
	else if (flag1 == 1 && flag2 == 0) {IP2[count] = IP2[N1];}      // connected to pixel to top only
	else if (flag1 == 0 && flag2 == 1) {IP2[count] = IP2[N2];}      // connected to pixel to left only
	else {                                                          // connected to both neighbors
	  IP2[count] = IP2[N1];
	  if (IP2[N1] != IP2[N2]) {                                     // add entry (IP2[N1]=IP2[N2]) to equivalency table 
	    
	    m1.insert(pair<int, int>(IP2[N1],IP2[N2]));
	    m1.insert(pair<int, int>(IP2[N2],IP2[N1]));
	    m1.insert(pair<int, int>(IP2[N1],IP2[N1]));
	    m1.insert(pair<int, int>(IP2[N2],IP2[N2]));
	    
	  }
	}
	
	if (count2 > MAX_OBJECTS) {
	  fprintf(stderr, "Too many objects in an image to keep track...\n");
	  exit(42);
	}
      }
      
    }
    
    // consolidate equivalency list
    
    for (count4 = 1; count4 < count2; count4++) { 
      assign(count4, count4);	
    }
    
    // run through image again and reassign labels
    
    map<int,int>::iterator p;
    map<int,int>::iterator p2;
    
    count4 = 1;
    
    for (count5 = 1; count5 < count2; count5++) {
      if (m2.count(count5) == 0) {
	m2.insert(pair<int, int>(count5, -1));	
      }
      
      p = m2.find(count5);
      
      if (m3.count(p->second) > 0 && p->second != -1) {
	p2 = m3.find(p->second);
	m4.insert(pair<int,int>(count5,p2->second));
      } else {
	m4.insert(pair<int,int>(count5,count4));
	m3.insert(pair<int,int>(p->second,count4));
	count4++;
      }
      
    }
    
    count5 = 0;
    
    for (count4 = 1; count4 < count2; count4++) {
      p = m4.find(count4);
      if (p->second > count5) {count5 = p->second;} 
    }
    
    // Iterate over the bwlabel image, replacing values with object numbers specified by the m4 map
    for(count = 0; count < (yrows * xcols); count++){
      if (IP2[count]>0) {
	p = m4.find(IP2[count]);
	IP2[count] = p->second;	  
      }
    }    
    fprintf(stderr, "Frame %d analyzed.  %d pixels met threshold.  %d objects found.\n", current_frame, count3, count5);
    fprintf(obj_count_outfile, "%d\n%d\n", current_frame, count5); 
    count = (int) floor(float(current_frame)/34);
    count2 = current_frame % 34;
    
    
    //
    //--------------------- END REGISTRATION -----------------------------
    
    //save object image just created
    save_image(curr_raw_fn, curr_obj_fn, (char*)IP2);
    
    //free memory
    free(raw_image);
    free(object_image);
    free(IP2);
  }
  fclose(obj_count_outfile);
  exit(0);
}  


//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);
}


int assign(int A, int B) {
  
  multimap<int,int>::iterator iter;
  multimap<int,int>::iterator lower;
  multimap<int,int>::iterator upper;
  
  lower = m1.lower_bound(A);
  upper = m1.upper_bound(A);
  
  for (iter = lower; iter != upper; iter++) {
    
    if (m2.count(iter->second) == 0) {      
      m2.insert(pair<int, int>(iter->second,B));
      assign(iter->second,B);
    }
  }
  
  return 0;
}


int save_image(char *load_image_fn, char *save_image_fn, char *save_image_data){
  
  //TIFF *image1;
  TIFF *image2;
  
  tsize_t stripSize;
  unsigned long imageOffset, result;
  int stripMax, stripCount;
  unsigned long count;
  
  // Open the WL TIFF image to read characteristics
  //  if((image1 = TIFFOpen(load_image_fn, "r")) == NULL){
  //fprintf(stderr, "Could not open incoming image\n");
  //exit(42);
  //}
  
  // Open the output TIFF image
  if((image2 = TIFFOpen(save_image_fn, "w")) == NULL){
    fprintf(stderr, "Could not open outgoing image\n");
    exit(42);
  }
  
  // We need to set some values for basic tags before we can add any data
  // Hard-coding this for now, ideally it would be flexible to different image sizes
  TIFFSetField(image2, TIFFTAG_IMAGEWIDTH, xcols);
  TIFFSetField(image2, TIFFTAG_IMAGELENGTH, yrows);
  TIFFSetField(image2, TIFFTAG_BITSPERSAMPLE, 16);
  TIFFSetField(image2, TIFFTAG_ROWSPERSTRIP, 4);
  TIFFSetField(image2, TIFFTAG_PHOTOMETRIC, PHOTOMETRIC_MINISBLACK);
  TIFFSetField(image2, TIFFTAG_PLANARCONFIG, PLANARCONFIG_CONTIG);
  TIFFSetField(image2, TIFFTAG_XRESOLUTION, 72.0);
  TIFFSetField(image2, TIFFTAG_YRESOLUTION, 72.0);
  TIFFSetField(image2, TIFFTAG_RESOLUTIONUNIT, RESUNIT_INCH);
  
  // Read in the number and size of strips from the other file
  stripSize = STRIPSIZE;
  stripMax = STRIPMAX;
  //  stripSize = TIFFStripSize (image1);
  //stripMax = TIFFNumberOfStrips (image1);
  imageOffset = 0;
 
  // Write the data
  for (stripCount = 0; stripCount < stripMax; stripCount++){
    if((result = TIFFWriteEncodedStrip (image2, stripCount, save_image_data + imageOffset, stripSize)) == -1){
      fprintf(stderr, "Write error on input strip number %d\n", stripCount);
      exit(42);
    }
    //fprintf(stderr, "%d\t%d\t%d\t%d\n", stripCount, imageOffset, stripSize, (short unsigned int)*(save_image_data + imageOffset)); 
    imageOffset += result;
  }
  
  //TIFFClose(image1);
  TIFFClose(image2);
  
  return 0;
}


char* load_image(char FILENAME[MAX_FULLPATH_LEN]){
  TIFF *image;
  tsize_t stripSize;
  unsigned long imageOffset, result;
  int stripMax, stripCount;
  unsigned long count;
  char* buffer;
  
  if((buffer = (char*) malloc(MAX_FILESIZE)) == NULL){
    fprintf(stderr, "Could not allocate enough memory for the image.\n");
    exit(42);
  }
  
  // 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,
				       buffer + imageOffset,
				       stripSize)) == -1){
      fprintf(stderr, "Read error on input strip number %d\n", stripCount);
      exit(42);
    }
    
    imageOffset += result;
  }
  
  // Close image
  TIFFClose(image);   
  return buffer;
}