#include <assert.h>
#include <lib/algo/flood_fill.h>
#include <lib/png.h>
#include <lib/seg/util.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <tiffio.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(const Mask *mask) {
  MaskData_t *buffer = mask->image[0];
  uint32_t width = mask->width;
  uint32_t height = mask->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];
    }
  }
}