import numpy
import theano
from theano import gof
import theano.tensor as T

theano.config.gcc.cxxflags="<NEED TO ENTER PATH TO INCLUDES HERE>"
class wall_log_op(gof.Op):
    __props__ = ()

    def make_node(self, value, n, m, N, odds):
        value_ =T.as_tensor_variable(numpy.array(value, dtype=numpy.int32))
        n_ =T.as_tensor_variable(numpy.array(n, dtype=numpy.int32))
        m_ =T.as_tensor_variable(numpy.array(m, dtype=numpy.int32))
        N_ =T.as_tensor_variable(numpy.array(N, dtype=numpy.int32))
        odds_ = T.as_tensor_variable(numpy.array(odds, dtype=numpy.double))

        output_var = odds_.type()

        return gof.Apply(self, [value_, n_, m_, N_, odds_], [output_var])

    def c_code_cache_version(self):
        return (1, 0)

    def c_headers(self):
        return (
        ["<time.h>", "<randomc.h>", "<mersenne.cpp>", "<stocc.h>", "<stoc1.cpp>", "<stoc3.cpp>", "<userintf.cpp>",
         "<wnchyppr.cpp>", "<fnchyppr.cpp>"])

    def c_code(self, node, name, inp, out, sub):

        value, n, m, N, odds = inp
        logp, = out

        # Extract the dtypes of the inputs and outputs storage to
        # be able to declare pointers for those dtypes in the C
        # code.
        dtype_value = node.inputs[0].dtype

        dtype_n = node.inputs[1].dtype
        dtype_m = node.inputs[2].dtype
        dtype_N = node.inputs[3].dtype
        dtype_odds = node.inputs[4].dtype
        dtype_logp = node.outputs[0].dtype


        itemsize_value = numpy.dtype(dtype_value).itemsize

        itemsize_n = numpy.dtype(dtype_n).itemsize
        itemsize_m = numpy.dtype(dtype_m).itemsize
        itemsize_N = numpy.dtype(dtype_N).itemsize
        itemsize_odds = numpy.dtype(dtype_odds).itemsize
        itemsize_logp = numpy.dtype(dtype_logp).itemsize

        fail = sub['fail']

        c_code = """

        // Validate that the output storage exists and has the same
        // dimension as x.
        if (NULL == %(logp)s  ||
            PyArray_DIMS(%(logp)s)[0] != PyArray_DIMS(%(value)s)[0])
        {

            /* Reference received to invalid output variable.
            Decrease received reference's ref count and allocate new
            output variable */

            Py_XDECREF(%(logp)s);
            %(logp)s = (PyArrayObject*)PyArray_EMPTY(1,
                                                PyArray_DIMS(%(value)s),
                                                PyArray_TYPE(%(odds)s),
                                                0);
            if (!%(logp)s) {
                %(fail)s;
            }

        }


        {

            npy_%(dtype_logp)s* logp_data_ptr =
                            (npy_%(dtype_logp)s*)PyArray_DATA(%(logp)s);

            npy_%(dtype_value)s* value_data_ptr =
                            (npy_%(dtype_value)s*)PyArray_DATA(%(value)s);
            npy_%(dtype_n)s* n_data_ptr =
                            (npy_%(dtype_n)s*)PyArray_DATA(%(n)s);

            npy_%(dtype_m)s* m_data_ptr =
                            (npy_%(dtype_m)s*)PyArray_DATA(%(m)s);

            npy_%(dtype_N)s* N_data_ptr =
                            (npy_%(dtype_N)s*)PyArray_DATA(%(N)s);

            npy_%(dtype_odds)s* odds_data_ptr =
                            (npy_%(dtype_odds)s*)PyArray_DATA(%(odds)s);


            int value_stride = PyArray_STRIDES(%(value)s)[0] / %(itemsize_value)s;

            int n_stride = PyArray_STRIDES(%(n)s)[0] / %(itemsize_n)s;

            int m_stride = PyArray_STRIDES(%(m)s)[0] / %(itemsize_m)s;
            int N_stride = PyArray_STRIDES(%(N)s)[0] / %(itemsize_N)s;
            int odds_stride = PyArray_STRIDES(%(odds)s)[0] / %(itemsize_odds)s;
            int logp_stride = PyArray_STRIDES(%(logp)s)[0] / %(itemsize_logp)s;
            int value_dim = PyArray_DIMS(%(value)s)[0];



            for(int i=0; i < value_dim; i++)
            {
            CWalleniusNCHypergeometric wnch(
                n_data_ptr[i * n_stride],
                m_data_ptr[i * m_stride],
                N_data_ptr[i * N_stride],
                odds_data_ptr[i * odds_stride],
                                1.E-8);
                logp_data_ptr[i * logp_stride] =  (log(wnch.probability(value_data_ptr[i * value_stride])));

            }


        }
        """

        return c_code % locals()