#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_tree *tree) {
  uint64_t size = 1;
  vieter_heap_node **stack =
      malloc(((uint64_t)1 << tree->order) * sizeof(vieter_heap_node *));
  stack[0] = tree->root;

  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->next_largest_order != NULL) {
      stack[size] = node->next_largest_order;
      size++;
    }

    vieter_heap_node_free(node);
  }

  free(stack);
  free(tree);
}

vieter_heap_tree *vieter_heap_tree_init(vieter_heap_node *root,
                                        vieter_heap_tree *next, uint8_t order) {
  vieter_heap_tree *tree = malloc(sizeof(vieter_heap_tree));

  tree->root = root;
  tree->next = next;
  tree->order = order;

  return tree;
}

void vieter_heap_tree_swap(vieter_heap_tree *t1, vieter_heap_tree *t2) {
  vieter_heap_tree temp = {
      .order = t1->order, .root = t1->root, .next = t1->next};

  t1->order = t2->order;
  t1->root = t2->root;
  t1->next = t2->next;

  t2->order = temp.order;
  t2->root = temp.root;
  t2->next = temp.next;
}

vieter_heap_tree *vieter_heap_tree_merge_same_order(vieter_heap_tree *tree_a,
                                                    vieter_heap_tree *tree_b) {
  vieter_heap_tree *new_tree;

  if (tree_a->root->key <= tree_b->root->key) {
    new_tree = tree_a;
    tree_a->root->next_largest_order = tree_a->root->largest_order;
    tree_a->root->largest_order = tree_b->root;

    free(tree_b);
  } else {
    new_tree = tree_b;
    tree_b->root->next_largest_order = tree_b->root->largest_order;
    tree_b->root->largest_order = tree_a->root;

    free(tree_a);
  }

  new_tree->order++;

  return new_tree;
}

vieter_heap_tree *vieter_heap_tree_merge(vieter_heap_tree *tree_a,
                                         vieter_heap_tree *tree_b) {
  vieter_heap_tree *tree, *target, *out;

  if (tree_a->order <= tree_b->order) {
    target = tree_a;
    tree = tree_b;
  } else {
    target = tree_b;
    tree = tree_a;
  }

  vieter_heap_tree *next_tree, *next_target;
  vieter_heap_tree *previous_target = NULL;

  while (target != NULL && tree != NULL) {
    if (target->order == tree->order) {
      next_tree = tree->next;
      next_target = target->next;

      target = vieter_heap_tree_merge_same_order(target, tree);

      target->next = next_target;

      // 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 (target->next != NULL && target->next->order == target->order) {
        next_target = target->next->next;
        target = vieter_heap_tree_merge_same_order(target, target->next);
        target->next = next_target;
      }

      if (previous_target != NULL) {
        previous_target->next = target;
      } else {
        out = target;
      }

      tree = next_tree;
    } else if (target->order > tree->order) {
      next_tree = tree->next;

      if (previous_target == NULL) {
        previous_target = tree;
        out = tree;
      } else {
        previous_target->next = tree;
      }

      tree->next = target;
      tree = next_tree;
    } else {
      if (previous_target == NULL) {
        out = target;
      }

      previous_target = target;
      target = target->next;
    }
  }

  // Append final part of tree to target
  if (target == NULL) {
    previous_target->next = tree;
  }

  return out;
}

vieter_heap_tree *vieter_heap_tree_pop(void **out, vieter_heap_tree *tree) {
  vieter_heap_tree *tree_before_smallest = NULL;
  vieter_heap_tree *previous_tree = NULL;
  vieter_heap_tree *tree_out = tree;

  uint64_t smallest_key = tree->root->key;

  while (tree->next != NULL) {
    previous_tree = tree;
    tree = tree->next;

    if (tree->root->key < smallest_key) {
      smallest_key = tree->root->key;
      *out = tree->root->data;
      tree_before_smallest = previous_tree;
    }
  }

  if (tree_before_smallest != NULL) {
    previous_tree->next = tree->next;
  } else {
    tree_out = tree_out->next;
  }

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

    return NULL;
  }

  uint8_t old_order = tree->order;

  vieter_heap_node *node = tree->root->largest_order;
  free(tree->root);

  previous_tree = vieter_heap_tree_init(node, NULL, old_order - 1);

  uint8_t i = 2;
  while (node->next_largest_order != NULL) {
    node = node->next_largest_order;
    previous_tree = vieter_heap_tree_init(node, previous_tree, old_order - i);

    i++;
  }

  return vieter_heap_tree_merge(tree_out, previous_tree);
}