//------------------------------------------------------------------------------
/// \file
/// \brief Provides functionality to do spatial queries on a set of triangles.
/// Uses the boost RTree for queries.
/// \ingroup geometry
/// \copyright (C) Copyright Aquaveo 2018. Distributed under FreeBSD License
/// (See accompanying file LICENSE or https://aqaveo.com/bsd/license.txt)
//------------------------------------------------------------------------------

//----- Included files ---------------------------------------------------------

// 1. Precompiled header

// 2. My header
#include <xmsgrid/geometry/GmTriSearch.h>

// 3. Standard Library Headers

// 4. External Library Headers
#pragma warning(push)
#pragma warning(disable : 4512) // boost code: no assignment operator
#include <boost/geometry.hpp>
#include <boost/geometry/geometries/polygon.hpp>
#include <boost/geometry/index/rtree.hpp>
#include <boost/unordered_map.hpp>
#include <boost/iterator/counting_iterator.hpp>
#pragma warning(pop)

// 5. Shared Headers
#include <xmscore/math/math.h>
#include <xmscore/stl/utility.h>
#include <xmsgrid/geometry/GmBoostTypes.h>
#include <xmsgrid/geometry/geoms.h>
#include <xmsgrid/geometry/GmPtSearch.h>
#include <xmscore/misc/xmstype.h>
#include <xmscore/misc/XmError.h>
#include <xmscore/misc/XmConst.h>

// 6. Non-shared Headers

//----- Forward declarations ---------------------------------------------------

//----- External globals -------------------------------------------------------

//----- Namespace declaration --------------------------------------------------
namespace xms
{
namespace bg = boost::geometry;
namespace bgi = bg::index;

typedef size_t value;
typedef bgi::quadratic<16> qRtree;

//----- Constants / Enumerations -----------------------------------------------

//----- Classes / Structs ------------------------------------------------------

class GmTriSearchImpl;

#define BARYTOL 1e-6 ///< This is an odd duck tolerance

/// \brief provides indexing for the spatial index
class idx_tri
{
public:
  typedef const GmBstBox3d result_type; ///< created to follow boost examples
  /// \ brief constructor
  /// \param[in] a_pts Point locations
  /// \param[in] a_tris Triangles referencing point indexes
  /// \param[in] a_tol Tolerance for expanding triangle extents
  explicit idx_tri(const Pt3d* a_pts, const int* a_tris, double a_tol)
  : m_pts(a_pts)
  , m_tris(a_tris)
  , m_tol(a_tol)
  {
  }
  /// \brief calculates triangle extents based on a triangle index
  /// \param i: index.
  /// \return result_type.
  result_type operator()(size_t i) const
  {
    i *= 3;
    const Pt3d& p0(m_pts[m_tris[i]]);
    const Pt3d& p1(m_pts[m_tris[i + 1]]);
    const Pt3d& p2(m_pts[m_tris[i + 2]]);
    GmBstBox3d b;
    b.max_corner() = b.min_corner() = p0;
    if (b.min_corner().x > p1.x)
      b.min_corner().x = p1.x;
    if (b.min_corner().x > p2.x)
      b.min_corner().x = p2.x;
    if (b.max_corner().x < p1.x)
      b.max_corner().x = p1.x;
    if (b.max_corner().x < p2.x)
      b.max_corner().x = p2.x;

    if (b.min_corner().y > p1.y)
      b.min_corner().y = p1.y;
    if (b.min_corner().y > p2.y)
      b.min_corner().y = p2.y;
    if (b.max_corner().y < p1.y)
      b.max_corner().y = p1.y;
    if (b.max_corner().y < p2.y)
      b.max_corner().y = p2.y;
    b.min_corner() -= m_tol;
    b.max_corner() += m_tol;
    return b;
  }
  const Pt3d* m_pts; ///< point locations
  const int* m_tris; ///< triangles referencing point indexes
  double m_tol;      ///< tolerance used to expand extents of triangle
};

/// \brief Structure for Barycentric coordinate calculations
struct BarycentricVals
{
  Pt3d orig;      ///< original point location
  double coef[6]; ///< barycentric coefficients
  int dir;        ///< direction based on the xyz's of the points making the triangle
};

////////////////////////////////////////////////////////////////////////////////
/// \brief Implementation of GmTriSearch
class GmTriSearchImpl : public GmTriSearch
{
public:
  GmTriSearchImpl();
  virtual ~GmTriSearchImpl();
  virtual void TrisToSearch(BSHP<std::vector<Pt3d>> a_pts, BSHP<std::vector<int>> a_tris) override;
  virtual void SetPtActivity(DynBitset& a_activity) override;
  virtual void SetTriActivity(DynBitset& a_activity) override;
  /// \brief Activity of the points based on the triangle activity
  /// \return boost::dynamic_bitset of size_t
  virtual DynBitset GetPtActivity() const override { return PointActivityFromTriActivity(); }
  virtual DynBitset GetTriActivity() const override;
  bool ActiveTri(int a_idx) const;
  virtual int TriContainingPt(const Pt3d& a_pt) override;
  virtual void TriEnvelopsContainingPt(const Pt3d& a_pt, std::vector<int>& a_tris) override;
  virtual void TriEnvelopesOverlap(const Pt3d& a_pMin,
                                   const Pt3d& a_pMax,
                                   std::vector<int>& a_tris) override;

  virtual bool InterpWeights(const Pt3d& a_pt,
                             std::vector<int>& a_idxs,
                             std::vector<double>& a_wts) override;
  virtual bool InterpWeightsTriangleIdx(const Pt3d& a_pt,
                                        int& a_triangleIdx,
                                        std::vector<int>& a_idxs,
                                        std::vector<double>& a_wts) override;
  virtual const BSHP<VecPt3d> GetPoints() const override;
  virtual const BSHP<VecInt>  GetTriangles() const override;

  virtual std::string ToString() const override;

  void CreateRTree();

  void PointIdxesFromTriIdx(int a_triIdx, int a_ptIdxes[3]);
  int FindTriangle(const Pt3d& a_pt, int ix[3], Pt3d& weights);
  void GetTriBarycentricVals(int a_idx, int a_ix[3], BarycentricVals& a_b);

  DynBitset PointActivityFromTriActivity() const;
  BSHP<GmPtSearch> CreatePtSearch();

  Pt3d m_min; ///< mininum extents of all points
  Pt3d m_max; ///< maximum extents of all points

  BSHP<std::vector<Pt3d>> m_pts; ///< point locations
  BSHP<std::vector<int>> m_tris; ///< triangles that reference point indexes
  DynBitset m_triActivity;       ///< activity of the triangles

  /// we create a cache of these for faster interp the second time we hit the
  /// triangle
  boost::unordered_map<size_t, BarycentricVals> m_cache;
  /// spatial index of the triangle extents
  bgi::rtree<value, qRtree, idx_tri>* m_rTree;
};

//----- Internal functions -----------------------------------------------------

//------------------------------------------------------------------------------
/// \brief Converts from Cartesian coordinates to Barycentric coordinates (local
/// triangle coords).
/// \param a_pt1: Point 1.
/// \param a_pt2: Point 2.
/// \param a_pt3: Point 3.
/// \param a_b: Barycentric values.
//------------------------------------------------------------------------------
void iCartToBary(const Pt3d& a_pt1, const Pt3d& a_pt2, const Pt3d& a_pt3, BarycentricVals& a_b)
{
  Pt3d a(a_pt1.x, a_pt1.y, 0.0);
  Pt3d b(a_pt2.x, a_pt2.y, 0.0);
  Pt3d c(a_pt3.x, a_pt3.y, 0.0);
  Pt3d norm(0, 0, 1);
  gmBaryPrepare(&a, &b, &c, &norm, &a_b.orig, a_b.coef, &a_b.dir, true);
} // iCartToBary
//------------------------------------------------------------------------------
/// \brief Calculates weights at the 3 points of a triangle given the location
/// "a_p" and the Barycentric coords of the triangle.
/// \param a_p: A point
/// \param a_b: Barycentric values.
/// \param weights: Weights.
//------------------------------------------------------------------------------
void iGetBarycentricCoords(const Pt3d& a_p, BarycentricVals& a_b, Pt3d& weights)
{
  Pt3d p(a_p.x, a_p.y, 0.0);
  gmCartToBary(&p, &a_b.orig, a_b.coef, a_b.dir, &weights);
} // iGetBarycentricCoords

//----- Class / Function definitions -------------------------------------------

//------------------------------------------------------------------------------
/// \brief Creates an TriSearch class
/// \return GmTriSearch.
//------------------------------------------------------------------------------
BSHP<GmTriSearch> GmTriSearch::New()
{
  BSHP<GmTriSearch> ptr(new GmTriSearchImpl());
  return ptr;
} // TriSearch::New
//------------------------------------------------------------------------------
/// \brief Constructor
//------------------------------------------------------------------------------
GmTriSearch::GmTriSearch()
{
} // TriSearch::TriSearch
//------------------------------------------------------------------------------
/// \brief Destructor
//------------------------------------------------------------------------------
GmTriSearch::~GmTriSearch()
{
} // TriSearch::~TriSearch

////////////////////////////////////////////////////////////////////////////////
/// \class GmTriSearchImpl
/// \brief Provides functionality to do spatial queries on a set of triangles.
/// Uses the boost RTree for queries.
////////////////////////////////////////////////////////////////////////////////
//------------------------------------------------------------------------------
/// \brief Constructor
//------------------------------------------------------------------------------
GmTriSearchImpl::GmTriSearchImpl()
: m_min(XM_FLT_HIGHEST)
, m_max(XM_FLT_LOWEST)
, m_pts(new std::vector<Pt3d>())
, m_tris(new std::vector<int>())
, m_triActivity()
, m_rTree(nullptr)
{
} // GmTriSearchImpl::GmTriSearchImpl
//------------------------------------------------------------------------------
/// \brief Destructor
//------------------------------------------------------------------------------
GmTriSearchImpl::~GmTriSearchImpl()
{
  if (m_rTree)
    delete (m_rTree);
} // GmTriSearchImpl::~GmTriSearchImpl
//------------------------------------------------------------------------------
/// \brief Adds the triangles to the class
/// \param a_pts Array of point locations.
/// \param a_tris Array of triangles that references the a_pts array. This
/// array will have a size that is a multiple of 3. The first 3 locations in
/// array represent the first triangle and will have indices that correspond to
/// locations in the a_pts array.
//------------------------------------------------------------------------------
void GmTriSearchImpl::TrisToSearch(BSHP<std::vector<Pt3d>> a_pts, BSHP<std::vector<int>> a_tris)
{
  // remove existing rtree and cache
  if (m_rTree)
    delete (m_rTree);
  m_cache.clear();

  m_pts = a_pts;
  m_tris = a_tris;
  CreateRTree();
} // GmTriSearchImpl::TrisToSearch
//------------------------------------------------------------------------------
/// \brief Creates the RTree of the triangles
//------------------------------------------------------------------------------
void GmTriSearchImpl::CreateRTree()
{
  if (m_pts->empty() || m_tris->empty())
    return;

  m_min = XM_FLT_HIGHEST;
  m_max = XM_FLT_LOWEST;
  for (size_t i = 0; i < m_pts->size(); ++i)
  {
    Pt3d& pt = (*m_pts)[i];
    if (pt.x < m_min.x)
      m_min.x = pt.x;
    if (pt.y < m_min.y)
      m_min.y = pt.y;
    if (pt.z < m_min.z)
      m_min.z = pt.z;
    if (pt.x > m_max.x)
      m_max.x = pt.x;
    if (pt.y > m_max.y)
      m_max.y = pt.y;
    if (pt.z > m_max.z)
      m_max.z = pt.z;
  }
  double tol = sqr(m_max.x - m_min.x) + sqr(m_max.y - m_min.y);
  tol = sqrt(tol) / 1e9;

  qRtree rtree;
  size_t nTri = m_tris->size() / 3;
  idx_tri getter(&(*m_pts)[0], &(*m_tris)[0], tol);
  m_rTree = new bgi::rtree<size_t, qRtree, idx_tri>(
    boost::counting_iterator<size_t>(0), boost::counting_iterator<size_t>(nTri), rtree, getter);
} // GmTriSearchImpl::CreateRTree
//------------------------------------------------------------------------------
/// \brief Modifies the activity bitset of the points.
/// \param a_activity bitset of the activity of the points
//------------------------------------------------------------------------------
void GmTriSearchImpl::SetPtActivity(DynBitset& a_activity)
{
  XM_ENSURE_TRUE_NO_ASSERT(a_activity.size() == m_pts->size());

  // set triangle activity from point activity
  m_triActivity.clear();
  m_triActivity.resize(m_tris->size(), 1);
  for (size_t i = 0; i < m_tris->size(); i += 3)
  {
    if (!a_activity.test((*m_tris)[i]) || !a_activity.test((*m_tris)[i + 1]) ||
        !a_activity.test((*m_tris)[i + 2]))
    {
      m_triActivity[i] = m_triActivity[i + 1] = m_triActivity[i + 2] = 0;
    }
  }
} // GmTriSearchImpl::SetPtActivity
//------------------------------------------------------------------------------
/// \brief Modifies the activity bitset of the triangles.
/// \param a_activity bitset of the activity of the triangles
//------------------------------------------------------------------------------
void GmTriSearchImpl::SetTriActivity(DynBitset& a_activity)
{
  if (a_activity.empty())
  {
    m_triActivity = a_activity;
    return;
  }

  XM_ENSURE_TRUE_NO_ASSERT(a_activity.size() == m_tris->size() / 3);

  m_triActivity.reset();
  m_triActivity.resize(m_tris->size(), 1);
  size_t idx;
  for (size_t i = 0; i < a_activity.size(); ++i)
  {
    idx = i * 3;
    m_triActivity[idx + 0] = a_activity[i];
    m_triActivity[idx + 1] = a_activity[i];
    m_triActivity[idx + 2] = a_activity[i];
  }
} // GmTriSearchImpl::SetTriActivity
//------------------------------------------------------------------------------
/// \brief Gets the activity bitset of the triangles.
/// \return bitset of the activity of the triangles
//------------------------------------------------------------------------------
DynBitset GmTriSearchImpl::GetTriActivity() const
{
  DynBitset rval;
  if (m_triActivity.empty() || m_tris->empty())
    return rval;
  rval.resize(m_tris->size() / 3, 1);
  for (size_t i = 0; i < rval.size(); ++i)
  {
    size_t idx = i * 3;
    rval[i] = m_triActivity[idx];
  }
  return rval;
} // GmTriSearchImpl::GetTriActivity
//------------------------------------------------------------------------------
/// \brief Tests if a triangle is active. If any of the points of the tri
/// are not active then the tri is not active
/// \param a_idx Triangle index
/// \return true if the triangle a_idx is active.
//------------------------------------------------------------------------------
bool GmTriSearchImpl::ActiveTri(int a_idx) const
{
  if (m_triActivity.empty())
    return true;
  if (!m_triActivity.test(a_idx))
    return false;
  return true;
} // GmTriSearchImpl::ActiveTri
//------------------------------------------------------------------------------
/// \brief Find the triangle containing the point
/// \param a_pt Location used to find a triangle.
/// \return Index of triangle containing a_pt. If XM_NONE is returned then no
/// triangle contained the point.
//------------------------------------------------------------------------------
int GmTriSearchImpl::TriContainingPt(const Pt3d& a_pt)
{
  int ix[3];
  Pt3d weights;
  return FindTriangle(a_pt, ix, weights);
} // GmTriSearchImpl::TriContainingPt
//------------------------------------------------------------------------------
/// \brief Find all triangles whose envelope contains the point.
/// \param a_pt Location used to find a triangle.
/// \param a_tris The indices to triangles whose envelop contains the point
//------------------------------------------------------------------------------
void GmTriSearchImpl::TriEnvelopsContainingPt(const Pt3d& a_pt, VecInt& a_tris)
{
  int idx;
  a_tris.resize(0);
  Pt3d p(a_pt.x, a_pt.y, 0);
  std::vector<value> vals;
  m_rTree->query(bgi::intersects(p), std::back_inserter(vals));
  for (size_t i = 0; i < vals.size(); ++i)
  {
    idx = (int)(vals[i] * 3);
    a_tris.push_back(idx);
  }
} // GmTriSearchImpl::TriEnvelopsContainingPt
//------------------------------------------------------------------------------
/// \brief Find all triangles whose envelope overlaps the envelope defined by
/// a_pMin and a_pMax.
/// \param a_pMin min x,y location of the box
/// \param a_pMax max x,y location of the box
/// \param a_tris The indices to triangles whose envelope intersects with the
/// input envelope.
//------------------------------------------------------------------------------
void GmTriSearchImpl::TriEnvelopesOverlap(const Pt3d& a_pMin, const Pt3d& a_pMax, VecInt& a_tris)
{
  GmBstBox3d bx;
  bx.max_corner() = a_pMax;
  bx.max_corner().z = 0;
  bx.min_corner() = a_pMin;
  bx.min_corner().z = 0;
  int idx;
  a_tris.resize(0);
  std::vector<value> vals;
  // m_rTree->query(bgi::overlaps(bx), std::back_inserter(vals));
  m_rTree->query(bgi::intersects(bx), std::back_inserter(vals));
  for (size_t i = 0; i < vals.size(); ++i)
  {
    idx = (int)(vals[i] * 3);
    a_tris.push_back(idx);
  }
} // GmTriSearchImpl::TriContainingPt
//------------------------------------------------------------------------------
/// \brief Use the stored triangles to get interpolation weights for a point.
/// Returns false if the point is outside the triangles.
/// \param a_pt Location that is interpolated to.
/// \param a_idxs Vector of triangle point indices filled in by this method.
/// \param a_wts Vector of triangle point weights filled in by this method.
/// \return false if a_pt is not inside of any triangle
//------------------------------------------------------------------------------
bool GmTriSearchImpl::InterpWeights(const Pt3d& a_pt, VecInt& a_idxs, VecDbl& a_wts)
{
  int triangleIdx;
  bool result = InterpWeightsTriangleIdx(a_pt, triangleIdx, a_idxs, a_wts);
  return result;
} // GmTriSearchImpl::InterpWeights
//------------------------------------------------------------------------------
/// \brief Use the stored triangles to get interpolation weights for a point.
/// Returns false if the point is outside the triangles.
/// \param a_pt Location that is interpolated to.
/// \param a_triangleIdx Found triangles index in triangle array
///                      (zero based count*3).
/// \param a_idxs Vector of triangle point indices filled in by this method.
/// \param a_wts Vector of triangle point weights filled in by this method.
/// \return false if a_pt is not inside of any triangle
//------------------------------------------------------------------------------
bool GmTriSearchImpl::InterpWeightsTriangleIdx(const Pt3d& a_pt, int& a_triangleIdx,
                                               VecInt& a_idxs, VecDbl& a_wts)
{
  int ix[3];
  Pt3d weights;
  a_triangleIdx = FindTriangle(a_pt, ix, weights);
  if (a_triangleIdx != XM_NONE)
  {
    a_wts.resize(0);
    a_wts.push_back(weights.x);
    a_wts.push_back(weights.y);
    a_wts.push_back(weights.z);
    a_idxs.assign(ix, ix + 3);
    return true;
  }
  else
  {
    a_idxs.clear();
    a_wts.clear();
  }
  return false;
} // GmTriSearchImpl::InterpWeightsTriangleIdx
//------------------------------------------------------------------------------
/// \brief Returns the points
/// \return constant shared pointer to the vector of points.
//------------------------------------------------------------------------------
const BSHP<VecPt3d> GmTriSearchImpl::GetPoints() const
{
  return m_pts;
} // GmTriSearchImpl::GetPoints
//------------------------------------------------------------------------------
/// \brief Returns the triangles
/// \return constant shared pointer to the vector of triangles.
//------------------------------------------------------------------------------
const BSHP<VecInt>  GmTriSearchImpl::GetTriangles() const
{
  return m_tris;
} // GmTriSearchImpl::GetTriangles
//------------------------------------------------------------------------------
/// \brief Write the internals to a string.
/// \return The string.
//------------------------------------------------------------------------------
std::string GmTriSearchImpl::ToString() const
{
  std::stringstream ss;
  ss << m_min << "=min" << m_max << "=max"
     << "\n";

  if (m_pts)
    VecToStream(ss, *m_pts, "pts");
  if (m_tris)
    VecToStream(ss, *m_tris, "tris");
  VecToStream(ss, m_triActivity, "triActivity");
  return ss.str();
} // GmTriSearchImpl::ToString
//------------------------------------------------------------------------------
/// \brief Fills in the point indices based on the triangle index.
/// \param a_triIdx The triangle index in the m_tris vector.
/// \param a_ptIdxes Array of point indices filled by this method. The indices
/// correspond to locations in the m_pts vector.
//------------------------------------------------------------------------------
void GmTriSearchImpl::PointIdxesFromTriIdx(int a_triIdx, int a_ptIdxes[3])
{
  a_ptIdxes[0] = a_ptIdxes[1] = a_ptIdxes[2] = -1;
  if (a_triIdx < 0 || a_triIdx + 2 >= (int)m_tris->size())
  {
    std::stringstream msg;
    msg << "Index out of range. File: " << __FILE__ << ", Line:" << __LINE__;
    XM_LOG(xmlog::debug, msg.str());
    return;
  }
  a_ptIdxes[0] = (*m_tris)[a_triIdx + 0];
  a_ptIdxes[1] = (*m_tris)[a_triIdx + 1];
  a_ptIdxes[2] = (*m_tris)[a_triIdx + 2];
} // GmTriSearchImpl::PointIdxesFromTriIdx
//------------------------------------------------------------------------------
/// \brief Find the triangle containing the location a_pt.
/// \param a_pt The location to interpolate to.
/// \param ix The point indices from the triangle containing a_pt.
/// \param weights The weights associated with the points ix.
/// \return index to the triangle containg a_pt.
//------------------------------------------------------------------------------
int GmTriSearchImpl::FindTriangle(const Pt3d& a_pt, int ix[3], Pt3d& weights)
{
  if (!m_tris || m_tris->empty() || !m_pts || m_pts->empty())
  {
    XM_LOG(xmlog::error, "Unable to find triangle; no points or triangles exist.");
    return XM_NONE;
  }
  Pt3d p(a_pt.x, a_pt.y, 0);
  std::vector<value> vals;
  m_rTree->query(bgi::intersects(p), std::back_inserter(vals));
  if (vals.empty())
    return XM_NONE;

  size_t idx;
  for (size_t i = 0; i < vals.size(); ++i)
  {
    idx = vals[i] * 3;
    // if any point is inactive then skip this triangle
    if (!ActiveTri((int)idx))
      continue;
    ix[0] = (*m_tris)[idx];
    ix[1] = (*m_tris)[idx + 1];
    ix[2] = (*m_tris)[idx + 2];

    GmBstPoly3d poly;
    for (size_t j = 0; j < 3; ++j)
    {
      bg::exterior_ring(poly).push_back((*m_pts)[ix[j]]);
    }
    bg::exterior_ring(poly).push_back((*m_pts)[ix[0]]);
    if (bg::covered_by(a_pt, poly))
    {
      BarycentricVals b;
      GetTriBarycentricVals((int)idx, ix, b);
      iGetBarycentricCoords(p, b, weights);
      return (int)idx;
    }
  }
  // additional checks with some tolerancing
  for (size_t i = 0; i < vals.size(); ++i)
  {
    idx = vals[i] * 3;
    // if any point is inactive then skip this triangle
    if (!ActiveTri((int)idx))
      continue;
    ix[0] = (*m_tris)[idx];
    ix[1] = (*m_tris)[idx + 1];
    ix[2] = (*m_tris)[idx + 2];
    BarycentricVals b;
    GetTriBarycentricVals((int)idx, ix, b);
    iGetBarycentricCoords(p, b, weights);
    // if (GTEQ_TOL(weights.x, 0, BARYTOL) && GTEQ_TOL(weights.y, 0, BARYTOL) &&
    //    GTEQ_TOL(weights.z, 0, BARYTOL)) {
    if (weights.x >= -BARYTOL && weights.y >= -BARYTOL && weights.z >= -BARYTOL)
    {
      return (int)idx;
    }
  }

  return XM_NONE;
} // GmTriSearchImpl::FindTriangle
//------------------------------------------------------------------------------
/// \brief Gets the barycentric values for a triangle
/// \param a_idx Triangle index
/// \param a_ix Point indexes that make up triangle
/// \param a_b The barycentric values
//------------------------------------------------------------------------------
void GmTriSearchImpl::GetTriBarycentricVals(int a_idx, int a_ix[3], BarycentricVals& a_b)
{
  if (m_cache.find(a_idx) != m_cache.end())
  {
    a_b = m_cache[a_idx];
  }
  else
  {
    iCartToBary((*m_pts)[a_ix[0]], (*m_pts)[a_ix[1]], (*m_pts)[a_ix[2]], a_b);
    m_cache[a_idx] = a_b;
  }
} // GmTriSearchImpl::GetTriBarycentricVals
//------------------------------------------------------------------------------
/// \brief Gets the activity of the points from the triangles
/// \return bitset of the activity of the points.
//------------------------------------------------------------------------------
DynBitset GmTriSearchImpl::PointActivityFromTriActivity() const
{
  DynBitset r;
  if (m_triActivity.empty())
    return r;

  // all points inactive, if any triangle attached to point is active then
  // the point will be active
  r.resize(m_pts->size(), 0);
  for (size_t i = 0; i < m_tris->size(); i += 3)
  {
    if (ActiveTri((int)i))
    {
      r[(*m_tris)[i + 0]] = 1;
      r[(*m_tris)[i + 1]] = 1;
      r[(*m_tris)[i + 2]] = 1;
    }
  }
  return r;
} // GmTriSearchImpl::PointActivityFromTriActivity
//------------------------------------------------------------------------------
/// \brief creates the rtree of the points
/// \return GmPtSearch.
//------------------------------------------------------------------------------
BSHP<GmPtSearch> GmTriSearchImpl::CreatePtSearch()
{
  BSHP<GmPtSearch> pts = GmPtSearch::New(true);
  pts->PtsToSearch(m_pts);
  DynBitset b = PointActivityFromTriActivity();
  pts->SetActivity(b);
  return pts;
} // GmTriSearchImpl::CreatePtSearch

} // namespace xms

#ifdef CXX_TEST
////////////////////////////////////////////////////////////////////////////////

#include <xmsgrid/geometry/GmTriSearch.t.h>

#include <xmscore/testing/TestTools.h>

#include <xmsgrid/triangulate/TrTin.h>
#include <xmsgrid/triangulate/TrTin.t.h>

using namespace xms;

////////////////////////////////////////////////////////////////////////////////
/// \class TriSearchUnitTests
/// \brief tester for the TriSearch class
////////////////////////////////////////////////////////////////////////////////
//------------------------------------------------------------------------------
/// \brief tests creating the class
//------------------------------------------------------------------------------
void TriSearchUnitTests::testCreateClass()
{
  boost::shared_ptr<GmTriSearch> pts = GmTriSearch::New();
  TS_ASSERT(pts);
}
//------------------------------------------------------------------------------
/// \brief tests calculation of interpolation weights from triangles
//------------------------------------------------------------------------------
void TriSearchUnitTests::testInterpWeights()
{
  boost::shared_ptr<GmTriSearch> tris = GmTriSearch::New();
  BSHP<std::vector<Pt3d>> p(new std::vector<Pt3d>());
  *p = {{0, 0, 0}, {1, 1, 1}, {1, 0, 2}, {0, 1, 2}, {.5, 1.5, 1}};
  BSHP<std::vector<int>> t(new std::vector<int>());
  *t = {0, 2, 1, 0, 1, 3, 3, 1, 4};
  tris->TrisToSearch(p, t);

  std::vector<int> idxs, baseIdx;
  std::vector<double> wt, baseWt;
  Pt3d pt(.5, .2, 0);
  TS_ASSERT(tris->InterpWeights(pt, idxs, wt));
  baseIdx = {0, 2, 1};
  baseWt = {.5, .3, .2};
  TS_ASSERT_EQUALS_VEC(baseIdx, idxs);
  TS_ASSERT_DELTA_VEC(baseWt, wt, 1e-7);

  int triangleIdx;
  TS_ASSERT(tris->InterpWeightsTriangleIdx({.25, .75, 0}, triangleIdx, idxs, wt));
  TS_ASSERT_EQUALS(3, triangleIdx);
  baseIdx = {0, 1, 3};
  TS_ASSERT_EQUALS_VEC(baseIdx, idxs);
  baseWt = {.25, .25, .5};
  TS_ASSERT_DELTA_VEC(baseWt, wt, 1e-7);

  TS_ASSERT(!tris->InterpWeightsTriangleIdx({0, 1.25, 0}, triangleIdx, idxs, wt));
  TS_ASSERT_EQUALS(XM_NONE, triangleIdx);
  TS_ASSERT_EQUALS_VEC(VecInt(), idxs);
  TS_ASSERT_DELTA_VEC(VecDbl(), wt, 1e-7);
} // TriSearchUnitTests::testInterpWeights
//------------------------------------------------------------------------------
/// \brief tests setting the point activity
//------------------------------------------------------------------------------
void TriSearchUnitTests::testPtActivity()
{
  BSHP<GmTriSearch> tris = GmTriSearch::New();
  BSHP<std::vector<Pt3d>> p(new std::vector<Pt3d>());
  *p = {{0, 0, 0}, {1, 1, 1}, {1, 0, 2}, {0, 1, 2}, {.5, 1.5, 1}};
  BSHP<std::vector<int>> t(new std::vector<int>());
  *t = {0, 2, 1, 0, 1, 3, 3, 1, 4};
  tris->TrisToSearch(p, t);

  DynBitset act_wrongSize(6);
  // won't do anything
  tris->SetPtActivity(act_wrongSize);

  Pt3d pt(.5, .2, 0);
  TS_ASSERT_EQUALS(0, tris->TriContainingPt(pt));

  DynBitset act(5);
  act.flip();
  act.set(1, false);
  tris->SetPtActivity(act);

  TS_ASSERT_EQUALS(XM_NONE, tris->TriContainingPt(pt));
} // TriSearchUnitTests::testPtActivity
//------------------------------------------------------------------------------
/// \brief tests setting the triangle activity
//------------------------------------------------------------------------------
void TriSearchUnitTests::testTriActivity()
{
  BSHP<GmTriSearch> tris = GmTriSearch::New();
  BSHP<std::vector<Pt3d>> p(new std::vector<Pt3d>());
  *p = {{0, 0, 0}, {1, 1, 1}, {1, 0, 2}, {0, 1, 2}, {.5, 1.5, 1}};
  BSHP<std::vector<int>> t(new std::vector<int>());
  *t = {0, 2, 1, 0, 1, 3, 3, 1, 4};
  tris->TrisToSearch(p, t);

  DynBitset act_wrongSize(4);
  // won't do anything
  tris->SetTriActivity(act_wrongSize);

  Pt3d pt(.5, .2, 0), pt2(.5, .5, 0);
  TS_ASSERT_EQUALS(0, tris->TriContainingPt(pt));
  // pt2 is right on the boundary between first two triangles
  TS_ASSERT_EQUALS(0, tris->TriContainingPt(pt2));

  DynBitset act(3);
  act.flip();
  act.set(0, false);
  tris->SetTriActivity(act);

  TS_ASSERT_EQUALS(XM_NONE, tris->TriContainingPt(pt));
  TS_ASSERT_EQUALS(3, tris->TriContainingPt(pt2));

  // setting to empty activity turns on everything
  act.clear();
  tris->SetTriActivity(act);
  TS_ASSERT_EQUALS(0, tris->TriContainingPt(pt));
  TS_ASSERT_EQUALS(0, tris->TriContainingPt(pt2));
} // TriSearchUnitTests::testPtActivity
//------------------------------------------------------------------------------
/// \brief A particular test case from SMS that was included
//------------------------------------------------------------------------------
void TriSearchUnitTests::testSmsCase1()
{
  Pt3d pt(-31.459823375717541, 29.927133417260336, 0);
  GmTriSearchImpl tris;
  BSHP<std::vector<Pt3d>> p(new std::vector<Pt3d>());
  *p = {{-20.150000000000002, 46.579999999999998, 7},
        {-41.100000000000001, 30.370000000000001, 8},
        {-19.550000000000001, 29.379999999999999, 9}};
  BSHP<std::vector<int>> t(new std::vector<int>());
  *t = {2, 0, 1};
  tris.TrisToSearch(p, t);
  TS_ASSERT_EQUALS(0, tris.TriContainingPt(pt));
} // TriSearchUnitTests::testCase1
//------------------------------------------------------------------------------
/// \brief tests a point that touches a triangle
//------------------------------------------------------------------------------
void TriSearchUnitTests::testTouch()
{
  Pt3d pt(.5, .5, 0);
  GmTriSearchImpl tris;
  BSHP<std::vector<Pt3d>> p(new std::vector<Pt3d>());
  *p = {{0, 0, 7}, {1, 0, 8}, {1, 1, 9}};
  BSHP<std::vector<int>> t(new std::vector<int>());
  *t = {2, 0, 1};
  tris.TrisToSearch(p, t);
  TS_ASSERT_EQUALS(0, tris.TriContainingPt(pt));
} // TriSearchUnitTests::testTouch
//------------------------------------------------------------------------------
/// \brief tests overlaping triangles
//------------------------------------------------------------------------------
void TriSearchUnitTests::testTriEnvelopesOverlap()
{
  BSHP<TrTin> t = trBuildTestTin();
  GmTriSearchImpl ts;
  BSHP<VecInt> tPtr(new VecInt());
  tPtr->reserve(t->Triangles().size());
  for (const auto& tt : t->Triangles())
    tPtr->push_back((int)tt);
  ts.TrisToSearch(t->PointsPtr(), tPtr);

  Pt3d pMin(9, 4, 0), pMax(15, 10, 0);
  VecInt tIdxes;
  ts.TriEnvelopesOverlap(pMin, pMax, tIdxes);
  std::sort(tIdxes.begin(), tIdxes.end());
  VecInt baseIdxes = {3, 6, 9, 18, 21, 27};
  TS_ASSERT_EQUALS_VEC(baseIdxes, tIdxes);
} // TriSearchUnitTests::testTriEnvelopesOverlap

#endif // CXX_TEST


/*
 * This program is free software; you can redistribute it and/or
 * modify it under the terms of the GNU General Public License
 * as published by the Free Software Foundation; either version 2
 * of the License, or (at your option) any later version.
 *
 * This program is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU General Public License for more details.
 *
 * You should have received a copy of the GNU General Public License
 * along with this program; if not, write to the Free Software Foundation,
 * Inc., 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301, USA.
 *
 * The Original Code is Copyright (C) 2015 Blender Foundation.
 * All rights reserved.
 */

#ifndef __EIGEN3_SVD_C_API_CC__
#define __EIGEN3_SVD_C_API_CC__

/* Eigen gives annoying huge amount of warnings here, silence them! */
#if defined(__GNUC__) && !defined(__clang__)
#  pragma GCC diagnostic ignored "-Wlogical-op"
#endif

#ifdef __EIGEN3_SVD_C_API_CC__ /* quiet warning */
#endif

#include <Eigen/Core>
#include <Eigen/Dense>
#include <Eigen/SVD>

#include "svd.h"

using Eigen::JacobiSVD;

using Eigen::NoQRPreconditioner;

using Eigen::ComputeThinU;
using Eigen::ComputeThinV;

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

using Eigen::Matrix4f;

void EIG_svd_square_matrix(const int size, const float *matrix, float *r_U, float *r_S, float *r_V)
{
  /* Since our matrix is squared, we can use thinU/V. */
  unsigned int flags = (r_U ? ComputeThinU : 0) | (r_V ? ComputeThinV : 0);

  /* Blender and Eigen matrices are both column-major. */
  JacobiSVD<MatrixXf, NoQRPreconditioner> svd(Map<MatrixXf>((float *)matrix, size, size), flags);

  if (r_U) {
    Map<MatrixXf>(r_U, size, size) = svd.matrixU();
  }

  if (r_S) {
    Map<VectorXf>(r_S, size) = svd.singularValues();
  }

  if (r_V) {
    Map<MatrixXf>(r_V, size, size) = svd.matrixV();
  }
}

#endif /* __EIGEN3_SVD_C_API_CC__ */


#include <ros/ros.h>
#include <sensor_msgs/PointCloud2.h>
#include <sensor_msgs/point_cloud_conversion.h>
#include <pcl_conversions/pcl_conversions.h>
#include <pcl/point_cloud.h>
#include <pcl/point_types.h>
#include <pcl/filters/passthrough.h>
#include <pcl/common/centroid.h>
#include <pcl_ros/transforms.h>
#include <pcl/filters/extract_indices.h>
#include <pcl/common/common.h>
#include <pcl/segmentation/extract_polygonal_prism_data.h>
#include <boost/make_shared.hpp>
#include <tf_conversions/tf_eigen.h>
#include <tf/transform_listener.h>
#include <tf/transform_broadcaster.h>

// parameters names
const std::string CLOUD_DESCRIPTIONS="point_cloud_descriptions";
const std::string FRAME_ID = "frame_id";

// topics
const std::string POINT_CLOUD_TOPIC = "generated_cloud";

// constants

class GeneratePointCloud
{
	typedef pcl::PointCloud<pcl::PointXYZ> Cloud;
	typedef pcl::PointCloud<pcl::PointXYZRGBA> CloudRGBA;

	enum ShapeType
	{
		BOX = 1,
		CYLINDER = 2,
		SPHERE = 3,
		CONE = 4
	};

	struct Description
	{
		tf::Vector3 size;
		int shape_type;
		tf::Transform transform;
		double resolution;
	};

	public:
		GeneratePointCloud()
		{
		}

		bool init()
		{
			XmlRpc::XmlRpcValue list;
			ros::NodeHandle ph("~");

			// loading frame id
			if(!ph.getParam(FRAME_ID,frame_id_))
			{
				ROS_ERROR_STREAM("failed to load frame id");
			}

			// parameter numeric_fiels
			int shape;
			std::map<std::string,double> numeric_fields;
			numeric_fields.insert(std::make_pair("x",0));
			numeric_fields.insert(std::make_pair("y",0));
			numeric_fields.insert(std::make_pair("z",0));
			numeric_fields.insert(std::make_pair("rx",0));
			numeric_fields.insert(std::make_pair("ry",0));
			numeric_fields.insert(std::make_pair("rz",0));
			numeric_fields.insert(std::make_pair("l",0));
			numeric_fields.insert(std::make_pair("w",0));
			numeric_fields.insert(std::make_pair("h",0));
			numeric_fields.insert(std::make_pair("resolution",0));

			// loading point cloud descriptions
			if(ph.getParam(CLOUD_DESCRIPTIONS,list))
			{
				if(list.getType() == XmlRpc::XmlRpcValue::TypeArray)
				{
					for(int j = 0; j < list.size(); j++)
					{

						XmlRpc::XmlRpcValue &entry = list[j];

						// parse numeric fields
						std::map<std::string,double>::iterator i;
						for(i = numeric_fields.begin() ; i != numeric_fields.end(); i++)
						{
							if(entry.hasMember(i->first))
							{
								double val = static_cast<double>(entry[i->first]);
								numeric_fields[i->first] = val;
							}
							else
							{
								ROS_ERROR_STREAM("Point Cloud description entry is missing field '"<<i->first<<"'");
								return false;
							}
						}

						// shape
						if(entry.hasMember("shape"))
						{
							shape = static_cast<int>(entry["shape"]);
						}
						else
						{
							return false;
						}

						// populating structure
						Description d;
						d.size = tf::Vector3(numeric_fields["l"],numeric_fields["w"],numeric_fields["h"]);
						d.transform = tf::Transform(
								tf::Quaternion(numeric_fields["ry"],numeric_fields["rx"],numeric_fields["rz"]),
								tf::Vector3(numeric_fields["x"],numeric_fields["y"],numeric_fields["z"]));
						d.resolution = numeric_fields["resolution"];
						d.shape_type = shape;

						cloud_descriptions_.push_back(d);
					}

				}
			}

			return true;
		}

		void run()
		{
			ros::NodeHandle nh;
			ros::Publisher cloud_publisher = nh.advertise<sensor_msgs::PointCloud2>(POINT_CLOUD_TOPIC,1);

			if(init())
			{
				generate_cloud();

				sensor_msgs::PointCloud2 msg;
				pcl::toROSMsg(full_color_cloud_,msg);

				ros::Duration loop_duration(0.4f);
				while(ros::ok())
				{
					msg.header.stamp = ros::Time::now() - loop_duration;
					cloud_publisher.publish(msg);
					loop_duration.sleep();

				}
			}

		}

	protected:

		void generate_cloud()
		{
			full_cloud_.clear();
			for(unsigned int i = 0; i < cloud_descriptions_.size() ; i++)
			{
				Cloud points;
				Description &desc = cloud_descriptions_[i];

				switch(desc.shape_type)
				{
				case ShapeType::BOX:

					create_box(desc,points);
					break;

				case ShapeType::CYLINDER:
					create_cylinder(desc,points);
					break;

				default:
					ROS_WARN_STREAM("shape type "<<desc.shape_type<<" not implemented.  Defaulting to box");
					create_box(desc,points);
					break;
				}

				// transforming points
				Eigen::Affine3d eigen3d;
				tf::transformTFToEigen(desc.transform,eigen3d);
				pcl::transformPointCloud(points,points,Eigen::Affine3f(eigen3d));

				// concatenating points
				full_cloud_ +=points;

				// copying to color cloud
				pcl::copyPointCloud(full_cloud_,full_color_cloud_);
			}

			full_cloud_.header.frame_id = frame_id_;
			full_color_cloud_.header.frame_id = frame_id_;
		}

		void create_rectangular_patch(tf::Vector3 start,tf::Vector3 end,double res,Cloud& patch)
		{
			int count_x = (end.x() - start.x())/res + 1;
			int count_y = (end.y() - start.y())/res +1;

			pcl::PointXYZ p;
			p.z = 0;
			patch.resize(count_x * count_y);
			int counter = 0;
			for(int i = 0; i < count_x ; i++)
			{
				p.x = start.x() + res*i;
				for(int j = 0 ; j < count_y;j++)
				{
					p.y = start.y() + res*j;
					patch.points[counter] = p;
					counter++;
				}
			}
		}

		void create_box(const Description& desc,Cloud& box_points)
		{
			// face points
			Cloud top,bottom,front,rear,left,right,temp;
			tf::Vector3 start,end;

			// transforms
			tf::Transform t;
			Eigen::Affine3d eigen3d;

			// ================================ create top and bottom patches ===============================
			start = tf::Vector3(-desc.size.x()/2,-desc.size.y()/2,0);
			end = tf::Vector3(desc.size.x()/2,desc.size.y()/2,0);
			create_rectangular_patch(start,end,desc.resolution,temp);

			// transform cloud to top
			t = tf::Transform(tf::Quaternion::getIdentity(),tf::Vector3(0,0,0.5f*desc.size.z()));
			tf::transformTFToEigen(t,eigen3d);
			pcl::transformPointCloud(temp,top,Eigen::Affine3f(eigen3d));

			// transform cloud to bottom
			t = tf::Transform(tf::Quaternion::getIdentity(),tf::Vector3(0,0,0.5f*-desc.size.z()));
			tf::transformTFToEigen(t,eigen3d);
			pcl::transformPointCloud(temp,bottom,Eigen::Affine3f(eigen3d));

			// concatenate
			box_points += top;
			box_points += bottom;

			// ================================ create front and rear patches ===============================
			start = tf::Vector3(-desc.size.x()/2,-desc.size.z()/2,0);
			end = tf::Vector3(desc.size.x()/2,desc.size.z()/2,0);
			create_rectangular_patch(start,end,desc.resolution,temp);

			// transform cloud to front
			t = tf::Transform(tf::Quaternion(tf::Vector3(1,0,0),M_PI_2)
				,tf::Vector3(0,0.5f*desc.size.y(),0));
			tf::transformTFToEigen(t,eigen3d);
			pcl::transformPointCloud(temp,front,Eigen::Affine3f(eigen3d));

			// transform cloud to rear
			t = tf::Transform(tf::Quaternion(tf::Vector3(1,0,0),M_PI_2)
							,tf::Vector3(0,0.5f*-desc.size.y(),0));
			tf::transformTFToEigen(t,eigen3d);
			pcl::transformPointCloud(temp,rear,Eigen::Affine3f(eigen3d));

			box_points += front;
			box_points += rear;

			// ================================ create left and right patches ===============================
			start = tf::Vector3(-desc.size.z()/2,-desc.size.y()/2,0);
			end = tf::Vector3(desc.size.z()/2,desc.size.y()/2,0);
			create_rectangular_patch(start,end,desc.resolution,temp);

			// transform cloud to left
			t = tf::Transform(tf::Quaternion(tf::Vector3(0,1,0),M_PI_2)
				,tf::Vector3(0.5f*desc.size.x(),0,0));
			tf::transformTFToEigen(t,eigen3d);
			pcl::transformPointCloud(temp,left,Eigen::Affine3f(eigen3d));

			// transform cloud to right
			t = tf::Transform(tf::Quaternion(tf::Vector3(0,1,0),M_PI_2)
				,tf::Vector3(0.5f*-desc.size.x(),0,0));
			tf::transformTFToEigen(t,eigen3d);
			pcl::transformPointCloud(temp,right,Eigen::Affine3f(eigen3d));

			box_points += left;
			box_points += right;

		}

		void create_cylinder(const Description& desc,Cloud& points)
		{
			Cloud top,bottom,middle,temp,ring;
			double radius = desc.size.x(); // along x axis
			double lenght = desc.size.z(); // along z axis

			// transforms
			tf::Transform t;
			Eigen::Affine3d eigen3d;

			// ================================ create top and bottom patches ===============================
			create_circle(radius,desc.resolution,temp);

			// transform cloud to top
			t = tf::Transform(tf::Quaternion::getIdentity(),tf::Vector3(0,0,0.5f*lenght));
			tf::transformTFToEigen(t,eigen3d);
			pcl::transformPointCloud(temp,top,Eigen::Affine3f(eigen3d));

			// transform cloud to bottom
			t = tf::Transform(tf::Quaternion::getIdentity(),tf::Vector3(0,0,0.5f*-lenght));
			tf::transformTFToEigen(t,eigen3d);
			pcl::transformPointCloud(temp,bottom,Eigen::Affine3f(eigen3d));

			// concatenate
			points += top;
			points += bottom;

			// ================================ create middle body ===============================

			double dtheta = desc.resolution/radius;
			int num_points_per_ring = static_cast<int>(2*M_PI/dtheta);
			ring.resize(num_points_per_ring);
			pcl::PointXYZ p;
			for(int i = 0; i < num_points_per_ring;i++)
			{
				p.x = radius*std::cos(dtheta*i);
				p.y = radius*std::sin(dtheta*i);
				p.z = 0;
				ring[i] = p;
			}

			// creting all rings
			tf::Vector3 start(0,0,-lenght*0.5f);
			int num_rings = static_cast<int>(lenght/desc.resolution);
			for(int i = 1; i < num_rings;i++)
			{
				temp.clear();
				t = tf::Transform(tf::Quaternion::getIdentity(),start + tf::Vector3(0,0,i*desc.resolution));
				tf::transformTFToEigen(t,eigen3d);
				pcl::transformPointCloud(ring,temp,Eigen::Affine3f(eigen3d));
				middle+= temp;
			}

			points +=middle;

		}

		void create_circle(double radius,double res, Cloud& circle)
		{
			// creating rectagular patch
			Cloud rect;
			tf::Vector3 start = tf::Vector3(-radius,-radius,0);
			tf::Vector3 end = tf::Vector3(radius,radius,0);
			create_rectangular_patch(start,end,res,rect);

			// creating circular polygon
			int num_points = 20;
			double dtheta = 2*M_PI/num_points;
			Cloud circle_boundary;
			circle_boundary.resize(num_points+1);
			for(int i = 0; i < num_points;i++)
			{
				pcl::PointXYZ p;
				p.x = radius*std::cos(dtheta*i);
				p.y = radius*std::sin(dtheta*i);
				p.z = 0;
				circle_boundary.points[i] = p;
			}
			circle_boundary.points[num_points] = circle_boundary.points[0]; // closing boundary

			// removing points outside boundary
			pcl::ExtractIndices<pcl::PointXYZ> extract;
			pcl::PointIndices::Ptr inliers(new pcl::PointIndices());
			pcl::ExtractPolygonalPrismData<pcl::PointXYZ> prism;
			prism.setInputCloud(rect.makeShared());
			prism.setInputPlanarHull(circle_boundary.makeShared());
			prism.setHeightLimits(-1,1);
			prism.setViewPoint(0,0,1);
			prism.segment(*inliers);

			extract.setInputCloud(rect.makeShared());
			extract.setIndices(inliers);
			extract.setNegative(false);
			extract.filter(circle);

		}

	protected:

		std::vector<Description> cloud_descriptions_;
		std::string frame_id_;
		Cloud full_cloud_;
		CloudRGBA full_color_cloud_;
		float resolution_;

};

int main(int argc,char** argv)
{
	ros::init(argc,argv,"generate_point_cloud_node");

	GeneratePointCloud g;
	g.run();
	return 0;
}


#include <limits>
#include <fstream>
#include <vector>
#include <Eigen/Core>
#include <pcl/point_types.h>
#include <pcl/point_cloud.h>
#include <pcl/io/pcd_io.h>
#include <pcl/kdtree/kdtree_flann.h>
#include <pcl/filters/passthrough.h>
#include <pcl/filters/voxel_grid.h>
#include <pcl/features/normal_3d.h>
#include <pcl/features/fpfh.h>
#include <pcl/registration/ia_ransac.h>

class FeatureCloud
{
  public:
    // A bit of shorthand
    typedef pcl::PointCloud<pcl::PointXYZ> PointCloud;
    typedef pcl::PointCloud<pcl::Normal> SurfaceNormals;
    typedef pcl::PointCloud<pcl::FPFHSignature33> LocalFeatures;
    typedef pcl::search::KdTree<pcl::PointXYZ> SearchMethod;

    FeatureCloud () :
      search_method_xyz_ (new SearchMethod),
      normal_radius_ (0.02f),
      feature_radius_ (0.02f)
    {}

    ~FeatureCloud () {}

    // Process the given cloud
    void
    setInputCloud (PointCloud::Ptr xyz)
    {
      xyz_ = xyz;
      processInput ();
    }

    // Load and process the cloud in the given PCD file
    void
    loadInputCloud (const std::string &pcd_file)
    {
      xyz_ = PointCloud::Ptr (new PointCloud);
      pcl::io::loadPCDFile (pcd_file, *xyz_);
      processInput ();
    }

    // Get a pointer to the cloud 3D points
    PointCloud::Ptr
    getPointCloud () const
    {
      return (xyz_);
    }

    // Get a pointer to the cloud of 3D surface normals
    SurfaceNormals::Ptr
    getSurfaceNormals () const
    {
      return (normals_);
    }

    // Get a pointer to the cloud of feature descriptors
    LocalFeatures::Ptr
    getLocalFeatures () const
    {
      return (features_);
    }

  protected:
    // Compute the surface normals and local features
    void
    processInput ()
    {
      computeSurfaceNormals ();
      computeLocalFeatures ();
    }

    // Compute the surface normals
    void
    computeSurfaceNormals ()
    {
      normals_ = SurfaceNormals::Ptr (new SurfaceNormals);

      pcl::NormalEstimation<pcl::PointXYZ, pcl::Normal> norm_est;
      norm_est.setInputCloud (xyz_);
      norm_est.setSearchMethod (search_method_xyz_);
      norm_est.setRadiusSearch (normal_radius_);
      norm_est.compute (*normals_);
    }

    // Compute the local feature descriptors
    void
    computeLocalFeatures ()
    {
      features_ = LocalFeatures::Ptr (new LocalFeatures);

      pcl::FPFHEstimation<pcl::PointXYZ, pcl::Normal, pcl::FPFHSignature33> fpfh_est;
      fpfh_est.setInputCloud (xyz_);
      fpfh_est.setInputNormals (normals_);
      fpfh_est.setSearchMethod (search_method_xyz_);
      fpfh_est.setRadiusSearch (feature_radius_);
      fpfh_est.compute (*features_);
    }

  private:
    // Point cloud data
    PointCloud::Ptr xyz_;
    SurfaceNormals::Ptr normals_;
    LocalFeatures::Ptr features_;
    SearchMethod::Ptr search_method_xyz_;

    // Parameters
    float normal_radius_;
    float feature_radius_;
};

class TemplateAlignment
{
  public:

    // A struct for storing alignment results
    struct Result
    {
      float fitness_score;
      Eigen::Matrix4f final_transformation;
      EIGEN_MAKE_ALIGNED_OPERATOR_NEW
    };

    TemplateAlignment () :
      min_sample_distance_ (0.05f),
      max_correspondence_distance_ (0.01f*0.01f),
      nr_iterations_ (500)
    {
      // Intialize the parameters in the Sample Consensus Intial Alignment (SAC-IA) algorithm
      sac_ia_.setMinSampleDistance (min_sample_distance_);
      sac_ia_.setMaxCorrespondenceDistance (max_correspondence_distance_);
      sac_ia_.setMaximumIterations (nr_iterations_);
    }

    ~TemplateAlignment () {}

    // Set the given cloud as the target to which the templates will be aligned
    void
    setTargetCloud (FeatureCloud &target_cloud)
    {
      target_ = target_cloud;
      sac_ia_.setInputTarget (target_cloud.getPointCloud ());
      sac_ia_.setTargetFeatures (target_cloud.getLocalFeatures ());
    }

    // Add the given cloud to the list of template clouds
    void
    addTemplateCloud (FeatureCloud &template_cloud)
    {
      templates_.push_back (template_cloud);
    }

    // Align the given template cloud to the target specified by setTargetCloud ()
    void
    align (FeatureCloud &template_cloud, TemplateAlignment::Result &result)
    {
      sac_ia_.setInputCloud (template_cloud.getPointCloud ());
      sac_ia_.setSourceFeatures (template_cloud.getLocalFeatures ());

      pcl::PointCloud<pcl::PointXYZ> registration_output;
      sac_ia_.align (registration_output);

      result.fitness_score = (float) sac_ia_.getFitnessScore (max_correspondence_distance_);
      result.final_transformation = sac_ia_.getFinalTransformation ();
    }

    // Align all of template clouds set by addTemplateCloud to the target specified by setTargetCloud ()
    void
    alignAll (std::vector<TemplateAlignment::Result, Eigen::aligned_allocator<Result> > &results)
    {
      results.resize (templates_.size ());
      for (size_t i = 0; i < templates_.size (); ++i)
      {
        align (templates_[i], results[i]);
      }
    }

    // Align all of template clouds to the target cloud to find the one with best alignment score
    int
    findBestAlignment (TemplateAlignment::Result &result)
    {
      // Align all of the templates to the target cloud
      std::vector<Result, Eigen::aligned_allocator<Result> > results;
      alignAll (results);

      // Find the template with the best (lowest) fitness score
      float lowest_score = std::numeric_limits<float>::infinity ();
      int best_template = 0;
      for (size_t i = 0; i < results.size (); ++i)
      {
        const Result &r = results[i];
        if (r.fitness_score < lowest_score)
        {
          lowest_score = r.fitness_score;
          best_template = (int) i;
        }
      }

      // Output the best alignment
      result = results[best_template];
      return (best_template);
    }

  private:
    // A list of template clouds and the target to which they will be aligned
    std::vector<FeatureCloud> templates_;
    FeatureCloud target_;

    // The Sample Consensus Initial Alignment (SAC-IA) registration routine and its parameters
    pcl::SampleConsensusInitialAlignment<pcl::PointXYZ, pcl::PointXYZ, pcl::FPFHSignature33> sac_ia_;
    float min_sample_distance_;
    float max_correspondence_distance_;
    int nr_iterations_;
};

// Align a collection of object templates to a sample point cloud
int
main (int argc, char **argv)
{
  if (argc < 3)
  {
    printf ("No target PCD file given!\n");
    return (-1);
  }

  // Load the object templates specified in the object_templates.txt file
  std::vector<FeatureCloud> object_templates;
  std::ifstream input_stream (argv[1]);
  object_templates.resize (0);
  std::string pcd_filename;
  while (input_stream.good ())
  {
    std::getline (input_stream, pcd_filename);
    if (pcd_filename.empty () || pcd_filename.at (0) == '#') // Skip blank lines or comments
      continue;

    FeatureCloud template_cloud;
    template_cloud.loadInputCloud (pcd_filename);
    object_templates.push_back (template_cloud);
  }
  input_stream.close ();

  // Load the target cloud PCD file
  pcl::PointCloud<pcl::PointXYZ>::Ptr cloud (new pcl::PointCloud<pcl::PointXYZ>);
  pcl::io::loadPCDFile (argv[2], *cloud);

  // Preprocess the cloud by...
  // ...removing distant points
  const float depth_limit = 1.0;
  pcl::PassThrough<pcl::PointXYZ> pass;
  pass.setInputCloud (cloud);
  pass.setFilterFieldName ("z");
  pass.setFilterLimits (0, depth_limit);
  pass.filter (*cloud);

  // ... and downsampling the point cloud
  const float voxel_grid_size = 0.005f;
  pcl::VoxelGrid<pcl::PointXYZ> vox_grid;
  vox_grid.setInputCloud (cloud);
  vox_grid.setLeafSize (voxel_grid_size, voxel_grid_size, voxel_grid_size);
  //vox_grid.filter (*cloud); // Please see this http://www.pcl-developers.org/Possible-problem-in-new-VoxelGrid-implementation-from-PCL-1-5-0-td5490361.html
  pcl::PointCloud<pcl::PointXYZ>::Ptr tempCloud (new pcl::PointCloud<pcl::PointXYZ>);
  vox_grid.filter (*tempCloud);
  cloud = tempCloud;

  // Assign to the target FeatureCloud
  FeatureCloud target_cloud;
  target_cloud.setInputCloud (cloud);

  // Set the TemplateAlignment inputs
  TemplateAlignment template_align;
  for (size_t i = 0; i < object_templates.size (); ++i)
  {
    template_align.addTemplateCloud (object_templates[i]);
  }
  template_align.setTargetCloud (target_cloud);

  // Find the best template alignment
  TemplateAlignment::Result best_alignment;
  int best_index = template_align.findBestAlignment (best_alignment);
  const FeatureCloud &best_template = object_templates[best_index];

  // Print the alignment fitness score (values less than 0.00002 are good)
  printf ("Best fitness score: %f\n", best_alignment.fitness_score);

  // Print the rotation matrix and translation vector
  Eigen::Matrix3f rotation = best_alignment.final_transformation.block<3,3>(0, 0);
  Eigen::Vector3f translation = best_alignment.final_transformation.block<3,1>(0, 3);

  printf ("\n");
  printf ("    | %6.3f %6.3f %6.3f | \n", rotation (0,0), rotation (0,1), rotation (0,2));
  printf ("R = | %6.3f %6.3f %6.3f | \n", rotation (1,0), rotation (1,1), rotation (1,2));
  printf ("    | %6.3f %6.3f %6.3f | \n", rotation (2,0), rotation (2,1), rotation (2,2));
  printf ("\n");
  printf ("t = < %0.3f, %0.3f, %0.3f >\n", translation (0), translation (1), translation (2));

  // Save the aligned template for visualization
  pcl::PointCloud<pcl::PointXYZ> transformed_cloud;
  pcl::transformPointCloud (*best_template.getPointCloud (), transformed_cloud, best_alignment.final_transformation);
  pcl::io::savePCDFileBinary ("output.pcd", transformed_cloud);

  return (0);
}

Aligning object templates to a point cloud/template_alignment.cpp


/*************************************************************************
 *
 *              Author: b51
 *                Mail: b51live@gmail.com
 *            FileName: SuperPointFeatureDetector.cc
 *
 *          Created On: Fri 27 Sep 2019 08:30:41 AM UTC
 *     Licensed under The MIT License [see LICENSE for details]
 *
 ************************************************************************/

#include "SuperPointFeatureDetector.h"

#include <gflags/gflags.h>
#include <glog/logging.h>
#include <Eigen/Core>
#include <cmath>

DEFINE_string(model, "/home/ubuntu/Models",
              "Path to pytorch model");
DEFINE_double(conf_thresh, 0.015,
              "Detector confidence threshold (default: 0.015)");
DEFINE_int32(H, 120, "Net input image height");
DEFINE_int32(W, 160, "Net input image width");
DEFINE_int32(nms_dist, 4, "Non Maximum Suppression (NMS) distance (default: 4).");
DEFINE_int32(CUDA, 0, "Default GPU device");

//Size of each output cell, Keep this fixed
const int CELL = 8;

SuperPointFeatureDetector::SuperPointFeatureDetector()
    : width_(0),
      height_(0),
      max_number_of_features_(0),
      conf_thresh_(0.),
      nms_dist_(0),
      is_initialized_(false),
      module_(nullptr) {}

void SuperPointFeatureDetector::Init(int width, int height,
                                     int max_number_of_features) {
  CHECK_GT(width, 0);
  CHECK_GT(height, 0);
  CHECK_GT(max_number_of_features, 0);
  Init(width, height, max_number_of_features, FLAGS_model,
       FLAGS_conf_thresh, FLAGS_nms_dist);
}

void SuperPointFeatureDetector::Init(int width, int height,
                                     int max_number_of_features,
                                     std::string model_path, float conf_thresh,
                                     int nms_dist) {
  CHECK_GT(width, 0);
  CHECK_GT(height, 0);
  CHECK_GT(max_number_of_features, 0);
  width_ = width;
  height_ = height;
  max_number_of_features_ = max_number_of_features;
  model_path_ = model_path;
  conf_thresh_ = conf_thresh;
  nms_dist_ = nms_dist;

  image_.create(height_, width_, CV_8UC1);

  try {
    // TODO(b51): Make shared_ptr normal useable
    module_ = torch::jit::load(model_path.c_str());
    // Load module on tx2
    // module_ = std::make_shared<torch::jit::script::Module>(
    //    torch::jit::load(model_path.c_str()));
    module_->to(torch::Device(torch::kCUDA, FLAGS_CUDA));
  } catch (const c10::Error& e) {
    LOG(FATAL) << "Error loading the model from " << model_path;
  }
  VLOG(3) << "SuperPoint detector parameters:";
  VLOG(3) << "      max number of features: " << max_number_of_features_;
  VLOG(3) << "      model loaded: " << model_path;
  is_initialized_ = true;
}

void SuperPointFeatureDetector::NMSFast(
    const Eigen::Matrix3Xf& pts, int width, int height, int dist_thresh,
    Eigen::Matrix3Xf& filtered_pts) {
  Eigen::MatrixXf sorted_descend_pts;
  sorted_descend_pts.resize(pts.rows(), pts.cols());

  // Sort with heatmap value in descending order
  std::vector<Eigen::Vector3f> vec;
  for (int i = 0; i < pts.cols(); i++)
    vec.push_back(pts.col(i));
  std::sort(vec.begin(), vec.end(),
            [](Eigen::Vector3f const& v1, Eigen::Vector3f const& v2) {
              return v1[2] > v2[2];
            });
  for (int i = 0; i < pts.cols(); i++)
    sorted_descend_pts.col(i) = vec[i];

  // Pad the border of the grid, so that we can NMS points near the border.
  int pad = dist_thresh;
  int paded_width = width + 2 * pad;
  int paded_height = height + 2 * pad;
  Eigen::MatrixXf grid = Eigen::MatrixXf::Zero(paded_width, paded_height);
#pragma omp parallel for
  for (int i = 0; i < sorted_descend_pts.cols(); i++) {
    grid(std::round(sorted_descend_pts.col(i)[0] + pad),
         std::round(sorted_descend_pts.col(i)[1] + pad)) = 1;
  }

  // Iterate through points, highest to lowest conf, suppress neighborhood.
  int count = 0;
  int pad_2 = 2 * pad;
  filtered_pts = Eigen::Matrix3Xf::Zero(3, sorted_descend_pts.cols());
  for (int i = 0; i < sorted_descend_pts.cols(); i++) {
    // Account for top and left padding.
    int x = std::round(sorted_descend_pts.col(i)[0] + pad);
    int y = std::round(sorted_descend_pts.col(i)[1] + pad);
    if (grid(x, y) == 1) {
      // Block of size(p, q), starting at (i, j)
      // dynamic-size block expression: matrix.block(i, j, p, q)
      // https://eigen.tuxfamily.org/dox/group__TutorialBlockOperations.html
      grid.block(x - pad, y - pad, pad_2 + 1, pad_2 + 1) =
          Eigen::MatrixXf::Zero(pad_2 + 1, pad_2 + 1);
      grid(x, y) = -1;
      // Ignore points along border
      if ((x - pad) < pad or (x - pad >= width - pad) or
          (y - pad) < pad or (y - pad >= height - pad))
        continue;
      filtered_pts.col(count++) =
          Eigen::Vector3f(x - pad, y - pad, sorted_descend_pts.col(i)[2]);
    }
  }
  // resize the matrix to rows x cols while leaving old values untouched
  filtered_pts.conservativeResize(Eigen::NoChange, count);
}

void SuperPointFeatureDetector::Detect(const cv::Mat& image,
                                       Eigen::Matrix3Xf& filtered_pts,
                                       torch::Tensor& descriptors) {
  int H = image.rows;
  int W = image.cols;
  // convert image to float in range [0.0, 1.0]
  std::vector<torch::jit::IValue> inputs;

  // convert cv mat to tensor and transport to cuda
  at::Tensor image_tensor =
      torch::from_blob(image.data, {1, 1, H, W}, at::kFloat);
  inputs.emplace_back(image_tensor.to(torch::Device(torch::kCUDA, FLAGS_CUDA)));

  // inference
  auto outputs = module_->forward(inputs).toTuple();
  torch::Tensor semi = outputs->elements()[0].toTensor();
  torch::Tensor coarse_desc = outputs->elements()[1].toTensor();

  // Deal with feature points locations
  semi = semi.squeeze();
  torch::Tensor dense = semi.exp();
  dense = dense / (dense.sum(0) + 0.00001);
  torch::Tensor nodust = dense.slice(0, 0, dense.size(0) - 1);
  int Hc = H / CELL;
  int Wc = W / CELL;
  nodust = nodust.transpose(0, 1).transpose(1, 2);
  torch::Tensor heatmap = nodust.reshape({Hc, Wc, CELL, CELL});
  heatmap = heatmap.transpose(1, 2);
  heatmap = heatmap.reshape({Hc * CELL, Wc * CELL});
  torch::Tensor heatmap_cpu = heatmap.to(torch::kCPU, /*non_blocking=*/true);

  // Eigen::Map<Eigen::MatrixXf> eig_heatmap(heatmap_cpu.data<float>(),
  //                                         heatmap.size(0), heatmap.size(1));
  cv::Mat mat_heatmap(heatmap_cpu.size(0), heatmap_cpu.size(1), CV_32F,
                      heatmap_cpu.data<float>());
  /**
   *  Filtered to binary mat first and than get non zero value coordinates
   *  cv::threshold(src, dst, threshold, max_binary_value, threshold_type)
   *  threshold_type: 0, Binary
   *                  1, Binary Inverted
   *                  2, Threshold Truncated
   *                  3, Threshold to Zero
   *                  4, Threshold to Zero Inverted
   */
  cv::Mat bin_mat;
  cv::threshold(mat_heatmap, bin_mat, conf_thresh_, 1, 0);
  bin_mat.convertTo(bin_mat, CV_8UC1);
  cv::Mat coordinates;
  cv::findNonZero(bin_mat, coordinates);

  Eigen::Matrix3Xf pts;
  pts.resize(3, coordinates.total());
#pragma omp parallel for
  for (size_t i = 0; i < coordinates.total(); i++) {
    pts.block<3, 1>(0, i) = Eigen::Vector3f(
        coordinates.at<cv::Point>(i).x, coordinates.at<cv::Point>(i).y,
        mat_heatmap.at<float>(coordinates.at<cv::Point>(i)));
  }
  // nms fast
  NMSFast(pts, W, H, nms_dist_, filtered_pts);

  VLOG(4) << "filtered points size: " << filtered_pts.cols();
  int number_of_features = filtered_pts.cols();

  // Process descriptor
  // Interpolate into descriptor map using 2D point locations.
  cv::Mat mat_measurements(2, number_of_features, CV_32F);
  for (int i = 0; i < number_of_features; i++) {
    // grid_sampler needs grid range between [-1., 1.]
    mat_measurements.at<float>(0, i) =
        (filtered_pts.col(i)[0] / (float(W) / 2.)) - 1.;
    mat_measurements.at<float>(1, i) =
        (filtered_pts.col(i)[1] / (float(H) / 2.)) - 1.;
    /*  For debug display
    Eigen::Vector3f p = filtered_pts.col(i);
    LOG(INFO) << "x: " << p[0] << " y: " << p[1] << " value: " << p[2];
    cv::circle(image, cv::Point(p[0], p[1]), 2, cv::Scalar(0, 255, 255));

    cv::imwrite("key_points.jpg", image);
    exit(0);
    */
  }
  torch::Tensor sample_pts =
      torch::from_blob(mat_measurements.data, {2, number_of_features},
                       at::kFloat)
          .to(torch::Device(torch::kCUDA, FLAGS_CUDA));
  sample_pts = sample_pts.transpose(0, 1);
  sample_pts = sample_pts.view({1, 1, -1, 2});

  /**
   *  grid_sampler(at::Tensor input, at::Tensor grid, interpolation_mode,
   *  padding_mode)
   *  grid: range needs to be [-1., 1]
   *  interpolation_mode:
   *      at::native::detail
   *      enum class GridSamplerInterpolation {Bilinear, Nearest};
   *  padding_mode
   *      at::native::detail
   *      enum class GridSamplerPadding {Zeros, Border, Reflection}
   */
  descriptors = at::grid_sampler(
      coarse_desc, sample_pts.to(torch::Device(torch::kCUDA, FLAGS_CUDA)), 0,
      0);
  descriptors = descriptors.reshape({coarse_desc.size(1), -1});
  descriptors = descriptors / (descriptors.norm(c10::nullopt, 0).view({1, -1}));
  descriptors = descriptors.to(torch::kCPU, true);
}

void SuperPointFeatureDetector::Detect(
    const cv::Mat& image, Eigen::Matrix2Xf* current_measurements,
    FeatureDescriptorf* current_feature_descriptors) {
  image.copyTo(image_);
  cv::Mat input_image;
  double H_scale = 1.;
  double W_scale = 1.;
  // resize to net input

  cv::resize(image, input_image, cv::Size(FLAGS_W, FLAGS_H), 0, 0,
             cv::INTER_AREA);
  H_scale = image.rows / FLAGS_H;
  W_scale = image.cols / FLAGS_W;

  // Superpoint net only accept gray image
  if (input_image.channels() > 1) {
    cv::cvtColor(input_image, input_image, CV_BGR2GRAY);
  }
  // normalize to [0.0, 1.0]
  cv::Mat normed_image;
  input_image.convertTo(normed_image, CV_32F, 1.0 / 255, 0);

  Eigen::Matrix3Xf filtered_pts;
  torch::Tensor descriptors;
  Detect(normed_image, filtered_pts, descriptors);
  // LOG(INFO) << " descriptor 0: " << descriptors.slice(1, 0, 1);

  int number_of_features = filtered_pts.cols();
  int descriptor_size = descriptors.size(0);
  current_measurements->resize(Eigen::NoChange, number_of_features);
  current_measurements->block(0, 0, 2, number_of_features) =
      filtered_pts.block(0, 0, 2, number_of_features);

  // restor to original size
  current_measurements->row(0) = current_measurements->row(0) * W_scale;
  current_measurements->row(1) = current_measurements->row(1) * H_scale;

  CHECK_EQ(number_of_features, descriptors.size(1));

  cv::Mat mat_descriptors(descriptors.size(0), descriptors.size(1), CV_32F,
                          descriptors.data<float>());
  // Make mat_descriptors width = descriptor_size (256 here)
  //                      height = number_features
  mat_descriptors = mat_descriptors.t();

  current_feature_descriptors->Configure(descriptor_size,
                                         number_of_features);
  current_feature_descriptors->Resize(number_of_features);
  if (mat_descriptors.isContinuous()) {
    // if continuous
    memcpy(current_feature_descriptors->descriptor(0), mat_descriptors.data,
           descriptor_size * number_of_features * sizeof(float));
  } else {
    for (int i = 0; i < number_of_features; ++i) {
      memcpy(current_feature_descriptors->descriptor(i),
             &mat_descriptors.at<float>(i, 0), descriptor_size * sizeof(float));
    }
  }
  // LOG(INFO) << current_feature_descriptors->Size();
  // LOG(INFO) << current_feature_descriptors->DescriptorSize();
  // Eigen::Map<Eigen::Matrix<float, 1, 256> > for_display(
  //     current_feature_descriptors->descriptor(0), 1, 256);
  // LOG(INFO) << for_display;
}


// Copyright (c) 2020 Graphcore Ltd. All rights reserved.
#include <algorithm>
#include <cassert>
#include <cmath>
#include <exception>
#include <fstream>
#include <iomanip>
#include <random>

#include <poplibs_support/Algorithm.hpp>
#include <poplibs_support/CTCInferenceDefs.hpp>
#include <poplibs_support/LogArithmetic.hpp>
#include <poplibs_test/CTCInference.hpp>
#include <poplibs_test/CTCLoss.hpp>
#include <poplibs_test/CTCUtil.hpp>
#include <poplibs_test/Embedding.hpp>
#include <poplibs_test/MatrixTransforms.hpp>
#include <poplibs_test/Util.hpp>
#include <poputil/exceptions.hpp>

#include <boost/multi_array.hpp>
#include <boost/program_options.hpp>
#include <boost/test/tools/floating_point_comparison.hpp>

#include <boost/random.hpp>

namespace po = boost::program_options;
namespace br = boost::random;

using namespace poplibs_support;
using namespace poplibs_test::ctc;
using namespace poplibs_test;
using namespace poplibs_test::util;

// Default tolerances used in tests
#define FLOAT_REL_TOL 0.01
#define HALF_REL_TOL 0.1
#define FLOAT_ABS_TOL 1e-6
#define HALF_ABS_TOL 1e-5

// Print a sequence, inserting `-` for the blank class
void print(const std::string &prefix, const std::vector<unsigned> &idx,
           unsigned blankClass, bool verbose = true) {
  if (!verbose) {
    return;
  }
  std::cout << "\n" << prefix << " ";
  for (auto &i : idx) {
    if (i == blankClass) {
      std::cout << "- ";
    } else {
      std::cout << i << " ";
    }
  }
  std::cout << "\n";
}

// Print the matrix `in`, using `paddedSequence` as row labels
template <typename FPType>
void print(const std::string &prefix, const boost::multi_array<FPType, 2> &in,
           const std::vector<unsigned> &paddedSequence, unsigned blankClass,
           bool verbose = true) {
  if (!verbose) {
    return;
  }
  std::cout << "\n" << prefix << "\n          ";
  for (unsigned i = 0; i < in[0].size(); i++) {
    std::cout << "         t" << i;
  }

  for (unsigned i = 0; i < in.size(); i++) {
    std::cout << "\nIndex:" << i << "  ";
    if (paddedSequence[i] == blankClass) {
      std::cout << "- ";
    } else {
      std::cout << paddedSequence[i] << " ";
    }
    for (unsigned j = 0; j < in[i].size(); j++) {
      std::cout << std::setw(10) << std::setprecision(4) << in[i][j] << ",";
    }
  }
  std::cout << "\n";
}

// Print the matrix `in`
template <typename FPType>
void print(const std::string &prefix, const boost::multi_array<FPType, 2> &in,
           unsigned blankClass, bool verbose = true) {
  if (!verbose) {
    return;
  }
  std::cout << "\n" << prefix << "\n        ";
  for (unsigned i = 0; i < in[0].size(); i++) {
    std::cout << "         t" << i;
  }

  for (unsigned i = 0; i < in.size(); i++) {
    if (i == blankClass) {
      std::cout << "\nIndex:-  ";
    } else {
      std::cout << "\nIndex:" << i << "  ";
    }
    for (unsigned j = 0; j < in[i].size(); j++) {
      std::cout << std::setw(10) << std::setprecision(4) << in[i][j] << ",";
    }
  }
  std::cout << "\n";
}
// Print the matrix `in` to export for comparison
template <typename FPType>
void exportPrint(const std::string &prefix,
                 const boost::multi_array<FPType, 2> &in) {
  std::cout << prefix << "\n";
  for (unsigned i = 0; i < in[0].size(); i++) {
    std::cout << "[[";
    for (unsigned j = 0; j < in.size(); j++) {
      std::cout << std::fixed << std::setw(10) << std::setprecision(4)
                << in[j][i] << ",";
    }
    std::cout << "]],\n";
  }
}

template <typename FPType>
boost::multi_array<FPType, 2> toLinear(const boost::multi_array<FPType, 2> &in,
                                       FPType scale) {
  boost::multi_array<FPType, 2> out(boost::extents[in.size()][in[0].size()]);
  for (unsigned i = 0; i < in.size(); i++) {
    for (unsigned j = 0; j < in[i].size(); j++) {
      out[i][j] = -scale * std::exp(in[i][j]);
    }
  }
  return out;
}
template <typename FPType>
boost::multi_array<FPType, 2> scale(const boost::multi_array<FPType, 2> &in,
                                    FPType scale) {
  boost::multi_array<FPType, 2> out(boost::extents[in.size()][in[0].size()]);
  for (unsigned i = 0; i < in.size(); i++) {
    for (unsigned j = 0; j < in[i].size(); j++) {
      out[i][j] = -scale * in[i][j];
    }
  }
  return out;
}
// Struct and function to return the test inputs
template <typename FPType> struct InputSequence {
  boost::multi_array<FPType, 2> input;
  std::vector<unsigned> extendedLabel;
  unsigned blankClass;
};

template <typename FPType>
void testTraining(const InputSequence<FPType> &test, bool normalise,
                  bool verbose) {
  auto blankClass = test.blankClass;
  auto paddedSequence = test.extendedLabel;
  auto validTimesteps = test.input.shape()[1];

  std::cout << "\n";

  // This can be used to copy into python/pytorch implementation for
  // comparison
  print("Log Softmax in", test.input, blankClass);
  boost::multi_array<FPType, 2> logSequence(
      boost::extents[paddedSequence.size()][test.input.shape()[1]]);
  poplibs_test::embedding::multiSlice(test.input, paddedSequence, logSequence);
  print("Sequence(log)", logSequence, paddedSequence, blankClass, verbose);

  auto alphaLog =
      alpha(logSequence, paddedSequence, blankClass, validTimesteps);
  print("Alphas(log)", alphaLog, paddedSequence, blankClass, verbose);
  print("Alphas", log::exp(alphaLog), paddedSequence, blankClass, verbose);
  auto lastTimeIndex = alphaLog.shape()[1] - 1;
  auto prodSum = log::add(alphaLog[alphaLog.size() - 1][lastTimeIndex],
                          alphaLog[alphaLog.size() - 2][lastTimeIndex]);

  auto loss = -prodSum;
  if (normalise) {
    loss = loss / test.extendedLabel.size();
  }
  std::cout << std::setprecision(6) << "\nP(sequence)=" << std::exp(prodSum)
            << "\nloss:" << loss << "\n";

  auto betaLog = beta(logSequence, paddedSequence, blankClass, validTimesteps);
  print("Betas(log)", betaLog, paddedSequence, blankClass, verbose);
  print("Betas", log::exp(betaLog), paddedSequence, blankClass, verbose);
  auto gradient =
      grad(logSequence, test.input, alphaLog, betaLog, paddedSequence,
           test.input.shape()[0], blankClass, validTimesteps, false);
  print("Gradient(log)", gradient, verbose);
  print("Gradient",
        toLinear(gradient, static_cast<FPType>(1 / std::exp(prodSum))),
        blankClass);
  std::cout << "\n";
}

template <typename FPType>
void testAllPathsInference(const InputSequence<FPType> &test,
                           bool useLogArithmetic, bool verbose) {
  unsigned timeSteps = test.input.shape()[1];
  unsigned sequenceLength = test.input.shape()[0];
  auto allPaths = findAllInputPaths(timeSteps, sequenceLength);
  print("AllInputPaths:", allPaths, UINT_MAX, verbose);
  auto allProbabilities =
      findAllInputPathProbabilities(test.input, allPaths, useLogArithmetic);
  auto allOutputPaths = inputToOutputPath(allPaths, test.blankClass);
  print("AllOutputPaths:", allOutputPaths, UINT_MAX, verbose);
  auto [paths, probabilities, instances] =
      mergePaths(allOutputPaths, allProbabilities, useLogArithmetic);
  if (verbose) {
    std::cout << "\nPaths:" << paths.size() << " probs:" << probabilities.size()
              << "\n";
  }

  // Find the index of the best output path
  auto max = std::max_element(probabilities.begin(), probabilities.end());
  const auto bestIndex = max - probabilities.begin();
  if (verbose) {
    std::cout << "Best has index:" << bestIndex << " Out sequence:";
  }
  std::vector<unsigned> bestPath;
  for (unsigned i = 0; i < allOutputPaths[paths[bestIndex]].size(); i++) {
    if (verbose) {
      std::cout << "B:" << allOutputPaths[paths[bestIndex]][i] << "\n";
    }
    if (allOutputPaths[paths[bestIndex]][i] == popnn::ctc_infer::voidSymbol) {
      break;
    }
    bestPath.push_back(allOutputPaths[paths[bestIndex]][i]);
  }
  print(bestPath, test.blankClass);
  if (useLogArithmetic) {
    std::cout << "Log(P) = " << (*max) << "\n";
    std::cout << "P = " << exp(*max) << "\n";
  } else {
    std::cout << "P = " << (*max) << "\n";
  }
  if (verbose) {
    for (unsigned i = 0; i < probabilities.size(); i++) {
      std::cout << "Path:";
      for (unsigned j = 0; j < allOutputPaths[paths[i]].size(); j++) {
        if (allOutputPaths[paths[i]][j] == popnn::ctc_infer::voidSymbol) {
          break;
        }
        std::cout << allOutputPaths[paths[i]][j] << ",";
      }
      std::cout << "    From paths:";
      for (unsigned j = 0; j < instances[i].size(); j++) {
        std::cout << instances[i][j] << ",";
      }
      std::cout << "    P:" << probabilities[i] << "\n";
    }
  }
  bestPath = extendedLabels(bestPath, test.blankClass, false);
  if (test.extendedLabel != bestPath) {
    throw std::logic_error("Incorrect result");
  }
}
template <typename FPType>
void testBeamSearchInference(
    const InputSequence<FPType> &test, bool useLogArithmetic,
    unsigned beamwidth, boost::optional<std::vector<unsigned>> expectedSequence,
    boost::optional<double> expectedLogProb, bool verbose = false) {
  auto [output, probability] =
      infer<FPType>(test.input, test.blankClass, beamwidth, 1, useLogArithmetic,
                    verbose)
          .at(0);

  std::cout << "Output:" << std::endl;
  print(output);
  std::cout << "P = " << std::exp(probability) << std::endl;
  std::cout << "Log(P) = " << probability << std::endl;

  if (expectedSequence) {
    if (!std::equal(output.begin(), output.end(), expectedSequence->begin())) {
      throw std::logic_error("Incorrect sequence");
    }
  }

  if (expectedLogProb) {
    bool matchesModel =
        checkIsClose<double>(probability, *expectedLogProb, FLOAT_REL_TOL);
    if (!matchesModel) {
      throw std::logic_error("Incorrect probability for sequence");
    }
  }
}

template <typename FPType>
InputSequence<FPType> parseInput(const std::vector<double> &input,
                                 const std::vector<unsigned> &shape,
                                 unsigned blankClass) {
  const auto numClassesIncBlank = shape.at(0);
  const auto tSize = shape.at(1);
  boost::multi_array<FPType, 2> reshapedInput(
      boost::extents[numClassesIncBlank][tSize]);
  for (unsigned i = 0; i < input.size(); i++) {
    *(reshapedInput.data() + i) = input[i];
  }
  return {reshapedInput, {}, blankClass};
}

int main(int argc, char **argv) {
  // Default input parameters.
  bool inference = false;
  bool useLogArithmetic = false;
  bool allPathsMethod = false;
  bool verbose = false;
  bool normalise = false;

  unsigned testNumber = 0;
  unsigned beamWidth = 3;
  unsigned randomTestLength = 15;
  unsigned blankClass = 0;
  unsigned numClassesIncBlank = 4;
  boost::optional<unsigned> seed = boost::none;
  boost::optional<unsigned> baseSequenceLength = boost::none;

  VectorOption<double> input;
  ShapeOption<unsigned> inputShape;
  boost::optional<VectorOption<unsigned>> expectedSequence;
  boost::optional<double> expectedLogProb;

  po::options_description desc("Options");
  // clang-format off
  desc.add_options()
    ("help", "Produce help message")
    ("inference", po::value(&inference)->default_value(inference),
     "Test inference")
    ("beam-width", po::value(&beamWidth)->default_value(beamWidth),
     "Beam width for beam search (not all paths inference)")
    ("input", po::value<VectorOption<double>>(&input),
     "1-D array of softmax input")
    ("input-shape", po::value<ShapeOption<unsigned>>(&inputShape),
     "dimensions to reshape the input")
    ("expectedSequence", po::value<boost::optional<VectorOption<unsigned>>>(&expectedSequence),
     "Expected sequence")
    ("expectedLogProb", po::value<boost::optional<double>>(&expectedLogProb),
     "Expected log probability")
    ("seed", po::value(&seed),
     "If random data provide a seed, if not one is chosen and displayed")
    ("test-length", po::value(&randomTestLength)->default_value(randomTestLength),
     "Test length (t) for random test sequences")
    ("sequence-length", po::value(&baseSequenceLength),
     "Sequence length for which to increase probability for random test input"
     " Defaults to --test-length/2")
    ("num-classes", po::value(&numClassesIncBlank)->default_value(numClassesIncBlank),
     "Number of classes (including blank) to use for random test sequences")
    ("blank-class", po::value(&blankClass)->default_value(blankClass),
     "Index of the blank class")
    ("test", po::value(&testNumber)->default_value(testNumber),
     "Index of hand coded test example")
    ("all-paths", po::value(&allPathsMethod)->default_value(allPathsMethod),
     "Use all paths inference method for reference")
    ("normalise", po::value(&normalise)->default_value(normalise),
     "Normalise the loss")
    ("log", po::value(&useLogArithmetic)->default_value(useLogArithmetic),
     "Compute using log arithmetic")
    ("verbose", po::value(&verbose)->default_value(verbose),
     "More debug print");
  // clang-format on

  po::variables_map vm;
  try {
    po::store(po::parse_command_line(argc, argv, desc), vm);
    if (vm.count("help")) {
      std::cout << desc << "\n\n";
      return 1;
    }

  } catch (std::exception &e) {
    std::cerr << "error parsing command line: " << e.what() << "\n";
    return 1;
  }
  // Needed to set default arguments.
  po::notify(vm);
  if (!baseSequenceLength) {
    baseSequenceLength = ceildiv(randomTestLength, 2u);
  }

  if (useLogArithmetic && !inference) {
    throw std::logic_error("Log arithmetic only supported with inference");
  }

  if (!seed) {
    seed = std::random_device{}();
  }
  std::cout << "Using seed: " << *seed << "\n";
  RandomUtil rand{*seed};

  const auto test = [&]() -> InputSequence<double> {
    if (input.val.empty()) {
      auto [inputProbs, label] = getRandomTestInput<double>(
          *baseSequenceLength, randomTestLength, randomTestLength,
          numClassesIncBlank, blankClass, !useLogArithmetic, rand);
      std::cout << "Input sequence: ";
      print(label, blankClass);
      auto extendedLabel = extendedLabels(label, blankClass, true);
      return {inputProbs, extendedLabel, blankClass};
    } else {
      auto test = parseInput<double>(input.val, inputShape, blankClass);
      test.input = log::log(test.input);
      return test;
    }
  }();

  if (verbose) {
    if (useLogArithmetic) {
      exportPrint("Log Softmax in", test.input);
    }
    print("Test sequence:", test.extendedLabel, test.blankClass);
    std::cout << "The blank symbol is: " << test.blankClass << "\n";
    if (useLogArithmetic) {
      print("Log Softmax input", test.input, test.blankClass);
    } else {
      print("Softmax input", test.input, test.blankClass);
    }
  }
  if (inference) {
    if (allPathsMethod) {
      testAllPathsInference<double>(test, useLogArithmetic, verbose);
    } else {
      testBeamSearchInference<double>(
          test, useLogArithmetic, beamWidth,
          expectedSequence
              ? boost::optional<std::vector<unsigned>>(expectedSequence->val)
              : boost::none,
          expectedLogProb, verbose);
    }
  } else {
    testTraining<double>(test, normalise, verbose);
  }

  return 0;
}


#include <iostream>
#include <vector>
#include <string>
#include "jefflib.h"
#include <boost/tokenizer.hpp>
#include <map>
#include <set>
#include <iomanip>
#include <thread>
#include <chrono>

using namespace std;
using namespace boost;

struct stepInfo_t {
    set<char> prereqs;
    bool workerAssigned = false;
    int clock = 0;
};

void ClearScreen()
{
    cout << string( 100, '\n' );
}

void printSteps(map<char, stepInfo_t > &steps){
    ClearScreen();
    for(const auto item : steps){
        //cout << item.first << ": " << item.second.size() << endl;
        cout << item.first;

        if(item.second.workerAssigned){
            cout << " w ";
        }
        else {
            cout << " . ";
        }

        cout << setfill('0') << setw(2) << item.second.clock;
        cout << " [";
        for(auto prereq : item.second.prereqs){
            cout << prereq;
        }

        cout << "]" << endl;
    }
    cout << endl;
}

int main()
{
    vector<string> vect;
    if(GetStringInput(vect)){
        cout << "Got data!" << endl;
        cout << endl;
    }
    else {
        cout << "Failed to read input :( " << endl;
        return -1;
    }

    map<char, stepInfo_t > steps;

    // .....5..............................36
    // Step V must be finished before step M can begin.
    for (const string line : vect) {
        char c = line[36];
        char p = line[5];
        //cout << c << " deps " << p << endl;
        steps[c].prereqs.insert(p);

        // So, items without prereqs don't have dependencies, so make sure they get
        // have entries created as well...
        steps[p];
    }

    string dummy;
    getline(cin, dummy);

    //cout << "Working..." << endl;
    int elapsedTime = 0;
    int workers = 5;
    while(!steps.empty()){
        printSteps(steps);
        cout << "Elapsed time: " << elapsedTime << endl;

        this_thread::sleep_for (std::chrono::milliseconds(100));

        // Find ready step
        char ready = ' ';

        // Assing workers to any ready steps
        //cout << "Assigning workers..." << endl;
        for(auto& item : steps){
            if(item.second.prereqs.size() == 0 && !item.second.workerAssigned && workers > 0){
                ready = item.first;
                item.second.workerAssigned = true;
                workers--;
                // ASCII A = 65, so always subtract 4
                item.second. clock = int(item.first) - 4;
            }
        }

        // Decrement clocks
        //cout << "Decrementing clocks..." << endl;
        for(auto& item : steps){
            if(item.second.clock >  0){
                item.second. clock -= 1;
            }
        }

        // If any steps are done, remove them and free up workers!
        //cout << "Checking for completed steps, size = " << steps.size() << endl;
        for(auto& item : steps){
            if(item.second.workerAssigned && item.second.clock == 0){
                workers++;
                // Remove step from other step's dependendencies
                char ready = item.first;
                printSteps(steps);
                //cout << "Removing " << ready << " from deps" << endl;
                for(auto& item : steps){
                    auto it = item.second.prereqs.find(ready);
                    if(it != item.second.prereqs.end()){
                        //cout << "Erasing " << ready << ", Num erased: " <<  item.second.erase(*it) << endl;
                        item.second.prereqs.erase(*it);
                    }
                }

                //Remove step from list of steps
                //cout << "Removing " << ready << " from steps" << endl;
                auto it = steps.find(ready);
                if(it != steps.end()){
                    //cout << "Erasing " << ready;
                    steps.erase(it);
                    //cout << "... Erased!" << endl;

                    /*
                    // HACKING
                    if(ready == 'M'){
                        cout << "Removed M. A clock: " << steps['A'].clock << endl;
                        cout << "Elapsed time: " << elapsedTime << endl;
                    }
                    */
                }
                else {
                    cout << "Hmm, couldn't remove " << ready << "??" << endl;
                }
            }
            //cout << "Continuing loop..." << endl;
        }

        // Increment overall elapsed time
        //cout << "Incrementing time from " << elapsedTime << endl;
        elapsedTime++;
    }

    cout << "Elapsed time: " << elapsedTime << endl;
}


#include "aikido/constraint/dart/TSR.hpp"

#include <cmath>
#include <random>
#include <stdexcept>
#include <vector>

#include <boost/format.hpp>
#include <dart/common/Console.hpp>
#include <dart/math/Geometry.hpp>

#include "aikido/common/memory.hpp"

using aikido::statespace::SE3;
using boost::format;
using boost::str;

namespace aikido {
namespace constraint {
namespace dart {

class TSRSampleGenerator : public SampleGenerator
{
public:
  TSRSampleGenerator(const TSRSampleGenerator&) = delete;
  TSRSampleGenerator(TSRSampleGenerator&& other) = delete;
  TSRSampleGenerator& operator=(const TSRSampleGenerator& other) = delete;
  TSRSampleGenerator& operator=(TSRSampleGenerator&& other) = delete;
  virtual ~TSRSampleGenerator() = default;

  // Documentation inherited.
  statespace::ConstStateSpacePtr getStateSpace() const override;

  /// Return a transform sampled from this TSR.
  ///
  /// This function uses the provided RNG to create a sample `Tw_s` from the
  /// `Bw` bounds matrix of this TSR, and returns the result:
  /// `T0_w * Tw_s * Tw_e`.
  ///
  /// \param[in] rng Random number generator from which to sample
  /// \return a transform within the bounds of this TSR.
  bool sample(statespace::StateSpace::State* _state) override;

  // Documentation inherited.
  bool canSample() const override;

  // Documentation inherited.
  int getNumSamples() const override;

private:
  // For internal use only.
  TSRSampleGenerator(
      std::unique_ptr<common::RNG> _rng,
      std::shared_ptr<statespace::SE3> _stateSpace,
      const Eigen::Isometry3d& _T0_w,
      const Eigen::Matrix<double, 6, 2>& _Bw,
      const Eigen::Isometry3d& _Tw_e);

  std::unique_ptr<common::RNG> mRng;

  std::shared_ptr<statespace::SE3> mStateSpace;

  /// Transformation from origin frame into "wiggle" frame.
  Eigen::Isometry3d mT0_w;

  /// Bounds on "wiggling" in `x, y, z, roll, pitch, yaw`.
  Eigen::Matrix<double, 6, 2> mBw;

  /// Transformation from "wiggle" frame into end frame.
  Eigen::Isometry3d mTw_e;

  // True for point TSR.
  bool mPointTSR;

  // True if point TSR and has already been sampled.
  bool mPointTSRSampled;

  friend class TSR;

public:
  EIGEN_MAKE_ALIGNED_OPERATOR_NEW
};

//==============================================================================
TSR::TSR(
    std::unique_ptr<common::RNG> _rng,
    const Eigen::Isometry3d& _T0_w,
    const Eigen::Matrix<double, 6, 2>& _Bw,
    const Eigen::Isometry3d& _Tw_e,
    double _testableTolerance)
  : mT0_w(_T0_w)
  , mBw(_Bw)
  , mTw_e(_Tw_e)
  , mTestableTolerance(_testableTolerance)
  , mRng(std::move(_rng))
  , mStateSpace(std::make_shared<SE3>())
{
  validate();
}

//==============================================================================
TSR::TSR(
    const Eigen::Isometry3d& _T0_w,
    const Eigen::Matrix<double, 6, 2>& _Bw,
    const Eigen::Isometry3d& _Tw_e,
    double _testableTolerance)
  : mT0_w(_T0_w)
  , mBw(_Bw)
  , mTw_e(_Tw_e)
  , mTestableTolerance(_testableTolerance)
  , mRng(std::unique_ptr<common::RNG>(
        new common::RNGWrapper<std::default_random_engine>(0)))
  , mStateSpace(std::make_shared<SE3>())
{
  validate();
}

//==============================================================================
TSR::TSR(const TSR& other)
  : mT0_w(other.mT0_w)
  , mBw(other.mBw)
  , mTw_e(other.mTw_e)
  , mTestableTolerance(other.mTestableTolerance)
  , mRng(other.mRng->clone())
  , mStateSpace(std::make_shared<SE3>())
{
  validate();
}

//==============================================================================
TSR::TSR(TSR&& other)
  : mT0_w(other.mT0_w)
  , mBw(other.mBw)
  , mTw_e(other.mTw_e)
  , mTestableTolerance(other.mTestableTolerance)
  , mRng(std::move(other.mRng))
  , mStateSpace(std::make_shared<SE3>())
{
  validate();
}

//==============================================================================
TSR& TSR::operator=(const TSR& other)
{

  mT0_w = other.mT0_w;
  mBw = other.mBw;
  mTw_e = other.mTw_e;
  mTestableTolerance = other.mTestableTolerance;
  mRng = other.mRng->clone();

  // Intentionally don't assign StateSpace.

  return *this;
}

//==============================================================================
TSR& TSR::operator=(TSR&& other)
{
  mT0_w = std::move(other.mT0_w);
  mBw = std::move(other.mBw);
  mTw_e = std::move(other.mTw_e);
  mTestableTolerance = other.mTestableTolerance;
  mRng = std::move(other.mRng);
  mStateSpace = std::move(other.mStateSpace);

  return *this;
}

//==============================================================================
std::shared_ptr<const statespace::StateSpace> TSR::getStateSpace() const
{
  return mStateSpace;
}

//==============================================================================
std::shared_ptr<statespace::SE3> TSR::getSE3() const
{
  return mStateSpace;
}

//==============================================================================
std::unique_ptr<SampleGenerator> TSR::createSampleGenerator() const
{
  validate();

  if (!mRng)
    throw std::invalid_argument("Random generator is nullptr.");

  for (int i = 0; i < 6; ++i)
  {
    for (int j = 0; j < 2; ++j)
    {
      if (!std::isfinite(mBw(i, j)))
        throw std::invalid_argument(
            str(format("Sampling requires finite bounds. Bw[%d, %d] is %f.") % i
                % j % mBw(i, j)));
    }
  }

  if (mBw.col(0) == mBw.col(1))
  {
    dtwarn << "[TSR::createSampleGenerator] This is a point TSR that represents"
              " a zero measure set in SE(3). The SampleGenerator<Isometry3d>"
              " returned by this function will sample the same pose infinitely"
              " many times.\n";
  }

  return std::unique_ptr<TSRSampleGenerator>(
      new TSRSampleGenerator(mRng->clone(), mStateSpace, mT0_w, mBw, mTw_e));
}

//==============================================================================
bool TSR::isSatisfied(
    const statespace::StateSpace::State* _s, TestableOutcome* outcome) const
{
  auto defaultOutcomeObject
      = dynamic_cast_or_throw<DefaultTestableOutcome>(outcome);

  Eigen::VectorXd dist;
  getValue(_s, dist);

  bool isSatisfiedResult = dist.norm() < mTestableTolerance;
  if (defaultOutcomeObject)
    defaultOutcomeObject->setSatisfiedFlag(isSatisfiedResult);
  return isSatisfiedResult;
}

//==============================================================================
std::unique_ptr<TestableOutcome> TSR::createOutcome() const
{
  return std::unique_ptr<TestableOutcome>(new DefaultTestableOutcome);
}

//==============================================================================
void TSR::validate() const
{
  // Assertion checks for min, max on bounds
  for (int i = 0; i < 6; i++)
  {
    if (mBw(i, 0) > mBw(i, 1))
    {
      throw std::logic_error(
          str(format("Lower bound exceeds upper bound on dimension %d: %f > %f")
              % i % mBw(i, 0) % mBw(i, 1)));
    }
  }
}

//==============================================================================
void TSR::setRNG(std::unique_ptr<common::RNG> rng)
{
  mRng = std::move(rng);
}

//==============================================================================
std::size_t TSR::getConstraintDimension() const
{
  return 6;
}

//==============================================================================
void TSR::getValue(
    const statespace::StateSpace::State* _s, Eigen::VectorXd& _out) const
{
  using SE3 = statespace::SE3;
  using SE3State = SE3::State;

  auto se3state = static_cast<const SE3State*>(_s);
  Eigen::Isometry3d se3 = se3state->getIsometry();

  using TransformTraits = Eigen::TransformTraits;

  Eigen::Isometry3d T0_w_inv = mT0_w.inverse(TransformTraits::Isometry);
  Eigen::Isometry3d Tw_e_inv = mTw_e.inverse(TransformTraits::Isometry);
  Eigen::Isometry3d Tw_s = T0_w_inv * se3 * Tw_e_inv;

  Eigen::Vector3d translation = Tw_s.translation();
  Eigen::Vector3d eulerOrig = ::dart::math::matrixToEulerZYX(Tw_s.linear());
  Eigen::Vector3d eulerZYX = eulerOrig.reverse();

  _out.resize(6);

  for (int i = 0; i < 3; ++i)
  {
    if (translation(i) < mBw(i, 0))
      _out(i) = std::abs(translation(i) - mBw(i, 0));

    else if (translation(i) > mBw(i, 1))
      _out(i) = std::abs(translation(i) - mBw(i, 1));

    else
      _out(i) = 0;
  }

  for (int i = 3; i < 6; ++i)
  {
    // Find n such that: 2*n*pi <= mBw(i, 0) < 2*(n+1)*pi
    int n = mBw(i, 0) / (2 * M_PI);

    // Map eulerZYX(i-3) to [2*n*pi, 2*(n+1)*pi)
    double angle = M_PI * 2 * n + eulerZYX(i - 3);

    // check if angle is within bound
    if ((angle >= mBw(i, 0) && angle <= mBw(i, 1))
        || (angle + M_PI * 2 >= mBw(i, 0) && angle + M_PI * 2 <= mBw(i, 1))
        || (angle - M_PI * 2 >= mBw(i, 0) && angle - M_PI * 2 <= mBw(i, 1)))
    {
      _out(i) = 0;
      continue;
    }

    // Take min-distance between angle and either side of bound
    if (angle < mBw(i, 0))
      _out(i) = std::min(mBw(i, 0) - angle, angle - (mBw(i, 1) - 2 * M_PI));

    else if (mBw(i, 1) < angle)
      _out(i) = std::min(angle - mBw(i, 1), mBw(i, 0) + 2 * M_PI - angle);
  }
}

//==============================================================================
void TSR::getJacobian(
    const statespace::StateSpace::State* _s, Eigen::MatrixXd& _out) const
{
  using SE3 = statespace::SE3;
  using SE3State = SE3::State;

  _out.resize(6, 6);

  auto se3state = static_cast<const SE3State*>(_s);
  Eigen::Isometry3d se3 = se3state->getIsometry();

  Eigen::Vector6d twist = ::dart::math::logMap(se3);

  Eigen::Vector6d posit(twist), negat(twist);

  static constexpr double eps = 1e-5;

  auto se3posit = mStateSpace->createState();
  auto se3negat = mStateSpace->createState();

  // Finite Differencing.
  for (int i = 0; i < 6; ++i)
  {
    posit(i) = twist(i) + eps;
    negat(i) = twist(i) - eps;

    se3posit.setIsometry(::dart::math::expMap(posit));
    se3negat.setIsometry(::dart::math::expMap(negat));

    Eigen::VectorXd positValue, negatValue;
    getValue(se3posit, positValue);
    getValue(se3negat, negatValue);

    Eigen::Vector6d diff = positValue - negatValue;
    _out.col(i) = diff / (2 * eps);

    posit(i) = twist(i);
    negat(i) = twist(i);
  }
}

//==============================================================================
std::vector<ConstraintType> TSR::getConstraintTypes() const
{
  return std::vector<ConstraintType>(6, ConstraintType::INEQUALITY);
}

//==============================================================================
bool TSR::project(
    const statespace::StateSpace::State* /*_s*/,
    statespace::StateSpace::State* /*_out*/) const
{
  // TODO
  return false;
}

//==============================================================================
TSRSampleGenerator::TSRSampleGenerator(
    std::unique_ptr<common::RNG> _rng,
    std::shared_ptr<SE3> _stateSpace,
    const Eigen::Isometry3d& _T0_w,
    const Eigen::Matrix<double, 6, 2>& _Bw,
    const Eigen::Isometry3d& _Tw_e)
  : mRng(std::move(_rng))
  , mStateSpace(std::move(_stateSpace))
  , mT0_w(_T0_w)
  , mBw(_Bw)
  , mTw_e(_Tw_e)
{
  if (!mRng)
  {
    throw std::invalid_argument("Random generator is empty.");
  }

  if (mBw.col(0) == mBw.col(1))
    mPointTSR = true;
  else
    mPointTSR = false;

  mPointTSRSampled = false;
}

//==============================================================================
statespace::ConstStateSpacePtr TSRSampleGenerator::getStateSpace() const
{
  return mStateSpace;
}

//==============================================================================
bool TSRSampleGenerator::sample(statespace::StateSpace::State* _state)
{
  if (mPointTSR && mPointTSRSampled)
    return false;

  using statespace::SE3;

  Eigen::Vector3d translation;
  Eigen::Vector3d angles;

  if (mPointTSR)
  {
    translation = mBw.block(0, 0, 3, 1);
    angles = mBw.block(3, 0, 3, 1);

    mPointTSRSampled = true;
  }
  else
  {
    std::vector<std::uniform_real_distribution<double> > distributions;
    for (int i = 0; i < 6; i++)
      distributions.emplace_back(mBw(i, 0), mBw(i, 1));

    for (int i = 0; i < 3; i++)
      translation(i) = distributions.at(i)(*mRng);

    for (int i = 0; i < 3; i++)
      angles(i) = distributions.at(i + 3)(*mRng);
  }

  Eigen::Isometry3d Tw_s;
  Tw_s.setIdentity();
  Tw_s.translation() = translation;
  Tw_s.linear() = ::dart::math::eulerZYXToMatrix(angles.reverse());

  Eigen::Isometry3d T0_s(mT0_w * Tw_s * mTw_e);
  mStateSpace->setIsometry(static_cast<SE3::State*>(_state), T0_s);

  return true;
}

//==============================================================================
double TSR::getTestableTolerance()
{
  return mTestableTolerance;
}

//==============================================================================
void TSR::setTestableTolerance(double _testableTolerance)
{
  mTestableTolerance = _testableTolerance;
}

//==============================================================================
bool TSRSampleGenerator::canSample() const
{
  if (mPointTSR && mPointTSRSampled)
    return false;

  return true;
}

//==============================================================================
int TSRSampleGenerator::getNumSamples() const
{
  if (mPointTSR && !mPointTSRSampled)
    return 1;

  if (mPointTSR && mPointTSRSampled)
    return 0;

  return NO_LIMIT;
}

} // namespace dart
} // namespace constraint
} // namespace aikido