#include <errno.h>
#include <stdlib.h>

#include "utils.h"
#include "xvec.h"

int xvec_init(xvec_t* v, int count, int elsize)
{
    if ((v == NULL) || (count <= 0) || (elsize <= 0)) {
        return -EINVAL;
    }

    v->data = (void*)calloc(count, elsize);
    if (v->data == NULL) {
        return -ENOMEM;
    }

    v->count  = count;
    v->elsize = elsize;

    return v->count;
}

void xvec_fini(xvec_t* v)
{
    if (v == NULL) {
        return;
    }

    FREEUP(v->data);
    memset(v, 0x0, sizeof(*v));
}

int xvec_zero(xvec_t* v)
{
    if (v == NULL) {
        return -EINVAL;
    }

    if (v->data == NULL) {
        return -EINVAL;
    }

    memset(v->data, 0x0, v->count * v->elsize);

    return 1;
}

int xvec_resize(xvec_t* v, int new_size)
{
    void* p = NULL;

    if (v == NULL) {
        return -EINVAL;
    }

    if (new_size <= v->count) {
        /* do thing for shrink now */
        return 0;
    }

    p = (void*)realloc(v->data, new_size * v->elsize);
    if (p == NULL) {
        return -ENOMEM;
    }

    v->data = p;
    /* zero fill the extended space */
    memset(((char*)p) + v->count * v->elsize, 0x0,
           (new_size - v->count) * v->elsize);
    v->count = new_size;

    return 0;
}

#define X_OP(op) X_OP_##op
#define X_OP_add +=
#define X_OP_sub -=

#define DEFINE_VEC_OP(op, typ)                                                 \
    static inline int vec_##typ##_##op(typ* v1, typ* v2, int count)            \
    {                                                                          \
        int i;                                                                 \
        for (i = 0; i < count; i++) {                                          \
            v1[i] X_OP(op) v2[i];                                              \
        }                                                                      \
        return 0;                                                              \
    }

DEFINE_VEC_OP(add, int)
DEFINE_VEC_OP(sub, int)
DEFINE_VEC_OP(add, float)
DEFINE_VEC_OP(sub, float)

int xvec_op(xvec_t* v1, xvec_t* v2, xvec_op_t op)
{
    int rc = 0;

    if ((v1 == NULL) || (v2 == NULL)) {
        return -EINVAL;
    }

    if (v1->elsize != v2->elsize) {
        return -EINVAL;
    }

    switch (op) {
    case XVEC_ADD_INT:
    case XVEC_SUB_INT:
        if (v1->elsize != sizeof(int)) {
            return -EINVAL;
        }
        break;
    case XVEC_ADD_FLOAT:
    case XVEC_SUB_FLOAT:
        if (v1->elsize != sizeof(float)) {
            return -EINVAL;
        }
        break;
    }

    if (v1->count < v2->count) {
        rc = xvec_resize(v1, v2->count);

        if (rc < 0) {
            return rc;
        }
    }

    switch (op) {
    case XVEC_ADD_INT:
        rc = vec_int_add((int*)v1->data, (int*)v2->data, v2->count);
        break;
    case XVEC_SUB_INT:
        rc = vec_int_sub((int*)v1->data, (int*)v2->data, v2->count);
        break;
    case XVEC_ADD_FLOAT:
        rc = vec_float_add((float*)v1->data, (float*)v2->data, v2->count);
        break;
    case XVEC_SUB_FLOAT:
        rc = vec_float_sub((float*)v1->data, (float*)v2->data, v2->count);
        break;
    }

    return rc;
}

int xvec_dup(xvec_t** dst, xvec_t* src)
{
    xvec_t* v = NULL;

    v    = (xvec_t*)calloc(1, sizeof(*v));
    *dst = v;

    if (v == NULL) {
        return -ENOMEM;
    }

    v->data = (void*)calloc(src->count, src->elsize);
    if (v->data == NULL) {
        return -ENOMEM;
    }

    memcpy(v->data, src->data, src->elsize * src->count);
    v->elsize = src->elsize;
    v->count  = src->count;

    return v->count;
}

int xvec_put(xvec_t* v, int idx, void* data)
{
    int   rc;
    char* p;

    if ((v == NULL) || (idx < 0) || (data == NULL)) {
        return -EINVAL;
    }

    rc = xvec_resize(v, idx + 1);
    if (rc < 0) {
        return rc;
    }

    p = (char*)v->data;
    memcpy(p + v->elsize * idx, data, v->elsize);

    return 1;
}

int xvec_get(xvec_t* v, int idx, void* data)
{
    char* p = NULL;

    if ((v == NULL) || (idx < 0) || (data == NULL)) {
        return -EINVAL;
    }

    if (idx >= v->count) {
        memset(data, 0x0, v->elsize);
        return 0;
    }

    p = (char*)v->data;
    memcpy(data, p + v->elsize * idx, v->elsize);

    return 1;
}

int xvec_idval_init(xvec_idval_t* v, int initsize)
{
    int rc;

    if ((v == NULL) || (initsize <= 0)) {
        return -EINVAL;
    }

    memset(v, 0x0, sizeof(*v));
    rc = xvec_init(&(v->idxs), initsize, sizeof(int));
    if (rc < 0) {
        return rc;
    }

    rc = xvec_init(&(v->vals), initsize, sizeof(float));
    if (rc < 0) {
        return rc;
    }

    v->count = 0;

    return 0;
}

void xvec_idval_fini(xvec_idval_t* v)
{
    if (v == NULL) {
        return;
    }

    xvec_fini(&(v->idxs));
    xvec_fini(&(v->vals));

    memset(v, 0x0, sizeof(*v));
    return;
}

int xvec_idval_store(xvec_idval_t* v, int idx, float val)
{
    int rc;
    int pos;

    if ((v == NULL) || (idx < 0)) {
        return -EINVAL;
    }

    pos = v->count;
    rc  = xvec_put(&(v->idxs), pos, &idx);
    if (rc < 0) {
        return rc;
    }

    rc = xvec_put(&(v->vals), pos, &val);
    if (rc < 0) {
        return rc;
    }

    v->count = pos + 1;

    return v->count;
}

int xvec_idval_store_idx(xvec_idval_t* v, int idx)
{
    int pos;
    int rc;

    if ((v == NULL) || (idx < 0)) {
        return -EINVAL;
    }

    pos = v->count;
    rc  = xvec_put(&(v->idxs), pos, &idx);
    if (rc < 0) {
        return rc;
    }

    v->count = pos + 1;
    return v->count;
}

int xvec_fill_idval(xvec_idval_t* v, xvec_t* src)
{
    int i;

    if ((v == NULL) || (src == NULL)) {
        return -EINVAL;
    }

    if (v->count < 0) {
        return -EINVAL;
    }

    if (src->elsize != sizeof(float)) {
        return -EINVAL;
    }

    for (i = 0; i < v->count; i++) {
        int    idx  = ((int*)v->idxs.data)[i];
        float* vals = (float*)v->vals.data;

        if (idx < 0) {
            vals[i] = 0.0;
            continue;
        }

        if (idx >= src->count) {
            vals[i] = 0.0;
        } else {
            vals[i] = ((float*)src->data)[idx];
        }
    }

    return v->count;
}

#define DEFINE_XVEC_IDVAL_OP(op)                                               \
    static inline int xvec_idval_##op(xvec_t* d, xvec_idval_t* v)              \
    {                                                                          \
        int rc, maxidx;                                                        \
        int i;                                                                 \
                                                                               \
        maxidx = 0;                                                            \
        for (i = 0; i < v->count; i++) {                                       \
            if (maxidx < ((int*)v->idxs.data)[i]) {                            \
                maxidx = ((int*)v->idxs.data)[i];                              \
            }                                                                  \
        }                                                                      \
                                                                               \
        rc = xvec_resize(d, maxidx + 1);                                       \
        if (rc < 0) {                                                          \
            return rc;                                                         \
        }                                                                      \
                                                                               \
        for (i = 0; i < v->count; i++) {                                       \
            int idx = ((int*)v->idxs.data)[i];                                 \
                                                                               \
            if (idx < 0) {                                                     \
                continue;                                                      \
            }                                                                  \
                                                                               \
            ((float*)d->data)[idx] X_OP(op)((float*)v->vals.data)[i];          \
        }                                                                      \
                                                                               \
        return v->count;                                                       \
    }

DEFINE_XVEC_IDVAL_OP(add)
DEFINE_XVEC_IDVAL_OP(sub)

int xvec_idval(xvec_t* d, xvec_idval_t* v, xvec_op_t op)
{
    if ((d == NULL) || (v == NULL)) {
        return -EINVAL;
    }

    if (d->elsize != sizeof(float)) {
        return -EINVAL;
    }

    if ((v->idxs.elsize != sizeof(int)) || (v->vals.elsize != sizeof(float))) {
        return -EINVAL;
    }

    switch (op) {
    case XVEC_ADD_INT:
    case XVEC_SUB_INT:
        return -EINVAL;
    case XVEC_ADD_FLOAT:
        return xvec_idval_add(d, v);
    case XVEC_SUB_FLOAT:
        return xvec_idval_sub(d, v);
    }

    return -EINVAL;
}