/****************************************************************************
 * Copyright (c) 2017-2021 by the ArborX authors                            *
 * All rights reserved.                                                     *
 *                                                                          *
 * This file is part of the ArborX library. ArborX is                       *
 * distributed under a BSD 3-clause license. For the licensing terms see    *
 * the LICENSE file in the top-level directory.                             *
 *                                                                          *
 * SPDX-License-Identifier: BSD-3-Clause                                    *
 ****************************************************************************/

#include <boost/test/unit_test.hpp>

#include <algorithm>
#include <iostream>
#include <random>
#include <tuple>

#include "Search_UnitTestHelpers.hpp"
// clang-format off
#include "ArborXTest_TreeTypeTraits.hpp"
// clang-format on

// We need a forward-declaration for NVCC, see below.
namespace ArborX
{
namespace Experimental
{
template <int k>
struct KDOP;
}
} // namespace ArborX

BOOST_AUTO_TEST_SUITE(ManufacturedSolution)

namespace tt = boost::test_tools;

BOOST_AUTO_TEST_CASE_TEMPLATE(structured_grid, TreeTypeTraits,
                              TreeTypeTraitsList)
{
  using Tree = typename TreeTypeTraits::type;
  using ExecutionSpace = typename TreeTypeTraits::execution_space;
  using DeviceType = typename TreeTypeTraits::device_type;

  // FIXME_NVCC we see inexplainable test failures with NVCC and KDOP<18> and
  // KDOP<26> here.
#ifdef __NVCC__
  if (std::is_same<typename Tree::bounding_volume_type,
                   ArborX::Experimental::KDOP<18>>::value ||
      std::is_same<typename Tree::bounding_volume_type,
                   ArborX::Experimental::KDOP<26>>::value)
    return;
#endif

  float Lx = 100.0;
  float Ly = 100.0;
  float Lz = 100.0;
  int nx = 11;
  int ny = 11;
  int nz = 11;
  int n = nx * ny * nz;
  float hx = Lx / (nx - 1);
  float hy = Ly / (ny - 1);
  float hz = Lz / (nz - 1);

  std::function<int(int, int, int)> ind = [nx, ny](int i, int j, int k) {
    return i + j * nx + k * (nx * ny);
  };

  Kokkos::View<ArborX::Box *, DeviceType> bounding_boxes("bounding_boxes", n);
  auto bounding_boxes_host = Kokkos::create_mirror_view(bounding_boxes);

  for (int i = 0; i < nx; ++i)
    for (int j = 0; j < ny; ++j)
      for (int k = 0; k < nz; ++k)
      {
        ArborX::Point p{{i * hx, j * hy, k * hz}};
        bounding_boxes_host[ind(i, j, k)] = {p, p};
      }
  Kokkos::deep_copy(bounding_boxes, bounding_boxes_host);

  Tree const tree(ExecutionSpace{}, bounding_boxes);

  std::vector<int> offset_ref(n + 1);
  std::vector<int> indices_ref;

  // (i) use the same objects for the queries as the objects used in BVH
  // construction
  // i-2  i-1  i  i+1
  //
  //  o    o   o   o   j+1
  //          ---
  //  o    o | x | o   j
  //          ---
  //  o    o   o   o   j-1
  //
  //  o    o   o   o   j-2
  //
  indices_ref.resize(n);
  std::iota(offset_ref.begin(), offset_ref.end(), 0);
  std::iota(indices_ref.begin(), indices_ref.end(), 0);

  Kokkos::View<decltype(ArborX::intersects(ArborX::Box{})) *, DeviceType>
      queries("queries", n);
  Kokkos::parallel_for("fill_queries",
                       Kokkos::RangePolicy<ExecutionSpace>(0, n),
                       KOKKOS_LAMBDA(int i) {
                         queries(i) = ArborX::intersects(bounding_boxes(i));
                       });
  ARBORX_TEST_QUERY_TREE(ExecutionSpace{}, tree, queries,
                         make_reference_solution(indices_ref, offset_ref));

  // (ii) use bounding boxes that intersects with first neighbors
  //
  // i-2  i-1  i  i+1
  //
  //  o    x---x---x   j+1
  //       |       |
  //  o    x   x   x   j
  //       |       |
  //  o    x---x---x   j-1
  //
  //  o    o   o   o   j-2
  //
  std::vector<std::set<int>> ref(n);
  for (int i = 0; i < nx; ++i)
    for (int j = 0; j < ny; ++j)
      for (int k = 0; k < nz; ++k)
      {
        int const index = ind(i, j, k);
        // bounding box around nodes of the structured grid will
        // intersect with neighboring nodes
        bounding_boxes_host[index] = {
            {{(i - 1) * hx, (j - 1) * hy, (k - 1) * hz}},
            {{(i + 1) * hx, (j + 1) * hy, (k + 1) * hz}}};
        // fill in reference solution to check against the collision
        // list computed during the tree traversal
        if ((i > 0) && (j > 0) && (k > 0))
          ref[index].emplace(ind(i - 1, j - 1, k - 1));
        if ((i > 0) && (k > 0))
          ref[index].emplace(ind(i - 1, j, k - 1));
        if ((i > 0) && (j < ny - 1) && (k > 0))
          ref[index].emplace(ind(i - 1, j + 1, k - 1));
        if ((i > 0) && (j > 0))
          ref[index].emplace(ind(i - 1, j - 1, k));
        if (i > 0)
          ref[index].emplace(ind(i - 1, j, k));
        if ((i > 0) && (j < ny - 1))
          ref[index].emplace(ind(i - 1, j + 1, k));
        if ((i > 0) && (j > 0) && (k < nz - 1))
          ref[index].emplace(ind(i - 1, j - 1, k + 1));
        if ((i > 0) && (k < nz - 1))
          ref[index].emplace(ind(i - 1, j, k + 1));
        if ((i > 0) && (j < ny - 1) && (k < nz - 1))
          ref[index].emplace(ind(i - 1, j + 1, k + 1));

        if ((j > 0) && (k > 0))
          ref[index].emplace(ind(i, j - 1, k - 1));
        if (k > 0)
          ref[index].emplace(ind(i, j, k - 1));
        if ((j < ny - 1) && (k > 0))
          ref[index].emplace(ind(i, j + 1, k - 1));
        if (j > 0)
          ref[index].emplace(ind(i, j - 1, k));
        if (true) // NOLINT
          ref[index].emplace(ind(i, j, k));
        if (j < ny - 1)
          ref[index].emplace(ind(i, j + 1, k));
        if ((j > 0) && (k < nz - 1))
          ref[index].emplace(ind(i, j - 1, k + 1));
        if (k < nz - 1)
          ref[index].emplace(ind(i, j, k + 1));
        if ((j < ny - 1) && (k < nz - 1))
          ref[index].emplace(ind(i, j + 1, k + 1));

        if ((i < nx - 1) && (j > 0) && (k > 0))
          ref[index].emplace(ind(i + 1, j - 1, k - 1));
        if ((i < nx - 1) && (k > 0))
          ref[index].emplace(ind(i + 1, j, k - 1));
        if ((i < nx - 1) && (j < ny - 1) && (k > 0))
          ref[index].emplace(ind(i + 1, j + 1, k - 1));
        if ((i < nx - 1) && (j > 0))
          ref[index].emplace(ind(i + 1, j - 1, k));
        if (i < nx - 1)
          ref[index].emplace(ind(i + 1, j, k));
        if ((i < nx - 1) && (j < ny - 1))
          ref[index].emplace(ind(i + 1, j + 1, k));
        if ((i < nx - 1) && (j > 0) && (k < nz - 1))
          ref[index].emplace(ind(i + 1, j - 1, k + 1));
        if ((i < nx - 1) && (k < nz - 1))
          ref[index].emplace(ind(i + 1, j, k + 1));
        if ((i < nx - 1) && (j < ny - 1) && (k < nz - 1))
          ref[index].emplace(ind(i + 1, j + 1, k + 1));
      }
  Kokkos::deep_copy(bounding_boxes, bounding_boxes_host);

  indices_ref.resize(0);
  for (int i = 0; i < n; ++i)
  {
    std::copy(ref[i].begin(), ref[i].end(), std::back_inserter(indices_ref));
    offset_ref[i + 1] = indices_ref.size();
  }

  Kokkos::parallel_for("fill_first_neighbors_queries",
                       Kokkos::RangePolicy<ExecutionSpace>(0, n),
                       KOKKOS_LAMBDA(int i) {
                         queries[i] = ArborX::intersects(bounding_boxes[i]);
                       });
  ARBORX_TEST_QUERY_TREE(ExecutionSpace{}, tree, queries,
                         make_reference_solution(indices_ref, offset_ref));

  // (iii) use random points
  //
  // i-1      i      i+1
  //
  //  o       o       o   j+1
  //         -------
  //        |       |
  //        |   +   |
  //  o     | x     | o   j
  //         -------
  //
  //  o       o       o   j-1
  //
  std::default_random_engine generator;
  std::uniform_int_distribution<> dist_x(0, nx - 1);
  std::uniform_int_distribution<> dist_y(0, ny - 1);
  std::uniform_int_distribution<> dist_z(0, nz - 1);
  std::uniform_real_distribution<double> dist_shift(-0.45, 0.45);

  // The generation is a bit convoluted to avoid a situation where a centroid
  // of a box falls on any of the lattice planes, resulting in multiple
  // collisions. As a workaround, we generate the centroids of boxes within
  // 0.45 x grid step of a lattice point.
  std::iota(offset_ref.begin(), offset_ref.end(), 0);
  indices_ref.resize(n);
  for (int l = 0; l < n; ++l)
  {
    auto const i = dist_x(generator);
    auto const j = dist_y(generator);
    auto const k = dist_z(generator);

    auto const x = (i + dist_shift(generator)) * hx;
    auto const y = (j + dist_shift(generator)) * hy;
    auto const z = (k + dist_shift(generator)) * hz;
    bounding_boxes_host(l) = {{{x - 0.5 * hx, y - 0.5 * hy, z - 0.5 * hz}},
                              {{x + 0.5 * hx, y + 0.5 * hy, z + 0.5 * hz}}};

    // Save the indices for the check
    indices_ref[l] = ind(i, j, k);
  }
  Kokkos::deep_copy(bounding_boxes, bounding_boxes_host);

  Kokkos::parallel_for("fill_first_neighbors_queries",
                       Kokkos::RangePolicy<ExecutionSpace>(0, n),
                       KOKKOS_LAMBDA(int i) {
                         queries[i] = ArborX::intersects(bounding_boxes[i]);
                       });
  ARBORX_TEST_QUERY_TREE(ExecutionSpace{}, tree, queries,
                         make_reference_solution(indices_ref, offset_ref));
}

BOOST_AUTO_TEST_SUITE_END()


#ifndef INCLUDED_scheme_nest_maps_SphereQuad_HH
#define INCLUDED_scheme_nest_maps_SphereQuad_HH

#include <Eigen/Dense>
#include "scheme/util/SimpleArray.hh"
#include <boost/static_assert.hpp>
#include <iostream>
#include <vector>

namespace scheme {
namespace nest {
namespace pmap {

	using namespace Eigen;
	using std::cout;
	using std::endl;

	namespace quadsphere_data {
		double * get_quadsphere_cells(){
			static double quadsphere_cells[9*6] = {
				 1, 0, 0,
	 			 0, 1, 0,
	 			 0, 0, 1,

				 1, 0, 0,
				 0, 0,-1,
				 0, 1, 0,

				 0, 0, 1,
				 0, 1, 0,
				-1, 0, 0,

				 1, 0, 0,
				 0, 0, 1,
				 0,-1, 0,

				 0, 0,-1,
				 0, 1, 0,
				 1, 0, 0,

				 1, 0, 0,
				 0,-1, 0,
				 0, 0,-1,
			};
			return quadsphere_cells;
		}
	}

	template<
		int DIM,
		class Value=util::SimpleArray<DIM,double>,
		class Index=uint64_t,
		class Float=double
	>
	struct SphereQuad {
		BOOST_STATIC_ASSERT_MSG(DIM==2,"SphereQuad DIM must be == 2");

		static int const DIMENSION = DIM;
		typedef Value ValueType ;
		typedef Float FloatType ;
		typedef Index IndexType ;
		typedef util::SimpleArray<DIM,Index> Indices;
		typedef util::SimpleArray<DIM,Float> Params;

		///@brief sets value to parameters without change
		///@return false iff invalid parameters
		bool params_to_value(
			Params const & params,
			Index cell_index,
			Index resl,
			Value & value
		) const {
			Map< Matrix<double,3,3> > const rot_to_cell(
				const_cast<double*>( quadsphere_data::get_quadsphere_cells() + 9*cell_index ));
			Vector3d vec( params[0]*2.0-1.0 , params[1]*2.0-1.0 , 1.0 );
			vec = rot_to_cell * vec;
			vec = vec / vec.norm();
			value[0] = vec[0];
			value[1] = vec[1];
			value[2] = vec[2];
			return true;
		}

		///@brief sets params/cell_index from value
		///@note necessary for value lookup and neighbor lookup
		bool value_to_params(
			Value const & value,
			Index resl,
			Params & params,
			Index & cell_index
		) const {
			Matrix<double,3,1> tmpval( value[0], value[1], value[2] );
			tmpval = tmpval / tmpval.norm();
			// cout << tmpval.transpose() << endl;

			Map< Matrix<double,3,6>, 0, OuterStride<9> > const
				centers( quadsphere_data::get_quadsphere_cells()+6 );
			// cout << centers << endl;
			Matrix<double,1,6> dotprods = centers.transpose() * tmpval;

			double highest = -9e9;
			for(size_t i = 0; i < 6; ++i){
				if( dotprods[i] > highest ){
					highest = dotprods[i];
					cell_index = i;
				}
			}
			// cout << "DOTS " << dotprods << " " << cell_index << " val: " << tmpval.transpose() << endl;
			Map< Matrix<double,3,3> > const rot_to_cell(
				quadsphere_data::get_quadsphere_cells() + 9*cell_index );
			tmpval = rot_to_cell.transpose() * tmpval;
			// cout << "unit cell val: " << tmpval.transpose() << endl;

			params[0] = (tmpval[0]/highest+1.0)/2.0;
			params[1] = (tmpval[1]/highest+1.0)/2.0;
			// cout << "parm: " << params.transpose() << endl;
			// cout << "PRM " << params.transpose() << " " << highest << endl;
			assert( 0.0 <= params[0] && params[0] <= 1.0);
			assert( 0.0 <= params[1] && params[1] <= 1.0);
			return true;
		}

		///@brief get parameter space repr of Value for particular cell
		///@note necessary only for neighbor lookup
		void value_to_params_for_cell(
			Value const & value,
			Params & params
		) const ;

		///@brief return the cell_index of neighboring cells within radius of value
		///@note delta parameter is in "Parameter Space"
		template<class OutIter>
		void get_neighboring_cells(Value const & value, Float radius, OutIter out) const;

		///@brief aka covering radius max distance from bin center to any value within bin
		Float bin_circumradius(Index resl) const {
			double const delta = 1.0/(double)(1ul<<resl);
			/// these are correct with a mapping spreading out the corners
			// Vector3d pworst = Vector3d(1,1,1) - delta*Vector3d(2.0,0,0);
			// Vector3d pNest0 = Vector3d(1,1,1) - delta*Vector3d(1.0,1.0,0);
			/// with no mapping, the center is probably the worst
			// cout << resl << " " << delta << endl;
			Vector3d pworst = Vector3d(1,0,0) - delta*Vector3d(0,1,1);
			Vector3d pNest0 = Vector3d(1,0,0);
			pworst = pworst / pworst.norm();
			pNest0 = pNest0 / pNest0.norm();
			return (pworst-pNest0).norm() * 1.001; // fudge factor
		}

		///@brief maximum distance from the bin center which must be within the bin
		Float bin_inradius(Index resl) const {
			// double const delta = 1.0/(double)(1ul<<resl);
			// Vec pworst = Vec(1,1,1);// - delta*Vec(2.0,0,0);
			// Vec pNest0 = Vec(1,1,1) - delta*Vec(1.0,1.0,0);
			// cube_to_sphere(pworst);
			// cube_to_sphere(pNest0);
			// pworst.normalize();
			// pNest0.normalize();
			// return pworst.distance(pNest0) * 0.5999; // should be half of curcumradius based on geometry
		}

		///@brief cell size
		Index num_cells() const { return 6; }
	};

}
}
}

#endif


#ifndef TEST_UNIT_MATH_PRIM_MAT_VECTORIZE_EXPECT_PRIM_VECTOR_VALUE_HPP
#define TEST_UNIT_MATH_PRIM_MAT_VECTORIZE_EXPECT_PRIM_VECTOR_VALUE_HPP

#include <math/prim/mat/vectorize/expect_val_eq.hpp>
#include <Eigen/Dense>
#include <vector>

template <typename F>
void expect_prim_vector_value() {
  using Eigen::VectorXd;
  using std::vector;

  std::vector<double> valid_inputs = F::valid_inputs();

  VectorXd b = VectorXd::Map(valid_inputs.data(), valid_inputs.size());
  VectorXd fb = F::template apply<VectorXd>(b);
  EXPECT_EQ(b.size(), fb.size());
  for (int i = 0; i < fb.size(); ++i) {
    expect_val_eq(F::apply_base(b(i)), fb(i));
  }

  vector<VectorXd> d;
  d.push_back(b);
  d.push_back(b);
  vector<VectorXd> fd = F::template apply<vector<VectorXd> >(d);
  EXPECT_EQ(d.size(), fd.size());
  for (size_t i = 0; i < fd.size(); ++i) {
    EXPECT_EQ(d[i].size(), fd[i].size());
    for (int j = 0; j < fd[i].size(); ++j) {
      expect_val_eq(F::apply_base(d[i](j)), fd[i](j));
    }
  }
}
#endif


// Copyright (c) 2016
// Author: Chrono Law
#include <std.hpp>
//using namespace std;

#include <boost/core/ignore_unused.hpp>
#include <boost/type_traits.hpp>

#include <boost/mpl/placeholders.hpp>
#include <boost/mpl/apply_wrap.hpp>
#include <boost/mpl/apply.hpp>

using namespace boost;
using namespace boost::mpl;

///////////////////////////////////////

void case1()
{
    typedef apply_wrap2<_1, int, char>::type t1;
    typedef apply<_3, int, char, float>::type t2;

    assert((is_same<t1, int>::value));
    assert((is_same<t2, float>::value));

}

///////////////////////////////////////

#include <boost/mpl/bind.hpp>

struct func1
{
    template<typename T1>
    struct apply
    {   typedef T1 type;};
};

void case2()
{
    typedef bind<func1, _1> f1;

    typedef apply<f1, int>::type data1;
    assert((is_same<data1, int>::value));

    typedef apply<bind<func1, float>>::type data2;
    assert((is_same<data2, float>::value));
}

///////////////////////////////////////

#include <boost/mpl/int.hpp>
#include <boost/mpl/arithmetic.hpp>
#include <boost/mpl/lambda.hpp>

void case3()
{
    typedef apply_wrap2<lambda<plus<_,_>>::type, int_<1>, int_<2>>::type data0;

    assert((data0::value == 3));

    typedef apply<plus<_,_>,
            int_<1>, int_<2>>::type data1;
    assert((data1::value == 3));

    typedef bind<lambda<plus<_,_>>::type, _1,_2> f1;
    typedef apply<f1, int_<1>, int_<2>>::type data2;
    assert((data2::value == 3));

}

///////////////////////////////////////

#include <boost/mpl/range_c.hpp>
#include <boost/mpl/copy.hpp>
#include <boost/mpl/inserter.hpp>

void case4()
{
    typedef range_c<int, 1, 11> rc;
    typedef mpl::copy<rc,
            mpl::inserter<int_<0>, plus<_,_> >
                >::type sum;
    std::cout << sum::value;

}

///////////////////////////////////////
#include <boost/mpl/vector.hpp>
#include <boost/mpl/count_if.hpp>
#include <boost/mpl/find_if.hpp>
#include <boost/mpl/sizeof.hpp>
#include <boost/mpl/front.hpp>
#include <boost/mpl/min_element.hpp>
#include <boost/mpl/max_element.hpp>

void case5()
{
    typedef mpl::vector<float, double,
        char, int, long> vec;

    assert((mpl::count_if<vec,
        is_float<_> >::value==2));
    assert((is_same<
            deref<mpl::find_if<vec,
            is_same<_, char>>::type>::type,
            char>::value));

    typedef lambda<less<sizeof_<_1>, sizeof_<_2>> > Comp;

    typedef deref<min_element<vec, Comp>::type>::type mint;
    assert((is_same<mint, char>::value));

    typedef deref<max_element<vec, Comp>::type>::type maxt;
    assert((is_same<maxt, double>::value));
}

///////////////////////////////////////

#include <boost/mpl/vector_c.hpp>
#include <boost/mpl/copy_if.hpp>
#include <boost/mpl/replace_if.hpp>
#include <boost/mpl/remove_if.hpp>
#include <boost/mpl/sort.hpp>
#include <boost/mpl/unique.hpp>
#include <boost/mpl/modulus.hpp>
#include <boost/mpl/equal_to.hpp>
#include <boost/mpl/greater.hpp>

void case6()
{
    typedef vector_c<int, 5,3,7,2,6,4,2> vec;

    typedef copy_if<vec,
        equal_to<modulus<_1, int_<2> > ,
                int_<0>> >::type vec2;

    typedef replace_if<vec2,
        equal_to<_1, int_<6> >, int_<10>>::type vec3;

    typedef remove_if<vec3,
        greater<_1, int_<5>>>::type vec4;

    typedef sort<vec4, less<_,_>>::type vec5;

    typedef unique<vec5, equal_to<_,_>>::type vec6;

    boost::ignore_unused<vec6>();
}

///////////////////////////////////////

int main()
{
    std::cout << "hello lambda" << std::endl;

    case1();
    case2();
    case3();
    case4();
    case5();
    case6();
}


#include <stan/math/prim/mat.hpp>
#include <stan/math/prim/scal.hpp>
#include <test/unit/math/prim/prob/vector_rng_test_helper.hpp>
#include <test/unit/math/prim/prob/util.hpp>
#include <boost/random/mersenne_twister.hpp>
#include <boost/math/distributions.hpp>
#include <gtest/gtest.h>
#include <limits>
#include <vector>

class BetaProportionTestRig : public VectorRealRNGTestRig {
 public:
  BetaProportionTestRig()
      : VectorRealRNGTestRig(10000, 10, {0.3, 0.4, 0.5, 0.6, 0.7}, {1, 2, 3},
                             {-2.5, -1.7, -0.1, 0.0}, {-3, -2, -1, 0},
                             {0.35, 0.5, 0.9, 1.7, 2.1, 4.1}, {1, 2, 3, 4},
                             {-2.7, -1.5, -0.5, 0.0}, {-3, -2, -1, 0}) {}

  template <typename T1, typename T2, typename T3, typename T_rng>
  auto generate_samples(const T1& mu, const T2& kappa, const T3&,
                        T_rng& rng) const {
    return stan::math::beta_proportion_rng(mu, kappa, rng);
  }

  std::vector<double> generate_quantiles(double mu, double kappa,
                                         double) const {
    // transform from location and precision parameterization
    // into shape1 (alpha) and shape2 (beta) parameterization
    double alpha = mu * kappa;
    double beta = kappa - alpha;
    std::vector<double> quantiles;
    double K = stan::math::round(2 * std::pow(N_, 0.4));
    boost::math::beta_distribution<> dist(alpha, beta);

    for (int i = 1; i < K; ++i) {
      double frac = i / K;
      quantiles.push_back(quantile(dist, frac));
    }
    quantiles.push_back(std::numeric_limits<double>::max());

    return quantiles;
  }
};

TEST(ProbDistributionsBetaProportion, errorCheck) {
  check_dist_throws_real_first_argument(BetaProportionTestRig());
}

TEST(ProbDistributionsBetaProportion, distributionTest) {
  check_quantiles_real_first_argument(BetaProportionTestRig());
}

TEST(ProbDistributionsBetaProportion, error_check) {
  boost::random::mt19937 rng;
  EXPECT_NO_THROW(stan::math::beta_proportion_rng(0.5, 3.0, rng));
  EXPECT_NO_THROW(stan::math::beta_proportion_rng(1e-10, 1e-10, rng));
  EXPECT_THROW(stan::math::beta_proportion_rng(2.0, -1.0, rng),
               std::domain_error);
  EXPECT_THROW(stan::math::beta_proportion_rng(-2.0, 1.0, rng),
               std::domain_error);
  EXPECT_THROW(stan::math::beta_proportion_rng(-2.0, -1.0, rng),
               std::domain_error);
  EXPECT_THROW(stan::math::beta_proportion_rng(stan::math::positive_infinity(),
                                               1.0, rng),
               std::domain_error);
  EXPECT_THROW(
      stan::math::beta_proportion_rng(2, stan::math::positive_infinity(), rng),
      std::domain_error);
}

TEST(ProbDistributionsBetaProportion, chiSquareGoodnessFitTest) {
  boost::random::mt19937 rng;
  int N = 10000;
  int K = stan::math::round(2 * std::pow(N, 0.4));

  double mu = 0.5;     // location
  double kappa = 3.0;  // precision

  std::vector<double> samples;
  for (int i = 0; i < N; ++i) {
    samples.push_back(stan::math::beta_proportion_rng(mu, kappa, rng));
  }

  // transform from location and precision parameterization
  // into shape1 (alpha) and shape2 (beta) parameterization
  double alpha = mu * kappa;
  double beta = kappa - alpha;

  // Generate quantiles from boost's beta distribution
  boost::math::beta_distribution<> dist(alpha, beta);
  std::vector<double> quantiles;
  for (int i = 1; i < K; ++i) {
    double frac = static_cast<double>(i) / K;
    quantiles.push_back(quantile(dist, frac));
  }
  quantiles.push_back(std::numeric_limits<double>::max());

  // Assert that they match
  assert_matches_quantiles(samples, quantiles, 1e-6);
}

TEST(ProbDistributionsBetaProportion, chiSquareGoodnessFitTest2) {
  boost::random::mt19937 rng;
  int N = 10000;
  int K = stan::math::round(2 * std::pow(N, 0.4));

  double mu = 0.3;     // location
  double kappa = 0.5;  // precision

  std::vector<double> samples;
  for (int i = 0; i < N; ++i) {
    samples.push_back(stan::math::beta_proportion_rng(mu, kappa, rng));
  }

  // transform from location and precision parameterization
  // into shape1 (alpha) and shape2 (beta) parameterization
  double alpha = mu * kappa;
  double beta = kappa - alpha;

  // Generate quantiles from boost's beta distribution
  boost::math::beta_distribution<> dist(alpha, beta);
  std::vector<double> quantiles;
  for (int i = 1; i < K; ++i) {
    double frac = static_cast<double>(i) / K;
    quantiles.push_back(quantile(dist, frac));
  }
  quantiles.push_back(std::numeric_limits<double>::max());

  // Assert that they match
  assert_matches_quantiles(samples, quantiles, 1e-6);
}


/* vim: set tabstop=4 expandtab shiftwidth=4 softtabstop=4: */

/**
 * \file libs/numeric/ublasx/test/rot90.cpp
 *
 * \brief Test suite for the \c rot90 operation.
 *
 * \author Marco Guazzone (marco.guazzone@gmail.com)
 *
 * <hr/>
 *
 * Copyright (c) 2011, Marco Guazzone
 *
 * 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)
 */

#include <boost/numeric/ublas/io.hpp>
#include <boost/numeric/ublas/matrix.hpp>
#include <boost/numeric/ublas/vector.hpp>
#include <boost/numeric/ublasx/operation/rot90.hpp>
#include <cmath>
#include <complex>
#include <cstddef>
#include "libs/numeric/ublasx/test/utils.hpp"

namespace ublas = ::boost::numeric::ublas;
namespace ublasx = ::boost::numeric::ublasx;

static const double tol = 1.0e-5;

BOOST_UBLASX_TEST_DEF( test_real_vector )
{
    BOOST_UBLASX_DEBUG_TRACE( "Test Case: Real - Vector" );

    typedef double value_type;
    typedef std::size_t size_type;
    typedef ublas::vector<value_type> vector_type;

    const size_type n(4);

    vector_type v(n);

    v(0) = 1;
    v(1) = 2;
    v(2) = 3;
    v(3) = 4;

    vector_type res;
    vector_type expect_res(n);

    // k=1
    res = ublasx::rot90(v, 1);
    BOOST_UBLASX_DEBUG_TRACE( "v = " << v );
    BOOST_UBLASX_DEBUG_TRACE( "rot90(v) = " << res );
    expect_res = v;
    BOOST_UBLASX_TEST_CHECK_VECTOR_CLOSE( res, expect_res, n, tol );

    // k=2
    res = ublasx::rot90(v, 2);
    BOOST_UBLASX_DEBUG_TRACE( "v = " << v );
    BOOST_UBLASX_DEBUG_TRACE( "rot90(v) = " << res );
    expect_res.clear();
    for (size_type i = 0; i < n; ++i)
    {
        expect_res(n-i-1) = v(i);
    }
    BOOST_UBLASX_TEST_CHECK_VECTOR_CLOSE( res, expect_res, n, tol );

    // k=3
    res = ublasx::rot90(v, 3);
    BOOST_UBLASX_DEBUG_TRACE( "v = " << v );
    BOOST_UBLASX_DEBUG_TRACE( "rot90(v) = " << res );
    expect_res.clear();
    for (size_type i = 0; i < n; ++i)
    {
        expect_res(n-i-1) = v(i);
    }
    BOOST_UBLASX_TEST_CHECK_VECTOR_CLOSE( res, expect_res, n, tol );

    // k=4
    res = ublasx::rot90(v, 4);
    BOOST_UBLASX_DEBUG_TRACE( "v = " << v );
    BOOST_UBLASX_DEBUG_TRACE( "rot90(v,4) = " << res );
    expect_res = v;
    BOOST_UBLASX_TEST_CHECK_VECTOR_CLOSE( res, expect_res, n, tol );
}

BOOST_UBLASX_TEST_DEF( test_real_matrix )
{
    BOOST_UBLASX_DEBUG_TRACE( "Test Case: Real - Matrix" );

    typedef double value_type;
    typedef std::size_t size_type;
    typedef ublas::matrix<value_type> matrix_type;

    const size_type nr(2);
    const size_type nc(3);

    matrix_type A(nr,nc);

    A(0,0) = 1; A(0,1) = 2; A(0,2) = 3;
    A(1,0) = 4; A(1,1) = 5; A(1,2) = 6;

    matrix_type R;
    matrix_type expect_R(nr,nc);

    // k=1
    R = ublasx::rot90(A);
    BOOST_UBLASX_DEBUG_TRACE( "A = " << A );
    BOOST_UBLASX_DEBUG_TRACE( "rot90(A,1) = " << R );
    expect_R.resize(nc, nr, false);
    expect_R.clear();
    for (size_type r = 0; r < nr; ++r)
    {
        for (size_type c = 0; c < nc; ++c)
        {
            expect_R(nc-c-1,r) = A(r,c);
        }
    }
    BOOST_UBLASX_TEST_CHECK_MATRIX_CLOSE( R, expect_R, nc, nr, tol );

    // k=2
    R = ublasx::rot90(A, 2);
    BOOST_UBLASX_DEBUG_TRACE( "A = " << A );
    BOOST_UBLASX_DEBUG_TRACE( "rot90(A,2) = " << R );
    expect_R.resize(nr, nc, false);
    expect_R.clear();
    for (size_type r = 0; r < nr; ++r)
    {
        for (size_type c = 0; c < nc; ++c)
        {
            expect_R(nr-r-1,nc-c-1) = A(r,c);
        }
    }
    BOOST_UBLASX_TEST_CHECK_MATRIX_CLOSE( R, expect_R, nr, nc, tol );

    // k=3
    R = ublasx::rot90(A, 3);
    BOOST_UBLASX_DEBUG_TRACE( "A = " << A );
    BOOST_UBLASX_DEBUG_TRACE( "rot90(A,3) = " << R );
    expect_R.resize(nc, nr, false);
    expect_R.clear();
    for (size_type r = 0; r < nr; ++r)
    {
        for (size_type c = 0; c < nc; ++c)
        {
            expect_R(c,nr-r-1) = A(r,c);
        }
    }
    BOOST_UBLASX_TEST_CHECK_MATRIX_CLOSE( R, expect_R, nc, nr, tol );

    // k=4
    R = ublasx::rot90(A, 4);
    BOOST_UBLASX_DEBUG_TRACE( "A = " << A );
    BOOST_UBLASX_DEBUG_TRACE( "rot90(A,4) = " << R );
    expect_R = A;
    BOOST_UBLASX_TEST_CHECK_MATRIX_CLOSE( R, expect_R, nr, nc, tol );
}

int main()
{

    BOOST_UBLASX_DEBUG_TRACE("Test Suite: 'rot90' operation");

    BOOST_UBLASX_TEST_BEGIN();

    BOOST_UBLASX_TEST_DO( test_real_vector );
//  BOOST_UBLASX_TEST_DO( test_complex_vector );
    BOOST_UBLASX_TEST_DO( test_real_matrix );
//  BOOST_UBLASX_TEST_DO( test_complex_matrix );

    BOOST_UBLASX_TEST_END();
}


#include <iostream>
#include <boost/random/mersenne_twister.hpp>
#include <boost/random/discrete_distribution.hpp>

using namespace std;

typedef boost::mt19937::result_type rv_t;

int main( void ) {

  boost::mt19937 gen;
  assert( sizeof( rv_t ) == 4 );
  uint64_t val = (uint64_t(gen()) << 32) + gen();
  cout << val << endl;

#if 0
  string ret;
  uint32_t num_bytes = 3;
  for (int i = 0; i < ( (num_bytes-1) / sizeof( rv_t ) +1 ) ; ++i) {
    const rv_t rv = gen();
    for( uint32_t b = 0; b < sizeof( rv_t ); ++b ) {
      ret.push_back( uint8_t(rv >>(b*8)) );
    }
  }
  assert( ret.size() >= num_bytes );
  ret.resize( num_bytes );
  for( string::const_iterator i = ret.begin(); i != ret.end(); ++i ) {
    cout << uint32_t( uint8_t( *i ) ) << endl;
  }
#endif
}


#include <basalt/utils/assert.h>
#include <basalt/utils/time_utils.hpp>

#include <fstream>
#include <iomanip>

#include <fmt/format.h>
#include <Eigen/Dense>
#include <nlohmann/json.hpp>

namespace basalt {

using namespace fmt::literals;

// compute the median of an eigen vector
// Note: Changes the order of elements in the vector!
// Note: For even sized vectors we don't return the mean of the middle two, but
// simply the second one as is.
template <class Scalar, int Rows>
Scalar median_non_const(Eigen::Matrix<Scalar, Rows, 1>& vec) {
  static_assert(Rows != 0);
  if constexpr (Rows < 0) {
    BASALT_ASSERT(vec.size() >= 1);
  }
  int n = vec.size() / 2;
  std::nth_element(vec.begin(), vec.begin() + n, vec.end());
  return vec(n);
}

template <class Scalar, int N>
Scalar variance(const Eigen::Matrix<Scalar, N, 1>& vec) {
  static_assert(N != 0);
  const Eigen::Matrix<Scalar, N, 1> centered = vec.array() - vec.mean();
  return centered.squaredNorm() / Scalar(vec.size());
}

ExecutionStats::Meta& ExecutionStats::add(const std::string& name,
                                          double value) {
  auto [it, new_item] = stats_.try_emplace(name);
  if (new_item) {
    order_.push_back(name);
    it->second.data_ = std::vector<double>();
  }
  std::get<std::vector<double>>(it->second.data_).push_back(value);
  return it->second;
}

ExecutionStats::Meta& ExecutionStats::add(const std::string& name,
                                          const Eigen::VectorXd& value) {
  auto [it, new_item] = stats_.try_emplace(name);
  if (new_item) {
    order_.push_back(name);
    it->second.data_ = std::vector<Eigen::VectorXd>();
  }
  std::get<std::vector<Eigen::VectorXd>>(it->second.data_).push_back(value);
  return it->second;
}

ExecutionStats::Meta& ExecutionStats::add(const std::string& name,
                                          const Eigen::VectorXf& value) {
  Eigen::VectorXd x = value.cast<double>();
  return add(name, x);
}

void ExecutionStats::merge_all(const ExecutionStats& other) {
  for (const auto& name : other.order_) {
    const auto& meta = other.stats_.at(name);
    std::visit(
        [&](auto& data) {
          for (auto v : data) {
            add(name, v);
          }
        },
        meta.data_);
    stats_.at(name).set_meta(meta);
  }
}

namespace {  // helper
// ////////////////////////////////////////////////////////////////////////////
// overloads for generic lambdas
// See also: https://stackoverflow.com/q/55087826/1813258
template <class... Ts>
struct overload : Ts... {
  using Ts::operator()...;
};
template <class... Ts>
overload(Ts...) -> overload<Ts...>;
}  // namespace

void ExecutionStats::merge_sums(const ExecutionStats& other) {
  for (const auto& name : other.order_) {
    const auto& meta = other.stats_.at(name);
    std::visit(overload{[&](const std::vector<double>& data) {
                          Eigen::Map<const Eigen::VectorXd> map(data.data(),
                                                                data.size());
                          add(name, map.sum());
                        },
                        [&](const std::vector<Eigen::VectorXd>& data) {
                          UNUSED(data);
                          // TODO: for now no-op
                        }},
               meta.data_);
    stats_.at(name).set_meta(meta);
  }
}

void ExecutionStats::print() const {
  for (const auto& name : order_) {
    const auto& meta = stats_.at(name);

    std::visit(
        overload{
            [&](const std::vector<double>& data) {
              Eigen::Map<const Eigen::VectorXd> map(data.data(), data.size());

              // create a copy for median computation
              Eigen::VectorXd vec = map;

              if (meta.format_ == "ms") {
                // convert seconds to milliseconds
                vec *= 1000;
              }

              int count = vec.size();
              // double sum = vec.sum();
              double mean = vec.mean();
              double stddev = std::sqrt(variance(vec));
              double max = vec.maxCoeff();
              double min = vec.minCoeff();

              // double median = median_non_const(vec);

              if (meta.format_ == "count") {
                std::cout << "{:20} ({:>4}):{: 8.1f}+-{:.1f} [{}, {}]\n"_format(
                    name, count, mean, stddev, min, max);
              } else if (meta.format_ != "none") {
                std::cout
                    << "{:20} ({:>4}):{: 8.2f}+-{:.2f} [{:.2f}, {:.2f}]\n"_format(
                           name, count, mean, stddev, min, max);
              }
            },
            [&](const std::vector<Eigen::VectorXd>& data) {
              int count = data.size();
              std::cout << "{:20} ({:>4})\n"_format(name, count);
            }},
        meta.data_);
  }
}

bool ExecutionStats::save_json(const std::string& path) const {
  using json = nlohmann::json;
  json result;

  for (const auto& name : order_) {
    const auto& meta = stats_.at(name);

    std::visit(
        overload{[&](const std::vector<double>& data) { result[name] = data; },
                 [&](const std::vector<Eigen::VectorXd>& data) {
                   std::vector<int> indices;
                   std::vector<double> values;
                   for (const auto& v : data) {
                     indices.push_back(int(values.size()));
                     values.insert(values.end(), v.begin(), v.end());
                   }
                   std::string name_values = std::string(name) + "__values";
                   std::string name_indices = std::string(name) + "__index";
                   result[name_indices] = indices;
                   result[name_values] = values;
                 }},
        meta.data_);
  }

  constexpr bool save_as_json = false;
  constexpr bool save_as_ubjson = true;

  // save json
  if (save_as_json) {
    std::ofstream ofs(path);

    if (!ofs.is_open()) {
      std::cerr << "Could not save ExecutionStats to {}.\n"_format(path);
      return false;
    }

    ofs << std::setw(4) << result;  //!< pretty printing
    // ofs << result;  //!< no pretty printing

    std::cout << "Saved ExecutionStats to {}.\n"_format(path);
  }

  // save ubjson
  if (save_as_ubjson) {
    std::string ubjson_path =
        path.substr(0, path.find_last_of('.')) + ".ubjson";
    std::ofstream ofs(ubjson_path, std::ios_base::binary);

    if (!ofs.is_open()) {
      std::cerr << "Could not save ExecutionStats to {}.\n"_format(ubjson_path);
      return false;
    }

    json::to_ubjson(result, ofs);

    std::cout << "Saved ExecutionStats to {}.\n"_format(ubjson_path);
  }

  return true;
}

}  // namespace basalt


/* Copyright 2016 Ramakrishnan Kannan */
#ifndef DISTNTF_DISTAUNTF_HPP_
#define DISTNTF_DISTAUNTF_HPP_

#include <armadillo>
#include <string>
#include <vector>
#include "common/distutils.hpp"
#include "common/ntf_utils.hpp"
#include "dimtree/ddt.hpp"
#include "distntf/distntfmpicomm.hpp"
#include "distntf/distntftime.hpp"

/**
 * Tensor A of size is M1 x M2 x... x Mn is distributed among
 * P1 x P2 x ... x Pn grid of P processors. That means, every
 * processor has (M1/P1) x (M2/P2) x ... x (Mn/Pn) tensor as
 * m_input_tensor. Similarly every process own a portion of the
 * factors as H(i,pi) of size (Mi/Pi) x k
 * and collects from its neighbours as H(i,p) as (Mi/P) x k
 * H(i,p) and m_input_tensor can perform local MTTKRP. The
 * local MTTKRP's are reduced scattered for local NNLS.
 */

// #define DISTNTF_VERBOSE 1

namespace planc {

#define TENSOR_LOCAL_DIM (m_input_tensor.dimensions())
#define TENSOR_LOCAL_NUMEL (m_input_tensor.numel())

class DistAUNTF {
 protected:
  // communication related variables
  const NTFMPICommunicator &m_mpicomm;
  // NLS solve sizes
  UVEC m_nls_sizes;
  UVEC m_nls_idxs;
  // local ncp factors
  NCPFactors m_local_ncp_factors;
  NCPFactors m_local_ncp_factors_t;
  // mttkrp related variables
  MAT *ncp_mttkrp_t;
  MAT *ncp_local_mttkrp_t;
  // hadamard of the global_grams
  MAT global_gram;

  virtual MAT update(int current_mode) = 0;

 private:
  const Tensor &m_input_tensor;
  NCPFactors m_gathered_ncp_factors;
  NCPFactors m_gathered_ncp_factors_t;
  // mttkrp related variables
  MAT *ncp_krp;
  // gram related variables.
  MAT factor_local_grams;    // U in the algorithm.
  MAT *factor_global_grams;  // G in the algorithm

  // NTF related variable.
  const unsigned int m_low_rank_k;
  const unsigned int m_modes;
  const algotype m_updalgo;
  const UVEC m_global_dims;
  const UVEC m_factor_local_dims;
  unsigned int m_num_it;
  unsigned int current_mode;
  FVEC m_regularizers;
  bool m_compute_error;
  bool m_enable_dim_tree;
  unsigned int m_current_it;
  double m_rel_error;

  // needed for acceleration algorithms.
  bool m_accelerated;
  std::vector<bool> m_stale_mttkrp;
  // stats
  DistNTFTime time_stats;

  // computing error related;
  double m_global_sqnorm_A;
  MAT hadamard_all_grams;

  DenseDimensionTree *kdt;

  /**
   * do the local syrk only for the current updated factor
   * and all reduce only for the current updated factor.
   * computes G^(current_mode)
   * @param[in] current_mode
   */
  void update_global_gram(const int current_mode) {
    // computing U
    MPITIC;  // gram
    // force a ssyrk instead of gemm.
    MAT H = m_local_ncp_factors.factor(current_mode);
    factor_local_grams = H.t() * H;

    double temp = MPITOC;  // gram
    this->time_stats.compute_duration(temp);
    this->time_stats.gram_duration(temp);
    factor_global_grams[current_mode].zeros();
    // Computing G.
    MPITIC;  // allreduce gram
    MPI_Allreduce(factor_local_grams.memptr(),
                  factor_global_grams[current_mode].memptr(),
                  this->m_low_rank_k * this->m_low_rank_k, MPI_DOUBLE, MPI_SUM,
                  MPI_COMM_WORLD);
    temp = MPITOC;  // allreduce gram
    applyReg(this->m_regularizers(current_mode * 2),
             this->m_regularizers(current_mode * 2 + 1),
             &(factor_global_grams[current_mode]));
    this->time_stats.communication_duration(temp);
    this->time_stats.allreduce_duration(temp);
  }

  /**
   * Apply L1 and L2 regularization on the gram of hadamards
   * @param[in] L1 - L1 Regularization parameters as a float
   * @param[in] L2 - L2 Regularization parameters as a float
   * @param[out] AtA - L1/L2 regularization applied gram matrix
   */

  void applyReg(float lambda_l2, float lambda_l1, MAT *AtA) {
    // Frobenius norm regularization
    if (lambda_l2 > 0) {
      MAT identity = arma::eye<MAT>(this->m_low_rank_k, this->m_low_rank_k);
      (*AtA) = (*AtA) + 2 * lambda_l2 * identity;
    }

    // L1 - norm regularization
    if (lambda_l1 > 0) {
      MAT onematrix = arma::ones<MAT>(this->m_low_rank_k, this->m_low_rank_k);
      (*AtA) = (*AtA) + 2 * lambda_l1 * onematrix;
    }
  }

  /**
   * This iterates over all grams leaving out current mode
   * and finds hadamard of the grams
   *
   * @param[in] current_mode.
   */
  void gram_hadamard(unsigned int current_mode) {
    global_gram.ones();
    MPITIC;  // gram hadamard
    for (unsigned int i = 0; i < m_modes; i++) {
      if (i != current_mode) {
        //%= element-wise multiplication
        global_gram %= factor_global_grams[i];
      }
    }
    double temp = MPITOC;  // gram hadamard
    this->time_stats.compute_duration(temp);
    this->time_stats.gram_duration(temp);
  }
  /**
   * factor matrices all gather only on the current updated factor
   * It performs the gather over the slice communicator
   * @param[in] current_mode
   */

  //
  void gather_ncp_factor(const int current_mode) {
    m_gathered_ncp_factors_t.factor(current_mode).zeros();
    // Had this comment for debugging memory corruption in all_gather
    // DISTPRINTINFO("::ncp_krp::" << ncp_krp[current_mode].memptr()
    //               << "::size::" << ncp_krp[current_mode].n_rows
    //               << "x" << ncp_krp[current_mode].n_cols
    //               << "::m_gathered_ncp_factors_t::"
    //               << m_gathered_ncp_factors_t.factor(current_mode).memptr()
    //               << "::diff from recvcnt::"
    //               << m_gathered_ncp_factors_t.factor(current_mode).memptr()
    //               - recvcnt * 8);

    MPI_Comm current_slice_comm = this->m_mpicomm.slice(current_mode);
    int slice_size;
    MPI_Comm_size(current_slice_comm, &slice_size);

    // int sendcnt = m_local_ncp_factors.factor(current_mode).n_elem;
    int sendcnt = m_nls_sizes[current_mode] * m_low_rank_k;

    // int recvcnt = m_local_ncp_factors.factor(current_mode).n_elem;
    std::vector<int> recvgathercnt(slice_size, 0);
    std::vector<int> recvgatherdispl(slice_size, 0);

    int dimsize = m_factor_local_dims[current_mode];
    for (int i = 0; i < slice_size; i++) {
      recvgathercnt[i] = itersplit(dimsize, slice_size, i) * m_low_rank_k;
      recvgatherdispl[i] = startidx(dimsize, slice_size, i) * m_low_rank_k;
    }

#ifdef DISTNTF_VERBOSE
    MPI_Comm current_fiber_comm = this->m_mpicomm.fiber(current_mode);
    int fiber_size;

    MPI_Comm_size(current_fiber_comm, &fiber_size);
    MPI_Comm_size(current_slice_comm, &slice_size);
    DISTPRINTINFO("::current_mode::"
                  << current_mode << "::fiber comm size::" << fiber_size
                  << "::my_global_rank::" << MPI_RANK << "::my_slice_rank::"
                  << this->m_mpicomm.slice_rank(current_mode)
                  << "::my_fiber_rank::"
                  << this->m_mpicomm.fiber_rank(current_mode)
                  << "::sendcnt::" << sendcnt << "::gathered factor size::"
                  << m_gathered_ncp_factors_t.factor(current_mode).n_elem);
#endif
    MPITIC;  // allgather tic
    MPI_Allgatherv(m_local_ncp_factors_t.factor(current_mode).memptr(), sendcnt,
                   MPI_DOUBLE,
                   m_gathered_ncp_factors_t.factor(current_mode).memptr(),
                   &recvgathercnt[0], &recvgatherdispl[0], MPI_DOUBLE,
                   // todo:: check whether it is slice or fiber while running
                   // and debugging the code.
                   current_slice_comm);
    // current_slice_comm);
    double temp = MPITOC;  // allgather toc
    this->time_stats.communication_duration(temp);
    this->time_stats.allgather_duration(temp);
#ifdef DISTNTF_VERBOSE
    DISTPRINTINFO("sent local factor::"
                  << std::endl
                  << m_local_ncp_factors_t.factor(current_mode) << std::endl
                  << " gathered factor::" << std::endl
                  << m_gathered_ncp_factors_t.factor(current_mode));
#endif
    // keep gather_ncp_factors_t consistent.
    MPITIC;  // transpose tic
    m_gathered_ncp_factors.set(
        current_mode, m_gathered_ncp_factors_t.factor(current_mode).t());
    temp = MPITOC;  // transpose toc
    this->time_stats.compute_duration(temp);
    this->time_stats.trans_duration(temp);
  }

  /**
   * It perform the mttkrp of the current_mode. That is., it determines
   * the KRP leaving out the current mode and matrix multiplies with the
   * current_mode NCP factor. Alternatively, we use the dimension trees for
   * this.
   * @param[in] current_mode
   */
  void distmttkrp(const int &current_mode) {
    double temp;
    if (!this->m_enable_dim_tree) {
      MPITIC;  // krp tic
      m_gathered_ncp_factors.krp_leave_out_one(current_mode,
                                               &ncp_krp[current_mode]);
      temp = MPITOC;  // krp toc
      this->time_stats.compute_duration(temp);
      this->time_stats.krp_duration(temp);
    }

    if (this->m_enable_dim_tree) {
      double multittv_time = 0;
      double mttkrp_time = 0;
      kdt->in_order_reuse_MTTKRP(current_mode,
                                 ncp_mttkrp_t[current_mode].memptr(), false,
                                 multittv_time, mttkrp_time);
      this->time_stats.compute_duration(multittv_time);
      this->time_stats.compute_duration(mttkrp_time);
      this->time_stats.multittv_duration(multittv_time);
      this->time_stats.mttkrp_duration(mttkrp_time);

    } else {
      MPITIC;  // mttkrp tic
      m_input_tensor.mttkrp(current_mode, ncp_krp[current_mode],
                            &ncp_mttkrp_t[current_mode]);
      temp = MPITOC;  // mttkrp toc
      this->time_stats.compute_duration(temp);
      this->time_stats.mttkrp_duration(temp);
    }
    // verify if the dimension tree output matches with the classic one
    // MAT kdt_ncp_mttkrp_t = ncp_mttkrp_t[current_mode];
    // bool same_mttkrp = arma::approx_equal(kdt_ncp_mttkrp_t,
    // ncp_mttkrp_t[current_mode], "absdiff", 1e-3); PRINTROOT("kdt vs
    // mttkrp_t::" << same_mttkrp); MAT ncp_mttkrp =
    // ncp_mttkrp_t[current_mode].t(); same_mttkrp =
    // arma::approx_equal(kdt_ncp_mttkrp_t, ncp_mttkrp, "absdiff", 1e-3);
    // PRINTROOT("kdt vs mttkrp::" << same_mttkrp);
    // PRINTROOT("kdt mttkrp::" << kdt_ncp_mttkrp_t);
    // PRINTROOT("classic mttkrp_t::" << ncp_mttkrp_t[current_mode]);

    MPI_Comm current_slice_comm = this->m_mpicomm.slice(current_mode);
    int slice_size;
    int slice_rank;
    MPI_Comm_size(current_slice_comm, &slice_size);
    slice_rank = this->m_mpicomm.slice_rank(current_mode);
    std::vector<int> recvmttkrpsize(slice_size);
    int dimsize = m_factor_local_dims[current_mode];
    for (int i = 0; i < slice_size; i++) {
      recvmttkrpsize[i] = itersplit(dimsize, slice_size, i) * m_low_rank_k;
    }
#ifdef DISTNTF_VERBOSE
    MPI_Comm current_fiber_comm = this->m_mpicomm.fiber(current_mode);
    int fiber_size;
    MPI_Comm_size(current_fiber_comm, &fiber_size);
    DISTPRINTINFO("::current_mode::"
                  << current_mode << "::slice comm size::" << slice_size
                  << "::fiber comm size::" << fiber_size
                  << "::my_global_rank::" << MPI_RANK << "::my_slice_rank::"
                  << this->m_mpicomm.slice_rank(current_mode)
                  << "::my_fiber_rank::"
                  << this->m_mpicomm.fiber_rank(current_mode)
                  << "::mttkrp_size::" << ncp_mttkrp_t[current_mode].n_elem
                  << "::local_mttkrp_size::"
                  << ncp_local_mttkrp_t[current_mode].n_elem);
#endif
    ncp_local_mttkrp_t[current_mode].zeros();
    MPITIC;  // reduce_scatter mttkrp
    MPI_Reduce_scatter(ncp_mttkrp_t[current_mode].memptr(),
                       ncp_local_mttkrp_t[current_mode].memptr(),
                       &recvmttkrpsize[0], MPI_DOUBLE, MPI_SUM,
                       current_slice_comm);
    temp = MPITOC;  // reduce_scatter mttkrp
    this->time_stats.communication_duration(temp);
    this->time_stats.reducescatter_duration(temp);
#ifdef DISTNTF_VERBOSE
    DISTPRINTINFO(ncp_mttkrp_t[current_mode]);
    DISTPRINTINFO(ncp_local_mttkrp_t[current_mode]);
#endif
    this->m_stale_mttkrp[current_mode] = false;
  }

  void allocateMatrices() {
    // allocate matrices.
    if (!m_enable_dim_tree) {
      ncp_krp = new MAT[m_modes];
    }
    ncp_mttkrp_t = new MAT[m_modes];
    ncp_local_mttkrp_t = new MAT[m_modes];
    factor_global_grams = new MAT[m_modes];
    factor_local_grams.zeros(this->m_low_rank_k, this->m_low_rank_k);
    global_gram.ones(this->m_low_rank_k, this->m_low_rank_k);
    UWORD current_size = 0;
    for (unsigned int i = 0; i < m_modes; i++) {
      current_size = TENSOR_LOCAL_NUMEL / TENSOR_LOCAL_DIM[i];
      if (!m_enable_dim_tree) {
        ncp_krp[i] = arma::zeros(current_size, this->m_low_rank_k);
      }
      ncp_mttkrp_t[i] = arma::zeros(this->m_low_rank_k, TENSOR_LOCAL_DIM[i]);
      ncp_local_mttkrp_t[i] = arma::zeros(m_local_ncp_factors.factor(i).n_cols,
                                          m_local_ncp_factors.factor(i).n_rows);
      factor_global_grams[i] =
          arma::zeros(this->m_low_rank_k, this->m_low_rank_k);
    }
  }

  void freeMatrices() {
    for (unsigned int i = 0; i < m_modes; i++) {
      if (!m_enable_dim_tree) {
        ncp_krp[i].clear();
      }
      ncp_mttkrp_t[i].clear();
      ncp_local_mttkrp_t[i].clear();
      factor_global_grams[i].clear();
    }
    if (!m_enable_dim_tree) {
      delete[] ncp_krp;
    }
    delete[] ncp_mttkrp_t;
    delete[] ncp_local_mttkrp_t;
    delete[] factor_global_grams;
  }

  void reportTime(const double temp, const std::string &reportstring) {
    double mintemp, maxtemp, sumtemp;
    MPI_Allreduce(&temp, &maxtemp, 1, MPI_DOUBLE, MPI_MAX, MPI_COMM_WORLD);
    MPI_Allreduce(&temp, &mintemp, 1, MPI_DOUBLE, MPI_MIN, MPI_COMM_WORLD);
    MPI_Allreduce(&temp, &sumtemp, 1, MPI_DOUBLE, MPI_SUM, MPI_COMM_WORLD);
    PRINTROOT(reportstring  // << "::dims::" << this->m_global_dims.t()
              << "::k::" << this->m_low_rank_k << "::SIZE::" << MPI_SIZE
              << "::algo::" << this->m_updalgo << "::root::" << temp
              << "::min::" << mintemp << "::avg::" << (sumtemp) / (MPI_SIZE)
              << "::max::" << maxtemp);
  }

  /**
   * Updates the current_mode of the NCP factors with the given
   * factor matrix. It appropriately normalizes, updates lambda,
   * updates global_gram and performs gather_ncp_factor to prepare
   * for the next iteration.
   * @param[in] current_mode
   * @param[in] factor matrix
   */

  void update_factor_mode(const unsigned int current_mode, const MAT &factor) {
    m_local_ncp_factors.set(current_mode, factor);
    m_local_ncp_factors.distributed_normalize(current_mode);
    MAT factor_t = m_local_ncp_factors.factor(current_mode).t();
    m_local_ncp_factors_t.set(current_mode, factor_t);
    m_local_ncp_factors_t.set_lambda(m_local_ncp_factors.lambda());
    // line 13 and 14
    update_global_gram(current_mode);
    // line 15
    gather_ncp_factor(current_mode);
    if (this->m_enable_dim_tree) {
      kdt->set_factor(m_gathered_ncp_factors_t.factor(current_mode).memptr(),
                      current_mode);
    }
    for (unsigned int mode = 0; mode < this->m_modes; mode++) {
      if (mode != current_mode) this->m_stale_mttkrp[mode] = true;
    }
  }

  virtual void accelerate() {}

  void generateReport() {
    MPI_Barrier(MPI_COMM_WORLD);
    this->reportTime(this->time_stats.duration(), "total_d");
    this->reportTime(this->time_stats.communication_duration(), "total_comm");
    this->reportTime(this->time_stats.compute_duration(), "total_comp");
    this->reportTime(this->time_stats.allgather_duration(), "total_allgather");
    this->reportTime(this->time_stats.allreduce_duration(), "total_allreduce");
    this->reportTime(this->time_stats.reducescatter_duration(),
                     "total_reducescatter");
    this->reportTime(this->time_stats.gram_duration(), "total_gram");
    this->reportTime(this->time_stats.krp_duration(), "total_krp");
    this->reportTime(this->time_stats.mttkrp_duration(), "total_mttkrp");
    this->reportTime(this->time_stats.multittv_duration(), "total_multittv");
    this->reportTime(this->time_stats.nnls_duration(), "total_nnls");
    if (this->m_compute_error) {
      this->reportTime(this->time_stats.err_compute_duration(),
                       "total_err_compute");
      this->reportTime(this->time_stats.err_compute_duration(),
                       "total_err_communication");
    }
  }

 public:
  DistAUNTF(const Tensor &i_tensor, const int i_k, algotype i_algo,
            const UVEC &i_global_dims, const UVEC &i_local_dims,
            const UVEC &i_nls_sizes, const UVEC &i_nls_idxs,
            const NTFMPICommunicator &i_mpicomm)
      : m_mpicomm(i_mpicomm),
        m_nls_sizes(i_nls_sizes),
        m_nls_idxs(i_nls_idxs),
        m_local_ncp_factors(i_nls_sizes, i_k, false),
        m_local_ncp_factors_t(i_nls_sizes, i_k, true),
        m_input_tensor(i_tensor),
        m_gathered_ncp_factors(i_tensor.dimensions(), i_k, false),
        m_gathered_ncp_factors_t(i_tensor.dimensions(), i_k, true),
        m_low_rank_k(i_k),
        m_modes(m_input_tensor.modes()),
        m_updalgo(i_algo),
        m_global_dims(i_global_dims),
        m_factor_local_dims(i_local_dims),
        time_stats(0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0) {
    this->m_compute_error = false;
    this->m_enable_dim_tree = false;
    this->m_accelerated = false;
    this->m_num_it = 30;
    this->m_rel_error = 1.0;
    // randomize again. otherwise all the process and factors
    // will be same.
    m_local_ncp_factors.randu(149 * i_mpicomm.rank() + 103);
    m_local_ncp_factors.distributed_normalize();
    for (unsigned int i = 0; i < this->m_modes; i++) {
      MAT current_factor = arma::trans(m_local_ncp_factors.factor(i));
      m_local_ncp_factors_t.set(i, current_factor);
      this->m_stale_mttkrp.push_back(true);
    }
    m_local_ncp_factors_t.set_lambda(m_local_ncp_factors.lambda());
    m_gathered_ncp_factors.trans(m_gathered_ncp_factors_t);
    allocateMatrices();
    double normA = i_tensor.norm();
    MPI_Allreduce(&normA, &this->m_global_sqnorm_A, 1, MPI_DOUBLE, MPI_SUM,
                  MPI_COMM_WORLD);

    DISTPRINTINFO("::NLS Solve Sizes::"
                  << m_nls_sizes << "::NLS start indices::" << m_nls_idxs);
  }
  ~DistAUNTF() {
    freeMatrices();
    if (this->m_enable_dim_tree) {
      delete kdt;
    }
  }
  /// Returns number of iterations
  void num_iterations(const int i_n) { this->m_num_it = i_n; }
  /// Returns the numbers of modes of the tensor
  size_t modes() const { return this->m_modes; }
  /// Low Rank
  size_t rank() const { return this->m_low_rank_k; }
  /// L1 and L2 Regularization for every mode
  void regularizers(const FVEC i_regs) { this->m_regularizers = i_regs; }
  /// Sets whether to compute the error or not
  void compute_error(bool i_error) {
    this->m_compute_error = i_error;
    hadamard_all_grams =
        arma::ones<MAT>(this->m_low_rank_k, this->m_low_rank_k);
  }
  /// Returns the lambda of the NCP factors
  VEC lambda() { return m_local_ncp_factors.lambda(); }
  /// Returns the current outer iteration of the computeNTF
  int current_it() const { return this->m_current_it; }
  /// Returns the current error
  double current_error() const { return this->m_rel_error; }
  /// MTTKRP can be computed with or without dimension trees. Dimtree is
  /// default.
  void dim_tree(bool i_dim_tree) {
    this->m_enable_dim_tree = i_dim_tree;
    if (this->m_enable_dim_tree) {
      if (this->ncp_krp != NULL) {
        for (unsigned int i = 0; i < m_modes; i++) {
          ncp_krp[i].clear();
        }
        delete[] ncp_krp;
      }
    }
  }
  /// Does the algorithm need acceleration?
  void accelerated(const bool &set_acceleration) {
    this->m_accelerated = set_acceleration;
    this->m_compute_error = true;
  }

  bool is_stale_mttkrp(const int &current_mode) const {
    return this->m_stale_mttkrp[current_mode];
  }

  /**
   * This function will completely reset all the factors
   * and the state of AUNTF. Use it with caution!!!
   * It is the responsibility of the caller to preserve
   * the current factor states and other relevant computations
   */

  void reset(const NCPFactors &new_factors, bool trans = false) {
    if (!trans) {
      for (unsigned int i = 0; i < m_modes; i++) {
        update_factor_mode(i, new_factors.factor(i));
      }
    } else {
      for (unsigned int i = 0; i < m_modes; i++) {
        update_factor_mode(i, new_factors.factor(i).t());
      }
    }
    m_local_ncp_factors.set_lambda(new_factors.lambda());
    m_local_ncp_factors_t.set_lambda(new_factors.lambda());
  }

  // Preferrably call this after the computeNTF().
  // This is right now called to save the factor matrices.
  /**
   * Returns the factor matrix by collected it across all the processors.
   * Don't call this during the computeNTF loop. It is right now
   * called after all the iterations to dump the factor matrix.
   * @param[in] mode of the factor matrix to be collected across processors
   * @param[in] factor_matrix. To be allocated and freed by the caller.
   */
  void factor(int mode, double *factor_matrix) {
    gather_ncp_factor(mode);
    int sendcnt = m_gathered_ncp_factors_t.factor(mode).n_elem;
    int fiber_size = this->m_mpicomm.proc_grids()[mode];
    int global_size = this->m_global_dims[mode];
    std::vector<int> recvcnts(fiber_size, 0);
    std::vector<int> displs(fiber_size, 0);

    DISTPRINTINFO("Collecting mode::" << mode << "::sendcnt::" << sendcnt
                                      << "::fiber_size::" << fiber_size
                                      << "::global_size::" << global_size);

    // int dimsize = m_factor_local_dims[current_mode];
    for (int i = 0; i < fiber_size; i++) {
      recvcnts[i] = itersplit(global_size, fiber_size, i) * m_low_rank_k;
      displs[i] = startidx(global_size, fiber_size, i) * m_low_rank_k;
    }

    MPI_Gatherv(m_gathered_ncp_factors_t.factor(mode).memptr(), sendcnt,
                MPI_DOUBLE, factor_matrix, &recvcnts[0], &displs[0], MPI_DOUBLE,
                // todo:: check whether it is slice or fiber while running
                // and debugging the code.
                0, this->m_mpicomm.fiber(mode));
  }

  /// The main computeNTF loop
  void computeNTF() {
    // initialize everything.
    // line 3,4,5 of the algorithm
    for (unsigned int i = 1; i < m_modes; i++) {
      update_global_gram(i);
      gather_ncp_factor(i);
    }
    if (this->m_enable_dim_tree) {
      // Determine optimial split when given mode ordering.
      // product of left dimensions \approx product of right dimensions.
      size_t split_criteria = arma::prod(m_input_tensor.dimensions());
      split_criteria = std::round(std::sqrt(split_criteria));
      UVEC temp_cum_prod = arma::cumprod(m_input_tensor.dimensions());
      int split_mode = 0;
      while (temp_cum_prod(split_mode) < split_criteria) {
        split_mode++;
      }
      PRINTROOT("KDT Split Mode::" << split_mode
                                   << "::split criteria::" << split_criteria
                                   << "::cum prod::" << std::endl
                                   << temp_cum_prod << std::endl);

      // check to see if split mode is left or right.
      if (split_mode > 0) {
        size_t current_left = temp_cum_prod(split_mode);
        size_t good_criteria = temp_cum_prod(temp_cum_prod.n_rows - 1) /
                               temp_cum_prod(split_mode - 1);
        if (current_left > good_criteria) split_mode--;
        PRINTROOT("KDT Split Mode::"
                  << split_mode << "::split criteria::" << split_criteria
                  << "::numerator::" << temp_cum_prod(temp_cum_prod.n_rows - 1)
                  << "::good_criteria::" << good_criteria
                  << "::current_left::" << current_left << std::endl
                  << "::cum prod::" << std::endl
                  << temp_cum_prod << std::endl);
      }
      kdt = new DenseDimensionTree(m_input_tensor, m_gathered_ncp_factors,
                                   split_mode);
    }
#ifdef DISTNTF_VERBOSE
    DISTPRINTINFO("local factor matrices::");
    this->m_local_ncp_factors.print();
    DISTPRINTINFO("local factor matrices transpose::");
    this->m_local_ncp_factors_t.print();
    DISTPRINTINFO("gathered factor matrices::");
    this->m_gathered_ncp_factors.print();
#endif
    for (this->m_current_it = 0; this->m_current_it < m_num_it;
         this->m_current_it++) {
      MAT unnorm_factor;
      for (unsigned int current_mode = 0; current_mode < m_modes;
           current_mode++) {
        // line 9 and 10 of the algorithm
        if (is_stale_mttkrp(current_mode)) distmttkrp(current_mode);
        // line 11 of the algorithm
        gram_hadamard(current_mode);
        // line 12 of the algorithm
#ifdef DISTNTF_VERBOSE
        DISTPRINTINFO("local factor matrix::"
                      << this->m_local_ncp_factors.factor(current_mode));
        DISTPRINTINFO("gathered factor matrix::");
        this->m_gathered_ncp_factors.print();
        PRINTROOT("global_grams::" << std::endl << this->global_gram);
        DISTPRINTINFO("mttkrp::");
        this->ncp_local_mttkrp_t[current_mode].print();
#endif
        MPITIC;  // nnls_tic
        MAT factor = update(current_mode);
        double temp = MPITOC;  // nnls_toc
        this->time_stats.compute_duration(temp);
        this->time_stats.nnls_duration(temp);
#ifdef DISTNTF_VERBOSE
        DISTPRINTINFO("it::" << this->m_current_it << "::mode::" << current_mode
                             << std::endl
                             << factor);
#endif
        if (m_compute_error && current_mode == this->m_modes - 1) {
          unnorm_factor = factor;
        }
        update_factor_mode(current_mode, factor.t());
      }
      if (m_compute_error) {
        double temp_err = computeError(unnorm_factor, this->m_modes - 1);
        this->m_rel_error = temp_err;
        double iter_time = this->time_stats.compute_duration() +
                           this->time_stats.communication_duration();
        PRINTROOT("Iter::" << this->m_current_it << "::k::"
                           << this->m_low_rank_k << "::SIZE::" << MPI_SIZE
                           << "::algo::" << this->m_updalgo << "::time::"
                           << iter_time << "::relative_error::" << temp_err);
      }
      if (this->m_accelerated) {
        // there is a acceleration possible. call accelerate method
        // in the derived class.
        accelerate();
      }
      PRINTROOT("completed it::" << this->m_current_it);
    }
    generateReport();
  }
  /**
   * This is naive way of computing error with local rank-k tensor
   * rel_err = \f$ sqrt(max(\|X\|_F^2 + \lambda^T * Hadamard\; of\; all\; gram *
               \lambda - 2 * innerprod(X,F_{kten}),0))/\|X\|;
   * Typically called after the update of the last mode.
   */

  double computeError(const MAT &unnorm_factor, int mode) {
    // rel_Error = sqrt(max(init.nr_X^2 + lambda^T * Hadamard of all gram *
    // lambda - 2 * innerprod(X,F_kten),0))/init.nr_X;
    MPITIC;  // err compute
    hadamard_all_grams = global_gram % factor_global_grams[mode];
    VEC local_lambda = m_local_ncp_factors.lambda();
    ROWVEC temp_vec = local_lambda.t() * hadamard_all_grams;
    double sq_norm_model = arma::dot(temp_vec, local_lambda);
    // double sq_norm_model = arma::norm(hadamard_all_grams, "fro");
    // sum of the element-wise dot product between the local mttkrp and
    // the factor matrix
    double inner_product = arma::dot(ncp_local_mttkrp_t[mode], unnorm_factor);
    double temp = MPITOC;  // err compute
    this->time_stats.compute_duration(temp);
    this->time_stats.err_compute_duration(temp);
    double all_inner_product;
    MPITIC;  // err comm
    MPI_Allreduce(&inner_product, &all_inner_product, 1, MPI_DOUBLE, MPI_SUM,
                  MPI_COMM_WORLD);
    temp = MPITOC;  // err comm
    this->time_stats.communication_duration(temp);
    this->time_stats.err_communication_duration(temp);
#ifdef DISTNTF_VERBOSE
    DISTPRINTINFO("local_lambda::" << local_lambda);
    DISTPRINTINFO("local_inner_product::" << inner_product << std::endl);
    PRINTROOT("norm_A_sq :: "
              << this->m_global_sqnorm_A << "::model_norm_sq::" << sq_norm_model
              << "::global_inner_product::" << all_inner_product << std::endl);
#endif
    double squared_err =
        this->m_global_sqnorm_A + sq_norm_model - 2 * all_inner_product;
    if (squared_err < 0) {
      PRINTROOT("computed error is negative due to round off");
      PRINTROOT("norm_A_sq :: "
                << this->m_global_sqnorm_A
                << "::model_norm_sq::" << sq_norm_model
                << "::global_inner_product::" << all_inner_product
                << "::squared_err::" << squared_err << std::endl);
    }
    return std::sqrt(std::abs(squared_err) / this->m_global_sqnorm_A);
  }
  /**
   * This is used to evaluate during acceleration stage. If the accelerated
   * factors are better than the current factors, we accept the acceleration.
   * Hence, we pass the accelerated factors to compute error.
   */
  double computeError(const NCPFactors &new_factors_t, const int mode) {
    // rel_Error = sqrt(max(init.nr_X^2 + lambda^T * Hadamard of all gram *
    // lambda - 2 * innerprod(X,F_kten),0))/init.nr_X;
    // Reset with new factors and compute error on mode 0
    reset(new_factors_t, true);
    distmttkrp(mode);
    gram_hadamard(mode);
    hadamard_all_grams = global_gram % factor_global_grams[mode];
    VEC local_lambda = m_local_ncp_factors.lambda();
    MAT unnorm_factor =
        arma::diagmat(local_lambda) * new_factors_t.factor(mode);
    ROWVEC temp_vec = local_lambda.t() * hadamard_all_grams;
    double sq_norm_model = arma::dot(temp_vec, local_lambda);
    // double sq_norm_model = arma::norm(hadamard_all_grams, "fro");
    // sum of the element-wise dot product between the local mttkrp and
    // the factor matrix
    double inner_product = arma::dot(ncp_local_mttkrp_t[mode], unnorm_factor);
    double all_inner_product;
    MPI_Allreduce(&inner_product, &all_inner_product, 1, MPI_DOUBLE, MPI_SUM,
                  MPI_COMM_WORLD);
    double squared_err =
        this->m_global_sqnorm_A + sq_norm_model - 2 * all_inner_product;
    if (squared_err < 0) {
      PRINTROOT("computed error is negative due to round off");
      PRINTROOT("norm_A_sq :: "
                << this->m_global_sqnorm_A
                << "::model_norm_sq::" << sq_norm_model
                << "::global_inner_product::" << all_inner_product
                << "::squared_err::" << squared_err << std::endl);
    }
    return std::sqrt(std::abs(squared_err) / this->m_global_sqnorm_A);
  }
};  // class DistAUNTF
}  // namespace planc
#endif  // DISTNTF_DISTAUNTF_HPP_


// slideshow.cpp : This file contains the 'main' function. Program execution begins and ends there.
//

#include <iostream>
#include <fstream>
#include <sstream>
#include <string>
#include <vector>
#include <unordered_map>
#include <algorithm>
#include <map>
#include <list>
#include <omp.h>
#include <chrono>
#include <functional>

#include <Eigen/Sparse>

struct ScopedTimer
{
	std::chrono::high_resolution_clock::time_point t0_;
	std::string title_;

	ScopedTimer(std::string title)
		: t0_(std::chrono::high_resolution_clock::now())
		, title_(title)
	{}

	~ScopedTimer()
	{
		auto t1 = std::chrono::high_resolution_clock::now();
		auto nanos = std::chrono::duration_cast<std::chrono::seconds>(t1 - t0_).count();

		std::cout << title_ << " took " << nanos << "s to complete" << std::endl;
	}
};

enum OrientationMode
{
	Horizontal,
	Vertical
};

struct Stats
{
	uint64_t min_tags_count = -1;
	uint64_t max_tags_count = 0;
	uint64_t avg_tags_count = 0;
};

struct Photo
{
	uint32_t idx;
	OrientationMode orientation;
	std::vector<uint32_t> tags;
};

bool photoLess(const Photo& p1, const Photo& p2)
{
	return p1.tags.size() < p2.tags.size();
}

void BuildAdjacencyMatrix(std::vector<Photo> &photos, Stats &s, Eigen::SparseMatrix<uint8_t> &mat)
{
	const uint64_t photos_count = photos.size();

	std::vector<uint8_t> intersections_per_photo(photos_count, 0);
	std::vector<uint64_t> intersected_tags_count(s.max_tags_count, 0);

	mat.reserve(photos.size() * 4);

#pragma omp parallel for shared(photos, mat) num_threads(8)
	for (int i = 0; i < photos_count - 1; ++i)
	{
		const auto& curent_photo = photos[i];
		for (int k = i + 1; k < photos_count; ++k)
		{
			const auto& matching_photo = photos[k];
			std::vector<uint32_t> tmp;
			std::set_intersection(curent_photo.tags.begin(), curent_photo.tags.end(),
								  matching_photo.tags.begin(), matching_photo.tags.end(),
								  std::back_inserter(tmp));

			const uint64_t intersected_tags = tmp.size();
			if (intersected_tags == 0)
				continue;

			++intersected_tags_count[intersected_tags];
			++intersections_per_photo[i];
			++intersections_per_photo[k];

			uint64_t score = std::min(std::min(intersected_tags, curent_photo.tags.size()), matching_photo.tags.size());

#pragma omp critical
			{
				mat.insert(i, k) = score;
				mat.insert(k, i) = score;
			}
		}
	}
	mat.makeCompressed();

	auto minmaxIt = std::minmax_element(intersections_per_photo.begin(), intersections_per_photo.end());
	std::cout << "min intersections: " << *minmaxIt.first << " max intersections: " << *minmaxIt.second << std::endl;
}

template <typename ResultIt>
uint64_t BuildRoute(Photo current_photo, std::vector<Photo> &horizontal_photos, ResultIt resultIterator, const bool untillTheEnd = true)
{
	uint64_t length_of_slides = 0;
	while (true)
	{

		std::vector<uint32_t> tmp;
		auto nextPhotoIt = std::find_if(horizontal_photos.begin(), horizontal_photos.end(), [&](const Photo &p) {
			tmp.clear();
			std::set_intersection(current_photo.tags.begin(), current_photo.tags.end(),
								  p.tags.begin(), p.tags.end(),
								  std::back_inserter(tmp));

			const uint64_t intersected_tags = tmp.size();
			return intersected_tags > 0;
		});

		if (nextPhotoIt != horizontal_photos.end())
		{
			resultIterator = nextPhotoIt->idx;
			current_photo = *nextPhotoIt;
			++length_of_slides;

			horizontal_photos.erase(nextPhotoIt);

			if (!untillTheEnd)
				break;
		}
		else
		{
			break;
		}
	}

	return length_of_slides;
}

std::vector<Photo>::iterator Closest(std::vector<Photo> &photos, uint8_t tag_count)
{
	auto lowIt = std::lower_bound(photos.begin(), photos.end(), tag_count, [&](const Photo &p, uint8_t tags) {return p.tags.size() < tags; });
	auto highIt = std::upper_bound(photos.begin(), photos.end(), tag_count, [&](uint8_t tags, const Photo &p) {return p.tags.size() > tags; });

	if (lowIt != photos.end() && highIt != photos.end())
	{
		auto lowDiff = tag_count - lowIt->tags.size();
		auto hiDiff = highIt->tags.size() - tag_count;

		return lowDiff < hiDiff ? lowIt : highIt;
	}
	else if (lowIt != photos.end())
		return lowIt;
	else if (highIt != photos.end())
		return highIt;

	return photos.begin();
}

Photo PickNextPhoto(Stats &s, std::vector<Photo> &horizontal_photos, std::vector<Photo> &vertical_photos)
{
	if (!horizontal_photos.empty())
	{
		auto current_it = Closest(horizontal_photos, s.avg_tags_count);
		auto current_photo = *current_it;
		horizontal_photos.erase(current_it);

		std::cout << "Path start tag length = " << current_photo.tags.size() << std::endl;

		return current_photo;
	}
}

std::list<uint64_t> Solve(Stats &s, std::vector<Photo> &horizontal_photos, std::vector<Photo> &vertical_photos)
{
	std::sort(horizontal_photos.begin(), horizontal_photos.end(), photoLess);
	const uint64_t photos_count = horizontal_photos.size();

	uint64_t current_index = 0;
	std::vector<uint32_t> tmp;

	std::list<uint64_t> final_results;
	uint64_t length_of_slides = 1;

	bool reverse_order = false;
	while (!horizontal_photos.empty() || !vertical_photos.empty())
	{
		auto current_photo = PickNextPhoto(s, horizontal_photos, vertical_photos);

		std::list<uint64_t> current_results;
		current_results.push_back(current_photo.idx);

		uint64_t route_length = 1;
		{
			route_length += BuildRoute(current_photo, horizontal_photos, std::back_inserter(current_results), true);
			route_length += BuildRoute(current_photo, horizontal_photos, std::front_inserter(current_results), true);
		}

		std::cout << "Length of path = " << route_length << std::endl;

		final_results.splice(final_results.end(), current_results);

	}

	return std::move(final_results);
}

int main(int argc, char **argv)
{
	ScopedTimer overall("Whole Process");

	if (argc != 3)
		std::cout << "Invalid parameter!\n";

	std::vector<Photo> vertical_photos;
	std::vector<Photo> horizontal_photos;

	Stats s;
	{
		ScopedTimer preprocess("Preprocess");

		std::ifstream inputFileStream(argv[1]);
		int count;
		inputFileStream >> count;
		inputFileStream.ignore(1, '\n');

		std::string line;
		vertical_photos.reserve(count);
		horizontal_photos.reserve(count);

		std::unordered_map<std::string, uint32_t> tag_ids;
		uint32_t current_tag_id = 0;
		for (int i = 0; i < count; i++)
		{
			std::getline(inputFileStream, line);

			std::istringstream iss(line);
			char orientation = 0;
			iss >> orientation;
			int tag_count = 0;
			iss >> tag_count;

			Photo p = {};
			p.orientation = orientation == 'V' ? OrientationMode::Vertical : OrientationMode::Horizontal;
			p.idx = i;
			p.tags.reserve(tag_count);

			for (int j = 0; j < tag_count; ++j)
			{
				std::string tag;
				iss >> tag;

				auto it = tag_ids.find(tag);
				if (it == tag_ids.end())
				{
					it = tag_ids.insert(std::make_pair(tag, current_tag_id++)).first;
				}
				p.tags.push_back(it->second);
			}
			sort(p.tags.begin(), p.tags.end(), std::less<uint32_t>());

			if (tag_count < s.min_tags_count)
				s.min_tags_count = tag_count;
			if (tag_count > s.max_tags_count)
				s.max_tags_count = tag_count;
			s.avg_tags_count += tag_count;

			if (p.orientation == OrientationMode::Vertical)
				vertical_photos.emplace_back(std::move(p));
			else
				horizontal_photos.emplace_back(std::move(p));
		}
		s.avg_tags_count /= vertical_photos.size() + horizontal_photos.size();
	}

	std::cout << "min_tags: " << s.min_tags_count << "max_tags: " << s.max_tags_count << "avg_tags: " << s.avg_tags_count << std::endl;
	std::list<uint64_t> res;
	{
		ScopedTimer solve("Solve");
		res = Solve(s, horizontal_photos, vertical_photos);
	}

	ScopedTimer write_res("Write results");

	std::ofstream output_file(argv[2]);
	output_file << res.size() << "\n";
	std::ostream_iterator<uint64_t> output_iterator(output_file, "\n");
	std::copy(res.begin(), res.end(), output_iterator);
}


#include <NTL/vec_GF2E.h>
#include <rank/knowledge_of_valid_opening/knowledge_of_valid_opening.h>
#include <rank/rank_commitment/rank_commitment.h>
#include <utils/utils.h>
#include "linear_relations.h"

void rank::linear_relations::initalized_commitments_and_responses(
        commitments_t *commitments,
        responses_t *responses) {

    commitments->s0.SetLength(I);
    commitments->s1.SetLength(I);
    commitments->s2.SetLength(I);

    commitments->c0.SetLength(I);
    commitments->c1.SetLength(I);
    commitments->c2.SetLength(I);

    responses->r1.SetLength(I);
    responses->r2.SetLength(I);
    responses->r3.SetLength(I);
}

void rank::linear_relations::generate_revealed_values(
        revealed_values_t *revealed_values) {

    revealed_values->P.SetLength(I);
    revealed_values->Q.SetLength(I);

    ::utils::generate_random_binary_matrix(
            revealed_values->matrices.x_0,
            MI,
            MI);
    ::utils::generate_random_binary_matrix(
            revealed_values->matrices.x_1,
            MI,
            MI);

    for(int i = 0; i < I; i++) {
        utils::generate_random_square_invertible_matrix(
                revealed_values->P[i],
                EN);
        utils::generate_random_square_invertible_matrix(
                revealed_values->Q[i],
                EM);
    }
}

void rank::linear_relations::generate_private_key(
        private_key_t *private_key,
        const linear_relation_matrices_t *matrices) {

    private_key->m.kill();

    private_key->e.kill();
    private_key->e.SetLength(I);

    for(int i = 0; i < I; i++) {
        private_key->e[i].kill();
        private_key->e[i].SetLength(EN);
        ::utils::generate_vector_of_specific_rank(
                private_key->e[i],
                EN,
                RHO);
    }

    for(int i = 0; i < 2; i++) {
        private_key->m.append(NTL::vec_GF2());
        private_key->m[i] = NTL::random_vec_GF2(MI);
    }
    private_key->m.append((private_key->m[0] * matrices->x_0) + (private_key->m[1] * matrices->x_1));
}

void rank::linear_relations::generate_public_key(
        public_key_t *public_key,
        const private_key_t *private_key) {

    public_key->commitments.SetLength(I);

    utils::generate_random_matrix_gf2e(
            public_key->G,
            K,
            EN);

    for(int i = 0; i < I; i++) {
        NTL::vec_GF2 s;
        utils::generate_random_binary_vector(
                s,
                PI);

        rank::rank_commitment::generate_commitment(
                public_key->commitments[i],
                s,
                private_key->m[i],
                public_key->G,
                private_key->e[i]);
    }
}

void rank::linear_relations::generate_random_values(
        random_values_t *random_values,
        const linear_relation_matrices_t *matrices) {

    for (int i = 0; i < I; i++) {
        random_values->u.append(NTL::random_vec_GF2(PI));
        random_values->f.append(NTL::random_vec_GF2E(EN));
    }

    for (int i = 0; i < I - 1; i++) {
        random_values->v.append(NTL::random_vec_GF2(MI));
    }

    auto v_2 = (random_values->v[0] * matrices->x_0) + (random_values->v[1] * matrices->x_1);
    random_values->v.append(v_2);

    random_values->u_v = random_values->u;

    for(int i = 0; i < I; i++) {
        NTL::append(random_values->u_v[i], random_values->v[i]);
    }
}

void rank::linear_relations::generate_commitments_and_responses(
        commitments_t *commitments,
        responses_t *responses,
        const random_values_t *random_values,
        const revealed_values_t *revealed_values,
        const private_key_t *private_key,
        const public_key_t *public_key) {

    for (int i = 0; i < I; i++) {
        rank::knowledge_of_valid_opening::generate_commitment_0_and_response_0(
                commitments->c0[i],
                commitments->s0[i],
                responses->r1[i],
                utils::gf2e_from_vec_gf2(random_values->u_v[i]),
                random_values->f[i],
                revealed_values->P[i],
                revealed_values->Q[i],
                public_key->G);

        rank::knowledge_of_valid_opening::generate_commitment_1_and_response_1(
                commitments->c1[i],
                commitments->s1[i],
                responses->r2[i],
                random_values->f[i],
                revealed_values->P[i],
                revealed_values->Q[i],
                public_key->G);

        rank::knowledge_of_valid_opening::generate_commitment_2_and_response_2(
                commitments->c2[i],
                commitments->s2[i],
                responses->r3[i],
                private_key->e[i],
                random_values->f[i],
                revealed_values->P[i],
                revealed_values->Q[i],
                public_key->G);
    }
}

int rank::linear_relations::verify_0(
        const NTL::Vec<NTL::vec_GF2> &s0,
        const NTL::Vec<NTL::vec_GF2> &s1,
        const NTL::Vec<NTL::vec_GF2E> &c0,
        const NTL::Vec<NTL::vec_GF2E> &c1,
        const NTL::Vec<NTL::vec_GF2E> &r1,
        const NTL::Vec<NTL::vec_GF2E> &r2,
        const NTL::Vec<NTL::mat_GF2> &P,
        const NTL::Vec<NTL::mat_GF2> &Q,
        const NTL::mat_GF2 &x_0,
        const NTL::mat_GF2 &x_1,
        const public_key_t *public_key) {

    for (int i = 0; i < I; i++) {
        auto _r1 = utils::encode(P[i], Q[i], r1[i], MI);
        if (rank::rank_commitment::verify_proof(
                c0[i],
                s0[i],
                _r1,
                public_key->G) != 0) {
            std::cout << "Verification 2 for ch = 0 failed" << std::endl;
            return 1;
        }
    }

    for (int i = 0; i < I; i++) {
        auto _r2 = utils::encode_2(r2[i], MI);
        if (rank::rank_commitment::verify_proof(
                c1[i],
                s1[i],
                _r2,
                public_key->G) != 0) {
            std::cout << "Verification 2 for ch = 0 failed" << std::endl;
            return 1;
        }
    }

    NTL::Vec<NTL::vec_GF2E> results;
    for (int i = 0; i < I; i++) {
        NTL::vec_GF2E pi_inv_r2;
        rank::knowledge_of_valid_opening::calculate_pi_inv(
                pi_inv_r2,
                P[i],
                Q[i],
                r2[i]);

        NTL::vec_GF2E result;
        if (utils::solve_equation(
                result,
                public_key->G,
                r1[i] + pi_inv_r2) != 0) {
            std::cout << "No Solutions" << std::endl;
            return 1;
        }
        results.append(result);
    }

    NTL::Vec<NTL::vec_GF2E> b;
    b.SetLength(I);
    for (int i = 0; i < I; i++) {
        NTL::VectorCopy(b[i], results[i], K);
    }

    auto b0_gf2 = utils::gf2_from_gf2e(b[0]);
    auto b1_gf2 = utils::gf2_from_gf2e(b[1]);
    auto b2_gf2 = utils::gf2_from_gf2e(b[2]);

    NTL::vec_GF2 b0_gf2_s, b1_gf2_s, b2_gf2_s;
    for (int j = PI; j < PI + MI; j++) {
        b0_gf2_s.append(b0_gf2[j]);
        b1_gf2_s.append(b1_gf2[j]);
        b2_gf2_s.append(b2_gf2[j]);
    }

    if (NTL::IsZero(b2_gf2_s + ((b0_gf2_s * x_0) + (b1_gf2_s * x_1)))) {
        return 0;
    } else {
        std::cout << "Verify 0 failed" << std::endl;
        std::cout << b2_gf2_s << std::endl;
        std::cout << (b0_gf2_s * x_0) + (b1_gf2_s * x_1) << std::endl;
        return 1;
    }
}

int rank::linear_relations::verify_1(
        const NTL::Vec<NTL::vec_GF2> &s0,
        const NTL::Vec<NTL::vec_GF2> &s2,
        const NTL::Vec<NTL::vec_GF2E> &c0,
        const NTL::Vec<NTL::vec_GF2E> &c2,
        const NTL::Vec<NTL::vec_GF2E> &r1,
        const NTL::Vec<NTL::vec_GF2E> &r3,
        const NTL::Vec<NTL::mat_GF2> &P,
        const NTL::Vec<NTL::mat_GF2> &Q,
        const NTL::mat_GF2 &x_0,
        const NTL::mat_GF2 &x_1,
        const public_key_t *public_key) {

    for (int i = 0; i < I; i++) {
        auto _r1 = utils::encode(P[i], Q[i], r1[i], MI);
        if (rank::rank_commitment::verify_proof(
                c0[i],
                s0[i],
                _r1,
                public_key->G) != 0) {
            std::cout << "Verification 2 for ch = 0 failed" << std::endl;
            return 1;
        }
    }

    for (int i = 0; i < I; i++) {
        auto _r3 = utils::encode_2(r3[i], MI);
        if (rank::rank_commitment::verify_proof(
                c2[i],
                s2[i],
                _r3,
                public_key->G) != 0) {
            std::cout << "Verification 2 for ch = 0 failed" << std::endl;
            return 1;
        }
    }

    NTL::Vec<NTL::vec_GF2E> results;
    for (int i = 0; i < I; i++) {
        NTL::vec_GF2E pi_inv_r3;
        rank::knowledge_of_valid_opening::calculate_pi_inv(
                pi_inv_r3,
                P[i],
                Q[i],
                r3[i]);

        NTL::vec_GF2E result;
        if (utils::solve_equation(
                result,
                public_key->G,
                r1[i] + pi_inv_r3 + public_key->commitments[i]) != 0) {
            std::cout << "No Solutions" << std::endl;
            return 1;
        }
        results.append(result);
    }

    NTL::Vec<NTL::vec_GF2E> d;
    d.SetLength(I);
    for (int i = 0; i < I; i++) {
        NTL::VectorCopy(d[i], results[i], K);
    }

    auto d0_gf2 = utils::gf2_from_gf2e(d[0]);
    auto d1_gf2 = utils::gf2_from_gf2e(d[1]);
    auto d2_gf2 = utils::gf2_from_gf2e(d[2]);

    NTL::vec_GF2 d0_gf2_s, d1_gf2_s, d2_gf2_s;
    for (int j = 0; j < MI; j++) {
        d0_gf2_s.append(d0_gf2[j + PI]);
        d1_gf2_s.append(d1_gf2[j + PI]);
        d2_gf2_s.append(d2_gf2[j + PI]);
    }

    if (NTL::IsZero(d2_gf2_s + ((d0_gf2_s * x_0) + (d1_gf2_s * x_1)))) {
        return 0;
    } else {
        std::cout << "Verify 1 failed" << std::endl;
        std::cout << d2_gf2_s << std::endl;
        std::cout << (d0_gf2_s * x_0) + (d1_gf2_s * x_1) << std::endl;
        return 1;
    }
}

int rank::linear_relations::verify_2(
        const NTL::Vec<NTL::vec_GF2> &s1,
        const NTL::Vec<NTL::vec_GF2> &s2,
        const NTL::Vec<NTL::vec_GF2E> &c1,
        const NTL::Vec<NTL::vec_GF2E> &c2,
        const NTL::Vec<NTL::vec_GF2E> &r2,
        const NTL::Vec<NTL::vec_GF2E> &r3,
        const public_key_t *public_key) {

    for (int i = 0; i < I; i++) {
        auto _r2 = utils::encode_2(r2[i], MI);
        if (rank::rank_commitment::verify_proof(
                c1[i],
                s1[i],
                _r2,
                public_key->G)) {
            std::cout << "Verification 1 for ch = 2 failed" << std::endl;
            return 1;
        }
    }

    for (int i = 0; i < I; i++) {
        auto _r3 = utils::encode_2(r3[i], MI);
        if (rank::rank_commitment::verify_proof(
                c2[i],
                s2[i],
                _r3,
                public_key->G)) {
            std::cout << "Verification 2 for ch = 2 failed" << std::endl;
            return 1;
        }
    }

    for(int i = 0; i < I; i++) {
        auto r2_plus_r3 = r2[i] + r3[i];

        NTL::mat_GF2 M;
        M.SetDims(EN, EM);

        for (int j = 0; j < EN; j++) {
            for (int z = 0; z < EM; z++) {
                if (NTL::IsZero(r2_plus_r3[j]._GF2E__rep[z])) {
                    M[j][z] = NTL::GF2(0);
                } else {
                    M[j][z] = NTL::GF2(1);
                }
            }
        }

        auto _M = M;

        auto res = NTL::gauss(_M);
        if (res != RHO) {
            std::cout << "Matrix weight not equal to RHO" << std::endl;
            return 1;
        }
    }

    return 0;
}


// Copyright 2015-2018 Hans Dembinski
//
// 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)

//[ guide_fill_histogram

#include <boost/histogram.hpp>
#include <cassert>
#include <functional>
#include <numeric>
#include <utility>
#include <vector>

int main() {
  using namespace boost::histogram;

  auto h = make_histogram(axis::integer<>(0, 3), axis::regular<>(2, 0.0, 1.0));

  // fill histogram, number of arguments must be equal to number of axes,
  // types must be convertible to axis value type (here integer and double)
  h(0, 0.2); // increase a cell value by one
  h(2, 0.5); // increase another cell value by one

  // fills from a tuple are also supported; passing a tuple of wrong size
  // causes an error at compile-time or an assertion at runtime in debug mode
  auto xy = std::make_tuple(1, 0.3);
  h(xy);

  // functional-style processing is also supported, generate some data...
  std::vector<std::tuple<int, double>> input_data;
  input_data.emplace_back(0, 0.8);
  input_data.emplace_back(2, 0.4);
  input_data.emplace_back(5, 0.7); // goes into overflow bin of first axis

  // std::for_each takes the functor by value, a reference wrapper avoid copies
  std::for_each(input_data.begin(), input_data.end(), std::ref(h));

  // once histogram is filled, access cells using operator[] or at(...)
  // - operator[] can only accept a single argument in the current version of C++,
  //   it is convenient when you have a 1D histogram
  // - at(...) can accept several values, so use this by default
  // - underflow bins are at index -1, overflow bins at index `size()`
  // - passing an invalid index triggers a std::out_of_range exception
  assert(h.at(0, 0) == 1);
  assert(h.at(0, 1) == 1);
  assert(h.at(1, 0) == 1);
  assert(h.at(1, 1) == 0);
  assert(h.at(2, 0) == 1);
  assert(h.at(2, 1) == 1);
  assert(h.at(-1, -1) == 0); // underflow for axis 0 and 1
  assert(h.at(-1, 0) == 0);  // underflow for axis 0, normal bin for axis 1
  assert(h.at(-1, 2) == 0);  // underflow for axis 0, overflow for axis 1
  assert(h.at(3, 1) == 1);   // overflow for axis 0, normal bin for axis 1

  // iteration over values works, but see next example for a better way
  // - iteration using begin() and end() includes under- and overflow bins
  // - iteration order is an implementation detail and may change in future versions
  const double sum = std::accumulate(h.begin(), h.end(), 0.0);
  assert(sum == 6);
}

//]


#ifndef ALM_ROBOT_MODEL_H
#define ALM_ROBOT_MODEL_H

#include <Configuration.h> // Package Path
#include <avatar_locomanipulation/enable_pinocchio_with_hpp_fcl.h> // Enable HPP FCL
// Multibody
#include "pinocchio/multibody/model.hpp"
#include "pinocchio/multibody/data.hpp"
#include "pinocchio/multibody/geometry.hpp"
// Algorithms
#include "pinocchio/algorithm/kinematics.hpp"
#include "pinocchio/algorithm/kinematics-derivatives.hpp"
#include "pinocchio/algorithm/geometry.hpp"
#include "pinocchio/algorithm/model.hpp"
#include "pinocchio/algorithm/jacobian.hpp"
#include "pinocchio/algorithm/frames.hpp" // Jacobian Frame Computation
#include "pinocchio/algorithm/aba.hpp" // Articulated Body Algorithm
#include "pinocchio/algorithm/crba.hpp" // Composite Rigid Body Algorithm
#include "pinocchio/algorithm/rnea.hpp" // Recursive Newton-Euler Algorithm
#include "pinocchio/algorithm/center-of-mass.hpp" // Center-of-Mass related algorithms
#include "pinocchio/algorithm/center-of-mass-derivatives.hpp"
// Parsers
#include "pinocchio/parsers/urdf.hpp"
#include "pinocchio/parsers/srdf.hpp"
// Spatial
#include "pinocchio/spatial/fcl-pinocchio-conversions.hpp"

// Utils
// #include "pinocchio/utils/timer.hpp"

// Standard
#include <math.h>
#include <map>
#include <iostream>
#include <boost/shared_ptr.hpp>

// Macros
#define VAL_MODEL_NUM_FLOATING_JOINTS 7 // 3 for x,y,z and 4 for qx, qy, qz, qw
#define VAL_MODEL_JOINT_INDX_OFFSET 2 //pinocchio attaches a universe joint and a root joint that we need to remove

class RobotModel{
private:

  // private temporary index value holders. Prevents memory allocation on runtime
  pinocchio::FrameIndex tmp_frame_index;
  pinocchio::JointIndex tmp_joint_index;
  void buildPinocchioModel(const std::string & filename);
  void buildPinocchioGeomModel(const std::string & filename, const std::string & meshDir);
  void commonInitialization(bool geom_data_flag=false);

  bool updateGeomWithKinematics = false;

  bool srdf_bool = false;

public:

  std::string srdf_filename;
  pinocchio::Model model;
  pinocchio::GeometryModel geomModel;
  pinocchio::fcl::CollisionResult result;
  pinocchio::fcl::DistanceResult dresult;
  std::vector<pinocchio::fcl::Contact> contacts;

  std::unique_ptr<pinocchio::Data> data;
  std::unique_ptr<pinocchio::GeometryData> geomData;

  Eigen::MatrixXd A; // inertia matrix
  Eigen::MatrixXd Ainv; // inverse of the inertia matrix
  Eigen::MatrixXd C; // coriolis matrix
  Eigen::VectorXd g; // gravity vector

  Eigen::VectorXd q_current;
  Eigen::VectorXd q_lower_pos_limit;
  Eigen::VectorXd q_upper_pos_limit;
  std::vector<std::string> joint_names;

  Eigen::Vector3d x_com; // com position
  Eigen::Vector3d xdot_com; // com velocity
  Eigen::Vector3d xddot_com; // com acceleration

  Eigen::Matrix3Xd J_com; // Center of mass Jacobian
  Eigen::Matrix3Xd Jdot_com; // Center of mass Jacobian Derivative

  RobotModel();
  RobotModel(const std::string & filename);
  RobotModel(const std::string & filename, const std::string & meshDir);
  RobotModel(const std::string & filename, const std::string & meshDir, const std::string & srdf);

  ~RobotModel();

  void enableUpdateGeomOnKinematicsUpdate(bool enable);

  void common_initialization(bool geom_data_flag=false);

  /* updateGeometry
  Input: a vector of configuration with dimension model.nq to update the kinematics.
  */
  void updateGeometry(const Eigen::VectorXd & q_update);

  /* updateFullKinematics
  Input: a vector of configuration with dimension model.nq to update the kinematics.
  */
  void updateFullKinematics(const Eigen::VectorXd & q_update);

  // updates the kinematic derivatives
  void updateKinematicsDerivatives(const Eigen::VectorXd & q_update, const Eigen::VectorXd & qdot_update, const Eigen::VectorXd & qddot_update);

  // updates the
  void updateJointJacobiansDerivatives(const Eigen::VectorXd & q_update, const Eigen::VectorXd & qdot_update);

  /* get6DTaskJacobian
  Input: the frame name.
  Output: the 6D task jacobian (dimension 6 x model.nv) expressed in world frame
          with the linear components first: [(dx/dq)^T , (dy/dq)^T, (dz/dq)^T]^T
          then the rotational components: [(wx/dq)^T , (wy/dq)^T, (wz/dq)^T]^T
  */
  void get6DTaskJacobian(const std::string & frame_name, Eigen::MatrixXd & J_out);

  /* get6DTaskJacobianDot
  Input: the frame name.
  Output: the 6D task jacobian dot (6 x model.nv) expressed in world frame.
  */
  void get6DTaskJacobianDot(const std::string & frame_name, Eigen::MatrixXd & Jdot_out);

  /* get6DTaskJacobianLocal
  Input: the frame name.
  Output: the 6D task jacobian (dimension 6 x model.nv) expressed in the local frame
          with the linear components first: [(dx/dq)^T , (dy/dq)^T, (dz/dq)^T]^T
          then the rotational components: [(wx/dq)^T , (wy/dq)^T, (wz/dq)^T]^T
  */
  void get6DTaskJacobianLocal(const std::string & frame_name, Eigen::MatrixXd & J_out);

  /* get6DTaskJacobianDotLocal
  Input: the frame name.
  Output: the 6D task jacobian dot (6 x model.nv) expressed in world frame.
  */
  void get6DTaskJacobianDotLocal(const std::string & frame_name, Eigen::MatrixXd & Jdot_out);

  /* getFrameWorldPose
  Input: the frame name.
  Output: the position and orientation of the frame with respect to world.
  */

  void getFrameWorldPose(const std::string & name, Eigen::Vector3d & pos, Eigen::Quaternion<double> & ori);

  /* getJointIndex
  Input std::string name
  Output: The joint index in configuration space
      For Valkyrie, the first 7 joints are floating base joints. 3 for the linear and 4 for the rotation described as a quaternion.
      Per C++ standard, the index starts at 0.

      Thus, for a "leftHipYaw" input, which is the first joint in the kinematic chain,
      this function should return 7.
  */
  int getJointIndex(const std::string & name);

  /* getJointIndexNoFloatingJoints
  Input std::string name
  Output: The joint index in configuration space
      unlike getJointIndex, this one does not include the floating joints

      Thus, for a "leftHipYaw" input, which is the first joint in the kinematic chain,
      this function should return 0.
  */
  int getJointIndexNoFloatingJoints(const std::string & name);

  /* getDimQ();
  Output: The dimension of the configuration space
  */
  int getDimQ();
  /* getDimQdot();
  Output: The dimension of the configuration space
  */
  int getDimQdot();

  // Tangent space forward integration given q_start, qdot*dt
  void forwardIntegrate(const Eigen::VectorXd & q_start, const Eigen::VectorXd & qdotDt, Eigen::VectorXd & q_post);

  // Computes the inertia matrix dim(model.nv x model.nv) given configuration and stores the result.
  void computeInertiaMatrix(const Eigen::VectorXd & q);
  // Computes the inverse of the inertia matrix dim(model.nv x model.nv) given configuration and stores the result.
  void computeInertiaMatrixInverse(const Eigen::VectorXd & q);
  // Computes the coriolis matrix dim(model.nv x model.nv) given configuration and velocity and stores the result.
  void computeCoriolisMatrix(const Eigen::VectorXd & q, const Eigen::VectorXd & qdot);
  // Computes the gravity vector dim(model.nv) given configuration and stores the result
  void computeGravityVector(const Eigen::VectorXd & q);

  // CoM positions, velocities, and accelerations are expressed in world frame.
  // Computes the CoM position given q
  void computeCoMPos(const Eigen::MatrixBase<Eigen::VectorXd> & q);
  // Computes the CoM position and velocity given q and qdot
  void computeCoMPosVel(const Eigen::MatrixBase<Eigen::VectorXd> & q, const Eigen::MatrixBase<Eigen::VectorXd> & qdot);
  // Computes the CoM position, velocity, and acceleration given q, qdot, and qddot
  void computeCoMPosVelAcc(const Eigen::MatrixBase<Eigen::VectorXd> & q, const Eigen::MatrixBase<Eigen::VectorXd> & qdot, const Eigen::MatrixBase<Eigen::VectorXd> & qddot);

  // Computes the Jacobian and the position of the Center of Mass of the robot expressed in world frame
  //  Assumes that pinocchio::forwardKinematics() has been called.
  //  the computation also runs a computeJointJacobians() routine.
  //  value is stored in J_com;
  void computeCoMJacobian();
  // Computes the Jacobian and the position of the Center of Mass of the robot expressed in world frame
  // the computation also runs a computeJointJacobians() routine.
  // value is stored in J_com;
  void computeCoMJacobian(const Eigen::MatrixBase<Eigen::VectorXd> & q);

  // Computes the value of Jdot_com. Assumes that updateKinematicsDerivatives was called.
  void computeCoMJacobianDot(const Eigen::MatrixBase<Eigen::VectorXd> & q, const Eigen::MatrixBase<Eigen::VectorXd> & qdot);

  void printJointNames();
  void printFrameNames();

};

#endif