/*
 * segmentation.cu
 *
 *  Created on: Oct 30, 2019
 *      Author: lutton
 */

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

#include "segmentation.h"


/** edge steps depending on direction:
 * 	0: vertical
 * 	1: horizontal
 * 	2: up-down
 * 	3: horizontal-vertical +ve diagonal
 * 	4: horizontal-vertical -ve diagonal
 * 	5: vertical-up-down +ve diagonal
 * 	6: vertical-up-down -ve diagonal
 * 	7: horizontal-up-down +ve diagonal
 * 	8: horizontal-up-down -diagonal
 */
const int r_step[9] = {0,1,0,1, 1,0, 0,1, 1};
const int c_step[9] = {1,0,0,1,-1,1, 1,0, 0};
const int z_step[9] = {0,0,1,0, 0,1,-1,1,-1};


// gets dimensions struct with reduced image size based on seed position
int set_image_dimensions(unsigned char *seeds, struct image_dimensions *dims0, struct image_dimensions *dims1, int max_step){
	int R0 = dims0->R;
	int C0 = dims0->C;
	int Z0 = dims0->Z;
	int c0 = C0;
	int r0 = R0;
	int z0 = Z0;
	int c1 = 0;
	int r1 = 0;
	int z1 = 0;

	for(int r = 1; r < R0 - 1; r++){
		for(int c = 1; c < C0 - 1; c++){
			for(int z = 1; z < Z0 - 1; z++){
				int p = z * R0 * C0 + r * C0 + c;
				if(seeds[p] != 0){
					r0 = (r - max_step < r0) ? r - max_step : r0;
					c0 = (c - max_step < c0) ? c - max_step : c0;
					z0 = (z - max_step < z0) ? z - max_step : z0;
					r1 = (r + max_step > r1) ? r + max_step : r1;
					c1 = (c + max_step > c1) ? c + max_step : c1;
					z1 = (z + max_step > z1) ? z + max_step : z1;
				}
			}
		}
	}

	// new volume inclusive of bounds
	int R1 = r1 - r0 + 1;
	int C1 = c1 - c0 + 1;
	int Z1 = z1 - z0 + 1;

	// reset boundaries where bounding box is too short.
	if(C1 < 10){
		C1 = C0;
		c0 = 0;
		c1 = C0;
	}
	if(R1 < 10){
		R1 = R0;
		r0 = 0;
		r1 = R0;
	}
	if(Z1 < 10){
		Z1 = Z0;
		z0 = 0;
		z1 = Z0;
	}

	// TODO: optimise these values based on image dimensions
	int threads_per_block_x = 16;
	int threads_per_block_y = 8;
	int threads_per_block_z = 8;

	// increase size so that each dimension is divisible by the block size along that axis.
	int R = (int) ceil(R1 * 1.0 / threads_per_block_y) * threads_per_block_y;
	int C = (int) ceil(C1 * 1.0 / threads_per_block_x) * threads_per_block_x;
	int Z = (int) ceil(Z1 * 1.0 / threads_per_block_z) * threads_per_block_z;
	int N = R * C * Z;

	// adjust bounds to center the volume
	r0 = r0 - (R - R1) / 2;
	r1 = r0 + R;
	c0 = c0 - (C - C1) / 2;
	c1 = c0 + C;
	z0 = z0 - (Z - Z1) / 2;
	z1 = z0 + Z;

	dims1->c0 = c0;
	dims1->c1 = c1;
	dims1->r0 = r0;
	dims1->r1 = r1;
	dims1->z0 = z0;
	dims1->z1 = z1;

	dims1->N = N;
	dims1->R = R;
	dims1->C = C;
	dims1->Z = Z;

	dims1->threads_per_block_x = threads_per_block_x;
	dims1->threads_per_block_y = threads_per_block_y;
	dims1->threads_per_block_z = threads_per_block_z;

	dims1->num_blocks_x = C / threads_per_block_x;
	dims1->num_blocks_y = R / threads_per_block_y;
	dims1->num_blocks_z = Z / threads_per_block_z;

	dims1->edge_direction_count = dims0->edge_direction_count;

	dims1->E = N * dims1->edge_direction_count;

	return 0;
}

// copies the input image to a location in the output image specified by the dims structs
void * copy_image(void *img_in, struct image_dimensions *dims0, struct image_dimensions *dims1, im_t im_type){

	float fill_float = 0.0;
	unsigned char fill_char = 0;

	void *img_out;

	// get offset of image 1 to image 0
	int c0 = dims1->c0 - dims0->c0;
	int r0 = dims1->r0 - dims0->r0;
	int z0 = dims1->z0 - dims0->z0;

	int R = dims1->R;
	int C = dims1->C;
	int Z = dims1->Z;
	int N = dims1->N;

	int R0 = dims0->R;
	int C0 = dims0->C;
	int Z0 = dims0->Z;

	switch(im_type){
	case FLOAT_IM :
		img_out = malloc(sizeof(float) * N);
		break;

	case BYTE_IM :
		img_out = malloc(sizeof(char) * N);
		break;
	}

	//copy image data into resized image
	for (int r = 0; r < R; r++)
		for (int c = 0; c < C; c++)
			for (int z = 0; z < Z; z++){
				int p = z * R * C + r * C + c;
				int pc0 = c0 + c;
				int pr0 = r0 + r;
				int pz0 = z0 + z;
				int in_bounds = (
						(0 <= pc0 && pc0 < C0) &&
						(0 <= pr0 && pr0 < R0) &&
						(0 <= pz0 && pz0 < Z0));
				int p0 = (in_bounds) ? pz0 * R0 * C0 + pr0 * C0 + pc0 : 0;
				switch(im_type){
				case FLOAT_IM :
					((float *)img_out)[p] = (in_bounds) ? ((float *)img_in)[p0] : fill_float;
					break;
				case BYTE_IM :
					((unsigned char *)img_out)[p] = (in_bounds) ? ((unsigned char *)img_in)[p0] : fill_char;
					break;
				}
			}
	return img_out;
}

int set_neighbours(int *h_NR, int *h_NL, int *h_NR2, int *h_NL2, int *h_NRs, int *h_NLs, struct image_dimensions *dims, int s){

	int R = dims->R;
	int C = dims->C;
	int Z = dims->Z;
	int N = dims->N;

	int n,e;
	int r1,c1,z1;
	int r2,c2,z2;

	for(int r = 0; r < R; r++){
		for(int c = 0; c < C; c++){
			for(int z = 0; z < Z; z++){
				for(int d = 0; d < dims->edge_direction_count; d++){
					n = z * C * R + r * C + c;
					e = n + d * N;
					r1 = (r + r_step[d] + R) % R;
					c1 = (c + c_step[d] + C) % C;
					z1 = (z + z_step[d] + Z) % Z;
					h_NR[e] = z1 * R * C + r1 * C + c1;
					if(d < 3){
						r1 = (r - r_step[d] + R) % R;
						c1 = (c - c_step[d] + C) % C;
						z1 = (z - z_step[d] + Z) % Z;
						h_NL[e] = z1 * R * C + r1 * C + c1;

						r1 = (r + 2 * r_step[d] + R) % R;
						c1 = (c + 2 * c_step[d] + C) % C;
						z1 = (z + 2 * z_step[d] + Z) % Z;
						h_NR2[e] = z1 * R * C + r1 * C + c1;

						r1 = (r - 2 * r_step[d] + R) % R;
						c1 = (c - 2 * c_step[d] + C) % C;
						z1 = (z - 2 * z_step[d] + Z) % Z;
						h_NL2[e] = z1 * R * C + r1 * C + c1;

						// for curvature, periodic boundaries pose an issue with gradient calculation
						// to avoid this issue, we take the gradient in the direction of the boundary
						// at the nearest point away from the boundary.
						r1 = r + s * r_step[d];
						r2 = r - s * r_step[d];
						if(r1 > R - 1){
							r1 = R - 1;
							r2 = r1 - 2 * s * r_step[d];
						} else if(r2 < 0){
							r2 = 0;
							r1 = r2 + 2 * s * r_step[d];
						}

						c1 = c + s * c_step[d];
						c2 = c - s * c_step[d];
						if(c1 > C - 1){
							c1 = C - 1;
							c2 = c1 - 2 * s * c_step[d];
						} else if(c2 < 0){
							c2 = 0;
							c1 = c2 + 2 * s * c_step[d];
						}

						z1 = z + s * z_step[d];
						z2 = z - s * z_step[d];
						if(z1 > Z - 1){
							z1 = Z - 1;
							z2 = z1 - 2 * s * z_step[d];
						} else if(z2 < 0){
							z2 = 0;
							z1 = z2 + 2 * s * z_step[d];
						}

						h_NRs[e] = z1 * R * C + r1 * C + c1;
						h_NLs[e] = z2 * R * C + r2 * C + c2;

					}
				}
			}
		}
	}
	return 0;
}

// sets the edge weights based on image data and input parameters
int set_edges(float *h_w2, int *h_NR, float *h_img,
		struct segmentation_parameters *params, struct image_dimensions *dims
		){
	int N = dims->N;
	int R = dims->R;
	int C = dims->C;
	int Z = dims->Z;
	int n,e;
	for(int r = 0; r < R; r++)
		for(int c = 0; c < C; c++)
			for(int z = 0; z < Z; z++){
				for(int d = 0; d < dims->edge_direction_count; d++){
					n = z * R * C + r * C + c;
					e = n + d * N;
					// check boundary
					int rR = r + r_step[d];
					int cR = c + c_step[d];
					int zR = z + z_step[d];
					if(h_NR[e] == zR * R * C + rR * C + cR){
						// gradient weighting
						h_w2[e] = exp(-params->beta * pow(h_img[n] - h_img[h_NR[e]], 2));
						// distance weighting
						float dist = 1.0;
						if(d < 3){
							dist = 1.0;
						} else if(3 <= d && d <= 8){ //scale for any in-plane diagonal (sqrt(2)^2)
							dist = sqrt(2.0);
						}
						// multiply by exp(alpha) to normalize weights
						h_w2[e] = h_w2[e] * exp(params->alpha * (1.0 - dist * dist));
					} else {
						// periodic with 0 weight.
						h_w2[e] = 0;
					}
				}
			}
	return 0;
}

int halve_dimensions(unsigned char *seed_data, struct image_dimensions *dims0, struct image_dimensions *dims1){
	int R0 = dims0->R;
	int C0 = dims0->C;
	int Z0 = dims0->Z;
	long rC=0,cC=0,zC=0,nC=0;
	for(int r = 0; r < R0; r++)
		for(int c = 0; c < C0; c++)
			for(int z = 0; z < Z0; z++){
				int p = z * R0 * C0 + r * C0 + c;
				if(seed_data[p] == 255){
					rC += r;
					cC += c;
					zC += z;
					nC ++;
				}
			}
	rC = rC / nC - dims1->r0;
	cC = cC / nC - dims1->c0;
	zC = zC / nC - dims1->z0;
	/*
	dims1->R = (int) (dims1->threads_per_block_y * ceil(dims1->R / (2.0 * dims1->threads_per_block_y)));
	dims1->C = (int) (dims1->threads_per_block_x * ceil(dims1->C / (2.0 * dims1->threads_per_block_x)));
	dims1->Z = (int) (dims1->threads_per_block_z * ceil(dims1->Z / (2.0 * dims1->threads_per_block_z)));
	dims1->N = dims1->R * dims1->C * dims1->Z;
	dims1->E = dims1->edge_direction_count * dims1->N;
	*/
	dims1->R = dims1->R / 2;
	dims1->C = dims1->C / 2;
	dims1->Z = dims1->Z / 2;
	dims1->N = dims1->N / 8;
	dims1->E = dims1->E / 8;
	dims1->threads_per_block_x = dims1->threads_per_block_x/2;
	dims1->threads_per_block_y = dims1->threads_per_block_y/2;
	dims1->threads_per_block_z = dims1->threads_per_block_z/2;
	if(rC > dims1->R)
		dims1->r0 = dims1->r0 + dims1->R;
	else
		dims1->r1 = dims1->r1 - dims1->R;
	if(cC > dims1->C)
		dims1->c0 = dims1->c0 + dims1->C;
	else
		dims1->c1 = dims1->c1 - dims1->C;
	if(zC > dims1->Z)
		dims1->z0 = dims1->z0 + dims1->Z;
	else
		dims1->z1 = dims1->z1 - dims1->Z;
	return 0;
}

int set_seed_indicator(float *h_seed_indicator, unsigned char *h_seed_img, int N){
	for(int i = 0; i < N; i++){
		// BG seeds have value 0
		if(h_seed_img[i] == 0){
			h_seed_indicator[i] = 0;
		// FG seeds have value 255
		}else if(h_seed_img[i] == 255){
			h_seed_indicator[i] = 0;
		// non-seeds are points which do not have seed values
		} else {
			h_seed_indicator[i] = 1;
		}
	}
	return 0;
}

int set_v0(float * h_v0, unsigned char *h_seed_img, int N){
	for(int i = 0; i < N; i++){
		// BG seeds have value 0
		if(h_seed_img[i] == 0){
			h_v0[i] = 0;
		// FG seeds have value 255
		}else if(h_seed_img[i] == 255){
			h_v0[i] = 1;
		// non-seeds are points which do not have seed values
		} else {
			h_v0[i] = 0.5;
		}
	}
	return 0;
}