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

#include <raylib.h>

#include <lib/lib.h>
#include <lib/bool.h>
#include <lib/dir.h>
#include <lib/file.h>
#include <lib/time.h>
#include <lib/color.h>
#include <lib/algo/flood_fill.h>
#include <lib/seg/util.h>

#define OFFSET 16

#define N_DILATIONS 3

struct Result {
  void* data;
  bool_t success;
};

struct MaskData {
  uint16_t label;
  size_t area;
  size_t perimeter;
};

struct MaskData* create_mask_data(uint16_t label) {
  struct MaskData *data = (struct MaskData*)malloc(sizeof(struct MaskData));
  data->label = label;
  data->area = 0;
  data->perimeter = 0;
  return data;
}

struct AVLNode {
  void* data;
  bool_t (*compare)(void*, void*);
  struct AVLNode* left;
  struct AVLNode* right;
  uint8_t height;
};

uint8_t get_height(struct AVLNode* node)
{
  if (node == NULL) {
    return 0;
  }
  return node->height;
}

bool_t compare_labels(struct MaskData* left, struct MaskData* right)
{
  return left->label < right->label;
}

struct AVLNode* create_avl_mask_node(struct MaskData* data)
{
  struct AVLNode* node = (struct AVLNode*)malloc(sizeof(struct AVLNode));
  if (node == NULL) {
    return NULL;
  }
  node->data = data;
  node->compare = (bool_t (*)(void*,void*))&compare_labels;
  node->left = NULL;
  node->right = NULL;
  node->height = 1; // Leaf initially
  return node;
}

uint8_t max_height(uint8_t a, uint8_t b)
{
  return (a > b) ? a : b;
}

ssize_t get_balance_factor(struct AVLNode* node)
{
  if (node == NULL) {
    return 0;
  }
  return get_height(node->left) - get_height(node->right);
}

struct AVLNode* right_rotate(struct AVLNode* parent)
{
  struct AVLNode* child1 = parent->left;
  struct AVLNode* child2 = child1->right;

  child1->right = parent;
  parent->left = child2;

  parent->height = max_height(get_height(parent->left), get_height(parent->right)) + 1;
  child1->height = max_height(get_height(child1->left), get_height(child1->right)) + 1;
  return child1;
}

struct AVLNode* left_rotate(struct AVLNode* parent)
{
  struct AVLNode* child1 = parent->right;
  struct AVLNode* child2 = child1->left;

  child1->left = parent;
  parent->right = child2;

  parent->height = max_height(get_height(parent->left), get_height(parent->right)) + 1;
  child1->height = max_height(get_height(child1->left), get_height(child1->right)) + 1;
  return child1;
}

struct Result insert_mask(struct AVLNode* node, struct MaskData* data)
{
  struct Result result;
  // 1. Standard BST insertion
  if (node == NULL) {
    return (struct Result) {create_avl_mask_node(data), TRUE};
  }

  struct MaskData *node_data = (struct MaskData*)node->data;
  if (node->compare(data, node_data)) {
    result = insert_mask(node->left, data);
    if (!result.success) {
      fprintf(stderr, "Failed to insert!");
      return result;
    }
    node->left = (struct AVLNode*)result.data;
  } else if (node->compare(node->data, data)) {
    result = insert_mask(node->right, data);
    if (!result.success) {
      fprintf(stderr, "Failed to insert!");
      return result;
    }
    node->right = (struct AVLNode*)result.data;
  } else {
    return (struct Result) {node, FALSE};
  }

  // 2. Update height of the ancestor node
  node->height = 1 + max_height(get_height(node->left), get_height(node->right));

  ssize_t balance = get_balance_factor(node);

  // 4. If the node becomes unbalanced

  // LeftLeft
  if ((balance > 1) && node->compare(data, node->left->data)) {
    return (struct Result) {right_rotate(node), TRUE};
  }
  // RightRight
  if ((balance < -1) && node->compare(node->right->data, data)) {
    return (struct Result) {left_rotate(node), TRUE};
  }
  // LeftRight
  if ((balance > 1) && node->compare(node->left->data, data)) {
    return (struct Result) {right_rotate(node), TRUE};
  }
  // RightLeft
  if ((balance < -1) && node->compare(data,node->right->data)) {
    return (struct Result) {left_rotate(node), TRUE};
  }
  return (struct Result) {node, TRUE};
}

struct AVLNode* insert_mask_alloc(struct AVLNode* node, uint16_t label)
{
  struct MaskData* data = create_mask_data(label);
  struct Result result = insert_mask(node, data);
  if (!result.success) {
    free(data);
  }
  return (struct AVLNode*)result.data;
}

void print_label(struct AVLNode* root) {
  if (root != NULL) {
    print_label(root->left);
    struct MaskData* data = root->data;
    printf("%d: (%zu, %zu) ", data->label, data->area, data->perimeter);
    print_label(root->right);
  }
}

void free_avl_tree(struct AVLNode* root) {
  if (root != NULL) {
    free_avl_tree(root->left);
    free_avl_tree(root->right);
    free(root);
  }
}

void free_avl_tree_nodes(struct AVLNode* root) {
  if (root != NULL) {
    free_avl_tree_nodes(root->left);
    free_avl_tree_nodes(root->right);
    if (root->data != NULL) {
      free(root->data);
    }
    free(root);
  }
}

bool_t increase_label_area(struct AVLNode* root, uint16_t label) {
  if (root == NULL) {
    return FALSE;
  }
  struct MaskData* data = (struct MaskData*)root->data;
  if (data->label == label) {
    data->area++;
  }
  else if (data->label > label) {
    increase_label_area(root->left, label);
  }
  else if (data->label < label) {
    increase_label_area(root->right, label);
  }
  return TRUE;
}

bool_t increase_label_perimeter(struct AVLNode* root, uint16_t label) {
  if (root == NULL) {
    return FALSE;
  }
  struct MaskData* data = (struct MaskData*)root->data;
  if (data->label == label) {
    data->perimeter++;
  }
  else if (data->label > label) {
    increase_label_perimeter(root->left, label);
  }
  else if (data->label < label) {
    increase_label_perimeter(root->right, label);
  }
  return TRUE;
}

int main(int argc, char** argv)
{
  struct AVLNode* root = NULL;
  root = insert_mask_alloc(root, 2);
  root = insert_mask_alloc(root, 5);
  root = insert_mask_alloc(root, 1);
  root = insert_mask_alloc(root, 3);
  root = insert_mask_alloc(root, 4);
  printf("Inorder traversal of AVL tree: ");
  print_label(root);
  printf("\n");
  free_avl_tree_nodes(root);
  //-----------------------------------------------
  //-LIST-FILES-IN-DIRECTORY-----------------------
  //-----------------------------------------------
  char** file_list = NULL;
  uint32_t width, height;
  uint16_t starting_label = 1;
  uint16_t *masks = NULL;
  if (argc > 1) {
    if (is_directory(argv[1])) {
      file_list = list_directory(argv[1]);
      if (file_list) {
	size_t index = 0;
	while (1) {
	  char* fname = file_list[index];
	  if (fname == NULL) {
	    break;
	  }
	  if (is_tif_ext(fname) == FALSE) {
	    free(file_list[index++]);
	    continue;
	  }
	  char* fpath = full_path(argv[1], fname);
	  printf("Loading %s...\n", fpath);
	  uint16_t *file_labels = tif_to_labels(fpath, &width, &height, &starting_label);
	  masks = combine_masks(masks, file_labels, width, height);
	  free(file_labels);
	  free(fpath);
	  free(file_list[index++]);
	}
	free(file_list);
      }
    }
  }
  if (masks == NULL) {
    fprintf(stderr, "No masks found!\n");
    return 1;
  }
  printf("%u labels found\n", starting_label-1);
  printf("Mask dimensions: %u %u\n", width, height);
  TIME(ts_start);
  for (uint16_t count = 0; count < N_DILATIONS; count++) {
    uint16_t *new_labels = dilate(masks, width, height);
    free(masks);
    masks = new_labels;
  }
  TIME(ts_end);
  printf("Dilation took %f ms\n", 1000*diff_time(&ts_end, &ts_start));
  //free(masks);
  //-----------------------------------------------

  //-----------------------------------------------
  //-RAYLIB-INIT
  //-----------------------------------------------
  SetTraceLogLevel(LOG_ERROR);
  SetConfigFlags(FLAG_WINDOW_RESIZABLE);
  const char* gui_title = "Image Manip - Useful for segmentations!";
  InitWindow(SCREEN_WIDTH, SCREEN_HEIGHT, gui_title);

  //-----------------------------------------------
  // (When treating this as RGBA, last bits should ensure opaque)
  //  This assumes 4096 (2^12) > labels
  for (size_t y = 0; y < height; y++) {
    for (size_t x = 0; x < width; x++) {
      /// RGBA channels: Move labels to RGB
      masks[x + y*width] <<= 4;
      masks[x + y*width] |= 0x000F;
    }
  }
  //-----------------------------------------------

  //-----------------------------------------------
  //-RAYLIB-IMAGE-STRUCTURING----------------------
  //-----------------------------------------------
  Image RaylibImage;
  RaylibImage.width  = width;
  RaylibImage.height = height;
  RaylibImage.mipmaps = 1;
  // Use Contiguous Labels
  RaylibImage.data = masks;
  RaylibImage.format = PIXELFORMAT_UNCOMPRESSED_R4G4B4A4;
  //-----------------------------------------------

  // Image to a Texture
  Texture2D RaylibTexture = LoadTextureFromImage(RaylibImage);

  // Scale the image to the viewport
  /// Source Rectangle: Original Size
  Rectangle sourceRec = { 0.0f, 0.0f, (float)width, (float)height };
  /// Destination Rectangle: Transformed Size
  Rectangle destRec = { 0.0f, 0.0f, (float)SCREEN_WIDTH, (float)(SCREEN_HEIGHT-OFFSET) };
  /// Location to begin drawing
  Vector2 origin = { (float)0, (float)-OFFSET };

  // Raylib boilerplate
  SetTargetFPS(60);
  Camera2D camera = { 0 };
  camera.zoom = 1.0f;

  // GUI Loop
  while (!WindowShouldClose()) {
    //-----------------------------------------------
    //-DRAWING---------------------------------------
    //-----------------------------------------------
    BeginDrawing();
    ClearBackground(RAYWHITE);
    BeginMode2D(camera);
    EndMode2D();
    DrawText("Image Manip", 0, 0, OFFSET, DARKGRAY);
    DrawTexturePro(RaylibTexture, sourceRec, destRec, origin, (float)0,
                   RAYWHITE);
    /*
    uint32_t x = 0x49, y = 0x4A;
    uint32_t dx = 0x69 - x, dy = 0x6E - y;
    x = (SCREEN_WIDTH*x)/width;
    y = SCREEN_HEIGHT-((SCREEN_HEIGHT-OFFSET)*y)/height;
    dx = (SCREEN_WIDTH*dx)/width;
    dy = SCREEN_HEIGHT-((SCREEN_HEIGHT-OFFSET)*dy)/height;
    DrawRectangleGradientH(x, y, dx, dy, BLUE, PURPLE);
    */
    EndDrawing();
    //-----------------------------------------------
  }

  if (masks != NULL) {
    for (size_t y = 0; y < height; y++) {
      for (size_t x = 0; x < width; x++) {
	/// Restore labels from RGBA
	masks[x + y*width] &= 0xFFF0;
	masks[x + y*width] >>= 4;
      }
    }
    write_array("../out.bin", masks, width*height*sizeof(uint16_t));
    free(masks);
  }
  CloseWindow();
  return 0;
}