// -*- mode: C -*-

#include <stdio.h>
#include <stdlib.h>
#include <string.h>

#include <sys/stat.h>
#include <sys/types.h>

#include <errno.h>
#include <getopt.h>
#include <libgen.h>
#include <unistd.h>
#include <float.h>

#include "pgmio.h"
#include "segmentation.h"

#define PATH_MAX 1024
#define OPTSTR "i:s:o:b:k:m:g:h"
#define USAGE_FMT "%s [-i inputdir] [-s seedsdir] [-o outputdir] [-b beta] [-k kappa] [-m method] [-g gtdir] [-h]"
#define DEFAULT_PROGNAME "???"

extern int errno;
extern char *optarg;
extern int opterr, optind;

enum {DO_OPTIMIZE,DO_PAIRED,DO_BATCH};

typedef struct {
	char			 image_dir[PATH_MAX];
	char			 seeds_dir[PATH_MAX];
	char			 gt_dir[PATH_MAX];
	char			 output_dir[PATH_MAX];
	float			 beta[128];
	int				 b_count;
	float			 kappa[128];
	int				 k_count;
	int				 method;
} options_t;


void usage(char *progname, int opt);
int process_images(options_t *options);
int normalize_image(float *img_data, int n);


int main(int argc, char *argv[]) {

	int opt;
	options_t *options = (options_t *) malloc(sizeof(options_t));

	options->image_dir[0] = '\0';
	options->seeds_dir[0] = '\0';
	options->gt_dir[0] = '\0';
	options->output_dir[0] = '\0';

	getcwd(options->output_dir,sizeof(options->output_dir));

	options->b_count = 0;
	options->k_count = 0;
	int n = -1;
	int m = -1;
	float b_max = -1;
	float k_max = -1;

	options->method = DO_PAIRED;

	while((opt = getopt(argc, argv, OPTSTR)) != EOF){
		switch(opt) {
		case 'i':
			strcpy(options->image_dir, optarg);
			break;
		case 's':
			strcpy(options->seeds_dir, optarg);
			break;
		case 'g':
			strcpy(options->gt_dir, optarg);
			break;
		case 'o':
			strcpy(options->output_dir, optarg);
			break;
		case 'b':
			options->beta[options->b_count] = strtof(optarg, NULL);
			options->b_count++;
			break;
		case 'B':
			b_max = strtof(optarg, NULL);
			break;
		case 'n':
			n = atoi(optarg);
			break;
		case 'k':
			options->kappa[options->k_count] = strtof(optarg, NULL);
			options->k_count++;
			break;
		case 'K':
			k_max = strtof(optarg, NULL);
			break;
		case 'm':
			m = atoi(optarg);
			break;
		case 'p':
			if(strcmp(optarg,"optimize") == 0){
				options->method = DO_OPTIMIZE;
			} else if(strcmp(optarg,"batch") == 0){
				options->method = DO_BATCH;
			} else if(strcmp(optarg, "paired") == 0){
				options->method = DO_PAIRED;
			} else {
				printf("Unknown processing method: %s\n", optarg);
				return 1;
			}
			break;
		case 'h':
		default:
			usage(basename(argv[0]), opt);
			break;
		}
	}
	/* check that core values have been set */
	if(strlen(options->image_dir) == 0){
		printf("input image directory required.\n");
		return 1;
	} else if(strlen(options->seeds_dir) == 0){
		printf("seeds directory required.\n");
		return 1;
	} else if(options->b_count == 0){
		printf("Please enter at least one beta value.\n");
		return 1;
	} else if(options->k_count == 0){
		printf("Please enter at least one kappa value.\n");
	}

	/* determine if range of beta values is to be used */
	if(n > 0 && b_max > 0.0){
		/* add range of values to the beta list */
		int b_n = options->b_count - 1;
		float b0 = options->beta[b_n];
		float b_step = (b_max - b0) / (n-1);
		for(int i = 1; i <= n; i++){
			options->beta[b_n + i] = b0 + i * b_step;
		}
		options->b_count = b_n + n;
	}

	/* determine if range of kappa values is to be used */
	if(m > 0 && k_max > -0.0001){
		/* add range of values to the kappa list */
		int k_n = options->k_count - 1;
		float k0 = options->kappa[k_n];
		float k_step = (k_max - k0) / (m-1);
		for(int i = 1; i <= m; i++){
			options->kappa[k_n + i] = k0 + i * k_step;
		}
		options->k_count = k_n + m;
	}

	if(options->method == DO_PAIRED){
		/* paired option only works with equal number of b and k values */
		if(options->b_count != options->k_count){
			printf("Paired value processing must have an equal number of beta and kappa values.\n");
			return 1;
		}
	} else if(options->method == DO_OPTIMIZE){
		/* also need ground truth for optimization */
		if(strlen(options->gt_dir) == 0){
			printf("ground truth image directory required for optimization.\n");
			return 1;
		} else if(options->gt_dir[strlen(options->gt_dir)-1] != '/'){
			strcat(options->seeds_dir,"/");
		}
	}

	/* make sure all directories have leading '/' for using strcat to generate file paths */
	if(options->image_dir[strlen(options->image_dir)-1] != '/')
		strcat(options->image_dir,"/");
	if(options->seeds_dir[strlen(options->seeds_dir)-1] != '/')
		strcat(options->seeds_dir,"/");
	if(options->output_dir[strlen(options->output_dir)-1] != '/')
		strcat(options->output_dir,"/");

	printf("Image dir:\t%s\n",options->image_dir);
	printf("Seeds dir:\t%s\n",options->seeds_dir);
	printf("GT dir:\t%s\n",options->gt_dir);
	printf("Output dir:\t%s\n",options->output_dir);
	printf("Processing method number: %d\n",options->method);
	printf("%d beta values:", options->b_count);
	for(int i = 0; i < options->b_count; i++){
		printf("\t%.0f", options->beta[i]);
	}
	printf("\n");
	printf("%d kappa values:", options->k_count);
	for(int i = 0; i < options->k_count; i++){
		printf("\t%.3f", options->kappa[i]);
	}
	printf("\n");
	return process_images(options);
}

void usage(char *progname, int opt){
	fprintf(stderr,USAGE_FMT, progname?progname:DEFAULT_PROGNAME);
	exit(EXIT_FAILURE);
}

int process_images(options_t *options){
	// get image files to be processed§
	file_node_t *image_file_head = get_pgm_files(options->image_dir);
	if(image_file_head == NULL){
		printf("Could not find any image files.\n");
		return 1;
	}
	file_node_t *seeds_file_head = get_pgm_files(options->seeds_dir);
	if(seeds_file_head == NULL){
		printf("Could not find any seeds files.\n");
		return 1;
	}
	file_node_t *gt_file_head = NULL;
	if(options->method == DO_OPTIMIZE){
		gt_file_head = get_pgm_files(options->gt_dir);
		if(gt_file_head == NULL){
			printf("Could not find any ground truth files.\n");
			return 1;
		}
	}

	/* parameters */
	struct segmentation_parameters *params;
	struct image_dimensions *dims0;
	struct image_dimensions *dims1;

	/* processing method parameters */
	int max_it = 20;

	/* parameters for CERW segmentation */
	params = (segmentation_parameters *)malloc(sizeof(segmentation_parameters));

	/* fixed parameters */
	int edge_direction_count = 9; // half the neighbourhood size
	params->T = 50000 + 1; // Maximum number of timesteps, we use relative error to detect when to break
	params->test_interval = 1000; // number of timesteps between testing relative error
	params->test_time_start = 0; // earliest timestep to test relative error
	params->curv_interval = 10; // number of timesteps between updating the curvature
	params->curv_time_start = 10010; //earliest timestep to compute curvature
	params->re_stop = 5e-5; // threshold for relative error at which the simulation stops
	params->Ddt = 0.9/(2.0 * edge_direction_count); // time step, stable for values of kappa < 18
	params->t_norm = 0.01; // threshold for the size of gradient when computing curvature
	params->grad_step = 3; // step size when computing gradients for curvature
	params->max_threads_per_block = 1024;
	/* parameters for future implementations */
	params->augment_seeds = 0;
	params->v_T = 1.0;

	/* image dimension parameters */
	dims0 = (image_dimensions *)malloc(sizeof(image_dimensions));
	dims1 = (image_dimensions *)malloc(sizeof(image_dimensions));

	/* image data arrays */
	float *img_data; // raw intensity data from the image file, used to weight edges
	unsigned char *seed_data; // foreground seeds, given as positive at seed locations and zero elsewhere
	unsigned char *img_char; // temporary array when data imported is in 8-bit format
	unsigned char *gt_data; // ground truth data, non-zero for foreground
	// cropped image arrays
	float *h_img; // raw intensity data from the image file, optimally cropped and padded to give a whole number of blocks in each direction
	unsigned char *h_seed_img; // has value 1 for FG, -1 for BG, and zero elsewhere
	unsigned char *h_gt; // ground truth image for comparison

	// loop over real image list
	file_node_t *image_file_node = image_file_head;
	while(image_file_node != NULL){
		printf("Processing image %s...\n", image_file_node->file_name);

		char img_path[PATH_MAX] = "";
		char seed_path[PATH_MAX] = "";
		char gt_path[PATH_MAX] = "";
		char seg_path[1024] = "";
		char op_path[1024] = "";

		// image processing requires naming such that each input file name is contained in
		// exactly one seeds file (and gt file if optimization is used).
		/* read images */
		// input image file
		int name_length = strlen(image_file_node->file_name);
		char img_name_base[256] = "";
		// remove ".pgm" extension and optional leading t
		strncpy(img_name_base, &image_file_node->file_name[1], name_length - 5);
		strcpy(img_path, image_file_node->file_path);

		// find corresponding seed file
		file_node_t *seeds_file_node = seeds_file_head;
		while(strstr(seeds_file_node->file_name, img_name_base) == NULL){
			seeds_file_node = seeds_file_node->next;
			if(seeds_file_node == NULL)
				break;
		}

		// check the seed file has been found
		if(seeds_file_node == NULL){
			printf("Error: could not find seeds file for %s\n", image_file_node->file_name);
			image_file_node = image_file_node->next;
			continue;
		}else {
			printf("Using seeds file: %s\n", seeds_file_node->file_name);
			strcat(strcpy(seed_path, options->seeds_dir), seeds_file_node->file_name);
		}


		if(options->method == DO_OPTIMIZE){
			char gt_name[256] = "";
			// find corresponding gt file
			file_node_t *gt_file_node = gt_file_head;
			while(strstr(gt_file_node->file_name,img_name_base) == NULL){
				gt_file_node = gt_file_node->next;
				if(gt_file_node == NULL)
					break;
			}

			if(gt_file_node == NULL){
				printf("Error: could not find gt file for %s\n", image_file_node->file_name);
				image_file_node = image_file_node->next;
				continue;
			} else {
				printf("Using gt file: %s\n", gt_file_node->file_name);
				strcat(strcpy(gt_path, options->gt_dir), gt_file_node->file_name);
			}
		}

		// optimization sequence measurements
		strcat(strcat(strcpy(op_path,options->output_dir),img_name_base),"_JACCARD.txt");
		// output directory for files (one per image)
		strcat(strcat(strcpy(seg_path,options->output_dir),img_name_base),"/");
		struct stat st = {0};
		if(stat(seg_path, &st) == -1){
			mkdir(seg_path, 0777);
		}

		// get input dimensions
		dims0->N = get_pgm_size(img_path);
		dims0->R = get_pgm_dimension(img_path, DIM_ROWS);
		dims0->C = get_pgm_dimension(img_path, DIM_COLUMNS);
		dims0->Z = get_pgm_dimension(img_path, DIM_SLICES);
		dims0->edge_direction_count = edge_direction_count;
		dims0->c0 = 0;
		dims0->r0 = 0;
		dims0->z0 = 0;

		img_data = (float *) malloc(sizeof(float) * dims0->N);

		// read input images
		printf("Loading data...\n");
		img_char = pgm_read(img_path);
		if(img_char == NULL){
			printf("Error: could not load image data.\n");
			image_file_node = image_file_node->next;
			continue;
		}
		printf("Loading seeds...\n");
		seed_data = pgm_read(seed_path);
		if(seed_data == NULL){
			printf("Error: could not load seed data.\n");
			image_file_node = image_file_node->next;
			continue;
		}
		if(options->method == DO_OPTIMIZE){
			printf("Loading ground truth...\n");
			gt_data = pgm_read(gt_path);
			if(gt_data == NULL){
				printf("Error: could not load ground truth data.\n");
				image_file_node = image_file_node->next;
				continue;
			}
		}
		// convert to float
		for(int i = 0; i < dims0->N; i++){
			img_data[i] = img_char[i];
		}

		// get optimized image dimensions
		set_image_dimensions(seed_data, dims0, dims1, params->grad_step);

		// allocate host memory
		h_img = (float *) malloc(dims1->N * sizeof(float));
		h_seed_img = (unsigned char *) malloc(dims1->N * sizeof(char));
		h_gt = (unsigned char *) malloc(dims1->N * sizeof(char));

		// copy input images to resized images
		h_img = (float *) copy_image(img_data, dims0, dims1, FLOAT_IM);
		h_seed_img = (unsigned char *) copy_image(seed_data, dims0, dims1, BYTE_IM);
		if(options->method == DO_OPTIMIZE){
			h_gt = (unsigned char *) copy_image(gt_data, dims0, dims1, BYTE_IM);
		}

		// normalize image intensities to [0,1]
		normalize_image(h_img, dims1->N);

		// all processing and writing of outputs performed by functions in batch_segment.c
		if(options->method == DO_OPTIMIZE){
			// optimization saves intermediates in resized volumes to save space
			// ground truth saved in this form so comparison can be made
			printf("Saving ground truth...\n");
			int out_dims[3] = {dims1->C, dims1->R, dims1->Z};
			char gt_out_path[1024] = "";
			strcat(strcpy(gt_out_path, seg_path), "GT.pgm");
			pgm_write(h_gt, gt_out_path, out_dims);

			printf("Optimizing...\n");
			optimize_parameters(h_img,h_seed_img,h_gt,dims1,params,seg_path,op_path,max_it);
		}
		if(options->method == DO_BATCH){
			printf("Segmenting range...\n");
			segment_range(h_img, h_seed_img, dims1, dims1, params, seg_path, img_name_base, options->beta, options->kappa, options->b_count, options->k_count);
		}
		if(options->method == DO_PAIRED){
			printf("Segmenting paired parameter values...\n");
			segment_paired(h_img, h_seed_img, dims0, dims1, params, seg_path, img_name_base, options->beta, options->kappa, options->b_count);
		}

		free(h_img);
		free(h_seed_img);
		free(seed_data);
		free(img_data);
		free(img_char);

		// move to next image file
		image_file_node = image_file_node->next;
	}

	free(params);
	free(dims0);
	free(dims1);

    return 0;
}

// Normalizes the values in img_data to lie within [0,1]
int normalize_image(float *img_data, int n){
	float max_img = 0;
	float min_img = FLT_MAX;
	for (int i = 0; i < n; i++){
	  	if(img_data[i] > max_img){
	  		max_img = img_data[i];
	  	}
	   	if(img_data[i] < min_img)
	   		min_img = img_data[i];
	}
	float scale = max_img - min_img;
	//Normalize to [0,1]
	for (int i = 0; i < n; i++){
	   	img_data[i] = (img_data[i] - min_img) / scale;
	}
	return 0;
}