#include<stdio.h>
#include<stdlib.h>
#include<fcntl.h>
#include<unistd.h>
#define STB_IMAGE_WRITE_IMPLEMENTATION
#define STB_IMAGE_IMPLEMENTATION
#include<stb/stb_image.h>
#include<stb/stb_image_write.h>
#include<sys/time.h>
#include<pthread.h>
#include "scolor.h"

pthread_mutex_t m;

unsigned char blend(unsigned char a, unsigned char b, float ratio){
	return b * ratio + a * (1 - ratio);
}

unsigned char pix_average(unsigned char* pix){
	int sum = pix[0];
	sum += pix[1];
	sum += pix[2];
	return sum/3;
}

void apply_color(unsigned char* pix, unsigned char r, unsigned char g, unsigned char b){
	pix[0] = r;
	pix[1] = g;
	pix[2] = b;
}

unsigned char mortar_threshold = 160;
size_t array_length;
unsigned char *image, *output;
int width, height;

char ismortarcolor(unsigned char* pix){
	if(pix < image)
		return 0;
	if(pix > (image + array_length))
		return 0;
	return pix[0] > mortar_threshold && pix[1] > mortar_threshold && pix[2] > mortar_threshold;
}

struct neighbor_node {
	size_t idx;
	struct neighbor_node *next;
};

// Changed this to reject indices that aren't valid or were already filled
void push_neighbor(size_t idx, struct neighbor_node **stack){
	if(idx > width*height)
		return;
	if(output[idx*3] == 254)
		return;
	struct neighbor_node *new_neighbor = malloc(sizeof(struct neighbor_node));
	new_neighbor->next = *stack;
	new_neighbor->idx = idx; // Forgot that part in class
	*stack = new_neighbor;
}

void free_neighbors(struct neighbor_node* place){
	if(!place)
		return;	
	free_neighbors(place->next);
	free(place);
}

size_t pop_neighbor(struct neighbor_node **stack){
	size_t toreturn = (*stack)->idx;
	struct neighbor_node *tofree = *stack;
	*stack = tofree->next;
	free(tofree);
	return toreturn;
}

size_t fill_counts[72];
size_t flood_fill(size_t idx, size_t thread_idx){
//	pthread_mutex_lock(&m);
 	fill_counts[thread_idx] = 0;
 	if(output[idx * 3] == 254)
		return 0;
	struct neighbor_node *neighbors = 0;
	size_t na[] = {idx + width, idx + 1, idx - 1, idx - width, 0};
	for(size_t *np = na; *np; np++)
		push_neighbor(*np, &neighbors);
	while(neighbors){
		size_t cn = pop_neighbor(&neighbors);
		if(output[cn * 3] == 254 && output[cn * 3 + 2] > thread_idx * 3 % 255) {
			size_t other_thread_idx = output[cn * 3 + 2] / 3;
			fill_counts[other_thread_idx] += fill_counts[thread_idx];
//			printf("Thread %d stopping flood fill to defer to thread %d\n", thread_idx, output[cn * 3 + 2] / 3);
			return 0;
		}
		if(output[cn * 3] != 254 && image[cn] < 175){
			output[cn * 3] = 254;
			output[cn * 3 + 2] = thread_idx * 3 % 255;
			fill_counts[thread_idx]++;
			size_t na[] = {cn + width, cn - 1, cn + 1, cn - width, 0};
			for(size_t *np = na; *np; np++)
				push_neighbor(*np, &neighbors);	
		}
	}
	if(fill_counts[thread_idx])
		printf("Finished fill with %lu pixels\n", fill_counts[thread_idx]);
//	pthread_mutex_unlock(&m);
	return fill_counts[thread_idx];
}

size_t un_flood_fill(size_t idx){
	size_t neighbors[] = {idx + 1, idx - 1, idx + width, idx - width, 0};
	size_t fill_size = 0;
	for(size_t *np = neighbors; *np; np++){
		if(*np > width*height)
			continue;
		if(output[*np * 3 + 2] != 254 && image[*np] < 175){
			output[*np * 3 + 2] = 254;
			output[*np * 3 + 1] = 254;
			fill_size++;
			fill_size += un_flood_fill(*np);
		}
	}
	return fill_size;
}

void print_square(size_t i){
	for(size_t row = 0; row < 20; row++) {
		for(size_t col = 0; col < 20; col++) {
			output[3 * (i + col + width * row) + 2] = 255;
			output[3 * (i + col + width * row) + 1] = 0;
		}
	}
}

struct limits {
	size_t start;
	size_t stop;
	size_t idx;
};
size_t flood_fills_done = 0;
void* apply_filter(void* p){
	struct limits *l = p;
//	printf("Starting at %lu, stopping at %lu\n", l->start, l->stop);
	size_t fill_size;
	for(size_t i = l->start; i < l->stop; i += 1){
		if(output[3*i + 2] != 255)
			output[3*i + 1] = image[i];
//		output[3*i + 2] = image[i];
		if(image[i] < 175) { // This condition was backwards for some reason
			if(fill_size = flood_fill(i, l->idx)){
				if(fill_size > 700) {
					flood_fills_done++;
					print_square(i);
					printf("Found object at index %lu (%lu, %lu)\n", i, i%width, i/width);
				}
			}
		} else {
			output[3*i + 2] = 200;
		}
		
	}

}

double timediff(struct timeval* start, struct timeval* end) {
	double s_diff = end->tv_sec - start->tv_sec;
	double us_diff = end->tv_usec - start->tv_usec;
	s_diff += us_diff / 1000000;
	return s_diff;
}

int main(int argc, char ** argv){
	image = stbi_load("nutsandbolts.jpg", &width, &height, 0, 1); 
	array_length = width * height * 1; // 1 channel
	if(!image)
		puts(RED("No image, giving up!"));
	else
		printf(PURPLE("Image size:  %d by %d\n"), width, height);

	output = malloc(array_length * 3);
	memset(output, 0, array_length * 3);

	struct timeval start, end;
	gettimeofday(&start, NULL);
	int thread_count = 1;
	if(argc == 2)
		thread_count = atoi(argv[1]);
	pthread_t threads[thread_count];
	struct limits parameters[thread_count];
	for(int i = 0; i < thread_count; i++){
		parameters[i].start = i * (array_length / thread_count);
		parameters[i].stop = (i + 1) * (array_length / thread_count);
		parameters[i].idx = i;
		pthread_create(&threads[i], 0, apply_filter, &parameters[i]);
	}
	for(int i = 0; i < thread_count; i++){
		pthread_join(threads[i], 0);
	}
	gettimeofday(&end, NULL);
	printf("Counting time: %lf\n",  timediff(&start, &end));		
	printf("Found %d objects\n", flood_fills_done);

	puts(GREEN("Done applying filter, saving image as done.png"));
	stbi_write_png("done.png", width, height, 3, output, 0);
	puts(GREEN("Saved image as done.png"));
	free(output);
	return 0;
}