/*
 * batch_segment.c

 *
 *  Created on: Jan 21, 2020
 *      Author: lutton
 */

#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <sys/stat.h>
#include <dirent.h>
#include <math.h>
#include <sys/types.h>
#include <unistd.h>

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

float get_hausdorff(unsigned char *b1, unsigned char *b2, int t1, int t2, struct image_dimensions *dims);
float get_dist_above(unsigned char *b1, int t1, int p, float d_max, struct image_dimensions *dims);
float get_jaccard(float *h_v, unsigned char *h_gt, int N);
int compare_beta(const void* a, const void* b);

int optimize_parameters(
		float *h_img,
		unsigned char *h_seed_img,
		unsigned char *h_gt,
		struct image_dimensions *dims,
		struct segmentation_parameters *params,
		char output_dir[],
		char op_path[],
		int max_it
		){

	int N = dims->N;

	float *jacc; // Jaccard scores
	float *haus; // Hausdorff distances for segmentation optimization
	float *beta; // array of beta values used in optimization
	float *kappa; // array of kappa values used in optimization

	float *j_out; // Jaccard scores to be written
	float *h_out; // Hausdorff distances to be written
	float *b_out; // beta values to be written
	float *k_out; // kappa values to be written

	float *h_v0;
	float *h_seed_indicator;

	beta = (float *) malloc(max_it * sizeof(float));
	jacc = (float *) malloc(max_it * sizeof(float));
	haus = (float *) malloc(max_it * sizeof(float));
	kappa = (float *) malloc(max_it * sizeof(float));

	b_out = (float *) malloc(max_it * max_it * sizeof(float));
	k_out = (float *) malloc(max_it * max_it * sizeof(float));
	j_out = (float *) malloc(max_it * max_it * sizeof(float));
	h_out = (float *) malloc(max_it * max_it * sizeof(float));

	h_v0 = (float *) malloc(N * sizeof(float));
	h_seed_indicator = (float *) malloc(N * sizeof(float));

	params->kappa = 0.0;

	for(int i = 0; i < max_it * max_it; i++){
		b_out[i] = 0.0;
		k_out[i] = 0.0;
		j_out[i] = 0.0;
		h_out[i] = N;
	}
	for(int i = 0; i < max_it; i++){
		beta[i] = 0;
		jacc[i] = 0;
		haus[i] = N;
	}

	int out_dims[3] = {dims->C, dims->R, dims->Z};
	char gt_path[1024] = "";
	strcat(strcpy(gt_path, output_dir), "GT.pgm");
	pgm_write(h_gt, gt_path, out_dims);


	set_seed_indicator(h_seed_indicator, h_seed_img, N);
	set_v0(h_v0, h_seed_img, N);

	optimize_beta(
			h_img,h_seed_indicator,h_gt,h_v0,
			dims,params,output_dir,jacc,haus,beta,max_it);
	float j_max = 0;
	float b_max = 0;
	for(int i = 0; i < max_it; i++){
		if(jacc[i] > j_max){
			j_max = jacc[i];
			b_max = beta[i];
		}
	}
	int i_out = 0;
	// store scores for beta above optimal and kappa=0
	// these beta values are unlikely to improve with curvature
	for(int i = 0; i < max_it; i++){
		if(beta[i] > b_max){
			b_out[i_out] = beta[i];
			j_out[i_out] = jacc[i];
			h_out[i_out] = haus[i];
			k_out[i_out] = 0.0;
			i_out++;
		}
	}

	// sort beta in descending order
	qsort(beta, max_it, sizeof(float),compare_beta);
	float k0 = 0.0;
	float k1 = 0.6;
	for(int i_b = 0; i_b < max_it; i_b++){
		if(beta[i_b] > b_max || beta[i_b] < 1)
			continue;
		params->beta = beta[i_b];
		kappa[0] = k0;
		kappa[1] = k1;
		for(int i = 2; i < max_it; i++){
			kappa[i] = 0;
			jacc[i] = 0;
			haus[i] = N;
		}
		optimize_kappa(h_img,h_seed_indicator,h_gt,h_v0,dims,params,output_dir,jacc,haus,kappa,max_it);
		float k_max = 0;
		int h_min = N * 10;
		float j_max = 0.0;
		for(int i = 0; i < max_it; i++){
			int h_floor = (int) floor(haus[i] * 10);
			if(h_floor < h_min){
				h_min = h_floor;
				k_max = kappa[i];
				j_max = 0.0;
			}
			if(h_floor == h_min){
				if(jacc[i] > j_max){
					j_max = jacc[i];
					k_max = kappa[i];
				}
			}
		}
		// increase minimum kappa for faster optimization
		k0 = k_max;
		printf("Optimal kappa: %.3f, haus=%.1f, jacc=%.4f\n",k_max,(h_min/10.0),j_max);
		// store results
		for(int i = 0; i < max_it; i++){
			if(haus[i] > N-1)
				continue;
			b_out[i_out] = params->beta;
			k_out[i_out] = kappa[i];
			j_out[i_out] = jacc[i];
			h_out[i_out] = haus[i];
			i_out++;
		}
	}
	// save results
	FILE *f_op;
	f_op = fopen(op_path,"w+");
	fprintf(f_op,"beta\t kappa\t Hausdorff\t Jaccard\n");
	for(int i = 0; i < max_it*max_it; i++){
		if(b_out[i] > 0){
			fprintf(f_op,"%.0f\t %.3f\t %.2f\t %.4f\n",b_out[i],k_out[i],h_out[i],j_out[i]);
		}
	}
	fclose(f_op);

	free(h_v0);
	free(beta);
	free(kappa);
	free(jacc);
	free(haus);
	free(b_out);
	free(k_out);
	free(j_out);
	free(h_out);

	return 0;
}

int optimize_beta(
		float *h_img,
		float *h_seed_indicator,
		unsigned char *h_gt,
		float *h_v0,
		struct image_dimensions *dims,
		struct segmentation_parameters *params,
		char output_dir[],
		float *jacc,
		float *haus,
		float *beta,
		int max_it
		                ) {

	unsigned char * h_v_temp;

	float *h_v;
	float *h_w2; // forward weights
	int *h_NR; //array of node indices where nodes are "right" of an edge
	int *h_NL; // array of "left neighbour" node indices, used for curvature only
	int *h_NR2; // array of "right second neighbour" node indices, used for curvature only
	int *h_NL2; // array of "left second neighbour" node indices, used for curvature only
	int *h_NRs; // array of "right <step>th neighbour" node indices, used for curvature only
	int *h_NLs; // array of "left <step>th neighbour" node indices, used for curvature only

    printf("Setting initial beta values...\n");
	beta[0] = 800;
	beta[1] = 1300;

	printf("Getting dimensions...\n");
	int N = dims->N;
	int E = dims->E;

	int out_dims[3] = {dims->C, dims->R, dims->Z};

	// allocate host memory
	printf("allocating memory..\n");
	h_v_temp = (unsigned char *) malloc(N * sizeof(char));
	h_v = (float *) malloc(N * sizeof(float));
	h_w2 = (float *) malloc(E * sizeof(float));
	h_NR = (int *) malloc(E * sizeof(int));
	h_NL = (int *) malloc(3 * N * sizeof(int));
	h_NR2 = (int *) malloc(3 * N * sizeof(int));
	h_NL2 = (int *) malloc(3 * N * sizeof(int));
	h_NRs = (int *) malloc(3 * N * sizeof(int));
	h_NLs = (int *) malloc(3 * N * sizeof(int));

	int r_max = 20;

	// set neighbour arrays
	printf("Setting neighbours...\n");
	set_neighbours(h_NR, h_NL, h_NR2, h_NL2, h_NRs, h_NLs, dims, params->grad_step);

	int b_asc = 0;
	int b_des = 0;
	int i0,i1,i2,i3;

	for(int i_j = 0; i_j < max_it; i_j++){
		params->beta = beta[i_j];
		printf("Testing RW with beta=%.0f\n", beta[i_j]);

		// set output path to store temp data
		char output_name[256] = "";
		sprintf(output_name, "p_map_b%.0f_k%.3f.pgm", params->beta, params->kappa);
		char output_path[1024] = "";
		strcat(strcpy(output_path, output_dir), output_name);
		printf("Output path: %s\n", output_path);

		struct stat st = {0};
		if(stat(output_path, &st) == 0){
			h_v_temp = pgm_read(output_path);
			for(int i = 0; i < N; i++)
				h_v[i] = h_v_temp[i] * 1.0 / 255;
		} else {
			params->alpha = params->beta * 1.0 / (255 * 255); // distance weighting parameter
			set_edges(h_w2, h_NR, h_img, params, dims);
			for(int i = 0; i < N; i++)
				h_v[i] = h_v0[i];
			run_segmentation(h_v,h_w2,h_NR, h_NL,h_NR2, h_NL2,h_NRs, h_NLs,h_seed_indicator,params, dims);
			// write segmentation with same dimensions as input
			for(int i = 0; i < N; i++)
				h_v_temp[i] = (unsigned char) (h_v[i] * 255);
			pgm_write(h_v_temp, output_path, out_dims);
		}
		jacc[i_j] = get_jaccard(h_v, h_gt, N);

		haus[i_j] = get_hausdorff(h_v_temp, h_gt, 127, 0, dims);

		if(i_j == 1){
			if(jacc[i_j] > jacc[i_j-1]){
				beta[i_j+1] = beta[i_j] + beta[i_j-1];
				b_asc = 1;
			} else {
				beta[i_j+1] = beta[i_j] - beta[i_j-1];
				b_des = 1;
			}
		} else if(1 < i_j && i_j < max_it){
			if(b_asc == 1){
				if(jacc[i_j] > jacc[i_j-1]){
					beta[i_j+1] = beta[i_j] + beta[i_j-1];
				} else {
					b_asc = 0;
					i0 = i_j-2;
					i1 = i_j;
					i2 = i_j-1;
					i3 = i_j+1;
					beta[i3] = beta[i0] + beta[i1] - beta[i2];
				}
			} else if(b_des == 1){
				if(jacc[i_j] > jacc[i_j-1]){
					beta[i_j+1] = beta[i_j-1] - beta[i_j];
				} else {
					b_des = 0;
					i0 = i_j;
					if(i_j > 2){
						i1 = i_j-2;
						i2 = i_j-1;
					} else {
						i1 = i_j-1;
						i2 = i_j-2;
					}
					i3 = i_j+1;
					beta[i3] = beta[i0] + beta[i1] - beta[i2];
				}
			} else {
				if(jacc[i3] > jacc[i2]){
					i0 = i2;
					i2 = i3;
					i3 = i_j+1;
					beta[i3] = beta[i0] + beta[i1] - beta[i2];
				} else {
					i1 = i3;
					i3 = i2;
					i2 = i_j+1;
					beta[i2] = beta[i0] + beta[i1] - beta[i3];
				}
			}
		}

		printf("Current optimization sequence:\n");
		for(int i = 0; i <= i_j; i++){
			printf("%.0f, %.2f, %.4f\n", beta[i], haus[i], jacc[i]);
		}

		if(abs(beta[i_j+1]-beta[i_j]) < 100){
			beta[i_j] = 0;
			break;
		}
	}
	free(h_v);
	free(h_w2);
	free(h_NR);
	free(h_NL);
	free(h_NR2);
	free(h_NL2);
	free(h_NRs);
	free(h_NLs);

    return 0;
}

int optimize_kappa(
		float *h_img,
		float *h_seed_indicator,
		unsigned char *h_gt,
		float *h_v0,
		struct image_dimensions *dims,
		struct segmentation_parameters *params,
		char output_dir[],
		float *jacc,
		float *haus,
		float *kappa,
		int max_it
		                ) {

	unsigned char * h_v_temp;

	float inv_phi = (sqrt(5) - 1) / 2;
	float inv_phi2 = inv_phi * inv_phi;

	float *h_v;
	float *h_w2; // forward weights
	int *h_NR; //array of node indices where nodes are "right" of an edge
	int *h_NL; // array of "left neighbour" node indices, used for curvature only
	int *h_NR2; // array of "right second neighbour" node indices, used for curvature only
	int *h_NL2; // array of "left second neighbour" node indices, used for curvature only
	int *h_NRs; // array of "right <step>th neighbour" node indices, used for curvature only
	int *h_NLs; // array of "left <step>th neighbour" node indices, used for curvature only

	int i0 = 0,i1 = 1,i2 = 2,i3 = 3, i_start = 3;
	float h = kappa[1] - kappa[0];
	kappa[2] = kappa[0] + h * inv_phi2;
	kappa[3] = kappa[0] + h * inv_phi;
	// default maximum kappa is 0.6
	if(kappa[1] < 0.001){
		printf("Setting initial kappa values...\n");
		kappa[0] = 0.0;
		kappa[1] = 0.6;
		h = kappa[1] - kappa[0];
		kappa[2] = kappa[0] + h * inv_phi2;
		kappa[3] = kappa[0] + h * inv_phi;
	} else if(kappa[0] > 0){
		kappa[4] = kappa[3];
		i3 = 4;
		kappa[3] = kappa[2];
		i2 = 3;
		kappa[2] = kappa[1];
		i1 = 2;
		kappa[1] = kappa[0];
		i0 = 1;
		kappa[0] = 0.0;
		i_start = 4;
	}

	printf("Getting dimensions...\n");
	int N = dims->N;
	int E = dims->E;

	int out_dims[3] = {dims->C, dims->R, dims->Z};
	char gt_path[1024] = "";
	strcat(strcpy(gt_path, output_dir), "GT.pgm");
	pgm_write(h_gt, gt_path, out_dims);

	// allocate host memory
	printf("allocating memory..\n");
	h_v_temp = (unsigned char *) malloc(N * sizeof(char));
	h_v = (float *) malloc(N * sizeof(float));
	h_w2 = (float *) malloc(E * sizeof(float));
	h_NR = (int *) malloc(E * sizeof(int));
	h_NL = (int *) malloc(3 * N * sizeof(int));
	h_NR2 = (int *) malloc(3 * N * sizeof(int));
	h_NL2 = (int *) malloc(3 * N * sizeof(int));
	h_NRs = (int *) malloc(3 * N * sizeof(int));
	h_NLs = (int *) malloc(3 * N * sizeof(int));

	int r_max = 20;

	// set neighbour arrays
	printf("Setting neighbours...\n");
	set_neighbours(h_NR, h_NL, h_NR2, h_NL2, h_NRs, h_NLs, dims, params->grad_step);
	set_edges(h_w2, h_NR, h_img, params, dims);

	for(int i_j = 0; i_j < max_it; i_j++){
		params->kappa = kappa[i_j];
		printf("Testing RW with kappa=%.3f\n", kappa[i_j]);

		// set output path to store temp data
		char output_name[256] = "";
		sprintf(output_name, "p_map_b%.0f_k%.3f.pgm", params->beta, params->kappa);
		char output_path[1024] = "";
		strcat(strcpy(output_path, output_dir), output_name);
		printf("Output path: %s\n", output_path);

		struct stat st = {0};
		if(stat(output_path, &st) == 0){
			h_v_temp = pgm_read(output_path);
			for(int i = 0; i < N; i++)
				h_v[i] = h_v_temp[i] * 1.0 / 255;
		} else {
			for(int i = 0; i < N; i++)
				h_v[i] = h_v0[i];
			run_segmentation(h_v,h_w2,h_NR, h_NL,h_NR2, h_NL2,h_NRs, h_NLs,h_seed_indicator,params, dims);
			// write segmentation with same dimensions as input
			for(int i = 0; i < N; i++)
				h_v_temp[i] = (unsigned char) (h_v[i] * 255);
			pgm_write(h_v_temp, output_path, out_dims);
		}
		jacc[i_j] = get_jaccard(h_v, h_gt, N);

		if(jacc[i_j] > 0.5){
			haus[i_j] = get_hausdorff(h_v_temp, h_gt, 127, 0, dims);
		} else {
			haus[i_j] = N;
		}

		if(i_j == 0){
			// initialise v0 for k=0 for faster optimization
			for(int i = 0; i < N; i++){
				h_v0[i] = h_v[i];
			}
		} else if(i_start <= i_j && i_j < max_it){
			h = h * inv_phi;
			int f_h3 = (int) roundf(haus[i3] * 10);
			int f_h2 = (int) roundf(haus[i2] * 10);
			if(f_h3 < f_h2 || (f_h3 == f_h2 && jacc[i3] > jacc[i2])){
				i0 = i2;
				i2 = i3;
				i3 = i_j+1;
				kappa[i3] = kappa[i0] + h * inv_phi;
			} else {
				i1 = i3;
				i3 = i2;
				i2 = i_j+1;
				kappa[i2] = kappa[i0] + h * inv_phi2;
			}
		}

		printf("Current optimization sequence:\n");
		for(int i = 0; i <= i_j; i++){
			printf("%.0f, %.3f, %.2f, %.4f\n", params->beta, kappa[i], haus[i], jacc[i]);
		}

		if(h < 0.01){
			break;
		}
	}
	free(h_v);
	free(h_w2);
	free(h_NR);
	free(h_NL);
	free(h_NR2);
	free(h_NL2);
	free(h_NRs);
	free(h_NLs);

    return 0;
}

float get_hausdorff(unsigned char *b1, unsigned char *b2, int t1, int t2, struct image_dimensions *dims){
	int N = dims->N;
	float d_max = 0;
	for(int i = 0; i < N; i++){
		if(b1[i] <= t1){
			if(b2[i] > t2){
				d_max = get_dist_above(b1, t1, i, d_max, dims);

			}
		}
		/*
		 else if(b2[i] <= t2){
			d_max = get_dist_above(b2, t2, i, d_max, dims);
		}
		*/
	}
	return d_max;
}

float get_dist_above(unsigned char *b1, int t1, int p, float d_max, struct image_dimensions *dims){
	int N = dims->N;
	int R = dims->R;
	int C = dims->C;
	int Z = dims->Z;
	int d0 = (int) floor(d_max)-1;
	int d02 = d0 * d0;
	for(int d1 = d0 +1; d1 < N; d1++){
		int d12 = d1 * d1;
		for(int dr = -d1; dr <= d1; dr++){
			int d1r = (int) floor(sqrt(d12 - dr*dr));
			for(int dc = -d1r; dc <= d1r; dc++){
				int d1c = (int) floor(sqrt(d12 - dr*dr - dc*dc));
				int d0c = (int) floor(sqrt(fmax(d02 - dr*dr - dc*dc,0)));
				for(int dz = d0c; dz <= d1c; dz++){
					int p1 = p + dz * R * C + dr * C + dc;
					if(0 <= p1 && p1 < N){
						if(b1[p1] > t1){
							if(d1 == d0 + 1)
								return d_max;
							else
								return sqrt(dr * dr + dc * dc + dz * dz);

						}
					}
				}
				for(int dz = -d1c; dz <= -d0c; dz++){
					int p1 = p + dz * R * C + dr * C + dc;
					if(0 <= p1 && p1 < N){
						if(b1[p1] > t1){
							if(d1 == d0 + 1)
								return d_max;
							else
								return sqrt(dr * dr + dc * dc + dz * dz);
						}
					}
				}
			}
			d02 = d12;
		}
	}
	return d_max;
}


float get_jaccard(float *h_v, unsigned char *h_gt, int N){
	long j_inter = 0;
	long j_union = 0;
	for(int i = 1; i < N; i++){
		if(h_v[i] > 0.5){
			j_union++;
			if(h_gt[i] > 0){
				j_inter++;
			}
		} else if(h_gt[i] > 0){
			j_union++;
		}
	}
	return (float) (j_inter * 1.0 / j_union);
}

int compare_beta(const void* a, const void* b){
	// inputs are floats
	float d_f = ( *(float*)b - *(float*)a);
	return (int) d_f;
}


int segment_range(
		float *h_img,
		unsigned char *h_seed_img,
		struct image_dimensions *dims0,
		struct image_dimensions *dims1,
		struct segmentation_parameters *params,
		char output_dir[],
		char img_name_base[],
		float *beta,
		float *kappa,
		int b_sz,
		int k_sz
		){

	int N = dims1->N;
	int E = dims1->E;

	// node values
	float *h_seed_indicator;
	float *h_v;
	float *h_v0;
	// edge values
	float *h_w2;
	int *h_NR;
	int *h_NL;
	int *h_NR2;
	int *h_NL2;
	int *h_NRs;
	int *h_NLs;
	// output values
	float *v_out;
	unsigned char *v_out_char;

	h_seed_indicator = (float *) malloc(N * sizeof(float));
	h_v = (float *) malloc(N * sizeof(float));
	h_v0 = (float *) malloc(N * sizeof(float));

	h_w2 = (float *) malloc(E * sizeof(float));
	// only right node used for diffusion component, so requires values for all edges
	h_NR = (int *) malloc(E * sizeof(int));
	// neighbors for curvature, along with first 3N values of NR
	h_NL = (int *) malloc(3 * N * sizeof(int));
	h_NR2 = (int *) malloc(3 * N * sizeof(int));
	h_NL2 = (int *) malloc(3 * N * sizeof(int));
	h_NRs = (int *) malloc(3 * N * sizeof(int));
	h_NLs = (int *) malloc(3 * N * sizeof(int));

	v_out = (float *) malloc(dims0->N * sizeof(float));
	v_out_char = (unsigned char *) malloc(dims0->N * sizeof(char));

	// get connectivity of voxels
	set_neighbours(h_NR, h_NL, h_NR2, h_NL2, h_NRs, h_NLs, dims1, params->grad_step);
	// mark points that don't change over time
	set_seed_indicator(h_seed_indicator, h_seed_img, N);

	for(int i_b = 0; i_b < b_sz; i_b++){
		params->beta = beta[i_b]; // gradient weighting parameter
		params->alpha = params->beta * 1.0 / (255 * 255); // distance weighting parameter
		// initialize v0
		set_v0(h_v0, h_seed_img, N);
		// generate edge weights
		set_edges(h_w2, h_NR, h_img, params, dims1);

		// run initial random walker segmentation
		params->kappa = 0.0;
		char output_name[256] = "";
		sprintf(output_name, "%s_SEG_b%.0f_k%.3f.pgm",
				img_name_base, params->beta, params->kappa);
		printf("Output file name: %s\n", output_name);
		char output_path[1024] = "";
		strcat(strcpy(output_path, output_dir), output_name);

		// run the random walker to get initial conditions
		struct stat st = {0};
		if(stat(output_path, &st) == 0){
			v_out_char = pgm_read(output_path);
			for(int i = 0; i < dims0->N; i++)
				v_out[i] = v_out_char[i] * 1.0 / 255;
			h_v0 = copy_image(v_out, dims0, dims1, FLOAT_IM);
		} else {
			// the random walker
			run_segmentation(
					h_v0, h_w2, h_NR, h_NL,
					h_NR2, h_NL2, h_NRs, h_NLs,
					h_seed_indicator, params, dims1);

			// seed augmentation currently not in use
			if(params->augment_seeds == 1){
				printf("Augmenting seeds...\n");
				for(int p = 0; p < N; p++){
					if(h_v0[p] > params->v_T){
						h_v0[p] = 1.0;
						h_seed_indicator[p] = 0;
					}
				}
				// update v0
				run_segmentation(
						h_v0, h_w2, h_NR, h_NL,
						h_NR2, h_NL2, h_NRs, h_NLs,
						h_seed_indicator, params, dims1);
			}
			v_out = (float *) copy_image(h_v0, dims1, dims0, FLOAT_IM);
			for(int i = 0; i < dims0->N; i++)
				v_out_char[i] = (unsigned char) (255 * v_out[i]);
			// write segmentation with same dimensions as input
			int out_dims[3] = {dims0->C, dims0->R, dims0->Z};
			pgm_write(v_out_char, output_path, out_dims);
		}
		for(int i_k = 0; i_k < k_sz; i_k++){
			params->kappa = kappa[i_k]; // curvature coefficient
			// initialize v
			for(int i = 0; i < N; i++)
				h_v[i] = h_v0[i];

			// set output path
			strncpy(output_name, "",sizeof(output_path));
			sprintf(output_name, "%s_SEG_b%.0f_k%.3f.pgm",
					img_name_base, params->beta, params->kappa);
			printf("Output file name: %s\n", output_name);
			strncpy(output_path, "", sizeof(output_path));
			strcat(strcpy(output_path, output_dir), output_name);

			if(stat(output_path, &st) != 0){
				///////////////////////////////////////////////
				// GPU processing

				// the curvature-enhanced random walker
				run_segmentation(
						h_v, h_w2, h_NR, h_NL,
						h_NR2, h_NL2, h_NRs, h_NLs,
						h_seed_indicator, params, dims1);

				// write the output probabilities as an 8-bit pgm image
				v_out = (float *) copy_image(h_v, dims1, dims0, FLOAT_IM);
				for(int i = 0; i < dims0->N; i++)
					v_out_char[i] = (unsigned char) (255 * v_out[i]);
				// write segmentation with same dimensions as input
				int out_dims[3] = {dims0->C, dims0->R, dims0->Z};
				pgm_write(v_out_char, output_path, out_dims);
			}
		}

	}

	free(h_seed_indicator);
	free(h_v);
	free(h_v0);
	free(h_w2);
	free(h_NR);
	free(h_NL);
	free(h_NR2);
	free(h_NL2);
	free(h_NRs);
	free(h_NLs);
	free(v_out);
	free(v_out_char);

	return 0;
}

int segment_paired(
		float *h_img,
		unsigned char *h_seed_img,
		struct image_dimensions *dims0,
		struct image_dimensions *dims1,
		struct segmentation_parameters *params,
		char output_dir[],
		char img_name_base[],
		float *beta,
		float *kappa,
		int b_sz
		){

	int N = dims1->N;
	int E = dims1->E;

	// node values
	float *h_seed_indicator;
	float *h_v;
	float *h_v0;
	// edge values
	float *h_w2;
	int *h_NR;
	int *h_NL;
	int *h_NR2;
	int *h_NL2;
	int *h_NRs;
	int *h_NLs;
	// output values
	float *v_out;
	unsigned char *v_out_char;

	h_seed_indicator = (float *) malloc(N * sizeof(float));
	h_v = (float *) malloc(N * sizeof(float));
	h_v0 = (float *) malloc(N * sizeof(float));

	h_w2 = (float *) malloc(E * sizeof(float));
	// only right node used for diffusion component, so requires values for all edges
	h_NR = (int *) malloc(E * sizeof(int));
	// neighbors for curvature, along with first 3N values of NR
	h_NL = (int *) malloc(3 * N * sizeof(int));
	h_NR2 = (int *) malloc(3 * N * sizeof(int));
	h_NL2 = (int *) malloc(3 * N * sizeof(int));
	h_NRs = (int *) malloc(3 * N * sizeof(int));
	h_NLs = (int *) malloc(3 * N * sizeof(int));

	v_out = (float *) malloc(dims0->N * sizeof(float));
	v_out_char = (unsigned char *) malloc(dims0->N * sizeof(char));

	// get connectivity of voxels
	set_neighbours(h_NR, h_NL, h_NR2, h_NL2, h_NRs, h_NLs, dims1, params->grad_step);
	// mark points that don't change over time
	set_seed_indicator(h_seed_indicator, h_seed_img, N);

	for(int i_b = 0; i_b < b_sz; i_b++){
		params->beta = beta[i_b]; // gradient weighting parameter
		params->alpha = params->beta * 1.0 / (255 * 255); // distance weighting parameter
		// initialize v0
		set_v0(h_v0, h_seed_img, N);
		// generate edge weights
		set_edges(h_w2, h_NR, h_img, params, dims1);

		// run initial random walker segmentation
		params->kappa = 0.0;
		char output_name[256] = "";
		sprintf(output_name, "%s_SEG_b%.0f_k%.3f.pgm",
				img_name_base, params->beta, params->kappa);
		printf("Output file name: %s\n", output_name);
		char output_path[1024] = "";
		strcat(strcpy(output_path, output_dir), output_name);

		// run the random walker to get initial conditions
		struct stat st = {0};
		if(stat(output_path, &st) == 0){
			v_out_char = pgm_read(output_path);
			for(int i = 0; i < dims0->N; i++)
				v_out[i] = v_out_char[i] * 1.0 / 255;
			h_v0 = copy_image(v_out, dims0, dims1, FLOAT_IM);
		} else {
			// the random walker
			run_segmentation(
					h_v0, h_w2, h_NR, h_NL,
					h_NR2, h_NL2, h_NRs, h_NLs,
					h_seed_indicator, params, dims1);

			// augment seeds if required
			if(params->augment_seeds == 1){
				printf("Augmenting seeds...\n");
				for(int p = 0; p < N; p++){
					if(h_v0[p] > params->v_T){
						h_v0[p] = 1.0;
						h_seed_indicator[p] = 0;
					}
				}
				// update v0
				run_segmentation(
						h_v0, h_w2, h_NR, h_NL,
						h_NR2, h_NL2, h_NRs, h_NLs,
						h_seed_indicator, params, dims1);
			}
			v_out = (float *) copy_image(h_v0, dims1, dims0, FLOAT_IM);
			for(int i = 0; i < dims0->N; i++)
				v_out_char[i] = (unsigned char) (255 * v_out[i]);
			// write segmentation with same dimensions as input
			int out_dims[3] = {dims0->C, dims0->R, dims0->Z};
			pgm_write(v_out_char, output_path, out_dims);
		}
		params->kappa = kappa[i_b]; // curvature coefficient
		// initialize v
		for(int i = 0; i < N; i++)
			h_v[i] = h_v0[i];

		// set output path
		strncpy(output_name, "",sizeof(output_path));
		sprintf(output_name, "%s_SEG_b%.0f_k%.3f.pgm",
				img_name_base, params->beta, params->kappa);
		printf("Output file name: %s\n", output_name);
		strncpy(output_path, "", sizeof(output_path));
		strcat(strcpy(output_path, output_dir), output_name);

		if(stat(output_path, &st) != 0){
			///////////////////////////////////////////////
			// GPU processing

			// the curvature-enhanced random walker
			run_segmentation(
					h_v, h_w2, h_NR, h_NL,
					h_NR2, h_NL2, h_NRs, h_NLs,
					h_seed_indicator, params, dims1);

			// write the output probabilities as an 8-bit pgm image
			v_out = (float *) copy_image(h_v, dims1, dims0, FLOAT_IM);
			for(int i = 0; i < dims0->N; i++)
				v_out_char[i] = (unsigned char) (255 * v_out[i]);
			// write segmentation with same dimensions as input
			int out_dims[3] = {dims0->C, dims0->R, dims0->Z};
			pgm_write(v_out_char, output_path, out_dims);
		}
	}

	free(h_seed_indicator);
	free(h_v);
	free(h_v0);
	free(h_w2);
	free(h_NR);
	free(h_NL);
	free(h_NR2);
	free(h_NL2);
	free(h_NRs);
	free(h_NLs);
	free(v_out);
	free(v_out_char);

	return 0;
}