#ifdef NERV_GENERIC_MMATRIX
#include "../matrix.h"
#include "../../lib/matrix/generic/matrix.h"
#include "../../lib/matrix/generic/elem_type.h"
#define MATRIX_DATA_WRITE(L, data, idx, val) (data[idx] = val)
#define MATRIX_DATA_READ(L, data, idx) (data[idx])
#define MATRIX_INIT(L) host_matrix_(init)(L)
#define MATRIX_BASE_TNAME nerv_matrix_host_tname
#define NERV_GENERIC_MATRIX
#include "../../lib/common.h"
#include "../../lib/cblas.h"
#include "../../lib/matrix/generic/mmatrix.h"
#include "../../io/chunk_file.h"
#include <string.h>

#define BLAS_OP_N CblasNoTrans
static int nerv_matrix_(lua_get_blas_op)(char ch) {
    return (ch == 'T' || ch == 't') ? CblasTrans : CblasNoTrans;
}

static int nerv_matrix_(lua_get_elem)(lua_State *L) {
    Matrix *self = luaT_checkudata(L, 1, nerv_matrix_(tname));
    int idx = luaL_checkinteger(L, 2);
    if (idx < 0 || idx >= self->nmax)
        nerv_error(L, "index must be within range [0, %d)", self->nmax);
    lua_pushnumber(L, MATRIX_ELEM_PTR(self)[idx]);
    return 1;
}

static int nerv_matrix_(lua_set_elem)(lua_State *L) {
    Matrix *self = luaT_checkudata(L, 1, nerv_matrix_(tname));
    int idx = luaL_checkinteger(L, 2);
    MATRIX_ELEM v = luaL_checknumber(L, 3);
    if (idx < 0 || idx >= self->nmax)
        nerv_error(L, "index must be within range [0, %d)", self->nmax);
    MATRIX_ELEM_PTR(self)[idx] = v;
    return 0;
}

static const luaL_Reg nerv_matrix_(extra_methods)[];
static void host_matrix_(init)(lua_State *L) {
    luaN_append_methods(L, nerv_matrix_(extra_methods));
#ifdef MMATRIX_INIT
    MMATRIX_INIT(L);
#endif
}

#include "matrix.c"

static int nerv_matrix_(lua_load)(lua_State *L) {
    Status status;
    MATRIX_CONTEXT *context;
    MATRIX_GET_CONTEXT(L, 2);
    ChunkData *cdp = luaT_checkudata(L, 1, nerv_chunk_data_tname);
    Matrix *self = nerv_matrix_(load)(cdp, context, &status);
    NERV_LUA_CHECK_STATUS(L, status);
    luaT_pushudata(L, self, nerv_matrix_(tname));
    return 1;
}

static int nerv_matrix_(lua_save)(lua_State *L) {
    Status status;
    MATRIX_CONTEXT *context;
    MATRIX_GET_CONTEXT(L, 3);
    ChunkFile *cfp = luaT_checkudata(L, 2,
                            nerv_chunk_file_handle_tname);
    Matrix *self = luaT_checkudata(L, 1, nerv_matrix_(tname));
    nerv_matrix_(save)(self, cfp, context, &status);
    NERV_LUA_CHECK_STATUS(L, status);
    return 0;
}

static int nerv_matrix_(lua_copy_fromh)(lua_State *L) {
    Status status;
    MATRIX_CONTEXT *context;
    MATRIX_GET_CONTEXT(L, 6);
    Matrix *a = luaT_checkudata(L, 1, nerv_matrix_(tname));
    const Matrix *b = luaT_checkudata(L, 2, nerv_matrix_(tname));
    int nargs = lua_gettop(L);
    int b_begin = nargs > 2 ? luaL_checkinteger(L, 3) : 0;
    int b_end = nargs > 3 ? luaL_checkinteger(L, 4) : b->nrow;
    int a_begin = nargs > 4 ? luaL_checkinteger(L, 5) : 0;
    nerv_matrix_(copy_fromh)(a, b, a_begin, b_begin, b_end, context, &status);
    NERV_LUA_CHECK_STATUS(L, status);
    return 0;
}

static int nerv_matrix_(lua_copy_rows_fromh_by_idx)(lua_State *L)
{
    Status status;
    MATRIX_CONTEXT *context;
    MATRIX_GET_CONTEXT(L, 5);
    Matrix *a = luaT_checkudata(L, 1, nerv_matrix_(tname));
    const Matrix *b = luaT_checkudata(L, 2, nerv_matrix_(tname));
    const Matrix *idx = luaT_checkudata(L, 3, nerv_matrix_(tname));
    int b_begin = lua_gettop(L) > 3 ? luaL_checkinteger(L, 4) : 0;
    nerv_matrix_(copy_rows_fromh_by_idx)(a, b, idx, b_begin, context, &status);
    NERV_LUA_CHECK_STATUS(L, status);
    return 0;
}

static const luaL_Reg nerv_matrix_(extra_methods)[] = {
    {"colsum", nerv_matrix_(lua_colsum)},
    {"colsame", nerv_matrix_(lua_colsame)},
    {"rowsum", nerv_matrix_(lua_rowsum)},
    {"rowmax", nerv_matrix_(lua_rowmax)},
    {"rowmax_idx", nerv_matrix_(lua_rowmax_idx)},
    {"trans", nerv_matrix_(lua_trans)},
    {"decompress", nerv_matrix_(lua_decompress)},
    /* in-place calc */
    {"copy_fromh", nerv_matrix_(lua_copy_fromh)},
    /* alias for copy_from */
    {"copy_from", nerv_matrix_(lua_copy_fromh)},
    {"add", nerv_matrix_(lua_add)},
    {"mul", nerv_matrix_(lua_mul)},
    {"add_row", nerv_matrix_(lua_add_row)},
    {"clip", nerv_matrix_(lua_clip)},
    {"fill", nerv_matrix_(lua_fill)},
    {"diagonalize", nerv_matrix_(lua_diagonalize)},
    {"set_values_by_mask", nerv_matrix_(lua_set_values_by_mask)},
    {"sigmoid", nerv_matrix_(lua_sigmoid)},
    {"sigmoid_grad", nerv_matrix_(lua_sigmoid_grad)},
    {"tanh", nerv_matrix_(lua_tanh)},
    {"tanh_grad", nerv_matrix_(lua_tanh_grad)},
    {"relu", nerv_matrix_(lua_relu)},
    {"relu_grad", nerv_matrix_(lua_relu_grad)},
    {"softmax", nerv_matrix_(lua_softmax)},
    {"mul_elem", nerv_matrix_(lua_mul_elem)},
    {"log_elem", nerv_matrix_(lua_log_elem)},
    {"copy_rows_fromh_by_idx", nerv_matrix_(lua_copy_rows_fromh_by_idx)},
    {"copy_rows_from_by_idx", nerv_matrix_(lua_copy_rows_fromh_by_idx)},
    {"expand_frm", nerv_matrix_(lua_expand_frm)},
    {"rearrange_frm", nerv_matrix_(lua_rearrange_frm)},
    {"scale_rows_by_row", nerv_matrix_(lua_scale_rows_by_row)},
    {"scale_rows_by_col", nerv_matrix_(lua_scale_rows_by_col)},
    {"load", nerv_matrix_(lua_load)},
    {"save", nerv_matrix_(lua_save)},
    {NULL, NULL}
};

#endif