/*****************************************************************************
 -                                                                           -
 -                    POLONY SEQUENCER ACQUISITION SUITE                     -
 -                                Church Lab                                 -
 -                          Harvard Medical School                           -
 -                                                                           -
 -    Free software for running a Polony Sequencing automated microscope     -
 -                                                                           -
 - ========================================================================= -
 -                                                                           -
 -                                                                           -
 - autofocus.cpp                                                             -
 -                                                                           -
 -   Program to autofocus brightfield or fluorescence images of Polony beads -
 -                                                                           -
 -   Written by Greg Porreca, 02-15-2006                                     -
 -                                                                           -
 - Revised                                                                   -
 -                                                                           -
 -                                                                           -
 - This software may be used, modified, and distributed freely, but this     -
 - header may not be modified and must appear at the top of this file.       -
 -                                                                           -
 -                                                                           -
 *****************************************************************************/


#include "Autofocus.h"

#define _BRIGHTFIELD_EXPOSURETIME 0.02
#define _BRIGHTFIELD_SENSITIVITY 0
#define _BRIGHTFIELD_INTENS_THRESH 11000


Autofocus::Autofocus(Reporter* rep, CHamCam* cHamCam, CFilter* cFilter, CScope* cScope, CPifoc* cPifoc, Acquisition* acquis){
	r = rep;
	this_camera = cHamCam;
	this_filter = cFilter;
	this_te2000 = cScope;
	this_pifoc = cPifoc;
	acq = acquis;
	focus_axis = 0;
	exposure = 0;
	sensitivity = 0;
}


Autofocus::~Autofocus(){;}

void Autofocus::setFocusAxis(int f_axis){
	focus_axis = f_axis;
}

double Autofocus::autofocus_fl(int num_steps_coarse, double coarse_interval, double exposure, double sensitivity, int ex, int em){
	int i;
	double initial_zpos, start_focus_pos, curr_focus_pos;
	short unsigned int **image_ptr_stack;
	double max_complexity=0, max_complexity_zpos, curr_complexity;
	double *focus_pos;

	sprintf(log_string, "executing autofocus_fl with %d steps spaced %f microns apart (total travel of %.2f)", num_steps_coarse, coarse_interval, (num_steps_coarse - 1)*coarse_interval);
	r->log(log_string,1);

	if(!this_filter->SetPosition(2, ex, em)){
		r->error("set filter positions failed", 1);
	}


	if((image_ptr_stack = (short unsigned int**)malloc(num_steps_coarse * sizeof(short unsigned int*)))==NULL){
		r->error("allocate memory for autofocus pointer stack failed", 1);
	}
	
	if((focus_pos = (double*) malloc( (num_steps_coarse + (int) ((coarse_interval / 0.05)+1) ) * sizeof(double) ))==NULL){
		r->error("allocate memory for focus position list failed", 1);
	}

	if(focus_axis==1){//PIFOC USED
		initial_zpos = this_pifoc->GetPos();
	}
	else{
		initial_zpos = this_te2000->GetPos();
	}

	start_focus_pos = initial_zpos - (((num_steps_coarse - 1)/2)*coarse_interval);
	max_complexity_zpos = initial_zpos;

	sprintf(log_string, "----Initial autofocus plane: %f", initial_zpos);
	r->log(log_string, 3);

	if(!this_camera->QueueCapture()){
		r->error("queue capture of multiple frames failed", 0);
		return initial_zpos;
	}


	for(i=0; i< num_steps_coarse; i++){
		
		curr_focus_pos = start_focus_pos + (i*coarse_interval);
		sprintf(log_string, "---- %d %f", i, curr_focus_pos);
		r->log(log_string, 3);

		if(focus_axis==1){
			this_pifoc->SetPosition(curr_focus_pos);
			while(!(this_pifoc->ConfirmPosition())){;}
			*(focus_pos + i) = this_pifoc->GetPos();
		}
		else{
			this_te2000->SetPosition(curr_focus_pos);
			while(!(this_te2000->ConfirmPosition())){;}
			*(focus_pos + i) = this_te2000->GetPos();
		}

		r->log("----Acquire start", 5);
		if(!acq->acquire_expose_shutter(exposure, sensitivity)){
			sprintf(log_string, "Error acquiring image %d in fluorescence autofocus stack", i);
			r->error(log_string, 0);
		}
		else{
			image_ptr_stack[i] = acq->acquire_getdata(i);  
		}
		r->log("----Acquire end", 5);
	}
	
	max_complexity = 0;
	for(i=0; i<num_steps_coarse; i++){
		curr_complexity = compute_complexity(image_ptr_stack[i]);

		sprintf(log_string, "---- Contrast for position %02d (z plane %f) = %f\n", i, *(focus_pos + i), curr_complexity);
		r->log(log_string, 3);

		if(curr_complexity > max_complexity){
			sprintf(log_string, "CURR_CONTRAST %f > MAX_CONTRAST %f", curr_complexity, max_complexity);
			r->log(log_string, 4);
			max_complexity = curr_complexity;
			max_complexity_zpos = *(focus_pos + i);
		}
	}
	sprintf(log_string, "----Best position: %f, score: %f", max_complexity_zpos, max_complexity); 
	r->log(log_string, 3);

//WE'VE FOUND THE RIGHT POSITION IN THE STACK, NOW MOVE THERE
	if(focus_axis==1){
		this_pifoc->SetPosition(max_complexity_zpos);
		while(!(this_pifoc->ConfirmPosition())){;}
	}
	else{
		sprintf(log_string, "Setting focus to position %f (started at %f)", max_complexity_zpos, initial_zpos);
		r->log(log_string, 2);
		this_te2000->SetPosition(max_complexity_zpos);
		while(!(this_te2000->ConfirmPosition())){;}
	}

	free(image_ptr_stack);
	free(focus_pos);
	return max_complexity_zpos;
}



double Autofocus::compute_complexity(unsigned short int *curr_img){
	double delta_sum, delta_count, delta_mean;
	int i;

	delta_count = 0;
	delta_sum = 0;
	for(i=0; i<999999; i++){
		if(((i+1)%1000)){
			if(((i)%1000)){
				delta_sum = delta_sum + (double)abs( (*(curr_img + i) - *(curr_img + i + 1)));
				delta_count++;
			}
		}
	}

	if(delta_count > 0){
		delta_mean = delta_sum / delta_count;
	}
	else{
		r->error("--pixel_count==0", 0);
		delta_mean = 0;
	}

	sprintf(log_string, "--Contrast = %.0f", delta_mean);
	r->log(log_string, 2);

	return delta_mean;
}





double Autofocus::autofocus_brightfield(int num_steps_coarse, double coarse_interval){
	int i;
	double initial_zpos, start_focus_pos, curr_focus_pos;
	short unsigned int **image_ptr_stack;
	double max_complexity=0, max_complexity_zpos, curr_complexity;
	double *focus_pos;

	sprintf(log_string, "executing autofocus_brightfield with %d steps spaced %.3f microns apart (total travel of %.2f microns)", num_steps_coarse, coarse_interval, (num_steps_coarse - 1)*coarse_interval);
	r->log(log_string,1);

	if((image_ptr_stack = (short unsigned int**)malloc(num_steps_coarse * sizeof(short unsigned int*)))==NULL){
		r->error("allocate memory for autofocus pointer stack failed", 1);
	}
	
	if((focus_pos = (double*) malloc( (num_steps_coarse + (int) ((coarse_interval / 0.05)+1) ) * sizeof(double) ))==NULL){
		r->error("allocate memory for focus position list failed", 1);
	}

	if(focus_axis==1){//PIFOC USED
		initial_zpos = this_pifoc->GetPos();
	}
	else{
		initial_zpos = this_te2000->GetPos();
	}

	start_focus_pos = initial_zpos - (((num_steps_coarse - 1)/2)*coarse_interval);
	max_complexity_zpos = initial_zpos;

	sprintf(log_string, "----Initial autofocus plane: %f", initial_zpos);
	r->log(log_string, 3);

	if(!this_camera->QueueCapture()){
		r->error("queue capture of multiple frames failed", 0);
		return initial_zpos;
	}


	for(i=0; i< num_steps_coarse; i++){
		
		curr_focus_pos = start_focus_pos + (i*coarse_interval);
		sprintf(log_string, "---- %d %f", i, curr_focus_pos);
		r->log(log_string, 3);

		if(focus_axis==1){
			this_pifoc->SetPosition(curr_focus_pos);
			while(!(this_pifoc->ConfirmPosition())){;}
			*(focus_pos + i) = this_pifoc->GetPos();
		}
		else{
			this_te2000->SetPosition(curr_focus_pos);
			while(!(this_te2000->ConfirmPosition())){;}
			*(focus_pos + i) = this_te2000->GetPos();
		}

		r->log("----Acquire start", 5);
		if(!acq->acquire_expose_noshutter(_BRIGHTFIELD_EXPOSURETIME, _BRIGHTFIELD_SENSITIVITY)){
			sprintf(log_string, "Error acquiring image %d in brightfield autofocus stack", i);
			r->error(log_string, 0);
		}
		else{
			image_ptr_stack[i] = acq->acquire_getdata(i);  
		}
		r->log("----Acquire end", 5);
	}
	
	max_complexity = 0;
	for(i=0; i<num_steps_coarse; i++){
		curr_complexity = compute_complexity_autofocus(image_ptr_stack[i]);

		sprintf(log_string, "---- Contrast for position %02d (z plane %f) = %f\n", i, *(focus_pos + i), curr_complexity);
		r->log(log_string, 3);

		if(curr_complexity > max_complexity){
			sprintf(log_string, "CURR_CONTRAST %f > MAX_CONTRAST %f", curr_complexity, max_complexity);
			r->log(log_string, 4);
			max_complexity = curr_complexity;
			max_complexity_zpos = *(focus_pos + i);
		}
	}
	sprintf(log_string, "----Best position: %f, score: %f", max_complexity_zpos, max_complexity); 
	r->log(log_string, 3);

//WE'VE FOUND THE RIGHT POSITION IN THE STACK, NOW MOVE THERE
	if(focus_axis==1){
		this_pifoc->SetPosition(max_complexity_zpos);
		while(!(this_pifoc->ConfirmPosition())){;}
	}
	else{
		sprintf(log_string, "Setting focus to position %f (started at %f)", max_complexity_zpos, initial_zpos);
		r->log(log_string, 2);
		this_te2000->SetPosition(max_complexity_zpos);
		while(!(this_te2000->ConfirmPosition())){;}
	}

	free(image_ptr_stack);
	free(focus_pos);
	return max_complexity_zpos;
}



double Autofocus::compute_complexity_autofocus(short unsigned int* curr_img){
	double above_thresh = 0;
	double thresh = _BRIGHTFIELD_INTENS_THRESH;
	int i;
	for(i=0; i<1000000; i++){
		if((double)*(curr_img + i) > thresh) above_thresh++;
	}
	
	sprintf(log_string, "Pixels above threshold: %f",above_thresh);
	r->log(log_string, 3);

	return above_thresh;
}