#include <Python.h>
#include <numpy/arrayobject.h>
#include "lsh.h"
#include <math.h>

/* Docstrings */
static char module_docstring[] = "This module implements fast nearest-neighbor retrieval using LSH.";
static char lsh_docstring[] = "Calculate the closest neightbours with distances given a query window.";

/* Available functions */
static PyObject* lsh_lsh(PyObject *self, PyObject *args);

/* Module specification */
static PyMethodDef module_methods[] = { { "lsh", lsh_lsh, METH_VARARGS, lsh_docstring }, { NULL, NULL, 0, NULL } };

/* Initialize the module */
#if PY_MAJOR_VERSION >= 3
#define MOD_ERROR_VAL NULL
  #define MOD_SUCCESS_VAL(val) val
  #define MOD_INIT(name) PyMODINIT_FUNC PyInit_##name(void)
  #define MOD_DEF(ob, name, doc, methods) \
          static struct PyModuleDef moduledef = { \
            PyModuleDef_HEAD_INIT, name, doc, -1, methods, }; \
          ob = PyModule_Create(&moduledef);
#else
#define MOD_ERROR_VAL
#define MOD_SUCCESS_VAL(val)
#define MOD_INIT(name) void init##name(void)
#define MOD_DEF(ob, name, doc, methods) \
          ob = Py_InitModule3(name, methods, doc);
#endif

MOD_INIT(_lsh)
{
    PyObject *m;

    MOD_DEF(m, "_lsh", module_docstring,
            module_methods)

    if (m == NULL)
        return MOD_ERROR_VAL;

    import_array();

    return MOD_SUCCESS_VAL(m);
}

static PyObject* lsh_lsh(PyObject* self, PyObject* args) {
    PyObject* data_obj = NULL;
    PyObject* query_obj = NULL;
    PyObject* hash_obj = NULL;
    double **** hashFunctions;
    double ** weights = NULL;
    double*** data;
    double ** query;
    int * candidates;
    double * distances;
    int nrOfCandidates;
    double r;
    double a;
    double sd;
    PyArray_Descr *descr;
    descr = PyArray_DescrFromType(NPY_DOUBLE);
    PyArray_Descr *descr_float;
    descr_float = PyArray_DescrFromType(NPY_FLOAT64);

    /// Parse the input tuple
    if (!PyArg_ParseTuple(args, "OOddd|O", &data_obj, &query_obj, &r, &a, &sd, &hash_obj)) {
        return NULL;
    }

    /// Get the dimensions of the data and query
    int data_size = (long) PyArray_DIM(data_obj, 0);
    int channel_size = (long) PyArray_DIM(data_obj, 2);
    int query_size = (int) PyArray_DIM(query_obj, 0);

    /// Convert data, query and weights to C array
    npy_intp dims0[3];
    if (PyArray_AsCArray(&query_obj, (void **)&query, dims0, 2, descr) < 0 || PyArray_AsCArray(&data_obj, (void ***)&data, dims0, 3, descr) < 0) {
        PyErr_SetString(PyExc_TypeError, "error converting to c array");
        return NULL;
    }
    if (hash_obj != NULL)
    {
        printf("Using weights");
        if (PyArray_AsCArray(&hash_obj, (void **)&weights, dims0, 2, descr) < 0) {
            PyErr_SetString(PyExc_TypeError, "error converting weights to c array");
            return NULL;
        }
    }

    int K = ceil(log((log(0.5))/log(1-exp(-2*(0.1)*(0.1)*query_size)))/log((1-exp(-2*(0.1)*(0.1)*query_size))/(0.5)));
    int L = ceil(log(0.05)/(log(1-pow(1-exp(-2*(0.1)*(0.1)*query_size), K))));
    printf("K: %d\n", K);
    printf("L: %d\n", L);
    printf("Dim: %d\n", channel_size);

    /// Initialize output parameters
    hashFunctions = (double ****)malloc(L*sizeof(double***));
    for (int l=0;l<L;l++)
    {
        hashFunctions[l] = (double ***)malloc(K*sizeof(double**));
        for (int k=0;k<K;k++)
        {
            hashFunctions[l][k] = (double **)malloc(query_size*sizeof(double*));
            for (int t=0;t<query_size;t++)
            {
                hashFunctions[l][k][t] = (double *)malloc(channel_size*sizeof(double));
            }
        }
    }

    int status = lsh(data, data_size, query_size, channel_size, query, L, K, r, a, sd, candidates, distances, hashFunctions, weights, nrOfCandidates);
    if (status) {
        PyErr_SetString(PyExc_RuntimeError, "lsh could not allocate memory");
        return NULL;
    }

    npy_intp dimscandidates[1] = {nrOfCandidates};
    printf("Number of candidates: %d\n", nrOfCandidates);
    npy_intp dims4[4] = {L, K, query_size, channel_size};

//    PyArrayObject* numpy_candidates = (PyArrayObject*)PyArray_SimpleNewFromData(1, dimscandidates, NPY_INT, (void*)&candidates);
//    PyArrayObject* numpy_distances = (PyArrayObject*)PyArray_SimpleNewFromData(1, dimsdistance, NPY_DOUBLE, (void*)&distances);
    // https://github.com/suiyun0234/scipy-master/commit/da7dfc7aad8daa7a516e43f4c7001eea7c1a707e
    // https://github.com/fjean/pymeanshift/commit/1ba90da647342184ea7df378d7f21eba257a51d9
    PyArrayObject* numpy_candidates = (PyArrayObject*)PyArray_SimpleNew(1, dimscandidates, NPY_INT);
    memcpy(PyArray_DATA(numpy_candidates), candidates, dimscandidates[0]*sizeof(int));
    PyArrayObject* numpy_distances = (PyArrayObject*)PyArray_SimpleNew(1, dimscandidates, NPY_DOUBLE);
    memcpy(PyArray_DATA(numpy_distances), distances, dimscandidates[0]*sizeof(double));
    PyArrayObject* TEST = (PyArrayObject*)PyArray_SimpleNew(1, dimscandidates, NPY_DOUBLE);
    memcpy(PyArray_DATA(TEST), distances, dimscandidates[0]*sizeof(double));
//    PyArrayObject* numpy_hash_functions = (PyArrayObject*)PyArray_SimpleNew(4, dims4, NPY_DOUBLE);
//    double* numpy_hash_functions_data = (double*)PyArray_DATA(numpy_hash_functions);
//    for (int l=0;l<L;l++)
//    {
//        for (int k=0;k<K;k++)
//        {
//            for (int t=0;t<query_size;t++)
//            {
//                memcpy(numpy_hash_functions_data, hashFunctions[l][k][t], channel_size*sizeof(double));
//                numpy_hash_functions_data += channel_size;
//            }
//        }
//    }
    PyObject* ret = Py_BuildValue("NNN", PyArray_Return(numpy_candidates), PyArray_Return(numpy_distances), PyArray_Return(TEST));// PyArray_Return(numpy_hash_functions));
//    Py_XDECREF(data_obj);
//    Py_XDECREF(query_obj);
//    Py_XDECREF(hash_obj);
//    Py_XDECREF(weights);
//    Py_XDECREF(data);
//    Py_XDECREF(query);
//    Py_XDECREF(descr);
//    Py_XDECREF(descr_float);
//    Py_XDECREF(query);
//    Py_XDECREF(numpy_candidates);
//    Py_XDECREF(numpy_distances);
//    Py_XDECREF(TEST);
//    free(candidates);
//    free(distances);
//    for (int l=0;l<L;l++)
//    {
//        for (int k=0;k<K;k++)
//        {
//            for (int t=0;t<query_size;t++)
//            {
//                free(hashFunctions[l][k][t]);
//            }
//            free(hashFunctions[l][k]);
//        }
//        free(hashFunctions[l]);
//    }
//    free(hashFunctions);
    return ret;
}