#include <lib/seg/util.h>
#include <lib/algo/flood_fill.h>
#include <tiffio.h>
#include <assert.h>
#include <stdio.h>
#include <stdlib.h>

// Convert x,y coords to linear coordinate
size_t xy_to_coord(size_t x, size_t y, uint32_t width, uint32_t height)
{
  return x + y*width;
}

// Dilate masks by one 4-connected pixel
uint16_t* dilate(uint16_t* mask, uint32_t width, uint32_t height)
{
  uint16_t *new_mask = (uint16_t*)calloc(width*height,sizeof(uint16_t));
  for (size_t y = 0; y < height; y++) {
    for (size_t x = 0; x < width; x++) {
      size_t current_position = xy_to_coord(x, y, width, height);
      if (mask[current_position] != 0) {
	new_mask[current_position] = mask[current_position];
	continue;
      }
      size_t proposed_position;
      if (x != 0) {
	proposed_position = xy_to_coord(x-1, y, width, height);
	if (mask[proposed_position] != 0) {
	  new_mask[current_position] = mask[proposed_position];
	  continue;
	}
      }
      if ((x+1) != width) {
	proposed_position = xy_to_coord(x+1, y, width, height);
	if (mask[proposed_position] != 0) {
	  new_mask[current_position] = mask[proposed_position];
	  continue;
	}
      }
      if (y != 0) {
	proposed_position = xy_to_coord(x, y-1, width, height);
	if (mask[proposed_position] != 0) {
	  new_mask[current_position] = mask[proposed_position];
	  continue;
	}
      }
      if ((y+1) != height) {
	proposed_position = xy_to_coord(x, y+1, width, height);
	if (mask[proposed_position] != 0) {
	  new_mask[current_position] = mask[proposed_position];
	  continue;
	}
      }
    }
  }
  return new_mask;
}

// Combine Label Masks
// For all empty spaces in the destination, put the extra label if it exists
// Allocates an array if destination is unallocated
uint16_t* combine_masks(uint16_t *destination, uint16_t *extra_labels, uint32_t width, uint32_t height)
{
  if (destination == NULL) {
    destination = (uint16_t*)calloc(width*height, sizeof(uint16_t));
  }
  for (size_t y = 0; y < height; y++) {
    for (size_t x = 0; x < width; x++) {
      size_t coord = x + y*width;
      if (destination[coord] == 0) {
	destination[coord] = extra_labels[coord];
      }
    }
  }
  return destination;
}

// Process Tif File to Labels
//  width, height will be overwritten with image dimensions
//  starting_label_p will be incremented for each label found in the image
uint16_t* tif_to_labels(char* tif_file_name, uint32_t *width, uint32_t *height, uint16_t *starting_label_p)
{
  //-TIFF-IMAGE-OPEN-------------------------------
  TIFF *tif = TIFFOpen(tif_file_name, "r");
  if (!tif) {
    fprintf(stderr, "Failed to open TIFF file\n");
    return NULL;
  }

  //-TIFF-FIND-DIMENSIONS--------------------------
  size_t channels = 1;
  TIFFGetField(tif, TIFFTAG_IMAGEWIDTH, width);
  TIFFGetField(tif, TIFFTAG_IMAGELENGTH, height);

  tmsize_t STRIP_LENGTH = TIFFStripSize(tif);
  tmsize_t STRIP_COUNT = TIFFNumberOfStrips(tif);
  if ((*width)*(*height)*3 == STRIP_LENGTH*STRIP_COUNT) {
    channels = 3;
  } else if ((*width)*(*height)*4 == STRIP_LENGTH*STRIP_COUNT) {
    channels = 4;
  }

  //-TIFF-LOAD-DATA--------------------------------
  void* buffer = malloc(STRIP_LENGTH*sizeof(uint8_t));
  if (buffer == NULL) {
    fprintf(stderr, "Memory allocation error\n");
    TIFFClose(tif);
    return NULL;
  }
  uint8_t* image_data = calloc((*width)*(*height)*channels,sizeof(uint8_t));
  if (image_data == NULL) {
    fprintf(stderr, "Memory allocation error\n");
    free(buffer);
    TIFFClose(tif);
    return NULL;
  }
  for (size_t y = 0; y < STRIP_COUNT; y++) {
    tmsize_t strip_size = TIFFReadRawStrip(tif, y, buffer, STRIP_LENGTH);
    assert(strip_size == STRIP_LENGTH);
    for (size_t x = 0; x < STRIP_LENGTH; x++) {
      image_data[x + y*STRIP_LENGTH] = ((uint8_t*)buffer)[x];
    }
  }
  free(buffer);

  //-FLOOD-FILL-SEGMENTATION-----------------------
  //-CONTIGUOUS-REGION-FINDING---------------------
  uint16_t *labels = NULL;
  labels = (uint16_t*)calloc((*width)*(*height),sizeof(uint16_t));
  if (labels == NULL) {
    fprintf(stderr, "Memory allocation error\n");
    free(image_data);
    TIFFClose(tif);
    return NULL;
  }
  // Flood fill on each pixel
  //  Increase label for each success
  for (size_t y = 0; y < *height; y++) {
    for (size_t x = 0; x < *width; x++) {
      size_t coord = x + y*(*width);
      if(flood(image_data, labels, *width, *height, channels, x, y, &(image_data[coord*channels]), *starting_label_p)) {
        *starting_label_p += 1;
      }
    }
  }
  free(image_data);
  TIFFClose(tif);
  return labels;
}