#include "heap.h"

#include "../core/memory.h"

typedef struct BzHeapHead {
    i32 capacity;
    i32 size;
    i32 stride;
    i32 weightOffset;
} BzHeapHead;

#define HEAP_HEAD(heap) (((BzHeapHead *)heap) - 1)

#define HEAP_LEFT(i) ((i32)(i * 2 + 1))
#define HEAP_RIGHT(i) ((i32)(i * 2 + 2))
#define HEAP_PARENT(i) ((i32) ((i - 1) / 2))

static void heapSiftUp(BzHeapHead *head, void *heap, i32 idx);
static void heapSiftDown(BzHeapHead *head, void *heap, i32 idx);

void *_bzHeapCreate(i32 startCapacity, i32 stride, i32 weightOffset) {
    i32 numBytes = sizeof(BzHeapHead) + (startCapacity + 1) * stride;
    BzHeapHead *heap = bzAlloc(numBytes);
    heap[0] = (BzHeapHead) {
        .capacity=startCapacity,
        .size=0,
        .stride=stride,
        .weightOffset=weightOffset
    };

    return &heap[1];
}

void _bzHeapDestroy(void *heap) {
    bzFree(HEAP_HEAD(heap));
}
void _bzHeapClear(void *heap) {
    BzHeapHead *head = HEAP_HEAD(heap);
    head->size = 0;
}

i32 _bzHeapSize(void *heap) {
    return HEAP_HEAD(heap)->size;
}
bool _bzHeapIsEmpty(void *heap) {
    return HEAP_HEAD(heap)->size == 0;

}
i32 _bzHeapPop(void *heap) {
    BzHeapHead *head = HEAP_HEAD(heap);
    BZ_ASSERT(head->size > 0);

    // Move first to index capacity (for output)
    bzMemMove(((u8 *)heap) + head->capacity * head->stride, heap, head->stride);
    head->size--;
    if (head->size == 0) return 0;

    // Move last to index 0
    bzMemMove(heap, ((u8 *) heap) + head->size * head->stride, head->stride);

    heapSiftDown(head, heap, 0);

    return head->capacity;
}
void _bzHeapPush(void *heap) {
    BzHeapHead *head = HEAP_HEAD(heap);
    BZ_ASSERT(head->size < head->capacity);

    void *item = ((u8 *)heap) + head->size * head->stride;
    bzMemMove(((u8 *)heap) + head->size * head->stride, item, head->stride);

    head->size++;
    heapSiftUp(head, heap, head->size - 1);
}
i32 _bzHeapPushIdx(void *heap) {
    BzHeapHead *head = HEAP_HEAD(heap);
    BZ_ASSERT(head->size < head->capacity);
    return head->size;
}
void _bzHeapUpdate(void *heap, i32 idx) {
    BzHeapHead *head = HEAP_HEAD(heap);
    BZ_ASSERT(idx >= 0 && idx < head->size);

    heapSiftUp(head, heap, idx);
    heapSiftDown(head, heap, idx);
}

static void heapSwap(BzHeapHead *head, void *heap, i32 aIdx, i32 bIdx) {
    u8 *aItem = ((u8 *)heap) + aIdx * head->stride;
    u8 *bItem = ((u8 *)heap) + bIdx * head->stride;

    u8 tmp[head->stride];
    bzMemMove(tmp, aItem, head->stride);

    bzMemMove(aItem, bItem, head->stride);
    bzMemMove(bItem, tmp, head->stride);
}
static int heapCmp(BzHeapHead *head, void *heap, i32 lhs, i32 rhs) {
    int *aWeight = (i32 *) (((u8 *)heap) + lhs * head->stride + head->weightOffset);
    int *bWeight = (i32 *) (((u8 *)heap) + rhs * head->stride + head->weightOffset);
    return (*aWeight) - (*bWeight);
}
static void heapSiftUp(BzHeapHead *head, void *heap, i32 idx) {
    while (idx >= 0) {
        i32 parent = HEAP_PARENT(idx);
        if (heapCmp(head, heap, idx, parent) >= 0)
            break;
        heapSwap(head, heap, idx, parent);
        idx = parent;
    }

}
static void heapSiftDown(BzHeapHead *head, void *heap, i32 idx) {
    while (idx < head->size) {
        i32 l = HEAP_LEFT(idx);
        i32 r = HEAP_RIGHT(idx);
        i32 smallest = idx;
        if (l < head->size && heapCmp(head, heap, l, idx) < 0)
            smallest = l;
        if (r < head->size && heapCmp(head, heap, r, smallest) < 0)
            smallest = r;

        if (smallest == idx) break;

        heapSwap(head, heap, idx, smallest);
        idx = smallest;
    }
}