#include "pyfgc.h"

void init_coreset(py::module &m) {

py::class_<CSType>(m, "CoresetSampler")

.def(py::init<>())

.def("make_sampler", [](

CSType &cs, size_t ncenters, py::array costs, INA assignments, py::object weights, uint64_t seed, int sens_)

{

if(costs.ndim() != 1) throw std::invalid_argument("costs must be 1-dimensional");

std::fprintf(stderr, "Making sampler. %zu centers, costs has %zd items, sens = %d\n", ncenters, costs.shape(0), sens_);

const auto sens(static_cast<minicore::coresets::SensitivityMethod>(sens_));

py::buffer_info buf1 = costs.request();

const uint32_t *asnp = (const uint32_t *)assignments.request().ptr;

if(buf1.ndim != 1) throw std::runtime_error("buffer must have one dimension (reshape if necessary)");

float *wp = nullptr;

double *dp = nullptr;

if(!weights.is_none()) {

auto winf = py::cast<py::array>(weights).request();

if(winf.ndim != 1) throw std::invalid_argument("weightss must be 1-dimensional");

if(winf.format[0] == 'f')

wp = (float *)winf.ptr;

else if(winf.format[0] == 'd')

dp = (double *)winf.ptr;

else throw std::invalid_argument("Weights can only be double or float");

}

const py::ssize_t np = costs.shape(0);

switch(buf1.format[0]) {

case 'f': if(dp) {

cs.make_sampler(np, ncenters, (float *)buf1.ptr, asnp, dp, seed, sens); break;

} else {

cs.make_sampler(np, ncenters, (float *)buf1.ptr, asnp, wp, seed, sens); break;

}

case 'd': if(dp) {

cs.make_sampler(np, ncenters, (double *)buf1.ptr, asnp, dp, seed, sens); break;

} else {

cs.make_sampler(np, ncenters, (double *)buf1.ptr, asnp, wp, seed, sens); break;

}

default: throw std::invalid_argument("Costs is not double or float");

}

},

"Generates a coreset sampler given a set of costs, assignments, and, optionally, weights. This can be used to generate an index coreset",

py::arg("ncenters"), py::arg("costs"), py::arg("assignments"),

py::arg("weights") = py::none(), py::arg("seed") = 13, py::arg("sens")=0

).def("get_probs", [](CSType &cs) {

py::array_t<float> ret(cs.np_);

std::copy(cs.probs_.get(), cs.probs_.get() + cs.np_, (float *)ret.request().ptr);

return ret;

}, "Create a numpy array of sampling probabilities")

.def("sample", [](CSType &cs, py::ssize_t size, py::ssize_t seed, bool unique_samples, double eps) {

if(cs.sampler_ == nullptr) throw std::invalid_argument("Can't sample without created sampler. Call make_ampler");

if(seed == 0) seed = std::rand();

auto ret = cs.sample(size, seed, eps, unique_samples);

py::array_t<float> rf(size);

py::array_t<uint64_t> ri(size);

std::copy(ret.weights_.begin(), ret.weights_.end(), (float *)rf.request().ptr);

std::copy(ret.indices_.begin(), ret.indices_.end(), (uint64_t *)ri.request().ptr);

return py::make_tuple(rf, ri);

}, py::arg("size"), py::arg("seed") = 0, py::arg("unique") = false, py::arg("eps") = 0.1);

}