#include "ds.h"

#include <assert.h>

extern int memcmp(const void *ptr1, const void *ptr2, size_t num);
extern void *memcpy(void *destination, const void *source, size_t num);

/* Utilities for index calculation */

static RT_INLINE int Parent(int i) {
    return (i - 1) / 2;
}

static RT_INLINE int Left(int i) {
    return 2 * i + 1;
}

static RT_INLINE int Right(int i) {
    return 2 * i + 2;
}

static void Swap(rt_minheap *minheap, int i, int j, void *tmpv) {
    void *vi = (char *)minheap->values + (size_t)i * minheap->value_size;
    void *vj = (char *)minheap->values + (size_t)j * minheap->value_size;
    int tk   = minheap->keys[i];
    memcpy(tmpv, vi, minheap->value_size);
    minheap->keys[i] = minheap->keys[j];
    memcpy(vi, vj, minheap->value_size);
    minheap->keys[j] = tk;
    memcpy(vj, tmpv, minheap->value_size);
}

static void Heapify(rt_minheap *minheap, int start) {
    int i    = start;
    int size = (int)minheap->capacity;

    /* FIXME: Allocate on temp-arena */
    char tmpv[256];
    assert(minheap->value_size < sizeof(tmpv) && "tmpv should be allocated on a temporary arena.");

    while (1) {
        int min = i;
        if (Left(i) < size && minheap->keys[Left(i)] < minheap->keys[min])
            min = Left(i);
        if (Right(i) < size && minheap->keys[Right(i)] < minheap->keys[min])
            min = Right(i);
        if (min == i)
            break;
        Swap(minheap, i, min, tmpv);
    }
}

static void Decrease(rt_minheap *minheap, int i, int newkey) {
    minheap->keys[i] = newkey;
    /* FIXME: Allocate on temp-arena */
    char tmpv[256];
    assert(minheap->value_size < sizeof(tmpv) && "tmpv should be allocated on a temporary arena.");
    while (i > 0 && minheap->keys[i] < minheap->keys[Parent(i)]) {
        Swap(minheap, i, Parent(i), tmpv);
        i = Parent(i);
    }
}

static void Remove(rt_minheap *minheap, int i) {
    int last = minheap->size - 1;
    if (last < 0)
        return;
    /* FIXME: Allocate on temp-arena */
    char tmpv[256];
    assert(minheap->value_size < sizeof(tmpv) && "tmpv should be allocated on a temporary arena.");
    Swap(minheap, i, last, tmpv);
    minheap->size = last;
    if (i != last) {
        if (i == 0 || minheap->keys[i] > minheap->keys[Parent(i)])
            Heapify(minheap, i);
        else
            Decrease(minheap, i, minheap->keys[i]);
    }
}

/* Takes the arrays and re-orders the values to create a minheap. */
RT_DLLEXPORT rt_minheap
rtCreateMinheap(int *keys, void *values, size_t value_size, size_t capacity, int initial_size) {
    rt_minheap minheap = {
        .keys       = keys,
        .values     = values,
        .value_size = value_size,
        .capacity   = capacity,
        .size       = initial_size,
    };
    /* Start at the last non-leaf element */
    for (int i = initial_size / 2 - 1; i >= 0; --i)
        Heapify(&minheap, i);
    return minheap;
}

/* Copies the value with the smallest key to min_value */
RT_DLLEXPORT void rtMinheapPeek(const rt_minheap *minheap, void *min_value) {
    memcpy(min_value, minheap->values, minheap->value_size);
}

/* Copies the value with the smallest key to min_value and removes it from the heap */
RT_DLLEXPORT void rtMinheapPop(rt_minheap *minheap, void *min_value) {
    rtMinheapPeek(minheap, min_value);
    Remove(minheap, 0);
}

/* Pushes a new value into the minheap */
RT_DLLEXPORT void rtMinheapPush(rt_minheap *minheap, int key, const void *value) {
    int at            = minheap->size;
    void *v           = (char *)minheap->values + (size_t)at * minheap->value_size;
    minheap->keys[at] = key;
    memcpy(v, value, minheap->value_size);
    ++minheap->size;
    Decrease(minheap, at, key);
}

/* Changes the key of an existing value, or inserts it, if it's not found.
 *
 * Uses memcmp if cmp is NULL. */
RT_DLLEXPORT void
rtMinheapUpdate(rt_minheap *minheap, const void *value, int new_key, rt_minheap_cmp_fn *cmp) {
    if (!cmp)
        cmp = memcmp;
    for (int i = 0; i < minheap->size; ++i) {
        void *v = (char *)minheap->values + (size_t)i * minheap->value_size;
        if (cmp(v, value, minheap->value_size) == 0) {
            if (minheap->keys[i] > new_key) {
                Decrease(minheap, i, new_key);
            } else if (minheap->keys[i] < new_key) {
                Remove(minheap, i);
                rtMinheapPush(minheap, new_key, value);
            }
            return;
        }
    }
    rtMinheapPush(minheap, new_key, value);
}