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

// Suppress Tiff Warnings
void TiffDummyHandler(const char* module, const char* fmt, va_list ap)
{
    // ignore errors and warnings (or handle them your own way)
}

// Determine if coordinate is on a mask boundary
//  Assumes mask is (WxH)
bool_t is_on_mask_boundary(Mask* mask, size_t x, size_t y)
{
  MaskData_t current_value = mask->image[y][x];

  // Left neighbor
  if (x != 0) {
    if (mask->image[y][x-1] != current_value) {
      return TRUE;
    }
  }
  // Right neighbor
  if ((x+1) != mask->width) {
    if (mask->image[y][x+1] != current_value) {
      return TRUE;
    }
  }
  if (y != 0) {
    if (mask->image[y-1][x] != current_value) {
      return TRUE;
    }
  }
  if ((y+1) != mask->height) {
    if (mask->image[y+1][x] != current_value) {
      return TRUE;
    }
  }
  return FALSE;
}

// Dilate masks by one 4-connected pixel
MaskData_t* _dilate(const Mask* mask)
{
  uint32_t width = mask->width;
  uint32_t height = mask->height;
  Mask *new_mask = create_image_mask(width, height);
  for (size_t y = 0; y < height; y++) {
    for (size_t x = 0; x < width; x++) {
      if (mask->image[y][x] != 0) {
	new_mask->image[y][x] = mask->image[y][x];
	continue;
      }
      if (x != 0) {
	if (mask->image[y][x-1] != 0) {
	  new_mask->image[y][x] = mask->image[y][x-1];
	  continue;
	}
      }
      if ((x+1) != width) {
	if (mask->image[y][x+1] != 0) {
	  new_mask->image[y][x] = mask->image[y][x+1];
	  continue;
	}
      }
      if (y != 0) {
	if (mask->image[y-1][x] != 0) {
	  new_mask->image[y][x] = mask->image[y-1][x];
	  continue;
	}
      }
      if ((y+1) != height) {
	if (mask->image[y+1][x] != 0) {
	  new_mask->image[y][x] = mask->image[y+1][x];
	  continue;
	}
      }
    }
  }
  MaskData_t *ret_mask = new_mask->image[0];
  free(new_mask->image);
  free(new_mask);
  return ret_mask;
}

// Dilate masks by one 4-connected pixel
void dilate(Mask* mask)
{
  MaskData_t *new_mask = _dilate(mask);
  if (new_mask != NULL) {
    free(mask->image[0]);
    for (size_t y = 0; y < mask->height; y++) {
      mask->image[y] = &(new_mask[y*mask->width]);
    }
  }
}

// Erode masks by one 4-connected pixel
MaskData_t* _erode(const Mask* mask)
{
  uint32_t width = mask->width;
  uint32_t height = mask->height;
  Mask *new_mask = create_image_mask(width, height);
  memcpy(new_mask->image[0], mask->image[0], width*height*sizeof(MaskData_t));
  for (size_t y = 0; y < height; y++) {
    for (size_t x = 0; x < width; x++) {
      if (x != 0) {
	if (mask->image[y][x-1] == 0) {
	  new_mask->image[y][x] = 0;
	  continue;
	}
      }
      if ((x+1) != width) {
	if (mask->image[y][x+1] == 0) {
	  new_mask->image[y][x] = 0;
	  continue;
	}
      }
      if (y != 0) {
	if (mask->image[y-1][x] == 0) {
	  new_mask->image[y][x] = 0;
	  continue;
	}
      }
      if ((y+1) != height) {
	if (mask->image[y+1][x] == 0) {
	  new_mask->image[y][x] = 0;
	  continue;
	}
      }
    }
  }
  MaskData_t *ret_mask = new_mask->image[0];
  free(new_mask->image);
  free(new_mask);
  return ret_mask;
}

// Erode masks by one 4-connected pixel
void erode(Mask* mask)
{
  MaskData_t *new_mask = _erode(mask);
  if (new_mask != NULL) {
    free(mask->image[0]);
    for (size_t y = 0; y < mask->height; y++) {
      mask->image[y] = &(new_mask[y*mask->width]);
    }
  }
}

// Close up masks by N-pixels
MaskData_t* _closeup(Mask* mask, size_t num_pixels)
{
  Mask *new_mask = create_image_mask(mask->width, mask->height);
  memcpy(new_mask->image[0], mask->image[0],
         mask->width * mask->height * sizeof(MaskData_t));
  for (size_t count = 0; count < num_pixels; count++) {
    dilate(mask);
  }
  for (size_t count = 0; count < num_pixels; count++) {
    erode(mask);
  }
  for (size_t y = 0; y < mask->height; y++) {
    for (size_t x = 0; x < mask->width; x++) {
      if (mask->image[y][x] != 0) {
        if (new_mask->image[y][x] != mask->image[y][x]) {
	  new_mask->image[y][x] = mask->image[y][x];
        }
      }
    }
  }
  MaskData_t *ret_mask = new_mask->image[0];
  free(new_mask->image);
  free(new_mask);
  return ret_mask;
}

// Close up masks by N-pixels
//  Update pointer
void closeup(Mask* mask, size_t num_pixels)
{
  MaskData_t *new_mask = _closeup(mask, num_pixels);
  if (new_mask != NULL) {
    free(mask->image[0]);
    for (size_t y = 0; y < mask->height; y++) {
      mask->image[y] = &(new_mask[y*mask->width]);
    }
  }
}

// Combine Label Masks
// For all empty spaces in the destination, put the extra label if it exists
// Allocates an array if destination is unallocated
Mask* combine_masks(Mask *destination, Mask *extra_labels)
{
  if (destination == NULL) {
    destination = create_image_mask(extra_labels->width, extra_labels->height);
  }
  for (size_t y = 0; y < destination->height; y++) {
    for (size_t x = 0; x < destination->width; x++) {
      if (destination->image[y][x] == 0) {
	destination->image[y][x] = extra_labels->image[y][x];
      }
    }
  }
  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
Mask *tif_to_labels(char *tif_file_name, MaskData_t *starting_label_p)
{
  uint32_t width = 0;
  uint32_t height = 0;
  TIFFSetWarningHandler(TiffDummyHandler);
  //-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;
  TIFFGetField(tif, TIFFTAG_IMAGEWIDTH, &width);
  TIFFGetField(tif, TIFFTAG_IMAGELENGTH, &height);

  tmsize_t STRIP_LENGTH = TIFFStripSize(tif);
  tmsize_t STRIP_COUNT = TIFFNumberOfStrips(tif);
  channels = (STRIP_LENGTH*STRIP_COUNT)/(width*height);

  //-TIFF-LOAD-DATA--------------------------------
  void* buffer = malloc(STRIP_LENGTH*sizeof(ImageData_t));
  if (buffer == NULL) {
    fprintf(stderr, "Memory allocation error\n");
    TIFFClose(tif);
    return NULL;
  }
  Image *image = create_image(width, height, channels);
  if (image == 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->image[0][0][x + y*STRIP_LENGTH] = ((ImageData_t*)buffer)[x];
    }
  }
  free(buffer);

  //-FLOOD-FILL-SEGMENTATION-----------------------
  //-CONTIGUOUS-REGION-FINDING---------------------
  Mask *im_data = create_image_mask((size_t)width, (size_t)height);
  if (im_data == NULL) {
    fprintf(stderr, "Memory allocation error\n");
    free_image(image);
    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++) {
      if(flood(image->image[0][0], im_data->image[0], width, height, channels, x, y, image->image[y][x], *starting_label_p)) {
        *starting_label_p += 1;
      }
    }
  }
  free_image(image);
  TIFFClose(tif);
  return im_data;
  //return labels;
}

// Convert mask to bitmap
Bitmap* image_mask_data_to_bitmap(MaskData_t* buffer, uint32_t width, uint32_t height)
{
  Pixel* out_buffer = (Pixel*)calloc(width*height, sizeof(Pixel));
  if (out_buffer == NULL) {
    return NULL;
  }
  Bitmap* bitmap = (Bitmap*)malloc(sizeof(Bitmap));
  if (bitmap == NULL) {
    free(out_buffer);
    return NULL;
  }
  for (size_t y = 0; y < height; y++) {
    for (size_t x = 0; x < width; x++) {
      size_t coord = x + y*width;
      ImageData_t red   = (buffer[coord] & 0xF00) >> 4*2;
      ImageData_t green = (buffer[coord] & 0x0F0) >> 4*1;
      ImageData_t blue  = (buffer[coord] & 0x00F) >> 4*0;
      out_buffer[coord].red = red | (red << 4);
      out_buffer[coord].green = green | (green << 4);
      out_buffer[coord].blue = blue | (blue << 4);
    }
  }
  bitmap->image_buffer = out_buffer;
  bitmap->width = (size_t)width;
  bitmap->height = (size_t)height;
  return bitmap;
}

// Reduce a mask to the contiguous regions
MaskData_t* _reduce_contiguous_regions(MaskData_t* masks, uint32_t width, uint32_t height, MaskData_t* total_labels)
{
  MaskData_t starting_label = 1;
  MaskData_t* new_masks = (MaskData_t*)calloc(width*height, sizeof(MaskData_t));
  if (new_masks == NULL) {
    return NULL;
  }
  for (size_t y = 0; y < height; y++) {
    for (size_t x = 0; x < width; x++) {
      size_t coord = x + y*width;
      ImageData_t channels = 2; // MaskData_t = 2*uint8_t
      if (flood((ImageData_t*)masks, new_masks, width, height, channels, x, y, &(((ImageData_t*)masks)[coord*channels]), starting_label)) {
	starting_label++;
      }
    }
  }
  if (total_labels != NULL) {
    *total_labels = starting_label;
  }
  return new_masks;
}

// Reduce a mask to the contiguous regions
//  Automatically update pointer to contiguous mask
//  Freeing previous mask
void reduce_contiguous_regions(Mask* mask, MaskData_t* total_labels)
{
  if (mask == NULL) {
    return;
  }
  MaskData_t* new_masks = _reduce_contiguous_regions(mask->image[0], mask->width, mask->height, total_labels);
  if (new_masks != NULL) {
    free(mask->image[0]);
    for (size_t y = 0; y < mask->height; y++) {
      mask->image[y] = &new_masks[y*mask->width];
    }
  }
}