#include <boost/random.hpp>

#include <boost/python/module.hpp>
#include <boost/python/class.hpp>
#include <boost/python/object.hpp>

#include <numpy/arrayobject.h>

namespace py = boost::python;

template<class T>
struct Generator
{
    typedef T rng_type;
    typedef typename T::result_type result_type;

    static void seed(result_type value) {
        get_rng().seed(value);
    }

    static PyObject* uniform_smallint(int min, int max, py::object size) {
        boost::uniform_smallint<int> dist(min, max);
        return generate(dist, size, NPY_LONG);
    }

    static PyObject* uniform_int(int min, int max, py::object size) {
        boost::uniform_int<int> dist(min, max);
        return generate(dist, size, NPY_LONG);
    }

#if 0
    static PyObject* uniform_01(py::object size) {
        boost::uniform_01<int> dist;
        return generate(dist, size, NPY_LONG);
    }
#endif

    static PyObject* uniform_real(double min, double max, py::object size) {
        boost::uniform_real<double> dist(min, max);
        return generate(dist, size, NPY_DOUBLE);
    }

    static PyObject* bernoulli(double p, py::object size) {
        boost::bernoulli_distribution<double> dist(p);
        return generate(dist, size, NPY_DOUBLE);
    }

    static PyObject* geometric(double p, py::object size) {
        boost::geometric_distribution<double> dist(p);
        return generate(dist, size, NPY_DOUBLE);
    }

    static PyObject* triangle(double a, double b, double c, py::object size) {
        boost::triangle_distribution<double> dist(a, b, c);
        return generate(dist, size, NPY_DOUBLE);
    }

    static PyObject* exponential(double lambda, py::object size) {
        boost::exponential_distribution<double> dist(lambda);
        return generate(dist, size, NPY_DOUBLE);
    }

    static PyObject* normal(double mean, double sigma, py::object size) {
        boost::normal_distribution<double> dist(mean, sigma);
        return generate(dist, size, NPY_DOUBLE);
    }

    static PyObject* lognormal(double mean, double sigma, py::object size) {
        boost::lognormal_distribution<double> dist(mean, sigma);
        return generate(dist, size, NPY_DOUBLE);
    }

    static PyObject* uniform_on_sphere(int dim, py::object size) {
        boost::uniform_on_sphere<double> dist(dim);
        return generate(dist, size, NPY_DOUBLE);
    }

    template<class DistType>
    static PyObject* generate(DistType dist, py::object size, int typenum)
    {
        PyArray_Dims dims;
        if (!PyArray_IntpConverter(size.ptr(), &dims)) {
            throw py::error_already_set();
        }
        PyObject* arr = PyArray_SimpleNew(dims.len, dims.ptr, typenum);
        if (!arr) {
            throw py::error_already_set();
        }
        rng_type& rng = get_rng();
        boost::variate_generator<rng_type&, DistType> generator(rng, dist);
        double* data = (double*)PyArray_DATA(arr);
        std::generate(data, data + PyArray_SIZE(arr), generator);
        return arr;
    }

    static rng_type& get_rng() {
        static T rng;
        return rng;
    }
};

template<class T>
void wrap_generator(const char* name)
{
    py::class_<T, boost::noncopyable>(name, py::no_init)
        .def("seed", &Generator<T>::seed)
        .def("uniform_smallint", &Generator<T>::uniform_smallint,
             (py::arg("min"), py::arg("max"), py::arg("size")=py::handle<>(Py_None)))
        .def("uniform_int", &Generator<T>::uniform_int,
             (py::arg("min")=0, py::arg("max")=9, py::arg("size")=py::handle<>(Py_None)))
#if 0
        .def("uniform_01", &Generator<T>::uniform_01, py::arg("size")=py::handle<>(Py_None))
#endif
        .def("uniform_real", &Generator<T>::uniform_real,
             (py::arg("min")=0.0, py::arg("max")=1.0, py::arg("size")=py::handle<>(Py_None)))
        .def("bernoulli", &Generator<T>::bernoulli,
             (py::arg("p")=0.5, py::arg("size")=py::handle<>(Py_None)))
        .def("geometric", &Generator<T>::geometric,
             (py::arg("p")=0.5, py::arg("size")=py::handle<>(Py_None)))
        .def("triangle", &Generator<T>::triangle,
             (py::arg("a"), py::arg("b"), py::arg("c"), py::arg("size")=py::handle<>(Py_None)))
        .def("exponential", &Generator<T>::exponential,
             (py::arg("lambda")=1.0, py::arg("size")=py::handle<>(Py_None)))
        .def("normal", &Generator<T>::normal,
             (py::arg("mean")=0.0, py::arg("sigma")=1.0, py::arg("size")=py::handle<>(Py_None)))
        .def("lognormal", &Generator<T>::lognormal,
             (py::arg("mean")=1.0, py::arg("sigma")=1.0, py::arg("size")=py::handle<>(Py_None)))
#if 0
        .def("uniform_on_sphere", &Generator<T>::uniform_on_sphere,
             (py::arg("dim")=2, py::arg("size")=py::handle<>(Py_None)))
#endif
        .staticmethod("seed")
        .staticmethod("uniform_smallint")
        .staticmethod("uniform_int")
#if 0
        .staticmethod("uniform_01")
#endif
        .staticmethod("uniform_real")
        .staticmethod("bernoulli")
        .staticmethod("geometric")
        .staticmethod("triangle")
        .staticmethod("exponential")
        .staticmethod("normal")
        .staticmethod("lognormal")
#if 0
        .staticmethod("uniform_on_sphere")
#endif
        ;
}

BOOST_PYTHON_MODULE(_pyspkrec_random)
{
    import_array();
    wrap_generator<boost::minstd_rand>("minstd_rand");
    wrap_generator<boost::minstd_rand0>("minstd_rand0");
    wrap_generator<boost::rand48>("rand48");
#if 0
    wrap_generator<boost::ecuyer1988>("ecuyer1988");
#endif
    wrap_generator<boost::kreutzer1986>("kreutzer1986");
    wrap_generator<boost::hellekalek1995>("hellekalek1995");
    wrap_generator<boost::mt11213b>("mt11213b");
    wrap_generator<boost::mt19937>("mt19937");
#if 0
    wrap_generator<boost::lagged_fibonacci607>("");
    wrap_generator<boost::lagged_fibonacci1279>("");
    wrap_generator<boost::lagged_fibonacci2281>("");
    wrap_generator<boost::lagged_fibonacci3217>("");
    wrap_generator<boost::lagged_fibonacci4423>("");
    wrap_generator<boost::lagged_fibonacci9689>("");
    wrap_generator<boost::lagged_fibonacci19937>("");
    wrap_generator<boost::lagged_fibonacci23209>("");
    wrap_generator<boost::lagged_fibonacci44497>("");
#endif
}