/**
 * @file      MeshFitterImpl.cpp
 *
 * @brief     Implementation file for MeshFitter::Impl
 *
 * @author    Stefan Reinhold
 * @copyright Copyright (C) 2018 Stefan Reinhold  -- All Rights Reserved.
 *            You may use, distribute and modify this code under the terms of
 *            the AFL 3.0 license; see LICENSE for full license details.
 */

#include "MeshFitterImpl.h"

#include "CommonMath.h"
#include "DiscreteRangeDecorators.h"
#include "DisplacementOptimizer.h"
#include "EigenAdaptors.h"
#include "MeshAdaptors.h"
#include "MeshFitterHiddenState.h"
#include "MeshHelpers.h"
#include "Sampler.h"
#include "WeightedARAPFitter.h"

#include <Eigen/Core>
#include <Eigen/Geometry>
#include <gsl/gsl>

namespace CortidQCT {

using namespace Internal;

/**********************
 * Helper Functions
 */

// MARK: -
// MARK: Helper Function

/**
 * @brief Returns a matrix containing positions to sample a voxel volume at
 *
 * Let \f$N := |V|\f$ and let \f$M := |R|\f$, where \f$R\f$ represent
 * `model.samplingRange`, then the returned NMx3 matrix contains N*M positions,
 * where each M consecutive postitions represent a line through a vertex in `V`
 * sampled along the surface normal in `N`.
 *
 * @param V Nx3 matrix with vertex positions
 * @param N Nx3 matrix of per-vertex surface normals
 * @param model MeasurementModel instance
 * @return NMx3 matrix containing sampling positions
 */
template <class DerivedV, class DerivedN>
Eigen::Matrix<typename DerivedV::Scalar, Eigen::Dynamic, 3>
samplingPoints(Eigen::MatrixBase<DerivedV> const &V,
               Eigen::MatrixBase<DerivedN> const &N,
               MeasurementModel const &model) {
  using Scalar = typename DerivedV::Scalar;
  using Eigen::Dynamic;
  using Eigen::Index;
  using Eigen::Matrix;
  using gsl::narrow_cast;

  auto const t = Internal::discreteRangeElementVector(model.samplingRange);
  Matrix<Scalar, Dynamic, 3> samples(V.rows() * t.rows(), 3);

  for (auto i = 0; i < V.rows(); ++i) {
    auto const iStart = i * t.rows();
    samples.block(iStart, 0, t.rows(), 3).colwise() = -t;
    samples.block(iStart, 0, t.rows(), 3).array().rowwise() *= N.row(i).array();
    samples.block(iStart, 0, t.rows(), 3).rowwise() += V.row(i);
  }

  return samples;
}

/***********************************
 * MeshFitter::Impl Implementation
 */

// MARK: -
// MARK: MeshFitter::Impl Implementation

MeshFitter::Result MeshFitter::Impl::fit(VoxelVolume const &volume) const {

  auto state = init(volume);

  VertexMatrix<float> Vlast =
      Adaptor::vertexMap(state.deformedMesh).transpose();

  while (!state.converged) {

    fitOneIteration(state);

    auto V = Adaptor::vertexMap(state.deformedMesh);

    auto const diff = (V.transpose() - Vlast).norm() / V.norm();
    Vlast = V.transpose();

    auto disNorm = Adaptor::map(state.displacementVector).norm() /
                   static_cast<float>(state.displacementVector.size());

    std::cout << "Converged after iteration " << state.iteration << ": "
              << std::boolalpha << state.converged << " (" << diff << " | "
              << disNorm << " | " << state.nonDecreasing << ")" << std::endl;
  }

  state.success = true;

  return std::move(state);
}

MeshFitter::State MeshFitter::Impl::init(VoxelVolume const &volume) const {
  using Eigen::Dynamic;
  using Eigen::Index;
  using Eigen::Map;
  using Eigen::Matrix;
  using Eigen::MatrixXf;
  using Eigen::Vector3f;
  using Eigen::VectorXf;
  using gsl::narrow_cast;

  MeshFitter::State state;

  auto const &conf = fitter_.configuration;
  // auto const &model = fitter_.configuration.model;

  // Init result state
  state.referenceMesh = conf.referenceMesh;

  auto const nVertices = narrow_cast<Index>(conf.referenceMesh.vertexCount());

  // Apply initial transofrmation on the vertices of the reference mesh
  auto const translation = Adaptor::vec(conf.meshTranslation(volume));
  Vector3f const rotInRad = Adaptor::vec(conf.referenceMeshRotation) *
                            static_cast<float>(M_PI) / 180.f;
  auto V0 = Adaptor::vertexMap(state.referenceMesh);

  // Apply transformation to vertices of reference mesh
  V0 = ((Eigen::Translation<float, 3>{translation} *
         Eigen::AngleAxisf(rotInRad[0], Vector3f::UnitX()) *
         Eigen::AngleAxisf(rotInRad[1], Vector3f::UnitY()) *
         Eigen::AngleAxisf(rotInRad[2], Vector3f::UnitZ())) *
        Adaptor::vec(conf.referenceMeshScale).asDiagonal() * V0);

  // Init vertex normals
  state.referenceMesh.updatePerVertexNormals();
  // This will be removed in versino 2.0:
  state.vertexNormals.resize(narrow_cast<std::size_t>(nVertices));
  Adaptor::map(state.vertexNormals) =
      Adaptor::vertexNormalMap(state.referenceMesh);

  // Init deformed mesh with reference mesh
  state.deformedMesh = state.referenceMesh;

  // Init displacement vector
  state.displacementVector =
      std::vector<float>(narrow_cast<std::size_t>(nVertices), .0f);

  // Init weight vector
  state.weights = std::vector<float>(narrow_cast<std::size_t>(nVertices),
                                     1 / static_cast<float>(nVertices));

  // Init vertex normals
  state.vertexNormals.resize(narrow_cast<std::size_t>(nVertices));

  // Init hidden state
  state.hiddenState_ = std::make_unique<State::HiddenState>(
      volume, DisplacementOptimizer{conf},
      WeightedARAPFitter<float>{V0.transpose(), facetMatrix(conf.referenceMesh),
                                narrow_cast<float>(conf.sigmaE)},
      facetMatrix(conf.referenceMesh));

  // Init volume sampling positions
  auto const nSamples = conf.model.samplingRange.numElements() *
                        narrow_cast<std::size_t>(nVertices);
  state.volumeSamplingPositions.resize(nSamples);

  // Init volume samples
  state.volumeSamples.resize(nSamples);

  sampleVolume(state);

  return state;
} // namespace CortidQCT

void MeshFitter::Impl::fitOneIteration(MeshFitter::State &state) const {

  if (state.hiddenState_ == nullptr) {
    throw std::invalid_argument("Invalid state argument, call init() first!");
  }

  findOptimalDisplacements(state);
  findOptimalDeformation(state);
  sampleVolume(state);
  computeLogLikelihood(state);
  checkConvergence(state);

  ++state.iteration;
}

void MeshFitter::Impl::findOptimalDisplacements(
    MeshFitter::State &state) const {
  if (state.hiddenState_ == nullptr) {
    throw std::invalid_argument("Invalid state argument, call init() first!");
  }

  // Get maps to state variables
  auto const labels = Adaptor::labelMap(state.referenceMesh);
  auto const N = Adaptor::vertexNormalMap(state.deformedMesh);
  auto const volumeSamples = Adaptor::map(state.volumeSamples);
  auto optimalDisplacements = Adaptor::map(state.displacementVector);
  auto gamma = Adaptor::map(state.weights);

  // Find optimal displacements
  std::tie(optimalDisplacements, gamma) =
      state.hiddenState_->displacementOptimizer(
          N.transpose(), labels, volumeSamples, state.nonDecreasing,
          state.effectiveSigmaS);
}

void MeshFitter::Impl::findOptimalDeformation(MeshFitter::State &state) const {
  if (state.hiddenState_ == nullptr) {
    throw std::invalid_argument("Invalid state argument, call init() first!");
  }

  auto const N = Adaptor::vertexNormalMap(state.deformedMesh);
  auto const optimalDisplacements = Adaptor::map(state.displacementVector);
  auto const gamma = Adaptor::map(state.weights);
  auto V = Adaptor::vertexMap(state.deformedMesh);

  VertexMatrix<> const Y =
      (V - (N.array().rowwise() * optimalDisplacements.array().transpose())
               .matrix())
          .transpose();

  // Fit mesh
  V = state.hiddenState_->meshFitter.fit(Y, N.transpose(), gamma).transpose();
  // Update normals
  state.deformedMesh.updatePerVertexNormals();
  // This will be removed in v2.0:
  Adaptor::map(state.vertexNormals) =
      Adaptor::vertexNormalMap(state.deformedMesh);
}

void MeshFitter::Impl::sampleVolume(MeshFitter::State &state) const {

  using Eigen::Map;
  using Eigen::MatrixXf;
  using Eigen::VectorXf;

  if (state.hiddenState_ == nullptr) {
    throw std::invalid_argument("Invalid state argument, call init() first!");
  }

  auto const &conf = fitter_.configuration;
  auto const V = Adaptor::vertexMap(state.deformedMesh);
  auto const nVertices = V.cols();
  auto N = Adaptor::vertexNormalMap(state.deformedMesh);
  auto volumeSamplingPositions = Adaptor::map(state.volumeSamplingPositions);
  auto volumeSamples = Adaptor::map(state.volumeSamples);

  // Copmute new sampling positions
  volumeSamplingPositions =
      samplingPoints(V.transpose(), N.transpose(), conf.model).transpose();

  // Sample the volume
  auto const volumeSampler = VolumeSampler{
      state.hiddenState_->volume, conf.ignoreExteriorSamples
                                      ? std::numeric_limits<float>::quiet_NaN()
                                      : 0.f};
  volumeSampler(volumeSamplingPositions.transpose(), volumeSamples,
                fitter_.configuration.calibrationSlope,
                fitter_.configuration.calibrationIntercept);

  // Reorder samples
  state.hiddenState_->volumeSamplesMatrix =
      Map<MatrixXf const>{volumeSamples.data(),
                          volumeSamples.rows() / nVertices, nVertices}
          .transpose();

  volumeSamples = Map<VectorXf const>{
      state.hiddenState_->volumeSamplesMatrix.data(), volumeSamples.rows()};
}

void MeshFitter::Impl::computeLogLikelihood(MeshFitter::State &state) const {
  using Eigen::Map;
  using Eigen::VectorXf;
  using gsl::narrow_cast;

  if (state.hiddenState_ == nullptr) {
    throw std::invalid_argument("Invalid state argument, call init() first!");
  }

  auto const labels = Adaptor::labelMap(state.referenceMesh);
  auto const N = Adaptor::vertexNormalMap(state.deformedMesh);
  auto const volumeSamples = Adaptor::map(state.volumeSamples);

  auto const llVec =
      state.hiddenState_->displacementOptimizer.logLikelihoodVector(
          N.transpose(), labels, volumeSamples);
  auto const LL = llVec.sum();

  state.logLikelihood = LL;
  state.perVertexLogLikelihood.resize(narrow_cast<std::size_t>(llVec.size()));

  Map<VectorXf>{state.perVertexLogLikelihood.data(), llVec.rows(), 1} = llVec;
}

void MeshFitter::Impl::checkConvergence(MeshFitter::State &state) const {
  if (state.hiddenState_ == nullptr) {
    throw std::invalid_argument("Invalid state argument, call init() first!");
  }

  auto const optimalDisplacements = Adaptor::map(state.displacementVector);
  auto const &conf = fitter_.configuration;

  state.converged = state.iteration >= conf.maxIterations ||
                    (optimalDisplacements.norm() < 1e-3f);

  auto disNorm = optimalDisplacements.norm() / optimalDisplacements.rows();
  if (disNorm < state.minDisNorm) {
    state.minDisNorm = disNorm;
    state.nonDecreasing = 0;
  } else {
    ++state.nonDecreasing;
  }
}

} // namespace CortidQCT


#include "geodesic.hh"
#include "Geo/entity.hh"
#include "Geo/polynomial_solver.hh"
#include "Topology/topology.hh"
#include "Topology/impl.hh"
#include "Topology/iterator.hh"
#include "Utils/error_handling.hh"

#include <boost/math/tools/roots.hpp>

#include <bitset>
#include <functional>
#include <fstream>

#if 1

#include <map>
#include <queue>
#include <set>

namespace Offset
{
struct EdgeDistance;
using EdgeDistances = std::multimap<Topo::Wrap<Topo::Type::EDGE>, EdgeDistance>;
using EdgeAndDistance = EdgeDistances::value_type;

struct EdgeDistance
{
  EdgeDistance()
  {
    static int id_seq = 0;
    id_ = id_seq++;
  }
  ~EdgeDistance()
  {}

  void init(const Topo::Wrap<Topo::Type::FACE>& _f,
            const EdgeAndDistance* _parent)
  {
    if (x_ < 0)
    {
      x_ *= -1;
      y_ed_[0] *= -1;
      y_ed_[1] *= -1;
    }
    update_parameter();
    origin_face_ = _f;
    parent_ = _parent;
  }

  void update_parameter()
  {
    y_[0] = y_ed_[0] + (y_ed_[1] - y_ed_[0]) * param_range_[0];
    y_[1] = y_ed_[0] + (y_ed_[1] - y_ed_[0]) * param_range_[1];
    distance_range_.set_empty();
    distance_range_.add(sqrt(Geo::sq(x_) + Geo::sq(y_[0])) + dist0_);
    distance_range_.add(sqrt(Geo::sq(x_) + Geo::sq(y_[1])) + dist0_);
    if ((y_[0] > 0) != (y_[1] > 0))
      distance_range_.add(x_ + dist0_);
  }

  double get_distance(double _par,  double* _first_der = nullptr) const
  {
    auto dy = y_ed_[1] - y_ed_[0];
    auto y_part = y_ed_[0] + _par * dy;
    auto dist = sqrt(Geo::sq(x_) + Geo::sq(y_part));
    if (_first_der != nullptr)
      *_first_der = dy * y_part / dist;
    return dist0_ + dist;
  }

  size_t parameters(const double _dist, std::array<double, 2>& _pars)
  {
    if (!distance_range_.contain_close(_dist))
      return 0;
    auto y = std::sqrt(Geo::sq(_dist - dist0_) - Geo::sq(x_));
    size_t sol_nmbr = 0;
    auto get_solution = [this, &sol_nmbr, &_pars](double _y)
    {
      _pars[sol_nmbr] = (_y - y_ed_[0]) / (y_ed_[1] - y_ed_[0]);
      if (_pars[sol_nmbr] >= 0 && _pars[sol_nmbr] < 1)
        ++sol_nmbr;
    };
    get_solution(y);
    get_solution(-y);
    return sol_nmbr;
  }

  enum class Status { Keep = 0, Skip = 1, Remove = 2 };

  int id_;
  Topo::Wrap<Topo::Type::FACE> origin_face_; // Face used to get there
  double x_;
  double y_ed_[2];
  double dist0_ = 0;
  Geo::Interval<double> param_range_;  // Portion of edges
  Status status_ = Status::Keep;
  double y_[2];
  Geo::Interval<double> distance_range_;             // range of distances
  std::array<EdgeAndDistance*, 2> children_ = {nullptr};
  const EdgeAndDistance* parent_ = nullptr;
};

void check(const EdgeAndDistance& _edd, const Geo::Point& _orig)
{
  auto t = _edd.second.param_range_.interpolate(0.5);
  Topo::Iterator<Topo::Type::EDGE, Topo::Type::VERTEX> ev(_edd.first);
  Geo::Point pt, pt1;
  ev.get(0)->geom(pt);
  ev.get(1)->geom(pt1);
  pt = Geo::interpolate(pt, pt1, t);
  auto dist = Geo::angle(_orig, pt) * 5;
  if (!_edd.second.distance_range_.contain_close(dist)) dist; // std::cout << "error";
}

// Return 0 if _apr is a root, else
// < 0 ==> _a < _b
// > 0 ==> _b < _a

int find_minimum(const EdgeDistance& _a, const EdgeDistance& _b,
                 double& _par, Geo::Interval<double> _range)
{
  auto func = [&_a, &_b](const double& _t)
  {
    double der_a, der_b;
    auto val_a = _a.get_distance(_t, &der_a);
    auto val_b = _b.get_distance(_t, &der_b);
    return std::make_tuple(val_a - val_b, der_a - der_b);
  };
  boost::uintmax_t max_iter = 20;
  _par = boost::math::tools::newton_raphson_iterate(
    func, (_range[0] + _range[1]) / 2, _range[0], _range[1],
    std::numeric_limits<double>::digits - 4, max_iter);
  if (max_iter < 20)
    return 0;
  return std::get<0>(func(_range.interpolate(0.5))) > 0 ? 1 : -1;
}

struct GeodesicDistance: public IGeodesic
{
  Geo::Point origin_;
  // First smaller distances. Priority queue process first big elements.
  struct CompareEdgeDistancePtr
  {
    bool operator()(const EdgeAndDistance* _a, const EdgeAndDistance* _b) const
    {
      return _a->second.distance_range_[1] > _b->second.distance_range_[1];
    }
  };

  EdgeDistances edge_distances_;
  std::priority_queue<const EdgeAndDistance*, std::vector<const EdgeAndDistance*>,
    CompareEdgeDistancePtr> priority_queue_;

  bool compute(const Topo::Wrap<Topo::Type::VERTEX>& _v) override;

  bool find_graph(
    double _dist,
    std::vector<Geo::VectorD3>& _pts,
    std::vector<std::array<size_t, 2>>& _inds) override;

private:
  void advance(const EdgeAndDistance* _ed_span,
               const Topo::Wrap<Topo::Type::FACE>& _f);
  void merge(Topo::Wrap<Topo::Type::EDGE> _ed,
             EdgeDistance& _ed_dist);

  void insert_element(const Topo::Wrap<Topo::Type::EDGE> _ed,
                      EdgeDistance& _ed_dist)
  {
    auto it = edge_distances_.emplace(_ed, _ed_dist);
    check(*it, origin_);
    priority_queue_.push(&(*it));
  }
};

std::shared_ptr<IGeodesic> IGeodesic::make()
{
  return std::make_shared<GeodesicDistance>();
}

bool GeodesicDistance::compute(const Topo::Wrap<Topo::Type::VERTEX>& _v)
{
  _v->geom(origin_);
  Topo::Iterator<Topo::Type::VERTEX, Topo::Type::EDGE> vert_it(_v);
  for (auto e : vert_it)
  {
    EdgeDistance ed_dist;
    Topo::Iterator<Topo::Type::EDGE, Topo::Type::VERTEX> ed_vert_it(e);
    auto v1 = ed_vert_it.get(1);
    auto inv = v1 == _v;
    if (inv)
      v1 = ed_vert_it.get(0);
    Geo::Point pt;
    v1->geom(pt);
    ed_dist.x_ = 0;
    ed_dist.y_ed_[0] = 0;
    ed_dist.y_ed_[1] = Geo::length(pt - origin_);
    if (inv)
      std::swap(ed_dist.y_ed_[0], ed_dist.y_ed_[1]);
    ed_dist.param_range_ = { 0, 1 };
    ed_dist.init(Topo::Wrap<Topo::Type::FACE>(), nullptr);
    ed_dist.status_ = EdgeDistance::Status::Skip;
    insert_element(e, ed_dist);
  }
  Topo::Iterator<Topo::Type::VERTEX, Topo::Type::FACE> face_it(_v);
  for (auto f : face_it)
  {
    Topo::Iterator<Topo::Type::FACE, Topo::Type::EDGE> ed_it(f);
    Topo::Wrap<Topo::Type::EDGE> curr_ed;
    for (auto e : ed_it)
    {
      auto it = edge_distances_.equal_range(e);
      if (it.first == it.second)
      {
        curr_ed = e;
        break;
      }
    }
    if (curr_ed.get() == nullptr)
      continue;
    EdgeDistance ed_dist;
    Topo::Iterator<Topo::Type::EDGE, Topo::Type::VERTEX> ed_vert(curr_ed);
    Geo::Segment seg;
    ed_vert.get(0)->geom(seg[0]);
    ed_vert.get(1)->geom(seg[1]);
    double dist_sq, t;
    Geo::VectorD3 proj;
    if (!Gen::closest_point<double, 3, true>(seg, origin_, &proj, &t, &dist_sq))
      return false;
    ed_dist.x_ = sqrt(dist_sq);
    ed_dist.y_ed_[0] = Geo::length(seg[0] - proj);
    ed_dist.y_ed_[1] = Geo::length(seg[1] - proj);
    if (t < 1 && t > 0)
      ed_dist.y_ed_[ed_dist.y_ed_[1] < ed_dist.y_ed_[0]] *= -1;
    ed_dist.param_range_ = { 0, 1 };
    ed_dist.init(f, nullptr);
    insert_element(curr_ed, ed_dist);
  }
  while (!priority_queue_.empty())
  {
    auto edge_span = priority_queue_.top();
    priority_queue_.pop();
    if (edge_span->second.status_ != EdgeDistance::Status::Keep)
      continue;
    Topo::Iterator<Topo::Type::EDGE, Topo::Type::FACE> fe_it(edge_span->first);
    for (auto f : fe_it)
    {
      if (f != edge_span->second.origin_face_)
        advance(edge_span, f);
    }
  }
  for (auto it = edge_distances_.begin(); it != edge_distances_.end();)
  {
    if (it->second.status_ == EdgeDistance::Status::Remove)
      it = edge_distances_.erase(it);
    else
      ++it;
  }
  std::ofstream outf("c:/t/out.txt");
  outf << "V0 V1 Id Px Py Pz Qx Qy Qz x y0 y1 interval0 interval1 distance0  distance1 dist0 par_id \n";
  for (auto& es : edge_distances_)
  {
    Geo::Segment seg;
    es.first->geom(seg);
    Topo::Iterator<Topo::Type::EDGE, Topo::Type::VERTEX> ed_it(es.first);
    outf << ed_it.get(0)->id() << " " << ed_it.get(1)->id() << " ";
    outf << es.second.id_ << " " << seg[0] << " " << seg[1] << " ";
    outf << es.second.x_ << " ";
    outf << es.second.y_ed_[0] << " " << es.second.y_ed_[1] << " ";
    outf << es.second.param_range_[0] << " " << es.second.param_range_[1] << " ";
    outf << es.second.distance_range_[0] << " " << es.second.distance_range_[1] << " ";
    outf << es.second.dist0_ << " ";
    if (es.second.parent_ == nullptr)
      outf << -1;
    else
      outf << es.second.parent_->second.id_;
    outf  << std::endl;
  }
  outf << "SIZE = " << edge_distances_.size() << "\n";
  return true;
}

void GeodesicDistance::advance(
  const EdgeAndDistance* _parent, const Topo::Wrap<Topo::Type::FACE>& _f)
{
  Topo::Iterator<Topo::Type::EDGE, Topo::Type::VERTEX> it_ev_pa(_parent->first);
  Topo::Wrap<Topo::Type::VERTEX> verts[2] = { it_ev_pa.get(0), it_ev_pa.get(1) };
  Geo::VectorD3 p_parent[2];
  Geo::iterate_forw<2>::eval([&verts, &p_parent](int _i) { verts[_i]->geom(p_parent[_i]); });
  Geo::VectorD3 v_parent = p_parent[1] - p_parent[0];
  Topo::Iterator<Topo::Type::FACE, Topo::Type::EDGE> it_fe(_f);
  struct OtherEdge
  {
    Topo::Wrap<Topo::Type::EDGE> ed_;
    bool inv_;
    Geo::VectorD3 v_;
  };
  Gen::Segment<double, 2> parent_seg{ {
    { _parent->second.x_, _parent->second.y_[0] },
    { _parent->second.x_, _parent->second.y_[1] } } };
  Gen::Segment<double, 2> parent_trace_segs[2];
  for (int i = 0; i < 2; ++i)
  {
    parent_trace_segs[i][0] = {};
    parent_trace_segs[i][1] = parent_seg[i];
  }

  OtherEdge oth_eds[2];
  auto dy_par = _parent->second.y_ed_[1] - _parent->second.y_ed_[0];
  for (auto& fe : it_fe)
  {
    if (fe == _parent->first)
      continue;
    EdgeDistance new_edd;
    Topo::Iterator<Topo::Type::EDGE, Topo::Type::VERTEX> it_ev(fe);
    bool idx = it_ev.get(1) == verts[1] || it_ev.get(0) == verts[1];
    auto inv = it_ev.get(1) == verts[idx];
    oth_eds[idx].ed_ = fe;
    oth_eds[idx].inv_ = inv;
    Geo::VectorD3 p[2];
    Geo::iterate_forw<2>::eval([&it_ev, &p](int _i) { it_ev.get(_i)->geom(p[_i]); });
    auto v0 = v_parent;
    if (dy_par < 0) v0 *= -1.;
    v0 /= Geo::length(v0);
    auto v1 = oth_eds[idx].v_ = p[1] - p[0];
    if (inv) v1 *= -1.;
    auto dy = v0 * v1;
    //if (dy * dy_par < 0)
    //  continue;
    auto dx = Geo::length(v0 % v1);
    Gen::Segment<double, 2> v2 = { {
      { _parent->second.x_, _parent->second.y_ed_[idx] },
      { _parent->second.x_ + dx, _parent->second.y_ed_[idx] + dy } } };
    if (inv)
      std::swap(v2[0], v2[1]);
    auto dist_3d = Geo::length(p[0] - p[1]);

    auto compute_x_y_ed = [dist_3d](const Gen::Segment<double, 2>& _v, EdgeDistance& _new_ed)
    {
      Geo::VectorD2 proj;
      double dist_sq, par;
      Gen::closest_point<double, 2, true>(_v, Geo::VectorD2{},
                                          &proj, &par, &dist_sq);
      _new_ed.x_ = sqrt(dist_sq);
      _new_ed.y_ed_[0] = Geo::length(_v[0] - proj);
      _new_ed.y_ed_[1] = Geo::length(_v[1] - proj);
      if (par > 0 && par < 1)
        _new_ed.y_ed_[_new_ed.y_ed_[0] >_new_ed.y_ed_[1]] *= -1;
      THROW_IF((fabs(_new_ed.y_ed_[1] - _new_ed.y_ed_[0]) - dist_3d) > 1e-8,
               "Wrong y reange");
    };
    compute_x_y_ed(v2, new_edd);

    Geo::Interval<double> tmp;
    for (int i = 0; i < 2; ++i)
    {
      bool done = false;
      for (int j = 0; j < 2; ++j)
      {
        double t, dist_sq;
        Gen::closest_point<double, 2, true>(
          parent_trace_segs[i], v2[j], nullptr, &t, &dist_sq);
        done = dist_sq < Geo::epsilon_sq(v2[j]);
        if (done)
        {
          tmp.add(j);
          break;
        }
      }
      if (done)
        continue;

      double pars[2], dist_sq;
      Gen::closest_point<double, 2, true, true>(
        parent_trace_segs[i], v2, nullptr, pars, &dist_sq);
      auto dir = v2[1] - v2[0];
      double cr_pr = fabs(parent_trace_segs[i][1] % dir);
      if (cr_pr < 1e-12)
        pars[1] = Geo::length_square(v2[1]) > Geo::length_square(v2[0]) ? 1: 0;
      //auto val_sq = std::max(Geo::length_square(v2[0]), Geo::length_square(v2[1]));
      else if (pars[0] <= 1e-12)
      {
        if (pars[1] <= 0) pars[1] = 1;
        else if (pars[1] >= 1) pars[1] = 0;
        else
          THROW("Bad par");

      }
      else if (pars[0] <= 1 - 1e-12)
        continue;
      tmp.add(std::clamp(pars[1], 0., 1.));
    }

    if (!tmp.empty())
    {
      new_edd.param_range_ = tmp;
      new_edd.dist0_ = _parent->second.dist0_;
      new_edd.init(_f, _parent);
      merge(fe, new_edd);
    }

    size_t par_bndr_to_fill = idx;
    if (tmp.empty() &&
        (_parent->second.param_range_[par_bndr_to_fill] == par_bndr_to_fill))
    {
      EdgeDistance new_edd2;
      new_edd2.dist0_ = _parent->second.get_distance(static_cast<double>(par_bndr_to_fill));
      auto orig = Geo::VectorD2{ _parent->second.x_, _parent->second.y_ed_[par_bndr_to_fill] };
      Gen::Segment<double, 2> new_seg;
      new_seg[0] = v2[0] - orig;
      new_seg[1] = v2[1] - orig;
      compute_x_y_ed(new_seg, new_edd2);
      new_edd2.param_range_ = { 0., 1. };
      new_edd2.init(_f, _parent);
      merge(fe, new_edd2);
    }
    par_bndr_to_fill = !idx;
    size_t new_par = !inv;
    if (!tmp.empty() && tmp[new_par] != new_par &&
      (_parent->second.param_range_[par_bndr_to_fill] == par_bndr_to_fill))
    {
      EdgeDistance new_edd2;
      new_edd2.dist0_ = _parent->second.get_distance(static_cast<double>(par_bndr_to_fill));
      auto orig = Geo::VectorD2{ _parent->second.x_, _parent->second.y_ed_[par_bndr_to_fill] };
      Gen::Segment<double, 2> new_seg;
      new_seg[0] = v2[0] - orig;
      new_seg[1] = v2[1] - orig;
      compute_x_y_ed(new_seg, new_edd2);
      if (inv)
        new_edd2.param_range_ = Geo::Interval<double>(0., tmp[0]);
      else
        new_edd2.param_range_ = Geo::Interval<double>(tmp[1], 1.);
      new_edd2.init(_f, _parent);
      merge(fe, new_edd2);
    }

  }
  THROW_IF(!oth_eds[0].ed_.get() || !oth_eds[1].ed_.get(), "");
}

static std::bitset<2> intersect(EdgeDistance& _a, EdgeDistance& _b,
                                Geo::Interval<double> _extras[2])
{
  auto inters_par_range = _a.param_range_ * _b.param_range_;
  if (inters_par_range.length() < 1e-12)
    return 0;
  double root;
  auto a_minus_b = find_minimum(_a, _b, root, inters_par_range);
  std::bitset<2> res;
  if (a_minus_b > 0)
    res[0] = _a.param_range_.subtract(inters_par_range, _extras[0]);
  else if (a_minus_b < 0)
    res[1] = _b.param_range_.subtract(inters_par_range, _extras[1]);
  else
  {
    auto snap_to_boundary = [](double& _root, const Geo::Interval<double>& _interv)
    {
      for (int i = 0; i < 2; ++i)
        if (fabs(_root - _interv[i]) < 1e-8)
        {
          _root = _interv[i];
          return true;
        }
      return false;
    };
    snap_to_boundary(root, _a.param_range_) || snap_to_boundary(root, _b.param_range_);
    Geo::Interval<double> near_a(Geo::Interval<double>::min(), root);
    Geo::Interval<double> near_b(root, Geo::Interval<double>::max());
    if (_a.get_distance(root - 1) > _b.get_distance(root - 1))
      std::swap(near_a, near_b);

    near_b *= _b.param_range_;
    res[0] = _a.param_range_.subtract(near_b, _extras[0]);
    near_a *= _a.param_range_;
    res[1] = _b.param_range_.subtract(near_a, _extras[1]);
  }
  return res;
}

void GeodesicDistance::merge(Topo::Wrap<Topo::Type::EDGE> _ed,
                             EdgeDistance& _ed_dist)
{
  auto range = edge_distances_.equal_range(_ed);
  Topo::Iterator<Topo::Type::EDGE, Topo::Type::VERTEX> edv(_ed);
  if (edv.get(0)->id() == 395 && edv.get(1)->id() == 436)
  {
    std::cout << "Here we are";
  }
  std::vector<EdgeDistance*> old_ed_dists;
  for (auto it = range.first; it != range.second; ++it)
    old_ed_dists.push_back(&(it->second));
  std::vector<EdgeDistance> new_edds;
  new_edds.push_back(_ed_dist);
  for (auto old_edd : old_ed_dists)
  {
    //if (it->second.status_ != EdgeDistance::Status::Keep)
    if (old_edd->status_ == EdgeDistance::Status::Remove)
      continue;
    auto size = new_edds.size();
    for (size_t i = 0; i < size; ++i)
    {
      auto& new_edd = new_edds[i];
      Geo::Interval<double> extras[2];
      auto int_res = intersect(new_edd, *old_edd, extras);
      if (!int_res.any())
        continue;
      if (int_res[1])
      {
        if (old_edd->param_range_.empty())
        {
          std::function<void(EdgeDistance* _ed_dist)> skip_children =
            [&](EdgeDistance* _ed_dist)
          {
            _ed_dist->status_ = EdgeDistance::Status::Remove;
            for (auto child : _ed_dist->children_)
              if (child != nullptr)
                skip_children(&child->second);
          };
          skip_children(old_edd);
        }
        else
        {
          old_edd->update_parameter();
          if (!extras[1].empty())
          {
            auto split_piece = *old_edd;
            split_piece.param_range_ = extras[1];
            split_piece.update_parameter();
            insert_element(_ed, split_piece);
          }
        }
      }
      if (int_res[0])
      {
        new_edd.update_parameter();
        if (!extras[0].empty())
        {
          auto new_el = new_edds.emplace_back(new_edd);
          new_el.param_range_ = extras[0];
          new_el.update_parameter();
        }
      }
    }
  }
  for (auto& new_edd : new_edds)
    if (!new_edd.param_range_.empty())
      insert_element(_ed, new_edd);
}

bool GeodesicDistance::find_graph(
  double _dist,
  std::vector<Geo::VectorD3>& _pts,
  std::vector<std::array<size_t, 2>>& _inds)
{
  using FacePointMap = std::map<Topo::Wrap<Topo::Type::FACE>, std::vector<size_t>>;
  FacePointMap fp_map;
  for (auto& ed_dist : edge_distances_)
  {
    std::array<double, 2> pars;
    auto sol_num = ed_dist.second.parameters(_dist, pars);
    if (sol_num == 0)
      continue;
    Topo::Iterator<Topo::Type::EDGE, Topo::Type::VERTEX> vert_it(ed_dist.first);
    Geo::Segment seg;
    vert_it.get(0)->geom(seg[0]);
    vert_it.get(1)->geom(seg[1]);
    Topo::Iterator<Topo::Type::EDGE, Topo::Type::FACE> face_it(ed_dist.first);
    for (int i = 0; i < sol_num; ++i)
    {
      auto pt = Geo::evaluate(seg, pars[i]);
      for (auto f : face_it)
        fp_map[f].push_back(_pts.size());
      _pts.push_back(pt);
    }
  }
  for (auto& fp : fp_map)
  {
    for (auto it = fp.second.begin(); it != fp.second.end(); ++it)
      for (auto it1 = it; ++it1 != fp.second.end(); )
        _inds.push_back({ *it, *it1 });
  }
  return true;
}

} // namespace Offset

#endif


#include "Sh3FixedPoint.h"

#include <boost/multiprecision/cpp_int.hpp>
#include <cryptoTools/Common/BitVector.h>
namespace aby3
{
	template<typename T, Decimal D>
	fp<T, D> fp<T, D>::operator*(const fp<T, D>& rhs) const
	{
		boost::multiprecision::int128_t v0 = mValue, v1 = rhs.mValue;

		v0 = v0 * v1;
		v0 = v0 / (1ull << mDecimal);

		return {
			static_cast<i64>(v0),
			monostate{}
		};
	}

	//template<typename T, Decimal D>
	//std::ostream& operator<<(std::ostream& o, const fp<T, D>& f)
	//{
	//	auto mask = ((1ull << f.mDecimal) - 1);
	//	auto print = [mask](T v) {
	//		std::stringstream s;
	//		oc::BitVector bv;
	//		bv.append((u8*)& v, D);
	//		s << bv;

	//		s << ".";
	//		bv = {};
	//		bv.append((u8*)& v, 64 - D, D);
	//		s << bv;
	//		return s.str();
	//	};

	//	std::stringstream ss;
	//	//auto& ss = o;
	//	auto vv = (f.mValue / double(1 << f.mDecimal));
	//	o << std::setprecision(std::numeric_limits<long double>::digits10 + 1)
	//		<< vv << std::endl;
	//	o << print(f.mValue) << std::endl;

	//	u64 v = 0;
	//	if (f.mValue >= 0)
	//		v = f.mValue;
	//	else
	//	{
	//		ss << '-';
	//		v =-f.mValue;
	//	}

	//	ss << (v >> f.mDecimal) << ".";

	//	v &= mask;

	//	int i = 0;
	//	if (v)
	//	{
	//		o << "["<<i<<"]  " << print(v) << std::endl;
	//		while (v & mask)
	//		{
	//			++i;
	//			v *= 10;
	//			ss << (v >> f.mDecimal);
	//			o << "[" << i << "]  " << print(v) << std::endl;

	//			v &= mask;
	//			o << "[" << i << "]* " << print(v) << std::endl;
	//		}
	//	}
	//	else
	//	{
	//		ss << '0';
	//	}
	//	o << ss.str();

	//	return o;
	//}

	//template std::ostream& operator<<<i64, D8>(std::ostream& o, const fp<i64, D8>& f);
	//template std::ostream& operator<<<i64, D16>(std::ostream& o, const fp<i64, D16>& f);
	//template std::ostream& operator<<<i64, D32>(std::ostream& o, const fp<i64, D32>& f);

	template struct fp<i64, D0>;
	template struct fp<i64, D8>;
	template struct fp<i64, D16>;
	template struct fp<i64, D32>;

	template struct sf64<D0>;
	template struct sf64<D8>;
	template struct sf64<D16>;
	template struct sf64<D32>;

	template struct sf64Matrix<D0>;
	template struct sf64Matrix<D8>;
	template struct sf64Matrix<D16>;
	template struct sf64Matrix<D32>;

}


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

#include <math.h>
#include "Eigen/Dense"
#include <boost/math/special_functions/digamma.hpp>
#include <boost/math/special_functions/gamma.hpp>

using namespace Eigen;
using namespace std;

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

extern "C" {
    void sparseLocalStepSingleDoc(
        double* ElogLik_d_IN,
        double* alphaEbeta_IN,
        int nnzPerRow,
        int N,
        int K,
        int nCoordAscentIterLP,
        double convThrLP,
        int initProbsToEbeta,
        double* topicCount_d_OUT,
        double* spResp_data_OUT,
        int* spResp_colids_OUT,
        int d,
        int D,
        int* numIterVec_OUT,
        double* maxDiffVec_OUT
        );

    void sparseLocalStepSingleDoc_ActiveOnly(
        double* ElogLik_d_IN,
        double* wc_d_IN,
        double* alphaEbeta_IN,
        int nnzPerRow,
        int N,
        int K,
        int nCoordAscentIterLP,
        double convThrLP,
        int initProbsToEbeta,
        double* topicCount_d_OUT,
        double* spResp_data_OUT,
        int* spResp_colids_OUT,
        int d,
        int D,
        int* numIterVec_OUT,
        double* maxDiffVec_OUT,
        int doTrackELBO,
        double* elboVec_OUT,
        int numRestarts,
        int REVISE_FIRST,
        int REVISE_EVERY,
        int *rAcceptVec_IN, int* rTrialVec_IN,
        int verbose
        );

    void sparseLocalStepSingleDocWithWordCounts(
        double* wordcounts_d_IN,
        double* ElogLik_d_IN,
        double* alphaEbeta_IN,
        int nnzPerRow,
        int N,
        int K,
        int nCoordAscentIterLP,
        double convThrLP,
        int initProbsToEbeta,
        double* topicCount_d_OUT,
        double* spResp_data_OUT,
        int* spResp_colids_OUT
        );
}

// ======================================================== 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 calcELBOForSingleDoc_V2(
    Arr1D_d alphaEbeta,
    Arr1D_d topicCount_d,
    Arr1D_d spResp_data,
    Arr1D_i spResp_colids,
    Arr2D_d ElogLik_d,
    int K,
    int N,
    int nnzPerRow
    )
{
    double ELBO = 0.0;
    // Lalloc
    for (int k = 0; k < K; k++) {
        ELBO += boost::math::lgamma(topicCount_d(k) + alphaEbeta(k));
    }

    // Ldata and Lentropy
    for (int n = 0; n < N; n++) {
        for (int nzk = n * nnzPerRow; nzk < (n+1) * nnzPerRow; nzk++) {
            int k = spResp_colids(nzk);
            ELBO += spResp_data(nzk) * ElogLik_d(n,k);
            if (spResp_data(nzk) > 1e-9) {
                ELBO -= spResp_data(nzk) * log(spResp_data(nzk));
            }
        }
    }
    return ELBO;
}

double calcELBOForSingleDoc_V1(
    ExtArr1D_d alphaEbeta,
    ExtArr1D_d topicCount_d,
    Arr1D_d ElogProb_d,
    Arr1D_i activeTopics_d,
    double totalLogSumResp,
    double sum_gammalnalphaEbeta,
    int Kactive
    )
{
    // Compute Lalloc = \sum_k gammaln(\theta_dk)
    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 updateAssignments_FixPerTokenActiveSet(
    ExtArr2D_d ElogLik_d,
    ExtArr1D_d wc_d,
    ExtArr1D_d alphaEbeta,
    Arr1D_i activeTopics_d,
    ExtArr1D_d& topicCount_d, // & makes sure we can edit passed array in place
    Arr1D_d& ElogProb_d,
    Arr1D_d& logScores_n,
    ExtArr1D_d& spResp_data,
    ExtArr1D_i& spResp_colids,
    int N,
    int K,
    int Kactive,
    int nnzPerRow,
    int doTrackELBO
    )
{
    assert(nnzPerRow > 1);
    double totalLogSumResp = 0.0;
    // UPDATE ElogProb_d using input doc-topic counts
    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 to all zeros
    topicCount_d.fill(0);
    // UPDATE assignments, obeying sparsity constraint
    for (int n = 0; n < N; n++) {
        double sumResp_n = 0.0;
        int m = n * nnzPerRow;
        double maxScore_n;
        for (int ka = 0; ka < nnzPerRow; ka++) {
            int k = spResp_colids(m + ka);
            logScores_n(ka) = ElogProb_d(k) + ElogLik_d(n,k);
            if (ka == 0 || logScores_n(ka) > maxScore_n) {
                maxScore_n = logScores_n(ka);
            }
        }
        for (int nzk = 0; nzk < nnzPerRow; nzk++) {
            spResp_data(m + nzk) = \
                exp(logScores_n(nzk) - maxScore_n);
            sumResp_n += spResp_data(m + nzk);
        }

        for (int nzk = m; nzk < m + nnzPerRow; nzk++) {
            spResp_data(nzk) /= sumResp_n;
            topicCount_d(spResp_colids(nzk)) += \
                wc_d(n) * spResp_data(nzk);
        }
        if (doTrackELBO) {
            totalLogSumResp += wc_d(n) * (maxScore_n + log(sumResp_n));
        }
    } // end for loop over tokens n
    assert(abs(wc_d.sum() - topicCount_d.sum()) < .000001);
    return totalLogSumResp;
}

double updateAssignments_ActiveOnly(
    ExtArr2D_d ElogLik_d,
    ExtArr1D_d wc_d,
    ExtArr1D_d alphaEbeta,
    Arr1D_i activeTopics_d,
    ExtArr1D_d& topicCount_d, // & makes sure we can edit passed array in place
    Arr1D_d& ElogProb_d,
    Arr1D_d& logScores_n,
    Arr1D_d& tempScores_n,
    ExtArr1D_d& spResp_data,
    ExtArr1D_i& spResp_colids,
    int N,
    int K,
    int Kactive,
    int nnzPerRow,
    int iter,
    int initProbsToEbeta,
    int doTrackELBO
    )
{
    double totalLogSumResp = 0.0;

    // UPDATE ElogProb_d using input doc-topic counts
    if (iter == 0 and 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 to all zeros
    topicCount_d.fill(0);
    // UPDATE assignments, obeying sparsity constraint
    for (int n = 0; n < N; n++) {
        int m = n * nnzPerRow;
        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_d(n,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(m) = 1.0;
            spResp_colids(m) = argmax_n;
            topicCount_d(argmax_n) += wc_d(n);
            if (doTrackELBO) {
                totalLogSumResp += wc_d(n) * 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 nzk = m;
            double sumResp_n = 0.0;
            for (int ka = 0; ka < Kactive; ka++) {
                if (logScores_n(ka) >= pivotScore) {
                    spResp_data(nzk) = \
                        exp(logScores_n(ka) - maxScore_n);
                    spResp_colids(nzk) = activeTopics_d(ka);
                    sumResp_n += spResp_data(nzk);
                    nzk += 1;
                }
            }
            //assert(nzk - m == nnzPerRow);

            for (nzk = m; nzk < m + nnzPerRow; nzk++) {
                spResp_data(nzk) /= sumResp_n;
                topicCount_d(spResp_colids(nzk)) += \
                    wc_d(n) * spResp_data(nzk);
            }
            if (doTrackELBO) {
                totalLogSumResp += wc_d(n) * (maxScore_n + log(sumResp_n));
            }
        } // end if statement branch for nnz > 1
    } // end for loop over tokens n
    return totalLogSumResp;
}

void sparseLocalStepSingleDoc_ActiveOnly(
        double* ElogLik_d_IN,
        double* wc_d_IN,
        double* alphaEbeta_IN,
        int nnzPerRow,
        int N,
        int K,
        int nCoordAscentIterLP,
        double convThrLP,
        int initProbsToEbeta,
        double* topicCount_d_OUT,
        double* spResp_data_OUT,
        int* spResp_colids_OUT,
        int d,
        int D,
        int* numIterVec_OUT,
        double* maxDiffVec_OUT,
        int doTrackELBO,
        double* elboVec_OUT,
        int numRestarts,
        int REVISE_FIRST,
        int REVISE_EVERY,
        int* rAcceptVec_IN,
        int* rTrialVec_IN,
        int verbose
        )
{
    nCoordAscentIterLP = max(nCoordAscentIterLP + max(0, initProbsToEbeta), 1);

    // Unpack inputs, treated as fixed constants
    ExtArr2D_d ElogLik_d (ElogLik_d_IN, N, K);
    ExtArr1D_d wc_d (wc_d_IN, N);
    ExtArr1D_d alphaEbeta (alphaEbeta_IN, K);

    // Unpack outputs
    ExtArr1D_d spResp_data (spResp_data_OUT, N * nnzPerRow);
    ExtArr1D_i spResp_colids (spResp_colids_OUT, N * nnzPerRow);
    ExtArr1D_d topicCount_d (topicCount_d_OUT, K);

    ExtArr1D_i numIterVec (numIterVec_OUT, D);
    ExtArr1D_d maxDiffVec (maxDiffVec_OUT, D);
    ExtArr1D_d elboVec (elboVec_OUT, nCoordAscentIterLP);

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

    // 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);

    prevTopicCount_d.fill(-1);
    double maxDiff = N;
    int iter = 0;
    int Kactive = K;
    double ACTIVE_THR = 1e-9;
    double totalLogSumResp = 0.0;
    double sum_gammalnalphaEbeta = 0.0;
    if (doTrackELBO || numRestarts > 0) {
        for (int k = 0; k < K; k++) {
            sum_gammalnalphaEbeta += boost::math::lgamma(alphaEbeta(k));
        }
    }
    // Before any updates... All topics are active!
    for (int k = 0; k < K; k++) {
        activeTopics_d(k) = k;
    }
    if (verbose > 1) {
        printf("Initial topic counts \n");
        for (int k = 0; k < K; k++) {
            printf("%d:%6.1f ", k, topicCount_d(k));
        }
        printf("\n");
    }

    for (iter = 0; iter < nCoordAscentIterLP; iter++) {
        int doReviseActiveSet;
        if (iter >= REVISE_FIRST && Kactive <= nnzPerRow) {
            // Set of active docs is already very small,
            // so nothing big to gain from revising
            doReviseActiveSet = 0;
        } else if (iter < REVISE_FIRST || (iter - 1) % REVISE_EVERY == 0) {
            doReviseActiveSet = 1;
        } else {
            doReviseActiveSet = 0;
        }

        if (iter > 0 || initProbsToEbeta < 0) {
            if (doReviseActiveSet) {
                int newKactive = 0;
                int ia = 0; // index for spare inactive topics
                for (int a = 0; a < Kactive; a++) {
                    int k = activeTopics_d(a);
                    if (topicCount_d(k) > ACTIVE_THR) {
                        activeTopics_d(newKactive) = k;
                        prevTopicCount_d(k) = topicCount_d(k);
                        newKactive += 1;
                    } else if (newKactive < nnzPerRow - ia) {
                        spareActiveTopics_d(ia) = k;
                        ia += 1;
                    }
                }
                // If num topics above threshold is less than nnzPerRow,
                // We need to fill in with some spare empty topics.
                Kactive = newKactive;
                while (Kactive < nnzPerRow) {
                    int k = spareActiveTopics_d(Kactive - newKactive);
                    activeTopics_d(Kactive) = k;
                    Kactive++;
                }
            } else {
                for (int a = 0; a < Kactive; a++) {
                    int k = activeTopics_d(a);
                    prevTopicCount_d(k) = topicCount_d(k);
                }
            }
        }
        assert(Kactive >= nnzPerRow);
        assert(Kactive <= K);
        if (nnzPerRow == 1 || doReviseActiveSet) {
            totalLogSumResp = updateAssignments_ActiveOnly(
                ElogLik_d, wc_d, alphaEbeta, activeTopics_d,
                topicCount_d, ElogProb_d,
                logScores_n, tempScores_n,
                spResp_data, spResp_colids,
                N, K, Kactive, nnzPerRow,
                iter, initProbsToEbeta, doTrackELBO
                );
        } else {
            totalLogSumResp = updateAssignments_FixPerTokenActiveSet(
                ElogLik_d, wc_d, alphaEbeta, activeTopics_d,
                topicCount_d, ElogProb_d,
                logScores_n,
                spResp_data, spResp_colids,
                N, K, Kactive, nnzPerRow,
                doTrackELBO
                );
        }

        if (doTrackELBO) {
            elboVec(iter) = calcELBOForSingleDoc_V1(
                alphaEbeta, topicCount_d, ElogProb_d, activeTopics_d,
                totalLogSumResp, sum_gammalnalphaEbeta, Kactive);
            /*
            double elboV2 = calcELBOForSingleDoc_V2(
                alphaEbeta, topicCount_d, spResp_data, spResp_colids,
                ElogLik_d, K, N, nnzPerRow);

            double elboV1 = calcELBOForSingleDoc_V1(
                alphaEbeta, topicCount_d, ElogProb_d, activeTopics_d,
                totalLogSumResp, sum_gammalnalphaEbeta, Kactive);
            printf(" V1: %.6f\n V2: %.6f\n", elboV2, elboV1);
            */
        }

        if (verbose > 1) {
            printf("end of iter %3d Kactive %3d maxDiff %10.6f\n",
                iter, Kactive, maxDiff);
            for (int ka = 0; ka < Kactive; ka++) {
                int k = activeTopics_d(ka);
                printf("%d:%6.1f ", k, topicCount_d(k));
            }
            printf("\n");
        }

        // END ITERATION. Decide whether to quit early
        if (iter > 0 && iter % 5 == 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 (maxDiff <= convThrLP) {
                break;
            }
        }
    }

    // Figure out which topics are eligible for sparse restarts
    double smallThr = 1e-6;
    double curELBO = 0.0;
    for (int riter = 0; riter < numRestarts; riter++) {
        if (riter == 0) {
            totalLogSumResp = updateAssignments_ActiveOnly(
                ElogLik_d, wc_d, alphaEbeta, activeTopics_d,
                topicCount_d, ElogProb_d,
                logScores_n, tempScores_n,
                spResp_data, spResp_colids,
                N, K, Kactive, nnzPerRow,
                0, 0, 1);
            curELBO = calcELBOForSingleDoc_V1(
                alphaEbeta, topicCount_d, ElogProb_d, activeTopics_d,
                totalLogSumResp, sum_gammalnalphaEbeta, Kactive);
            // Remember the best-known topic-count vector!
            for (int k = 0; k < K; k++) {
                prevTopicCount_d(k) = topicCount_d(k);
            }
        }
        // SEARCH FOR SMALLEST TOPIC HAVE NOT YET TRIED YET
        int numAboveThr = 0;
        double minVal = N + 1.0; // topicCount_d must never have this value
        int minLoc = 0;
        for (int ka = 0; ka < Kactive; ka++) {
            int k = activeTopics_d(ka);
            if (topicCount_d(k) > smallThr) {
                numAboveThr += 1;
                if (topicCount_d(k) < minVal) {
                    minVal = topicCount_d(k);
                    minLoc = k;
                }
            }
        }
        if (numAboveThr == 1) {
            break;
        }

        if (verbose) {
            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");
        }
        smallThr = minVal;
        topicCount_d(minLoc) = 0.0;
        if (verbose) {
            printf(
                "RESTART: Set index %d to zero (%d left)\n",
                minLoc, numAboveThr - 1);
            for (int ka = 0; ka < Kactive; ka++) {
                int k = activeTopics_d(ka);
                printf("%02d:%06.2f ", k, topicCount_d(k));
            }
            printf("\n");
        }
        int NSTEP = 2;
        for (int step = 0; step < NSTEP; step++) {
            totalLogSumResp = updateAssignments_ActiveOnly(
                ElogLik_d, wc_d, alphaEbeta, activeTopics_d,
                topicCount_d, ElogProb_d,
                logScores_n, tempScores_n,
                spResp_data, spResp_colids,
                N, K, Kactive, nnzPerRow,
                step, 0, step == NSTEP-1);
        }
        // If the change is small, abandon this
        double propELBO;
        if (abs(prevTopicCount_d(minLoc) -topicCount_d(minLoc)) < 1e-6) {
            propELBO = curELBO;
        } else {
            propELBO = calcELBOForSingleDoc_V1(
                alphaEbeta, topicCount_d, ElogProb_d, activeTopics_d,
                totalLogSumResp, sum_gammalnalphaEbeta, Kactive);
        }

        if (verbose) {
            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(minLoc));
                printf(" afterCount: %.6f\n", topicCount_d(minLoc));
                printf("gainELBO % .6e  *** ACCEPTED \n", propELBO - curELBO);
            } else {
                printf("gainELBO % .6e      rejected \n", propELBO - curELBO);
            }
        }
        // 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);
            }
            rAcceptVec(0) += 1;
        } else {
            for (int ka = 0; ka < Kactive; ka++) {
                int k = activeTopics_d(ka);
                topicCount_d(k) = prevTopicCount_d(k);
            }
        }
        rTrialVec(0) += 1;
    }
    if (numRestarts > 0) {
        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
        updateAssignments_ActiveOnly(
            ElogLik_d, wc_d, alphaEbeta, activeTopics_d,
            topicCount_d, ElogProb_d,
            logScores_n, tempScores_n,
            spResp_data, spResp_colids,
            N, K, Kactive, nnzPerRow,
            0, 0, 0);
    }
    maxDiffVec(d) = maxDiff;
    numIterVec(d) = iter; // will already have +1 from last iter of for loop
}

void sparseLocalStepSingleDoc(
        double* ElogLik_d_IN,
        double* alphaEbeta_IN,
        int nnzPerRow,
        int N,
        int K,
        int nCoordAscentIterLP,
        double convThrLP,
        int initProbsToEbeta,
        double* topicCount_d_OUT,
        double* spResp_data_OUT,
        int* spResp_colids_OUT,
        int d,
        int D,
        int* numIterVec_OUT,
        double* maxDiffVec_OUT
        )
{
    // Unpack inputs, treated as fixed constants
    ExtArr2D_d ElogLik_d (ElogLik_d_IN, N, K);
    ExtArr1D_d alphaEbeta (alphaEbeta_IN, K);
    // Unpack outputs
    ExtArr1D_d spResp_data (spResp_data_OUT, N * nnzPerRow);
    ExtArr1D_i spResp_colids (spResp_colids_OUT, N * nnzPerRow);
    ExtArr1D_d topicCount_d (topicCount_d_OUT, K);

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

    // Temporary storage
    Arr1D_d ElogProb_d (K);
    Arr1D_d prevTopicCount_d (K);
    Arr1D_d logScores_n (K);
    Arr1D_d tempScores_n (K);

    prevTopicCount_d.fill(-1);
    double maxDiff = N;
    int iter = 0;

    for (iter = 0; iter < nCoordAscentIterLP + initProbsToEbeta; iter++) {

        if (iter == 0 and initProbsToEbeta == 1) {
            for (int k = 0; k < K; k++) {
                ElogProb_d(k) = log(alphaEbeta(k));
            }
        } else {
            for (int k = 0; k < K; k++) {
                ElogProb_d(k) = boost::math::digamma(
                    topicCount_d(k) + alphaEbeta(k));
            }
        }

        topicCount_d.fill(0);
        // Step over each data atom
        for (int n = 0; n < N; n++) {
            int m = n * nnzPerRow;
            int argmax_n = 0;
            logScores_n(0) = ElogProb_d(0) + ElogLik_d(n,0);
            double maxScore_n = logScores_n(0);
            for (int k = 1; k < K; k++) {
                logScores_n(k) = ElogProb_d(k) + ElogLik_d(n,k);
                if (logScores_n(k) > maxScore_n) {
                    maxScore_n = logScores_n(k);
                    argmax_n = k;
                }
            }
            if (nnzPerRow == 1) {
                spResp_data(m) = 1.0;
                spResp_colids(m) = argmax_n;
                // Update topicCount_d
                topicCount_d(argmax_n) += 1.0;
            } 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);
                int nzk = 0;
                double sumResp_n = 0.0;
                for (int k = 0; k < K; k++) {
                    if (logScores_n(k) >= pivotScore) {
                        spResp_data(m + nzk) = \
                            exp(logScores_n(k) - maxScore_n);
                        spResp_colids(m + nzk) = k;
                        sumResp_n += spResp_data(m + nzk);
                        nzk += 1;
                    }
                }
                // Normalize for doc-topic counts
                for (int nzk = 0; nzk < nnzPerRow; nzk++) {
                    spResp_data(m + nzk) /= sumResp_n;
                    topicCount_d(spResp_colids(m + nzk)) += \
                        spResp_data(m + nzk);
                }
            }
        }
        // END ITERATION. Decide whether to quit early
        if (iter > 0 && iter % 5 == 0) {
            double absDiff_k = 0.0;
            maxDiff = 0.0;
            for (int k = 0; k < K; k++) {
                absDiff_k = abs(prevTopicCount_d(k) - topicCount_d(k));
                if (absDiff_k > maxDiff) {
                    maxDiff = absDiff_k;
                }
                prevTopicCount_d(k) = topicCount_d(k); // copy over
            }
            if (maxDiff <= convThrLP) {
                break;
            }
        }
    }
    maxDiffVec(d) = maxDiff;
    numIterVec(d) = iter + 1;
}

void sparseLocalStepSingleDocWithWordCounts(
        double* wordcounts_d_IN,
        double* ElogLik_d_IN,
        double* alphaEbeta_IN,
        int nnzPerRow,
        int N,
        int K,
        int nCoordAscentIterLP,
        double convThrLP,
        int initProbsToEbeta,
        double* topicCount_d_OUT,
        double* spResp_data_OUT,
        int* spResp_colids_OUT
        )
{
    // Unpack inputs, treated as fixed constants
    ExtArr1D_d wc_d (wordcounts_d_IN, N);
    ExtArr2D_d ElogLik_d (ElogLik_d_IN, N, K);
    ExtArr1D_d alphaEbeta (alphaEbeta_IN, K);
    // Unpack outputs
    ExtArr1D_d spResp_data (spResp_data_OUT, N * nnzPerRow);
    ExtArr1D_i spResp_colids (spResp_colids_OUT, N * nnzPerRow);
    ExtArr1D_d topicCount_d (topicCount_d_OUT, K);
    // Temporary storage
    VectorXd ElogProb_d (K);
    VectorXd prevTopicCount_d (K);
    VectorXd logScores_n (K);
    VectorXd tempScores_n (K);
    prevTopicCount_d.fill(-1);
    int iter = 0;
    double maxDiff = 0.0;
    for (iter = 0; iter < nCoordAscentIterLP + initProbsToEbeta; iter++) {

        if (iter == 0 and initProbsToEbeta == 1) {
            for (int k = 0; k < K; k++) {
                ElogProb_d(k) = log(alphaEbeta(k));
            }
        } else {
            for (int k = 0; k < K; k++) {
                ElogProb_d(k) = boost::math::digamma(
                    topicCount_d(k) + alphaEbeta(k));
            }
        }

        topicCount_d.fill(0);
        // Step over each data atom
        for (int n = 0; n < N; n++) {
            int m = n * nnzPerRow;
            int argmax_n = 0;
            logScores_n(0) = ElogProb_d(0) + ElogLik_d(n,0);
            double maxScore_n = logScores_n(0);
            for (int k = 1; k < K; k++) {
                logScores_n(k) = ElogProb_d(k) + ElogLik_d(n,k);
                if (logScores_n(k) > maxScore_n) {
                    maxScore_n = logScores_n(k);
                    argmax_n = k;
                }
            }
            if (nnzPerRow == 1) {
                spResp_data(m) = 1.0;
                spResp_colids(m) = argmax_n;
                // Update topicCount_d
                topicCount_d(argmax_n) += wc_d(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);
                int nzk = 0;
                double sumResp_n = 0.0;
                for (int k = 0; k < K; k++) {
                    if (logScores_n(k) >= pivotScore) {
                        spResp_data(m + nzk) = \
                            exp(logScores_n(k) - maxScore_n);
                        spResp_colids(m + nzk) = k;
                        sumResp_n += spResp_data(m + nzk);
                        nzk += 1;
                    }
                }
                // Normalize for doc-topic counts
                for (int nzk = 0; nzk < nnzPerRow; nzk++) {
                    spResp_data(m + nzk) /= sumResp_n;
                    topicCount_d(spResp_colids(m + nzk)) += \
                        wc_d(n) * spResp_data(m + nzk);
                }
            }
        }
        // END ITERATION. Decide whether to quit early
        if (iter > 0 && iter % 5 == 0) {
            double absDiff_k = 0.0;
            maxDiff = 0.0;
            for (int k = 0; k < K; k++) {
                absDiff_k = abs(prevTopicCount_d(k) - topicCount_d(k));
                if (absDiff_k > maxDiff) {
                    maxDiff = absDiff_k;
                }
                prevTopicCount_d(k) = topicCount_d(k); // copy over
            }
            if (maxDiff <= convThrLP) {
                break;
            }
        }
    }
}