#include "pyfgc.h"

using blz::geomedian;

template<typename FT>

auto pygeomedian(py::array_t<FT, py::array::c_style | py::array::forcecast> data, py::object weights, FT eps) {

auto dbuf = data.request();

if(dbuf.ndim != 2) throw std::runtime_error("Expected 2-d array");

py::ssize_t ndret = dbuf.shape[1];

py::array_t<FT> ret(ndret);

auto rbuf = ret.request();

FT *retp = (FT *)rbuf.ptr;

FT *srcp = (FT *)dbuf.ptr;

blaze::CustomMatrix<FT, blaze::unaligned, blaze::unpadded> cm(srcp, dbuf.shape[0], dbuf.shape[1], dbuf.strides[0] / sizeof(FT));

blaze::CustomVector<FT, blaze::unaligned, blaze::unpadded, blaze::rowVector> cv(retp, dbuf.shape[1]);

if(weights.is_none()) {

geomedian(cm, cv, eps);

} else {

auto wa = py::cast<py::array>(weights);

auto wbi = wa.request();

void *wptr = wbi.ptr;

if(wbi.format.size() != 1) {

throw std::invalid_argument(std::string("unexpected format string ") + wbi.format);

}

if(wbi.shape[0] != dbuf.shape[0]) throw std::invalid_argument("Mismatched shapes");

switch(wbi.format.front()) {

case 'f': geomedian(cm, cv, static_cast<float *>(wptr), eps); break;

case 'd': geomedian(cm, cv, static_cast<double *>(wptr), eps); break;

case 'i': case 'I': case 'q': case 'Q': case 'l': case 'L': case 'h': case 'H': case 'b': case 'B':

switch(wbi.itemsize) {

case 1: geomedian(cm, cv, static_cast<uint8_t *>(wptr), eps); break;

case 2: geomedian(cm, cv, static_cast<uint16_t *>(wptr), eps); break;

case 4: geomedian(cm, cv, static_cast<uint32_t *>(wptr), eps); break;

case 8: geomedian(cm, cv, static_cast<uint64_t *>(wptr), eps); break;

default: throw std::invalid_argument("Not expected: anything");

}

break;

default: throw std::invalid_argument(std::string("Type for weights ") + wbi.format + "is unsupported. Supported: float, double, or {u,}int{16,32,64}.");

}

}

return ret;

}

void init_centroid(py::module &m) {

m.def("geomed", [](py::array data, py::object weights, double eps) {

auto dbi = data.request();

if(dbi.format.size() != 1u) throw std::invalid_argument(std::string("dbi format is the wrong size: ") + dbi.format);

if(dbi.shape.size() != 2) throw std::invalid_argument(std::string("dbi shape has the wrong number of dimensions. Expected 2, found ") + std::to_string(dbi.shape.size()));

switch(dbi.format.front()) {

case 'f': pygeomedian<float>(py::cast<py::array_t<float>>(data), weights, eps); break;

case 'd': pygeomedian<double>(py::cast<py::array_t<double>>(data), weights, eps); break;

default: throw std::invalid_argument(std::string("dbi format is not floating point.: ") + dbi.format);

}

}, "Compute geometric median for a 2-d array. Optional arguments:\n\n Weights: 1-d array, defaults to unweighted. Must be nonnegative.\n epsilon: parameter for termination (default: 0 / precision)\n",

py::arg("data"), py::arg("weights") = py::none(), py::arg("eps") = 0.);

}