#include "vieter_heap_tree.h"

vieter_heap_node *vieter_heap_node_init() {
  return calloc(1, sizeof(vieter_heap_node));
}

void vieter_heap_node_free(vieter_heap_node *node) { free(node); }

void vieter_heap_tree_free(vieter_heap_node *root) {
  if (root->order == 0) {
    goto end;
  }

  uint64_t size = 1;
  vieter_heap_node **stack =
      malloc(((uint64_t)1 << root->order) * sizeof(vieter_heap_node *));
  stack[0] = root->largest_order;

  vieter_heap_node *node;

  while (size > 0) {
    node = stack[size - 1];
    size--;

    if (node->largest_order != NULL) {
      stack[size] = node->largest_order;
      size++;
    }

    if (node->ptr.next_largest_order != NULL) {
      stack[size] = node->ptr.next_largest_order;
      size++;
    }

    vieter_heap_node_free(node);
  }

  free(stack);

end:
  vieter_heap_node_free(root);
}

vieter_heap_node *vieter_heap_tree_merge_same_order(vieter_heap_node *root_a,
                                                    vieter_heap_node *root_b) {
  vieter_heap_node *root, *child;

  if (root_a->key <= root_b->key) {
    root = root_a;
    child = root_b;
  } else {
    root = root_b;
    child = root_a;
  }

  child->ptr.next_largest_order = root->largest_order;
  root->largest_order = child;

  root->order++;

  return root;
}

vieter_heap_node *vieter_heap_tree_merge(vieter_heap_node *target_tree,
                                         vieter_heap_node *other_tree) {
  vieter_heap_node *out = target_tree;

  vieter_heap_node *next_other_tree, *next_target_tree;
  vieter_heap_node *previous_target_tree = NULL;

  while (target_tree != NULL && other_tree != NULL) {
    if (target_tree->order == other_tree->order) {
      next_other_tree = other_tree->ptr.next_tree;
      next_target_tree = target_tree->ptr.next_tree;

      target_tree = vieter_heap_tree_merge_same_order(target_tree, other_tree);

      target_tree->ptr.next_tree = next_target_tree;

      // If this merge produces a binomial tree whose size is already in
      // target, it will be the next target. Therefore, we can  merge target's
      // trees until we no longer have a duplicate depth.
      while (next_target_tree != NULL &&
             next_target_tree->order == target_tree->order) {
        next_target_tree = next_target_tree->ptr.next_tree;
        target_tree = vieter_heap_tree_merge_same_order(
            target_tree, target_tree->ptr.next_tree);
        target_tree->ptr.next_tree = next_target_tree;
      }

      if (previous_target_tree != NULL) {
        previous_target_tree->ptr.next_tree = target_tree;
      } else {
        out = target_tree;
      }

      other_tree = next_other_tree;
    } else if (target_tree->order > other_tree->order) {
      next_other_tree = other_tree->ptr.next_tree;

      if (previous_target_tree == NULL) {
        previous_target_tree = other_tree;
        out = other_tree;
      } else {
        previous_target_tree->ptr.next_tree = other_tree;

        // This single missing line right here broke this entire function for
        // nearly a week.
        previous_target_tree = other_tree;
      }

      other_tree->ptr.next_tree = target_tree;
      other_tree = next_other_tree;
    } else {
      if (previous_target_tree == NULL) {
        out = target_tree;
      }

      previous_target_tree = target_tree;
      target_tree = target_tree->ptr.next_tree;
    }
  }

  // Append final part of tree to target
  if (target_tree == NULL) {
    previous_target_tree->ptr.next_tree = other_tree;
  }

  return out;
}

vieter_heap_node *vieter_heap_tree_pop(void **out, vieter_heap_node *tree) {
  vieter_heap_node *tree_before_smallest = NULL;
  vieter_heap_node *previous_tree = NULL;
  vieter_heap_node *original_root = tree;

  uint64_t smallest_key = tree->key;

  while (tree->ptr.next_tree != NULL) {
    previous_tree = tree;
    tree = tree->ptr.next_tree;

    if (tree->key < smallest_key) {
      smallest_key = tree->key;
      tree_before_smallest = previous_tree;
    }
  }

  vieter_heap_node *tree_to_pop;

  if (tree_before_smallest != NULL) {
    tree_to_pop = tree_before_smallest->ptr.next_tree;
    tree_before_smallest->ptr.next_tree = tree_to_pop->ptr.next_tree;
  } else {
    tree_to_pop = original_root;
    original_root = original_root->ptr.next_tree;
  }

  *out = tree_to_pop->data;

  if (tree_to_pop->order == 0) {
    vieter_heap_tree_free(tree_to_pop);

    return original_root;
  }

  // Each child has a pointer to its sibling with the next largest order. If we
  // want to convert this list of children into their own tree, these pointers
  // have to be reversed.
  previous_tree = tree_to_pop->largest_order;
  vieter_heap_node_free(tree_to_pop);

  tree = previous_tree->ptr.next_largest_order;
  previous_tree->ptr.next_tree = NULL;

  vieter_heap_node *next_tree;

  while (tree != NULL) {
    next_tree = tree->ptr.next_largest_order;
    tree->ptr.next_tree = previous_tree;

    previous_tree = tree;
    tree = next_tree;
  }

  // original_root is zero if the heap only contained a single tree.
  if (original_root != NULL) {
    return vieter_heap_tree_merge(original_root, previous_tree);
  } else {
    return previous_tree;
  }
}