/* -*- c -*- */

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

#include "Python.h"
#include "numpy/arrayobject.h"
#include "numpy/ufuncobject.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_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;
    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**/

/*  The following lines were generated using a slightly modified
    version of code_generators/generate_umath.py and adding these
    lines to defdict:

defdict = {
'inner1d' :
    Ufunc(2, 1, None_,
        r'''inner on the last dimension and broadcast on the rest \n"
        "     \"(i),(i)->()\" \n''',
          TD('ld'),
          ),
'innerwt' :
    Ufunc(3, 1, None_,
          r'''inner1d with a weight argument \n"
          "     \"(i),(i),(i)->()\" \n''',
          TD('ld'),
          ),
}

*/

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 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);
}

/*
    End of auto-generated code.
*/



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;
}
