/* -*- c -*- */

/*
 *****************************************************************************
 **                            INCLUDES                                     **
 *****************************************************************************
 */
#define NPY_NO_DEPRECATED_API NPY_API_VERSION

#include "Python.h"
#include "numpy/arrayobject.h"
#include "numpy/ufuncobject.h"
#include "numpy/npy_math.h"

#include "npy_pycompat.h"

#include "npy_config.h"

/*
 *****************************************************************************
 **                            BASICS                                       **
 *****************************************************************************
 */

#define INIT_OUTER_LOOP_1       \
    npy_intp dN = *dimensions++;\
    npy_intp N_;                \
    npy_intp s0 = *steps++;

#define INIT_OUTER_LOOP_2       \
    INIT_OUTER_LOOP_1           \
    npy_intp s1 = *steps++;

#define INIT_OUTER_LOOP_3       \
    INIT_OUTER_LOOP_2           \
    npy_intp s2 = *steps++;

#define INIT_OUTER_LOOP_4       \
    INIT_OUTER_LOOP_3           \
    npy_intp s3 = *steps++;

#define BEGIN_OUTER_LOOP_2      \
    for (N_ = 0; N_ < dN; N_++, args[0] += s0, args[1] += s1) {

#define BEGIN_OUTER_LOOP_3      \
    for (N_ = 0; N_ < dN; N_++, args[0] += s0, args[1] += s1, args[2] += s2) {

#define BEGIN_OUTER_LOOP_4      \
    for (N_ = 0; N_ < dN; N_++, args[0] += s0, args[1] += s1, args[2] += s2, args[3] += s3) {

#define END_OUTER_LOOP  }


/*
 *****************************************************************************
 **                             UFUNC LOOPS                                 **
 *****************************************************************************
 */

char *inner1d_signature = "(i),(i)->()";

/**begin repeat

   #TYPE=LONG,DOUBLE#
   #typ=npy_long,npy_double#
*/

/*
 *  This implements the function
 *        out[n] = sum_i { in1[n, i] * in2[n, i] }.
 */

static void
@TYPE@_inner1d(char **args, npy_intp *dimensions, npy_intp *steps, void *NPY_UNUSED(func))
{
    INIT_OUTER_LOOP_3
    npy_intp di = dimensions[0];
    npy_intp i;
    npy_intp is1=steps[0], is2=steps[1];
    BEGIN_OUTER_LOOP_3
        char *ip1=args[0], *ip2=args[1], *op=args[2];
        @typ@ sum = 0;
        for (i = 0; i < di; i++) {
            sum += (*(@typ@ *)ip1) * (*(@typ@ *)ip2);
            ip1 += is1;
            ip2 += is2;
        }
        *(@typ@ *)op = sum;
    END_OUTER_LOOP
}

/**end repeat**/

char *innerwt_signature = "(i),(i),(i)->()";

/**begin repeat

   #TYPE=LONG,DOUBLE#
   #typ=npy_long,npy_double#
*/


/*
 *  This implements the function
 *        out[n] = sum_i { in1[n, i] * in2[n, i] * in3[n, i] }.
 */

static void
@TYPE@_innerwt(char **args, npy_intp *dimensions, npy_intp *steps, void *NPY_UNUSED(func))
{
    INIT_OUTER_LOOP_4
    npy_intp di = dimensions[0];
    npy_intp i;
    npy_intp is1=steps[0], is2=steps[1], is3=steps[2];
    BEGIN_OUTER_LOOP_4
        char *ip1=args[0], *ip2=args[1], *ip3=args[2], *op=args[3];
        @typ@ sum = 0;
        for (i = 0; i < di; i++) {
            sum += (*(@typ@ *)ip1) * (*(@typ@ *)ip2) * (*(@typ@ *)ip3);
            ip1 += is1;
            ip2 += is2;
            ip3 += is3;
        }
        *(@typ@ *)op = sum;
    END_OUTER_LOOP
}

/**end repeat**/

char *matrix_multiply_signature = "(m,n),(n,p)->(m,p)";

/**begin repeat

   #TYPE=FLOAT,DOUBLE,LONG#
   #typ=npy_float,npy_double,npy_long#
*/

/*
 *  This implements the function
 *        out[k, m, p] = sum_n { in1[k, m, n] * in2[k, n, p] }.
 */

static void
@TYPE@_matrix_multiply(char **args, npy_intp *dimensions, npy_intp *steps, void *NPY_UNUSED(func))
{
    /* no BLAS is available */
    INIT_OUTER_LOOP_3
    npy_intp dm = dimensions[0];
    npy_intp dn = dimensions[1];
    npy_intp dp = dimensions[2];
    npy_intp m,n,p;
    npy_intp is1_m=steps[0], is1_n=steps[1], is2_n=steps[2], is2_p=steps[3],
         os_m=steps[4], os_p=steps[5];
    npy_intp ib1_n = is1_n*dn;
    npy_intp ib2_n = is2_n*dn;
    npy_intp ib2_p = is2_p*dp;
    npy_intp ob_p  = os_p *dp;
    if (dn == 0) {
        /* No operand, need to zero the output */
        BEGIN_OUTER_LOOP_3
            char *op=args[2];
            for (m = 0; m < dm; m++) {
                for (p = 0; p < dp; p++) {
                    *(@typ@ *)op = 0;
                    op  +=  os_p;
                }
                op  +=  os_m - ob_p;
            }
        END_OUTER_LOOP
        return;
    }
    BEGIN_OUTER_LOOP_3
        char *ip1=args[0], *ip2=args[1], *op=args[2];
        for (m = 0; m < dm; m++) {
            for (n = 0; n < dn; n++) {
                @typ@ val1 = (*(@typ@ *)ip1);
                for (p = 0; p < dp; p++) {
                    if (n == 0) *(@typ@ *)op = 0;
                    *(@typ@ *)op += val1 * (*(@typ@ *)ip2);
                    ip2 += is2_p;
                    op  +=  os_p;
                }
                ip2 -= ib2_p;
                op  -=  ob_p;
                ip1 += is1_n;
                ip2 += is2_n;
            }
            ip1 -= ib1_n;
            ip2 -= ib2_n;
            ip1 += is1_m;
            op  +=  os_m;
        }
    END_OUTER_LOOP
}

/**end repeat**/

char *euclidean_pdist_signature = "(n,d)->(p)";

/**begin repeat

   #TYPE=FLOAT,DOUBLE#
   #typ=npy_float,npy_double#
   #sqrt_func=sqrtf,sqrt#
*/

/*
 *  This implements the function
 *        out[j*(2*n-3-j)+k-1] = sum_d { (in1[j, d] - in1[k, d])^2 }
 *  with 0 < k < j < n, i.e. computes all unique pairwise euclidean distances.
 */

static void
@TYPE@_euclidean_pdist(char **args, npy_intp *dimensions, npy_intp *steps,
                       void *NPY_UNUSED(func))
{
    INIT_OUTER_LOOP_2
    npy_intp len_n = *dimensions++;
    npy_intp len_d = *dimensions++;
    npy_intp stride_n = *steps++;
    npy_intp stride_d = *steps++;
    npy_intp stride_p = *steps;

    assert(len_n * (len_n - 1) / 2 == *dimensions);

    BEGIN_OUTER_LOOP_2
        const char *data_this = (const char *)args[0];
        char *data_out = args[1];
        npy_intp n;
        for (n = 0; n < len_n; ++n) {
            const char *data_that = data_this + stride_n;
            npy_intp nn;
            for (nn = n + 1; nn < len_n; ++nn) {
                const char *ptr_this = data_this;
                const char *ptr_that = data_that;
                @typ@ out = 0;
                npy_intp d;
                for (d = 0; d < len_d; ++d) {
                    const @typ@ delta = *(const @typ@ *)ptr_this -
                                        *(const @typ@ *)ptr_that;
                    out += delta * delta;
                    ptr_this += stride_d;
                    ptr_that += stride_d;
                }
                *(@typ@ *)data_out = npy_@sqrt_func@(out);
                data_that += stride_n;
                data_out += stride_p;
            }
            data_this += stride_n;
        }
    END_OUTER_LOOP
}

/**end repeat**/


static PyUFuncGenericFunction inner1d_functions[] = { LONG_inner1d, DOUBLE_inner1d };
static void * inner1d_data[] = { (void *)NULL, (void *)NULL };
static char inner1d_signatures[] = { NPY_LONG, NPY_LONG, NPY_LONG, NPY_DOUBLE, NPY_DOUBLE, NPY_DOUBLE };
static PyUFuncGenericFunction innerwt_functions[] = { LONG_innerwt, DOUBLE_innerwt };
static void * innerwt_data[] = { (void *)NULL, (void *)NULL };
static char innerwt_signatures[] = { NPY_LONG, NPY_LONG, NPY_LONG, NPY_LONG, NPY_DOUBLE, NPY_DOUBLE, NPY_DOUBLE, NPY_DOUBLE };
static PyUFuncGenericFunction matrix_multiply_functions[] = { LONG_matrix_multiply, FLOAT_matrix_multiply, DOUBLE_matrix_multiply };
static void *matrix_multiply_data[] = { (void *)NULL, (void *)NULL, (void *)NULL };
static char matrix_multiply_signatures[] = { NPY_LONG, NPY_LONG, NPY_LONG,  NPY_FLOAT, NPY_FLOAT, NPY_FLOAT,  NPY_DOUBLE, NPY_DOUBLE, NPY_DOUBLE };

static PyUFuncGenericFunction euclidean_pdist_functions[] =
                            { FLOAT_euclidean_pdist, DOUBLE_euclidean_pdist };
static void *eucldiean_pdist_data[] = { (void *)NULL, (void *)NULL };
static char euclidean_pdist_signatures[] = { NPY_FLOAT, NPY_FLOAT,
                                             NPY_DOUBLE, NPY_DOUBLE };


static void
addUfuncs(PyObject *dictionary) {
    PyObject *f;

    f = PyUFunc_FromFuncAndDataAndSignature(inner1d_functions, inner1d_data,
                    inner1d_signatures, 2, 2, 1, PyUFunc_None, "inner1d",
                    "inner on the last dimension and broadcast on the rest \n"
                    "     \"(i),(i)->()\" \n",
                    0, inner1d_signature);
    PyDict_SetItemString(dictionary, "inner1d", f);
    Py_DECREF(f);
    f = PyUFunc_FromFuncAndDataAndSignature(innerwt_functions, innerwt_data,
                    innerwt_signatures, 2, 3, 1, PyUFunc_None, "innerwt",
                    "inner1d with a weight argument \n"
                    "     \"(i),(i),(i)->()\" \n",
                    0, innerwt_signature);
    PyDict_SetItemString(dictionary, "innerwt", f);
    Py_DECREF(f);
    f = PyUFunc_FromFuncAndDataAndSignature(matrix_multiply_functions,
                    matrix_multiply_data, matrix_multiply_signatures,
                    3, 2, 1, PyUFunc_None, "matrix_multiply",
                    "matrix multiplication on last two dimensions \n"
                    "     \"(m,n),(n,p)->(m,p)\" \n",
                    0, matrix_multiply_signature);
    PyDict_SetItemString(dictionary, "matrix_multiply", f);
    Py_DECREF(f);
    f = PyUFunc_FromFuncAndDataAndSignature(euclidean_pdist_functions,
                    eucldiean_pdist_data, euclidean_pdist_signatures,
                    2, 1, 1, PyUFunc_None, "euclidean_pdist",
                    "pairwise euclidean distance on last two dimensions \n"
                    "     \"(n,d)->(p)\" \n",
                    0, euclidean_pdist_signature);
    PyDict_SetItemString(dictionary, "euclidean_pdist", f);
    Py_DECREF(f);
    f = PyUFunc_FromFuncAndDataAndSignature(inner1d_functions, inner1d_data,
                    inner1d_signatures, 2, 2, 1, PyUFunc_None, "inner1d_no_doc",
                    NULL,
                    0, inner1d_signature);
    PyDict_SetItemString(dictionary, "inner1d_no_doc", f);
    Py_DECREF(f);
}


static PyObject *
UMath_Tests_test_signature(PyObject *NPY_UNUSED(dummy), PyObject *args)
{
    int nin, nout;
    PyObject *signature, *sig_str;
    PyObject *f;
    int core_enabled;

    if (!PyArg_ParseTuple(args, "iiO", &nin, &nout, &signature)) return NULL;


    if (PyString_Check(signature)) {
        sig_str = signature;
    } else if (PyUnicode_Check(signature)) {
        sig_str = PyUnicode_AsUTF8String(signature);
    } else {
        PyErr_SetString(PyExc_ValueError, "signature should be a string");
        return NULL;
    }

    f = PyUFunc_FromFuncAndDataAndSignature(NULL, NULL, NULL,
        0, nin, nout, PyUFunc_None, "no name",
        "doc:none",
        1, PyString_AS_STRING(sig_str));
    if (sig_str != signature) {
        Py_DECREF(sig_str);
    }
    if (f == NULL) return NULL;
    core_enabled = ((PyUFuncObject*)f)->core_enabled;
    Py_DECREF(f);
    return Py_BuildValue("i", core_enabled);
}

static PyMethodDef UMath_TestsMethods[] = {
    {"test_signature",  UMath_Tests_test_signature, METH_VARARGS,
     "Test signature parsing of ufunc. \n"
     "Arguments: nin nout signature \n"
     "If fails, it returns NULL. Otherwise it will returns 0 for scalar ufunc "
     "and 1 for generalized ufunc. \n",
     },
    {NULL, NULL, 0, NULL}        /* Sentinel */
};

#if defined(NPY_PY3K)
static struct PyModuleDef moduledef = {
        PyModuleDef_HEAD_INIT,
        "umath_tests",
        NULL,
        -1,
        UMath_TestsMethods,
        NULL,
        NULL,
        NULL,
        NULL
};
#endif

#if defined(NPY_PY3K)
#define RETVAL m
PyMODINIT_FUNC PyInit_umath_tests(void)
#else
#define RETVAL
PyMODINIT_FUNC
initumath_tests(void)
#endif
{
    PyObject *m;
    PyObject *d;
    PyObject *version;

#if defined(NPY_PY3K)
    m = PyModule_Create(&moduledef);
#else
    m = Py_InitModule("umath_tests", UMath_TestsMethods);
#endif
    if (m == NULL)
        return RETVAL;

    import_array();
    import_ufunc();

    d = PyModule_GetDict(m);

    version = PyString_FromString("0.1");
    PyDict_SetItemString(d, "__version__", version);
    Py_DECREF(version);

    /* Load the ufunc operators into the module's namespace */
    addUfuncs(d);

    if (PyErr_Occurred()) {
        PyErr_SetString(PyExc_RuntimeError,
                        "cannot load umath_tests module.");
    }

    return RETVAL;
}