/*
 * Copyright (C) 2014 Open Source Robotics Foundation
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 *
*/
#include <string.h>
#include <boost/algorithm/string.hpp>
#include <ignition/math/MassMatrix3.hh>
#include <ignition/math/Pose3.hh>
#include <ignition/math/Vector3Stats.hh>

#include "gazebo/msgs/msgs.hh"
#include "gazebo/physics/physics.hh"
#include "gazebo/test/ServerFixture.hh"
#include "gazebo/test/helper_physics_generator.hh"

using namespace gazebo;

const double g_tolerance = 1e-4;

class PhysicsLinkTest : public ServerFixture,
                        public testing::WithParamInterface<const char*>
{
  /// \brief Test force adding functions.
  /// \param[in] _physicsEngine Type of physics engine to use.
  public: void AddForce(const std::string &_physicsEngine);

  /// \brief Use AddLinkForce on the given direction and then the opposite
  /// direction so they cancel out.
  /// \param[in] _world World pointer.
  /// \param[in] _link Link pointer.
  /// \param[in] _force Force expressed in link frame.
  /// \param[in] _offset Offset expressed in link frame, defaults to link
  /// origin.
  public: void AddLinkForceTwoWays(physics::WorldPtr _world,
      physics::LinkPtr _link, ignition::math::Vector3d _force,
      ignition::math::Vector3d _offset = ignition::math::Vector3d::Zero);

  /// \brief Test GetWorldAngularMomentum.
  /// \param[in] _physicsEngine Type of physics engine to use.
  public: void GetWorldAngularMomentum(const std::string &_physicsEngine);

  /// \brief Test GetWorldEnergy* functions.
  /// \param[in] _physicsEngine Type of physics engine to use.
  public: void GetWorldEnergy(const std::string &_physicsEngine);

  /// \brief Test Link::GetWorldInertia* functions.
  /// \param[in] _physicsEngine Physics engine to use.
  public: void GetWorldInertia(const std::string &_physicsEngine);

  /// \brief Test wrench subscriber.
  /// \param[in] _physicsEngine Type of physics engine to use.
  public: void OnWrenchMsg(const std::string &_physicsEngine);

  /// \brief Test velocity setting functions.
  /// \param[in] _physicsEngine Type of physics engine to use.
  public: void SetVelocity(const std::string &_physicsEngine);
};

/////////////////////////////////////////////////
void PhysicsLinkTest::AddLinkForceTwoWays(physics::WorldPtr _world,
    physics::LinkPtr _link, ignition::math::Vector3d _force,
    ignition::math::Vector3d _offset)
{
  // Get state before adding force
  ignition::math::Vector3d linearVelWorld0 = _link->WorldCoGLinearVel();
  ignition::math::Vector3d angularVelWorld0 = _link->WorldAngularVel();
  ignition::math::Pose3d poseWorld0 = _link->WorldPose();

  // Add Link Force
  if (_offset == ignition::math::Vector3d::Zero)
    _link->AddLinkForce(_force);
  else
    _link->AddLinkForce(_force, _offset);

  double dt = _world->Physics()->GetMaxStepSize();
  _world->Step(1);

  int moreThanOneStep = 2;

  // Check force and torque (at CoG?) in world frame
  ignition::math::Vector3d forceWorld = poseWorld0.Rot().RotateVector(_force);
  EXPECT_EQ(forceWorld, _link->WorldForce());

  ignition::math::Vector3d worldOffset = poseWorld0.Rot().RotateVector(
      _offset - _link->GetInertial()->CoG());
  ignition::math::Vector3d torqueWorld = worldOffset.Cross(forceWorld);
  EXPECT_EQ(torqueWorld, _link->WorldTorque());

  // Check acceleration in world frame
  ignition::math::Vector3d oneStepLinearAccel =
      forceWorld/_link->GetInertial()->Mass();
  EXPECT_EQ(oneStepLinearAccel, _link->WorldLinearAccel());

  // Compute angular accel by multiplying world torque
  // by inverse of world inertia matrix.
  // In this case, the gyroscopic coupling terms are zero
  // since the model is a unit box.
  ignition::math::Vector3d oneStepAngularAccel =
      _link->WorldInertiaMatrix().Inverse() * torqueWorld;
  EXPECT_EQ(oneStepAngularAccel, _link->WorldAngularAccel());

  // Check velocity in world frame
  ignition::math::Vector3d oneStepLinearVel = linearVelWorld0 +
    dt*oneStepLinearAccel;
  EXPECT_EQ(oneStepLinearVel, _link->WorldCoGLinearVel());

  ignition::math::Vector3d oneStepAngularVel = angularVelWorld0 +
    dt*oneStepAngularAccel;
  EXPECT_EQ(oneStepAngularVel, _link->WorldAngularVel());

  // Step forward and check again
  _world->Step(moreThanOneStep);

  // Check that force and torque are zero
  EXPECT_EQ(ignition::math::Vector3d::Zero, _link->WorldForce());
  EXPECT_EQ(ignition::math::Vector3d::Zero, _link->WorldTorque());

  // Check that acceleration is zero
  EXPECT_EQ(ignition::math::Vector3d::Zero, _link->WorldLinearAccel());
  EXPECT_EQ(ignition::math::Vector3d::Zero, _link->WorldAngularAccel());

  // Check that velocity hasn't changed
  EXPECT_EQ(oneStepLinearVel, _link->WorldCoGLinearVel());
  EXPECT_EQ(oneStepAngularVel, _link->WorldAngularVel());

  // Add opposing force in link frame and check that link is back to initial
  // velocity
  if (_offset == ignition::math::Vector3d::Zero)
    _link->AddLinkForce(-_force);
  else
    _link->AddLinkForce(-_force, _offset);

  _world->Step(moreThanOneStep);
  EXPECT_EQ(ignition::math::Vector3d::Zero, _link->WorldForce());
  EXPECT_EQ(ignition::math::Vector3d::Zero, _link->WorldTorque());
  EXPECT_EQ(linearVelWorld0, _link->WorldCoGLinearVel());
  EXPECT_EQ(angularVelWorld0, _link->WorldAngularVel());
  EXPECT_EQ(ignition::math::Vector3d::Zero, _link->WorldLinearAccel());
  EXPECT_EQ(ignition::math::Vector3d::Zero, _link->WorldAngularAccel());
}

/////////////////////////////////////////////////
void PhysicsLinkTest::AddForce(const std::string &_physicsEngine)
{
  // TODO bullet, dart and simbody currently fail this test
  if (_physicsEngine != "ode")
  {
    gzerr << "Aborting AddForce test for Bullet, DART and Simbody. "
          << "See issues #1476, #1477, and #1478."
          << std::endl;
    return;
  }

  Load("worlds/blank.world", true, _physicsEngine);
  physics::WorldPtr world = physics::get_world("default");
  ASSERT_TRUE(world != NULL);

  // check the physics engine
  physics::PhysicsEnginePtr physics = world->Physics();
  ASSERT_TRUE(physics != NULL);
  EXPECT_EQ(physics->GetType(), _physicsEngine);
  double dt = physics->GetMaxStepSize();
  EXPECT_GT(dt, 0);

  // disable gravity
  world->SetGravity(ignition::math::Vector3d::Zero);

  // Spawn a box
  ignition::math::Vector3d size(1, 1, 1);
  SpawnBox("box", size, ignition::math::Vector3d::Zero,
      ignition::math::Vector3d::Zero, false);
  physics::ModelPtr model = world->ModelByName("box");
  ASSERT_TRUE(model != NULL);
  physics::LinkPtr link = model->GetLink();
  ASSERT_TRUE(link != NULL);

  // Check that link is at rest
  EXPECT_EQ(ignition::math::Vector3d::Zero, link->WorldLinearVel());
  EXPECT_EQ(ignition::math::Vector3d::Zero, link->WorldAngularVel());
  EXPECT_EQ(ignition::math::Vector3d::Zero, link->WorldLinearAccel());
  EXPECT_EQ(ignition::math::Vector3d::Zero, link->WorldAngularAccel());

  // Add force at link frame
  gzdbg << "World == link == inertial frames, no offset" << std::endl;
  EXPECT_EQ(ignition::math::Pose3d::Zero, link->WorldPose());
  EXPECT_EQ(ignition::math::Pose3d::Zero, link->WorldInertialPose());
  this->AddLinkForceTwoWays(world, link, ignition::math::Vector3d(1, 20, 31));

  gzdbg << "World != link == inertial frames, no offset" << std::endl;
  model->SetLinkWorldPose(ignition::math::Pose3d(
        ignition::math::Vector3d(2, 3, 4),
        ignition::math::Quaterniond(0, IGN_PI/2.0, 1)), link);
  EXPECT_NE(ignition::math::Pose3d::Zero, link->WorldPose());
  EXPECT_EQ(link->WorldPose(), link->WorldInertialPose());
  this->AddLinkForceTwoWays(world, link, ignition::math::Vector3d(-1, 10, 5));

  gzdbg << "World == link == inertial frames, with offset" << std::endl;
  model->SetLinkWorldPose(ignition::math::Pose3d::Zero, link);
  EXPECT_EQ(ignition::math::Pose3d::Zero, link->WorldPose());
  EXPECT_EQ(ignition::math::Pose3d::Zero, link->WorldInertialPose());
  this->AddLinkForceTwoWays(world, link, ignition::math::Vector3d(5, 4, 3),
      ignition::math::Vector3d(-2, 1, 0));

  gzdbg << "World == link != inertial frames, no offset" << std::endl;
  model->SetLinkWorldPose(ignition::math::Pose3d::Zero, link);
  ignition::math::Pose3d inertialPose = ignition::math::Pose3d(
      ignition::math::Vector3d(1, 5, 8),
      ignition::math::Quaterniond(IGN_PI/3.0, IGN_PI*1.5, IGN_PI/4));
  link->GetInertial()->SetCoG(inertialPose);
  EXPECT_EQ(ignition::math::Pose3d::Zero, link->WorldPose());
  EXPECT_EQ(inertialPose, link->WorldInertialPose());
  this->AddLinkForceTwoWays(world, link, ignition::math::Vector3d(1, 2, 1));

  gzdbg << "World != link != inertial frames, with offset" << std::endl;
  model->SetLinkWorldPose(ignition::math::Pose3d(
        ignition::math::Vector3d(5, 10, -4),
        ignition::math::Quaterniond(0, IGN_PI/2.0, IGN_PI/6)), link);
  inertialPose = ignition::math::Pose3d(ignition::math::Vector3d(0, -5, 10),
      ignition::math::Quaterniond(0, 2.0*IGN_PI, IGN_PI/3));
  link->GetInertial()->SetCoG(inertialPose);
  this->AddLinkForceTwoWays(world, link, ignition::math::Vector3d(1, 2, 1),
      ignition::math::Vector3d(-2, 0.5, 1));

  gzdbg << "World != link != inertial frames, with offset and initial vel"
      << std::endl;
  model->SetLinkWorldPose(ignition::math::Pose3d(
        ignition::math::Vector3d(-1.5, 0.8, 3),
        ignition::math::Quaterniond(-IGN_PI/4.5, IGN_PI/3.0, IGN_PI*1.2)),
        link);
  inertialPose = ignition::math::Pose3d(ignition::math::Vector3d(1, 0, -5.6),
      ignition::math::Quaterniond(IGN_PI/9, 0, IGN_PI*3));
  link->GetInertial()->SetCoG(inertialPose);
  link->SetLinearVel(ignition::math::Vector3d(2, -0.1, 5));
  link->SetAngularVel(ignition::math::Vector3d(-IGN_PI/10, 0, 0.0001));
  this->AddLinkForceTwoWays(world, link, ignition::math::Vector3d(-3, 2.5, -15),
      ignition::math::Vector3d(-6, -1, -0.2));
}

/////////////////////////////////////////////////
// GetWorldAngularMomentum:
// Spawn box and verify Link::GetWorldAngularMomentum functions
// Make dimensions unequal and give angular velocity that causes
// gyroscopic tumbling.
void PhysicsLinkTest::GetWorldAngularMomentum(const std::string &_physicsEngine)
{
  // Load a blank world (no ground plane)
  Load("worlds/blank.world", true, _physicsEngine);
  auto world = physics::get_world("default");
  ASSERT_TRUE(world != NULL);

  // Verify physics engine type
  auto physics = world->Physics();
  ASSERT_TRUE(physics != NULL);
  EXPECT_EQ(physics->GetType(), _physicsEngine);

  // disable gravity
  world->SetGravity(ignition::math::Vector3d::Zero);

  physics::ModelPtr model;
  {
    // Box size
    const double dx = 0.1;
    const double dy = 0.4;
    const double dz = 0.9;
    const double mass = 10.0;

    msgs::Model msgModel;
    msgModel.set_name(this->GetUniqueString("model"));
    msgs::AddBoxLink(msgModel, mass, ignition::math::Vector3d(dx, dy, dz));
    model = this->SpawnModel(msgModel);
  }
  ASSERT_TRUE(model != NULL);

  // inertia matrix, recompute if dimensions change
  const double Ixx = 0.80833333;
  const double Iyy = 0.68333333;
  const double Izz = 0.14166667;
  const ignition::math::Matrix3d I0(
      Ixx, 0.0, 0.0 , 0.0, Iyy, 0.0 , 0.0, 0.0, Izz);

  // Since Ixx > Iyy > Izz,
  // angular velocity with large y component
  // will cause gyroscopic tumbling
  const ignition::math::Vector3d w0(1e-3, 1.5e0, 1.5e-2);
  model->SetAngularVel(w0);

  // Get link and verify inertia and initial velocity
  auto link = model->GetLink();
  ASSERT_TRUE(link != NULL);
  ASSERT_EQ(w0, link->WorldAngularVel());
  ASSERT_EQ(I0, link->WorldInertiaMatrix());

  // Compute initial angular momentum
  const auto H0((I0 * w0));
  ASSERT_EQ(H0, link->WorldAngularMomentum());
  const double H0mag = H0.Length();

  ignition::math::Vector3Stats angularMomentumError;
  const std::string stat("maxAbs");
  EXPECT_TRUE(angularMomentumError.InsertStatistic(stat));
  const int steps = 5000;
  for (int i = 0; i < steps; ++i)
  {
    world->Step(1);
    auto H = link->WorldAngularMomentum();
    angularMomentumError.InsertData((H - H0) / H0mag);
  }
  EXPECT_LT(angularMomentumError.Mag().Map()[stat], g_tolerance * 10);

  RecordProperty("engine", _physicsEngine);
  this->Record("angularMomentumError", angularMomentumError);
}

/////////////////////////////////////////////////
void PhysicsLinkTest::GetWorldEnergy(const std::string &_physicsEngine)
{
  Load("worlds/empty.world", true, _physicsEngine);
  physics::WorldPtr world = physics::get_world("default");
  ASSERT_TRUE(world != NULL);

  // check the physics engine
  physics::PhysicsEnginePtr physics = world->Physics();
  ASSERT_TRUE(physics != NULL);
  EXPECT_EQ(physics->GetType(), _physicsEngine);
  double dt = physics->GetMaxStepSize();
  EXPECT_GT(dt, 0);

  // Get gravity magnitude
  double g = world->Gravity().Length();

  // Spawn a box
  double z0 = 10.0;
  ignition::math::Vector3d size(1, 1, 1);
  ignition::math::Vector3d pos0(0, 0, z0 + size.Z() / 2);
  SpawnBox("box", size, pos0, ignition::math::Vector3d::Zero, false);
  physics::ModelPtr model = world->ModelByName("box");
  ASSERT_TRUE(model != NULL);
  physics::LinkPtr link = model->GetLink();
  ASSERT_TRUE(link != NULL);

  // Get initial energy
  double energy0 = link->GetWorldEnergy();
  EXPECT_NEAR(link->GetWorldEnergyKinetic(), 0, g_tolerance);

  double totalTime = sqrt(2*z0/g)*0.95;
  unsigned int stepSize = 10;
  unsigned int steps = floor(totalTime / (dt*stepSize));
  for (unsigned int i = 0; i < steps; ++i)
  {
    world->Step(stepSize);
    double energy = link->GetWorldEnergy();
    EXPECT_NEAR(energy / energy0, 1.0, g_tolerance * 10);
  }
}

/////////////////////////////////////////////////
// GetWorldInertia:
// Spawn boxes and verify Link::GetWorldInertia* functions
void PhysicsLinkTest::GetWorldInertia(const std::string &_physicsEngine)
{
  // Load a blank world (no ground plane)
  Load("worlds/blank.world", true, _physicsEngine);
  auto world = physics::get_world("default");
  ASSERT_TRUE(world != NULL);

  // Verify physics engine type
  auto physics = world->Physics();
  ASSERT_TRUE(physics != NULL);
  EXPECT_EQ(physics->GetType(), _physicsEngine);

  // disable gravity
  world->SetGravity(ignition::math::Vector3d::Zero);

  // Box size
  const double dx = 1.0;
  const double dy = 4.0;
  const double dz = 9.0;
  const double mass = 10.0;
  const double angle = IGN_PI / 3.0;

  const unsigned int testCases = 5;
  for (unsigned int i = 0; i < testCases; ++i)
  {
    // Use msgs::AddBoxLink
    msgs::Model msgModel;
    ignition::math::Pose3d modelPose, linkPose, inertialPose;

    msgModel.set_name(this->GetUniqueString("model"));
    msgs::AddBoxLink(msgModel, mass, ignition::math::Vector3d(dx, dy, dz));
    modelPose.Pos().X() = i * dz;
    modelPose.Pos().Z() = dz;

    // i=0: rotated model pose
    //  expect inertial pose to match model pose
    if (i == 0)
    {
      modelPose.Rot().Euler(0.0, 0.0, angle);
    }
    // i=1: rotated link pose
    //  expect inertial pose to match link pose
    else if (i == 1)
    {
      linkPose.Rot().Euler(0.0, 0.0, angle);
    }
    // i=2: rotated inertial pose
    //  expect inertial pose to differ from link pose
    else if (i == 2)
    {
      inertialPose.Rot().Euler(0.0, 0.0, angle);
    }
    // i=3: off-diagonal terms in inertia matrix
    //  expect inertial pose to differ from link pose
    else if (i == 3)
    {
      ignition::math::MassMatrix3d m;
      EXPECT_TRUE(m.SetFromBox(mass, ignition::math::Vector3d(dx, dy, dz),
          ignition::math::Quaterniond(0, 0, angle)));
      msgs::Set(msgModel.mutable_link(0)->mutable_inertial(), m);
    }
    // i=4: offset inertial pose
    //  expect inertial pose to differ from link pose
    else if (i == 4)
    {
      inertialPose.Pos().Set(1, 1, 1);
    }

    {
      auto msgLink = msgModel.mutable_link(0);
      auto msgInertial = msgLink->mutable_inertial();

      msgs::Set(msgModel.mutable_pose(), modelPose);
      msgs::Set(msgLink->mutable_pose(), linkPose);
      msgs::Set(msgInertial->mutable_pose(), inertialPose);
    }

    auto model = this->SpawnModel(msgModel);
    ASSERT_TRUE(model != NULL);

    auto link = model->GetLink();
    ASSERT_TRUE(link != NULL);

    EXPECT_EQ(model->WorldPose(), modelPose);
    EXPECT_EQ(link->WorldPose(), linkPose + modelPose);
    // only check inertial position
    // inertial rotation can vary (bullet for example)
    EXPECT_EQ(link->WorldInertialPose().Pos(),
              (inertialPose + linkPose + modelPose).Pos());

    // i=0: rotated model pose
    //  expect inertial pose to match model pose
    if (i == 0)
    {
      EXPECT_EQ(model->WorldPose(), link->WorldInertialPose());
    }
    // i=1: rotated link pose
    //  expect inertial pose to match link pose
    else if (i == 1)
    {
      EXPECT_EQ(link->WorldPose(), link->WorldInertialPose());
    }
    // i=2: rotated inertial pose
    // i=3: off-diagonal inertia matrix terms
    //  expect center-of-mass positions to match
    else if (i == 2 || i == 3)
    {
      EXPECT_EQ(link->WorldPose().Pos(),
                link->WorldInertialPose().Pos());
    }
    // i=4: offset inertial pose
    //  expect inertial pose to differ from link pose
    else if (i == 4)
    {
      EXPECT_EQ(link->WorldPose().Pos() + inertialPose.Pos(),
                link->WorldInertialPose().Pos());
    }

    // Expect rotated inertia matrix
    ignition::math::Matrix3d inertia = link->WorldInertiaMatrix();
    if (i == 4)
    {
      EXPECT_NEAR(inertia(0, 0), 80.8333, 1e-4);
      EXPECT_NEAR(inertia(1, 1), 68.3333, 1e-4);
      EXPECT_NEAR(inertia(2, 2), 14.1667, 1e-4);
      for (int row = 0; row < 3; ++row)
        for (int col = 0; col < 3; ++col)
          if (row != col)
            EXPECT_NEAR(inertia(row, col), 0.0, g_tolerance);
    }
    else
    {
      EXPECT_NEAR(inertia(0, 0), 71.4583, 1e-4);
      EXPECT_NEAR(inertia(1, 1), 77.7083, 1e-4);
      EXPECT_NEAR(inertia(2, 2), 14.1667, 1e-4);
      EXPECT_NEAR(inertia(0, 1),  5.4126, 1e-4);
      EXPECT_NEAR(inertia(1, 0),  5.4126, 1e-4);
      EXPECT_NEAR(inertia(0, 2), 0, g_tolerance);
      EXPECT_NEAR(inertia(2, 0), 0, g_tolerance);
      EXPECT_NEAR(inertia(1, 2), 0, g_tolerance);
      EXPECT_NEAR(inertia(2, 1), 0, g_tolerance);
    }

    // For 0-3, apply torque and expect equivalent response
    if (i <= 3)
    {
      for (int step = 0; step < 50; ++step)
      {
        link->SetTorque(ignition::math::Vector3d(100, 0, 0));
        world->Step(1);
      }
      if (_physicsEngine.compare("dart") == 0)
      {
        gzerr << "Dart fails this portion of the test (#1090)" << std::endl;
      }
      else
      {
        ignition::math::Vector3d vel = link->WorldAngularVel();
        EXPECT_NEAR(vel.X(),  0.0703, g_tolerance);
        EXPECT_NEAR(vel.Y(), -0.0049, g_tolerance);
        EXPECT_NEAR(vel.Z(),  0.0000, g_tolerance);
      }
    }
  }
}

/////////////////////////////////////////////////
void PhysicsLinkTest::OnWrenchMsg(const std::string &_physicsEngine)
{
  // TODO bullet, dart and simbody currently fail this test
  if (_physicsEngine != "ode")
  {
    gzerr << "Aborting OnWrenchMsg test for Bullet, DART and Simbody. "
          << "Because of issues #1476, #1477, and #1478."
          << std::endl;
    return;
  }

  Load("worlds/blank.world", true, _physicsEngine);
  physics::WorldPtr world = physics::get_world("default");
  ASSERT_TRUE(world != NULL);

  // check the physics engine
  physics::PhysicsEnginePtr physics = world->Physics();
  ASSERT_TRUE(physics != NULL);
  EXPECT_EQ(physics->GetType(), _physicsEngine);
  double dt = physics->GetMaxStepSize();
  EXPECT_GT(dt, 0);

  // disable gravity
  world->SetGravity(ignition::math::Vector3d::Zero);

  // Spawn a box
  ignition::math::Vector3d size(1, 1, 1);
  SpawnBox("box", size, ignition::math::Vector3d::Zero,
                        ignition::math::Vector3d::Zero, false);
  physics::ModelPtr model = world->ModelByName("box");
  ASSERT_TRUE(model != NULL);
  physics::LinkPtr link = model->GetLink();
  ASSERT_TRUE(link != NULL);

  // Check that link is at rest
  EXPECT_EQ(ignition::math::Vector3d::Zero, link->WorldLinearVel());
  EXPECT_EQ(ignition::math::Vector3d::Zero, link->WorldAngularVel());
  EXPECT_EQ(ignition::math::Vector3d::Zero, link->WorldLinearAccel());
  EXPECT_EQ(ignition::math::Vector3d::Zero, link->WorldAngularAccel());
  EXPECT_EQ(ignition::math::Vector3d::Zero, link->WorldForce());
  EXPECT_EQ(ignition::math::Vector3d::Zero, link->WorldTorque());

  // Publish wrench message
  std::string topicName = "~/" + link->GetScopedName() + "/wrench";
  boost::replace_all(topicName, "::", "/");
  transport::PublisherPtr wrenchPub =
    this->node->Advertise<msgs::Wrench>(topicName);

  msgs::Wrench msg;

  std::vector<ignition::math::Vector3d> forces;
  std::vector<ignition::math::Vector3d> torques;
  std::vector<ignition::math::Vector3d> forceOffsets;

  // Only force
  forces.push_back(ignition::math::Vector3d(1, 0, 0));
  torques.push_back(ignition::math::Vector3d::Zero);
  forceOffsets.push_back(ignition::math::Vector3d::Zero);

  // Only force, with an offset
  forces.push_back(ignition::math::Vector3d(5.2, 0.1, 10));
  torques.push_back(ignition::math::Vector3d::Zero);
  forceOffsets.push_back(ignition::math::Vector3d(2.1, 1, -0.6));

  // Only torque
  forces.push_back(ignition::math::Vector3d::Zero);
  torques.push_back(ignition::math::Vector3d(-0.2, 5, 0));
  forceOffsets.push_back(ignition::math::Vector3d::Zero);

  // All fields set
  forces.push_back(ignition::math::Vector3d(5, 6, -0.9));
  torques.push_back(ignition::math::Vector3d(-0.2, 5, 0));
  forceOffsets.push_back(ignition::math::Vector3d(-1, -4, -0.8));

  for (unsigned int i = 0; i < forces.size(); ++i)
  {
    gzdbg << "Testing force: " << forces[i].X() << ", "
                               << forces[i].Y() << ", "
                               << forces[i].Z() <<
                   " torque: " << torques[i].X() << ", "
                               << torques[i].Y() << ", "
                               << torques[i].Z() <<
             " force offset: " << forceOffsets[i].X() << ", "
                               << forceOffsets[i].Y() << ", "
                               << forceOffsets[i].Z() << std::endl;

    // Publish message
    msgs::Set(msg.mutable_force(), forces[i]);
    msgs::Set(msg.mutable_torque(), torques[i]);
    // Leave optional field unset if it's zero
    if (forceOffsets[i] != ignition::math::Vector3d::Zero)
      msgs::Set(msg.mutable_force_offset(), forceOffsets[i]);

    wrenchPub->Publish(msg);

    // Calculate expected values
    ignition::math::Vector3d forceWorld = forces[i];
    ignition::math::Vector3d worldOffset =
      forceOffsets[i] - link->GetInertial()->CoG();
    ignition::math::Vector3d torqueWorld =
      worldOffset.Cross(forces[i]) + torques[i];

    // Wait for message to be received
    while (link->WorldForce() != forceWorld ||
           link->WorldTorque() != torqueWorld)
    {
      world->Step(1);
      common::Time::MSleep(1);
    }

    // Check force and torque (at CoG?) in world frame
    EXPECT_EQ(link->WorldForce(), forceWorld);
    EXPECT_EQ(link->WorldTorque(), torqueWorld);

    // Reset link's physics states
    link->ResetPhysicsStates();
    link->SetWorldPose(ignition::math::Pose3d());
    world->Step(1);

    EXPECT_EQ(ignition::math::Vector3d::Zero, link->WorldLinearVel());
    EXPECT_EQ(ignition::math::Vector3d::Zero, link->WorldAngularVel());
    EXPECT_EQ(ignition::math::Vector3d::Zero, link->WorldLinearAccel());
    EXPECT_EQ(ignition::math::Vector3d::Zero, link->WorldAngularAccel());
    EXPECT_EQ(ignition::math::Vector3d::Zero, link->WorldForce());
    EXPECT_EQ(ignition::math::Vector3d::Zero, link->WorldTorque());
  }
}

/////////////////////////////////////////////////
void PhysicsLinkTest::SetVelocity(const std::string &_physicsEngine)
{
  Load("worlds/blank.world", true, _physicsEngine);
  physics::WorldPtr world = physics::get_world("default");
  ASSERT_TRUE(world != NULL);

  // check the physics engine
  physics::PhysicsEnginePtr physics = world->Physics();
  ASSERT_TRUE(physics != NULL);
  EXPECT_EQ(physics->GetType(), _physicsEngine);
  double dt = physics->GetMaxStepSize();
  EXPECT_GT(dt, 0);

  // disable gravity
  world->SetGravity(ignition::math::Vector3d::Zero);

  // Spawn a box
  ignition::math::Vector3d size(1, 1, 1);
  ignition::math::Vector3d pos0(0, 0, 1);
  SpawnBox("box", size, pos0, ignition::math::Vector3d::Zero, false);
  physics::ModelPtr model = world->ModelByName("box");
  ASSERT_TRUE(model != NULL);
  physics::LinkPtr link = model->GetLink();
  ASSERT_TRUE(link != NULL);

  // Set upward velocity and check
  ignition::math::Vector3d vel(0, 0, 1);
  link->SetLinearVel(vel);
  world->Step(1);
  EXPECT_EQ(vel, link->WorldLinearVel());
  EXPECT_EQ(ignition::math::Vector3d::Zero, link->WorldAngularVel());

  // Step forward and check velocity again
  world->Step(44);
  double time = world->SimTime().Double();
  EXPECT_EQ(vel, link->WorldLinearVel());
  EXPECT_EQ(ignition::math::Vector3d::Zero, link->WorldAngularVel());

  // check position
  ignition::math::Vector3d pos = link->WorldPose().Pos();
  EXPECT_EQ(pos0 + time*vel, pos);

  // Set velocity to zero
  link->SetLinearVel(ignition::math::Vector3d::Zero);
  world->Step(1);
  EXPECT_EQ(ignition::math::Vector3d::Zero, link->WorldLinearVel());
  EXPECT_EQ(ignition::math::Vector3d::Zero, link->WorldAngularVel());
  EXPECT_EQ(pos0 + time*vel, pos);

  // Start translating and rotating
  vel.Set(1, 1, 0);
  ignition::math::Vector3d vel2(0, 2.0, 0);
  link->SetLinearVel(vel);
  link->SetAngularVel(vel2);

  // Step once
  world->Step(1);
  EXPECT_EQ(vel, link->WorldLinearVel());
  EXPECT_EQ(vel2, link->WorldAngularVel());

  // test linear velocity at specific point in space
  ignition::math::Vector3d offset(0, 0, -0.5);
  ignition::math::Vector3d vel3 = link->WorldLinearVel(offset,
      ignition::math::Quaterniond::Identity);
  EXPECT_NEAR(vel3.X(), 0.0, g_tolerance);
  EXPECT_NEAR(vel3.Y(), 1.0, g_tolerance);
  EXPECT_NEAR(vel3.Z(), 0.0, g_tolerance);

  // check rotation
  ignition::math::Vector3d rpy = link->WorldPose().Rot().Euler();
  EXPECT_NEAR(rpy.X(), 0.0, g_tolerance);
  EXPECT_NEAR(rpy.Y(), vel2.Y()*dt, g_tolerance);
  EXPECT_NEAR(rpy.Z(), 0.0, g_tolerance);
}

/////////////////////////////////////////////////
TEST_P(PhysicsLinkTest, AddForce)
{
  AddForce(GetParam());
}

/////////////////////////////////////////////////
TEST_P(PhysicsLinkTest, GetWorldAngularMomentum)
{
  GetWorldAngularMomentum(GetParam());
}

/////////////////////////////////////////////////
TEST_P(PhysicsLinkTest, GetWorldEnergy)
{
  GetWorldEnergy(GetParam());
}

/////////////////////////////////////////////////
TEST_P(PhysicsLinkTest, GetWorldInertia)
{
  GetWorldInertia(GetParam());
}

/////////////////////////////////////////////////
TEST_P(PhysicsLinkTest, OnWrenchMsg)
{
  OnWrenchMsg(GetParam());
}

/////////////////////////////////////////////////
TEST_P(PhysicsLinkTest, SetVelocity)
{
  SetVelocity(GetParam());
}

INSTANTIATE_TEST_CASE_P(PhysicsEngines, PhysicsLinkTest,
                        PHYSICS_ENGINE_VALUES);

/////////////////////////////////////////////////
int main(int argc, char **argv)
{
  ::testing::InitGoogleTest(&argc, argv);
  return RUN_ALL_TESTS();
}


/**
 * \file metric_space_search.hpp
 *
 * This library provides a class that implements a Dynamic Vantage-Point Tree (DVP-Tree) that
 * allows for O(logN) time nearest-neighbor queries in a metric-space. A DVP-tree is essentially
 * a generalization of a search tree which only requires the space to have a metric which
 * respects the triangular inequality.
 *
 * \author Sven Mikael Persson <mikael.s.persson@gmail.com>
 * \date April 2011
 */

/*
 *    Copyright 2011 Sven Mikael Persson
 *
 *    THIS SOFTWARE IS DISTRIBUTED UNDER THE TERMS OF THE GNU GENERAL PUBLIC LICENSE v3 (GPLv3).
 *
 *    This file is part of ReaK.
 *
 *    ReaK is free software: you can redistribute it and/or modify
 *    it under the terms of the GNU General Public License as published by
 *    the Free Software Foundation, either version 3 of the License, or
 *    (at your option) any later version.
 *
 *    ReaK is distributed in the hope that it will be useful,
 *    but WITHOUT ANY WARRANTY; without even the implied warranty of
 *    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 *    GNU General Public License for more details.
 *
 *    You should have received a copy of the GNU General Public License
 *    along with ReaK (as LICENSE in the root folder).
 *    If not, see <http://www.gnu.org/licenses/>.
 */

#ifndef REAK_METRIC_SPACE_SEARCH_HPP
#define REAK_METRIC_SPACE_SEARCH_HPP

#include <boost/bind.hpp>
#include <boost/lambda/lambda.hpp>
#include <boost/graph/graph_concepts.hpp>
#include <boost/property_map/property_map.hpp>
#include <boost/graph/adjacency_list.hpp>
#include <boost/graph/topology.hpp>
#include <boost/graph/properties.hpp>

#include <map>
#include <vector>
#include "metric_space_concept.hpp"

namespace boost {

  enum vertex_mu_distance_t { vertex_mu_distance };

  BOOST_INSTALL_PROPERTY(vertex, mu_distance);

};

/**
 * This class template is a callable class that can be used to choose the best
 * vantage-point to use out of a set of points. In theory, the best vantage-point
 * is the one which deviates the most from the other points in the set, however,
 * this functor will approximately select that point by searching for it only
 * in a random subset of the given range of points.
 * \tparam RandomNumberGenerator The random number generator type to be use to obtain the randomness needed in the search.
 */
template <typename RandomNumberGenerator = boost::minstd_rand>
class random_best_vp_chooser {
  private:
    unsigned int m_divider;
    RandomNumberGenerator m_rand;
  public:

    /**
     * Default construction.
     * \param aDivider The divider of the set (determines the fraction of the points to search), default is 10.
     */
    random_best_vp_chooser(unsigned int aDivider = 10) : m_divider(aDivider), m_rand(std::time(0)) { };
    /**
     * Parametrized construction.
     * \param aRand The random number generator to use.
     * \param aDivider The divider of the set (determines the fraction of the points to search), default is 10.
     */
    random_best_vp_chooser(const RandomNumberGenerator& aRand, unsigned int aDivider = 10) : m_divider(aDivider), m_rand(aRand) { };

    /**
     * This call-operator will choose a vantage-point from within the given range.
     * \tparam RandomAccessIter A random-access iterator type that can describe the point-range.
     * \tparam Topology The topology type on which the points can reside, should model the MetricSpaceConcept.
     * \tparam PositionMap The property-map type that can map the vertex descriptors (which should be the value-type of the iterators) to a point (position).
     * \param aBegin The start of the range of vertices.
     * \param aEnd The end of the range of vertices (one element past the end).
     * \param aSpace The topology on which the points reside.
     * \param aPosition The property-map used to obtain the positions from the vertices.
     * \return A random-access iterator to the chosen vantage-point.
     */
    template <typename RandomAccessIter, typename Topology, typename PositionMap>
    RandomAccessIter operator() (RandomAccessIter aBegin, RandomAccessIter aEnd, const Topology& aSpace, PositionMap aPosition)
    {
      // BOOST_CONCEPT_ASSERT((MetricSpaceConcept<Topology>)); // not necessary - just need distance() function
      typedef typename metric_topology_traits<Topology>::point_type Point;
      RandomAccessIter best_pt = aEnd;
      double best_dev = -1;
      //std::cout << "Number of loops: " << (aEnd - aBegin) / m_divider + 1 << std::endl;
      int numLoops = (aEnd - aBegin) / m_divider + 1;

      for(unsigned int i=0; i < (aEnd - aBegin) / m_divider + 1;++i) {
        std::cout << i << " out of " << numLoops << std::endl;
	RandomAccessIter current_pt = aBegin + (m_rand() % (aEnd - aBegin));
	double current_mean = 0.0;
	double current_dev = 0.0;
	Point current_vp = get(aPosition, *current_pt);
	for(unsigned int j=0; aBegin + j != aEnd; ++j) {
	  double dist = aSpace.distance(current_vp, get(aPosition, *(aBegin + j)));
	  current_mean = (current_mean * j + dist) / (j + 1);
	  current_dev = (current_dev * j + dist * dist) / (j + 1);
	};
	double current_var = current_dev - current_mean * current_mean;
	if(current_var < 0) current_var = 0.0;
	current_dev = std::sqrt(current_var);

	if(current_dev > best_dev) {
	  best_pt = current_pt;
	  best_dev = current_dev;
	};
      };
      return best_pt;
    };
};

/**
 * This class implements a Dynamic Vantage-Point Tree (DVP-Tree) that
 * allows for O(logN) time nearest-neighbor queries in a metric-space. A DVP-tree is essentially
 * a generalization of a search tree which only requires the space to have a metric which
 * respects the triangular inequality.
 * \tparam Key The key type for the tree, essentially the key value is the vertex descriptor type.
 * \tparam Topology The topology type on which the points can reside, should model the MetricSpaceConcept.
 * \tparam PositionMap The property-map type that can map the vertex descriptors (which should be the value-type of the iterators) to a point (position).
 * \tparam Arity The arity of the tree, e.g., 2 means a binary-tree.
 * \tparam VPChooser The functor type to use to choose the vantage-point out of a set of vertices.
 */
template <typename Key,
          typename Topology,
          typename PositionMap,
	  unsigned int Arity = 2,
	  typename VPChooser = random_best_vp_chooser<> >
class dvp_tree
{
  public:
    //BOOST_CONCEPT_ASSERT((MetricSpaceConcept<Topology>)); // This actually is not necessary, all that is required is a distance() function.

    typedef typename metric_topology_traits<Topology>::point_type point_type;
    typedef typename metric_topology_traits<Topology>::point_difference_type point_difference_type;
    typedef double distance_type;

  private:
    typedef boost::property< boost::vertex_index_t, Key,
            boost::property< boost::vertex_mu_distance_t, distance_type, boost::no_property > > vertex_properties;

    typedef boost::no_property edge_properties;

    typedef boost::adjacency_list< boost::listS, boost::listS, boost::bidirectionalS,
                                   vertex_properties,
	  	                   edge_properties,
		                   boost::listS> tree_indexer;

    typedef boost::adjacency_list_traits<boost::listS,boost::listS,boost::bidirectionalS,boost::listS>::vertex_descriptor vertex_type;
    typedef boost::adjacency_list_traits<boost::listS,boost::listS,boost::bidirectionalS,boost::listS>::edge_descriptor edge_type;
    typedef typename boost::graph_traits<tree_indexer>::out_edge_iterator out_edge_iter;
    typedef typename boost::graph_traits<tree_indexer>::in_edge_iterator in_edge_iter;

    tree_indexer m_tree;
    vertex_type m_root;
    typename boost::property_map< tree_indexer, boost::vertex_index_t >::type m_key;
    typename boost::property_map< tree_indexer, boost::vertex_mu_distance_t >::type m_mu;
    const Topology& m_space;
    PositionMap m_position;
    VPChooser m_vp_chooser;

    //non-copyable.
    dvp_tree(const dvp_tree<Key,Topology,PositionMap,Arity,VPChooser>&);
    dvp_tree<Key,Topology,PositionMap,Arity,VPChooser>& operator=(const dvp_tree<Key,Topology,PositionMap,Arity,VPChooser>&);

    static bool closer(std::map<Key,distance_type>& m, const Key& k1, const Key& k2) {
      return m[k1] < m[k2];
    };

    void construct_node(vertex_type aNode, typename std::vector<Key>::iterator aBegin, typename std::vector<Key>::iterator aEnd, std::map<Key,distance_type>& aDistMap) {
      typedef typename std::vector<Key>::iterator KeyIter;
      KeyIter vp_ind = m_vp_chooser(aBegin, aEnd, m_space, m_position);
      point_type vp_pt = get(m_position, *vp_ind);
      for(KeyIter it = aBegin; it != aEnd; ++it) {
	aDistMap[*it] = m_space.distance(vp_pt, get(m_position, *it));
      };
      std::sort(aBegin,aEnd,boost::bind(closer,aDistMap,_1,_2));
      put(m_key, aNode, *aBegin);
      aBegin++;
      if((aEnd - aBegin) < static_cast<int>(Arity)) {
	for(KeyIter it = aBegin; it != aEnd; ++it) {
	  vertex_type k = boost::add_vertex(m_tree);
	  put(m_key, k, *it);
	  put(m_mu, k, aDistMap[*it]);
	  boost::add_edge(aNode,k,m_tree);
	};
      } else {
	for(unsigned int i=Arity;i>=1;--i) {
	  vertex_type k = boost::add_vertex(m_tree);
	  boost::add_edge(aNode,k,m_tree);
	  unsigned int num_children = (aEnd - aBegin) / i;
	  put(m_mu, k, aDistMap[*(aBegin + (num_children-1))]);
	  construct_node(k,aBegin,aBegin + num_children,aDistMap);
	  aBegin = aBegin + num_children;
	};
      };
    };

    void find_nearest_impl(const point_type& aPoint, distance_type& aSigma, vertex_type aNode, std::multimap<distance_type, Key>& aList, std::size_t K) const {
      using namespace boost;
      typedef typename std::multimap<distance_type, Key>::value_type ListType;
      Key current_key = get(m_key, aNode);
      distance_type current_dist = m_space.distance(aPoint, get(m_position, current_key));
      if(current_dist < aSigma) { //is the vantage point within current search bound? Yes...
        aList.insert(ListType(current_dist, current_key)); //then add the vantage point to the NN list.
	if(aList.size() > K) { //are there too many nearest neighbors? Yes...
	  aList.erase((++aList.rbegin()).base()); //delete last element to keep aList with K elements
	  aSigma = aList.rbegin()->first; //distance of the last element is now the search bound aSigma.
	};
      };
      out_edge_iter ei,ei_end;
      //first, locate the partition in which aPoint is:
      if(out_degree(aNode,m_tree) == 0)
	return;
      for(tie(ei,ei_end) = out_edges(aNode,m_tree); ei != ei_end; ++ei) {
	if(current_dist < get(m_mu, target(*ei,m_tree)))
	  break;
      };
      if(ei == ei_end) --ei; //back-track if the end was reached.
      find_nearest_impl(aPoint,aSigma,target(*ei,m_tree),aList,K); //search in the most likely node.
      out_edge_iter ei_left = ei;
      out_edge_iter ei_right = ei; ++ei_right;
      tie(ei,ei_end) = out_edges(aNode,m_tree); //find the bounds again (start and end).
      while(true) {
	if(ei_left == ei) {
	  out_edge_iter ei_rightleft = ei_right; --ei_rightleft;
	  while((ei_right != ei_end) && (get(m_mu,target(*ei_rightleft,m_tree)) < current_dist + aSigma)) {
	    find_nearest_impl(aPoint,aSigma,target(*ei_right,m_tree),aList,K);
	    ++ei_rightleft; ++ei_right;
	  };
	  break;
	} else if(ei_right == ei_end) {
	  out_edge_iter ei_leftleft = ei_left;
	  while((ei_left != ei) && (get(m_mu,target(*(--ei_leftleft),m_tree)) > current_dist - aSigma)) {
	    find_nearest_impl(aPoint,aSigma,target(*ei_leftleft,m_tree),aList,K);
	    --ei_left;
	  };
	  break;
	} else {
	  out_edge_iter ei_leftleft = ei_left; --ei_leftleft;
	  distance_type d1 = get(m_mu,target(*ei_leftleft,m_tree)) + aSigma - current_dist; //greater than 0 if ei_leftleft should be searched.
	  out_edge_iter ei_rightleft = ei_right; --ei_rightleft;
	  distance_type d2 = get(m_mu,target(*ei_rightleft,m_tree)) - aSigma - current_dist; //less than 0 if ei_right should be searched.
	  if(d1 + d2 > 0) { //this means that ei_leftleft's boundary is closer to aPoint.
            find_nearest_impl(aPoint,aSigma,target(*ei_leftleft,m_tree),aList,K);
	    ei_left = ei_leftleft;
	  } else {
	    find_nearest_impl(aPoint,aSigma,target(*ei_right,m_tree),aList,K);
	    ++ei_right;
	  };
	};
      };

    };

    vertex_type get_leaf(const point_type& aPoint, vertex_type aNode) {
      using namespace boost;
      Key current_key = get(m_key, aNode);
      distance_type current_dist = m_space.distance(aPoint, get(m_position, current_key));
      out_edge_iter ei,ei_end;
      //first, locate the partition in which aPoint is:
      if(out_degree(aNode,m_tree) == 0)
	return aNode;
      vertex_type result = aNode;
      for(tie(ei,ei_end) = out_edges(aNode,m_tree); ei != ei_end; ++ei) {
	result = target(*ei,m_tree);
	if(current_dist < get(m_mu, result))
	  break;
      };
      return get_leaf(aPoint,result);
    };

    vertex_type get_key(Key aVertex, const point_type& aPoint, vertex_type aNode) {
      using namespace boost;
      Key current_key = get(m_key, aNode);
      if(current_key == aVertex) return aNode;
      distance_type current_dist = m_space.distance(aPoint, get(m_position, current_key));
      out_edge_iter ei,ei_end;
      //first, locate the partition in which aPoint is:
      if(out_degree(aNode,m_tree) == 0)
	return vertex_type();
      vertex_type result = aNode;
      for(tie(ei,ei_end) = out_edges(aNode,m_tree); ei != ei_end; ++ei) {
	result = target(*ei,m_tree);
	if(current_dist < get(m_mu, result))
	  break;
      };
      return get_key(aVertex,aPoint,result);
    };

    void update_mu_upwards(const point_type& aPoint, vertex_type aNode) {
      using namespace boost;
      if(aNode == m_root) return;
      vertex_type parent = source(*(in_edges(aNode,m_tree).first),m_tree);
      distance_type dist = m_space.distance(aPoint,get(m_position,get(m_key,parent)));
      if(dist > get(m_mu,aNode))
	put(m_mu,aNode,dist);
      update_mu_upwards(aPoint,parent);
    };

    bool is_leaf_node(vertex_type aNode) {
      using namespace boost;
      if(out_degree(aNode,m_tree) == 0) return true;
      out_edge_iter ei,ei_end;
      for(tie(ei,ei_end) = out_edges(aNode,m_tree); ei != ei_end; ++ei) {
	if(out_degree(target(*ei,m_tree),m_tree) != 0)
	  return false;
      };
      return true;
    };

    bool is_node_full(vertex_type aNode, int& depth_limit) {
      using namespace boost;
      if(depth_limit < 0)
	return false;
      if((out_degree(aNode,m_tree) == 0) && (depth_limit == 0))
	return true;
      --depth_limit;
      if((out_degree(aNode,m_tree) == 0) || (out_degree(aNode,m_tree) < Arity))
	return false;
      out_edge_iter ei,ei_end;
      if(is_leaf_node(aNode)) {
	if(depth_limit == 0)
	  return true;
	else
	  return false;
      };
      for(tie(ei,ei_end) = out_edges(aNode,m_tree); ei != ei_end; ++ei) {
	int new_depth_limit = depth_limit;
	if(!is_node_full(target(*ei,m_tree),new_depth_limit)) {
	  depth_limit = new_depth_limit;
	  return false;
	};
      };
      return true;
    };

    void collect_keys(std::vector<Key>& aList, vertex_type aNode) {
      using namespace boost;
      aList.push_back(get(m_key, aNode));
      out_edge_iter ei,ei_end;
      for(tie(ei,ei_end) = out_edges(aNode,m_tree); ei != ei_end; ++ei)
	collect_keys(aList,target(*ei,m_tree));
    };

    template <typename Predicate>
    void collect_keys(std::vector<Key>& aList, vertex_type aNode, Predicate aShouldAdd) {
      using namespace boost;
      Key k = get(m_key, aNode);
      if(aShouldAdd(k))
        aList.push_back(k);
      out_edge_iter ei,ei_end;
      for(tie(ei,ei_end) = out_edges(aNode,m_tree); ei != ei_end; ++ei)
	collect_keys(aList,target(*ei,m_tree));
    };

    void clear_node(vertex_type aNode) {
      using namespace boost;
      if(out_degree(aNode,m_tree) == 0) return;
      out_edge_iter ei,ei_end;
      std::vector<vertex_type> children;
      children.reserve(out_degree(aNode,m_tree));
      for(tie(ei,ei_end) = out_edges(aNode,m_tree); ei != ei_end; ++ei) {
	vertex_type v = target(*ei,m_tree);
	clear_node(v);
	children.push_back(v);
      };
      for(std::vector<vertex_type>::iterator it = children.begin(); it != children.end(); ++it) {
	remove_edge(aNode, *it,m_tree);
	remove_vertex(*it,m_tree);
      };
    };

  public:

    /**
     * Construct the DVP-tree from a graph, topology and property-map.
     * \tparam Graph The graph type on which the vertices are taken from, should model the boost::VertexListGraphConcept.
     * \param g The graph from which to take the vertices.
     * \param aSpace The topology on which the positions of the vertices reside.
     * \param aPosition The property-map that can be used to obtain the positions of the vertices.
     * \param aVPChooser The vantage-point chooser functor (policy class).
     */
    template <typename Graph>
    dvp_tree(const Graph& g,
	     const Topology& aSpace,
	     PositionMap aPosition,
	     VPChooser aVPChooser = VPChooser()) :
	     m_tree(), m_root(),
             m_key(boost::get(boost::vertex_index,m_tree)),
             m_mu(boost::get(boost::vertex_mu_distance,m_tree)),
             m_space(aSpace), m_position(aPosition), m_vp_chooser(aVPChooser) {
      using namespace boost;
      if(num_vertices(g) == 0) return;

      m_root = add_vertex(m_tree);
      typename graph_traits<Graph>::vertex_iterator vi,vi_end;
      tie(vi,vi_end) = vertices(g);
      std::vector<Key> v(vi,vi_end); //Copy the list of vertices to random access memory.
      std::map<Key,distance_type> dist_map;
      construct_node(m_root, v.begin(), v.end(), dist_map);
    };

    /**
     * Construct the DVP-tree from a range, topology and property-map.
     * \tparam ForwardIterator The forward-iterator type from which the vertices can be obtained.
     * \param aBegin The start of the range from which to take the vertices.
     * \param aEnd The end of the range from which to take the vertices (one-past-last).
     * \param aSpace The topology on which the positions of the vertices reside.
     * \param aPosition The property-map that can be used to obtain the positions of the vertices.
     * \param aVPChooser The vantage-point chooser functor (policy class).
     */
    template <typename ForwardIterator>
    dvp_tree(ForwardIterator aBegin,
	     ForwardIterator aEnd,
	     const Topology& aSpace,
	     PositionMap aPosition,
	     VPChooser aVPChooser = VPChooser()) :
	     m_tree(), m_root(),
             m_key(boost::get(boost::vertex_index,m_tree)),
             m_mu(boost::get(boost::vertex_mu_distance,m_tree)),
             m_space(aSpace), m_position(aPosition), m_vp_chooser(aVPChooser) {
      using namespace boost;
      if(aBegin == aEnd) return;

      m_root = add_vertex(m_tree);
      std::vector<Key> v(aBegin,aEnd); //Copy the list of vertices to random access memory.
      std::map<Key,distance_type> dist_map;
      construct_node(m_root, v.begin(), v.end(), dist_map);
    };

    /**
     * Checks if the DVP-tree is empty.
     * \return True if the DVP-tree is empty.
     */
    bool empty() const { return (boost::num_vertices(m_tree) == 0); };
    /**
     * Returns the size of the DVP-tree (the number of vertices it contains.
     * \return The size of the DVP-tree (the number of vertices it contains.
     */
    std::size_t size() const { return boost::num_vertices(m_tree); };

    /**
     * Inserts a key-value (vertex).
     * \param u The vertex to be added to the DVP-tree.
     */
    void insert(Key u) {
      using namespace boost;
      if(num_vertices(m_tree) == 0) {
	m_root = add_vertex(m_tree);
	put(m_mu,m_root,0.0);
	put(m_key,m_root,u);
	return;
      };
      point_type u_pt = get(m_position, u);
      vertex_type u_realleaf = get_leaf(u_pt,m_root);
      if(u_realleaf == m_root) { //if the root is the leaf, it requires special attention since no parent exists.
	vertex_type u_node = add_vertex(m_tree);
	add_edge(u_realleaf,u_node,m_tree);
	put(m_key, u_node, u);
	put(m_mu, u_node, m_space.distance(u_pt,get(m_position,get(m_key,u_realleaf))));
	update_mu_upwards(u_pt,u_realleaf);
	return;
      };
      vertex_type u_leaf = source(*(in_edges(u_realleaf,m_tree).first),m_tree);
      if((out_degree(u_leaf,m_tree) < Arity) || (!is_leaf_node(u_leaf))) {
	// leaf node is not full of children, an additional child can be added
	//  (must be reconstructed to keep ordering, but this is a trivial operation O(Arity)).
	//OR
	// if leaf is not really a leaf, then it means that this sub-tree is definitely not balanced and not full either,
	//  then all the Keys ought to be collected and u_leaf ought to be reconstructed.
	std::vector<Key> key_list;
	collect_keys(key_list,u_leaf);
	key_list.push_back(u);
	clear_node(u_leaf);
	std::map<Key,distance_type> dist_map;
	construct_node(u_leaf, key_list.begin(), key_list.end(), dist_map);
	update_mu_upwards(u_pt,u_leaf);
      } else {
	//if it is a full-leaf, then this is a leaf node, and it is balanced but full,
	// we should then find a non-full parent.
	vertex_type p = u_leaf;
	int actual_depth_limit = 1;
	int last_depth_limit = actual_depth_limit;
  	while((p != m_root) && (is_node_full(p,last_depth_limit))) {
	  p = source(*(in_edges(p,m_tree).first),m_tree);
	  last_depth_limit = ++actual_depth_limit;
	};
	bool is_p_full = false;
	if(p == m_root)
	  is_p_full = is_node_full(p,last_depth_limit);
	if((!is_p_full) && (last_depth_limit >= 0)) {
	  //this means that we can add our key to the sub-tree of p and reconstruct from there.
	  std::vector<Key> key_list;
	  collect_keys(key_list,p);
	  key_list.push_back(u);
	  clear_node(p);
	  std::map<Key,distance_type> dist_map;
	  construct_node(p, key_list.begin(), key_list.end(), dist_map);
	  update_mu_upwards(u_pt,p);
	} else {
	  //this means that either the root node is full or there are branches of the tree that are deeper than u_realleaf,
	  // and thus, in either case, u_realleaf should be expanded.
	  p = add_vertex(m_tree);
	  add_edge(u_realleaf,p,m_tree);
	  put(m_key, p, u);
	  put(m_mu, p, m_space.distance(u_pt,get(m_position,get(m_key,u_realleaf))));
	  update_mu_upwards(u_pt,u_realleaf);
	};
      };
    };
    /**
     * Inserts a range of key-values (vertices).
     * \tparam ForwardIterator A forward-iterator type that can be used to obtain the vertices.
     * \param aBegin The start of the range from which to take the vertices.
     * \param aEnd The end of the range from which to take the vertices (one-past-last).
     */
    template <typename ForwardIterator>
    void insert(ForwardIterator aBegin, ForwardIterator aEnd) {
      if(boost::num_vertices(m_tree) == 0) return;
      std::for_each(aBegin,aEnd,boost::bind(&dvp_tree<Key,Topology,PositionMap,Arity,VPChooser>::insert_and_find_nearest,this,_1));
      //TODO: There's got to be a better way to insert many elements (most likely a similar strategy to the erase multiple function).
    };

    /**
     * Erases the given vertex from the DVP-tree.
     * \param u The vertex to be removed from the DVP-tree.
     */
    void erase(Key u) {
      using namespace boost;
      if(num_vertices(m_tree) == 0) return;
      point_type u_pt = get(m_position, u);
      vertex_type u_node = get_key(u, u_pt, m_root);
      if(u_node == vertex_type()) return;
      out_edge_iter ei,ei_end;
      std::vector<Key> key_list;
      for(tie(ei,ei_end) = out_edges(u_node,m_tree); ei != ei_end; ++ei)
	collect_keys(key_list,target(*ei,m_tree));
      clear_node(u_node);
      std::map<Key,distance_type> dist_map;
      construct_node(u_node, key_list.begin(), key_list.end(), dist_map);
    };

    /**
     * Erases the given vertex-range from the DVP-tree.
     * \tparam ForwardIterator A forward-iterator type that can be used to obtain the vertices.
     * \param aBegin The start of the range from which to take the vertices to be erased.
     * \param aEnd The end of the range from which to take the vertices to be erased (one-past-last).
     */
    template <typename ForwardIterator>
    void erase(ForwardIterator aBegin, ForwardIterator aEnd) {
      using namespace boost;
      if(num_vertices(m_tree) == 0) return;
      typedef std::list< std::pair< vertex_type,std::vector<Key> > > key_listing;
      key_listing key_lists; //will hold a list of unique nodes and all their non-erased
      for(;aBegin != aEnd;++aBegin) {
	point_type u_pt = get(m_position, *aBegin);
	vertex_type u_node = get_key(*aBegin, u_pt, m_root);
	if(u_node == vertex_type()) continue;
	key_lists.push_back( std::make_pair(u_node, std::vector<Key>()) );
	for(typename key_listing::iterator it = key_lists.begin(); ((it != key_lists.end()) && (it->first != u_node)); ++it) {
	  Key v = get(m_key,it->first);
	  if(get_key(v,get(m_position,v),u_node) == it->first) {
	    //this means that v is a child of u_node, thus, it's key_list should be merged with u_nodes list.
	    key_lists.back().second.insert(key_lists.back().second.end(),it->second.begin(),it->second.end());
	    remove_edge(source(*(in_edges(v,m_tree).first),m_tree),v,m_tree);
	    remove_vertex(v,m_tree);
	    key_lists.erase(it--);
	  };
	};
	out_edge_iter ei,ei_end;
	for(tie(ei,ei_end) = out_edges(u_node,m_tree); ei != ei_end; ++ei)
	  collect_keys(key_lists.back().second,target(*ei,m_tree),bind<bool>(lambda::unlambda(_1 == std::find(aBegin,aEnd,_2)),aEnd,_1));
	clear_node(u_node);
      };
      //at this point, there should be a list of nodes to be reconstructed along with a list of keys that they should contain.
      for(typename key_listing::iterator it = key_lists.begin(); it != key_lists.end(); ++it) {
	std::map<Key,distance_type> dist_map;
	construct_node(it->first,it->second.begin(),it->second.end(),dist_map);
      };
    };

    /**
     * Clears the DVP-tree.
     */
    void clear() {
      m_tree.clear();
      m_root = vertex_type();
    };

    /**
     * Finds the nearest neighbor to a given position.
     * \param aPoint The position from which to find the nearest-neighbor of.
     * \return The vertex in the DVP-tree that is closest to the given point.
     */
    Key find_nearest(const point_type& aPoint) const {
      using namespace boost;
      if(num_vertices(m_tree) == 0) return Key();
      std::multimap<distance_type,Key> m;
      distance_type sig = std::numeric_limits<distance_type>::infinity();
      find_nearest_impl(aPoint,sig,m_root,m,1);
      return m.begin()->second;
    };

    /**
     * Finds the K nearest-neighbors to a given position.
     * \param aPoint The position from which to find the nearest-neighbors.
     * \param aList Stores, as output, a map of all the K nearest-neighbors to aPoint, the map gives the distance and vertex pairs.
     * \param K The number of nearest-neighbors.
     */
    void find_nearest(const point_type& aPoint, std::multimap<distance_type, Key>& aList, std::size_t K) const {
      using namespace boost;
      if(num_vertices(m_tree) == 0) return;
      aList.clear();
      distance_type sig = std::numeric_limits<distance_type>::infinity();
      find_nearest_impl(aPoint,sig,m_root,aList,K);
    };

    /**
     * Finds the nearest-neighbors to a given position within a given range (radius).
     * \param aPoint The position from which to find the nearest-neighbors.
     * \param aList Stores, as output, a map of all the nearest-neighbors to aPoint, the map gives the distance and vertex pairs.
     * \param R The maximum distance value for the nearest-neighbors.
     */
    void find_in_range(const point_type& aPoint, std::multimap<distance_type, Key>& aList, distance_type R) const {
      using namespace boost;
      if(num_vertices(m_tree) == 0) return;
      find_nearest_impl(aPoint,R,m_root,aList,num_vertices(m_tree));
    };

    /**
     * Takes a vertex, finds its nearest-neighbor and then it adds it to the DVP-tree.
     * \param aVertex The vertex to be added to the DVP-tree.
     * \return The nearest-neighbor of the given vertex.
     */
    Key insert_and_find_nearest(Key aVertex) {
      Key result = find_nearest(get(m_position,aVertex));
      insert(aVertex);
    };

    /**
     * Takes a vertex, finds its K nearest-neighbors and then it adds it to the DVP-tree.
     * \param aVertex The vertex to be added to the DVP-tree.
     * \param aList Stores, as output, a map of all the K nearest-neighbors to aVertex, the map gives the distance and vertex pairs.
     * \param K The number of nearest-neighbors.
     */
    void insert_and_find_nearest(Key aVertex, std::list<Key>& aList, std::size_t K) {
      find_nearest(get(m_position,aVertex),aList,K);
      insert(aVertex);
    };

    /**
     * Takes a vertex, finds its nearest-neighbors within a range and then it adds it to the DVP-tree.
     * \param aVertex The vertex to be added to the DVP-tree.
     * \param aList Stores, as output, a map of all the nearest-neighbors to aVertex, the map gives the distance and vertex pairs.
     * \param R The maximum distance value for the nearest-neighbors.
     */
    void insert_and_find_in_range(Key aVertex, std::list<Key>& aList, distance_type R) {
      find_in_range(get(m_position,aVertex),aList,R);
      insert(aVertex);
    };

};

template <typename Graph, typename DVPTree>
struct multi_dvp_tree_search {
  typename std::map<Graph*, DVPTree*> graph_tree_map;
  multi_dvp_tree_search() : graph_tree_map() { };

  template <typename Topology, typename PositionMap>
  typename boost::graph_traits<Graph>::vertex_descriptor operator()(const typename boost::property_traits<PositionMap>::value_type& p,
                                                                    Graph& g, const Topology& space, PositionMap position) {
    typename std::map<Graph*,DVPTree*>::iterator it = graph_tree_map.find(&g);
    if((it != graph_tree_map.end()) && (it->second))
      return it->second->find_nearest(p);
    else
      return typename boost::graph_traits<Graph>::vertex_descriptor();
  };
};

#endif


#ifndef __TRACKBALL_HPP__
#define __TRACKBALL_HPP__

#include <vsim/env/scene_fwd.hpp>
#include <Eigen/Geometry>

using vsim::Camera ;

// adapted from https://github.com/rsmz/trackball/

class TrackBall
{
public:

    TrackBall();
    ~TrackBall();

    enum CameraMotionType {
        NONE, ARC, FIRSTPERSON, PAN, ROLL, ZOOM
    } ;

    static Eigen::Vector3f X, Y, Z ;

    Eigen::Vector3f computePointOnSphere(float px, float py);
    Eigen::Quaternionf computeRotationBetweenVectors(const Eigen::Vector3f & start,  const Eigen::Vector3f& stop) ;

    CameraMotionType getMotionLeftClick();
    CameraMotionType getMotionMiddleClick();
    CameraMotionType getMotionRightClick();
    CameraMotionType getMotionScroll();

    void setScrollDirection(bool up);
    void setCamera(Camera *c, const Eigen::Vector3f &eye, const Eigen::Vector3f &center, const Eigen::Vector3f &up);
    void setZoomScale(float zoom_scale) ;
	void setClickPoint(double x, double y);
	void setLeftClicked(bool value);
	void setMiddleClicked(bool value);
	void setMotionLeftClick(CameraMotionType motion);
	void setMotionMiddleClick(CameraMotionType motion);
	void setMotionRightClick(CameraMotionType motion);
    void setMotionScroll(CameraMotionType motion);
	void setRightClicked(bool value);
	void setScreenSize(float width, float height);
	void setSpeed(float s);
	void update();

protected:
    char clickQuadrant(float x, float y);
    Eigen::Vector3f computeCameraEye();
    Eigen::Vector3f computeCameraUp();
    Eigen::Vector3f computePan();
    void drag();
    void drag(bool isClicked, CameraMotionType motion);
    void dragArc();
    // Change eye position and up direction while keeping center point static.
    void dragArcCamera();
    void dragFirstPerson();
    void dragZoom();
    // Simulate zoom by moving camera along viewing eye direction.
    void dragZoomCamera();
    // Move camera focal center position with static up and eye direction.
    void dragPan();
    // Roll about eye direction.
    void rollCamera();
    void freezeTransform();
    void scroll();
    void updateCameraEyeUp(bool eye, bool up);

private:
    Camera *camera_;
    CameraMotionType camera_motion_left_click_;
    CameraMotionType camera_motion_middle_click_;
    CameraMotionType camera_motion_right_click_;
    CameraMotionType camera_motion_scroll_;
    float cx_, cy_, pcx_, pcy_;
    float height_, pan_scale_;
    bool is_dragging_, is_left_click_, is_middle_click_, is_right_click_, is_scrolling_ ;
    float roll_scale_, roll_sum_, speed_ ;

    Eigen::Quaternionf rotation_, rotation_sum_ ;
    Eigen::Vector3f start_, stop_ ;
    float translation_length_, width_, zoom_sum_, zoom_scale_ ;
    Eigen::Vector3f center_, eye_, up_ ;

};

#endif

/*
    LICENSE BEGIN

    trackball - A 3D view interactor for C++ programs.
    Copyright (C) 2016  Remik Ziemlinski

    This program is free software: you can redistribute it and/or modify
    it under the terms of the GNU General Public License as published by
    the Free Software Foundation, either version 3 of the License, or
    (at your option) any later version.

    This program is distributed in the hope that it will be useful,
    but WITHOUT ANY WARRANTY; without even the implied warranty of
    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
    GNU General Public License for more details.

    You should have received a copy of the GNU General Public License
    along with this program.  If not, see <http://www.gnu.org/licenses/>.

    LICENSE END
*/


#include "generate_data/clip_by_plane.h"

#include "generate_data/triangle_subset.h"
#include "intersect/triangle.h"
#include "intersect/triangle_impl.h"
#include "lib/array_vec.h"

#include <boost/geometry.hpp>

#include "dbg.h"

namespace generate_data {
ATTR_PURE_NDEBUG TriangleSubset
clip_by_plane(const Eigen::Vector3d &normal, double plane_threshold,
              const intersect::TriangleGen<double> &tri, bool debug) {
  ArrayVec<unsigned, 3> included;
  ArrayVec<double, 3> precomputed_vals;
  ArrayVec<unsigned, 3> excluded_or_on_edge;
  unsigned excluded_count = 0;
  if (debug) {
    dbg(normal);
    dbg(plane_threshold);
  }

  std::array<double, 3> vals_all;

  for (unsigned i = 0; i < 3; ++i) {
    double dotted = normal.dot(tri.vertices[i]);
    double value = plane_threshold - dotted;
    if (debug) {
      dbg(tri.vertices[i]);
      dbg(value);
    }
    vals_all[i] = value;
    if (value > -1e-6) {
      excluded_or_on_edge.push_back(i);
      if (value > 1e-6) {
        ++excluded_count;
      }
    } else {
      precomputed_vals.push_back(value);
      included.push_back(i);
    }
  }

  debug_assert(excluded_or_on_edge.size() + included.size() == 3);

  // if all are on edge, we include
  if (excluded_count == 0) {
    return {tag_v<TriangleSubsetType::All>, {}};
  }
  if (included.empty()) {
    return {tag_v<TriangleSubsetType::None>, {}};
  }

  if (debug) {
    dbg(excluded_count);
    dbg(unsigned(included.size()));
  }

  auto make_point = [&](double value, unsigned idx) -> Eigen::Vector2d {
    if (idx == 2) {
      return {0., value};
    } else {
      debug_assert(idx == 1);
      return {value, 0.};
    }
  };

  auto edge_points = [&](unsigned included_idx, unsigned end_idx) {
    unsigned origin_idx = included[included_idx];
    Eigen::Vector3d origin = tri.vertices[origin_idx];
    Eigen::Vector3d ray = tri.vertices[end_idx] - origin;

    // NOTE: ray.dot(normal) can't be zero (outside of some serious floating
    // point issues) because that would imply that either all points are
    // included or all are excluded.
    double prop = precomputed_vals[included_idx] / ray.dot(normal);
    Eigen::Vector3d endpoint = prop * ray + origin;

    // debug_assert(-1e-2 <= prop);
    // debug_assert(1. + 1e-2 >= prop);
    prop = std::clamp(prop, 0., 1.);
    debug_assert(std::abs(endpoint.dot(normal) - plane_threshold) < 1e-6);

    if (origin_idx == 0) {
      return make_point(prop, end_idx);
    } else if (end_idx == 0) {
      return make_point(1 - prop, origin_idx);
    } else {
      auto baryocentric_vals = tri.interpolation_values(endpoint);
      Eigen::Vector2d ret{baryocentric_vals[1], baryocentric_vals[2]};
      [[maybe_unused]] auto baryo_point =
          tri.baryo_to_point({ret.x(), ret.y()});
      debug_assert((endpoint - baryo_point).norm() < 1e-6);

      return ret;
    }
  };

  ArrayVec<Eigen::Vector2d, 4> points;

  auto add_point = [&](const Eigen::Vector2d &p) { points.push_back(p); };
  if (included.size() == 2) {
    add_point(edge_points(0, excluded_or_on_edge[0]));
    add_point(edge_points(1, excluded_or_on_edge[0]));
  } else {
    debug_assert(included.size() == 1);
    add_point(edge_points(0, excluded_or_on_edge[0]));
    add_point(edge_points(0, excluded_or_on_edge[1]));
  }

  auto get_point_in_baryo = [&](unsigned idx) -> Eigen::Vector2d {
    if (idx == 0) {
      return {0., 0.};
    } else {
      return make_point(1., idx);
    }
  };

  for (unsigned i : included) {
    add_point(get_point_in_baryo(i));
  }

  Eigen::Vector2d centroid{0, 0};
  for (const auto &p : points) {
    centroid += p;
  }
  centroid /= points.size();

  ArrayVec<std::tuple<BaryoPoint, double>, 4> points_angle;

  for (const auto &p : points) {
    Eigen::Vector2d vec = p - centroid;
    points_angle.push_back({{p.x(), p.y()}, std::atan2(vec.y(), vec.x())});
  }

  // sort by angle to ensure valid poly
  std::sort(points_angle.begin(), points_angle.end(),
            [](const auto &l, const auto &r) {
              return std::get<1>(l) > std::get<1>(r);
            });

  TriPolygon poly;
  for (const auto &p : points_angle) {
    boost::geometry::append(poly, std::get<0>(p));
  }
  boost::geometry::append(poly, std::get<0>(points_angle[0]));

  double area = boost::geometry::area(poly);
  if (debug) {
    dbg(area);
  }
  if (area <= 1e-12) {
    return {tag_v<TriangleSubsetType::None>, {}};
  } else if (area > 0.5 - 1e-12) {
    return {tag_v<TriangleSubsetType::All>, {}};
  }

  debug_assert(boost::geometry::is_valid(poly));

  return {tag_v<TriangleSubsetType::Some>, poly};
}
} // namespace generate_data


#include <filesystem>
#include <fstream>
#include <iostream>

#include <boost/accumulators/accumulators.hpp>
#include <boost/accumulators/statistics.hpp>

#include "algorithms/multiplicative_dijkstra.hpp"

#include "parsers/std_mutable_landscape_parser.hpp"
#include "parsers/std_restoration_plan_parser.hpp"

#include "helper.hpp"
#include "instances_helper.hpp"

#include "precomputation/trivial_reformulation.hpp"

using Accumulator = boost::accumulators::accumulator_set<
    double, boost::accumulators::features<boost::accumulators::tag::mean,
                                          boost::accumulators::tag::median>>;

int main() {
    std::ofstream data_log("output/data.log");
    data_log << std::fixed << std::setprecision(6);
    data_log << "ball_diameter "
             << "average_percentage_of_nodes "
             << "median_percentage_of_nodes "
             << "average_percentage_of_contribution "
             << "median_percentage_of_contribution " << std::endl;

    constexpr int nb_partitions = 1000;
    std::array<Accumulator, nb_partitions> percent_node_accs;
    std::array<Accumulator, nb_partitions> percent_contribution_accs;

    auto prob_to_partition = [&](double prob) {
        if(prob == 0) return nb_partitions - 1;
        return static_cast<int>((nb_partitions) * (1 - prob));
    };

    using Graph = MutableLandscape::Graph;
    using Node = MutableLandscape::Node;
    using NodeIt = MutableLandscape::NodeIt;

    Instance raw_instance = make_instance_biorevaix_level_2_v7(6);
    Instance instance = trivial_reformulate(std::move(raw_instance));
    const MutableLandscape & landscape = instance.landscape;

    const Graph & graph = landscape.getNetwork();

    for(NodeIt s(graph); s != lemon::INVALID; ++s) {
        int last_partition = 0;
        int node_count_sum = 0;
        double contribution_sum = 0;

        const auto sorted_pairs = Helper::computeDistancePairs(landscape, s);
        const double contribution_max = std::transform_reduce(
            sorted_pairs.begin(), sorted_pairs.end(), 0.0, std::plus<double>(),
            [&](const auto & p) {
                return landscape.getQuality(s) * landscape.getQuality(p.first) *
                       p.second;
            });

        for(const auto & [t, p_st] : sorted_pairs) {
            const int current_partition = prob_to_partition(p_st);
            if(current_partition > last_partition) {
                for(; last_partition < current_partition; ++last_partition) {
                    percent_node_accs[last_partition](
                        static_cast<double>(node_count_sum) /
                        sorted_pairs.size() * 100);
                    percent_contribution_accs[last_partition](
                        contribution_sum / contribution_max * 100);
                }
            }

            ++node_count_sum;
            contribution_sum +=
                landscape.getQuality(s) * landscape.getQuality(t) * p_st;
        }
    }

    for(int i = 0; i < nb_partitions; ++i) {
        const double diameter = 1 - (static_cast<double>(i) / nb_partitions);
        data_log << diameter << " "
                 << boost::accumulators::mean(percent_node_accs[i]) << " "
                 << boost::accumulators::median(percent_node_accs[i]) << " "
                 << boost::accumulators::mean(percent_contribution_accs[i])
                 << " "
                 << boost::accumulators::median(percent_contribution_accs[i])
                 << std::endl;
    }

    return EXIT_SUCCESS;
}


#ifndef _DECONVOLVE_HPP_
#define _DECONVOLVE_HPP_

#include <functional>
#ifdef NDEBUG
#define BOOST_DISABLE_ASSERTS
#endif
#include <boost/multi_array.hpp>

namespace deconvolution {

template <int D>
using Array = boost::multi_array<double, D>;
template <int D>
using LinearSystem = std::function<Array<D>(const Array<D>&)>;

template <int D> class Regularizer;

struct DeconvolveParams {
    double dataSmoothing {0.03};
    double maxSmoothing {1000.0};
    double minSmoothing { 10.0 };
    int maxIterations { 1000 };
    double admmRho { 10.0 };
    double admmConvergenceNorm { 1.0 };
    int lbfgsIters { 20 };
};

struct DeconvolveStats {
    double iterTime = 0;
    double regularizerTime = 0;
    double unaryTime = 0;
    double dataTime = 0;
};

template <int D>
using ProgressCallback = std::function<void(const Array<D>& x, double dual, double primalData, double primalReg, double smoothing)>;
/*
 * Solve a linear inverse system of the form
 * min_x |y - Hx|_2^2 + R(x)
 * where H is a linear system, R is a regularizer, and y are some given
 * observables.
 */
template <int D>
Array<D> Deconvolve(const Array<D>& y, const LinearSystem<D>& H, const LinearSystem<D>& Ht, Regularizer<D>& R, ProgressCallback<D>& pc, DeconvolveParams& params, DeconvolveStats& s);

template <int D>
Array<D> DeconvolveADMM(const Array<D>& y, const LinearSystem<D>& H, const LinearSystem<D>& Ht, Regularizer<D>& R, ProgressCallback<D>& pc, DeconvolveParams& params, DeconvolveStats& s);

template <int D>
Array<D> DeconvolvePrimal(const Array<D>& y, const LinearSystem<D>& H, const LinearSystem<D>& Ht, Regularizer<D>& R, ProgressCallback<D>& pc, DeconvolveParams& params, DeconvolveStats& s, const Array<D>& initX);

} // namespace deconvolution

#endif


/*=========================================================================

  Program:   Visualization Toolkit
  Module:    vtkPointCloudLinearProjector.cxx
  Author: Pierre Guilbert

  Copyright (c) Ken Martin, Will Schroeder, Bill Lorensen
  All rights reserved.
  See Copyright.txt or http://www.kitware.com/Copyright.htm for details.

  This software is distributed WITHOUT ANY WARRANTY; without even
  the implied warranty of MERCHANTABILITY or FITNESS FOR A PARTICULAR
  PURPOSE.  See the above copyright notice for more information.

=========================================================================*/

// LOCAL
#include "vtkPointCloudLinearProjector.h"
#include "vtkEigenTools.h"

// STD
#include <cmath>
#include <fstream>
#include <iostream>
#include <sstream>

// VTK
#include <vtkImageData.h>
#include <vtkInformation.h>
#include <vtkInformationVector.h>
#include <vtkObjectFactory.h>
#include <vtkPointData.h>
#include <vtkPolyData.h>
#include <vtkStreamingDemandDrivenPipeline.h>

// BOOST
#include <boost/algorithm/string.hpp>

// Eigen
#include <Eigen/Dense>

// Implementation of the New function
vtkStandardNewMacro(vtkPointCloudLinearProjector)

//------------------------------------------------------------------------------
int vtkPointCloudLinearProjector::FillInputPortInformation(int port, vtkInformation* info)
{
  if (port == 0)
  {
    info->Set(vtkDataObject::DATA_TYPE_NAME(), "vtkPolyData");
    return 1;
  }
  return 0;
}

//-----------------------------------------------------------------------------
int vtkPointCloudLinearProjector::RequestInformation(vtkInformation *vtkNotUsed(request),
                                                vtkInformationVector **vtkNotUsed(inputVector),
                                                vtkInformationVector *outputVector)
{
  if (this->Resolution[0] <= 0 || this->Resolution[1] <= 0)
  {
    vtkWarningMacro("Resolution must be positive, not " << this->Resolution[0] << "x" << this->Resolution[1] << ".")
    return VTK_ERROR;
  }

  vtkInformation* outInfo = outputVector->GetInformationObject(0);
  outInfo->Set(vtkStreamingDemandDrivenPipeline::WHOLE_EXTENT(),
               0, this->Resolution[0] - 1,
               0, this->Resolution[1] - 1,
               0, 0);

  outInfo->Set(vtkDataObject::ORIGIN(), this->Origin, 3);
  outInfo->Set(vtkDataObject::SPACING(), this->Spacing, 3);

  if (this->ExportAsChar)
  {
    vtkDataObject::SetPointDataActiveScalarInfo(outInfo, VTK_UNSIGNED_CHAR, 1);
  }
  else
  {
    vtkDataObject::SetPointDataActiveScalarInfo(outInfo, VTK_DOUBLE, 1);
  }
  return VTK_OK;
}

//------------------------------------------------------------------------------
int vtkPointCloudLinearProjector::RequestData(vtkInformation* vtkNotUsed(request),
  vtkInformationVector** inputVector,
  vtkInformationVector* outputVector)
{
  // Get the input
  vtkPolyData* input = vtkPolyData::GetData(inputVector[0]->GetInformationObject(0));

  double point[3];
  Eigen::Vector3d X;
  vtkDataArray* values = this->GetInputArrayToProcess(0, inputVector);

  // Express the data in an other reference frame to align the new
  // Z-axis with a settled direction (typically, the gravity acceleration
  // vector)
  auto transformedPoints = vtkSmartPointer<vtkPoints>::New();
  transformedPoints->SetNumberOfPoints(input->GetNumberOfPoints());
  // Transform the input polydata
  for (vtkIdType pointIndex = 0; pointIndex < input->GetNumberOfPoints(); ++pointIndex)
  {
    input->GetPoint(pointIndex, point);
    X << point[0], point[1], point[2];
    X = this->Projector * X;
    transformedPoints->SetPoint(pointIndex, X.data());
  }

  // Get the point cloud bounding box parameters
  double boundingBox[6];
  transformedPoints->Modified();
  transformedPoints->GetBounds(boundingBox);

  double pointRangeX = boundingBox[1] - boundingBox[0];
  double pointRangeY = boundingBox[3] - boundingBox[2];

  if (pointRangeX <= 0 || pointRangeY <= 0)
  {
    return VTK_ERROR;
  }

  this->Origin[0] = boundingBox[0];
  this->Origin[1] = boundingBox[2];

  this->Spacing[0] = pointRangeX / static_cast<double>(this->Resolution[0]);
  this->Spacing[1] = pointRangeY / static_cast<double>(this->Resolution[1]);

  double scaleX = (this->Resolution[0] - 1) / (pointRangeX);
  double scaleY = (this->Resolution[1] - 1) / (pointRangeY);

  // Get the output image and fill with zeros
  vtkSmartPointer<vtkImageData> image = vtkSmartPointer<vtkImageData>::New();
  image->SetSpacing(this->Spacing);
  // Is it not possible to orient a vtkImageData within a vtkImageAlgorithm
  // filter in ParaView 5.4, which LidarView is currently based on. There are
  // commits to add a "DirectionMatrix" to vtkImageData (commit
  // 55c0d1cbdf1a152eec74d538af9fa4e17617c379) but this is from April 2019.
  //
  // TODO
  // Use DirectionMatrix when it becomes available to align the projection to
  // the point cloud along the projection axis.
  image->SetOrigin(this->Origin);
  image->SetDimensions(this->Resolution[0], this->Resolution[1], 1);
  if (this->ExportAsChar)
  {
    image->AllocateScalars(VTK_UNSIGNED_CHAR, 1);
  }
  else
  {
    image->AllocateScalars(VTK_DOUBLE, 1);
  }
  // Compute distribution about the height of the data lying in a pixel
  std::vector<std::vector<double> > perPixelDistribution(this->Resolution[0] * this->Resolution[1]);
  for (vtkIdType pointIndex = 0; pointIndex < input->GetNumberOfPoints(); ++pointIndex)
  {
    transformedPoints->GetPoint(pointIndex, point);
    int xPixelCoord = std::floor((point[0] - boundingBox[0]) * scaleX);
    int yPixelCoord = std::floor((point[1] - boundingBox[2]) * scaleY);
    perPixelDistribution[xPixelCoord + this->Resolution[0] * yPixelCoord].push_back(
      this->HeightMap ? point[2] : values->GetTuple1(pointIndex)
    );
  }

  // fill the image
  double valueRange[2];
  if (this->HeightMap)
  {
    valueRange[0] = boundingBox[4];
    valueRange[1] = boundingBox[5];
  }
  else
  {
    values->GetRange(valueRange);
  }

  double valueShift = (this->ExportAsChar || this->ShiftToZero) ? valueRange[0] : 0.0;
  double valueScale = valueRange[1] - valueRange[0];
  for (unsigned int y = 0; y < this->Resolution[1]; ++y)
  {
    for (unsigned int x = 0; x < this->Resolution[0]; ++x)
    {
      unsigned int imageIndex = this->Resolution[0] * y + x;
      double value = 0.0;
      // if the pixel is empty, skip it
      if (perPixelDistribution[imageIndex].size() > 0)
      {
        std::sort(perPixelDistribution[imageIndex].begin(), perPixelDistribution[imageIndex].end());
        unsigned int rankIndex = std::floor((perPixelDistribution[imageIndex].size() - 1) * this->RankPercentile);
        value = perPixelDistribution[imageIndex][rankIndex];
      }
      if (this->ExportAsChar)
      {
        value -= valueShift;
        value = std::round((value / valueScale) * 0xff);
      }
      else if (this->ShiftToZero)
      {
        value -= valueShift;
      }
      // Use unscaled values to match input values and ranges.
      image->SetScalarComponentFromDouble(x, y, 0, 0, value);
    }
  }

  int neigh = this->MedianFilterWidth;
  if (this->ShouldMedianFilter)
  {
    vtkSmartPointer<vtkImageData> tempImage = vtkSmartPointer<vtkImageData>::New();
    tempImage->DeepCopy(image);
    int dimX = static_cast<int>(this->Resolution[0]);
    int dimY = static_cast<int>(this->Resolution[1]);
    for (int x = 0; x < dimX; ++x)
    {
      for (int y = 0; y < dimY; ++y)
      {
        if (tempImage->GetScalarComponentAsDouble(x, y, 0, 0) == 0)
        {
          continue;
        }
        std::vector<double> neighborhoodValues;
        unsigned int minU = std::max(0, x - neigh);
        unsigned int maxU = std::min(dimX - 1, x + neigh);
        unsigned int minV = std::max(0, y - neigh);
        unsigned int maxV = std::min(dimY - 1, y + neigh);

        for (unsigned int u = minU; u <= maxU; ++u)
        {
          for (unsigned int v = minV; v <= maxV; ++v)
          {
            neighborhoodValues.push_back(tempImage->GetScalarComponentAsDouble(u, v, 0, 0));
          }
        }
        std::sort(neighborhoodValues.begin(), neighborhoodValues.end());
        double medianValue = neighborhoodValues[neighborhoodValues.size() / 2];
        image->SetScalarComponentFromDouble(x, y, 0, 0, medianValue);
      }
    }
  }

  vtkImageData* outputImage = vtkImageData::GetData(outputVector->GetInformationObject(0));
  outputImage->ShallowCopy(image);
  // Use the input name to share the color scale.
  if (this->HeightMap)
  {
    outputImage->GetPointData()->GetScalars()->SetName("DistanceToPlane");
  }
  else
  {
    outputImage->GetPointData()->GetScalars()->SetName(values->GetName());
  }
  this->Modified();
  return VTK_OK;
}

//------------------------------------------------------------------------------
void vtkPointCloudLinearProjector::SetPlaneNormal(double w0, double w1, double w2)
{
  // Here we will construct a new base of R3 using
  // the normal of the plane as the Z-axis. To proceed,
  // we will compute the rotation that map ez toward n
  // and its axis being cross(ez, n)
  Eigen::Vector3d ez(0, 0, 1);
  Eigen::Vector3d n(w0, w1, w2);

  // check that the plane normal is not the null pointer
  if (n.norm() < std::numeric_limits<float>::epsilon())
  {
    vtkGenericWarningMacro("The plane normal cannot be the null vector");
    return;
  }
  n.normalize();

  this->Modified();

  double cosAngle = ez.dot(n);
  Eigen::Vector3d rawAxis = ez.cross(n);
  double rawAxisMagnitude = rawAxis.norm();
  // Check if ez and n are colinear.
  if (rawAxisMagnitude < std::numeric_limits<float>::epsilon())
  {
    // Rotate 180 deg around the x axis if the projection is along negative z.
    if (cosAngle < 0)
    {
      this->ChangeOfBasis << 1,0,0,  0,-1,0,  0,0,-1;
    }
    // Reset it to the identity matrix otherwise.
    else
    {
      this->ChangeOfBasis = Eigen::Matrix3d::Identity();
    }
  }
  else
  {
    Eigen::Vector3d unitAxis = rawAxis / rawAxisMagnitude;
    double sinAngle = rawAxis.dot(unitAxis);
    double angle = std::atan2(sinAngle, cosAngle);
    Eigen::Matrix3d R(Eigen::AngleAxisd(angle, unitAxis));
    this->ChangeOfBasis = R;
  }
  this->Projector = this->DiagonalizedProjector * this->ChangeOfBasis.inverse();
}


#include <memory>
#include <random>
#include <vector>

#include <Eigen/Core>
#include <gtest/gtest.h>

#include "test/utils.hpp"

#include "ldaplusplus/Parameters.hpp"
#include "ldaplusplus/events/ProgressEvents.hpp"
#include "ldaplusplus/em/MultinomialSupervisedEStep.hpp"
#include "ldaplusplus/em/MultinomialSupervisedMStep.hpp"

using namespace Eigen;
using namespace ldaplusplus;

// T will be available as TypeParam in TYPED_TEST functions
template <typename T>
class TestMultinomialMaximizationStep : public ParameterizedTest<T> {};

TYPED_TEST_CASE(TestMultinomialMaximizationStep, ForFloatAndDouble);

TYPED_TEST(TestMultinomialMaximizationStep, Maximization) {
    // Build the corpus
    std::mt19937 rng;
    rng.seed(0);
    MatrixXi X(100, 50);
    VectorXi y(50);
    std::uniform_int_distribution<> class_generator(0, 5);
    std::exponential_distribution<> words_generator(0.1);
    for (int d=0; d<50; d++) {
        for (int w=0; w<100; w++) {
            X(w, d) = static_cast<int>(words_generator(rng));
        }
        y(d) = class_generator(rng);
    }

    // Create the corpus and the model
    auto corpus = std::make_shared<corpus::EigenClassificationCorpus>(X, y);
    MatrixX<TypeParam> beta = MatrixX<TypeParam>::Random(10, 100);
    beta.array() -= beta.minCoeff();
    beta.array().rowwise() /= beta.array().colwise().sum();
    auto model = std::make_shared<parameters::SupervisedModelParameters<TypeParam> >(
        VectorX<TypeParam>::Constant(10, 0.1),
        beta,
        MatrixX<TypeParam>::Constant(10, 6, 1. / 6)
    );

    em::MultinomialSupervisedEStep<TypeParam> e_step(10, 1e-2, 2);
    em::MultinomialSupervisedMStep<TypeParam> m_step(2);

    for (size_t i=0; i<corpus->size(); i++) {
        m_step.doc_m_step(
            corpus->at(i),
            e_step.doc_e_step(
                corpus->at(i),
                model
            ),
            model
        );
    }

    std::vector<TypeParam> progress;
    m_step.get_event_dispatcher()->add_listener(
        [&progress](std::shared_ptr<events::Event> event) {
            if (event->id() == "MaximizationProgressEvent") {
                auto prog_ev = std::static_pointer_cast<events::MaximizationProgressEvent<TypeParam> >(event);
                progress.push_back(prog_ev->likelihood());
            }
        }
    );

    m_step.m_step(
        model
    );

    ASSERT_EQ(1, progress.size());
    ASSERT_GT(0, progress[0]);
}