/**
 * @author Andre Anjos <andre.anjos@idiap.ch>
 * @date Tue 29 Apr 09:26:02 2014 CEST
 *
 * @brief Bindings for Cost functions
 *
 * Copyright (C) 2011-2014 Idiap Research Institute, Martigny, Switzerland
 */

#define BOB_LEARN_MLP_MODULE
#include <bob.blitz/cppapi.h>
#include <bob.blitz/cleanup.h>
#include <bob.extension/defines.h>
#include <bob.learn.mlp/api.h>
#include <bob.learn.activation/api.h>
#include <structmember.h>
#include <boost/function.hpp>
#include <boost/bind.hpp>

/*************************************
 * Implementation of Cost base class *
 *************************************/

PyDoc_STRVAR(s_cost_str, BOB_EXT_MODULE_PREFIX ".Cost");

PyDoc_STRVAR(s_cost_doc,
"A base class for evaluating the performance cost.\n\
\n\
This is the base class for all concrete (C++ only) loss\n\
function implementations. You cannot instantiate objects\n\
of this type directly, use one of the derived classes.\n\
");

static int PyBobLearnCost_init
(PyBobLearnCostObject* self, PyObject* args, PyObject* kwds) {

  PyErr_Format(PyExc_NotImplementedError, "cannot instantiate objects of type `%s', use one of the derived classes", Py_TYPE(self)->tp_name);
  return -1;

}

int PyBobLearnCost_Check(PyObject* o) {
  return PyObject_IsInstance(o, reinterpret_cast<PyObject*>(&PyBobLearnCost_Type));
}

static PyObject* PyBobLearnCost_RichCompare
(PyBobLearnCostObject* self, PyObject* other, int op) {

  if (!PyBobLearnCost_Check(other)) {
    PyErr_Format(PyExc_TypeError, "cannot compare `%s' with `%s'",
        Py_TYPE(self)->tp_name, Py_TYPE(other)->tp_name);
    return 0;
  }

  auto other_ = reinterpret_cast<PyBobLearnCostObject*>(other);

  switch (op) {
    case Py_EQ:
      if (self->cxx->str() == other_->cxx->str()) Py_RETURN_TRUE;
      Py_RETURN_FALSE;
      break;
    case Py_NE:
      if (self->cxx->str() != other_->cxx->str()) Py_RETURN_TRUE;
      Py_RETURN_FALSE;
      break;
    default:
      Py_INCREF(Py_NotImplemented);
      return Py_NotImplemented;
  }

}

#if PY_VERSION_HEX >= 0x03000000
#  define PYOBJECT_STR PyObject_Str
#else
#  define PYOBJECT_STR PyObject_Unicode
#endif

PyObject* PyBobLearnCost_Repr(PyBobLearnCostObject* self) {

  /**
   * Expected output:
   *
   * <bob.learn.linear.Cost [...]>
   */

  auto retval = PyUnicode_FromFormat("<%s [act: %s]>",
        Py_TYPE(self)->tp_name, self->cxx->str().c_str());

#if PYTHON_VERSION_HEX < 0x03000000
  if (!retval) return 0;
  PyObject* tmp = PyObject_Str(retval);
  Py_DECREF(retval);
  retval = tmp;
#endif

  return retval;

}

PyObject* PyBobLearnCost_Str(PyBobLearnCostObject* self) {
  return Py_BuildValue("s", self->cxx->str().c_str());
}

/**
 * Checks if a array `a1' and `a2' have a matching shape.
 */
static int have_same_shape (PyBlitzArrayObject* a1, PyBlitzArrayObject* a2) {

  if (a1->ndim != a2->ndim) return 0;

  for (Py_ssize_t k=0; k<a1->ndim; ++k) {
    if (a1->shape[k] != a2->shape[k]) return 0;
  }

  return 1;
}

static PyObject* apply_scalar(PyBobLearnCostObject* self, const char*,
    boost::function<double (double, double)> function,
    PyObject* args, PyObject* kwds) {

  static const char* const_kwlist[] = {"output", "target", 0};
  static char** kwlist = const_cast<char**>(const_kwlist);

  double output = 0.;
  double target = 0.;

  if (!PyArg_ParseTupleAndKeywords(args, kwds, "dd", kwlist,
        &output, &target)) return 0;

  return Py_BuildValue("d", function(output, target));

}

/**
 * Maps all elements of arr through function() into `result'
 */
static PyObject* apply_array(PyBobLearnCostObject* self, const char* fname,
    boost::function<double (double, double)> function,
    PyObject* args, PyObject* kwds) {

  static const char* const_kwlist[] = {"output", "target", "result", 0};
  static char** kwlist = const_cast<char**>(const_kwlist);

  PyBlitzArrayObject* output = 0;
  PyBlitzArrayObject* target = 0;
  PyBlitzArrayObject* result = 0;

  if (!PyArg_ParseTupleAndKeywords(args, kwds, "O&O&|O&", kwlist,
        &PyBlitzArray_Converter, &output,
        &PyBlitzArray_Converter, &target,
        &PyBlitzArray_OutputConverter, &result
        )) return 0;

  //protects acquired resources through this scope
  auto output_ = make_safe(output);
  auto target_ = make_safe(target);
  auto result_ = make_xsafe(result);

  if (output->type_num != NPY_FLOAT64) {
    PyErr_Format(PyExc_TypeError, "`%s.%s' only supports 64-bit float arrays for input array `output'", Py_TYPE(self)->tp_name, fname);
    return 0;
  }

  if (target->type_num != NPY_FLOAT64) {
    PyErr_Format(PyExc_TypeError, "`%s.%s' only supports 64-bit float arrays for input array `target'", Py_TYPE(self)->tp_name, fname);
    return 0;
  }

  if (result && result->type_num != NPY_FLOAT64) {
    PyErr_Format(PyExc_TypeError, "`%s.%s' only supports 64-bit float arrays for output array `result'", Py_TYPE(self)->tp_name, fname);
    return 0;
  }

  if (!have_same_shape(output, target)) {
    PyErr_Format(PyExc_RuntimeError, "`%s.%s' requires input arrays `output' and `target' to have the same shape, but you provided arrays with different shapes", Py_TYPE(self)->tp_name, fname);
    return 0;
  }

  if (result && !have_same_shape(output, result)) {
    PyErr_Format(PyExc_RuntimeError, "`%s.%s' requires output array `result' to have the same shape as input arrays `output' and `target', but you provided arrays with different shapes", Py_TYPE(self)->tp_name, fname);
    return 0;
  }

  /** if ``result`` was not pre-allocated, do it now **/
  if (!result) {
    result = (PyBlitzArrayObject*)PyBlitzArray_SimpleNew(NPY_FLOAT64, output->ndim, output->shape);
    result_ = make_safe(result);
  }

  switch (output->ndim) {
    case 1:
      {
        blitz::Array<double,1>& output_ =
          *PyBlitzArrayCxx_AsBlitz<double,1>(output);
        blitz::Array<double,1>& target_ =
          *PyBlitzArrayCxx_AsBlitz<double,1>(target);
        blitz::Array<double,1>& result_ =
          *PyBlitzArrayCxx_AsBlitz<double,1>(result);
        for (int k=0; k<output_.extent(0); ++k)
          result_(k) = function(output_(k), target_(k));
      }
      break;

    case 2:
      {
        blitz::Array<double,2>& output_ =
          *PyBlitzArrayCxx_AsBlitz<double,2>(output);
        blitz::Array<double,2>& target_ =
          *PyBlitzArrayCxx_AsBlitz<double,2>(target);
        blitz::Array<double,2>& result_ =
          *PyBlitzArrayCxx_AsBlitz<double,2>(result);
        for (int k=0; k<output_.extent(0); ++k)
          for (int l=0; l<output_.extent(1); ++l)
            result_(k,l) = function(output_(k,l), target_(k,l));
      }
      break;

    case 3:
      {
        blitz::Array<double,3>& output_ =
          *PyBlitzArrayCxx_AsBlitz<double,3>(output);
        blitz::Array<double,3>& target_ =
          *PyBlitzArrayCxx_AsBlitz<double,3>(target);
        blitz::Array<double,3>& result_ =
          *PyBlitzArrayCxx_AsBlitz<double,3>(result);
        for (int k=0; k<output_.extent(0); ++k)
          for (int l=0; l<output_.extent(1); ++l)
            for (int m=0; m<output_.extent(2); ++m)
              result_(k,l,m) = function(output_(k,l,m), target_(k,l,m));
      }
      break;

    case 4:
      {
        blitz::Array<double,4>& output_ =
          *PyBlitzArrayCxx_AsBlitz<double,4>(output);
        blitz::Array<double,4>& target_ =
          *PyBlitzArrayCxx_AsBlitz<double,4>(target);
        blitz::Array<double,4>& result_ =
          *PyBlitzArrayCxx_AsBlitz<double,4>(result);
        for (int k=0; k<output_.extent(0); ++k)
          for (int l=0; l<output_.extent(1); ++l)
            for (int m=0; m<output_.extent(2); ++m)
              for (int n=0; n<output_.extent(3); ++n)
                result_(k,l,m,n) = function(output_(k,l,m,n), target_(k,l,m,n));
      }
      break;

    default:
      PyErr_Format(PyExc_RuntimeError, "`%s.%s' only accepts 1, 2, 3 or 4-dimensional double arrays (not %" PY_FORMAT_SIZE_T "dD arrays)", Py_TYPE(self)->tp_name, fname, output->ndim);
    return 0;

  }

  return PyBlitzArray_NUMPY_WRAP(Py_BuildValue("O", result));

}

PyDoc_STRVAR(s_f_str, "f");
PyDoc_STRVAR(s_f_doc,
"o.f(output, target, [result]) -> result\n\
\n\
Computes the cost, given the current and expected outputs.\n\
\n\
Keyword arguments:\n\
\n\
output, ND array, float64 | scalar\n\
  Real output from the machine. May be a N-dimensional array\n\
  or a plain scalar.\n\
\n\
target, ND array, float64 | scalar\n\
  Target output you are training to achieve. The data type\n\
  and extents for this object must match that of ``target``.\n\
\n\
result (optional), ND array, float64\n\
  Where to place the result from the calculation. You can\n\
  pass this argument if the input are N-dimensional arrays.\n\
  Otherwise, it is an error to pass such a container. If the\n\
  inputs are arrays and an object for ``result`` is passed,\n\
  then its dimensions and data-type must match that of both\n\
  ``output`` and ``result``.\n\
\n\
Returns the cost as a scalar, if the input were scalars or\n\
as an array with matching size of ``output`` and ``target``\n\
otherwise.\n\
");

static PyObject* PyBobLearnCost_f
(PyBobLearnCostObject* self, PyObject* args, PyObject* kwds) {

  PyObject* arg = 0; ///< borrowed (don't delete)
  if (PyTuple_Size(args)) arg = PyTuple_GET_ITEM(args, 0);
  else {
    PyObject* tmp = PyDict_Values(kwds);
    auto tmp_ = make_safe(tmp);
    arg = PyList_GET_ITEM(tmp, 0);
  }

  if (PyBob_NumberCheck(arg))
    return apply_scalar(self, s_f_str,
        boost::bind(&bob::learn::mlp::Cost::f, self->cxx, _1, _2), args, kwds);

  return apply_array(self, s_f_str,
      boost::bind(&bob::learn::mlp::Cost::f, self->cxx, _1, _2), args, kwds);

}

PyDoc_STRVAR(s_f_prime_str, "f_prime");
PyDoc_STRVAR(s_f_prime_doc,
"o.f_prime(output, target, [result]) -> result\n\
\n\
Computes the derivative of the cost w.r.t. output.\n\
\n\
Keyword arguments:\n\
\n\
output, ND array, float64 | scalar\n\
  Real output from the machine. May be a N-dimensional array\n\
  or a plain scalar.\n\
\n\
target, ND array, float64 | scalar\n\
  Target output you are training to achieve. The data type\n\
  and extents for this object must match that of ``target``.\n\
\n\
result (optional), ND array, float64\n\
  Where to place the result from the calculation. You can\n\
  pass this argument if the input are N-dimensional arrays.\n\
  Otherwise, it is an error to pass such a container. If the\n\
  inputs are arrays and an object for ``result`` is passed,\n\
  then its dimensions and data-type must match that of both\n\
  ``output`` and ``result``.\n\
\n\
Returns the cost as a scalar, if the input were scalars or\n\
as an array with matching size of ``output`` and ``target``\n\
otherwise.\n\
");

static PyObject* PyBobLearnCost_f_prime
(PyBobLearnCostObject* self, PyObject* args, PyObject* kwds) {

  PyObject* arg = 0; ///< borrowed (don't delete)
  if (PyTuple_Size(args)) arg = PyTuple_GET_ITEM(args, 0);
  else {
    PyObject* tmp = PyDict_Values(kwds);
    auto tmp_ = make_safe(tmp);
    arg = PyList_GET_ITEM(tmp, 0);
  }

  if (PyBob_NumberCheck(arg))
    return apply_scalar(self, s_f_prime_str,
        boost::bind(&bob::learn::mlp::Cost::f_prime,
          self->cxx, _1, _2), args, kwds);

  return apply_array(self, s_f_prime_str,
      boost::bind(&bob::learn::mlp::Cost::f_prime,
        self->cxx, _1, _2), args, kwds);

}

PyDoc_STRVAR(s_error_str, "error");
PyDoc_STRVAR(s_error_doc,
"o.error(output, target, [result]) -> result\n\
\n\
Computes the back-propagated error for a given MLP ``output``\n\
layer.\n\
\n\
Computes the back-propagated error for a given MLP ``output``\n\
layer, given its activation function and outputs - i.e., the\n\
error back-propagated through the last layer neuron up to the\n\
synapse connecting the last hidden layer to the output layer.\n\
\n\
This implementation allows for optimization in the\n\
calculation of the back-propagated errors in cases where there\n\
is a possibility of mathematical simplification when using a\n\
certain combination of cost-function and activation. For\n\
example, using a ML-cost and a logistic activation function.\n\
\n\
Keyword arguments:\n\
\n\
output, ND array, float64 | scalar\n\
  Real output from the machine. May be a N-dimensional array\n\
  or a plain scalar.\n\
\n\
target, ND array, float64 | scalar\n\
  Target output you are training to achieve. The data type and\n\
  extents for this object must match that of ``target``.\n\
\n\
result (optional), ND array, float64\n\
  Where to place the result from the calculation. You can pass\n\
  this argument if the input are N-dimensional arrays.\n\
  Otherwise, it is an error to pass such a container. If the\n\
  inputs are arrays and an object for ``result`` is passed,\n\
  then its dimensions and data-type must match that of both\n\
  ``output`` and ``result``.\n\
\n\
Returns the cost as a scalar, if the input were scalars or as\n\
        an array with matching size of ``output`` and\n\
        ``target`` otherwise.\n\
");

static PyObject* PyBobLearnCost_error
(PyBobLearnCostObject* self, PyObject* args, PyObject* kwds) {

  PyObject* arg = 0; ///< borrowed (don't delete)
  if (PyTuple_Size(args)) arg = PyTuple_GET_ITEM(args, 0);
  else {
    PyObject* tmp = PyDict_Values(kwds);
    auto tmp_ = make_safe(tmp);
    arg = PyList_GET_ITEM(tmp, 0);
  }

  if (PyBob_NumberCheck(arg))
    return apply_scalar(self, s_error_str,
        boost::bind(&bob::learn::mlp::Cost::error, self->cxx, _1, _2), args, kwds);

  return apply_array(self, s_error_str,
      boost::bind(&bob::learn::mlp::Cost::error, self->cxx, _1, _2), args, kwds);

}

static PyMethodDef PyBobLearnCost_methods[] = {
  {
    s_f_str,
    (PyCFunction)PyBobLearnCost_f,
    METH_VARARGS|METH_KEYWORDS,
    s_f_doc
  },
  {
    s_f_prime_str,
    (PyCFunction)PyBobLearnCost_f_prime,
    METH_VARARGS|METH_KEYWORDS,
    s_f_prime_doc
  },
  {
    s_error_str,
    (PyCFunction)PyBobLearnCost_error,
    METH_VARARGS|METH_KEYWORDS,
    s_error_doc
  },
  {0} /* Sentinel */
};

static PyObject* PyBobLearnCost_new (PyTypeObject* type, PyObject*, PyObject*) {

  /* Allocates the python object itself */
  PyBobLearnCostObject* self = (PyBobLearnCostObject*)type->tp_alloc(type, 0);

  self->cxx.reset();

  return reinterpret_cast<PyObject*>(self);

}

PyObject* PyBobLearnCost_NewFromCost (boost::shared_ptr<bob::learn::mlp::Cost> p) {

  PyBobLearnCostObject* retval = (PyBobLearnCostObject*)PyBobLearnCost_new(&PyBobLearnCost_Type, 0, 0);

  retval->cxx = p;

  return reinterpret_cast<PyObject*>(retval);

}

PyTypeObject PyBobLearnCost_Type = {
    PyVarObject_HEAD_INIT(0, 0)
    s_cost_str,                               /* tp_name */
    sizeof(PyBobLearnCostObject),             /* tp_basicsize */
    0,                                        /* tp_itemsize */
    0,                                        /* tp_dealloc */
    0,                                        /* tp_print */
    0,                                        /* tp_getattr */
    0,                                        /* tp_setattr */
    0,                                        /* tp_compare */
    (reprfunc)PyBobLearnCost_Repr,            /* tp_repr */
    0,                                        /* tp_as_number */
    0,                                        /* tp_as_sequence */
    0,                                        /* tp_as_mapping */
    0,                                        /* tp_hash */
    (ternaryfunc)PyBobLearnCost_f,            /* tp_call */
    (reprfunc)PyBobLearnCost_Str,             /* tp_str */
    0,                                        /* tp_getattro */
    0,                                        /* tp_setattro */
    0,                                        /* tp_as_buffer */
    Py_TPFLAGS_DEFAULT | Py_TPFLAGS_BASETYPE, /* tp_flags */
    s_cost_doc,                               /* tp_doc */
    0,                                        /* tp_traverse */
    0,                                        /* tp_clear */
    (richcmpfunc)PyBobLearnCost_RichCompare,  /* tp_richcompare */
    0,                                        /* tp_weaklistoffset */
    0,                                        /* tp_iter */
    0,                                        /* tp_iternext */
    PyBobLearnCost_methods,                   /* tp_methods */
    0,                                        /* tp_members */
    0,                                        /* tp_getset */
    0,                                        /* tp_base */
    0,                                        /* tp_dict */
    0,                                        /* tp_descr_get */
    0,                                        /* tp_descr_set */
    0,                                        /* tp_dictoffset */
    (initproc)PyBobLearnCost_init,            /* tp_init */
    0,                                        /* tp_alloc */
    PyBobLearnCost_new,                       /* tp_new */
};

PyDoc_STRVAR(s_squareerror_str, BOB_EXT_MODULE_PREFIX ".SquareError");

PyDoc_STRVAR(s_squareerror_doc,
"SquareError(actfun) -> new SquareError functor\n\
\n\
Calculates the Square-Error between output and target.\n\
\n\
The square error is defined as follows:\n\
\n\
.. math::\n\
   J = \\frac{(\\hat{y} - y)^2}{2}\n\
\n\
where :math:`\\hat{y}` is the output estimated by your machine and\n\
:math:`y` is the expected output.\n\
\n\
Keyword arguments:\n\
\n\
actfun\n\
  The activation function object used at the last layer\n\
\n\
");

static int PyBobLearnSquareError_init
(PyBobLearnSquareErrorObject* self, PyObject* args, PyObject* kwds) {

  /* Parses input arguments in a single shot */
  static const char* const_kwlist[] = {"actfun", 0};
  static char** kwlist = const_cast<char**>(const_kwlist);

  PyObject* actfun = 0;

  if (!PyArg_ParseTupleAndKeywords(args, kwds, "O!", kwlist,
        &PyBobLearnActivation_Type, &actfun)) return -1;

  try {
    auto _actfun = reinterpret_cast<PyBobLearnActivationObject*>(actfun);
    self->cxx.reset(new bob::learn::mlp::SquareError(_actfun->cxx));
  }
  catch (std::exception& ex) {
    PyErr_SetString(PyExc_RuntimeError, ex.what());
  }
  catch (...) {
    PyErr_Format(PyExc_RuntimeError, "cannot create new object of type `%s' - unknown exception thrown", Py_TYPE(self)->tp_name);
  }

  self->parent.cxx = self->cxx;

  if (PyErr_Occurred()) return -1;

  return 0;

}

static void PyBobLearnSquareError_delete
(PyBobLearnSquareErrorObject* self) {

  self->parent.cxx.reset();
  self->cxx.reset();
  Py_TYPE(&self->parent)->tp_free((PyObject*)self);

}

PyTypeObject PyBobLearnSquareError_Type = {
    PyVarObject_HEAD_INIT(0, 0)
    s_squareerror_str,                        /*tp_name*/
    sizeof(PyBobLearnSquareErrorObject),      /*tp_basicsize*/
    0,                                        /*tp_itemsize*/
    (destructor)PyBobLearnSquareError_delete, /*tp_dealloc*/
    0,                                        /*tp_print*/
    0,                                        /*tp_getattr*/
    0,                                        /*tp_setattr*/
    0,                                        /*tp_compare*/
    0,                                        /*tp_repr*/
    0,                                        /*tp_as_number*/
    0,                                        /*tp_as_sequence*/
    0,                                        /*tp_as_mapping*/
    0,                                        /*tp_hash */
    0,                                        /*tp_call*/
    0,                                        /*tp_str*/
    0,                                        /*tp_getattro*/
    0,                                        /*tp_setattro*/
    0,                                        /*tp_as_buffer*/
    Py_TPFLAGS_DEFAULT | Py_TPFLAGS_BASETYPE, /*tp_flags*/
    s_squareerror_doc,                        /* tp_doc */
    0,		                                    /* tp_traverse */
    0,		                                    /* tp_clear */
    0,                                        /* tp_richcompare */
    0,		                                    /* tp_weaklistoffset */
    0,		                                    /* tp_iter */
    0,		                                    /* tp_iternext */
    0,                                        /* tp_methods */
    0,                                        /* tp_members */
    0,                                        /* tp_getset */
    0,                                        /* tp_base */
    0,                                        /* tp_dict */
    0,                                        /* tp_descr_get */
    0,                                        /* tp_descr_set */
    0,                                        /* tp_dictoffset */
    (initproc)PyBobLearnSquareError_init,     /* tp_init */
};

PyDoc_STRVAR(s_crossentropyloss_str, BOB_EXT_MODULE_PREFIX ".CrossEntropyLoss");

PyDoc_STRVAR(s_crossentropyloss_doc,
"CrossEntropyLoss(actfun) -> new CrossEntropyLoss functor\n\
\n\
Calculates the Cross Entropy Loss between output and target.\n\
\n\
The cross entropy loss is defined as follows:\n\
\n\
.. math::\n\
   J = - y \\cdot \\log{(\\hat{y})} - (1-y) \\log{(1-\\hat{y})}\n\
\n\
where :math:`\\hat{y}` is the output estimated by your machine and\n\
:math:`y` is the expected output.\n\
\n\
Keyword arguments:\n\
\n\
actfun\n\
  The activation function object used at the last layer. If you\n\
  set this to :py:class:`bob.learn.activation.Logistic`, a\n\
  mathematical simplification is possible in which\n\
  ``backprop_error()`` can benefit increasing the numerical\n\
  stability of the training process. The simplification goes\n\
  as follows:\n\
  \n\
  .. math::\n\
     b = \\delta \\cdot \\varphi'(z)\n\
  \n\
  But, for the cross-entropy loss: \n\
  \n\
  .. math::\n\
     \\delta = \\frac{\\hat{y} - y}{\\hat{y}(1 - \\hat{y})}\n\
  \n\
  and :math:`\\varphi'(z) = \\hat{y} - (1 - \\hat{y})`, so:\n\
  \n\
  .. math::\n\
     b = \\hat{y} - y\n\
\n\
");

static int PyBobLearnCrossEntropyLoss_init
(PyBobLearnCrossEntropyLossObject* self, PyObject* args, PyObject* kwds) {

  /* Parses input arguments in a single shot */
  static const char* const_kwlist[] = {"actfun", 0};
  static char** kwlist = const_cast<char**>(const_kwlist);

  PyObject* actfun = 0;

  if (!PyArg_ParseTupleAndKeywords(args, kwds, "O!", kwlist,
        &PyBobLearnActivation_Type, &actfun)) return -1;

  try {
    auto _actfun = reinterpret_cast<PyBobLearnActivationObject*>(actfun);
    self->cxx.reset(new bob::learn::mlp::CrossEntropyLoss(_actfun->cxx));
  }
  catch (std::exception& ex) {
    PyErr_SetString(PyExc_RuntimeError, ex.what());
  }
  catch (...) {
    PyErr_Format(PyExc_RuntimeError, "cannot create new object of type `%s' - unknown exception thrown", Py_TYPE(self)->tp_name);
  }

  self->parent.cxx = self->cxx;

  if (PyErr_Occurred()) return -1;

  return 0;

}

static void PyBobLearnCrossEntropyLoss_delete
(PyBobLearnCrossEntropyLossObject* self) {

  self->parent.cxx.reset();
  self->cxx.reset();
  Py_TYPE(&self->parent)->tp_free((PyObject*)self);

}

PyDoc_STRVAR(s_logistic_activation_str, "logistic_activation");
PyDoc_STRVAR(s_logistic_activation_doc,
"o.logistic_activation() -> bool\n\
\n\
Tells if this functor is set to operate together with a\n\
:py:class:`bob.learn.activation.Logistic` activation function.\n\
");

static PyObject* PyBobLearnCrossEntropyLoss_getLogisticActivation
(PyBobLearnCrossEntropyLossObject* self, void* /*closure*/) {
  if (self->cxx->logistic_activation()) Py_RETURN_TRUE;
  Py_RETURN_FALSE;
}

static PyGetSetDef PyBobLearnCrossEntropyLoss_getseters[] = {
    {
      s_logistic_activation_str,
      (getter)PyBobLearnCrossEntropyLoss_getLogisticActivation,
      0,
      s_logistic_activation_doc,
      0
    },
    {0}  /* Sentinel */
};

PyTypeObject PyBobLearnCrossEntropyLoss_Type = {
    PyVarObject_HEAD_INIT(0, 0)
    s_crossentropyloss_str,                        /*tp_name*/
    sizeof(PyBobLearnCrossEntropyLossObject),      /*tp_basicsize*/
    0,                                             /*tp_itemsize*/
    (destructor)PyBobLearnCrossEntropyLoss_delete, /*tp_dealloc*/
    0,                                             /*tp_print*/
    0,                                             /*tp_getattr*/
    0,                                             /*tp_setattr*/
    0,                                             /*tp_compare*/
    0,                                             /*tp_repr*/
    0,                                             /*tp_as_number*/
    0,                                             /*tp_as_sequence*/
    0,                                             /*tp_as_mapping*/
    0,                                             /*tp_hash */
    0,                                             /*tp_call*/
    0,                                             /*tp_str*/
    0,                                             /*tp_getattro*/
    0,                                             /*tp_setattro*/
    0,                                             /*tp_as_buffer*/
    Py_TPFLAGS_DEFAULT | Py_TPFLAGS_BASETYPE,      /*tp_flags*/
    s_crossentropyloss_doc,                        /* tp_doc */
    0,		                                         /* tp_traverse */
    0,		                                         /* tp_clear */
    0,                                             /* tp_richcompare */
    0,		                                         /* tp_weaklistoffset */
    0,		                                         /* tp_iter */
    0,		                                         /* tp_iternext */
    0,                                             /* tp_methods */
    0,                                             /* tp_members */
    PyBobLearnCrossEntropyLoss_getseters,          /* tp_getset */
    0,                                             /* tp_base */
    0,                                             /* tp_dict */
    0,                                             /* tp_descr_get */
    0,                                             /* tp_descr_set */
    0,                                             /* tp_dictoffset */
    (initproc)PyBobLearnCrossEntropyLoss_init,     /* tp_init */
};


// Distributed under the MIT License.
// See LICENSE.txt for details.

#pragma once

#include <boost/functional/hash.hpp>
#include <cstddef>
#include <limits>
#include <tuple>

#include "DataStructures/DataVector.hpp"
#include "DataStructures/FixedHashMap.hpp"
#include "DataStructures/Index.hpp"
#include "DataStructures/Tensor/Tensor.hpp"
#include "DataStructures/Variables.hpp"
#include "Domain/Structure/Direction.hpp"
#include "Domain/Structure/ElementId.hpp"
#include "Domain/Structure/MaxNumberOfNeighbors.hpp"
#include "Utilities/ContainerHelpers.hpp"
#include "Utilities/ErrorHandling/Assert.hpp"

namespace LinearSolver::Schwarz {

/// Identifies a subdomain region that overlaps with another element
template <size_t Dim>
using OverlapId = std::pair<Direction<Dim>, ElementId<Dim>>;

/// Data structure that can store the `ValueType` on each possible overlap of an
/// element-centered subdomain with its neighbors. Overlaps are identified by
/// their `OverlapId`.
template <size_t Dim, typename ValueType>
using OverlapMap =
    FixedHashMap<maximum_number_of_neighbors(Dim), OverlapId<Dim>, ValueType,
                 boost::hash<OverlapId<Dim>>>;

/*!
 * \brief The number of points that an overlap extends into the `volume_extent`
 *
 * In a dimension where an element has `volume_extent` points, the overlap
 * extent is the largest number under these constraints:
 *
 * - It is at most `max_overlap`.
 * - It is smaller than the `volume_extent`.
 *
 * This means the overlap extent is always smaller than the `volume_extent`. The
 * reason for this constraint is that we define the _width_ of the overlap as
 * the element-logical coordinate distance from the face of the element to the
 * first collocation point _outside_ the overlap extent. Therefore, even an
 * overlap region that covers the full element in width does not include the
 * collocation point on the opposite side of the element.
 *
 * Here's a few notes on the definition of the overlap extent and width:
 *
 * - A typical smooth weighting function goes to zero at the overlap width, so
 *   if the grid points located at the overlap width were included in the
 *   subdomain, their solutions would not contribute to the weighted sum of
 *   subdomain solutions.
 * - Defining the overlap width as the distance to the first point _outside_ the
 *   overlap extent makes it non-zero even for a single point of overlap into a
 *   Gauss-Lobatto grid (which has points located at the element face).
 * - Boundary contributions for many (but not all) discontinuous Galerkin
 *   schemes on Gauss-Lobatto grids are limited to the grid points on the
 *   element face, e.g. for a DG operator that is pre-multiplied by the mass
 *   matrix, or one where boundary contributions are lifted using the diagonal
 *   mass-matrix approximation. Not including the grid points facing away from
 *   the subdomain in the overlap allows to ignore that face altogether in the
 *   subdomain operator.
 */
size_t overlap_extent(size_t volume_extent, size_t max_overlap) noexcept;

/*!
 * \brief Total number of grid points in an overlap region that extends
 * `overlap_extent` points into the `volume_extents` from either side in the
 * `overlap_dimension`
 *
 * The overlap region has `overlap_extent` points in the `overlap_dimension`,
 * and `volume_extents` points in the other dimensions. The number of grid
 * points returned by this function is the product of these extents.
 */
template <size_t Dim>
size_t overlap_num_points(const Index<Dim>& volume_extents,
                          size_t overlap_extent,
                          size_t overlap_dimension) noexcept;

/*!
 * \brief Width of an overlap extending `overlap_extent` points into the
 * `collocation_points` from either side.
 *
 * The "width" of an overlap is the element-logical coordinate distance from the
 * element boundary to the first collocation point outside the overlap region in
 * the overlap dimension, i.e. the dimension perpendicular to the element face.
 * See `LinearSolver::Schwarz::overlap_extent` for details.
 *
 * This function assumes the `collocation_points` are mirrored around 0.
 */
double overlap_width(size_t overlap_extent,
                     const DataVector& collocation_points) noexcept;

/*!
 * \brief Iterate over grid points in a region that extends partially into the
 * volume
 *
 * Here's an example how to use this iterator:
 *
 * \snippet Test_OverlapHelpers.cpp overlap_iterator
 */
class OverlapIterator {
 public:
  template <size_t Dim>
  OverlapIterator(const Index<Dim>& volume_extents, size_t overlap_extent,
                  const Direction<Dim>& direction) noexcept;

  explicit operator bool() const noexcept;

  OverlapIterator& operator++();

  /// Offset into a DataVector that holds full volume data
  size_t volume_offset() const noexcept;

  /// Offset into a DataVector that holds data only on the overlap region
  size_t overlap_offset() const noexcept;

  void reset() noexcept;

 private:
  size_t size_ = std::numeric_limits<size_t>::max();
  size_t num_slices_ = std::numeric_limits<size_t>::max();
  size_t stride_ = std::numeric_limits<size_t>::max();
  size_t stride_count_ = std::numeric_limits<size_t>::max();
  size_t jump_ = std::numeric_limits<size_t>::max();
  size_t initial_offset_ = std::numeric_limits<size_t>::max();
  size_t volume_offset_ = std::numeric_limits<size_t>::max();
  size_t overlap_offset_ = std::numeric_limits<size_t>::max();
};

// @{
/// The part of the tensor data that lies within the overlap region
template <size_t Dim, typename DataType, typename... TensorStructure>
void data_on_overlap(const gsl::not_null<Tensor<DataType, TensorStructure...>*>
                         restricted_tensor,
                     const Tensor<DataType, TensorStructure...>& tensor,
                     const Index<Dim>& volume_extents,
                     const size_t overlap_extent,
                     const Direction<Dim>& direction) noexcept {
  for (OverlapIterator overlap_iterator{volume_extents, overlap_extent,
                                        direction};
       overlap_iterator; ++overlap_iterator) {
    for (size_t tensor_component = 0; tensor_component < tensor.size();
         ++tensor_component) {
      (*restricted_tensor)[tensor_component][overlap_iterator
                                                 .overlap_offset()] =
          tensor[tensor_component][overlap_iterator.volume_offset()];
    }
  }
}

template <size_t Dim, typename DataType, typename... TensorStructure>
Tensor<DataType, TensorStructure...> data_on_overlap(
    const Tensor<DataType, TensorStructure...>& tensor,
    const Index<Dim>& volume_extents, const size_t overlap_extent,
    const Direction<Dim>& direction) noexcept {
  Tensor<DataType, TensorStructure...> restricted_tensor{overlap_num_points(
      volume_extents, overlap_extent, direction.dimension())};
  data_on_overlap(make_not_null(&restricted_tensor), tensor, volume_extents,
                  overlap_extent, direction);
  return restricted_tensor;
}

namespace detail {
template <size_t Dim>
void data_on_overlap_impl(double* overlap_data, const double* volume_data,
                          size_t num_components,
                          const Index<Dim>& volume_extents,
                          size_t overlap_extent,
                          const Direction<Dim>& direction) noexcept;
}  // namespace detail

template <size_t Dim, typename OverlapTagsList, typename VolumeTagsList>
void data_on_overlap(
    const gsl::not_null<Variables<OverlapTagsList>*> overlap_data,
    const Variables<VolumeTagsList>& volume_data,
    const Index<Dim>& volume_extents, const size_t overlap_extent,
    const Direction<Dim>& direction) noexcept {
  constexpr size_t num_components =
      Variables<VolumeTagsList>::number_of_independent_components;
  ASSERT(volume_data.number_of_grid_points() == volume_extents.product(),
         "volume_data has wrong number of grid points.  Expected "
             << volume_extents.product() << ", got "
             << volume_data.number_of_grid_points());
  ASSERT(overlap_data->number_of_grid_points() ==
             overlap_num_points(volume_extents, overlap_extent,
                                direction.dimension()),
         "overlap_data has wrong number of grid points.  Expected "
             << overlap_num_points(volume_extents, overlap_extent,
                                   direction.dimension())
             << ", got " << overlap_data->number_of_grid_points());
  detail::data_on_overlap_impl(overlap_data->data(), volume_data.data(),
                               num_components, volume_extents, overlap_extent,
                               direction);
}

template <size_t Dim, typename TagsList>
Variables<TagsList> data_on_overlap(const Variables<TagsList>& volume_data,
                                    const Index<Dim>& volume_extents,
                                    const size_t overlap_extent,
                                    const Direction<Dim>& direction) noexcept {
  Variables<TagsList> overlap_data{overlap_num_points(
      volume_extents, overlap_extent, direction.dimension())};
  data_on_overlap(make_not_null(&overlap_data), volume_data, volume_extents,
                  overlap_extent, direction);
  return overlap_data;
}
// @}

namespace detail {
template <size_t Dim>
void add_overlap_data_impl(double* volume_data, const double* overlap_data,
                           size_t num_components,
                           const Index<Dim>& volume_extents,
                           size_t overlap_extent,
                           const Direction<Dim>& direction) noexcept;
}  // namespace detail

/// Add the `overlap_data` to the `volume_data`
template <size_t Dim, typename VolumeTagsList, typename OverlapTagsList>
void add_overlap_data(
    const gsl::not_null<Variables<VolumeTagsList>*> volume_data,
    const Variables<OverlapTagsList>& overlap_data,
    const Index<Dim>& volume_extents, const size_t overlap_extent,
    const Direction<Dim>& direction) noexcept {
  constexpr size_t num_components =
      Variables<VolumeTagsList>::number_of_independent_components;
  ASSERT(volume_data->number_of_grid_points() == volume_extents.product(),
         "volume_data has wrong number of grid points.  Expected "
             << volume_extents.product() << ", got "
             << volume_data->number_of_grid_points());
  ASSERT(overlap_data.number_of_grid_points() ==
             overlap_num_points(volume_extents, overlap_extent,
                                direction.dimension()),
         "overlap_data has wrong number of grid points.  Expected "
             << overlap_num_points(volume_extents, overlap_extent,
                                   direction.dimension())
             << ", got " << overlap_data.number_of_grid_points());
  detail::add_overlap_data_impl(volume_data->data(), overlap_data.data(),
                                num_components, volume_extents, overlap_extent,
                                direction);
}

// @{
/// Extend the overlap data to the full mesh by filling it with zeros outside
/// the overlap region
template <size_t Dim, typename ExtendedTagsList, typename OverlapTagsList>
void extended_overlap_data(
    const gsl::not_null<Variables<ExtendedTagsList>*> extended_data,
    const Variables<OverlapTagsList>& overlap_data,
    const Index<Dim>& volume_extents, const size_t overlap_extent,
    const Direction<Dim>& direction) noexcept {
  *extended_data = Variables<ExtendedTagsList>{volume_extents.product(), 0.};
  add_overlap_data(extended_data, overlap_data, volume_extents, overlap_extent,
                   direction);
}

template <size_t Dim, typename TagsList>
Variables<TagsList> extended_overlap_data(
    const Variables<TagsList>& overlap_data, const Index<Dim>& volume_extents,
    const size_t overlap_extent, const Direction<Dim>& direction) noexcept {
  Variables<TagsList> extended_data{volume_extents.product()};
  extended_overlap_data(make_not_null(&extended_data), overlap_data,
                        volume_extents, overlap_extent, direction);
  return extended_data;
}
// @}

}  // namespace LinearSolver::Schwarz


#ifndef TCM_CONSTANTS_HPP
#define TCM_CONSTANTS_HPP

#include <functional>
#include <type_traits>
#include <map>
#include <unordered_map>
#include <cmath>

#include <boost/program_options.hpp>
#include <boost/algorithm/string/erase.hpp>
#include <boost/algorithm/string/predicate.hpp>
#include <boost/range/adaptor/filtered.hpp>
#include <boost/range/adaptor/transformed.hpp>

///////////////////////////////////////////////////////////////////////////////
/// \file constants.hpp
/// \brief Defines the 'default' constants map.
///
/// \detail We implement some useful functions to manipulate _physical_
///         constants. These constants are represented as a symbol table
///         mapping #std::string to some numerical type _R representing real
///         field.
///
/// For the calculations, we're interested in the following constants:
/// |              Key             |     Notation     | Value(default) | Dimension |
/// | ---------------------------- | ---------------- | -------------- | --------- |
/// | `pi`                         | \f$\pi\f$        | 3.14159...     | 1         |
/// | `boltzmann-constant`         | \f$k_\text{B}\f$ | 8.61733...E-5  | eV        |
/// | `elementary-charge`          | \f$e\f$          | 1.60217...E-19 | C         |
/// | `planck-constant`            | \f$\hbar\f$      | 6.58212...E-16 | eV        |
/// | `self-interaction-potential` | \f$V_0\f$        | 15.78          | eV        |
/// | `temperature`                | \f$T\f$          | 300.0          | K         |
/// | `vacuum-permittivity`        | \f$\epsilon_0\f$ | 8.85419...E-12 | F/m       |
/// A table with these values can be easily obtained using the
/// #default_constants() function.
///
/// For different simulations different values of, for example, temperature
/// may be desired. To account for this case, we implement two additional
/// functions:
/// * #init_constants_options() simplifies the creation of
///   boost::program_options::options_description for the user to be able
///   to specify constants.
/// * #load_constants() allows the creation of constants table from
///   boost::program_options::variables_map, i.e. from command line arguments.
///
/// Here is an example how these functions can be used together:
/// \include src/constants_example.cpp
///////////////////////////////////////////////////////////////////////////////

namespace tcm {

//////////////////////////////////////////////////////////////////////////////
/// \brief Returns a table with default values for relevalnt physical
///        constants.

/// This function has the following behavior:
/// \snippet include/constants.hpp Default constants
///
/// \todo make this function return std::unordered_map in place of std::map.
/// Or even better, make the container type a template parameter.
//////////////////////////////////////////////////////////////////////////////
template <class _R>
auto default_constants() -> std::map<std::string, _R>
{
//! [Default constants]
	std::map<std::string, _R> constants =
		{ {"pi",                          M_PI}
		, {"boltzmann-constant",          8.6173303E-5}
		, {"chemical-potential",          0.4}
		, {"elementary-charge",           1.6021766208E-19}
		, {"planck-constant",             6.582119514E-16}
		, {"self-interaction-potential",  15.78}
		, {"temperature",                 300.0}
		, {"vacuum-permittivity",         8.854187817E-12}
		, {"tau",                         6.0E-3}
		};
	return constants;
//! [Default constants]
}

///////////////////////////////////////////////////////////////////////////////
/// \brief Simplifies the creation of #options_descriptions concerning
///        constants.

/// Creates a #boost::program_options::options_description called
/// _Simulation Constants_ with options of the form "in.constants.KEY" where
/// KEYs are the names of the constants such as `pi` or `planck-constant`.
/// Each of these options has a default value provided by the
/// #default_constants() function. Template parameter `_R` specifies the type
/// of the constants. It is used to create the default constants table and to
/// extract values from `any` where program_options stores the values.
///////////////////////////////////////////////////////////////////////////////
template <class _R>
auto init_constants_options() -> boost::program_options::options_description
{
	using namespace boost::program_options;
	options_description description{"Simulation Constants"};
	auto const cs = tcm::default_constants<_R>();
	description.add_options()
		( "in.constants.pi"
		, value<_R>()->default_value(cs.at("pi"))
		, "PI." )
		( "in.constants.boltzmann-constant"
		, value<_R>()->default_value(cs.at("boltzmann-constant"))
		, "Boltzmann constant." )
		( "in.constants.chemical-potential"
		, value<_R>()->default_value(cs.at("chemical-potential"))
		, "Chemical potential." )
		( "in.constants.elementary-charge"
		, value<_R>()->default_value(cs.at("elementary-charge"))
		, "Elementary charge." )
		( "in.constants.planck-constant"
		, value<_R>()->default_value(cs.at("planck-constant"))
		, "Planck constant" )
		( "in.constants.self-interaction-potential"
		, value<_R>()->default_value(cs.at("self-interaction-potential"))
		, "Self interaction Coulomb potential." )
		( "in.constants.temperature"
		, value<_R>()->default_value(cs.at("temperature"))
		, "Temperature" )
		( "in.constants.vacuum-permittivity"
		, value<_R>()->default_value(cs.at("vacuum-permittivity"))
		, "Vacuum permittivity." )
		( "in.constants.tau"
		, value<_R>()->default_value(cs.at("tau"))
		, "Relaxation time tau." );
	return description;
}

///////////////////////////////////////////////////////////////////////////////
/// \brief Creates a constants table from a
///        #boost::program_options::variables_map.

/// This function iterates over entries in \p vm, searches for all options
/// that are of the form "in.constants.*", converts the values from
/// boost::program_options::variable_value to _To and returns a table of
/// type _Map consisting only of those options.
/// \snippet include/constants.hpp Load constants
///////////////////////////////////////////////////////////////////////////////
template< class _To
        , class _From = _To
        , class _Map = std::unordered_map<std::string, _To>
        >
auto load_constants(boost::program_options::variables_map const& vm) -> _Map
{
	static_assert( std::is_same<typename _Map::key_type, std::string>::value
	             , "Wrong key_type: must be std::string." );
	static_assert( std::is_same<typename _Map::mapped_type, _To>::value
	             , "Wrong mapped_type: must be _To." );
//! [Load constants]
	constexpr char const prefix[] = "in.constants.";
	auto const is_constant = [](auto const& p) {
		return boost::algorithm::starts_with(p.first, prefix);
	};

	using value_type = boost::program_options::variables_map::value_type;
	auto const mk_constant = [](value_type const& p) {
		return std::make_pair(boost::algorithm::erase_head_copy(
			p.first, std::char_traits<char>::length(prefix)),
			p.second.as<_From>());
	};

	auto rng = vm | boost::adaptors::filtered(std::cref(is_constant))
	              | boost::adaptors::transformed(std::cref(mk_constant));
	return _Map{rng.begin(), rng.end()};
//! [Load constants]
}

///////////////////////////////////////////////////////////////////////////////
/// \brief Checks whether \p key is in \p constants map and throws if it is
///        not.

/// \snippet include/constants.hpp Require constant
///////////////////////////////////////////////////////////////////////////////
template<class _Map>
auto require( std::string const& func_name
            , _Map const& constants_map
            , std::string const& key ) -> void
{
//! [Require constant]
	if (!constants_map.count(key)) {
		throw std::runtime_error{ "Constant `" + key + "` is required to run "
		                          "`" + func_name + "`!" };
	}
//! [Require constant]
}

} // namespace tcm

#endif // TCM_CONSTANTS_HPP


// TopicModelLocalStepCPPX.cpp
// Define this symbol to enable runtime tests for allocations
#define EIGEN_RUNTIME_NO_MALLOC

#include <math.h>
#include <time.h>
#include "Eigen/Dense"
#include <boost/math/special_functions/digamma.hpp>
#include <boost/math/special_functions/gamma.hpp>
#include "fastexp.h"
using namespace Eigen;
using namespace std;

// ======================================================== Declare funcs
// ======================================================== visible externally

extern "C" {
    void sparseLocalStepManyDocs_ActiveOnly(
        double* alphaEbeta_IN,
        double* Eloglik_IN,
        double* word_count_IN,
        int* word_id_IN,
        int* doc_range_IN,
        int nnzPerRow,
        int Nall,
        int K,
        int D,
        int V,
        int nCoordAscentIterLP,
        double convThrLP,
        int initProbsToEbeta,
        double* topicCount_OUT,
        double* spResp_data_OUT,
        int* spResp_colids_OUT,
        int* numIterVec_OUT,
        double* maxDiffVec_OUT,
        int numRestarts,
        int* rAcceptVec_IN,
        int* rTrialVec_IN,
        int REVISE_FIRST,
        int REVISE_EVERY,
        int verbose
    );
}

// ======================================================== Custom Type Defs
// ========================================================
// Simple names for array types
typedef Matrix<double, Dynamic, Dynamic, RowMajor> Mat2D_d;
typedef Matrix<double, 1, Dynamic, RowMajor> Mat1D_d;
typedef Array<double, Dynamic, Dynamic, RowMajor> Arr2D_d;
typedef Array<double, 1, Dynamic, RowMajor> Arr1D_d;
typedef Array<int, 1, Dynamic, RowMajor> Arr1D_i;

// Simple names for array types with externally allocated memory
typedef Map<Mat2D_d> ExtMat2D_d;
typedef Map<Mat1D_d> ExtMat1D_d;
typedef Map<Arr2D_d> ExtArr2D_d;
typedef Map<Arr1D_d> ExtArr1D_d;
typedef Map<Arr1D_i> ExtArr1D_i;

double calcElapsedTime(timespec start_time, timespec end_time) {
    double diffSec = (double) end_time.tv_sec - start_time.tv_sec;
    double diffNano = (double) end_time.tv_nsec - start_time.tv_nsec;
    return diffSec + diffNano / 1.0e9;
}

struct LessThanFor1DArray {
    const double* xptr;

    LessThanFor1DArray(const double * xptrIN) {
        xptr = xptrIN;
    }

    bool operator()(int i, int j) {
        return xptr[i] < xptr[j];
    }

};

struct GreaterThanFor1DArray {
    const double* xptr;

    GreaterThanFor1DArray(const double * xptrIN) {
        xptr = xptrIN;
    }

    bool operator()(int i, int j) {
        return xptr[i] > xptr[j];
    }
};

struct Argsortable1DArray {
    double* xptr;
    int* iptr;
    int size;

    // Constructor
    Argsortable1DArray(double* xptrIN, int sizeIN) {
        xptr = xptrIN;
        size = sizeIN;
        iptr = new int[size];
        resetIndices(size);
    }

    // Helper method: reset iptr array to 0, 1, ... K-1
    void resetIndices(int cursize) {
        assert(cursize <= size);
        for (int i = 0; i < cursize; i++) {
            iptr[i] = i;
        }
    }

    void pprint() {
        for (int i = 0; i < size; i++) {
            printf("%03d:% 05.2f ",
                this->iptr[i],
                this->xptr[this->iptr[i]]);
        }
        printf("\n");
    }

    void argsort() {
        this->argsort_AscendingOrder();
    }

    void argsort_AscendingOrder() {
        std::sort(
            this->iptr,
            this->iptr + this->size,
            LessThanFor1DArray(this->xptr)
            );
    }

    void argsort_DescendingOrder() {
        std::sort(
            this->iptr,
            this->iptr + this->size,
            GreaterThanFor1DArray(this->xptr)
            );
    }

    void findSmallestL(int L) {
        assert(L >= 0);
        assert(L < this->size);
        std::nth_element(
            this->iptr,
            this->iptr + L,
            this->iptr + this->size,
            LessThanFor1DArray(this->xptr)
            );
    }

    void findLargestL(int L, int Kactive) {
        assert(L >= 0);
        assert(L <= Kactive);
        std::nth_element(
            this->iptr,
            this->iptr + Kactive - L,
            this->iptr + Kactive,
            LessThanFor1DArray(this->xptr)
            );
    }
};

void precomputeTopLRespForEachVocabTerm(
        int nnzPerRow,
        int V,
        int K,
        ExtArr2D_d & Eloglik,
        ExtArr1D_d & alphaEbeta,
        Arr1D_d & ElogProb,
        Arr1D_d & logScores_n,
        Arr1D_d & tempScores_n,
        Arr1D_d & termResp_data,
        Arr1D_i & termResp_colids
        )
{
    ElogProb = alphaEbeta.log();
    for (int v = 0; v < V; v++) {
        int argmax_n = -1;
        double maxScore_n;
        double curScore;
        for (int k = 0; k < K; k++) {
            curScore = ElogProb(k) + Eloglik(v, k);
            if (k == 0 || curScore > maxScore_n) {
                maxScore_n = curScore;
                if (nnzPerRow == 1) {
                    argmax_n = k;
                }
            }
            logScores_n(k) = curScore;
        } // end loop over K
        if (nnzPerRow == 1) {
            termResp_data(v) = 1.0;
            termResp_colids(v) = argmax_n;
        } else {
            // Find the top L entries in logScores_n
            // Copy current row over into a temp buffer
            std::copy(
                logScores_n.data(),
                logScores_n.data() + K,
                tempScores_n.data());
            // Sort the data in the temp buffer (in place)
            std::nth_element(
                tempScores_n.data(),
                tempScores_n.data() + K - nnzPerRow,
                tempScores_n.data() + K);
            // Walk thru this row and find the top "nnzPerRow" positions
            double pivotScore = tempScores_n(K - nnzPerRow);
            double sumResp_n = 0.0;
            int termResp_start = v * nnzPerRow;
            int termResp_id = termResp_start;
            for (int k = 0; k < K; k++) {
                if (logScores_n(k) >= pivotScore) {
                    termResp_data(termResp_id) = \
                        fastexp(logScores_n(k) - maxScore_n);
                    termResp_colids(termResp_id) = k;
                    sumResp_n += termResp_data(termResp_id);
                    termResp_id += 1;
                }
            }
            assert(termResp_id - v * nnzPerRow == nnzPerRow);
            for (termResp_id = v * nnzPerRow;
                    termResp_id < (v+1) * nnzPerRow;
                    termResp_id++) {
                termResp_data(termResp_id) /= sumResp_n;
            }

        }
    }
}

int updateActiveSetForDoc(
        int d,
        ExtArr2D_d & topicCount,
        Arr1D_i & activeTopics_d,
        Arr1D_i & spareActiveTopics_d,
        double ACTIVE_THR,
        int prevKactive,
        int nnzPerRow
        )
{
    int newKactive = 0;
    int ia = 0; // spare inactive topics
    for (int a = 0; a < prevKactive; a++) {
        int k = activeTopics_d(a);
        if (topicCount(d, k) > ACTIVE_THR) {
            activeTopics_d(newKactive) = k;
            newKactive += 1;
        } else if (newKactive < nnzPerRow - ia) {
            spareActiveTopics_d(ia) = k;
            ia += 1;
        }
    }

    int Kactive = newKactive;
    while (Kactive < nnzPerRow) {
        int k = spareActiveTopics_d(Kactive - newKactive);
        activeTopics_d(Kactive) = k;
        Kactive++;
    }
    return Kactive;
}

/*
 * Initialize activeTopics_d from provided topicCount
 * All K possible topics are available.
 */
int updateActiveSetForDocFromScratch(
        int d,
        ExtArr2D_d & topicCount,
        Arr1D_i & activeTopics_d,
        Arr1D_i & spareActiveTopics_d,
        double ACTIVE_THR,
        int K,
        int nnzPerRow
        )
{
    int newKactive = 0;
    int ia = 0; // spare inactive topics
    for (int k = 0; k < K; k++) {
        if (topicCount(d, k) > ACTIVE_THR) {
            activeTopics_d(newKactive) = k;
            newKactive += 1;
        } else if (newKactive < nnzPerRow - ia) {
            spareActiveTopics_d(ia) = k;
            ia += 1;
        }
    }

    int Kactive = newKactive;
    while (Kactive < nnzPerRow) {
        int k = spareActiveTopics_d(Kactive - newKactive);
        activeTopics_d(Kactive) = k;
        Kactive++;
    }
    return Kactive;
}

double calcELBOForDoc(
    int d,
    ExtArr1D_d & alphaEbeta,
    ExtArr2D_d & topicCount,
    Arr1D_d & ElogProb_d,
    Arr1D_i & activeTopics_d,
    double totalLogSumResp,
    double sum_gammalnalphaEbeta,
    int Kactive
    )
{
    double ELBO = sum_gammalnalphaEbeta;
    for (int ka = 0; ka < Kactive; ka++) {
        int k = activeTopics_d(ka);
        ELBO += (
            boost::math::lgamma(topicCount(d, k) + alphaEbeta(k))
            - boost::math::lgamma(alphaEbeta(k))
            - topicCount(d, k) * ElogProb_d(k)
            );
    }
    ELBO += totalLogSumResp;
    return ELBO;
}

double updateAssignmentsForDoc_ReviseActiveSetDupOK(
    int d,
    int start_d,
    int N_d,
    int nnzPerRow,
    int Kactive,
    ExtArr1D_d alphaEbeta,
    ExtArr2D_d Eloglik,
    ExtArr1D_d word_count,
    ExtArr1D_i word_id,
    ExtArr2D_d & topicCount,
    ExtArr1D_d & spResp_data,
    ExtArr1D_i & spResp_colids,
    Arr1D_i & activeTopics_d,
    Arr1D_d & ElogProb_d,
    Arr1D_d & logScores_n,
    Argsortable1DArray & logScoresHandler,
    int initProbsToEbeta,
    int doTrackELBO
    )
{
    double totalLogSumResp = 0.0;

    // Update ElogProb_d for active topics
    if (initProbsToEbeta == 1) {
        for (int ka = 0; ka < Kactive; ka++) {
            int k = activeTopics_d(ka);
            ElogProb_d(k) = log(alphaEbeta(k));
        }
    }  else {
        for (int ka = 0; ka < Kactive; ka++) {
            int k = activeTopics_d(ka);
            ElogProb_d(k) = boost::math::digamma(
                topicCount(d, k) + alphaEbeta(k));
        }
    }
    // RESET topicCounts for doc d
    topicCount.row(d).fill(0);

    // Update Resp_d for active topics
    // UPDATE assignments, obeying sparsity constraint
    for (int n = start_d; n < start_d + N_d; n++) {
        int spRind_dn_start = n * nnzPerRow;
        double w_ct = word_count(n);
        int w_id = word_id(n);
        int argmax_n = 0;
        double maxScore_n;
        for (int ka = 0; ka < Kactive; ka++) {
            int k = activeTopics_d(ka);
            logScores_n(ka) = ElogProb_d(k) + Eloglik(w_id, k);
            if (ka == 0 || logScores_n(ka) > maxScore_n) {
                maxScore_n = logScores_n(ka);
                if (nnzPerRow == 1) {
                    argmax_n = k;
                }
            }
        }
        if (nnzPerRow == 1) {
            spResp_data(spRind_dn_start) = 1.0;
            spResp_colids(spRind_dn_start) = argmax_n;
            topicCount(d, argmax_n) += w_ct;
            if (doTrackELBO) {
                totalLogSumResp += w_ct * maxScore_n;
            }
        } else {
            // Use handler's built-in findLargestL functionality
            logScoresHandler.resetIndices(Kactive);
            logScoresHandler.findLargestL(nnzPerRow, Kactive);

            // Read off the top L values into spR data structures
            int spRind_dn = spRind_dn_start;
            double sumResp_n = 0.0;
            for (int j = 0; j < nnzPerRow; j++) {
                int ka = logScoresHandler.iptr[Kactive - j - 1];
                spResp_data(spRind_dn) = fastexp(
                    logScoresHandler.xptr[ka] - maxScore_n);
                spResp_colids(spRind_dn) = activeTopics_d(ka);
                sumResp_n += spResp_data(spRind_dn);
                spRind_dn += 1;
            }
            for (spRind_dn = spRind_dn_start;
                    spRind_dn < spRind_dn_start + nnzPerRow; spRind_dn++) {
                spResp_data(spRind_dn) /= sumResp_n;
                topicCount(d, spResp_colids(spRind_dn)) += \
                    w_ct * spResp_data(spRind_dn);
            }
            if (doTrackELBO) {
                totalLogSumResp += w_ct * (maxScore_n + log(sumResp_n));
            }
        } // end if statement branch for nnz > 1
    } // end for loop over tokens in this doc
    return totalLogSumResp;
}

double updateAssignmentsForDoc_ReviseActiveSet(
    int d,
    int start_d,
    int N_d,
    int nnzPerRow,
    int Kactive,
    ExtArr1D_d alphaEbeta,
    ExtArr2D_d Eloglik,
    ExtArr1D_d word_count,
    ExtArr1D_i word_id,
    ExtArr2D_d & topicCount,
    ExtArr1D_d & spResp_data,
    ExtArr1D_i & spResp_colids,
    Arr1D_i & activeTopics_d,
    Arr1D_d & ElogProb_d,
    Arr1D_d & logScores_n,
    Arr1D_d & tempScores_n,
    int initProbsToEbeta,
    int doTrackELBO
    )
{
    double totalLogSumResp = 0.0;

    // Update ElogProb_d for active topics
    if (initProbsToEbeta == 1) {
        for (int ka = 0; ka < Kactive; ka++) {
            int k = activeTopics_d(ka);
            ElogProb_d(k) = log(alphaEbeta(k));
        }
    }  else {
        for (int ka = 0; ka < Kactive; ka++) {
            int k = activeTopics_d(ka);
            ElogProb_d(k) = boost::math::digamma(
                topicCount(d, k) + alphaEbeta(k));
        }
    }
    // RESET topicCounts for doc d
    topicCount.row(d).fill(0);

    // Update Resp_d for active topics
    // UPDATE assignments, obeying sparsity constraint
    for (int n = start_d; n < start_d + N_d; n++) {
        int spRind_dn_start = n * nnzPerRow;
        double w_ct = word_count(n);
        int w_id = word_id(n);
        int argmax_n = 0;
        double maxScore_n;
        for (int ka = 0; ka < Kactive; ka++) {
            int k = activeTopics_d(ka);
            logScores_n(ka) = ElogProb_d(k) + Eloglik(w_id, k);
            if (ka == 0 || logScores_n(ka) > maxScore_n) {
                maxScore_n = logScores_n(ka);
                if (nnzPerRow == 1) {
                    argmax_n = k;
                }
            }
        }
        if (nnzPerRow == 1) {
            spResp_data(spRind_dn_start) = 1.0;
            spResp_colids(spRind_dn_start) = argmax_n;
            topicCount(d, argmax_n) += w_ct;
            if (doTrackELBO) {
                totalLogSumResp += w_ct * maxScore_n;
            }
        } else {
            // Find the top L entries in logScores_n
            // Copy current row over into a temp buffer
            std::copy(
                logScores_n.data(),
                logScores_n.data() + Kactive,
                tempScores_n.data());
            // Sort the data in the temp buffer (in place)
            std::nth_element(
                tempScores_n.data(),
                tempScores_n.data() + Kactive - nnzPerRow,
                tempScores_n.data() + Kactive);
            // Walk thru this row and find the top "nnzPerRow" positions
            double pivotScore = tempScores_n(Kactive - nnzPerRow);

            int spRind_dn = spRind_dn_start;
            double sumResp_n = 0.0;
            for (int ka = 0; ka < Kactive; ka++) {
                if (logScores_n(ka) >= pivotScore) {
                    spResp_data(spRind_dn) = \
                        fastexp(logScores_n(ka) - maxScore_n);
                    spResp_colids(spRind_dn) = activeTopics_d(ka);
                    sumResp_n += spResp_data(spRind_dn);
                    spRind_dn += 1;
                }
            }
            assert(spRind_dn - spRind_dn_start == nnzPerRow);
            for (spRind_dn = spRind_dn_start;
                    spRind_dn < spRind_dn_start + nnzPerRow; spRind_dn++) {
                spResp_data(spRind_dn) /= sumResp_n;
                topicCount(d, spResp_colids(spRind_dn)) += \
                    w_ct * spResp_data(spRind_dn);
            }
            if (doTrackELBO) {
                totalLogSumResp += w_ct * (maxScore_n + log(sumResp_n));
            }
        } // end if statement branch for nnz > 1
    } // end for loop over tokens in this doc
    return totalLogSumResp;
}

/* Update spResp_data values, keeping spResp_colids fixed.
 *
 */
void updateAssignmentsForDoc_FixPerTokenActiveSet(
    int d,
    int start_d,
    int N_d,
    int nnzPerRow,
    int Kactive,
    ExtArr1D_d alphaEbeta,
    ExtArr2D_d Eloglik,
    ExtArr1D_d word_count,
    ExtArr1D_i word_id,
    ExtArr2D_d & topicCount,
    ExtArr1D_d & spResp_data,
    ExtArr1D_i & spResp_colids,
    Arr1D_i & activeTopics_d,
    Arr1D_d & ElogProb_d,
    Arr1D_d & logScores_n
    )
{
    assert(nnzPerRow > 1);
    //timespec start_time, end_time;
    //clock_gettime(CLOCK_PROCESS_CPUTIME_ID, &start_time);
    // Update ElogProb_d for active topics
    for (int ka = 0; ka < Kactive; ka++) {
        int k = activeTopics_d(ka);
        ElogProb_d(k) = boost::math::digamma(
            topicCount(d, k) + alphaEbeta(k));
    }
    //clock_gettime(CLOCK_PROCESS_CPUTIME_ID, &end_time);
    //double timeSpent_Digamma = calcElapsedTime(start_time, end_time);

    // RESET topicCounts for doc d
    topicCount.row(d).fill(0);

    //clock_gettime(CLOCK_PROCESS_CPUTIME_ID, &start_time);
    // Update Resp_d for active topics
    for (int n = start_d; n < start_d + N_d; n++) {
        int spRind_dn_start = n * nnzPerRow;
        double w_ct = word_count(n);
        int w_id = word_id(n);
        double maxScore_n;
        for (int ka = 0; ka < nnzPerRow; ka++) {
            int k = spResp_colids(spRind_dn_start + ka);
            logScores_n(ka) = ElogProb_d(k) + Eloglik(w_id, k);
            if (ka == 0 || logScores_n(ka) > maxScore_n) {
                maxScore_n = logScores_n(ka);
            }
        }

        double sumResp_n = 0.0;
        for (int ka = 0; ka < nnzPerRow; ka++) {
            spResp_data(spRind_dn_start + ka) = \
                fastexp(logScores_n(ka) - maxScore_n);
            sumResp_n += spResp_data(spRind_dn_start + ka);
        }

        for (int ka = 0; ka < nnzPerRow; ka++) {
            spResp_data(spRind_dn_start + ka) /= sumResp_n;
            topicCount(d, spResp_colids(spRind_dn_start + ka)) += \
                w_ct * spResp_data(spRind_dn_start + ka);
        }

    } // end for loop over tokens in this doc
    /*clock_gettime(CLOCK_PROCESS_CPUTIME_ID, &end_time);
    double timeSpent_Resp = calcElapsedTime(start_time, end_time);
    logScores_n(0) = timeSpent_Digamma;
    logScores_n(1) = timeSpent_Resp;
    */
}

double tryRestartsForDoc(
    int d,
    int start_d,
    int N_d,
    int nnzPerRow,
    int Kactive,
    ExtArr1D_d alphaEbeta,
    ExtArr2D_d Eloglik,
    ExtArr1D_d word_count,
    ExtArr1D_i word_id,
    ExtArr2D_d & topicCount,
    ExtArr1D_d & spResp_data,
    ExtArr1D_i & spResp_colids,
    Arr1D_i & activeTopics_d,
    Arr1D_d & prevTopicCount_d,
    Arr1D_d & ElogProb_d,
    Arr1D_d & logScores_n,
    Argsortable1DArray & logScoresHandler,
    //Arr1D_d & tempScores_n,
    ExtArr1D_i & rAcceptVec,
    ExtArr1D_i & rTrialVec,
    int numRestarts,
    double sum_gammalnalphaEbeta,
    int verbose
    )
{
    // Find active topics eligible for sparse restart
    int nAccept = 0;
    int nTrial = 0;

    double CANDIDATE_THR = 1e-10;
    double curELBO = 0.0;
    double totalLogSumResp = 0.0;
    prevTopicCount_d.fill(0);
    if (verbose > 1 && d < 1) {
        printf("\nSPARSE RESTARTS at doc %d!!\n", d);
    }
    if (verbose > 1 && d < 1) {
        for (int ka = 0; ka < Kactive; ka++) {
            int k = activeTopics_d(ka);
            printf(" %03d:%06.2f", k, topicCount(d,k));
        }
        printf("\n");
    }
    for (int riter = 0; riter < numRestarts; riter++) {
        // Seach for smallest topic we have yet to try
        double minVal = N_d + 1.0; // Will never be a value in topicCount
        int numAboveThr = 0;
        int chosenTopicID = -1;
        for (int ka = 0; ka < Kactive; ka++) {
            int k = activeTopics_d(ka);
            if (topicCount(d, k) > CANDIDATE_THR) {
                numAboveThr += 1;
                if (topicCount(d, k) < minVal) {
                    minVal = topicCount(d, k);
                    chosenTopicID = k;
                }
            }
        }

        // Need at least two topics with non-neglible size to try restarting.
        // Otherwise, we should just quit
        if (numAboveThr <= 1 || chosenTopicID < 0) {
            break;
        }

        if (riter == 0) {
            totalLogSumResp = updateAssignmentsForDoc_ReviseActiveSetDupOK(
                d, start_d, N_d, nnzPerRow, Kactive,
                alphaEbeta, Eloglik, word_count, word_id,
                topicCount, spResp_data, spResp_colids,
                activeTopics_d, ElogProb_d, logScores_n, logScoresHandler,
                0, 1);
            // ELBO for current configuration
            curELBO = calcELBOForDoc(d, alphaEbeta, topicCount,
                ElogProb_d, activeTopics_d,
                totalLogSumResp, sum_gammalnalphaEbeta, Kactive);
            // Remember the best-known topic-count vector!
            for (int ka = 0; ka < Kactive; ka++) {
                int k = activeTopics_d(ka);
                prevTopicCount_d(k) = topicCount(d, k);
            }
        }

        if (verbose > 1 && d < 1) {
            printf("START: best known counts. ELBO=%.5e \n", curELBO);
            for (int ka = 0; ka < Kactive; ka++) {
                int k = activeTopics_d(ka);
                printf("%02d:%06.2f ", k, topicCount(d, k));
            }
            printf("\n");
        }

        // Force chosen topic to zero
        topicCount(d, chosenTopicID) = 0.0;
        if (verbose > 1 && d < 1) {
            printf(
                "RESTART: Set index %d to zero (%d left)\n",
                chosenTopicID, numAboveThr - 1);
            for (int ka = 0; ka < Kactive; ka++) {
                int k = activeTopics_d(ka);
                printf("%02d:%06.2f ", k, topicCount(d, k));
            }
            printf("\n");
        }
        // Run inference forward from forced new location
        int NSTEP = 2;
        for (int step = 0; step < NSTEP; step++) {
            totalLogSumResp = updateAssignmentsForDoc_ReviseActiveSetDupOK(
                d, start_d, N_d, nnzPerRow, Kactive,
                alphaEbeta, Eloglik, word_count, word_id,
                topicCount, spResp_data, spResp_colids,
                activeTopics_d, ElogProb_d, logScores_n, logScoresHandler,
                0, step == NSTEP - 1);
        }
        // If the change is small, abandon current proposal
        double propELBO;
        if (abs(prevTopicCount_d(chosenTopicID) - topicCount(d, chosenTopicID)) < 1e-5) {
            propELBO = curELBO;
        } else {
            propELBO = calcELBOForDoc(d, alphaEbeta, topicCount,
                ElogProb_d, activeTopics_d,
                totalLogSumResp, sum_gammalnalphaEbeta, Kactive);
        }
        if (verbose > 1 && d < 1) {
            for (int ka = 0; ka < Kactive; ka++) {
                int k = activeTopics_d(ka);
                printf("%02d:%06.2f ", k, topicCount(d, k));
            }
            printf("\n");
            printf("propELBO % .6e\n", propELBO);
            printf(" curELBO % .6e\n", curELBO);
            if (propELBO > curELBO) {
                printf("beforeCount: %.6f\n", prevTopicCount_d(chosenTopicID));
                printf(" afterCount: %.6f\n", topicCount(d, chosenTopicID));
                printf("gainELBO % .6e ACCEPTED \n", propELBO - curELBO);
            } else {
                printf("gainELBO % .6e rejected \n", propELBO - curELBO);
            }
        }

        // Reset threshold for which topic to choose next
        CANDIDATE_THR = prevTopicCount_d(chosenTopicID);

        // If accepted, set current best doc-topic counts to latest proposal
        // Otherwise, reset the starting point for the next proposal.
        if (propELBO > curELBO) {
            curELBO = propELBO;
            for (int ka = 0; ka < Kactive; ka++) {
                int k = activeTopics_d(ka);
                prevTopicCount_d(k) = topicCount(d, k);
            }
            nAccept += 1;
        } else {
            for (int ka = 0; ka < Kactive; ka++) {
                int k = activeTopics_d(ka);
                topicCount(d, k) = prevTopicCount_d(k);
            }
        }
        nTrial += 1;

        if (verbose > 1 && d < 1) {
            printf("topicCount(d,k) = %.5e\n", topicCount(d,chosenTopicID));
            printf("prevTopicCount(d,k) = %.5e\n", prevTopicCount_d(chosenTopicID));
            printf("NEW THR = %.5e\n", CANDIDATE_THR);
        }

    } // end loop over restarts

    if (numRestarts > 0 && nTrial > 0) {
        rAcceptVec(0) += nAccept;
        rTrialVec(0) += nTrial;
        for (int ka = 0; ka < Kactive; ka++) {
            int k = activeTopics_d(ka);
            topicCount(d, k) = prevTopicCount_d(k);
        }
        // Final update! Make sure spResp reflects best topicCounts found
        updateAssignmentsForDoc_ReviseActiveSetDupOK(
            d, start_d, N_d, nnzPerRow, Kactive,
            alphaEbeta, Eloglik, word_count, word_id,
            topicCount, spResp_data, spResp_colids,
            activeTopics_d, ElogProb_d, logScores_n, logScoresHandler,
            0, 0);
    } // end if to synchronize at best known assignments

    return curELBO;
}

void sparseLocalStepManyDocs_ActiveOnly(
        double* alphaEbeta_IN,
        double* Eloglik_IN,
        double* word_count_IN,
        int* word_id_IN,
        int* doc_range_IN,
        int nnzPerRow,
        int Nall,
        int K,
        int D,
        int V,
        int nCoordAscentIterLP,
        double convThrLP,
        int initProbsToEbeta,
        double* topicCount_OUT,
        double* spResp_data_OUT,
        int* spResp_colids_OUT,
        int* numIterVec_OUT,
        double* maxDiffVec_OUT,
        int numRestarts,
        int* rAcceptVec_IN,
        int* rTrialVec_IN,
        int REVISE_FIRST,
        int REVISE_EVERY,
        int verbose
        )
{
    /*timespec start_time, end_time;
    double timeSpent_FindActive = 0.0;
    double timeSpent_Assign = 0.0;
    double timeSpent_AssignFixed = 0.0;
    double timeSpent_Restart = 0.0;
    double timeSpent_AssignFixed_Digamma = 0.0;
    double timeSpent_AssignFixed_Resp = 0.0;
    */
    nCoordAscentIterLP = max(nCoordAscentIterLP, 0);
    if (initProbsToEbeta == 1) {
        nCoordAscentIterLP += 1;
    }
    nCoordAscentIterLP = max(nCoordAscentIterLP, 1);

    if (nnzPerRow == 1) {
        REVISE_EVERY = 1;
    } else {
        REVISE_EVERY = max(REVISE_EVERY, 1);
    }
    REVISE_FIRST = max(REVISE_FIRST, 0);

    int CHECK_EVERY = 5; // Check convergence every X iterations.

    // Allocate temporary storage
    Arr1D_d ElogProb_d (K);
    Arr1D_d prevTopicCount_d (K);
    Arr1D_d logScores_n (K);
    Arr1D_d tempScores_n (K);
    Arr1D_i activeTopics_d (K);
    Arr1D_i spareActiveTopics_d (nnzPerRow);

    Arr1D_d termResp_data;
    Arr1D_i termResp_colids;

    Argsortable1DArray logScoresHandler = Argsortable1DArray(
        logScores_n.data(), K);

    if (initProbsToEbeta == 2) {
        termResp_data = Arr1D_d::Zero(V * nnzPerRow);
        termResp_colids = Arr1D_i::Zero(V * nnzPerRow);
    }

    // Disable all further memory allocation.
    // We do this to verify that we have no memory leaks!
    internal::set_is_malloc_allowed(false);

    // Unpack input arrays
    ExtArr1D_d alphaEbeta (alphaEbeta_IN, K);
    ExtArr2D_d Eloglik (Eloglik_IN, V, K);

    ExtArr1D_d word_count (word_count_IN, Nall);
    ExtArr1D_i word_id (word_id_IN, Nall);
    ExtArr1D_i doc_range (doc_range_IN, D+1);

    // Unpack output arrays
    ExtArr1D_d spResp_data (spResp_data_OUT, Nall * nnzPerRow);
    ExtArr1D_i spResp_colids (spResp_colids_OUT, Nall * nnzPerRow);
    ExtArr2D_d topicCount (topicCount_OUT, D, K);

    ExtArr1D_i numIterVec (numIterVec_OUT, D);
    ExtArr1D_d maxDiffVec (maxDiffVec_OUT, D);

    ExtArr1D_i rAcceptVec (rAcceptVec_IN, 1);
    ExtArr1D_i rTrialVec (rTrialVec_IN, 1);

    // Compute quantity used in sparse restart ELBO computation
    double sum_gammalnalphaEbeta = 0.0;
    if (numRestarts > 0) {
        for (int k = 0; k < K; k++) {
            sum_gammalnalphaEbeta += boost::math::lgamma(alphaEbeta(k));
        }
    }

    if (initProbsToEbeta == 2) {
        precomputeTopLRespForEachVocabTerm(
            nnzPerRow, V, K,
            Eloglik,
            alphaEbeta,
            ElogProb_d,
            logScores_n,
            tempScores_n,
            termResp_data,
            termResp_colids);
        if (verbose > 0) {
            printf("Precalculated Topics-by-term\n");
            for (int v = 0; v < 5; v++) {
                printf("term %06d ", v);
                for (int vid = v * nnzPerRow; vid < (v+1) * nnzPerRow; vid++) {
                    printf("%03d:%.2f ",
                        termResp_colids(vid), termResp_data(vid));
                }
                printf("\n");
            }
            for (int v = V-5; v < V; v++) {
                printf("term %06d ", v);
                for (int vid = v * nnzPerRow; vid < (v+1) * nnzPerRow; vid++) {
                    printf("%03d:%.2f ",
                        termResp_colids(vid), termResp_data(vid));
                }
                printf("\n");
            }
        }
    }

    // Visit each document and update its spResp (and corresponding topicCount)
    for (int d = 0; d < D; d++) {
        if (verbose == 2 && d < 1) {
            printf("\nSTANDARD INFERENCE AT DOC %d\n", d);
        }

        int start_d = doc_range(d);
        int N_d = doc_range(d+1) - doc_range(d);

        prevTopicCount_d.fill(0);
        double maxDiff = N_d;
        int iter = 0;
        double ACTIVE_THR = 1e-9;
        int doReviseActiveSet = 1;
        int Kactive = K;
        // Initialize activeTopics_d and topicCount_d
        if (initProbsToEbeta == 2) {
            // using the precomputed per-token resp values!
            for (int n = start_d; n < start_d + N_d; n++) {
                int wid = word_id(n);
                for (int a = wid * nnzPerRow; a < (wid+1) * nnzPerRow; a++) {
                    int k = termResp_colids(a);
                    topicCount(d, k) += word_count(n) * termResp_data(a);
                }
            }
            Kactive = updateActiveSetForDocFromScratch(
                d,
                topicCount,
                activeTopics_d,
                spareActiveTopics_d,
                ACTIVE_THR,
                K, nnzPerRow);
        } else if (initProbsToEbeta == 0) {
             Kactive = updateActiveSetForDocFromScratch(
                 d,
                 topicCount,
                 activeTopics_d,
                 spareActiveTopics_d,
                 ACTIVE_THR,
                 K, nnzPerRow);
        } else {
            // Initialize active set to ALL topics
            Kactive = K;
            for (int k = 0; k < K; k++) {
                activeTopics_d(k) = k;
            }
        }

        for (iter = 0; iter < nCoordAscentIterLP; iter++) {
            //clock_gettime(CLOCK_PROCESS_CPUTIME_ID, &start_time);
            // DETERMINE CURRENT ACTIVE SET
            if (iter > 0 && nnzPerRow > 1 && Kactive <= nnzPerRow) {
                // Set of active docs is already as small as can be,
                // so nothing to gain from revising
                doReviseActiveSet = 0;
            } else if (iter < REVISE_FIRST || (iter - 1) % REVISE_EVERY == 0) {
                doReviseActiveSet = 1;
            } else {
                doReviseActiveSet = 0;
            }
            if (iter > 0 && doReviseActiveSet) {
                Kactive = updateActiveSetForDoc(
                    d,
                    topicCount,
                    activeTopics_d,
                    spareActiveTopics_d,
                    ACTIVE_THR,
                    Kactive, nnzPerRow);
            }
            assert(Kactive >= nnzPerRow);
            assert(Kactive <= K);
            //clock_gettime(CLOCK_PROCESS_CPUTIME_ID, &end_time);
            //timeSpent_FindActive += calcElapsedTime(start_time, end_time);

            // COMPUTE ASSIGNMENTS FOR CURRENT ACTIVE SET
            if (doReviseActiveSet) {
                //clock_gettime(CLOCK_PROCESS_CPUTIME_ID, &start_time);
                updateAssignmentsForDoc_ReviseActiveSetDupOK(
                    d, start_d, N_d, nnzPerRow, Kactive,
                    alphaEbeta,
                    Eloglik,
                    word_count,
                    word_id,
                    topicCount,
                    spResp_data,
                    spResp_colids,
                    activeTopics_d,
                    ElogProb_d,
                    logScores_n,
                    logScoresHandler,
                    (initProbsToEbeta == 1) && (iter == 0),
                    0);
                //clock_gettime(CLOCK_PROCESS_CPUTIME_ID, &end_time);
                //timeSpent_Assign += calcElapsedTime(start_time, end_time);
            } else {
                //clock_gettime(CLOCK_PROCESS_CPUTIME_ID, &start_time);
                updateAssignmentsForDoc_FixPerTokenActiveSet(
                    d, start_d, N_d, nnzPerRow, Kactive,
                    alphaEbeta,
                    Eloglik,
                    word_count,
                    word_id,
                    topicCount,
                    spResp_data,
                    spResp_colids,
                    activeTopics_d,
                    ElogProb_d,
                    logScores_n
                    );
                /*clock_gettime(CLOCK_PROCESS_CPUTIME_ID, &end_time);
                timeSpent_AssignFixed += calcElapsedTime(start_time, end_time);
                timeSpent_AssignFixed_Digamma += logScores_n(0);
                timeSpent_AssignFixed_Resp += logScores_n(1);
                */
            }

            // END ITERATION. Decide whether to quit early
            // Find maximum difference from previous doc-topic count
            if (iter % CHECK_EVERY == 0 && iter > 0) {
                double absDiff_k = 0.0;
                maxDiff = 0.0;
                for (int ka = 0; ka < Kactive; ka++) {
                    int k = activeTopics_d(ka);
                    absDiff_k = abs(prevTopicCount_d(k) - topicCount(d, k));
                    if (absDiff_k > maxDiff) {
                        maxDiff = absDiff_k;
                    }
                }
            }

            if (verbose == 2 && d < 1) {
                printf("d %3d iter %3d doRevise %d Kactive %4d maxDiff %11.6f\n",
                    d, iter, doReviseActiveSet, Kactive, maxDiff);
                printf("  ");
                for (int ka = 0; ka < Kactive; ka++) {
                    int k = activeTopics_d(ka);
                    printf("%02d:%06.2f ", k, topicCount(d,k));
                }
                printf("\n");
            }
            if (maxDiff <= convThrLP) {
                if (verbose > 2) {
                    printf("EARLY EXIT! maxDiff < %11.6f\n", convThrLP);
                }
                break;
            }

            if (iter % CHECK_EVERY == CHECK_EVERY - 1) {
                // Make sure prevTopicCount vector is updated
                for (int ka = 0; ka < Kactive; ka++) {
                    int k = activeTopics_d(ka);
                    prevTopicCount_d(k) = topicCount(d, k);
                }
            }

        } // end loop over iterations at doc d
        maxDiffVec(d) = maxDiff;
        numIterVec(d) = iter; // iter will already be +1'd by last iter of loop

        if (numRestarts > 0) {
            //clock_gettime(CLOCK_PROCESS_CPUTIME_ID, &start_time);
            tryRestartsForDoc(
                d, start_d, N_d, nnzPerRow, Kactive,
                alphaEbeta,
                Eloglik,
                word_count,
                word_id,
                topicCount,
                spResp_data,
                spResp_colids,
                activeTopics_d,
                prevTopicCount_d,
                ElogProb_d,
                logScores_n,
                logScoresHandler,
                rAcceptVec,
                rTrialVec,
                numRestarts,
                sum_gammalnalphaEbeta,
                verbose
                );
           //clock_gettime(CLOCK_PROCESS_CPUTIME_ID, &end_time);
           //timeSpent_Restart += calcElapsedTime(start_time, end_time);

        }
    } // end loop over documents
    internal::set_is_malloc_allowed(true);

    if (verbose > 0) {
        printf(
            "Sparse Restarts: %d/%d accepted\n", rAcceptVec(0), rTrialVec(0));
    }

    /*
    double timeSpent_total = \
        timeSpent_FindActive + timeSpent_Assign + \
        timeSpent_AssignFixed + timeSpent_Restart;

    printf("FindActive   %8.3f sec  %.3f %%\n", timeSpent_FindActive,
        timeSpent_FindActive / timeSpent_total);
    printf("Assign:      %8.3f sec  %.3f %%\n", timeSpent_Assign,
        timeSpent_Assign / timeSpent_total);
    printf("AssignFixed: %8.3f sec  %.3f %%\n", timeSpent_AssignFixed,
        timeSpent_AssignFixed / timeSpent_total);
    printf("  -Digamma:  %8.3f sec\n", timeSpent_AssignFixed_Digamma);
    printf("  -Resp:     %8.3f sec\n", timeSpent_AssignFixed_Resp);
    printf("Restarts:    %8.3f sec  %.3f %%\n", timeSpent_Restart,
        timeSpent_Restart / timeSpent_total);
    printf("TOTAL:       %8.3f sec\n", timeSpent_total);
    */
}


//==============================================================================
//
//   (c) Copyright, 2010 University Corporation for Atmospheric Research (UCAR).
//       All rights reserved.
//       Do not copy or distribute without authorization.
//
//       File: $RCSfile: lambertGrid.cc,v $
//       Version: $Revision: 1.6 $  Dated: $Date: 2012-01-06 19:57:45 $
//
//==============================================================================

/**
 * @file cloudGrid.cc
 *
 * Convert lat-longs cloud grid coordinates
 *
 * @date 3/25/10
 */

// Include files
#include <stdio.h>
#include <string>
#include <vector>
#include <projects.h>
#include <boost/format.hpp>
#include "LambertConfigReader.hh"
#include "Proj4Wrap.hh"

using boost::format;
using std::string;
using std::vector;

// Functions

void usage(char *programName)
{
  fprintf(stderr, "Usage: %s configFile\n", programName);
}

int main(int argc, char **argv)
{
  printf("argc %d\n", argc);

  if (argc != 2)
    {
      usage(argv[0]);
      exit(2);
    }

  LambertConfigReader cfg(argv[1]);
  if (cfg.error != string(""))
    {
      printf("Error: configuration file error %s\n", cfg.error.c_str());
      return 1;
    }

  // Uses latitude where lambert conformal projection is true and
  // median aligned with cartesian y-axis
  string paramString = str(format("+proj=lcc +R=6371200 +lon_0=%1% +lat_0=%2% +lat_1=%3% +lat_2=%4%") % cfg.lov % cfg.latin1 % cfg.latin1 % cfg.latin2);

  printf("paramString %s\n", paramString.c_str());
  p4w::Proj4Wrap lambertProj(paramString, p4w::Proj4Wrap::LON_LAT_TYPE, cfg.lo1, cfg.la1, cfg.dx, cfg.dy);
  double xc;
  double yc;
  double lon;
  double lat;

  lambertProj.ll2xy(cfg.lo1, cfg.la1, &xc, &yc);
  printf("lower corner for %g %g: xc, yc: %g %g\n", cfg.lo1, cfg.la1, xc, yc);
  lambertProj.ll2xy(-57.383, 55.481, &xc, &yc);
  printf("upper corner for -57.383, 55.481 : xc, yc: %g %g\n", xc, yc);
  lambertProj.xy2ll(0, 0, &lon, &lat);
  printf("0,0 lon, lat: %g %g\n", lon, lat);
  lambertProj.xy2ll(1, 0, &lon, &lat);
  printf("1,0 lon, lat: %g %g\n", lon, lat);
  lambertProj.xy2ll(2, 0, &lon, &lat);
  printf("2,0 lon, lat: %g %g\n", lon, lat);
  lambertProj.xy2ll(cfg.nx-1, cfg.ny-1, &lon, &lat);
  printf("%d, %d, lon, lat: %g %g\n", cfg.nx-1, cfg.ny-1, lon, lat);
  /*
  */

  vector<float> xvec;
  vector<float> yvec;

  for (int i=0; i<cfg.ny; i++)
    {
      for (int j=0; j<cfg.nx; j++)
	{
	  lat = cfg.la1 + i * cfg.dy;
	  lon = cfg.lo1 + j * cfg.dx;

	  lambertProj.ll2xy(lon, lat, &xc, &yc);
	  xvec.push_back(xc);
	  yvec.push_back(yc);
	}
    }

  // Print cdl header

  printf("netcdf coord {\n");
  printf("dimensions:\n");
  printf("y = %d ;\n", cfg.ny);
  printf("x = %d ;\n", cfg.nx);
  printf("variables:\n");
  printf("float x(y, x) ;\n");
  printf("x:long_name = \"x coordinate\";\n") ;
  printf("float y(y, x) ;\n");
  printf("y:long_name = \"y coordinate\";\n") ;
  printf("data:\n\n");

  printf("x = \n");
  for (uint i=0; i<xvec.size()-1; i++)
    {
      printf("%.7f,\n", xvec[i]);
    }

  printf("%.7f;\n", xvec[xvec.size()-1]);

  printf("y = \n");

  for (uint i=0; i<yvec.size()-1; i++)
    {
      printf("%.7f,\n", yvec[i]);
    }
  printf("%.7f;\n", yvec[yvec.size()-1]);

  printf("}\n");
}


// Copyright (C) 2016-2018 T. Zachary Laine
//
// Distributed under the Boost Software License, Version 1.0. (See
// accompanying file LICENSE_1_0.txt or copy at
// http://www.boost.org/LICENSE_1_0.txt)
//[ self_evaluation
#include <boost/yap/expression.hpp>

#include <boost/optional.hpp>
#include <boost/hana/fold.hpp>
#include <boost/hana/maximum.hpp>

#include <algorithm>
#include <cassert>
#include <iostream>
#include <vector>

// A super-basic matrix type, and a few associated operations.
struct matrix
{
    matrix() : values_(), rows_(0), cols_(0) {}

    matrix(int rows, int cols) : values_(rows * cols), rows_(rows), cols_(cols)
    {
        assert(0 < rows);
        assert(0 < cols);
    }

    int rows() const { return rows_; }
    int cols() const { return cols_; }

    double operator()(int r, int c) const
    { return values_[r * cols_ + c]; }
    double & operator()(int r, int c)
    { return values_[r * cols_ + c]; }

private:
    std::vector<double> values_;
    int rows_;
    int cols_;
};

matrix operator*(matrix const & lhs, double x)
{
    matrix retval = lhs;
    for (int i = 0; i < retval.rows(); ++i) {
        for (int j = 0; j < retval.cols(); ++j) {
            retval(i, j) *= x;
        }
    }
    return retval;
}
matrix operator*(double x, matrix const & lhs) { return lhs * x; }

matrix operator+(matrix const & lhs, matrix const & rhs)
{
    assert(lhs.rows() == rhs.rows());
    assert(lhs.cols() == rhs.cols());
    matrix retval = lhs;
    for (int i = 0; i < retval.rows(); ++i) {
        for (int j = 0; j < retval.cols(); ++j) {
            retval(i, j) += rhs(i, j);
        }
    }
    return retval;
}

// daxpy() means Double-precision AX Plus Y.  This crazy name comes from BLAS.
// It is more efficient than a naive implementation, because it does not
// create temporaries.  The covnention of using Y as an out-parameter comes
// from FORTRAN BLAS.
matrix & daxpy(double a, matrix const & x, matrix & y)
{
    assert(x.rows() == y.rows());
    assert(x.cols() == y.cols());
    for (int i = 0; i < y.rows(); ++i) {
        for (int j = 0; j < y.cols(); ++j) {
            y(i, j) += a * x(i, j);
        }
    }
    return y;
}

template <boost::yap::expr_kind Kind, typename Tuple>
struct self_evaluating_expr;

template <boost::yap::expr_kind Kind, typename Tuple>
auto evaluate_matrix_expr(self_evaluating_expr<Kind, Tuple> const & expr);

// This is the primary template for our expression template.  If you assign a
// self_evaluating_expr to a matrix, its conversion operator transforms and
// evaluates the expression with a call to evaluate_matrix_expr().
template <boost::yap::expr_kind Kind, typename Tuple>
struct self_evaluating_expr
{
    operator auto() const;

    static const boost::yap::expr_kind kind = Kind;

    Tuple elements;
};

// This is a specialization of our expression template for assignment
// expressions.  The destructor transforms and evaluates via a call to
// evaluate_matrix_expr(), and then assigns the result to the variable on the
// left side of the assignment.
//
// In a production implementation, you'd need to have specializations for
// plus_assign, minus_assign, etc.
template <typename Tuple>
struct self_evaluating_expr<boost::yap::expr_kind::assign, Tuple>
{
    ~self_evaluating_expr();

    static const boost::yap::expr_kind kind = boost::yap::expr_kind::assign;

    Tuple elements;
};

struct use_daxpy
{
    // A plus-expression, which may be of the form double * matrix + matrix,
    // or may be something else.  Since our daxpy() above requires a mutable
    // "y", we only need to match a mutable lvalue matrix reference here.
    template <typename Tuple>
    auto operator()(
        boost::yap::expr_tag<boost::yap::expr_kind::plus>,
        self_evaluating_expr<boost::yap::expr_kind::multiplies, Tuple> const & expr,
        matrix & m)
    {
        // Here, we transform the left-hand side into a pair if it's the
        // double * matrix operation we're looking for.  Otherwise, we just
        // get a copy of the left side expression.
        //
        // Note that this is a bit of a cheat, done for clarity.  If we pass a
        // larger expression that happens to contain a double * matrix
        // subexpression, that subexpression will be transformed into a tuple!
        // In production code, this transform should probably only be
        // performed on an expression with all terminal members.
        auto lhs = boost::yap::transform(
            expr,
            [](boost::yap::expr_tag<boost::yap::expr_kind::multiplies>,
               double d,
               matrix const & m) {
                return std::pair<double, matrix const &>(d, m);
            });

        // If we got back a copy of expr above, just re-construct the
        // expression this function mathes; in other words, do not effectively
        // transform anything.  Otherwise, replace the expression matched by
        // this function with a call to daxpy().
        if constexpr (boost::yap::is_expr<decltype(lhs)>::value) {
            return expr + m;
        } else {
            return boost::yap::make_terminal(daxpy)(lhs.first, lhs.second, m);
        }
    }
};

// This is the heart of what self_evaluating_expr does.  If we had other
// optimizations/transformations we wanted to do, we'd put them in this
// function, either before or after the use_daxpy transformation.
template <boost::yap::expr_kind Kind, typename Tuple>
auto evaluate_matrix_expr(self_evaluating_expr<Kind, Tuple> const & expr)
{
    auto daxpy_form = boost::yap::transform(expr, use_daxpy{});
    return boost::yap::evaluate(daxpy_form);
}

template<boost::yap::expr_kind Kind, typename Tuple>
self_evaluating_expr<Kind, Tuple>::operator auto() const
{
    return evaluate_matrix_expr(*this);
}

template<typename Tuple>
self_evaluating_expr<boost::yap::expr_kind::assign, Tuple>::
    ~self_evaluating_expr()
{
    using namespace boost::hana::literals;
    boost::yap::evaluate(elements[0_c]) = evaluate_matrix_expr(elements[1_c]);
}

// In order to define the = operator with the semantics we want, it's
// convenient to derive a terminal type from a terminal instantiation of
// self_evaluating_expr.  Note that we could have written a template
// specialization here instead -- either one would work.  That would of course
// have required more typing.
struct self_evaluating :
    self_evaluating_expr<
        boost::yap::expr_kind::terminal,
        boost::hana::tuple<matrix>
    >
{
    self_evaluating() {}

    explicit self_evaluating(matrix m)
    { elements = boost::hana::tuple<matrix>(std::move(m)); }

    BOOST_YAP_USER_ASSIGN_OPERATOR(self_evaluating_expr, ::self_evaluating_expr);
};

BOOST_YAP_USER_BINARY_OPERATOR(plus, self_evaluating_expr, self_evaluating_expr)
BOOST_YAP_USER_BINARY_OPERATOR(minus, self_evaluating_expr, self_evaluating_expr)
BOOST_YAP_USER_BINARY_OPERATOR(multiplies, self_evaluating_expr, self_evaluating_expr)

int main()
{
    matrix identity(2, 2);
    identity(0, 0) = 1.0;
    identity(1, 1) = 1.0;

    // These are YAP-ified terminal expressions.
    self_evaluating m1(identity);
    self_evaluating m2(identity);
    self_evaluating m3(identity);

    // This transforms the YAP expression to use daxpy(), so it creates no
    // temporaries.  The transform happens in the destructor of the
    // assignment-expression specialization of self_evaluating_expr.
    m1 = 3.0 * m2 + m3;

    // Same as above, except that it uses the matrix conversion operator on
    // the self_evaluating_expr primary template, because here we're assigning
    // a YAP expression to a non-YAP-ified matrix.
    matrix m_result_1 = 3.0 * m2 + m3;

    // Creates temporaries and does not use daxpy(), because the A * X + Y
    // pattern does not occur within the expression.
    matrix m_result_2 = 3.0 * m2;

    return 0;
}
//]


//==============================================================================
//         Copyright 2003 - 2011 LASMEA UMR 6602 CNRS/Univ. Clermont II
//         Copyright 2009 - 2011 LRI    UMR 8623 CNRS/Univ Paris Sud XI
//
//          Distributed under the Boost Software License, Version 1.0.
//                 See accompanying file LICENSE.txt or copy at
//                     http://www.boost.org/LICENSE_1_0.txt
//==============================================================================
#ifndef BOOST_SIMD_ARITHMETIC_FUNCTIONS_SCALAR_VMX_FAST_RSQRT_HPP_INCLUDED
#define BOOST_SIMD_ARITHMETIC_FUNCTIONS_SCALAR_VMX_FAST_RSQRT_HPP_INCLUDED
#if defined(BOOST_SIMD_HAS_VMX_SUPPORT)

#include <boost/simd/arithmetic/functions/fast_rsqrt.hpp>
#include <boost/simd/include/functions/simd/multiplies.hpp>
#include <boost/simd/include/functions/simd/minus.hpp>

namespace boost { namespace simd { namespace ext
{
  BOOST_DISPATCH_IMPLEMENT          ( fast_rsqrt_
                                    , boost::simd::tag::vmx_
                                    , (A0)
                                    , (scalar_< single_<A0> >)
                                    )
  {
    typedef A0 result_type;

    BOOST_FORCEINLINE result_type operator()(float a0) const
    {
      float r;
      __asm__ ("frsqrte %0, %1"
              : "=f" (r)
              : "f" (a0)
              );

      r *= ((3.0f - (r * r) * a0) * 0.5f);
      r *= ((3.0f - (r * r) * a0) * 0.5f);

      return r;
    }
  };

  BOOST_DISPATCH_IMPLEMENT          ( fast_rsqrt_
                                    , boost::simd::tag::vmx_
                                    , (A0)
                                    , (scalar_< double_<A0> >)
                                    )
  {
    typedef A0 result_type;

    BOOST_FORCEINLINE result_type operator()(double a0) const
    {
      double r;
      __asm__ ("frsqrte %0, %1"
              : "=f" (r)
              : "f" (a0)
              );

      r *= ((3.0f - (r * r) * a0) * 0.5f);
      r *= ((3.0f - (r * r) * a0) * 0.5f);
      r *= ((3.0f - (r * r) * a0) * 0.5f);

      return r;
    }
  };
} } }

#endif

#endif