#include <frovedis.hpp>
#include <frovedis/matrix/pca.hpp>

#include <boost/program_options.hpp>

using namespace boost;
using namespace frovedis;
using namespace std;

void do_pca(const string& input_matrix,
            const string& pca_directions_file,
            const string& pca_scores_file,
            const string& eigen_values_file,
            const string& explained_variance_ratio_file,
            const string& singular_values_file,
            int k,
            bool to_standardize,
            bool binary) {
  time_spent t(DEBUG);
  rowmajor_matrix<double> matrix;
  if(binary) {
    matrix = make_rowmajor_matrix_loadbinary<double>(input_matrix);
  } else {
    matrix = make_rowmajor_matrix_load<double>(input_matrix);
  }
  t.show("load time: ");
  colmajor_matrix<double> pca_directions;
  colmajor_matrix<double> pca_scores;
  std::vector<double> eigen_values;
  std::vector<double> explained_variance_ratio;
  std::vector<double> singular_values;
  std::vector<double> mean; // not used
  double noise_variance; // not used

  pca(std::move(matrix), pca_directions, pca_scores, eigen_values,
      explained_variance_ratio, singular_values, mean, noise_variance,
      k, to_standardize);
  t.show("PCA: ");
  if(binary) {
    pca_directions.savebinary(pca_directions_file);
    if(pca_scores_file != "") {
      pca_scores.savebinary(pca_scores_file);
    }
    if(eigen_values_file != "") {
      make_dvector_scatter(eigen_values).savebinary(eigen_values_file);
    }
    if(explained_variance_ratio_file != "") {
      make_dvector_scatter(explained_variance_ratio).
        savebinary(explained_variance_ratio_file);
    }
    if(singular_values_file != "") {
      make_dvector_scatter(singular_values).
        savebinary(singular_values_file);
    }
  } else {
    pca_directions.save(pca_directions_file);
    if(pca_scores_file != "") {
      pca_scores.save(pca_scores_file);
    }
    if(eigen_values_file != "") {
      make_dvector_scatter(eigen_values).saveline(eigen_values_file);
    }
    if(explained_variance_ratio_file != "") {
      make_dvector_scatter(explained_variance_ratio).
        saveline(explained_variance_ratio_file);
    }
    if(singular_values_file != "") {
      make_dvector_scatter(singular_values).
        saveline(singular_values_file);
    }
  }
  t.show("save time: ");
}

int main(int argc, char* argv[]) {
  use_frovedis use(argc, argv);

  using namespace boost::program_options;

  options_description opt("option");
  opt.add_options()
    ("help,h", "print help")
    ("input,i", value<string>(), "input matrix")
    ("pca_directions,d", value<string>(),
     "PCA directions (components in scikit-learn)")
    ("pca_scores,c", value<string>(),
     "PCA score (no counter part in scikit-learn) [option]")
    ("eigen_values,e", value<string>(),
     "eigen values (explained_variance in scikit-learn) [option]")
    ("explained_variance_ratio,r", value<string>(),
     "explained variance ratio [option]")
    ("singular_values,s", value<string>(), "singular values [option]")
    ("to_standardize,t", "standardize the input [default: false]")
    ("k,k", value<int>(), "number of principal components to compute")
    ("verbose", "set loglevel DEBUG")
    ("verbose2", "set loglevel TRACE")
    ("binary,b", "use binary input/output");

  variables_map argmap;
  store(command_line_parser(argc,argv).options(opt).allow_unregistered().
        run(), argmap);
  notify(argmap);
  string input, pca_directions, pca_scores, eigen_values,
    explained_variance_ratio, singular_values;
  int k;
  bool binary = false;
  bool to_standardize = false;

  if(argmap.count("help")){
    cerr << opt << endl;
    return 1;
  }

  if(argmap.count("input")){
    input = argmap["input"].as<string>();
  } else {
    cerr << "input is not specified" << endl;
    cerr << opt << endl;
    return 1;
  }

  if(argmap.count("pca_directions")){
    pca_directions = argmap["pca_directions"].as<string>();
  } else {
    cerr << "file to store PCA directions is not specified" << endl;
    cerr << opt << endl;
    return 1;
  }

  if(argmap.count("pca_scores")){
    pca_scores = argmap["pca_scores"].as<string>();
  }

  if(argmap.count("eigen_values")){
    eigen_values = argmap["eigen_values"].as<string>();
  }

  if(argmap.count("explained_variance_ratio")){
    explained_variance_ratio = argmap["explained_variance_ratio"].as<string>();
  }

  if(argmap.count("singular_values")){
    singular_values = argmap["singular_values"].as<string>();
  }

  if(argmap.count("to_standardize")){
    to_standardize = true;
  }

  if(argmap.count("k")){
    k = argmap["k"].as<int>();
  } else {
    cerr << "number of singular value to compute is not specified" << endl;
    cerr << opt << endl;
    return 1;
  }

  if(argmap.count("binary")){
    binary = true;
  }

  if(argmap.count("verbose")){
    set_loglevel(DEBUG);
  }

  if(argmap.count("verbose2")){
    set_loglevel(TRACE);
  }

  try {
    do_pca(input, pca_directions, pca_scores, eigen_values,
           explained_variance_ratio, singular_values, k, to_standardize, binary);
  } catch (std::exception& e) {
    cerr << e.what() << endl;
    return 1;
  }
}


/******************************************************************************

  Computation of a convex hull in 3D for a point cloud.

  Copyright (c) 2013
  Dzmitry Hlindzich <hlindzich@gmail.com>
  All rights reserved.

  Redistribution and use in source and binary forms, with or without
  modification, are permitted provided that the following conditions
  are met:
  1.  Redistributions of source code must retain the above copyright
      notice, this list of conditions and the following disclaimer.
  2.  Redistributions in binary form must reproduce the above copyright
      notice, this list of conditions and the following disclaimer in the
      documentation and/or other materials provided with the distribution.

  THIS SOFTWARE IS PROVIDED BY THE AUTHOR AND CONTRIBUTORS "AS IS" AND
  ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
  IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
  ARE DISCLAIMED. IN NO EVENT SHALL THE AUTHOR OR CONTRIBUTORS BE LIABLE
  FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
  DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS
  OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION)
  HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT
  LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY
  OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF
  SUCH DAMAGE.

*******************************************************************************/

#ifndef CONVEX_HULL_3D_HPP_1CB416FB_1F43_4159_A19B_202B43247622_
#define CONVEX_HULL_3D_HPP_1CB416FB_1F43_4159_A19B_202B43247622_

#include <set>
#include <map>
#include <complex>
#include <algorithm>
#include <boost/assert.hpp>

#include "bo/core/vector.hpp"
#include "bo/core/mesh.hpp"
#include "bo/core/triangle.hpp"

namespace bo {
namespace surfaces {

// Computes the convex hull of the given points using the incremental algorithm (Michael Kallay,
// "The Complexity of Incremental Convex Hull Algorithms in Rd" Inf. Process. Lett. 19(4): 197 (1984)).
// The input points must be not coplanar.
template <typename RealType>
class IncrementalConvexHull3D
{
public:
    typedef bo::Vector<RealType, 3> Point3D;
    typedef std::vector<Point3D> Points3D;
    typedef bo::Triangle<Point3D> Face3D;
    typedef std::vector<Face3D> Faces3D;
    typedef bo::Mesh<RealType> Mesh;

    IncrementalConvexHull3D(const Points3D &points)
    {
        // Remove all duplicate points.
        std::set<Point3D> uniq(points.begin(), points.end());
        points_ = Points3D(uniq.begin(), uniq.end());

        // Compute the convex hull.
        initialize_convex_hull();
        expand_convex_hull();
    }

    // Returns the convex hull as a mesh.
    Mesh get_mesh()
    {
        Mesh mesh(faces_.size());

        // Create a reference map (Point3D -> indices of the mesh vertices).
        std::map<Point3D, std::size_t> mymap;

        // Fill in the mesh.
        for (typename FaceSet3D::const_iterator it = faces_.begin(); it != faces_.end(); ++it)
        {
            Face3D f = *it;

            // Add the vertices.
            for (std::size_t i = 0; i < 3; ++i)
            {
                Point3D p = f[i];

                // If the current vertex was not used before, add it.
                if (mymap.find(p) == mymap.end())
                {
                    std::size_t index = mesh.add_vertex(p);
                    mymap[p] = index;
                }
            }

            // Add the face.
            mesh.add_face(typename Mesh::Face(mymap[f.A()], mymap[f.B()], mymap[f.C()]));
        }

        return mesh;
    }

    // Returns the convex hull as the collection of faces.
    Faces3D get_faces()
    {
        return Faces3D(faces_.begin(), faces_.end());
    }

    // Computes the volume of the convex hull using the discrete case
    // of the Gauss-Ostrogradsky's Divergence theorem.
    RealType get_volume()
    {
        RealType v = 0;

        Point3D mass = centroid();

        for (typename FaceSet3D::const_iterator it = faces_.begin();
             it != faces_.end(); ++it)
        {
            Face3D f = *it;

            // Normal vector must be normalized.
            Point3D n = normal(f);
            n /= n.euclidean_norm();

            // Face barycenter.
            Point3D c = (f.A() + f.B() + f.C()) / 3;

            // Direct the face normal outside the volume.
            Point3D in_direction = mass - c;
            if (n * in_direction > 0)
                n = -n;

            RealType a = area(f);

            BOOST_ASSERT(a >= 0);

            v += (n * c) * a;
        }

        BOOST_ASSERT(v >= 0);

        return v / 3;
    }

private:

    // Faces comparator used for the set container.
    struct FaceCompare
    {
        bool operator()(const Face3D &f1, const Face3D &f2) const
        {
            typedef std::set<Point3D> PointSet3D;

            // Sort the vertices using the sorted set.
            PointSet3D fs1;
            fs1.insert(f1.A());
            fs1.insert(f1.B());
            fs1.insert(f1.C());

            PointSet3D fs2;
            fs2.insert(f2.A());
            fs2.insert(f2.B());
            fs2.insert(f2.C());

            BOOST_ASSERT(fs1.size() == fs2.size());

            typename PointSet3D::const_iterator it1 = fs1.begin();
            typename PointSet3D::const_iterator it2 = fs2.begin();

            while (it1 != fs1.end() && it2 != fs2.end())
            {
                if (*it1 < *it2)
                    return true;
                if (*it2 < *it1)
                    return false;

                ++it1;
                ++it2;
            }

            return false;
        }
    };

    typedef std::set<Face3D, FaceCompare> FaceSet3D;

    // Finds the initial tetrahedron.
    void initialize_convex_hull()
    {
        const RealType kEpsilon(0.001);

        if (points_.size() >= 4)
        {
            typedef typename Points3D::iterator Iterator;

            for (Iterator it1 = points_.begin(); it1 != points_.end() - 3; ++it1)
                for (Iterator it2 = it1 + 1; it2 != points_.end() - 2; ++it2)
                    for (Iterator it3 = it2 + 1; it3 != points_.end() - 1; ++it3)
                    {
                        Point3D v1 = *it2 - *it1;
                        Point3D v2 = *it3 - *it1;

                        // Only if non-collinear.
                        if (std::abs(std::abs(v1 * v2) -
                                     v1.euclidean_norm() * v2.euclidean_norm()) > kEpsilon)
                        {
                            Point3D c = v1.cross_product(v2);

                            for (Iterator it4 = it3 + 1; it4 != points_.end(); ++it4)
                            {
                                Point3D v3 = *it4 - *it1;

                                // Only if non-planar.
                                if (std::abs(c * v3) > kEpsilon)
                                {
                                    // Create faces of the initial tetrahedron.
                                    insert_tetrahedron(*it1, *it2, *it3, *it4);

                                    // Remove the points from the list.
                                    points_.erase(it4);
                                    points_.erase(it3);
                                    points_.erase(it2);
                                    points_.erase(it1);

                                    return;
                                }
                            }
                        }
                    }
        }
    }

    // Inserts faces of the tetrahedron.
    void insert_tetrahedron(const Point3D &p1, const Point3D &p2, const Point3D &p3, const Point3D &p4)
    {
        process_face(Face3D(p1, p2, p3));
        process_face(Face3D(p1, p2, p4));
        process_face(Face3D(p2, p3, p4));
        process_face(Face3D(p3, p1, p4));
    }

    // Inserts the face if it is not in the set yet, otherwise deletes it from the set.
    void process_face(const Face3D &f)
    {
        typename FaceSet3D::iterator it = faces_.find(f);

        if (it == faces_.end())
        {
            faces_.insert(f);
        }
        else
        {
            faces_.erase(it);
        }
    }

    // Incremental expansion of the convex hull.
    void expand_convex_hull()
    {
        const RealType kEpsilon(0.0001);

        if (faces_.size() > 3)
        {
            for (typename Points3D::const_iterator itp = points_.begin(); itp != points_.end(); ++itp)
            {
                Point3D p = *itp;

                Point3D c = centroid();

                Faces3D visible_faces;

                // Find the "visible" faces.
                for (typename FaceSet3D::const_iterator it = faces_.begin(); it != faces_.end(); ++it)
                {
                    Face3D f = *it;
                    Point3D v_c = c - f.A();
                    v_c = v_c / v_c.euclidean_norm();
                    Point3D v_p = p - f.A();
                    v_p = v_p / v_p.euclidean_norm();

                    Point3D n = normal(f);
                    n = n / n.euclidean_norm();

                    RealType dot_c = n * v_c;
                    RealType dot_p = n * v_p;

                    // If the face is "visible" from point p.
                    if (dot_c * dot_p < -kEpsilon)
                    {
                        visible_faces.push_back(f);
                    }
                }

                // Expand the hull and remove the "visible" faces.
                for (typename Faces3D::const_iterator it = visible_faces.begin();
                     it != visible_faces.end(); ++it)
                {
                    Face3D f = *it;

                    // Add new faces and remove the "visible" face.
                    insert_tetrahedron(f.A(), f.B(), f.C(), p);
                }
            }
        }
    }

    Point3D centroid()
    {
        Point3D center(0);
        std::size_t count = 0;

        for (typename FaceSet3D::const_iterator it = faces_.begin(); it != faces_.end(); ++it)
        {
            center += (it->A() + it->B() + it->C());
            count += 3;
        }

        return center / count;
    }

    inline Point3D normal(const Face3D &f)
    {
        return (f.B() - f.A()).cross_product(f.C() - f.A());
    }

    inline RealType area(const Face3D &f)
    {
        return normal(f).euclidean_norm() / 2;
    }

    Points3D points_;
    FaceSet3D faces_;
};

} // namespace surfaces
} // namespace bo

#endif // CONVEX_HULL_3D_HPP_1CB416FB_1F43_4159_A19B_202B43247622_


///////////////////////////////////////////////////////////////////////////////
//  Copyright 2014 Anton Bikineev
//  Copyright 2014 Christopher Kormanyos
//  Copyright 2014 John Maddock
//  Copyright 2014 Paul Bristow
//  Distributed under the Boost
//  Software License, Version 1.0. (See accompanying file
//  LICENSE_1_0.txt or copy at http://www.boost.org/LICENSE_1_0.txt)

#ifndef BOOST_MATH_HYPERGEOMETRIC_2F0_HPP
#define BOOST_MATH_HYPERGEOMETRIC_2F0_HPP

#include <boost/math/policies/policy.hpp>
#include <boost/math/policies/error_handling.hpp>
#include <boost/math/special_functions/detail/hypergeometric_series.hpp>
#include <boost/math/special_functions/laguerre.hpp>
#include <boost/math/special_functions/hermite.hpp>
#include <boost/math/tools/fraction.hpp>

namespace boost { namespace math { namespace detail {

   template <class T>
   struct hypergeometric_2F0_cf
   {
      //
      // We start this continued fraction at b on index -1
      // and treat the -1 and 0 cases as special cases.
      // We do this to avoid adding the continued fraction result
      // to 1 so that we can accurately evaluate for small results
      // as well as large ones.  See  http://functions.wolfram.com/07.31.10.0002.01
      //
      T a1, a2, z;
      int k;
      hypergeometric_2F0_cf(T a1_, T a2_, T z_) : a1(a1_), a2(a2_), z(z_), k(-2) {}
      typedef std::pair<T, T> result_type;

      result_type operator()()
      {
         ++k;
         if (k <= 0)
            return std::make_pair(z * a1 * a2, 1);
         return std::make_pair(-z * (a1 + k) * (a2 + k) / (k + 1), 1 + z * (a1 + k) * (a2 + k) / (k + 1));
      }
   };

   template <class T, class Policy>
   T hypergeometric_2F0_cf_imp(T a1, T a2, T z, const Policy& pol, const char* function)
   {
      using namespace boost::math;
      hypergeometric_2F0_cf<T> evaluator(a1, a2, z);
      boost::uintmax_t max_iter = policies::get_max_series_iterations<Policy>();
      T cf = tools::continued_fraction_b(evaluator, policies::get_epsilon<T, Policy>(), max_iter);
      policies::check_series_iterations<T>(function, max_iter, pol);
      return cf;
   }

   template <class T, class Policy>
   inline T hypergeometric_2F0_imp(T a1, T a2, const T& z, const Policy& pol, bool asymptotic = false)
   {
      //
      // The terms in this series go to infinity unless one of a1 and a2 is a negative integer.
      //
      using std::swap;
      BOOST_MATH_STD_USING

      static const char* const function = "boost::math::hypergeometric_2F0<%1%,%1%,%1%>(%1%,%1%,%1%)";

      if (z == 0)
         return 1;

      bool is_a1_integer = (a1 == floor(a1));
      bool is_a2_integer = (a2 == floor(a2));

      if (!asymptotic && !is_a1_integer && !is_a2_integer)
         return boost::math::policies::raise_overflow_error<T>(function, 0, pol);
      if (!is_a1_integer || (a1 > 0))
      {
         swap(a1, a2);
         swap(is_a1_integer, is_a2_integer);
      }
      //
      // At this point a1 must be a negative integer:
      //
      if(!asymptotic && (!is_a1_integer || (a1 > 0)))
         return boost::math::policies::raise_overflow_error<T>(function, 0, pol);
      //
      // Special cases first:
      //
      if (a1 == 0)
         return 1;
      if ((a1 == a2 - 0.5f) && (z < 0))
      {
         // http://functions.wolfram.com/07.31.03.0083.01
         int n = static_cast<int>(static_cast<boost::uintmax_t>(boost::math::lltrunc(-2 * a1)));
         T smz = sqrt(-z);
         return pow(2 / smz, -n) * boost::math::hermite(n, 1 / smz);
      }

      if (is_a1_integer && is_a2_integer)
      {
         if ((a1 < 1) && (a2 <= a1))
         {
            const unsigned int n = static_cast<unsigned int>(static_cast<boost::uintmax_t>(boost::math::lltrunc(-a1)));
            const unsigned int m = static_cast<unsigned int>(static_cast<boost::uintmax_t>(boost::math::lltrunc(-a2 - n)));

            return (pow(z, T(n)) * boost::math::factorial<T>(n, pol)) *
               boost::math::laguerre(n, m, -(1 / z), pol);
         }
         else if ((a2 < 1) && (a1 <= a2))
         {
            // function is symmetric for a1 and a2
            const unsigned int n = static_cast<unsigned int>(static_cast<boost::uintmax_t>(boost::math::lltrunc(-a2)));
            const unsigned int m = static_cast<unsigned int>(static_cast<boost::uintmax_t>(boost::math::lltrunc(-a1 - n)));

            return (pow(z, T(n)) * boost::math::factorial<T>(n, pol)) *
               boost::math::laguerre(n, m, -(1 / z), pol);
         }
      }

      if ((a1 * a2 * z < 0) && (a2 < -5) && (fabs(a1 * a2 * z) > 0.5))
      {
         // Series is alternating and maybe divergent at least for the first few terms
         // (until a2 goes positive), try the continued fraction:
         return hypergeometric_2F0_cf_imp(a1, a2, z, pol, function);
      }

      return detail::hypergeometric_2F0_generic_series(a1, a2, z, pol);
   }

} // namespace detail

template <class T1, class T2, class T3, class Policy>
inline typename tools::promote_args<T1, T2, T3>::type hypergeometric_2F0(T1 a1, T2 a2, T3 z, const Policy& /* pol */)
{
   BOOST_FPU_EXCEPTION_GUARD
      typedef typename tools::promote_args<T1, T2, T3>::type result_type;
   typedef typename policies::evaluation<result_type, Policy>::type value_type;
   typedef typename policies::normalise<
      Policy,
      policies::promote_float<false>,
      policies::promote_double<false>,
      policies::discrete_quantile<>,
      policies::assert_undefined<> >::type forwarding_policy;
   return policies::checked_narrowing_cast<result_type, Policy>(
      detail::hypergeometric_2F0_imp<value_type>(
         static_cast<value_type>(a1),
         static_cast<value_type>(a2),
         static_cast<value_type>(z),
         forwarding_policy()),
      "boost::math::hypergeometric_2F0<%1%>(%1%,%1%,%1%)");
}

template <class T1, class T2, class T3>
inline typename tools::promote_args<T1, T2, T3>::type hypergeometric_2F0(T1 a1, T2 a2, T3 z)
{
   return hypergeometric_2F0(a1, a2, z, policies::policy<>());
}

  } } // namespace boost::math

#endif // BOOST_MATH_HYPERGEOMETRIC_HPP


#include <catch2/catch.hpp>

#include "CollisionConstraints.hpp"
#include "CTCD.h"

#include <Eigen/Dense>
#include <finitediff.hpp>

const std::vector<std::string> constraintNames = {
    "VOLUME", "GRAPHICS", "NONSMOOTH_NEWMARK", "GAP_FUNCTION", "CMR",
    "VERSCHOOR"
};

TEST_CASE("Test Point-Triangle Collision Constraints",
    "[collision-constraints][sqp]")
{
    // point
    Eigen::Vector3d v0(0, 1, -0.5);
    // triangle = (v1, v2, v3) in counter-clockwise order
    Eigen::Vector3d v1(-1, 0, 1);
    Eigen::Vector3d v2(1, 0, 1);
    Eigen::Vector3d v3(0, 0, -1);

    // displacements
    double u0y = GENERATE(-1.1, 0.0, 1.1);
    Eigen::Vector3d u0(0, u0y, 0);
    double u1y = GENERATE(-1.1, 0.0, 1.1);
    Eigen::Vector3d u1(0, u1y, 0);

    bool intersects = (-u0y + u1y >= 1);

    double toi;
    bool ccd_intersects = CTCD::vertexFaceCTCD(
        v0, v1, v3, v2, v0 + u0, v1 + u1, v3 + u1, v2 + u1, 0, toi);
    REQUIRE(ccd_intersects == intersects);

    IPC::CollisionConstraintType constraintType = GENERATE(
        IPC::CollisionConstraintType::VOLUME,
        IPC::CollisionConstraintType::GRAPHICS,
        IPC::CollisionConstraintType::VERSCHOOR);
    double c;
    IPC::compute_collision_constraint(
        v0, v1, v3, v2,
        v0 + u0, v1 + u1, v3 + u1, v2 + u1,
        constraintType, /*is_edge_edge=*/false, toi,
        c);

    CAPTURE(u0y, u1y, c, constraintNames[constraintType]);
    CHECK((c < 0) == intersects);

    Eigen::Vector12d grad_c;
    IPC::compute_collision_constraint_gradient(
        v0, v1, v3, v2,
        v0 + u0, v1 + u1, v3 + u1, v2 + u1,
        constraintType, /*is_edge_edge=*/false, toi,
        grad_c);

    Eigen::VectorXd x(12);
    x.segment<3>(0) = v0 + u0;
    x.segment<3>(3) = v1 + u1;
    x.segment<3>(6) = v3 + u1; // Swap order to match compute_collision_constraint
    x.segment<3>(9) = v2 + u1; // Swap order to match compute_collision_constraint
    std::function<double(const Eigen::VectorXd&)> f = [&](const Eigen::VectorXd& x) {
        assert(x.size() == 12);
        double c;
        IPC::compute_collision_constraint(
            v0, v1, v3, v2,
            x.segment<3>(0), x.segment<3>(3), x.segment<3>(6), x.segment<3>(9),
            constraintType, /*is_edge_edge=*/false, toi,
            c);
        return c;
    };
    Eigen::VectorXd finite_grad_c(12);
    fd::finite_gradient(x, f, finite_grad_c);
    CAPTURE(grad_c.transpose(), finite_grad_c.transpose());
    CHECK(fd::compare_gradient(grad_c, finite_grad_c));
}

TEST_CASE("Test Edge-Edge Collision Constraints",
    "[collision-constraints][sqp]")
{
    // e0 = (v0, v1)
    Eigen::Vector3d v0(-1, -1, 0);
    Eigen::Vector3d v1(1, -1, 0);
    // e2 = (v2, v3)
    Eigen::Vector3d v2(0, 1, -1);
    Eigen::Vector3d v3(0, 1, 1);

    // displacements
    double y_displacement = GENERATE(-2, 0, 2);
    Eigen::Vector3d u0(0, y_displacement, 0);
    Eigen::Vector3d u1(0, -y_displacement, 0);

    bool intersects = y_displacement >= 1.0;
    double toi;
    bool ccd_intersects = CTCD::edgeEdgeCTCD(
        v0, v1, v2, v3, v0 + u0, v1 + u0, v2 + u1, v3 + u1, 0, toi);
    CAPTURE(y_displacement);
    REQUIRE(ccd_intersects == intersects);

    IPC::CollisionConstraintType constraintType = GENERATE(
        IPC::CollisionConstraintType::VOLUME,
        IPC::CollisionConstraintType::GRAPHICS,
        IPC::CollisionConstraintType::VERSCHOOR);

    double c;
    IPC::compute_collision_constraint(
        v0, v1, v2, v3,
        v0 + u0, v1 + u0, v2 + u1, v3 + u1,
        constraintType, /*is_edge_edge=*/true, toi,
        c);

    CAPTURE(y_displacement, c, constraintNames[constraintType]);
    CHECK((c < 0) == intersects);

    typedef Eigen::Matrix<double, 12, 1> Vector12d;
    Vector12d grad_c;
    IPC::compute_collision_constraint_gradient(
        v0, v1, v2, v3,
        v0 + u0, v1 + u0, v2 + u1, v3 + u1,
        constraintType, /*is_edge_edge=*/true, toi,
        grad_c);

    Eigen::VectorXd x(12);
    x.segment<3>(0) = v0 + u0;
    x.segment<3>(3) = v1 + u0;
    x.segment<3>(6) = v2 + u1;
    x.segment<3>(9) = v3 + u1;
    std::function<double(const Eigen::VectorXd&)> f = [&](const Eigen::VectorXd& x) {
        assert(x.size() == 12);
        double c;
        IPC::compute_collision_constraint(
            v0, v1, v2, v3,
            x.segment<3>(0), x.segment<3>(3), x.segment<3>(6), x.segment<3>(9),
            constraintType, /*is_edge_edge=*/true, toi,
            c);
        return c;
    };
    Eigen::VectorXd finite_grad_c(12);
    fd::finite_gradient(x, f, finite_grad_c);
    CAPTURE(grad_c.transpose(), finite_grad_c.transpose());
    CHECK(fd::compare_gradient(grad_c, finite_grad_c));
}

TEST_CASE("Barycentric coordinates", "[barycentric]")
{
    // Create random triangle
    Eigen::Vector3d a = Eigen::Vector3d::Random();
    Eigen::Vector3d b = Eigen::Vector3d::Random();
    Eigen::Vector3d c = Eigen::Vector3d::Random();
    // Create random point
    Eigen::Vector3d barycentric_coords = Eigen::Vector3d::Random();
    barycentric_coords(0) = 1 - barycentric_coords(1) - barycentric_coords(2);
    Eigen::Vector3d p = barycentric_coords(0) * a + barycentric_coords(1) * b + barycentric_coords(2) * c;

    Eigen::Vector3d computed_barycentric_coords;
    IPC::barycentric_coordinates(p.transpose(), a.transpose(), b.transpose(), c.transpose(), computed_barycentric_coords);
    Eigen::Vector3d computed_p = computed_barycentric_coords(0) * a + computed_barycentric_coords(1) * b + computed_barycentric_coords(2) * c;

    CAPTURE(barycentric_coords, computed_barycentric_coords);
    CHECK((computed_barycentric_coords - barycentric_coords).norm() < 1e-12);
    CAPTURE(p, computed_p);
    CHECK((computed_p - p).norm() < 1e-12);
}


#include <iostream>
#include <numeric>
#include <boost/assign.hpp>
#include "SketchRing.h"

using namespace std;
using namespace boost;

unsigned hashf(int value) {
    unsigned temp = ((value % 3451) + 1) * (value ^ 0x238bc1) + (1 + value * 13);
    return temp ^ std::hash<int>{}(value);
}

int SketchRing::size() const {
    return cms.begin()->second.size() * cms.size();
}

void SketchRing::update(unsigned int item, int diff) {
    auto hash   = hashf(item) % domainMax;
    auto bucket = buckets.find(hash)->second;
//    cout << item << " bucket:" << bucket << endl;
    cms.at(bucket).update(item, diff);
}

int SketchRing::estimate(unsigned int item) const {
    auto hash = hashf(item) % domainMax;
    return cms.at(buckets.find(hash)->second).estimate(item);
}

SketchRing::SketchRing(int width, int depth, int n,
                       Point domainMax) : buckets(make_pair(0, 0)), domainMax(domainMax) {
    auto interval = double(domainMax) / n;

    for (Point high = interval, low = 1; low < domainMax; low = high, high += interval) {

//        std::cout << "I:" << high << ":" << buckets.iterative_size() << std::endl;
        cms.emplace(piecewise_construct,
                    forward_as_tuple(low),
                    forward_as_tuple(width, depth, 1));
        buckets.insert(make_pair(icl::interval<Point>::type(low, high), low));
    }
}

auto erase(SketchRing::IntervalMap& buckets, SketchRing::IntervalMap::const_iterator it) {
    bool       isFirst  = it == buckets.begin();
    const auto erasedIv = it->first;

    auto prevBucket = isFirst ? buckets.rbegin()->second : (--it)->second;

    buckets.set(make_pair(erasedIv, prevBucket));

    return prevBucket;
}

void SketchRing::shrink() {
    // find least loaded sketch
    unsigned  minLoad = 0;
    auto      minIter = cms.begin();
    for (auto it      = cms.begin(); it != cms.end(); ++it) {
        auto load = it->second.load();
        if (load > minLoad) {
            minLoad = load;
            minIter = it;
        }
    }

    // now we have the min
    const auto& sketch = minIter->second;
    const auto location = minIter->first;
    const auto& hh = sketch.heavyHitters(.6);

    auto toAdd = erase(buckets, buckets.find(location));

    for (auto[hitter, freq]: hh)
        cms.at(toAdd).update(hitter, freq);

    cms.erase(minIter);
}

void SketchRing::expand() {
    throw std::bad_function_call();
    // find most loaded sketch
    unsigned  maxLoad = 0;
    auto      maxIter = cms.begin();
    for (auto it      = cms.begin(); it != cms.end(); ++it) {
        auto load = it->second.load();
        if (load > maxLoad) {
            maxLoad = load;
            maxIter = it;
        }
    }

    // now we have the max
    const auto& sketch = maxIter->second;
    const auto location = maxIter->first;
    const auto& hh = sketch.heavyHitters(.4);

    auto toAdd = erase(buckets, buckets.find(location));

    for (auto[hitter, freq]: hh)
        cms.at(toAdd).update(hitter, freq);

    cms.erase(maxIter);

    throw std::exception();
}

void SketchRing::resize(int n) {
    if (n <= 2) throw out_of_range("can't shrink below 2");

    while (cms.size() > n) {
        shrink();
    }
    while (cms.size() < n) {
        expand();
    }
}

std::vector<std::pair<int, int> > SketchRing::heavyHitters(int threshold) const {
    std::vector<std::pair<int, int> > res;

    for (const auto&[p, cm] : cms) {
        const auto& hh = cm.heavyHitters(threshold);
        res.insert(res.end(), hh.begin(), hh.end());
    }

    return res;
}

extern "C" {
#include "SketchRingC.h"

SK_type* SK_init(int width, int depth, int n, int domainMax) {
    return new SketchRing(width, depth, n, domainMax);
}

void SK_destroy(SK_type* sr) {
    delete sr;
}

int SK_size(const SK_type* sr) {
    return sr->size();
}

void SK_update(SK_type* sr, unsigned int item, int diff) {
    sr->update(item, diff);
}

int SK_estimate(const SK_type* sr, unsigned int item) {
    return sr->estimate(item);
}

void SK_resize(SK_type* sr, int n) {
    sr->resize(n);
}

//void HH_free(HH* hh) {
//    free(hh->freq);
//    free(hh->items);
//    *hh = {0, 0, 0};
//}
//
//HH SK_heavyHitters(const SK_type* sk, int threshold) {
//    HH res;
//    const auto& hh = sk->heavyHitters(threshold);
//
//    res.size = hh.size();
//    res.items = static_cast<int*>(malloc(sizeof(int) * res.size));
//    res.freq = static_cast<int*>(malloc(sizeof(int) * res.size));
//
//    for (int i = 0; i < res.size; ++i) {
//        res.items[i] = hh[i].first;
//        res.freq[i] = hh[i].second;
//    }
//
//    return res;
//}
int* SK_heavyHitters(const SK_type* sk, int threshold) {
    const auto& hh = sk->heavyHitters(threshold);
    unsigned size = hh.size();
    int* res = static_cast<int*>(calloc(size+1, (sizeof(int))));

    res[0] = size;

    for (int i = 1; i <= res[0]; ++i) {
        res[i] = hh[i-1].first;
    }

    return res;
}

}


#include "faasm/faasm.h"

#include <Eigen/Dense>

#include "ndpapi.h"

#include <stdio.h>
#include <string.h>

#include <regex>
#include <string_view>
#include <vector>

#include "util/kmeansrex.cpp"

using std::string_view;

string_view objKey;

uint32_t endswap32(uint32_t x)
{
    return ((x & 0xFFu) << 24u) | ((x & 0xFF00u) << 8u) |
           ((x & 0xFF0000u) >> 8u) | ((x & 0xFF000000u) >> 24u);
}

Eigen::MatrixXf pcaReducedData;

int work()
{
    uint32_t fetchedLength{};
    uint8_t* objData =
      __faasmndp_getMmap(reinterpret_cast<const uint8_t*>(objKey.data()),
                         objKey.size(),
                         1 * 1024 * 1024 * 1024,
                         &fetchedLength);
    if (objData == nullptr || fetchedLength < 64) {
        const string_view output{
            "FAILED - no object found with the given key"
        };
        faasmSetOutput(reinterpret_cast<const uint8_t*>(output.data()),
                       output.size());
        return 1;
    }
    const uint32_t magic = endswap32(*reinterpret_cast<uint32_t*>(objData));
    const uint32_t nImages =
      endswap32(*reinterpret_cast<uint32_t*>(objData + 4));
    const uint32_t nRows = endswap32(*reinterpret_cast<uint32_t*>(objData + 8));
    const uint32_t nCols =
      endswap32(*reinterpret_cast<uint32_t*>(objData + 12));
    const uint8_t* pixelDataStart = objData + 16;
    if (magic != 0x00000803) {
        const string_view output{ "FAILED - bad magic number" };
        faasmSetOutput(reinterpret_cast<const uint8_t*>(output.data()),
                       output.size());
        return 2;
    }
    const uint32_t pixelsPerImage = nRows * nCols;
    if (pixelsPerImage * nImages > fetchedLength - 16) {
        const string_view output{ "FAILED - missing image data" };
        faasmSetOutput(reinterpret_cast<const uint8_t*>(output.data()),
                       output.size());
        return 3;
    }
    Eigen::MatrixXf imageData(nImages, pixelsPerImage);
    const float fmul = 1.0f / 255.0f;
    for (uint32_t image = 0; image < nImages; image++) {
        const uint8_t* imageStart = pixelDataStart + (nRows * nCols) * image;
        for (uint32_t pixel = 0; pixel < nRows * nCols; pixel++) {
            imageData(image, pixel) = float(imageStart[pixel]) * fmul;
        }
    }
    // normalize data
    {
        Eigen::VectorXf mins = imageData.colwise().minCoeff();
        Eigen::VectorXf maxs = imageData.colwise().maxCoeff();
        Eigen::VectorXf scaleCoeff = maxs - mins;
        for (auto& val : scaleCoeff) {
            val = (val == 0) ? 1.0f : 1.0f / val;
        }
        imageData.rowwise() -= mins.transpose();
        imageData.array().rowwise() *= scaleCoeff.transpose().array();
    }
    // center on means
    {
        Eigen::VectorXf means = imageData.colwise().mean();
        imageData.rowwise() -= means.transpose();
    }
    Eigen::MatrixXf covarianceMat = imageData.adjoint() * imageData;
    covarianceMat = covarianceMat / (imageData.rows() - 1);
    Eigen::SelfAdjointEigenSolver<Eigen::MatrixXf> eigensolver(covarianceMat);
    Eigen::VectorXf eigenvalues = eigensolver.eigenvalues();
    eigenvalues /= eigenvalues.sum();
    Eigen::MatrixXf eigenvectors = eigensolver.eigenvectors();
    Eigen::MatrixXf pcaTransform = eigenvectors.rightCols(32);
    pcaReducedData = imageData * pcaTransform;
    return 0;
}

int main(int argc, char* argv[])
{
    long inputSz = faasmGetInputSize();
    std::vector<uint8_t> inputBuf(inputSz);
    faasmGetInput(inputBuf.data(), inputBuf.size());
    string_view inputStr(reinterpret_cast<char*>(inputBuf.data()),
                         inputBuf.size());
    if (inputStr.size() < 1) {
        const string_view output{
            "FAILED - no key/value pair. Usage: grep with input 'key regex'"
        };
        faasmSetOutput(reinterpret_cast<const uint8_t*>(output.data()),
                       output.size());
        return 0;
    }
    objKey = inputStr;

    if (__faasmndp_storageCallAndAwait(work) != 0) {
        return 1;
    }

    const int K = 10;
    const int n_features = pcaReducedData.cols();
    const int n_iters = 25;
    const int seed = 1234567;
    const int n_examples = pcaReducedData.rows();

    Eigen::ArrayXXd dblData = pcaReducedData.array().cast<double>();
    Eigen::ArrayXXd mu = Eigen::ArrayXXd::Zero(K, n_features);
    Eigen::ArrayXd z = Eigen::ArrayXd::Zero(n_examples);
    RunKMeans(dblData.data(),
              n_examples,
              n_features,
              K,
              n_iters,
              seed,
              "plusplus",
              mu.data(),
              z.data());

    std::string output;
    output.reserve(z.rows() * 4);
    for (auto& el : z) {
        output += std::to_string(int(el));
        output.push_back(' ');
    }
    output.push_back('\n');

    faasmSetOutput(reinterpret_cast<uint8_t*>(output.data()), output.size());
    return 0;
}


/*
 * @Author: tom: https://github.com/TOMsworkspace
 * @Date: 2021-09-03 15:32:52
 * @Last Modified by: tom: https://github.com/TOMsworkspace
 * @Last Modified time: 2021-09-03 19:17:13
 */

#include <algorithm>
#include <sstream>
#include <iostream>

#include <Eigen\\Dense>
using namespace std;
using namespace Eigen;

//#include <gl/GL.h>
//#include <glad/glad.h>
//#include <GLFW/glfw3.h>

//#include <irrklang/irrKlang.h>
//using namespace irrklang;

#include "dtkFemSimulation.h"

#include "GL/freeglut.h"

//#include "resource_manager.h"
//#include "sprite_renderer.h"

int dim = 2;
int n_node_x = 50;
int n_node_y = 6;
float node_mass = 1.0f;
int n_node = n_node_x * n_node_y;
int n_fem_element = (n_node_x - 1) * (n_node_y - 1) * 2;
float deltat = 3e-4;
float deltax = (1.0 / 32);

float Young_E = 1000.0f; /**< 杨氏模量 */
float Poisson_r = 0.3f; /**< 泊松比 [0 - 0.5] */
float Lame_parameter_1 = Young_E / (2 * (1 + Poisson_r));
float Lame_parameter_2 = Young_E * Poisson_r / ((1 + Poisson_r) * (1 - 2 * Poisson_r));
float element_v = 0.01f; /**< 微元体积 */

float radius = 0.05;

int iterate_time = 3;

inline int mesh(int i, int j) {
    return i * n_node_y + j;
}

dtkFemSimulation::dtkFemSimulation(unsigned int width, unsigned int height)
: dtkScene(width, height),
points(n_node),pre_points(n_node),points_v(n_node),points_force(n_node),B(n_fem_element),
MeshTable(n_fem_element,std::vector<int>(3,0)),sphere(dtk::dtkGraphicsKernel::Point2(0.5,0.25), radius),
total_energy(0),pre_total_energy(0),spherecenter(0.5,0.25)
{

}

Matrix2f dtkFemSimulation::compute_D(int i){
    int a = MeshTable[i][0];
    int b = MeshTable[i][1];
    int c = MeshTable[i][2];

    Matrix2f ans;
    ans(0,0) = points[a][0] - points[c][0];
    ans(0,1) = points[b][0] - points[c][0];
    ans(1,0) = points[a][1] - points[c][1];
    ans(1,1) = points[b][1] - points[c][1];
    return ans;
}

void dtkFemSimulation::compute_B(){
    for(int i = 0; i < n_fem_element; ++i){
        this->B[i] = compute_D(i).inverse();
        //cout << setw(6) << B[i] << endl;
    }
}

Matrix2f dtkFemSimulation::compute_P(int i){
    Matrix2f D = compute_D(i);
    Matrix2f F = D * B[i];

    Matrix2f F_T = F.transpose().inverse();

    float J = max(0.5f, F.determinant()); /**< 形变率 */

    return Lame_parameter_1 * (F - F_T) + Lame_parameter_2 * log(J) * F_T ;
    //Matrix2f ans = D * this->B[i];
}

void dtkFemSimulation::compute_total_energy(){

    //this->total_energy = 0.0f;
    for(int i = 0; i < n_fem_element; ++i){
        Matrix2f D = compute_D(i);
        Matrix2f F = D * B[i];

        //NeoHooken
        float I1 = (F * F.transpose()).trace();
        float J = max(0.2f, (float)F.determinant()); /**< 形变率 */

        //cout << J << endl;

        float element_energy_density = 0.5 * Lame_parameter_1 * (I1 - dim) - Lame_parameter_1 * log(J) + 0.5 * Lame_parameter_2 * log(J) * log(J);
        this->total_energy += element_energy_density * element_v;
    }
}

dtkFemSimulation::~dtkFemSimulation()
{
}

void dtkFemSimulation::Init()
{
    dtkScene::Init();
    //TODO: load shaders

    //TODO: configure shaders

    //TODO: load textures

    //TODO: set render-specific controls

    //TODO: configure Scene objects

    //build mesh
    for(int i = 0; i < n_node_x; ++i){
        for(int j = 0; j < n_node_y; ++j){
            int idx = mesh(i,j);
            //this->points[idx][0] = -14 + i * deltax * 0.5;
            //this->points[idx][1] = 8 + j * deltax * 0.5 + i * deltax * 0.05;

            this->points[idx][0] = 0.1f + i * deltax * 0.5f;
            this->points[idx][1] = 0.5f + j * deltax * 0.5f + i * deltax * 0.1f;
            this->points_v[idx][0] = 0.0f;
            this->points_v[idx][1] = -1.0f;
        }
    }

    //this->pre_points = points;

    for(int i = 0; i < n_node_x - 1; ++i ){
        for(int j = 0; j < n_node_y - 1; ++j){
            //element id
            int eidx = (i * (n_node_y - 1) + j) * 2;
            this->MeshTable[eidx][0] = mesh(i,j);
            this->MeshTable[eidx][1] = mesh(i + 1, j);
            this->MeshTable[eidx][2] = mesh(i,j + 1);

            eidx = (i * (n_node_y - 1) + j) * 2 + 1;
            this->MeshTable[eidx][0] = mesh(i,j + 1);
            this->MeshTable[eidx][1] = mesh(i + 1,j + 1);
            this->MeshTable[eidx][2] = mesh(i + 1,j);
        }
    }

    compute_B();

    //TODO: audio
}

void dtkFemSimulation::compute_force(){
    for(int i = 0; i < n_node ; ++i){
        this->points_force[i] = Vector2f(0.0f, - 10.0f * node_mass);
    }

    for(int i = 0; i < n_fem_element; ++i){

        Matrix2f P = compute_P(i);
        Matrix2f H = - element_v * (P * (this->B[i].transpose()));

        Vector2f h1 = Vector2f(H(0,0), H(1,0));
        Vector2f h2 = Vector2f(H(0,1), H(1,1));

        int a = this->MeshTable[i][0];
        int b = this->MeshTable[i][1];
        int c = this->MeshTable[i][2];

        this->points_force[a] += h1;
        this->points_force[b] += h2;
        this->points_force[c] += -(h1 + h2);
    }
}

void dtkFemSimulation::Update(float dt)
{
    //TODO: update objects
    //TODO: check for object collisions
    if(this->State == SCENE_ACTIVE){

        // 迭代多轮, 防止穿透
        for(int i = 0; i < iterate_time ; ++i){
            //this->pre_total_energy = total_energy;
            compute_total_energy();
            DoCollisions();
            compute_force();

            //cout << total_energy << endl;

            //float deltaU = this->total_energy - this->pre_total_energy;
            //deltaU = abs(deltaU) < 1e-9 ? 0 : deltaU;
            for(int i = 0; i < n_node; ++i){
                // update points

                //Vector2f deltaX = this->points[i] - this->pre_points[i];

                //Vector2f diffUtoX = Vector2f(deltaU / deltaX[0], deltaU / deltaX[1]);

               // Vector2f diffUtoX = Vector2f(abs(deltaX[0]) > 1e-4 ? deltaU / deltaX[0] : 0.0f, abs(deltaX[1]) > 1e-4 ? deltaU / deltaX[1] : 0.0f);

                //diffUtoX = Vector2f(0.0f,0.0f);

                //cout << diffUtoX << endl;

               // this->points_v[i] = (this->points_v[i] + ((- diffUtoX / node_mass) + Vector2f(0.0f, -10.0f)) * deltat) * exp(deltat * -6);

                this->points_v[i] = (this->points_v[i] + (this->points_force[i] / node_mass) * deltat) * exp(deltat * -3);
                //this->pre_points[i] = this->points[i];

                this->points[i] += deltat * this->points_v[i];
            }

           // this->pre_points = this->points;
           // this->pre_total_energy = this->total_energy;
        }
    }
}

void dtkFemSimulation::ProcessInput(float dt)
{
    dtkScene::ProcessInput(dt);
    //TODO: process input(keys)

}

void dtkFemSimulation::Render()
{
    //if(this->State == SCENE_ACTIVE){
        //TODO: draw circle

        Vector2f center = spherecenter;
        //Vector2f(this->sphere.center()[0], this->sphere.center()[1]);

        glColor3f(0x06 * 1.0 / 0xff, 0x85 * 1.0 / 0xff, 0x87 * 1.0 / 0xff);
        glBegin(GL_POLYGON);

        int n = 100;
        for (int i = 0; i < n; i++)
        {
            glVertex2f(center[0] + radius * cos(2 * dtk::dtkPI / n * i), center[1] + radius * sin(2 * dtk::dtkPI / n * i));
        }
        glEnd();

        //TODO: draw fem element(triangles here)
        glColor3f(0x4f * 1.0 / 0xff, 0xb9 * 1.0 / 0xff, 0x9f * 1.0 / 0xff);
        glBegin(GL_LINES);
        for(int i = 0; i < n_fem_element; ++i){
            for(int j = 0 ; j < 3; ++j){
                int a = this->MeshTable[i][j];
                int b = this->MeshTable[i][(j  + 1) % 3];

                //draw line from a to b;
                glVertex2f(this->points[a][0] , this->points[a][1] );
                glVertex2f(this->points[b][0], this->points[b][1] );
            }
        }
        //glEnd();

        glColor3f(1.0, 1.0, 1.0);
        //glBegin(GL_LINES);
        glVertex2f(0.0f, 0.2f);
        glVertex2f(1.0f, 0.2f);

        glVertex2f(1.0f, 0.85f);
        glVertex2f(1.0f, 0.2f);

        glVertex2f(1.0f, 0.85f);
        glVertex2f(0.0f, 0.85f);

        glVertex2f(0.0f, 0.85f);
        glVertex2f(0.0f, 0.2f);
        glEnd();
   // }
}

// collision detection

void dtkFemSimulation::DoCollisions()
{
    if(this->State == SCENE_ACTIVE){
        Vector2f center = spherecenter;
        //Vector2f(this->sphere.center()[0], this->sphere.center()[1]);
        float radius = this->sphere.squared_radius();

        for(int i = 0; i < n_node; ++i){
            //# Collide with sphere

            Vector2f dis = this->points[i] - center;
            if((float)(dis.dot(dis)) < radius * radius)
            {
                Vector2f normal = dis.normalized();

                this->points[i] = center + radius * normal;
                this->points_v[i] -=  (this->points_v[i].dot(normal)) *  normal;
            }

            // Collide with ground

            if(this->points[i][1] < 0.2f) {
                this->points[i][1] = 0.2f;
                this->points_v[i][1] = 0.0f;
            }

            if(this->points[i][1] > 0.9f) {
                this->points[i][1] = 0.9f;
                this->points_v[i][1] = 0.0f;
            }

            if(this->points[i][0] < 0.0f) {
                this->points[i][0] = 0.0f;
                this->points_v[i][0] = 0.0f;
            }

            if(this->points[i][0] > 1.0f) {
                this->points[i][0] = 1.0f;
                this->points_v[i][0] = 0.0f;
            }
        }

    }
}

void dtkFemSimulation::moveBall(int x, int y){

    //this->spherecenter = Vector2f((x) * 1.0 / 800 , (600 - y) * 1.0 / 600 );
}

float dtkFemSimulation::getEnergy(){
    return this->total_energy;
}


#ifndef EXPSUM_KERNEL_FUNCTIONS_BESSEL_KERNEL_HPP
#define EXPSUM_KERNEL_FUNCTIONS_BESSEL_KERNEL_HPP

#include <cassert>
#include <cmath>

#include <armadillo>

namespace expsum
{

template <typename T>
struct bessel_j_kernel
{
public:
    using size_type   = arma::uword;
    using real_type   = T;
    using vector_type = arma::Col<T>;
    using matrix_type = arma::Mat<T>;

private:
    vector_type exponent_;
    vector_type weight_;
    real_type v_;   // order of Bessel function Jv(x)
    real_type eps_; // tolerance

public:
    void compute(real_type v, real_type eps);

    size_type size() const
    {
        return exponent_.size();
    }

    const vector_type& exponents() const
    {
        return exponent_;
    }

    const vector_type& weights() const
    {
        return weight_;
    }

private:
};

template <typename T>
void bessel_j_kernel<T>::compute(real_type v, real_type eps)
{
}

} // namespace: expsum

#endif /* EXPSUM_KERNEL_FUNCTIONS_BESSEL_KERNEL_HPP */


////////////////////////////////////////////////////////////////////////////////
//  Copyright (c) 2011 Bryce Lelbach
//  Copyright (C) 2010 Scott McMurray
//
//  Distributed under the Boost Software License, Version 1.0. (See accompanying
//  file LICENSE_1_0.txt or copy at http://www.boost.org/LICENSE_1_0.txt)
////////////////////////////////////////////////////////////////////////////////

#if !defined(HPX_68441787_8A39_4DAC_9D9E_705B6FA19BCD)
#define HPX_68441787_8A39_4DAC_9D9E_705B6FA19BCD

#include <boost/cstdint.hpp>

namespace hpx { namespace util { namespace hardware
{

template <typename T, typename U>
bool has_bit_set(T value, U bit)
{ return (value & (1 << bit)) != 0; }

template <std::size_t N, typename T>
struct unbounded_shifter
{
    static T shl(T x) { return unbounded_shifter<N-1, T>::shl(T(x << 1)); }
    static T shr(T x) { return unbounded_shifter<N-1, T>::shr(T(x >> 1)); }
};

template <typename T>
struct unbounded_shifter<0, T>
{
    static T shl(T x) { return x; }
    static T shr(T x) { return x; }
};

template <std::size_t N, typename T>
T unbounded_shl(T x)
{ return unbounded_shifter<N, T>::shl(x); }

template <std::size_t N, typename T>
T unbounded_shr(T x)
{ return unbounded_shifter<N, T>::shr(x); }

template <std::size_t Low, std::size_t High, typename Result, typename T>
Result get_bit_range(T x)
{
    T highmask = unbounded_shl<High, T>(~T());
    T lowmask = unbounded_shl<Low, T>(~T());
    return static_cast<Result>
        (unbounded_shr<Low, T>(T(x & (lowmask ^ highmask))));
}

template <std::size_t Low, typename Result, typename T>
Result pack_bits(T x)
{ return unbounded_shl<Low, Result>(static_cast<Result>(x)); }

}}}

#endif // HPX_68441787_8A39_4DAC_9D9E_705B6FA19BCD


#pragma once

#include <pcl/point_types.h>
#include <pcl/features/normal_3d.h>
#include <pcl/io/pcd_io.h>
#include <pcl/io/ply_io.h>
#include <pcl/filters/voxel_grid.h>
#include <pcl/point_cloud.h>
#include <Eigen/Dense>

namespace {

Eigen::Vector3f getRealSortedEigenValues(Eigen::Matrix<std::complex<float>, 3, 1> eigenValues)
{
  Eigen::Vector3f realEigenValues;
  for (std::size_t i = 0; i < 3; i++){
    realEigenValues(i) = std::real(eigenValues(i));
  }
  if (realEigenValues(0) > realEigenValues(1))
    std::swap(realEigenValues(0), realEigenValues(1));
  if (realEigenValues(1) > realEigenValues(2))
    std::swap(realEigenValues(1), realEigenValues(2));
  if (realEigenValues(0) > realEigenValues(1))
    std::swap(realEigenValues(0), realEigenValues(1));
  // std::cout << "realEigenValues: " << realEigenValues << "\n";
  return realEigenValues;
}

} // namespace

namespace pcl {

// L: linearity, P: planarity, S: sphericality
struct PointXYZLPS
{
	PCL_ADD_POINT4D;
	// union
	// {
	// 	float coordinate[3];
	// 	struct
	// 	{
	// 		float x;
	// 		float y;
	// 		float z;
	// 	};
	// };
	float linearity;
	float planarity;
	float sphericity;
	EIGEN_MAKE_ALIGNED_OPERATOR_NEW;
} EIGEN_ALIGN16;

struct PointXYZILPS
{
    PCL_ADD_POINT4D;                  // preferred way of adding a XYZ+padding
    float intensity;
    float linearity;
	float planarity;
	float sphericity;
    EIGEN_MAKE_ALIGNED_OPERATOR_NEW   // make sure our new allocators are aligned
} EIGEN_ALIGN16;                    // enforce SSE padding for correct memory alignment

// struct PointXYZILPS
// {
// 	// PCL_ADD_POINT4D;
// 	union
// 	{
// 		float coordinate[3];
// 		struct
// 		{
// 			float x;
// 			float y;
// 			float z;
// 		};
// 	};
//     float intensity;
// 	float linearity;
// 	float planarity;
// 	float sphericity;
// 	EIGEN_MAKE_ALIGNED_OPERATOR_NEW;
// } EIGEN_ALIGN16;

template <typename PointInT, typename PointOutT>
class DescriptorEstimation : public Feature<PointInT, PointOutT>
{
    typedef typename DescriptorEstimation<PointInT, PointOutT>::PointCloudOut PointCloudOut;
    using Feature<PointInT, PointOutT>::input_;
    using Feature<PointInT, PointOutT>::search_parameter_;
    using Feature<PointInT, PointOutT>::search_radius_;
    typedef boost::shared_ptr<NormalEstimation<PointInT, PointOutT> > Ptr;
    typedef boost::shared_ptr<const NormalEstimation<PointInT, PointOutT> > ConstPtr;

protected:

    void computeFeature (PointCloudOut &output)
    {
        for (int i = 0; i < input_->size(); i++){
            const PointInT& searchPoint = input_->points[i];
            std::vector<int> pointIdxRadiusSearch;
            std::vector<float> pointRadiusSquaredDistance;
            // kdtree.radiusSearch (searchPoint, radius, pointIdxRadiusSearch, pointRadiusSquaredDistance);

            this->searchForNeighbors (i, search_parameter_, pointIdxRadiusSearch, pointRadiusSquaredDistance);

            if (pointIdxRadiusSearch.size() < 5) {
                continue;
            }
            typename pcl::PointCloud<PointInT>::Ptr neighbours (new pcl::PointCloud<PointInT>);
            for (std::size_t j = 0; j < pointIdxRadiusSearch.size (); ++j){
                neighbours->points.push_back(input_->points[ pointIdxRadiusSearch[j]]);
            }

            Eigen::Matrix3f covariance_matrix;

            // 16-bytes aligned placeholder for the XYZ centroid of a surface patch
            Eigen::Vector4f xyz_centroid;

            pcl::compute3DCentroid (*neighbours, xyz_centroid);
            // Compute the 3x3 covariance matrix
            pcl::computeCovarianceMatrix (*neighbours, xyz_centroid, covariance_matrix);
            // std::cout << "CCCC" << std::endl;
            // std::cout << "covariance_matrix: \n" << covariance_matrix << "\n";
            // std::cout << "xyz_centroid: \n" << xyz_centroid << "\n";

            Eigen::EigenSolver<Eigen::Matrix3f> eig(covariance_matrix);     // [vec val] = eig(A)
            // std::cout << "DDDD" << std::endl;
            Eigen::Matrix<std::complex<float>, 3, 1> eigenComplexValues = eig.eigenvalues();
            // std::cout << "EEEE" << std::endl;
            // Eigen::Matrix3f D = eig.pseudoEigenvalueMatrix();

            // int col_index, row_index;
            // std::cout << D.maxCoeff(&row_index, &col_index) << endl;
            // std::cout << row_index << " " << col_index << endl;
            // std::cout << "eigen value matrix: " << getRealSortedEigenValues(eigenValues) << "\n";
            Eigen::Vector3f eigenValues = getRealSortedEigenValues(eigenComplexValues);
            // std::cout << "FFFF" << std::endl;

            PointOutT pointWithFeature;

            pointWithFeature.x = searchPoint.x;
            pointWithFeature.y = searchPoint.y;
            pointWithFeature.z = searchPoint.z;
            pointWithFeature.linearity = 1 - eigenValues[1] / eigenValues[2];
            pointWithFeature.planarity = (eigenValues[1] - eigenValues[0]) / eigenValues[2];
            pointWithFeature.sphericity = eigenValues[0] / eigenValues[2];
            output.points.push_back(pointWithFeature);

            // float linearity = 1 - eigenValues[1] / eigenValues[2];
            // float planarity = (eigenValues[1] - eigenValues[0]) / eigenValues[2];
            // float sphericity = eigenValues[0] / eigenValues[2];
            // std::cout << "(" << linearity << ", " << planarity << ", " << sphericity << ")" << std::endl;
            // cloud->points[i].linearity = linearity;
            // cloud->points[i].planarity = planarity;
            // cloud->points[i].sphericity = sphericity;

        }
    }

}; // class DescriptorEstimation

} // namespace pcl

POINT_CLOUD_REGISTER_POINT_STRUCT(PointXYZLPS,// 注册点类型宏
	(float, x, x)
	(float, y, y)
	(float, z, z)
	(float, linearity, linearity)
	(float, planarity, planarity)
	(float, sphericity, sphericity)
)

POINT_CLOUD_REGISTER_POINT_STRUCT(PointXYZILPS,
	(float, x, x)
	(float, y, y)
	(float, z, z)
    (float, intensity, intensity)
	(float, linearity, linearity)
	(float, planarity, planarity)
	(float, sphericity, sphericity)
)

// pcl::NormalEstimation<PointTypeIO, PointTypeFull> ne;
// ne.setInputCloud (cloud_out);
// ne.setSearchMethod (search_tree);
// ne.setRadiusSearch (1);
// ne.compute (*cloud_with_normals);

typedef pcl::PointXYZ PointIn;
typedef pcl::PointXYZLPS PointOut;

int test_new(int argc, char** argv)
{
    pcl::PointCloud<PointIn>::Ptr cloud(new pcl::PointCloud<PointIn>);
    pcl::PointCloud<PointOut>::Ptr cloud_out(new pcl::PointCloud<PointOut>);

    pcl::io::loadPCDFile<PointIn> (argv[1], *cloud);
    // pcl::io::loadPLYFile<PointT> (argv[1], *cloud);

    pcl::VoxelGrid<PointIn> sor;
	sor.setInputCloud(cloud);
	sor.setLeafSize(0.1f, 0.1f, 0.1f);
	sor.filter(*cloud);

    // pcl::KdTreeFLANN<PointIn> kdtree;
    pcl::search::KdTree<PointIn>::Ptr search_tree (new pcl::search::KdTree<PointIn>);
    search_tree->setInputCloud (cloud);

    pcl::DescriptorEstimation<PointIn, PointOut> de;
    de.setInputCloud(cloud);
    de.setSearchMethod(search_tree);
    de.setRadiusSearch(1);
    de.compute (*cloud_out);

    cloud_out->width = 1;
    cloud_out->height = cloud_out->points.size();
    pcl::io::savePCDFile ("descriptors_output.pcd", *cloud_out);

    return 0;

}

int test_gt(int argc, char** argv)
{
    pcl::PointCloud<PointIn>::Ptr cloud (new pcl::PointCloud<PointIn>);
    pcl::PointCloud<PointOut>::Ptr cloudOut (new pcl::PointCloud<PointOut>);

    pcl::io::loadPCDFile<PointIn> (argv[1], *cloud);
    // pcl::io::loadPLYFile<PointT> (argv[1], *cloud);

    pcl::VoxelGrid<PointIn> sor;
	sor.setInputCloud(cloud);
	sor.setLeafSize(0.1f, 0.1f, 0.1f);
	sor.filter(*cloud);

    // Placeholder for the 3x3 covariance matrix at each surface patch
    Eigen::Matrix3f covariance_matrix;
    // 16-bytes aligned placeholder for the XYZ centroid of a surface patch
    Eigen::Vector4f xyz_centroid;

    // Estimate the XYZ centroid
    pcl::compute3DCentroid (*cloud, xyz_centroid);

    // Compute the 3x3 covariance matrix
    pcl::computeCovarianceMatrix (*cloud, xyz_centroid, covariance_matrix);

    // std::cout << "covariance_matrix: \n" << covariance_matrix << "\n";
    // std::cout << "xyz_centroid: \n" << xyz_centroid << "\n";

    pcl::KdTreeFLANN<PointIn> kdtree;

    kdtree.setInputCloud (cloud);

    float radius = 0.5f;

    std::cout << "Done" << std::endl;

    for (int i = 0; i < cloud->size(); i++){
        PointIn searchPoint = cloud->points[i];
        std::vector<int> pointIdxRadiusSearch;
        std::vector<float> pointRadiusSquaredDistance;
        kdtree.radiusSearch (searchPoint, radius, pointIdxRadiusSearch, pointRadiusSquaredDistance);

        pcl::PointCloud<PointIn>::Ptr neighbours (new pcl::PointCloud<PointIn>);
        for (std::size_t j = 0; j < pointIdxRadiusSearch.size (); ++j){
        neighbours->points.push_back(cloud->points[ pointIdxRadiusSearch[j]]);

        }
        if (neighbours->points.size() < 5) {
            continue;
        }
        pcl::compute3DCentroid (*neighbours, xyz_centroid);
        // std::cout << "BBBB" << std::endl;
        // Compute the 3x3 covariance matrix
        pcl::computeCovarianceMatrix (*neighbours, xyz_centroid, covariance_matrix);
        // std::cout << "CCCC" << std::endl;
        // std::cout << "covariance_matrix: \n" << covariance_matrix << "\n";
        // std::cout << "xyz_centroid: \n" << xyz_centroid << "\n";

        Eigen::EigenSolver<Eigen::Matrix3f> eig(covariance_matrix);     // [vec val] = eig(A)
        // std::cout << "DDDD" << std::endl;
        Eigen::Matrix<std::complex<float>, 3, 1> eigenComplexValues = eig.eigenvalues();
        // std::cout << "EEEE" << std::endl;
        // Eigen::Matrix3f D = eig.pseudoEigenvalueMatrix();

        // int col_index, row_index;
        // std::cout << D.maxCoeff(&row_index, &col_index) << endl;
        // std::cout << row_index << " " << col_index << endl;
        // std::cout << "eigen value matrix: " << getRealSortedEigenValues(eigenValues) << "\n";
        Eigen::Vector3f eigenValues = getRealSortedEigenValues(eigenComplexValues);
        // std::cout << "FFFF" << std::endl;

        PointOut pointWithFeature;

        pointWithFeature.x = searchPoint.x;
        pointWithFeature.y = searchPoint.y;
        pointWithFeature.z = searchPoint.z;
        pointWithFeature.linearity = 1 - eigenValues[1] / eigenValues[2];
        pointWithFeature.planarity = (eigenValues[1] - eigenValues[0]) / eigenValues[2];
        pointWithFeature.sphericity = eigenValues[0] / eigenValues[2];
        cloudOut->points.push_back(pointWithFeature);

        // float linearity = 1 - eigenValues[1] / eigenValues[2];
        // float planarity = (eigenValues[1] - eigenValues[0]) / eigenValues[2];
        // float sphericity = eigenValues[0] / eigenValues[2];
        // std::cout << "(" << linearity << ", " << planarity << ", " << sphericity << ")" << std::endl;
        // cloud->points[i].linearity = linearity;
        // cloud->points[i].planarity = planarity;
        // cloud->points[i].sphericity = sphericity;

    }

    cloudOut->width = 1;
    cloudOut->height = cloudOut->points.size();
    pcl::io::savePCDFile ("descriptors_output.pcd", *cloudOut);

    return 0;
}


#include "logit_sb_mixing.h"

#include <google/protobuf/stubs/casts.h>

#include <Eigen/Dense>
#include <memory>
#include <numeric>
#include <stan/math/prim.hpp>
#include <vector>

#include "mixing_prior.pb.h"
#include "mixing_state.pb.h"
#include "src/hierarchies/abstract_hierarchy.h"
#include "src/utils/proto_utils.h"
#include "src/utils/rng.h"

void LogitSBMixing::initialize() {
  if (prior == nullptr) {
    throw std::invalid_argument("Mixing prior was not provided");
  }
  auto priorcast = cast_prior();
  num_components = priorcast->num_components();
  initialize_state();
  acceptance_rates = Eigen::VectorXd::Zero(num_components - 1);
  n_iter = 0;
}

void LogitSBMixing::initialize_state() {
  auto priorcast = cast_prior();
  if (priorcast->has_normal_prior()) {
    Eigen::VectorXd prior_vec =
        bayesmix::to_eigen(priorcast->normal_prior().mean());
    dim = prior_vec.size();
    state.precision = stan::math::inverse_spd(
        bayesmix::to_eigen(priorcast->normal_prior().var()));
    if (dim != state.precision.cols()) {
      throw std::invalid_argument(
          "Hyperparameters dimensions are not consisent");
    }
    if (priorcast->step_size() <= 0) {
      throw std::invalid_argument("Step size parameter must be > 0");
    }

    state.regression_coeffs = Eigen::MatrixXd(dim, num_components - 1);
    for (int i = 0; i < num_components - 1; i++) {
      state.regression_coeffs.col(i) = prior_vec;
    }

  } else {
    throw std::invalid_argument("Unrecognized mixing prior");
  }
}

double LogitSBMixing::full_cond_lpdf(
    const Eigen::VectorXd &alpha, const unsigned int clust,
    const std::vector<unsigned int> &allocations) {
  auto priorcast = cast_prior();
  Eigen::VectorXd prior_mean =
      bayesmix::to_eigen(priorcast->normal_prior().mean());
  double like = -0.5 * ((alpha - prior_mean).transpose() * state.precision *
                        (alpha - prior_mean))(0);
  for (int i = 0; i < allocations.size(); i++) {
    if (allocations[i] >= clust) {
      double is_curr_clus = (allocations[i] == clust);
      double prob = sigmoid(covariates_ptr->row(i).dot(alpha));
      like += is_curr_clus * std::log(prob) +
              (1.0 - is_curr_clus) * std::log(1.0 - prob);
    }
  }
  return like;
}

Eigen::VectorXd LogitSBMixing::grad_full_cond_lpdf(
    const Eigen::VectorXd &alpha, const unsigned int clust,
    const std::vector<unsigned int> &allocations) {
  auto priorcast = cast_prior();
  Eigen::VectorXd prior_mean =
      bayesmix::to_eigen(priorcast->normal_prior().mean());
  Eigen::VectorXd grad = state.precision * (prior_mean - alpha);
  for (int i = 0; i < allocations.size(); i++) {
    if (allocations[i] >= clust) {
      double is_curr_clus = (allocations[i] == clust);
      double prob = sigmoid(covariates_ptr->row(i).dot(alpha));
      grad += (is_curr_clus - prob) * covariates_ptr->row(i);
    }
  }
  return grad;
}

void LogitSBMixing::update_state(
    const std::vector<std::shared_ptr<AbstractHierarchy>> &unique_values,
    const std::vector<unsigned int> &allocations) {
  n_iter += 1;
  // Langevin-Adjusted Metropolis-Hastings step
  unsigned int n = allocations.size();
  auto &rng = bayesmix::Rng::Instance().get();
  auto priorcast = cast_prior();
  Eigen::VectorXd prior_mean =
      bayesmix::to_eigen(priorcast->normal_prior().mean());
  double step = priorcast->step_size();
  double prop_var = std::sqrt(2.0 * step);
  // Loop over components
  for (int h = 0; h < num_components - 1; h++) {
    Eigen::VectorXd state_c = state.regression_coeffs.col(h);
    // Draw proposed state from its distribution
    Eigen::VectorXd prop_mean =
        state_c + step * grad_full_cond_lpdf(state_c, h, allocations);
    auto prop_covar = prop_var * Eigen::MatrixXd::Identity(dim, dim);
    Eigen::VectorXd state_prop =
        stan::math::multi_normal_rng(prop_mean, prop_covar, rng);
    // Compute acceptance ratio
    double full_cond_ratio = full_cond_lpdf(state_prop, h, allocations) -
                             full_cond_lpdf(state_c, h, allocations);
    double prop_ratio = (-0.5 / prop_var) *
                        ((state_prop - prop_mean).dot(state_prop - prop_mean) -
                         (state_c - prop_mean).dot(state_c - prop_mean));
    double log_accept_ratio = full_cond_ratio - prop_ratio;
    // Accept with probability ratio
    double p = stan::math::uniform_rng(0.0, 1.0, rng);
    if (p < std::exp(log_accept_ratio)) {
      state.regression_coeffs.col(h) = state_prop;
      acceptance_rates(h) += 1;
    }
  }
}

void LogitSBMixing::set_state_from_proto(
    const google::protobuf::Message &state_) {
  auto &statecast =
      google::protobuf::internal::down_cast<const bayesmix::MixingState &>(
          state_);
  state.regression_coeffs =
      bayesmix::to_eigen(statecast.log_sb_state().regression_coeffs());
}

void LogitSBMixing::write_state_to_proto(
    google::protobuf::Message *out) const {
  bayesmix::LogSBState state_;
  bayesmix::to_proto(state.regression_coeffs,
                     state_.mutable_regression_coeffs());
  google::protobuf::internal::down_cast<bayesmix::MixingState *>(out)
      ->mutable_log_sb_state()
      ->CopyFrom(state_);
}

Eigen::VectorXd LogitSBMixing::get_weights(
    const bool log, const bool propto,
    const Eigen::RowVectorXd &covariate /*= Eigen::RowVectorXd(0)*/) const {
  // Compute eta
  std::vector<double> eta(num_components);
  for (int h = 0; h < num_components - 1; h++) {
    eta[h] = covariate.dot(state.regression_coeffs.col(h));
  }
  eta[num_components - 1] = 1.0;
  // Compute cumulative sums of logarithms
  std::vector<double> cumsum(num_components + 1, 0.0);
  for (int h = 1; h < num_components + 1; h++) {
    cumsum[h] = cumsum[h - 1] + std::log(sigmoid(-eta[h - 1]));
  }
  // Compute weights
  Eigen::VectorXd logweights(num_components);
  for (int h = 0; h < num_components; h++) {
    logweights(h) = std::log(sigmoid(eta[h])) + cumsum[h];
  }
  if (log) {
    return logweights;
  } else {
    return logweights.array().exp();
  }
}


// Copyright 2018 The Simons Foundation, Inc. - All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//    http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

#ifndef NETKET_FULLCONNLAYER_HH
#define NETKET_FULLCONNLAYER_HH

#include <Eigen/Dense>
#include <complex>
#include <fstream>
#include <random>
#include <vector>
#include "Lookup/lookup.hpp"
#include "Utils/all_utils.hpp"
#include "abstract_layer.hpp"

namespace netket {

template <typename Activation, typename T>
class FullyConnected : public AbstractLayer<T> {
  using VectorType = typename AbstractLayer<T>::VectorType;
  using MatrixType = typename AbstractLayer<T>::MatrixType;

  Activation activation_;  // activation function class

  bool usebias_;

  int in_size_;        // input size
  int out_size_;       // output size
  int npar_;           // number of parameters in layer
  MatrixType weight_;  // Weight parameters, W(in_size x out_size)
  VectorType bias_;    // Bias parameters, b(out_size x 1)
                       // Note that input of this layer is also the output of
                       // previous layer

  std::size_t scalar_bytesize_;

 public:
  using StateType = typename AbstractLayer<T>::StateType;
  using LookupType = typename AbstractLayer<T>::LookupType;

  /// Constructor
  FullyConnected(const int input_size, const int output_size,
                 const bool use_bias = false)
      : activation_(),
        usebias_(use_bias),
        in_size_(input_size),
        out_size_(output_size) {
    Init();
  }

  explicit FullyConnected(const json &pars) : activation_() { Init(pars); }

  void Init(const json &pars) {
    in_size_ = FieldVal(pars, "Inputs");
    out_size_ = FieldVal(pars, "Outputs");

    usebias_ = FieldOrDefaultVal(pars, "UseBias", true);

    scalar_bytesize_ = sizeof(std::complex<double>);

    weight_.resize(in_size_, out_size_);
    bias_.resize(out_size_);

    npar_ = in_size_ * out_size_;

    if (usebias_) {
      npar_ += out_size_;
    } else {
      bias_.setZero();
    }
    std::string buffer = "";

    InfoMessage(buffer) << "Fully Connected Layer " << in_size_ << " --> "
                        << out_size_ << std::endl;
    InfoMessage(buffer) << "# # UseBias = " << usebias_ << std::endl;
  }

  void Init() {
    scalar_bytesize_ = sizeof(std::complex<double>);

    weight_.resize(in_size_, out_size_);
    bias_.resize(out_size_);

    npar_ = in_size_ * out_size_;

    if (usebias_) {
      npar_ += out_size_;
    } else {
      bias_.setZero();
    }

    std::string buffer = "";

    InfoMessage(buffer) << "Fully Connected Layer " << in_size_ << " --> "
                        << out_size_ << std::endl;
    InfoMessage(buffer) << "# # UseBias = " << usebias_ << std::endl;
  }

  void to_json(json &pars) const override {
    json layerpar;
    layerpar["Name"] = "FullyConnected";
    layerpar["UseBias"] = usebias_;
    layerpar["Inputs"] = in_size_;
    layerpar["Outputs"] = out_size_;
    layerpar["Bias"] = bias_;
    layerpar["Weight"] = weight_;

    pars["Machine"]["Layers"].push_back(layerpar);
  }

  void from_json(const json &pars) override {
    if (FieldExists(pars, "Weight")) {
      weight_ = pars["Weight"];
    } else {
      weight_.setZero();
    }
    if (FieldExists(pars, "Bias")) {
      bias_ = pars["Bias"];
    } else {
      bias_.setZero();
    }
  }

  void InitRandomPars(int seed, double sigma) override {
    VectorType par(npar_);

    netket::RandomGaussian(par, seed, sigma);

    SetParameters(par, 0);
  }

  int Npar() const override { return npar_; }

  int Ninput() const override { return in_size_; }

  int Noutput() const override { return out_size_; }

  void GetParameters(VectorType &pars, int start_idx) const override {
    int k = start_idx;

    if (usebias_) {
      std::memcpy(pars.data() + k, bias_.data(), out_size_ * scalar_bytesize_);
      k += out_size_;
    }

    std::memcpy(pars.data() + k, weight_.data(),
                in_size_ * out_size_ * scalar_bytesize_);
  }

  void SetParameters(const VectorType &pars, int start_idx) override {
    int k = start_idx;

    if (usebias_) {
      std::memcpy(bias_.data(), pars.data() + k, out_size_ * scalar_bytesize_);

      k += out_size_;
    }

    std::memcpy(weight_.data(), pars.data() + k,
                in_size_ * out_size_ * scalar_bytesize_);
  }

  void InitLookup(const VectorType &v, LookupType &lt,
                  VectorType &output) override {
    lt.resize(1);
    lt[0].resize(out_size_);

    Forward(v, lt, output);
  }

  void UpdateLookup(const VectorType &input,
                    const std::vector<int> &input_changes,
                    const VectorType &new_input, LookupType &theta,
                    const VectorType & /*output*/,
                    std::vector<int> &output_changes,
                    VectorType &new_output) override {
    const int num_of_changes = input_changes.size();
    if (num_of_changes == in_size_) {
      output_changes.resize(out_size_);
      new_output.resize(out_size_);
      Forward(new_input, theta, new_output);
    } else if (num_of_changes > 0) {
      UpdateTheta(input, input_changes, new_input, theta);
      output_changes.resize(out_size_);
      new_output.resize(out_size_);
      Forward(theta, new_output);
    } else {
      output_changes.resize(0);
      new_output.resize(0);
    }
  }

  void UpdateLookup(const Eigen::VectorXd &input,
                    const std::vector<int> &tochange,
                    const std::vector<double> &newconf, LookupType &theta,
                    const VectorType & /*output*/,
                    std::vector<int> &output_changes,
                    VectorType &new_output) override {
    const int num_of_changes = tochange.size();
    if (num_of_changes > 0) {
      UpdateTheta(input, tochange, newconf, theta);
      output_changes.resize(out_size_);
      new_output.resize(out_size_);
      Forward(theta, new_output);
    } else {
      output_changes.resize(0);
      new_output.resize(0);
    }
  }

  // Feedforward
  void Forward(const VectorType &prev_layer_output, LookupType &theta,
               VectorType &output) override {
    LinearTransformation(prev_layer_output, theta);
    NonLinearTransformation(theta, output);
  }

  // Feedforward Using lookup
  void Forward(const LookupType &theta, VectorType &output) override {
    // Apply activation function
    NonLinearTransformation(theta, output);
  }

  // Applies the linear transformation
  inline void LinearTransformation(const VectorType &input, LookupType &theta) {
    theta[0] = bias_;
    theta[0].noalias() += weight_.transpose() * input;
  }

  // Applies the nonlinear transformation
  inline void NonLinearTransformation(const LookupType &theta,
                                      VectorType &output) {
    activation_(theta[0], output);
  }

  // Updates theta given the input v, the change in the input (input_changes and
  // prev_input)
  inline void UpdateTheta(const VectorType &v,
                          const std::vector<int> &input_changes,
                          const VectorType &new_input, LookupType &theta) {
    const int num_of_changes = input_changes.size();
    for (int s = 0; s < num_of_changes; s++) {
      const int sf = input_changes[s];
      theta[0] += weight_.row(sf) * (new_input(s) - v(sf));
    }
  }

  // Updates theta given the previous input prev_input and the change in the
  // input (tochange and  newconf)
  inline void UpdateTheta(const VectorType &prev_input,
                          const std::vector<int> &tochange,
                          const std::vector<double> &newconf,
                          LookupType &theta) {
    const int num_of_changes = tochange.size();
    for (int s = 0; s < num_of_changes; s++) {
      const int sf = tochange[s];
      theta[0] += weight_.row(sf) * (newconf[s] - prev_input(sf));
    }
  }

  // Computes derivative.
  void Backprop(const VectorType &prev_layer_output,
                const VectorType &this_layer_output,
                const LookupType &this_layer_theta, const VectorType &dout,
                VectorType &din, VectorType &der, int start_idx) override {
    // After forward stage, m_z contains z = W' * in + b
    // Now we need to calculate d(L) / d(z) = [d(a) / d(z)] * [d(L) / d(a)]
    // d(L) / d(a) is computed in the next layer, contained in next_layer_data
    // The Jacobian matrix J = d(a) / d(z) is determined by the activation
    // function
    VectorType dLz(out_size_);
    activation_.ApplyJacobian(this_layer_theta[0], this_layer_output, dout,
                              dLz);

    // Now dLz contains d(L) / d(z)
    // Derivative for bias, d(L) / d(b) = d(L) / d(z)
    int k = start_idx;

    if (usebias_) {
      Eigen::Map<VectorType> der_b{der.data() + k, out_size_};

      der_b.noalias() = dLz;
      k += out_size_;
    }

    // Derivative for weights, d(L) / d(W) = [d(L) / d(z)] * in'
    Eigen::Map<MatrixType> der_w{der.data() + k, in_size_, out_size_};

    der_w.noalias() = prev_layer_output * dLz.transpose();

    // Compute d(L) / d_in = W * [d(L) / d(z)]
    din.noalias() = weight_ * dLz;
  }
};
}  // namespace netket

#endif


#include <tiny_math_types.h>

#define BOOST_AUTO_TEST_MAIN
#include <boost/test/auto_unit_test.hpp>
#include <boost/test/unit_test_suite.hpp>
#include <boost/test/floating_point_comparison.hpp>
#include <boost/test/test_tools.hpp>

BOOST_AUTO_TEST_SUITE(tiny_quaternion_rotate);

BOOST_AUTO_TEST_CASE(simple_test)
{
  typedef tiny::MathTypes<double> math_types;

  typedef math_types::value_traits     value_traits;
  typedef math_types::real_type        T;
  typedef math_types::vector3_type     V;
  typedef math_types::quaternion_type  Q;

  BOOST_CHECK(  Q::accessor::stride()     == 1);
  BOOST_CHECK(  Q::accessor::padding()    == 0);
  BOOST_CHECK(  Q::accessor::J_padded()   == 4);
  BOOST_CHECK(  Q::accessor::allocsize()  == 4);

  {
    // Set up a test rotation
    Q q;
    T const phi = value_traits::pi_half();
    V const m = V::make( 1.0, 0.0, 0.0 );
    q = Q::Ru( phi, m);
    BOOST_CHECK_CLOSE( tiny::norm( q ), value_traits::one(), 0.01 );
    V const n = V::make( 0.0, 1.0, 0.0 );
    // Try to rotate n-vector
    V const k = tiny::rotate( q, n);
    // Test if result is what we expect
    BOOST_CHECK( fabs( k(0) ) < 10e-10 );
    BOOST_CHECK( fabs( k(1) ) < 10e-10 );
    BOOST_CHECK( fabs( 1.0 - k(2) ) < 10e-10 );
  }
  {
    // Set up a test rotation
    Q q;
    T const phi = value_traits::pi_half();
    V const m = V::make( 1.0, 0.0, 0.0 );
    q = Q::Ru( phi, m);
    BOOST_CHECK_CLOSE( tiny::norm( q ), value_traits::one(), 0.01 );
    V const n = V::make( 0.0, 0.0, 1.0 );
    // Try to rotate n-vector
    V const k = tiny::rotate( q, n);
    // Test if result is what we expect
    BOOST_CHECK( fabs( k(0) ) < 10e-10 );
    BOOST_CHECK( fabs( 1.0 + k(1) ) < 10e-10 );
    BOOST_CHECK( fabs( k(2) ) < 10e-10 );
  }
  {
    // Set up a test rotation
    Q q;
    T const phi = value_traits::pi_half();
    V const m = V::make( 1.0, 0.0, 0.0 );
    q = Q::Ru( phi, m);
    BOOST_CHECK_CLOSE( tiny::norm( q ), value_traits::one(), 0.01 );
    V const n = V::make( 0.0, -1.0, 0.0 );
    // Try to rotate n-vector
    V const k = tiny::rotate( q, n);
    // Test if result is what we expect
    BOOST_CHECK( fabs( k(0) ) < 10e-10 );
    BOOST_CHECK( fabs( k(1) ) < 10e-10 );
    BOOST_CHECK( fabs( 1.0 + k(2) ) < 10e-10 );
  }
  {
    // Set up a test rotation
    Q q;
    T const phi = value_traits::pi_half();
    V const m = V::make( 1.0, 0.0, 0.0 );
    q = Q::Ru( phi, m);
    BOOST_CHECK_CLOSE( tiny::norm( q ), value_traits::one(), 0.01 );
    V const n = V::make( 0.0, 0.0, -1.0 );
    // Try to rotate n-vector
    V const k = tiny::rotate( q, n);
    // Test if result is what we expect
    BOOST_CHECK( fabs( k(0) ) < 10e-10 );
    BOOST_CHECK( fabs( 1.0 - k(1) ) < 10e-10 );
    BOOST_CHECK( fabs( k(2) ) < 10e-10 );
  }
  {
    // Set up a test rotation
    Q q;
    T const phi = value_traits::pi_half();
    V const m = V::make( 0.0, 1.0, 0.0 );
    q = Q::Ru( phi, m);
    BOOST_CHECK_CLOSE( tiny::norm( q ), value_traits::one(), 0.01 );
    V const n = V::make( 1.0, 0.0, 0.0 );
    // Try to rotate n-vector
    V const k = tiny::rotate( q, n);
    // Test if result is what we expect
    BOOST_CHECK( fabs( k(0) ) < 10e-10 );
    BOOST_CHECK( fabs( k(1) ) < 10e-10 );
    BOOST_CHECK( fabs( 1.0 + k(2) ) < 10e-10 );
  }
  {
    // Set up a test rotation
    Q q;
    T const phi = value_traits::pi_half();
    V const m = V::make( 0.0, 0.0, 1.0 );
    q = Q::Ru( phi, m);
    BOOST_CHECK_CLOSE( tiny::norm( q ), value_traits::one(), 0.01 );
    V const n = V::make( 1.0, 0.0, 0.0 );
    // Try to rotate n-vector
    V const k = tiny::rotate( q, n);
    // Test if result is what we expect
    BOOST_CHECK( fabs( k(0) ) < 10e-10 );
    BOOST_CHECK( fabs( 1.0 - k(1) ) < 10e-10 );
    BOOST_CHECK( fabs( k(2) ) < 10e-10 );
  }
}

BOOST_AUTO_TEST_SUITE_END();


#include "eigen-runtime.h"
#include <Eigen/Core>
#include <Eigen/Sparse>
#include <Eigen/LU>

using namespace Eigen;

template <class T>
Map< Matrix<T,Dynamic,Dynamic> > matrix(void* p, int r, int c) {
    return Map< Matrix<T,Dynamic,Dynamic> >((T*)p, r, c);
}

template <class T>
Map< Matrix<T,Dynamic,Dynamic> > matrix(const void* p, int r, int c) {
    return Map< Matrix<T,Dynamic,Dynamic> >((T*)p, r, c);
}

template <class T>
Map< SparseMatrix<T> > smatrix(void* val, void* outer,
    void* inner, int r, int c, int s) {
    return Map< SparseMatrix<T> >(r, c, s, (int*)outer, (int*)inner, (T*)val);
}

template <class T>
Map<SparseMatrix<T> > smatrix(const void* val, const void* outer,
    const void* inner, int r, int c, int s) {
    return Map< SparseMatrix<T> >(r, c, s, (int*)outer, (int*)inner, (T*)val);
}

// Matrix product
template <class T>
RET dd_mul( void* p, int r, int c,
    const void* p1, int r1, int c1,
    const void* p2, int r2, int c2)
{
    matrix<T>(p,r,c) = matrix<T>(p1,r1,c1) * matrix<T>(p2,r2,c2);
    return 0;
}
API(dd_mul, (int code,
    void* p, int r, int c,
    const void* p1, int r1, int c1,
    const void* p2, int r2, int c2), (p,r,c,p1,r1,c1,p2,r2,c2));

template <class T>
RET ds_mul( void* p, int r, int c,
    const void* p1, int r1, int c1,
    const void* val, const void* outer, const void* inner,
    int r2, int c2, int s)
{
    matrix<T>(p,r,c) = matrix<T>(p1,r1,c1) * smatrix<T>(val, outer, inner, r2, c2, s);
    return 0;
}
API(ds_mul, (int code,
    void* p, int r, int c,
    const void* p1, int r1, int c1,
    const void* val, const void* outer, const void* inner,
    int r2, int c2, int s), (p,r,c,p1,r1,c1,val,outer,inner,r2,c2,s));

template <class T>
RET sd_mul( void* p, int r, int c,
    const void* val, const void* outer, const void* inner,
    int r2, int c2, int s,
    const void* p1, int r1, int c1)
{
    matrix<T>(p,r,c) = smatrix<T>(val, outer, inner, r2, c2, s) * matrix<T>(p1,r1,c1);
    return 0;
}
API(sd_mul, (int code,
    void* p, int r, int c,
    const void* val, const void* outer, const void* inner, int r2, int c2, int s,
    const void* p1, int r1, int c1), (p,r,c,val,outer,inner,r2,c2,s,p1,r1,c1));

template <class T>
RET ss_mul( void** v, void* o, void** i, int r, int c, int* s,
    const void* v1, const void* o1, const void* i1, int r1, int c1, int s1,
    const void* v2, const void* o2, const void* i2, int r2, int c2, int s2)
{
    typedef Map<SparseMatrix<T> > MapSparseMatrix;
    MapSparseMatrix a(r1, c1, s1, (int*)o1, (int*)i1, (T*)v1);
    MapSparseMatrix b(r2, c2, s2, (int*)o2, (int*)i2, (T*)v2);
    SparseMatrix<T> M = (a * b).pruned();

    memcpy(o, M.outerIndexPtr(), (c+1) * sizeof(int));

    *s = M.nonZeros();
    T* p1 = (T*) malloc((*s) * sizeof(T));
    memcpy(p1, M.valuePtr(), (*s) * sizeof(T));
    *v = p1;
    int* p2 = (int*) malloc((*s) * sizeof(int));
    memcpy(p2, M.innerIndexPtr(), (*s) * sizeof(int));
    *i = p2;
    return 0;
}
API(ss_mul, (int code,
    void** v, void* o, void** i, int r, int c, int* s,
    const void* v1, const void* o1, const void* i1, int r1, int c1, int s1,
    const void* v2, const void* o2, const void* i2, int r2, int c2, int s2),
    (v,o,i,r,c,s,v1,o1,i1,r1,c1,s1,v2,o2,i2,r2,c2,s2));

template <class T>
RET ss_plus( void** v, void* o, void** i, int r, int c, int* s,
    const void* v1, const void* o1, const void* i1, int r1, int c1, int s1,
    const void* v2, const void* o2, const void* i2, int r2, int c2, int s2)
{
    typedef Map<SparseMatrix<T> > MapSparseMatrix;
    MapSparseMatrix a(r1, c1, s1, (int*)o1, (int*)i1, (T*)v1);
    MapSparseMatrix b(r2, c2, s2, (int*)o2, (int*)i2, (T*)v2);
    SparseMatrix<T> M = a + b;

    memcpy(o, M.outerIndexPtr(), (c+1) * sizeof(int));

    *s = M.nonZeros();
    T* p1 = (T*) malloc((*s) * sizeof(T));
    memcpy(p1, M.valuePtr(), (*s) * sizeof(T));
    *v = p1;
    int* p2 = (int*) malloc((*s) * sizeof(int));
    memcpy(p2, M.innerIndexPtr(), (*s) * sizeof(int));
    *i = p2;
    return 0;
}
API(ss_plus, (int code,
    void** v, void* o, void** i, int r, int c, int* s,
    const void* v1, const void* o1, const void* i1, int r1, int c1, int s1,
    const void* v2, const void* o2, const void* i2, int r2, int c2, int s2),
    (v,o,i,r,c,s,v1,o1,i1,r1,c1,s1,v2,o2,i2,r2,c2,s2));

template <class T>
RET ss_cmul( void** v, void* o, void** i, int r, int c, int* s,
    const void* v1, const void* o1, const void* i1, int r1, int c1, int s1,
    const void* v2, const void* o2, const void* i2, int r2, int c2, int s2)
{
    typedef Map<SparseMatrix<T> > MapSparseMatrix;
    MapSparseMatrix a(r1, c1, s1, (int*)o1, (int*)i1, (T*)v1);
    MapSparseMatrix b(r2, c2, s2, (int*)o2, (int*)i2, (T*)v2);

    SparseMatrix<T> M = a.cwiseProduct(b);

    memcpy(o, M.outerIndexPtr(), (c+1) * sizeof(int));

    *s = M.nonZeros();
    T* p1 = (T*) malloc((*s) * sizeof(T));
    memcpy(p1, M.valuePtr(), (*s) * sizeof(T));
    *v = p1;
    int* p2 = (int*) malloc((*s) * sizeof(int));
    memcpy(p2, M.innerIndexPtr(), (*s) * sizeof(int));
    *i = p2;
    return 0;
}
API(ss_cmul, (int code,
    void** v, void* o, void** i, int r, int c, int* s,
    const void* v1, const void* o1, const void* i1, int r1, int c1, int s1,
    const void* v2, const void* o2, const void* i2, int r2, int c2, int s2),
    (v,o,i,r,c,s,v1,o1,i1,r1,c1,s1,v2,o2,i2,r2,c2,s2));

template <class T>
RET sd_plus( void* p, int r, int c,
    const void* val, const void* outer, const void* inner,
    int r2, int c2, int s,
    const void* p1, int r1, int c1)
{
    typedef Map< Matrix<T,Dynamic,Dynamic> > MapMatrix;
    MapMatrix x((T*)p, r, c);
    x = matrix<T>(p1,r1,c1);
    x += smatrix<T>(val, outer, inner, r2, c2, s);
    return 0;
}
API(sd_plus, (int code,
    void* p, int r, int c,
    const void* val, const void* outer, const void* inner, int r2, int c2, int s,
    const void* p1, int r1, int c1), (p,r,c,val,outer,inner,r2,c2,s,p1,r1,c1));

#define UNOP(name) \
extern "C" RET __attribute__((noinline)) eigen_##name(int code, void* p, int r, int c, const void* p1, int r1, int c1) {\
        GUARD_START\
        switch (code) {\
            case 0: matrix<T0>(p,r,c) = matrix<T0>(p1,r1,c1).name(); break;\
            case 1: matrix<T1>(p,r,c) = matrix<T1>(p1,r1,c1).name(); break;\
            case 2: matrix<T2>(p,r,c) = matrix<T2>(p1,r1,c1).name(); break;\
            case 3: matrix<T3>(p,r,c) = matrix<T3>(p1,r1,c1).name(); break;\
        }\
        GUARD_END\
    }

UNOP(inverse);


#include <boost/random/student_t_distribution.hpp>


/**
 * Copyright (c) 2018, University Osnabrück
 * All rights reserved.
 *
 * Redistribution and use in source and binary forms, with or without
 * modification, are permitted provided that the following conditions are met:
 *     * Redistributions of source code must retain the above copyright
 *       notice, this list of conditions and the following disclaimer.
 *     * Redistributions in binary form must reproduce the above copyright
 *       notice, this list of conditions and the following disclaimer in the
 *       documentation and/or other materials provided with the distribution.
 *     * Neither the name of the University Osnabrück nor the
 *       names of its contributors may be used to endorse or promote products
 *       derived from this software without specific prior written permission.
 *
 * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND
 * ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED
 * WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
 * DISCLAIMED. IN NO EVENT SHALL University Osnabrück BE LIABLE FOR ANY
 * DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES
 * (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;
 * LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND
 * ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
 * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS
 * SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
 */

/**
 * EigenSVDPointAlign.cpp
 *
 *  @date Feb 21, 2014
 *  @author Thomas Wiemann
 */

#include <Eigen/SVD>

using namespace Eigen;

namespace lvr2
{

template<typename T, typename PointT>
T EigenSVDPointAlign<T, PointT>::alignPoints(
    SLAMScanPtr scan,
    Point3** neighbors,
    const Vec3& centroid_m,
    const Vec3& centroid_d,
    Mat4& align) const
{
    T error = 0.0;
    size_t pairs = 0;

    // Fill H matrix
    Mat3 H = Matrix3d::Zero();

    for (size_t i = 0; i < scan->numPoints(); i++)
    {
        if (neighbors[i] == nullptr)
        {
            continue;
        }

        Vec3 m = neighbors[i]->template cast<T>() - centroid_m;
        Vec3 d = scan->point(i).template cast<T>() - centroid_d;

        error += (m - d).squaredNorm();
        pairs++;

        // same as "H += m * d.transpose();" but faster
        for (int j = 0; j < 3; j++)
        {
            for (int k = 0; k < 3; k++)
            {
                H(j, k) += d[j] * m[k];
            }
        }
    }

    error = sqrt(error / (T)pairs);

    JacobiSVD<Mat3> svd(H, ComputeFullU | ComputeFullV);

    Mat3 U = svd.matrixU();
    Mat3 V = svd.matrixV();

    Mat3 R = V * U.transpose();

    align = Mat4::Identity();
    align.template block<3, 3>(0, 0) = R;

    // Calculate translation
    Vec3 translation = centroid_m - R * centroid_d;
    align.template block<3, 1>(0, 3) = translation;

    return error;
}

template<typename T, typename PointT>
T EigenSVDPointAlign<T, PointT>::alignPoints(
    PointPairVector& pairs,
    const Vec3& centroid_m,
    const Vec3& centroid_d,
    Mat4& align) const
{
    T error = 0.0;
    size_t n = pairs.size();

    // Fill H matrix
    Mat3 H = Mat3::Zero();

    for (size_t i = 0; i < n; i++)
    {
        Vec3 m = pairs[i].first.template cast<T>() - centroid_m;
        Vec3 d = pairs[i].second.template cast<T>() - centroid_d;

        error += (m - d).squaredNorm();

        // same as "H += m * d.transpose();" but faster
        for (int j = 0; j < 3; j++)
        {
            for (int k = 0; k < 3; k++)
            {
                H(j, k) += d[j] * m[k];
            }
        }
    }

    error = sqrt(error / (T)n);

    JacobiSVD<Mat3> svd(H, ComputeFullU | ComputeFullV);

    Mat3 U = svd.matrixU();
    Mat3 V = svd.matrixV();

    Mat3 R = V * U.transpose();

    align = Mat4::Identity();
    align.template block<3, 3>(0, 0) = R;

    // Calculate translation
    Vec3 translation = centroid_m - R * centroid_d;
    align.template block<3, 1>(0, 3) = translation;

    return error;
}

} // namespace lvr2


/* Copyright (C) 2021 Intel Corporation
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *  http://www.apache.org/licenses/LICENSE-2.0
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

#include "bgv_common.h"

#include <NTL/BasicThreadPool.h>
#include <helib/helib.h>
#include "../src/PrimeGenerator.h" // Private header

#include <benchmark/benchmark.h>
#include <iostream>

namespace {

static void helib_fft_forward(benchmark::State& state, Meta& meta)
{
  NTL::SetNumThreads(1);

  long N = meta.data->ea.size();
  long m = meta.data->context.getM();
  auto zms = meta.data->context.getZMStar();

  helib::PrimeGenerator prime_generator(49, m);

  long q = prime_generator.next();
  helib::Cmodulus cmod(zms, q, 0);

  NTL::ZZX poly;
  poly.SetLength(N);
  for (long i = 0; i < N; ++i)
    poly[i] = i;

  NTL::vec_long transformed;
  for (auto _ : state)
    cmod.FFT(transformed, poly);
}

static void helib_fft_inverse(benchmark::State& state, Meta& meta)
{
  NTL::SetNumThreads(1);

  long N = meta.data->ea.size();
  long m = meta.data->context.getM();
  auto zms = meta.data->context.getZMStar();

  helib::PrimeGenerator prime_generator(49, m);

  long q = prime_generator.next();
  helib::Cmodulus cmod(zms, q, 0);

  NTL::zz_pX inverse;
  inverse.SetLength(N);

  NTL::vec_long transformed(NTL::INIT_SIZE, N);
  for (long i = 0; i < N; ++i)
    transformed[i] = i;

  for (auto _ : state)
    cmod.iFFT(inverse, transformed);
}

Meta fn;
Params hexl_F4_params(/*m=*/16384, /*p=*/65537, /*r=*/1, /*qbits=*/5800);
HE_BENCH_CAPTURE(helib_fft_forward, hexl_F4_params, fn); //->Iterations(200);
HE_BENCH_CAPTURE(helib_fft_inverse, hexl_F4_params, fn); //->Iterations(200);

Params hexl_F3_params(/*m=*/16, /*p=*/257, /*r=*/1, /*qbits=*/5800);
HE_BENCH_CAPTURE(helib_fft_forward, hexl_F3_params, fn); //->Iterations(200);
HE_BENCH_CAPTURE(helib_fft_inverse, hexl_F3_params, fn); //->Iterations(200);

Params hexl_F3d2_params(/*m=*/512, /*p=*/257, /*r=*/1, /*qbits=*/5800);
HE_BENCH_CAPTURE(helib_fft_forward, hexl_F3d2_params, fn); //->Iterations(200);
HE_BENCH_CAPTURE(helib_fft_inverse, hexl_F3d2_params, fn); //->Iterations(200);

} // namespace