#include <I3Test.h>

#include <PROPOSAL-icetray/I3PropagatorServicePROPOSAL.h>
#include <PROPOSAL-icetray/SimplePropagator.h>
#include <dataclasses/physics/I3Particle.h>
#include <phys-services/I3SPRNGRandomService.h>

#include <boost/make_shared.hpp>

TEST_GROUP(Repeatablility);

static I3Particle make_particle()
{
    I3Particle p;
    p.SetPos(0, 0, 2e3);
    p.SetDir(0, 0);
    p.SetTime(0);
    p.SetType(I3Particle::MuMinus);
    p.SetLocationType(I3Particle::InIce);
    p.SetEnergy(1e6);

    return p;
}

TEST(SimplePropagator)
{
    I3RandomServicePtr rng(new I3SPRNGRandomService(1, 10000, 2));

    boost::shared_ptr<PROPOSAL::SimplePropagator> prop(
        new PROPOSAL::SimplePropagator(I3Particle::MuMinus, "ice", 5e2, -1));
    prop->SetRandomNumberGenerator(rng);
    double distance = 1e2;

    std::vector<std::vector<I3Particle> > daughters;
    std::vector<I3Particle> primaries;
    std::vector<I3FrameObjectPtr> states;
    for (int i = 0; i < 100; i++)
    {
        I3Particle p = make_particle();
        p.SetEnergy(std::pow(10, rng->Uniform(3, 6)));
        primaries.push_back(p);
        states.push_back(rng->GetState());
        boost::shared_ptr<std::vector<I3Particle> > d(new std::vector<I3Particle>);
        prop->propagate(p, distance, d);
        daughters.push_back(*d);
    }

    // Now, replay the simulation and ensure we get the same thing

    for (size_t i = 0; i < daughters.size(); i++)
    {
        // Throw out a random subset of events. The remainder should
        // be reproducible given the same sequence of random numbers,
        // unless of course the propagator keeps secret state.
        if (rng->Uniform() < 0.5)
            continue;
        rng->RestoreState(states[i]);
        boost::shared_ptr<std::vector<I3Particle> > d(new std::vector<I3Particle>);
        prop->propagate(primaries[i], distance, d);
        ENSURE_EQUAL(daughters[i].size(), d->size(), "Same number of particles must be produced");
        for (size_t j = 0; j < daughters[i].size(); j++)
        {
            I3Particle& p1 = daughters[i][j];
            I3Particle& p2 = (*d)[j];
            ENSURE_EQUAL(p1.GetType(), p2.GetType(), "Secondaries must have the same type");
            ENSURE_EQUAL(p1.GetTime(), p2.GetTime(), "Times must be completely identical");
            ENSURE_EQUAL(p1.GetEnergy(), p2.GetEnergy(), "Energies must completely identical");
        }
    }
}

TEST(PropagatorService)
{
    I3RandomServicePtr rng(new I3SPRNGRandomService(1, 10000, 1));
    I3FrameObjectPtr state = rng->GetState();

    PROPOSAL::I3PropagatorServicePROPOSALPtr prop(new PROPOSAL::I3PropagatorServicePROPOSAL);
    prop->SetRandomNumberGenerator(rng);

    I3PropagatorService::DiagnosticMapPtr frame(new I3PropagatorService::DiagnosticMap);
    // the dummy I3Frame makes compiler happy, but won't be used
    I3FramePtr dummy(new I3Frame());

    std::vector<std::vector<I3Particle> > daughters;
    for (int i = 0; i < 2; i++)
    {
        I3Particle p              = make_particle();
        std::vector<I3Particle> d = prop->Propagate(p, frame, dummy);
        daughters.push_back(d);
    }

    rng->RestoreState(state);

    for (size_t i = 0; i < daughters.size(); i++)
    {
        I3Particle p              = make_particle();
        std::vector<I3Particle> d = prop->Propagate(p, frame, dummy);
        ENSURE_EQUAL(daughters[i].size(), d.size());
        for (size_t j = 0; j < daughters[i].size(); j++)
        {
            I3Particle& p1 = daughters[i][j];
            I3Particle& p2 = d[j];
            ENSURE_EQUAL(p1.GetType(), p2.GetType());
            ENSURE_EQUAL(p1.GetTime(), p2.GetTime());
            ENSURE_EQUAL(p1.GetEnergy(), p2.GetEnergy());
        }
    }
}


// AbilitySpend.cpp
//
#include "StdAfx.h"
#include "AbilitySpend.h"
#include "XmlLib\SaxWriter.h"
#include <boost/static_assert.hpp>

#define DL_ELEMENT AbilitySpend

namespace
{
    const wchar_t f_saxElementName[] = L"AbilitySpend";
    DL_DEFINE_NAMES(AbilitySpend_PROPERTIES)

    const unsigned f_verCurrent = 1;

    size_t c_maxSpend = 10;
    int c_statIncreaseCosts[] = {1, 1, 1, 1, 1, 1, 2, 2, 3, 3};
    const unsigned arraySize = sizeof(c_statIncreaseCosts) / sizeof(*c_statIncreaseCosts);
    BOOST_STATIC_ASSERT(arraySize == 10);
}

AbilitySpend::AbilitySpend() :
    XmlLib::SaxContentElement(f_saxElementName, f_verCurrent)
{
    DL_INIT(AbilitySpend_PROPERTIES)
}

DL_DEFINE_ACCESS(AbilitySpend_PROPERTIES)

XmlLib::SaxContentElementInterface * AbilitySpend::StartElement(
        const XmlLib::SaxString & name,
        const XmlLib::SaxAttributes & attributes)
{
    XmlLib::SaxContentElementInterface * subHandler =
            SaxContentElement::StartElement(name, attributes);

    DL_START(AbilitySpend_PROPERTIES)

    return subHandler;
}

void AbilitySpend::EndElement()
{
    SaxContentElement::EndElement();
    DL_END(AbilitySpend_PROPERTIES)
    if (m_StrSpend > c_maxSpend)
    {
        SAXASSERT(false, DL_STRINGIZE(DL_ELEMENT) "::m_StrSpend had invalid value");
    }
    if (m_DexSpend > c_maxSpend)
    {
        SAXASSERT(false, DL_STRINGIZE(DL_ELEMENT) "::m_DexSpend had invalid value");
    }
    if (m_ConSpend > c_maxSpend)
    {
        SAXASSERT(false, DL_STRINGIZE(DL_ELEMENT) "::m_ConSpend had invalid value");
    }
    if (m_IntSpend > c_maxSpend)
    {
        SAXASSERT(false, DL_STRINGIZE(DL_ELEMENT) "::m_IntSpend had invalid value");
    }
    if (m_WisSpend > c_maxSpend)
    {
        SAXASSERT(false, DL_STRINGIZE(DL_ELEMENT) "::m_WisSpend had invalid value");
    }
    if (m_ChaSpend > c_maxSpend)
    {
        SAXASSERT(false, DL_STRINGIZE(DL_ELEMENT) "::m_ChaSpend had invalid value");
    }
    if (!m_hasUserSelectedSpend)
    {
        m_hasUserSelectedSpend = true;
    }
}

void AbilitySpend::Write(XmlLib::SaxWriter * writer) const
{
    writer->StartElement(ElementName());//, VersionAttributes());
    DL_WRITE(AbilitySpend_PROPERTIES)
    writer->EndElement();
}

size_t AbilitySpend::PointsSpent() const
{
    // sum how many build points have been spent on each ability
    size_t totalSpent = 0;
    totalSpent += GetAbilitySpendPoints(Ability_Strength);
    totalSpent += GetAbilitySpendPoints(Ability_Dexterity);
    totalSpent += GetAbilitySpendPoints(Ability_Constitution);
    totalSpent += GetAbilitySpendPoints(Ability_Intelligence);
    totalSpent += GetAbilitySpendPoints(Ability_Wisdom);
    totalSpent += GetAbilitySpendPoints(Ability_Charisma);
    return totalSpent;
}

bool AbilitySpend::CanSpendOnAbility(
        AbilityType ability) const
{
    bool canSpend = false;
    size_t currentSpend = GetAbilitySpend(ability);
    if (currentSpend < c_maxSpend)
    {
        // its not at max, but are there enough points left to
        // spend to the next level?
        int availablePoints = ((int)m_AvailableSpend - (int)PointsSpent());
        int nextCost = c_statIncreaseCosts[currentSpend];
        canSpend = (availablePoints >= nextCost); // true if enough points available
    }
    return canSpend;
}

bool AbilitySpend::CanRevokeSpend(
        AbilityType ability) const
{
    bool canRevokeSpend = false;
    size_t currentSpend = GetAbilitySpend(ability);
    // can always undo if any points spent
    canRevokeSpend = (currentSpend > 0);
    return canRevokeSpend;
}

size_t AbilitySpend::NextPointsSpentCost(
        AbilityType ability) const
{
    size_t nextCost = 0;
    size_t currentSpend = GetAbilitySpend(ability);
    if (currentSpend < c_maxSpend)
    {
        // its not at max, but are there enough points left to
        // spend to the next level?
        nextCost = c_statIncreaseCosts[currentSpend];
    }
    return nextCost;
}

void AbilitySpend::SpendOnAbility(
        AbilityType ability)
{
    switch (ability)
    {
    case Ability_Strength:
        ++m_StrSpend;
        break;
    case Ability_Dexterity:
        ++m_DexSpend;
        break;
    case Ability_Constitution:
        ++m_ConSpend;
        break;
    case Ability_Intelligence:
        ++m_IntSpend;
        break;
    case Ability_Wisdom:
        ++m_WisSpend;
        break;
    case Ability_Charisma:
        ++m_ChaSpend;
        break;
    default:
        ASSERT(FALSE);
        break;
    }
}

void AbilitySpend::RevokeSpendOnAbility(
        AbilityType ability)
{
    switch (ability)
    {
    case Ability_Strength:
        --m_StrSpend;
        break;
    case Ability_Dexterity:
        --m_DexSpend;
        break;
    case Ability_Constitution:
        --m_ConSpend;
        break;
    case Ability_Intelligence:
        --m_IntSpend;
        break;
    case Ability_Wisdom:
        --m_WisSpend;
        break;
    case Ability_Charisma:
        --m_ChaSpend;
        break;
    default:
        ASSERT(FALSE);
        break;
    }
}

size_t AbilitySpend::GetAbilitySpend(
        AbilityType ability) const
{
    size_t spent = c_maxSpend;  // assume
    switch (ability)
    {
    case Ability_Strength:
        spent = m_StrSpend;
        break;
    case Ability_Dexterity:
        spent = m_DexSpend;
        break;
    case Ability_Constitution:
        spent = m_ConSpend;
        break;
    case Ability_Intelligence:
        spent = m_IntSpend;
        break;
    case Ability_Wisdom:
        spent = m_WisSpend;
        break;
    case Ability_Charisma:
        spent = m_ChaSpend;
        break;
    default:
        ASSERT(FALSE);
        break;
    }
    return spent;
}

size_t AbilitySpend::GetAbilitySpendPoints(
        AbilityType ability) const
{
    size_t spentPoints = 0;
    size_t abilitySpend = GetAbilitySpend(ability);
    for (size_t i = 0; i < abilitySpend; ++i)
    {
        spentPoints += c_statIncreaseCosts[i];
    }
    return spentPoints;
}


/*
 * 	robot_server_dvrk.hpp
 *
 *	Author(s): Tamas D. Nagy
 *	Created on: 2016-11-07
 *
 *  Base class for dVRK robot arms, itself usable for
 *  MTMs and ECM.
 *
 */

#ifndef ROBOT_SERVER_DVRK_HPP_
#define ROBOT_SERVER_DVRK_HPP_

#include <iostream>
#include <sstream>
#include <fstream>
#include <vector>
#include <cmath>
#include <ros/ros.h>
#include <ros/package.h>
#include <std_msgs/String.h>
#include <sensor_msgs/JointState.h>
#include <geometry_msgs/TransformStamped.h>
#include <std_msgs/Float32.h>
#include <Eigen/Dense>
#include <Eigen/Geometry>
#include <cmath>
#include <irob_dvrk/arm_types.hpp>
#include <irob_general_robot/robot_server.hpp>

namespace saf {

class RobotServerDVRK: public RobotServer {

public:

  // Constants
  static const std::string READY;

protected:
  const ArmTypes arm_typ;

  // States
  std_msgs::String status;
  sensor_msgs::JointState measured_js;
  geometry_msgs::TransformStamped measured_cp;
  std_msgs::String error;
  std_msgs::String warning;

  // Subscribers
  ros::Subscriber status_sub;
  ros::Subscriber smeasured_js_sub;
  ros::Subscriber measured_cp_sub;
  ros::Subscriber error_sub;
  ros::Subscriber warning_sub;

  // Publishers
  ros::Publisher position_joint_pub;
  ros::Publisher position_cartesian_pub;

  void subscribeLowLevelTopics();
  void advertiseLowLevelTopics();

public:
  RobotServerDVRK(ros::NodeHandle, ros::NodeHandle, ArmTypes, std::string, bool);
  ~RobotServerDVRK();

  // Callbacks
  void resetPose(bool);
  void stop();
  void followTrajectory(Trajectory<ToolPose>);
  void moveJointAbsolute(sensor_msgs::JointState , double );

  void status_cb(const std_msgs::String);
  void error_cb(const std_msgs::String);
  void warning_cb(const std_msgs::String);
  void measured_js_cb(const sensor_msgs::JointStateConstPtr&);
  virtual void measured_cp_cb(
      const geometry_msgs::TransformStampedConstPtr&);

  void loadRegistration(std::string);

  double getJointStateCurrent(int);
  std::vector<double> getJointStateCurrent();
  Eigen::Vector3d getPositionCartesianCurrent();
  Eigen::Quaternion<double> getOrientationCartesianCurrent();
  ToolPose getPoseCurrent();

  //DVRK actions
  std::string getCurrentState();
  void moveCartesianRelative(Eigen::Translation3d, double = 0.01);
  virtual void moveCartesianAbsolute(ToolPose, double = 0.01);

  void recordTrajectory(Trajectory<Eigen::Vector3d>&);
  void recordTrajectory(Trajectory<ToolPose>&);
  void saveTrajectory(std::string);

  void checkErrors();
  void checkVelCartesian(const ToolPose&, const ToolPose&, double);
  void checkNaNCartesian(const ToolPose&);
  void checkVelJoint(const sensor_msgs::JointState&,
                     const std::vector<double>&, double);
  sensor_msgs::JointState maximizeVelJoint(const sensor_msgs::JointState&,
                     const std::vector<double>&, double);
  void checkNaNJoint(const sensor_msgs::JointState&);

};

}
#endif /* ROBOT_SERVER_DVRK_HPP_ */


//
// Copyright (c) 2019-2020 Kris Jusiak (kris at jusiak dot net)
//
// 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 <algorithm>
#include <boost/ut.hpp>
#include <string_view>
#include <tuple>
#include <type_traits>

class matcher : boost::ut::detail::op {
 public:
  matcher(bool result, const std::string& str) : result_{result}, str_{str} {}
  constexpr explicit operator bool() const { return result_; }
  friend auto operator<<(std::ostream& os, const matcher& self)
      -> std::ostream& {
    return (os << self.str_);
  }

 private:
  const bool result_{};
  const std::string str_{};
};

template <class... Ts>
class any_of {
 public:
  constexpr explicit any_of(Ts... ts) : ts_{ts...} {}

  constexpr auto operator==(std::common_type_t<Ts...> t) const {
    return std::apply([t](auto... args) { return eq(t, args...); }, ts_);
  }

 private:
  template <class T, class U, class... TArgs>
  static constexpr auto eq(const T& t, const U& u, const TArgs&... args) {
    using namespace boost::ut;
    if constexpr (sizeof...(args) > 0) {
      return (that % detail::value{u} == t) or eq(t, args...);
    } else {
      return (that % detail::value{u} == t);
    }
  }

  std::tuple<Ts...> ts_;
};

int main() {
  using namespace boost::ut;
  using namespace std::literals::string_view_literals;

  "matcher"_test = [] {
    constexpr auto is_between = [](auto lhs, auto rhs) {
      return [=](auto value) {
        return that % value >= lhs and that % value <= rhs;
      };
    };

    constexpr auto ends_with = [](const auto& arg, const auto& ext) {
      std::stringstream str{};
      str << '(' << arg << " ends with " << ext << ')';
      if (ext.size() > arg.size()) {
        return matcher{{}, str.str()};
      }
      return matcher{std::equal(ext.rbegin(), ext.rend(), arg.rbegin()),
                     str.str()};
    };

    auto value = 42;
    auto str = "example.test"sv;

    expect(is_between(1, 100)(value) and not is_between(1, 100)(0));
    expect(ends_with(str, ".test"sv));
    expect(any_of{1, 2, 3} == 2 or any_of{42, 43} == 44);
  };
}


//==================================================================================================
/*!
  @file

  @copyright 2016 NumScale SAS

  Distributed under the Boost Software License, Version 1.0.
  (See accompanying file LICENSE.md or copy at http://boost.org/LICENSE_1_0.txt)
*/
//==================================================================================================
#ifndef BOOST_SIMD_ARCH_COMMON_SIMD_CONSTANT_CONSTANT_VALUE_HPP_INCLUDED
#define BOOST_SIMD_ARCH_COMMON_SIMD_CONSTANT_CONSTANT_VALUE_HPP_INCLUDED

#if defined(BOOST_SIMD_DETECTED)
#include <boost/simd/meta/hierarchy/simd.hpp>
#include <boost/simd/detail/constant_traits.hpp>
#include <boost/simd/detail/dispatch/function/overload.hpp>
#include <boost/simd/detail/dispatch/as.hpp>
#include <boost/config.hpp>
#include <boost/simd/detail/nsm.hpp>

namespace boost { namespace simd { namespace ext
{
  namespace bd = boost::dispatch;
  namespace bs = boost::simd;

  BOOST_DISPATCH_OVERLOAD_FALLBACK( (typename Constant, typename T, typename X)
                                  , boost::dispatch::constant_value_<Constant>
                                  , bs::simd_
                                  , bd::target_< bs::pack_<bd::unspecified_<T>,X> >
                                  )
  {
    using value_t   = typename T::type::value_type;
    using scalar_t  = decltype(bd::functor<Constant>{}(bd::as_<value_t>()));
    using result_t  = typename T::type::template rebind<scalar_t>;

    BOOST_FORCEINLINE result_t operator()(T const&) const
    {
      return result_t{ bd::functor<Constant>{}(bd::as_<value_t>()) };
    }
  };
} } }

#endif

#endif


/*
 * OpusHandler.cpp
 *
 *  Created on: 5 oct. 2018
 *      Author: max
 */

#include <boost/asio.hpp>
#include <opus/opus_defines.h>
#include <OpusHandler.hpp>
#include <StreamHandler.hpp>
#include <cstdio>

namespace Babel {

	OpusHandler::OpusHandler() :
			_decodeSize(0), _encodeSize(0) {
		_encoder = opus_encoder_create(SoundDevice::sampleRate, SoundDevice::channelCount, OPUS_APPLICATION_VOIP, NULL);
		_decoder = opus_decoder_create(SoundDevice::sampleRate, SoundDevice::channelCount, NULL);

		this->_decodedData = new std::vector<float>(SoundDevice::channelCount * SoundDevice::sampleRate * SoundDevice::secTimeOfEachPacket);
		this->_encodedData = new std::vector<unsigned char>(SoundDevice::channelCount * SoundDevice::sampleRate * SoundDevice::secTimeOfEachPacket);
	}

	OpusHandler::~OpusHandler() {
	        opus_encoder_destroy(_encoder);
	        opus_decoder_destroy(_decoder);
	}

	std::vector<float> *OpusHandler::decode(std::vector<unsigned char> sound) {
		/*
		for (int i = 0; i < sound.size(); i++) {
			if (sound.at(i) != 0) {
				printf("BBBAAA%d\n", sound.at(i));
			}
		}*/

		//_decodedData->resize(sound.size());
		_decodeSize = opus_decode_float(_decoder, sound.data(), _encodeSize, _decodedData->data(), 480, 0); //* SoundDevice::channelCount;
		if (_decodeSize < 0) {
			throw("opus_decode_float : Couldn''t decode data.");
		}
		//printf("RET = %d\nSIZE = %lu\n", _decodeSize, _decodedData->size() );

		/*
		for (int i = 0; i < _decodedData->size(); i++) {
			if (_decodedData->at(i) != 0) {
				printf("DECODED = %f.\n", _decodedData->at(i));
			}
		}*/

		return _decodedData;
	}

	std::vector<unsigned char> *OpusHandler::encode(std::vector<float> sound) {

		printf("SIZE = %lu\n", sound.size());

		/*
		for (int i = 0; i < sound.size(); i++) {
			if (sound.at(i)) {
				printf("%f\n", sound.at(i));
			}
		}*/

		//_encodedData->resize(sound.size());
		_encodeSize = opus_encode_float(_encoder, sound.data(), 480, _encodedData->data(), sound.size());
		if (_encodeSize < 0) {
			throw("opus_encode_float : Couldn''t encode data.");
		}
		//printf("ret = %d\n", _encodeSize);
		/*

		for (int i = 0; i < _encodedData->size(); i++) {
			if (_encodedData->at(i) != 0) {
				printf("AAABBB = %d.\n", _encodedData->at(i));
			}
		}*/

		return _encodedData;
	}

} /* namespace Babel */


// Boost.Units - A C++ library for zero-overhead dimensional analysis and
// unit/quantity manipulation and conversion
//
// Copyright (C) 2003-2008 Matthias Christian Schabel
// Copyright (C) 2008 Steven Watanabe
//
// Distributed under the Boost Software License, Version 1.0. (See
// accompanying file LICENSE_1_0.txt or copy at
// http://www.boost.org/LICENSE_1_0.txt)

#ifndef BOOST_UNITS_CODATA_ELECTROMAGNETIC_CONSTANTS_HPP
#define BOOST_UNITS_CODATA_ELECTROMAGNETIC_CONSTANTS_HPP

///
/// \file
/// \brief CODATA recommended values of fundamental electromagnetic constants.
/// \details CODATA recommended values of the fundamental physical constants: NIST SP 961
///   CODATA 2006 values as of 2007/03/30
///

#include <boost/units/quantity.hpp>
#include <boost/units/static_constant.hpp>

#include <boost/units/systems/detail/constants.hpp>
#include <boost/units/systems/si/conductance.hpp>
#include <boost/units/systems/si/current.hpp>
#include <boost/units/systems/si/electric_charge.hpp>
#include <boost/units/systems/si/electric_potential.hpp>
#include <boost/units/systems/si/energy.hpp>
#include <boost/units/systems/si/frequency.hpp>
#include <boost/units/systems/si/magnetic_flux.hpp>
#include <boost/units/systems/si/magnetic_flux_density.hpp>
#include <boost/units/systems/si/resistance.hpp>

#include <boost/units/systems/si/codata/typedefs.hpp>

namespace boost {

namespace units {

namespace si {

namespace constants {

namespace codata {

// ELECTROMAGNETIC
/// elementary charge
BOOST_UNITS_PHYSICAL_CONSTANT(e,quantity<electric_charge>,1.602176487e-19*coulombs,4.0e-27*coulombs);
/// elementary charge to Planck constant ratio
BOOST_UNITS_PHYSICAL_CONSTANT(e_over_h,quantity<current_over_energy>,2.417989454e14*amperes/joule,6.0e6*amperes/joule);
/// magnetic flux quantum
BOOST_UNITS_PHYSICAL_CONSTANT(Phi_0,quantity<magnetic_flux>,2.067833667e-15*webers,5.2e-23*webers);
/// conductance quantum
BOOST_UNITS_PHYSICAL_CONSTANT(G_0,quantity<conductance>,7.7480917004e-5*siemens,5.3e-14*siemens);
/// Josephson constant
BOOST_UNITS_PHYSICAL_CONSTANT(K_J,quantity<frequency_over_electric_potential>,483597.891e9*hertz/volt,1.2e7*hertz/volt);
/// von Klitzing constant
BOOST_UNITS_PHYSICAL_CONSTANT(R_K,quantity<resistance>,25812.807557*ohms,1.77e-5*ohms);
/// Bohr magneton
BOOST_UNITS_PHYSICAL_CONSTANT(mu_B,quantity<energy_over_magnetic_flux_density>,927.400915e-26*joules/tesla,2.3e-31*joules/tesla);
/// nuclear magneton
BOOST_UNITS_PHYSICAL_CONSTANT(mu_N,quantity<energy_over_magnetic_flux_density>,5.05078324e-27*joules/tesla,1.3e-34*joules/tesla);

} // namespace codata

} // namespace constants

} // namespace si

} // namespace units

} // namespace boost

#endif // BOOST_UNITS_CODATA_ELECTROMAGNETIC_CONSTANTS_HPP


// CZI.
#include "CellGraph.hpp"
#include "color.hpp"
#include "CZI_ASSERT.hpp"
#include "deduplicate.hpp"
#include "iostream.hpp"
#include "iterator.hpp"
#include "SimilarPairs.hpp"
#include "timestamp.hpp"
using namespace ChanZuckerberg::ExpressionMatrix2;

// Boost libraries.
#include "boost_lexical_cast.hpp"
#include <boost/graph/graphviz.hpp>
#include <boost/graph/iteration_macros.hpp>
#include <boost/algorithm/string.hpp>
using boost::algorithm::split;
using boost::algorithm::is_any_of;

// Standard libraries.
#include "algorithm.hpp"
#include <chrono>
#include "fstream.hpp"
#include "set.hpp"
#include "stdexcept.hpp"
#include "utility.hpp"
#include "vector.hpp"
#include <limits>
#include <random>

CellGraph::CellGraph(
    const MemoryMapped::Vector<CellId>& cellSet, // The cell set to be used.
    const string& similarPairsName,              // The name of the SimilarPairs object to be used to create the graph.
    double similarityThreshold,                  // The minimum similarity to create an edge.
    size_t maxConnectivity                      // The maximum number of neighbors (k of the k-NN graph).
 )
{
    typedef SimilarPairs::Pair Pair;

    // Create the SimilarPairs object.
    const SimilarPairs similarPairs(similarPairsName, true);

    // Create a vertex for each cell in the cell set.
    for(const CellId cellId: cellSet) {
        const vertex_descriptor v = boost::add_vertex(CellGraphVertex(cellId), graph());
        vertexTable.insert(make_pair(cellId, v));
    }

    // Create the edges.
    // The similar pairs are sorted by decreasing similarity.
    vector< pair<vertex_descriptor, float > > pairs;
    for(const CellId cellId0: cellSet) {

        // Find the local cell id (in the cell set of the SimilarPairs object)
        // corresponding to this global cell id.
        // If the cell set of the SimilarPairs object does not contain this cell,
        // this is an invalid cell id and in that case we skip this cell.
        const CellId localCellId0 = similarPairs.getLocalCellId(cellId0);
        if(localCellId0 == invalidCellId) {
            continue;
        }

        // Locate the corresponding vertex.
        const vertex_descriptor v0 = vertexTable[cellId0];

        // Find the best up to k pairs such that the other vertex
        // is also in the cell set.
        pairs.clear();
        const Pair* begin = similarPairs.begin(localCellId0);
        const Pair* end = similarPairs.end(localCellId0);
        for(const Pair* p=begin; p!=end; ++p) {
            const float similarity = p->second;
            if(similarity < similarityThreshold) {
                break;
            }
            const CellId cellId1 = similarPairs.getGlobalCellId(p->first);
            const auto it1 = vertexTable.find(cellId1);
            if(it1 == vertexTable.end()) {
                continue;
            }
            pairs.push_back(make_pair(it1->second, similarity));
            if(pairs.size() == maxConnectivity) {
                break;
            }
        }

        // Add an edge for each pair we found.
        // Some of the edges may already exist (we added them when
        // processing cellId1).
        for(const pair<vertex_descriptor, float >& p: pairs) {
            const vertex_descriptor v1 = p.first;
            edge_descriptor e;
            bool edgeExists = false;
            tie(e, edgeExists) = boost::edge(v0, v1, graph());
            if(edgeExists) {
                continue;
            }
            boost::add_edge(v0, v1, CellGraphEdge(p.second), graph());
        }
    }

    // cout << "The cell similarity graph has " << boost::num_vertices(graph());
    // cout << " vertices and " << num_edges(graph()) << " edges." << endl;

}

// Write the graph in Graphviz format.
void CellGraph::write(const string& fileName) const
    {
    ofstream outputFileStream(fileName);
    if(!outputFileStream) {
        throw runtime_error("Error opening " + fileName);
    }
    write(outputFileStream);
}
void CellGraph::write(ostream& s) const
    {
    Writer writer(*this);
    boost::write_graphviz(s, graph(), writer, writer, writer,
        boost::get(&CellGraphVertex::cellId, graph()));
}

CellGraph::Writer::Writer(const Graph& graph) :
    graph(graph)
{
}

void CellGraph::Writer::operator()(std::ostream& s) const
{
    s << "tooltip=\"\";";
    s << "node [shape=point];";
}

// Write out a vertex of the cell graph.
void CellGraph::Writer::operator()(std::ostream& s, vertex_descriptor v) const
{
    // Get the vertex.
    const CellGraphVertex& vertex = graph[v];

    // Begin vertex attributes.
    s << "[";

    // Add a tooltip that shows the cell id.
    s << "tooltip=" << vertex.cellId;

    // End vertex attributes.
    s << "]";
}

void CellGraph::Writer::operator()(std::ostream& s, edge_descriptor e) const
{
    // Get the edge.
    const CellGraphEdge& edge = graph[e];

    // Begin edge attributes.
    s << "[";

    // Add a tooltip that shows the similarity.
    s.precision(2);
    s.setf(std::ios::fixed);
    s << "tooltip=\"" << edge.similarity << "\"";

    // End edge attributes.
    s << "]";
}

// Remove isolated vertices and returns\ the number of vertices that were removed
size_t CellGraph::removeIsolatedVertices()
{
    vector<vertex_descriptor> verticesToBeRemoved;
    BGL_FORALL_VERTICES(v, graph(), Graph) {
        if(out_degree(v, graph()) == 0) {
            verticesToBeRemoved.push_back(v);
        }
    }

    for(const vertex_descriptor v: verticesToBeRemoved) {
        const CellGraphVertex& vertex = graph()[v];
        vertexTable[vertex.cellId] = null_vertex();
        boost::remove_vertex(v, graph());
    }

    return verticesToBeRemoved.size();
}

// Use Graphviz to compute the graph layout and store it in the vertex positions.
void CellGraph::computeLayout()
{
    // Write the graph in Graphviz format.
    write("Graph.dot");

    // Run sfdp with output in Graphviz plain format.
    // See https://www.graphviz.org/doc/info/output.html#d:plain
    const int systemReturnCode = ::system("sfdp -O -Tplain -Goverlap=true -Gsmoothing=triangle Graph.dot");
    const int sfdpReturnCode = WEXITSTATUS(systemReturnCode);   // Man page for system is not super clear on this.
    if(sfdpReturnCode!=0 && sfdpReturnCode!=1) {    // Sfdp returns 1 if build without triangulation library.
        throw runtime_error("Error " +
            lexical_cast<string>(systemReturnCode) + " " +
            lexical_cast<string>(sfdpReturnCode) +
            " running sfdp.");
    }

    // Extract vertex positions from the output.
    ifstream file("Graph.dot.plain");
    string line;
    vector<string> tokens;
    while(true) {

        // Get a line.
        getline(file, line);
        if(!file) {
            break;
        }

        // Parse it.
        split(tokens, line, is_any_of(" "));

        // Only parse lines that describe vertices.
        CZI_ASSERT(tokens.size() >= 1);
        if(tokens[0] != "node") {
            continue;
        }
        CZI_ASSERT(tokens.size() >= 4);

        // Extract the positions for this vertex.
        try {
            const CellId cellId = lexical_cast<CellId>(tokens[1]);
            const vertex_descriptor v = vertexTable[cellId];
            CZI_ASSERT(v != null_vertex());
            CellGraphVertex& vertex = graph()[v];
            vertex.position[0] = lexical_cast<double>(tokens[2]);
            vertex.position[1] = lexical_cast<double>(tokens[3]);
        } catch(std::exception& e) {
            cout << "Error processing the following line of Graph.dot.plain:" << endl;
            cout << line << endl;
            throw;
        }
    }
}

// Compute minimum and maximum coordinates of all the vertices.
void CellGraph::computeCoordinateRange(
    double& xMin,
    double& xMax,
    double& yMin,
    double& yMax) const
{
    xMin = std::numeric_limits<double>::max();
    xMax = std::numeric_limits<double>::min();
    yMin = std::numeric_limits<double>::max();
    yMax = std::numeric_limits<double>::min();
    BGL_FORALL_VERTICES(v, graph(), Graph) {
        const CellGraphVertex& vertex = graph()[v];
        const double x = vertex.position[0];
        const double y = vertex.position[1];
        xMin = min(xMin, x);
        xMax = max(xMax, x);
        yMin = min(yMin, y);
        yMax = max(yMax, y);
    }

}

// Write the graph in svg format.
// This does not use Graphviz. It uses the graph layout stored in the vertices,
// and previously computed using Graphviz.
// The vertex coordinates are used without any transformation.
// The last argument specifies the color assigned to each vertex group.
// If empty, vertex groups are not used, and each vertex is drawn
// with its own color.
void CellGraph::writeSvg(
    ostream& s,
    bool hideEdges,
    double svgSizePixels,
    double xViewBoxCenter,
    double yViewBoxCenter,
    double viewBoxHalfSize,
    double vertexRadius,
    double edgeThickness,
    const map<int, string>& groupColors,
    const string& geneSetName   // Used for the cell URL
    ) const
{

    // Start the svg object.
    s <<
        "
<svg id=graphSvg width='" << svgSizePixels << "' height='" << svgSizePixels <<
        "' viewBox='" <<
        xViewBoxCenter-viewBoxHalfSize << " " << yViewBoxCenter-viewBoxHalfSize << " " <<
        2.*viewBoxHalfSize << " " << 2.*viewBoxHalfSize <<
        "'>";

    // Draw the edges first.
    // This makes it easier to see the vertices and their tooltips.
    if(!hideEdges) {
        s << "<g id=edges>";
        BGL_FORALL_EDGES(e, graph(), Graph){
            const vertex_descriptor v1 = source(e, graph());
            const vertex_descriptor v2 = target(e, graph());

            const CellGraphVertex& vertex1 = graph()[v1];
            const CellGraphVertex& vertex2 = graph()[v2];

            const double x1 = vertex1.position[0];
            const double y1 = vertex1.position[1];
            const double x2 = vertex2.position[0];
            const double y2 = vertex2.position[1];

            s << "<line x1='" << x1 << "' y1='" << y1 << "'";
            s << " x2='" << x2 << "' y2='" << y2 << "'";

            s << " style='stroke:";
            const string& color = graph()[e].color;
            if(color.empty()) {
                s << "black";
            } else {
                s << color;
            }
            s << ";stroke-width:" << edgeThickness << "' />";
        }
        s << "</g>";
    }

    // If the groupColors map is empty, vertex groups are not used,
    // and each vertex is drawn in its own color.
    // We still write two levels of groups, because the javascript code
    // to change vertex size expects that structure.
    if(groupColors.empty()) {
        s << "<g id=vertices><g>";
        BGL_FORALL_VERTICES(v, graph(), Graph) {
            const CellGraphVertex& vertex = graph()[v];
            const double x = vertex.position[0];
            const double y = vertex.position[1];
            s <<
                "<a xlink:href='cell?cellId=" << vertex.cellId << "&geneSetName=" << geneSetName << "'>"
                "<circle cx='" << x << "' cy='" << y << "' r='" << vertexRadius << "' stroke=none";
            if(!vertex.color.empty()) {
                s << " fill='" << vertex.color << "'";
            }
            s <<
                ">"
                "<title>Cell " << vertex.cellId << "</title></circle>"
                ""
                ;
        }
        s << "</g></g>";
    }

    // Otherwise, colors are specified for each group, not for each vertex.
    else {

        // Find the vertices in each group.
        vector< vector<vertex_descriptor> > groups;
        BGL_FORALL_VERTICES(v, graph(), Graph) {
            const CellGraphVertex& vertex = graph()[v];
            const size_t group = vertex.group;
            if(groups.size() <= group) {
                groups.resize(group+1);
            }
            groups[group].push_back(v);
        }

        // A circle at the center.
        // s << "<circle cx='" << svgSizePixels/2 << "' cy='" << svgSizePixels/2 << "' r='" << 10 << "' stroke='black' stroke-width='3' fill='red' />";

        // Draw the vertices, one group at a time.
        s << "<g id=vertices>";

        // Loop over all groups.
        for(int iGroup=0; iGroup<int(groups.size()); iGroup++) {
            string groupColor;
            const auto it = groupColors.find(iGroup);
            if(it == groupColors.end()) {
                groupColor = "black";
            } else {
                groupColor = it->second;
            }
            auto& group= groups[iGroup];
            s << "<g id=vertexGroup" << iGroup << " style='fill:" << groupColor << "'>";

            // Loop over all vertices of this group.
            for(const vertex_descriptor v: group) {
                const CellGraphVertex& vertex = graph()[v];

                const double x = vertex.position[0];
                const double y = vertex.position[1];
                s <<
                    "<a xlink:href='cell?cellId=" << vertex.cellId << "&geneSetName=" << geneSetName << "'>"
                    "<circle cx='" << x << "' cy='" << y << "' r='" << vertexRadius << "' stroke=none>"
                    "<title>Cell " << vertex.cellId << "</title></circle>"
                    ""
                    ;
            }
            s << "</g>";
        }
        s << "</g>";
    }

}

#if 1
// Clustering using the label propagation algorithm.
// The cluster each vertex is assigned to is stored in the clusterId data member of the vertex.
void CellGraph::labelPropagationClustering(
    ostream& out,
    size_t seed,                            // Seed for random number generator.
    size_t stableIterationCountThreshold,   // Stop after this many iterations without changes.
    size_t maxIterationCount                // Stop after this many iterations no matter what.
    )
{
    out << timestamp << "Clustering by label propagation begins." << endl;
    out << "Seed for random number generator is " << seed << "." << endl;
    out << "Will stop after " << stableIterationCountThreshold << " iterations without changes." << endl;
    out << "Maximum number of iterations is " << maxIterationCount << "." << endl;
    const auto t0 = std::chrono::steady_clock::now();

    // Set the cluster of each vertex equal to its cell id.
    BGL_FORALL_VERTICES(v, graph(), CellGraph) {
        CellGraphVertex& vertex = graph()[v];
        vertex.clusterId = vertex.cellId;
    }

    // Initialize the ClusterTable of each vertex.
    BGL_FORALL_VERTICES(v0, graph(), CellGraph) {
        CellGraphVertex& vertex0 = graph()[v0];
        ClusterTable& clusterTable0 = vertex0.clusterTable;
        clusterTable0.clear();
        BGL_FORALL_OUTEDGES(v0, e, graph(), CellGraph) {
            const vertex_descriptor v1 = target(e, graph());
            const CellGraphVertex& vertex1 = graph()[v1];
            const CellGraphEdge& edge = graph()[e];
            clusterTable0.addWeightQuick(vertex1.clusterId, edge.similarity);
        }
        clusterTable0.findBestCluster();
    }

    // Create the random number generator using the specified seed.
    std::mt19937 randomGenerator(seed);

    // Vector with all the vertices in the graph, in the same order as in the vertex table.
    vector<vertex_descriptor> allVertices;
    for(const auto& p : vertexTable) {
        const vertex_descriptor v = p.second;
        if(v != null_vertex()) {
            allVertices.push_back(p.second);
        }
    }

    // Vector to contain the vertices in the random order to be used at each iteration.
    vector<vertex_descriptor> shuffledVertices;

    // Counter of the number of stable iterations
    // (iterations without changes).
    size_t stableIterationCount = 0;

    // Iterate.
    out << timestamp << "Label propagation iteration begins." << endl;
    for(size_t iteration=0; iteration<maxIterationCount; iteration++) {
        // cout << "Begin iteration " << iteration << endl;
        const auto t0 = std::chrono::steady_clock::now();
        size_t changeCount = 0;

        // Create a random shuffle of the vertices, to be used for this iteration.
        shuffledVertices = allVertices;
        std::shuffle(shuffledVertices.begin(), shuffledVertices.end(), randomGenerator);

        // Process the vertices in the order determined by the random shuffle.
        for(const vertex_descriptor v0: shuffledVertices) {
            CZI_ASSERT(v0 != CellGraph::null_vertex());
            CellGraphVertex& vertex0 = graph()[v0];
            if(vertex0.clusterTable.isEmpty()) {
                continue;
            }

            // If the cluster is already consistent with the cluster table,
            // we don't need to do anything.
            const uint32_t bestClusterId = vertex0.clusterTable.bestCluster();
            if(vertex0.clusterId == bestClusterId) {
                continue;
            }

            // Change the cluster id of vertex0.
            const uint32_t oldClusterId = vertex0.clusterId;
            vertex0.clusterId = bestClusterId;
            ++changeCount;

            // Update the cluster table of its neighbors.
            BGL_FORALL_OUTEDGES(v0, e, graph(), CellGraph) {
                const vertex_descriptor v1 = target(e, graph());
                CellGraphVertex& vertex1 = graph()[v1];
                ClusterTable& clusterTable1 = vertex1.clusterTable;
                const CellGraphEdge& edge = graph()[e];
                clusterTable1.addWeight(bestClusterId, edge.similarity);
                clusterTable1.addWeight(oldClusterId, -edge.similarity);
            }
        }
        const auto t1 = std::chrono::steady_clock::now();
        const double t01 = 1.e-9 * double((std::chrono::duration_cast<std::chrono::nanoseconds>(t1 - t0)).count());
        out << "Iteration " << iteration << " took " << t01 << " s, made " << changeCount << " changes." << endl;

        // Update the number of stable iterations (iterations without changes).
        if(changeCount) {
            stableIterationCount = 0;
        } else {
            ++stableIterationCount;
        }

        // If we have done enough stable iterations, stop.
        if(stableIterationCount == stableIterationCountThreshold) {
            break;
        }
    }

    if(stableIterationCount == stableIterationCountThreshold) {
        out << "Terminating because the specified number of stable iterations was achieved." << endl;
    } else {
        out << "Terminating because the maximum number of iterations was reached." << endl;
    }

    // Compute the size of each cluster.
    map<uint32_t, size_t> clusterSize;    // Key=clusterId, Value=cluster size
    BGL_FORALL_VERTICES(v, graph(), CellGraph) {
        const uint32_t clusterId = graph()[v].clusterId;
        const auto it = clusterSize.find(clusterId);
        if(it == clusterSize.end()) {
            clusterSize.insert(make_pair(clusterId, 1));
        } else {
            ++(it->second);
        }
    }

    // Renumber the clusters beginning at 0 and in order of decreasing cluster size.
    vector< pair<size_t, uint32_t> > clusterSizeVector;   // first:cluster size, second: clusterId
    for(const auto& p: clusterSize) {
        clusterSizeVector.push_back(make_pair(p.second, p.first));
    }
    sort(clusterSizeVector.begin(), clusterSizeVector.end(), std::greater< pair<size_t, size_t> >());
    out << "Cluster sizes:";
    for(size_t newClusterId=0; newClusterId<clusterSizeVector.size(); newClusterId++) {
        const auto& p = clusterSizeVector[newClusterId];
        out << " " << p.first;
    }
    out << endl;
    map<uint32_t, uint32_t> clusterMap; // Key: old clusterId. Value: new clustyerId.
    for(uint32_t newClusterId=0; newClusterId<clusterSizeVector.size(); newClusterId++) {
        const uint32_t oldClusterId = clusterSizeVector[newClusterId].second;
        clusterMap.insert(make_pair(oldClusterId, newClusterId));
    }

    // Update the vertices to reflect the new cluster numbering.
    BGL_FORALL_VERTICES(v, graph(), CellGraph) {
        CellGraphVertex& vertex = graph()[v];
        vertex.clusterId = clusterMap[vertex.clusterId];
    }

    const auto t1 = std::chrono::steady_clock::now();
    const double t01 = 1.e-9 * double((std::chrono::duration_cast<std::chrono::nanoseconds>(t1 - t0)).count());
    out << timestamp << "Clustering by label propagation completed in " << t01 << " s." << endl;
}

#else
// Clustering using the label propagation algorithm.
// The cluster each vertex is assigned to is stored in the clusterId data member of the vertex.
void CellGraph::labelPropagationClustering(
    ostream& out,
    size_t seed,                            // Seed for random number generator.
    size_t stableIterationCountThreshold,   // Stop after this many iterations without changes.
    size_t maxIterationCount                // Stop after this many iterations no matter what.
    )
{
    out << timestamp << "Clustering by label propagation begins." << endl;
    out << "Seed for random number generator is " << seed << "." << endl;
    out << "Will stop after " << stableIterationCountThreshold << " iterations without changes." << endl;
    out << "Maximum number of iterations is " << maxIterationCount << "." << endl;

    // Set the cluster of each vertex equal to its cell id.
    BGL_FORALL_VERTICES(v, graph(), CellGraph) {
        CellGraphVertex& vertex = graph()[v];
        vertex.clusterId = vertex.cellId;
    }

    // Create the random number generator using the specified seed.
    std::mt19937 randomGenerator(seed);

    // Vector with all the vertices in the graph, in the same order as in the vertex table.
    vector<vertex_descriptor> allVertices;
    for(const auto& p : vertexTable) {
        const vertex_descriptor v = p.second;
        if(v != null_vertex()) {
            allVertices.push_back(p.second);
        }
    }

    // Vector to contain the vertices in the random order to be used at each iteration.
    vector<vertex_descriptor> shuffledVertices;

    // Counter of the number of stable iterations
    // (iterations without changes).
    size_t stableIterationCount = 0;

    // Iterate.
    for(size_t iteration=0; iteration<maxIterationCount; iteration++) {
        size_t changeCount = 0;

        // Create a random shuffle of the vertices, to be used for this iteration.
        shuffledVertices = allVertices;
        std::shuffle(shuffledVertices.begin(), shuffledVertices.end(), randomGenerator);

        // Process the vertices in the order determined by the random shuffle.
        for(const vertex_descriptor v0: shuffledVertices) {
            CZI_ASSERT(v0 != CellGraph::null_vertex());
            CellGraphVertex& vertex0 = graph()[v0];
            const size_t clusterId0 = vertex0.clusterId;

            // Map to contain the clusters of the neighbors of this vertex, each with its total weight.
            // This could be written for better performance.
            map<uint32_t, double> clusterMap;

            // Loop over the edges of this vertex.
            BGL_FORALL_OUTEDGES(v0, e, graph(), CellGraph) {
                const vertex_descriptor v1 = target(e, graph());
                const double weight = graph()[e].similarity;
                const uint32_t clusterId1 = graph()[v1].clusterId;

                const auto it = clusterMap.find(clusterId1);
                if(it == clusterMap.end()) {
                    clusterMap.insert(make_pair(clusterId1, weight));
                } else {
                    it->second += weight;
                }
            }
            if(clusterMap.empty()) {
                continue;   // There were no neighbors. Nothing to do.
            }

            // Find the cluster with the most weight.
            // Note that we don't bother to deal with ties explicitly.
            uint32_t bestClusterId = std::numeric_limits<uint32_t>::max();
            double bestWeight = -1;
            for(const auto& p: clusterMap) {
                const uint32_t clusterId = p.first;
                const double weight = p.second;
                if(weight > bestWeight) {
                    bestClusterId = clusterId;
                    bestWeight = weight;
                }
            }
            CZI_ASSERT(bestClusterId != std::numeric_limits<uint32_t>::max());

            // The best cluster id becomes the cluster id of vertex v0.
            if(bestClusterId != clusterId0) {
                vertex0.clusterId = bestClusterId;
                ++changeCount;
            }

        }
        out << "Iteration " << iteration << ": " << changeCount << " changes." << endl;

        // Scroll to the bottom of the page to show the line we just wrote.
        // But for some reason it does not work.
        // out << "" << endl;

        // Update the number of stable iterations (iterations without changes).
        if(changeCount) {
            stableIterationCount = 0;
        } else {
            ++stableIterationCount;
        }

        // If we have done enough stable iterations, stop.
        if(stableIterationCount == stableIterationCountThreshold) {
            break;
        }
    }

    if(stableIterationCount == stableIterationCountThreshold) {
        out << "Terminating because the specified number of stable iterations was achieved." << endl;
    } else {
        out << "Terminating because the maximum number of iterations was reached." << endl;
    }

    // Compute the size of each cluster.
    map<uint32_t, size_t> clusterSize;    // Key=clusterId, Value=cluster size
    BGL_FORALL_VERTICES(v, graph(), CellGraph) {
        const uint32_t clusterId = graph()[v].clusterId;
        const auto it = clusterSize.find(clusterId);
        if(it == clusterSize.end()) {
            clusterSize.insert(make_pair(clusterId, 1));
        } else {
            ++(it->second);
        }
    }

    // Renumber the clusters beginning at 0 and in order of decreasing cluster size.
    vector< pair<size_t, uint32_t> > clusterSizeVector;   // first:cluster size, second: clusterId
    for(const auto& p: clusterSize) {
        clusterSizeVector.push_back(make_pair(p.second, p.first));
    }
    sort(clusterSizeVector.begin(), clusterSizeVector.end(), std::greater< pair<size_t, size_t> >());
    out << "Cluster sizes:";
    for(size_t newClusterId=0; newClusterId<clusterSizeVector.size(); newClusterId++) {
        const auto& p = clusterSizeVector[newClusterId];
        out << " " << p.first;
    }
    out << endl;
    map<uint32_t, uint32_t> clusterMap; // Key: old clusterId. Value: new clustyerId.
    for(uint32_t newClusterId=0; newClusterId<clusterSizeVector.size(); newClusterId++) {
        const uint32_t oldClusterId = clusterSizeVector[newClusterId].second;
        clusterMap.insert(make_pair(oldClusterId, newClusterId));
    }

    // Update the vertices to reflect the new cluster numbering.
    BGL_FORALL_VERTICES(v, graph(), CellGraph) {
        CellGraphVertex& vertex = graph()[v];
        vertex.clusterId = clusterMap[vertex.clusterId];
    }

    out << timestamp << "Clustering by label propagation ends." << endl;
}
#endif

// Assign integer colors to groups.
// The same color can be used for multiple groups, but if two
// groups are joined by one or more edges they must have distinct colors.
// On return, colorTable[group] contains the integer color assigned to each group.
// This processes the groups in increasing order beginning at group 0,
// so it is best if the group numbers are all contiguous, starting at zero,
// and in decreasing size of group.
void CellGraph::assignColorsToGroups(vector<uint32_t>& colorTable)
{
    // Start with no colors assigned.
    colorTable.clear();

    // Find the vertices of each group.
    vector< vector<vertex_descriptor> > groupVertices;
    BGL_FORALL_VERTICES(v, graph(), CellGraph) {
        const uint32_t groupId = graph()[v].group;
        if(groupVertices.size() <= groupId) {
            groupVertices.resize(groupId + 1);
        }
        groupVertices[groupId].push_back(v);
    }
    const size_t groupCount = groupVertices.size();

    // Create the group graph.
    // Each vertex corresponds ot a group.
    typedef boost::adjacency_list<boost::setS, boost::vecS, boost::undirectedS> GroupGraph;
    GroupGraph groupGraph(groupCount);
    BGL_FORALL_EDGES(e, graph(), CellGraph) {
        const vertex_descriptor v0 = source(e, graph());
        const vertex_descriptor v1 = target(e, graph());
        const CellGraphVertex& vertex0 = graph()[v0];
        const CellGraphVertex& vertex1 = graph()[v1];
        const uint32_t group0 = vertex0.group;
        const uint32_t group1 = vertex1.group;
        if(group0 != group1) {
            boost::add_edge(group0, group1, groupGraph);
        }
    }

    // ofstream graphOut("GroupGraph.dot");
    // boost::write_graphviz(graphOut, groupGraph);

    // For each group, we have to look at edges to lowered numbered groups
    vector<uint32_t> adjacentColors;
    for(uint32_t group0=0; group0<groupCount; group0++) {
        adjacentColors.clear();
        // cout << "Already colored groups adjacent to group " << group0 << ":";
        BGL_FORALL_OUTEDGES(group0, e, groupGraph, GroupGraph) {
            const uint32_t group1 = uint32_t(target(e, groupGraph));
            if(group1 < group0) {
                // cout << " " << group1;
                adjacentColors.push_back(colorTable[group1]);
            }
        }
        // cout << endl;
        deduplicate(adjacentColors);
        // cout << "Already assigned colors adjacent to group " << group0 << ": ";
        // copy(adjacentColors.begin(), adjacentColors.end(), ostream_iterator<uint32_t>(cout, " "));
        // cout << endl;

        // Assign to this group the smallest integer color that does not appear
        // in the adjacent groups.
        for(uint32_t color=0; color<adjacentColors.size(); color++) {
            if(adjacentColors[color] != color) {
                colorTable.push_back(color);
                break;
            }
        }
        if(colorTable.size() == group0) {
            colorTable.push_back(uint32_t(adjacentColors.size()));
        }

        // cout << "Group " << group0 << " assigned color " << colorTable[group0] << endl;
    }
}


/*
 * Costmap2DConverterTest.cpp
 *
 *  Created on: Nov 23, 2016
 *      Authors: Peter Fankhauser, Gabriel Hottiger
 *      Institute: ETH Zurich, ANYbotics
 */

// Grid map
#include <grid_map_core/grid_map_core.hpp>
#include <grid_map_costmap_2d/grid_map_costmap_2d.hpp>

// Gtest
#include <gtest/gtest.h>

// Eigen
#include <Eigen/Core>

using namespace grid_map;

template <typename ConversionTable>
class TestCostmap2DConversion : public testing::Test {
 public:
  using TestValue = std::tuple<Position, unsigned char, double, bool>;

  //! Constructor
  TestCostmap2DConversion() : costmap2d_(8, 5, 1.0, 2.0, 3.0) { gridMap_.setGeometry(Length(8.0, 5.0), 1.0, Position(6.0, 5.5)); }

  //! Getter of test data
  std::vector<TestValue> getTestValues();

  //! Check that maps have same geometry
  void assertMapGeometry(const GridMap& gridMap, const costmap_2d::Costmap2D& costMap) {
    // Check map info.
    // Different conventions: Costmap2d returns the *centerpoint* of the last cell in the map.
    Length length = gridMap.getLength() - Length::Constant(0.5 * gridMap.getResolution());
    Length position = gridMap.getPosition() - 0.5 * gridMap.getLength().matrix();
    EXPECT_EQ(costMap.getSizeInMetersX(), length.x());
    EXPECT_EQ(costMap.getSizeInMetersY(), length.y());
    EXPECT_EQ(costMap.getSizeInCellsX(), gridMap.getSize()[0]);
    EXPECT_EQ(costMap.getSizeInCellsY(), gridMap.getSize()[1]);
    EXPECT_EQ(costMap.getResolution(), gridMap.getResolution());
    EXPECT_EQ(costMap.getOriginX(), position.x());
    EXPECT_EQ(costMap.getOriginY(), position.y());
  }

 protected:
  Costmap2DConverter<GridMap, ConversionTable> costmap2dConverter_;
  costmap_2d::Costmap2D costmap2d_;
  GridMap gridMap_;
};

//! Map type that has only mappings for cost map special values.
using Costmap2DSpecialTranslationTable = Costmap2DTranslationTable<3, 2, 1, 0>;
//! Map type that has larger intervals between special values.
using Costmap2DLargeIntervalsTranslationTable = Costmap2DTranslationTable<5, 500, 400, 100>;

// Test data for special translation
template <>
std::vector<TestCostmap2DConversion<Costmap2DSpecialTranslationTable>::TestValue>
TestCostmap2DConversion<Costmap2DSpecialTranslationTable>::getTestValues() {
  std::vector<TestValue> testValues;
  testValues.emplace_back(TestValue(Position(3.2, 5.1), costmap_2d::FREE_SPACE, 0.f, true));
  testValues.emplace_back(TestValue(Position(4.2, 4.1), costmap_2d::INSCRIBED_INFLATED_OBSTACLE, 1.f, true));
  testValues.emplace_back(TestValue(Position(6.2, 3.1), costmap_2d::LETHAL_OBSTACLE, 2.f, true));
  testValues.emplace_back(TestValue(Position(5.2, 7.8), costmap_2d::NO_INFORMATION, 3.f, true));
  // Check for grid map to costmap only.
  testValues.emplace_back(TestValue(Position(5.4, 6.8), costmap_2d::FREE_SPACE, -1.f, false));
  testValues.emplace_back(TestValue(Position(8.5, 4.5), costmap_2d::LETHAL_OBSTACLE, 4.f, false));
  return testValues;
}

// Test data for special translation
template <>
std::vector<TestCostmap2DConversion<Costmap2DLargeIntervalsTranslationTable>::TestValue>
TestCostmap2DConversion<Costmap2DLargeIntervalsTranslationTable>::getTestValues() {
  std::vector<TestValue> testValues;
  testValues.emplace_back(TestValue(Position(3.2, 5.1), costmap_2d::FREE_SPACE, 100.f, true));
  testValues.emplace_back(TestValue(Position(4.2, 4.1), costmap_2d::INSCRIBED_INFLATED_OBSTACLE, 400.f, true));
  testValues.emplace_back(TestValue(Position(6.2, 3.1), costmap_2d::LETHAL_OBSTACLE, 500.f, true));
  testValues.emplace_back(TestValue(Position(5.2, 7.8), costmap_2d::NO_INFORMATION, 5.f, true));
  testValues.emplace_back(TestValue(Position(3.4, 3.8), 253/11, 100.f + 300.f/11.f, true));
  // Check for grid map to costmap only.
  testValues.emplace_back(TestValue(Position(5.4, 6.8), costmap_2d::FREE_SPACE, 83.f, false));
  testValues.emplace_back(TestValue(Position(8.5, 4.5), costmap_2d::LETHAL_OBSTACLE, 504.f, false));
  testValues.emplace_back(TestValue(Position(4.7, 5.8), costmap_2d::INSCRIBED_INFLATED_OBSTACLE, 444.f, false));

  return testValues;
}

// Test data for direct translation
template <>
std::vector<TestCostmap2DConversion<Costmap2DDirectTranslationTable>::TestValue>
TestCostmap2DConversion<Costmap2DDirectTranslationTable>::getTestValues() {
  std::vector<TestValue> testValues;
  testValues.emplace_back(TestValue(Position(3.2, 5.1), costmap_2d::FREE_SPACE, costmap_2d::FREE_SPACE, true));
  testValues.emplace_back(
      TestValue(Position(4.2, 4.1), costmap_2d::INSCRIBED_INFLATED_OBSTACLE, costmap_2d::INSCRIBED_INFLATED_OBSTACLE, true));
  testValues.emplace_back(TestValue(Position(6.2, 3.1), costmap_2d::LETHAL_OBSTACLE, costmap_2d::LETHAL_OBSTACLE, true));
  testValues.emplace_back(TestValue(Position(5.2, 7.8), costmap_2d::NO_INFORMATION, costmap_2d::NO_INFORMATION, true));
  testValues.emplace_back(TestValue(Position(5.4, 6.8), 1, 1.0, true));
  testValues.emplace_back(TestValue(Position(8.5, 4.5), 97, 97.0, true));
  // Check for grid map to costmap only.
  testValues.emplace_back(TestValue(Position(4.7, 5.8), costmap_2d::FREE_SPACE, -30.f, false));
  testValues.emplace_back(TestValue(Position(3.4, 3.8), costmap_2d::LETHAL_OBSTACLE, 270.f, false));
  return testValues;
}

// Test data for century translation
template <>
std::vector<TestCostmap2DConversion<Costmap2DCenturyTranslationTable>::TestValue>
TestCostmap2DConversion<Costmap2DCenturyTranslationTable>::getTestValues() {
  std::vector<TestValue> testValues;
  testValues.emplace_back(TestValue(Position(3.2, 5.1), costmap_2d::FREE_SPACE, 0.f, true));
  testValues.emplace_back(TestValue(Position(4.2, 4.1), costmap_2d::INSCRIBED_INFLATED_OBSTACLE, 99.f, true));
  testValues.emplace_back(TestValue(Position(6.2, 3.1), costmap_2d::LETHAL_OBSTACLE, 100.f, true));
  testValues.emplace_back(TestValue(Position(5.2, 7.8), costmap_2d::NO_INFORMATION, -1.f, true));
  testValues.emplace_back(TestValue(Position(5.4, 6.8), 253/11, 99.f/11.f, true));
  // Check for grid map to costmap only.
  testValues.emplace_back(TestValue(Position(4.7, 5.8), costmap_2d::FREE_SPACE, -30.f, false));
  testValues.emplace_back(TestValue(Position(3.4, 3.8), costmap_2d::LETHAL_OBSTACLE, 270.f, false));
  return testValues;
}

using TranslationTableTestTypes = ::testing::Types<Costmap2DSpecialTranslationTable, Costmap2DLargeIntervalsTranslationTable,
                                                   Costmap2DDirectTranslationTable, Costmap2DCenturyTranslationTable>;

TYPED_TEST_CASE(TestCostmap2DConversion, TranslationTableTestTypes);

TYPED_TEST(TestCostmap2DConversion, initializeFromCostmap2d) {
  // Convert to grid map.
  GridMap gridMap;
  this->costmap2dConverter_.initializeFromCostmap2D(this->costmap2d_, gridMap);
  this->assertMapGeometry(gridMap, this->costmap2d_);
}

TYPED_TEST(TestCostmap2DConversion, initializeFromGridMap) {
  // Convert to Costmap2D.
  costmap_2d::Costmap2D costmap2d;
  this->costmap2dConverter_.initializeFromGridMap(this->gridMap_, costmap2d);
  this->assertMapGeometry(this->gridMap_, costmap2d);
}

TYPED_TEST(TestCostmap2DConversion, addLayerFromCostmap2d) {
  // Create grid map.
  const std::string layer("layer");
  GridMap gridMap;
  this->costmap2dConverter_.initializeFromCostmap2D(this->costmap2d_, gridMap);

  // Fill in test data to Costmap2D.
  for (const auto& testValue : this->getTestValues()) {
    if (std::get<3>(testValue)) {
      unsigned int xIndex, yIndex;
      ASSERT_TRUE(this->costmap2d_.worldToMap(std::get<0>(testValue).x(), std::get<0>(testValue).y(), xIndex, yIndex));
      this->costmap2d_.getCharMap()[this->costmap2d_.getIndex(xIndex, yIndex)] = std::get<1>(testValue);
    }
  }

  // Copy data.
  this->costmap2dConverter_.addLayerFromCostmap2D(this->costmap2d_, layer, gridMap);

  // Check data.
  for (const auto& testValue : this->getTestValues()) {
    if (std::get<3>(testValue)) {
      EXPECT_EQ(std::get<2>(testValue), gridMap.atPosition(layer, std::get<0>(testValue)));
    }
  }
}

TYPED_TEST(TestCostmap2DConversion, setCostmap2DFromGridMap) {
  // Create costmap2d.
  costmap_2d::Costmap2D costmap;
  this->costmap2dConverter_.initializeFromGridMap(this->gridMap_, costmap);

  // Fill in test data to grid map.
  const std::string layer("layer");
  this->gridMap_.add(layer);
  for (const auto& testValue : this->getTestValues()) {
    Index index;
    this->gridMap_.getIndex(std::get<0>(testValue), index);
    this->gridMap_.get(layer)(index(0), index(1)) = std::get<2>(testValue);
  }

  // Copy data.
  this->costmap2dConverter_.setCostmap2DFromGridMap(this->gridMap_, layer, costmap);

  // Check data.
  for (const auto& testValue : this->getTestValues()) {
    unsigned int xIndex, yIndex;
    ASSERT_TRUE(costmap.worldToMap(std::get<0>(testValue).x(), std::get<0>(testValue).y(), xIndex, yIndex));
    costmap.getCharMap()[costmap.getIndex(xIndex, yIndex)] = std::get<1>(testValue);
  }
}


#include <boost/numpy.hpp>
#include <cmath>
#include "mjc/core/slcp.h"
#include "macros.h"
#include <iostream>
#include "mjc/core/mj_render.h"
#include <boost/python/slice.hpp>

namespace bp = boost::python;
namespace bn = boost::numpy;

template<typename T>
bn::ndarray toNdarray1(const T* data, long dim0) {
  long dims[1] = {dim0};
  bn::ndarray out = bn::empty(1, dims, bn::dtype::get_builtin<T>());
  memcpy(out.get_data(), data, dim0*sizeof(T));
  return out;
}
template<typename T>
bn::ndarray toNdarray2(const T* data, long dim0, long dim1) {
  long dims[2] = {dim0,dim1};
  bn::ndarray out = bn::empty(2, dims, bn::dtype::get_builtin<T>());
  memcpy(out.get_data(), data, dim0*dim1*sizeof(T));
  return out;
}
template<typename T>
bn::ndarray toNdarray3(const T* data, long dim0, long dim1, long dim2) {
  long dims[3] = {dim0,dim1,dim2};
  bn::ndarray out = bn::empty(3, dims, bn::dtype::get_builtin<T>());
  memcpy(out.get_data(), data, dim0*dim1*dim2*sizeof(T));
  return out;
}

enum ContactType {
  ContactType_SPRING=0,
  ContactType_SLCP=1,
  ContactType_SPRINGTAU=2,
  ContactType_CONVEX=3
};

#define  NV (m_model->ndof)
#define  NQ (m_model->nqpos)
#define  NBODY (m_model->nbody)
#define  NL (m_model->nlmax)
#define  NC (m_model->ncmax)

const int Integrator_FWD=1;

class PyMJCWorld {

public:

    PyMJCWorld(const std::string& binfile);
    bn::ndarray Step(const bn::ndarray& x, const bn::ndarray& u);
    bn::ndarray StepMulti(const bn::ndarray& x, const bn::ndarray& u);
    bp::tuple StepMulti2(const bn::ndarray& x, const bn::ndarray& u, const bn::ndarray& done);
    bp::tuple StepJacobian(const bn::ndarray& x, const bn::ndarray& u);
    void Plot(const bn::ndarray& x);
    void SetActuatedDims(const bp::object& dims);
    void SetContactType(ContactType contactType) {
      m_contactType = contactType;
    }
    void SetTimestep(mjtNum dt) {
      m_model->timestep = dt;
    }
    bn::ndarray ComputeContacts(const bn::ndarray& x);
    bp::dict GetModel();
    void SetModel(bp::dict);
    bp::dict GetData(const bn::ndarray& x);

    bn::ndarray GetImage(const bn::ndarray& x);

    ~PyMJCWorld();
private:
    // PyMJCWorld(const PyMJCWorld&) {}

    void _PlotInit();
    mjtByte _StepLowLevel(const mjtByte flgDiff, const mjtNum* tu,
          mjtNum* contact_scratch, mjtNum *f_ctr,
          mjtNum *plam, mjtNum *pdist, mjtNum *pdebug, mjtNum *s,
          mjtNum *jac, mjtNum *f_j);

    void _PreProcess();
    mjtNum _ContactImpulse(mjtNum* contact_scratch, mjtNum *f_ctr, const mjtNum *qvel_next_user);

    void _SetState(const mjtNum* xdata) {mju_copy(m_data->qpos, (xdata), NQ); mju_copy(m_data->qvel, (xdata)+NQ, NV); }
    void _SetControl(const mjtNum* udata) {for (int i=0; i < m_actuatedDims.size(); ++i) m_u[m_actuatedDims[i]] = (udata)[i];}
    void _ComputeCOM(mjtNum* out);

    mjModel* m_model;
    mjData* m_data;
    mjRender* m_mjr;
    std::vector<mjtNum> m_u;
    std::vector<int> m_actuatedDims;
    ContactType m_contactType;
    int m_integrator;

    mjtByte m_flagPrehoc;

    mjtNum m_spring_kspring, m_spring_kdamp;
    mjtByte m_spring_scale;

    mjtNum m_st_tau, m_st_beta; // recovery time, softness

    mjtNum m_slcp_tau, m_slcp_beta, m_slcp_tau2;   // ? ? ?
    int m_slcp_iterCloud;
    int m_slcp_maxIter;
    mjtNum* m_slcp_betaVec; // XXX i'm not sure what this is.

    mjtByte m_flagPlanar;
    double *pDBG, *pXIN, *pUIN, *pXOUT, *pLAM, *pDIST, *pS, *pJAC, *pF_J;
    double *pPUSH_BODY, *pPUSH_POINT, *pPUSH_FORCE;
    int NS, NDS;
    int m_SZ_F, m_SZ_DIST, m_SZ_DISTX;
    // int FLG_CONTACT, FLG_FWD;
    double* FEAT;
    mjtNum* cscratch[1];

};

PyMJCWorld::PyMJCWorld(const std::string& binfile) :
  m_model(mj_loadModel(binfile.c_str())),
  m_data(mj_makeData(m_model)),
  m_mjr(NULL),
  m_u(NV, 0),
  m_actuatedDims(),

  m_contactType(ContactType_SPRING),
  m_integrator(Integrator_FWD),
  m_flagPrehoc(true),

  m_spring_kspring(0),
  m_spring_kdamp(0),
  m_spring_scale(false),

  m_st_tau(0.04), // depth of penetration
  m_st_beta(0.1),

  m_slcp_tau(0.04), // depth of penetration
  m_slcp_beta(0.001),
  m_slcp_tau2(0.04),
  m_slcp_iterCloud(2), // ??
  m_slcp_maxIter(10),
  m_slcp_betaVec(new mjtNum[(m_flagPlanar?2:4)*NC+NL]),

  m_flagPlanar(((m_model->jnt_type[0]==mjSLIDE)&&(m_model->jnt_type[1]==mjSLIDE)&&(m_model->jnt_type[2]==mjHINGE))),

  pDBG(NULL),
  pXIN(NULL),
  pUIN(NULL),
  pXOUT(NULL),
  pLAM(NULL),
  pDIST(NULL),
  pS(NULL),
  pJAC(NULL),
  pF_J(NULL),
  pPUSH_BODY(NULL),
  pPUSH_POINT(NULL),
  pPUSH_FORCE(NULL),

  NS(0),
  NDS(0),

  m_SZ_F(mjConDim*NC+NL),
  m_SZ_DIST(NC+NL),
  m_SZ_DISTX(4*NC+NL),

  FEAT(NULL)
{

    _PreProcess();

    printf("Damping[0]: %f\n", m_model->jnt_vel_damping[0]);
    printf("Gravity: %f %f %f\n", m_model->gravity[0], m_model->gravity[1], m_model->gravity[2]);
    printf("Viscosity: %f\n", m_model->viscosity);
    printf("Timestep: %f\n", m_model->timestep);

}

PyMJCWorld::~PyMJCWorld() {
  if (m_mjr) {
    m_mjr->Close();
    delete m_mjr;
  }
  mj_delete(m_model, m_data);
  delete[] m_slcp_betaVec;
}

#define MJTNUM_DTYPE bn::dtype::get_builtin<mjtNum>()
#define CHECK_MJTNUM_ARRAY(x,ndim) FAIL_IF_FALSE((x).get_dtype() == MJTNUM_DTYPE && (x).get_nd() == (ndim) && (x).get_flags() & bn::ndarray::C_CONTIGUOUS)
// #define ENSURE_MJTNUM_ARRAY(x) FAIL_IF_FALSE( x.get_dtype() ==  )

#define DO_STEP() do {_StepLowLevel(flgDiff, m_u.data(), cscratch[0], f_ctr, pLAM, pDIST, pDBG, s, jac, f_j);} while(0)

#define COPY_STATE(dest) do {mju_copy((dest), m_data->qpos, NQ); mju_copy((dest)+NQ, m_data->qvel, NV); } while(0)
#define INC_BY_STATE(dest) do {mju_add((dest), (dest), m_data->qpos, NQ); mju_add((dest)+NQ, (dest)+NQ, m_data->qvel, NV); } while(0)
#define DEC_BY_STATE(dest) do {mju_sub((dest), (dest), m_data->qpos, NQ); mju_sub((dest)+NQ, (dest)+NQ, m_data->qvel, NV); } while(0)
#define SCALE_INPLACE(p, scl, n) mju_scl((p), (p), (scl), (n))

#define PYPRINT(obj) do{std::string s1 = bp::extract<std::string>((obj).attr("__repr__")()); std::cout << s1 << std::endl;} while(0)

bn::ndarray PyMJCWorld::Step(const bn::ndarray& x, const bn::ndarray& u) {
  CHECK_MJTNUM_ARRAY(x, 1);
  CHECK_MJTNUM_ARRAY(u, 1);
  FAIL_IF_FALSE(u.shape(0) == m_actuatedDims.size());

  mjtNum* xdata = reinterpret_cast<mjtNum*>(x.get_data());
  mjtNum* udata = reinterpret_cast<mjtNum*>(u.get_data());

  _SetState(xdata);
  _SetControl(udata);

  mjtByte flgDiff=1;
  mjtNum *s=0, *jac=0, *f_j=0, *f_ctr=0;
  DO_STEP();

  bn::ndarray out = bn::empty(bp::make_tuple(NQ+NV), bn::dtype::get_builtin<mjtNum>());
  mjtNum* outdata = reinterpret_cast<mjtNum*>(out.get_data());
  COPY_STATE(outdata);
  return out;
}

bn::ndarray PyMJCWorld::StepMulti(const bn::ndarray& x, const bn::ndarray& u) {
  CHECK_MJTNUM_ARRAY(x, 2);
  CHECK_MJTNUM_ARRAY(u, 2);

  int n_steps = x.shape(0);
  FAIL_IF_FALSE (u.shape(0) == x.shape(0));

  mjtNum* xdata = reinterpret_cast<mjtNum*>(x.get_data());
  mjtNum* udata = reinterpret_cast<mjtNum*>(u.get_data());

  const int NX = NQ + NV, NU = m_actuatedDims.size();

  bn::ndarray out = bn::empty(bp::make_tuple(n_steps,NX), bn::dtype::get_builtin<mjtNum>());
  mjtNum* outdata = reinterpret_cast<mjtNum*>(out.get_data());

  for (int i=0; i < n_steps; ++i) {
    _SetState(xdata + NX*i);
    _SetControl(udata + NU*i);

    mjtByte flgDiff=1;
    mjtNum *s=0, *jac=0, *f_j=0, *f_ctr=0;
    DO_STEP();
    COPY_STATE(outdata + NX*i);
  }

  return out;
}

bp::tuple PyMJCWorld::StepMulti2(const bn::ndarray& x, const bn::ndarray& u, const bn::ndarray& done) {
  CHECK_MJTNUM_ARRAY(x, 2);
  CHECK_MJTNUM_ARRAY(u, 2);

  const int xcols = NQ+NV,
               ucols = m_actuatedDims.size(),
               ycols = NQ+NV,
               fcols=m_model->ndof,
               dcomcols = 3,
               distcols = m_model->nbody,
               kincols = m_model->ndof*7 + m_model->nbody*28;
  int n_steps = x.shape(0);

  // kincols is cdof thru qfrc_bias inclusive

  FAIL_IF_FALSE (u.shape(0) == x.shape(0));
  FAIL_IF_FALSE(u.shape(1) == ucols);
  FAIL_IF_FALSE(x.shape(1) == xcols);

  mjtNum* xdata = reinterpret_cast<mjtNum*>(x.get_data());
  mjtNum* udata = reinterpret_cast<mjtNum*>(u.get_data());

  uint8_t* donedata = reinterpret_cast<uint8_t*>(done.get_data());

  bn::ndarray yarray = bn::zeros(bp::make_tuple(n_steps,ycols), bn::dtype::get_builtin<mjtNum>());
  bn::ndarray farray = bn::zeros(bp::make_tuple(n_steps,fcols), bn::dtype::get_builtin<mjtNum>());
  bn::ndarray dcomarray = bn::zeros(bp::make_tuple(n_steps,dcomcols), bn::dtype::get_builtin<mjtNum>());
  bn::ndarray distarray = bn::zeros(bp::make_tuple(n_steps,distcols), bn::dtype::get_builtin<mjtNum>());
  bn::ndarray kinarray = bn::zeros(bp::make_tuple(n_steps,kincols), bn::dtype::get_builtin<mjtNum>());

  mjtNum* ydata = reinterpret_cast<mjtNum*>(yarray.get_data());
  mjtNum* fdata = reinterpret_cast<mjtNum*>(farray.get_data());
  mjtNum* dcomdata = reinterpret_cast<mjtNum*>(dcomarray.get_data());
  mjtNum* distdata = reinterpret_cast<mjtNum*>(distarray.get_data());
  mjtNum* kindata =  reinterpret_cast<mjtNum*>(kinarray.get_data());

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

    if (donedata[i]) continue;

    _SetState(xdata + xcols*i);
    _SetControl(udata + ucols*i);

    mjtByte flgDiff=1;
    mjtNum *s=0, *jac=0, *f_j=0, *f_ctr=0;

    mjtNum *plam = fdata + fcols*i; // note mju_copy(plam, d->lc_f, m->ndof);
    mjtNum *pdist = distdata + distcols*i;

    mjtNum comBefore[3];
    mj_kinematics(m_model, m_data);
    _ComputeCOM(comBefore); // XXX extra calculation

    _StepLowLevel(flgDiff, m_u.data(), cscratch[0], f_ctr, plam, pdist, pDBG, s, jac, f_j);

    COPY_STATE(ydata + ycols*i); // y
    mju_copy(kindata + kincols*i, m_data->cdof, kincols);
    // f is handled by steplowlevel
    mjtNum comAfter[3]; // dcom
    mj_kinematics(m_model, m_data);
    _ComputeCOM(comAfter); // XXX extra calculation
    mju_sub(dcomdata + dcomcols*i, comAfter, comBefore, 3);
    // dist is handled by steplowlevel

  }

  return bp::make_tuple(yarray, farray, dcomarray, distarray, kinarray);
}

void PyMJCWorld::_ComputeCOM(mjtNum* com) {
	 // XXX is idx 0 always the ground?
  mjtNum tot=0;
  com[0] = com[1] = com[2] = tot = 0;
  for(int i=1; i<m_model->nbody; i++ )
  {
    if( m_model->geom_type[i]==mjPLANE ) continue;
    com[0] += m_data->xpos[3*i+0]*m_model->body_mass[i];
    com[1] += m_data->xpos[3*i+1]*m_model->body_mass[i];
    com[2] += m_data->xpos[3*i+2]*m_model->body_mass[i];
    tot += m_model->body_mass[i];
  }
  com[0] /= tot;
  com[1] /= tot;
  com[2] /= tot;
}

bp::tuple PyMJCWorld::StepJacobian(const bn::ndarray& x, const bn::ndarray& u) {
    CHECK_MJTNUM_ARRAY(x, 1);
    CHECK_MJTNUM_ARRAY(u, 1);
    FAIL_IF_FALSE(u.shape(0) == m_actuatedDims.size());

    mjtNum* xdata = reinterpret_cast<mjtNum*>(x.get_data());
    mjtNum* udata = reinterpret_cast<mjtNum*>(u.get_data());

    for (int i=0; i < m_actuatedDims.size(); ++i) {
      m_u[m_actuatedDims[i]] = udata[i];
    }

    mjtByte flgDiff=1;
    mjtNum *s=0, *jac=0, *f_j=0, *f_ctr=0;

    const int NX = NQ+NV, NU = m_actuatedDims.size();
    bn::ndarray y =    bn::zeros(bp::make_tuple(NX), bn::dtype::get_builtin<mjtNum>());
    bn::ndarray dydx = bn::zeros(bp::make_tuple(NX, NX), bn::dtype::get_builtin<mjtNum>());
    bn::ndarray dydu = bn::zeros(bp::make_tuple(NU, NX), bn::dtype::get_builtin<mjtNum>());

    mjtNum* py = reinterpret_cast<mjtNum*>(y.get_data());
    mjtNum* pdydx = reinterpret_cast<mjtNum*>(dydx.get_data());
    mjtNum* pdydu = reinterpret_cast<mjtNum*>(dydu.get_data());

    mjtNum eps=1e-6;

    _SetState(xdata);
    _SetControl(udata);
    DO_STEP();
    COPY_STATE(py);
    for (int i=0; i < NX; ++i) { // XXX LOL using the fact that qpos and qvel are stored next to each other continguously

      _SetState(xdata);
      m_data->qpos[i] += eps;
      DO_STEP();
      INC_BY_STATE(pdydx+i*NX);

      _SetState(xdata);
      m_data->qpos[i] -= eps;
      DO_STEP();
      DEC_BY_STATE(pdydx+i*NX);

      SCALE_INPLACE(pdydx + i*NX, 1./(2.*eps), NX);

    }

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

      _SetState(xdata);
      m_u[m_actuatedDims[i]] = udata[i] + eps;
      DO_STEP();
      INC_BY_STATE(pdydu + i*NX);

      _SetState(xdata);
      m_u[m_actuatedDims[i]] = udata[i] - eps;
      DO_STEP();
      DEC_BY_STATE(pdydu + i*NX);

      m_u[m_actuatedDims[i]] = udata[i];

      SCALE_INPLACE(pdydu + i*NX, 1./(2.*eps), NX);

    }

  return bp::make_tuple(y, dydx.transpose(), dydu.transpose());
}

void PyMJCWorld::Plot(const bn::ndarray& x) {
    CHECK_MJTNUM_ARRAY(x, 1);
    if (!m_mjr) _PlotInit();

    mjtNum* xdata = reinterpret_cast<mjtNum*>(x.get_data());
    // mjtNum* udata = reinterpret_cast<mjtNum*>(u.get_data());

    _SetState(xdata);
    mj_kinematics(m_model, m_data);
    mj_global(m_model, m_data);

    m_mjr->Update(m_model, m_data);

}

bn::ndarray PyMJCWorld::GetImage(const bn::ndarray& x) {
  long dims[3]={480,640,4};
  bn::ndarray out = bn::empty(3, dims, bn::dtype::get_builtin<uint8_t>());
  Plot(x);
  m_mjr->SaveImage(out.get_data());
  using bp::_;
  using bp::slice;
  // using bp::slice;
  // return out;
  out = bp::extract<bn::ndarray>(out[bp::make_tuple(slice(_,_,-1),slice(),bp::make_tuple(2,1,0))].attr("copy")());
  return out;
}

void PyMJCWorld::SetActuatedDims(const bp::object& dims) {
  m_u.assign(NV, 0);
  int ndims = bp::len(dims);
  m_actuatedDims.clear();
  for (int i=0; i < ndims; ++i) {
    m_actuatedDims.push_back(bp::extract<int>(dims[i]));
  }
}

void PyMJCWorld::_PlotInit() {
  m_mjr = new mjRender();
  int x        = 0;
  int y        = 0;
  int width    = 640;
  int height   = 480;
  m_mjr->Open(m_model->name, x, y, width, height);
  m_mjr->Set(m_model, m_data);
}

void PyMJCWorld::_PreProcess() {

  // Original code:
  mj_solve1(m_model, m_data);
  mj_contactPrepare(m_model, m_data, 0);
  mj_solve2(m_model, m_data);
  m_data->nlim=m_model->nlmax;
  m_data->ncon=m_model->ncmax;

  // for (int i=0; i < sz_pyramid; ++i) m_slcp_betaVec[i] = 0; // siminit.m line 72
}

mjtNum PyMJCWorld::_ContactImpulse(mjtNum* contact_scratch, mjtNum *f_ctr, const mjtNum * /*qvel_next_user*/) {
  mjModel* m=m_model;
  mjData* d = m_data;

  mjtNum dbg, mu;
  int j, stepDbg;

  #if 0
  mjtNum *alpha, *kappaN, *kappaT, *kappaV, *contact_buffer, *P
  int sz_dist = nl+nc;
  #endif

  mjtNum *J, *A, *f, *v, *scratch;
  mjtByte inCloud;
  int nl    = d->nlim;
  int nc    = d->ncon;
  int sz = nl+mjConDim*nc;
  int sz_pyramid = (m_flagPlanar?2:4)*nc+nl;

   //no need to zero - contactPrepare did that already
   //mju_zero(d->lc_f, SZ_F);

   switch( m_contactType ) {
      case ContactType_SPRING:
         mju_mulMatVec(d->lc_v0, d->lc_J, d->qvel_next, sz, NV);
         mj_impulseSpring(m, d, m_spring_kspring, m_spring_kdamp, 1, 1, m_spring_scale);
         return 0;
      case ContactType_SPRINGTAU:
      case ContactType_SLCP:
         mju_mulMatVec(d->lc_v0, d->lc_J, d->qvel_next, sz, NV);
     if (m_contactType == ContactType_SPRINGTAU)
      mj_impulseTau(m, d, m_st_tau, m_st_beta);
     else
      mj_impulseTau2(m, d, m_slcp_tau, m_slcp_tau2, m_slcp_beta, m_slcp_maxIter);
     return 0; // This is a little hack - in this formulation, SLCP no longer exists, there is only springtau
         //if( FLG_CONTACT==CF_SPRINGTAU && !ITER_CLOUD ) return 0; //otherwise, continue with slcp
         f       = contact_scratch;
         v       = contact_scratch + sz_pyramid;
         J       = contact_scratch + sz_pyramid*2;
         A       = contact_scratch + sz_pyramid*2 + sz_pyramid*NV;
         scratch = contact_scratch + sz_pyramid*2 + sz_pyramid*NV + sz_pyramid*sz_pyramid;
         mju_mulMatVec(v, J, d->qvel_next, sz_pyramid, NV);
         sz_pyramid = convexify(m, d, m_flagPlanar, f, v, J, A);

         if( m_contactType==ContactType_SLCP ) { //if we came here from springTau, lc_f is already initialized
#ifdef _WIN32
            inCloud= (f_ctr && !_isnan(f_ctr[0]));
#else
            inCloud= (f_ctr && !std::isnan(f_ctr[0]));
#endif
            if (!inCloud) {
               // tol = 1e-15
               dbg=slcp_solve(sz_pyramid, A, v, f, m_slcp_betaVec, 1E-15, m_slcp_maxIter, scratch);
               if( dbg<0 )
                  return dbg;
               if (f_ctr) mju_copy(f_ctr, f, sz_pyramid);
            } else {
               mju_copy(f, f_ctr, sz_pyramid);
            }
         }
         for (j=0;j<m_slcp_iterCloud;j++) {
            mu=-1; //a small regularizing term to deal with rank-deficient A
            stepDbg = slcp_step(sz_pyramid, A, v, f, m_slcp_betaVec, scratch, scratch+sz_pyramid, &mu);
            if( stepDbg<0 ) return stepDbg-0.1234567;
         }
         deconvexify(m, d, m_flagPlanar, f); // assigns into d->lc_f

         //test: J*f should be equal to lc_J*lc_f:
         //mju_mulMatTVec(ds,J,f,LOCAL_sz_cl,NV);
         //mju_mulMatTVec(ds+SZ_F,d->lc_J,d->lc_f,SZ_F,NV);
         return dbg;

      case ContactType_CONVEX:
        NOTIMPLEMENTED;
        #if 0
         kappaN = contact_scratch;
         kappaT = contact_scratch + sz_dist;
         kappaV = contact_scratch + sz_dist*2;
         alpha  = contact_scratch + sz_dist*3;
         makeKappa(pkappal[0],pkappac[0], kappaN, d->lc_A, nl, nc);
         makeKappa(0         ,pkappac[1], kappaT, d->lc_A, nl, nc);
         makeKappa(pkappal[1],pkappac[2], kappaV, d->lc_A, nl, nc);
#ifdef _WIN32
         inCloud= (f_ctr && !_isnan(f_ctr[0]));
#else
         inCloud= (f_ctr && !std::isnan(f_ctr[0]));
#endif

         //initialize:
         mju_mulMatVec(d->lc_v0, d->lc_J, d->qvel_next, sz, NV);
         mj_impulseTau(m, d, TAU_SPRING, BETA_SPRING);

         if (FLG_FWD>0) {
            contact_buffer = contact_scratch + sz_dist*3 + sz;
            if (!inCloud) {
               dbg=mj_impulseConvex(d->lc_f, m, d, kappaN, kappaT, kappaV,
                       dmax, amin, amaxl, amaxc, marginF, marginV, 1e-15, MAX_ITER, contact_buffer);
               if( dbg<0 )
                  return dbg;
               if (f_ctr) mju_copy(f_ctr, d->lc_f, sz);
            } else {
               mju_copy(d->lc_f, f_ctr, sz);
            }
            makeAlpha(alpha, d->lc_dist, amin, amaxl, amaxc, dmax, d->lc_A, nl, nc);
            for (j=0;j<ITER_CLOUD;j++) {
               contact_buffer[sz]=0;
               stepDbg = convex_fwdIter(d->lc_f, m, nl, nc,
                       alpha, kappaN, kappaT, kappaV,
                       d->con_friction, d->lc_A, d->lc_v0, d->lc_vmin,
                       marginF, marginV, contact_buffer, contact_buffer+sz);
               if( stepDbg ) return stepDbg+0.7654321;
            }
         } else { //inv dyn
            P              = contact_scratch + sz_dist*3 + sz;
            contact_buffer = contact_scratch + sz_dist*3 + sz*2;
            mju_copy(d->qvel_next, qvel_next_user, NV);
            mju_mulMatVec(d->lc_v0, d->lc_J, d->qvel_next, sz, NV);
            if (!inCloud) {
               dbg=mj_impulseConvexInv(d->lc_f, m, d, kappaN, kappaT, kappaV,
                       dmax, amin, amaxl, amaxc, marginF, marginV, 1e-15, MAX_ITER, contact_buffer);
               if( dbg<0 )
                  return dbg;
               if (f_ctr) mju_copy(f_ctr, d->lc_f, sz);
            } else {
               mju_copy(d->lc_f, f_ctr, sz);
            }
            mju_zero(alpha, sz);
            makeP(P, d->lc_A, alpha, d->lc_v0, d->lc_vmin, nl, nc, kappaV, marginV);
            makeAlpha(alpha, d->lc_dist, amin, amaxl, amaxc, dmax, d->lc_A, nl, nc);
            for (j=0;j<ITER_CLOUD;j++) {
               contact_buffer[sz]=0;
               stepDbg=convex_invIter(d->lc_f, m, nl, nc,
                       alpha, kappaN, kappaT,
                       d->con_friction, d->lc_A, P,
                       marginF, contact_buffer, contact_buffer+sz);
               if( stepDbg ) return stepDbg+0.7654321;
            }
         }
         return dbg;
         #endif
   }

   PRINT_AND_THROW("UNREACHABLE");
   return 0;
}

mjtByte PyMJCWorld::_StepLowLevel(const mjtByte flgDiff, const mjtNum* tu,
        mjtNum* contact_scratch, mjtNum *f_ctr,
        mjtNum *plam, mjtNum *pdist, mjtNum *pdebug, mjtNum *s,
        mjtNum *jac, mjtNum *f_j) {

  mjModel* m=m_model;
  mjData* d=m_data;
   int i, j, k, n;
   mjtNum *ds = d->scratch;
   mjtNum dbg, total_mass;
   mjtByte needCoM;

   if (flgDiff) {
      // run kinematic computations
      mj_kinematics(m, d);
      mj_global(m, d);
      mj_crb(m, d);
      mj_factorM(m, d);
      if (m_SZ_F) {
         mj_contactPrepare(m, d, (tu?0:2));//(FLG_CONTACT==CF_SPRINGTAU && !ITER_CLOUD));
      }
   } else {
      // global does this, so if we're not running global, we should manually zero it out
      mju_zero(d->qfrc_ext, NV);
      // contactPrepare does this
      mju_zero(d->lc_f, m_SZ_F);
   }
   mj_rne(m, d, 0);
   mj_passive(m, d);

   if( s ) {
      needCoM = 0;
      i=0;
      // start with CoM:

      // CoM quat space:
      if( FEAT[0] ) i+=4;
      // CoM position:
      for( k=0; k<3; k++ ) {
         if( FEAT[1+k] != 0.0 ) {
            s[i]=d->com[k];
            i++;
            needCoM++;
         }
      }
      // CoM velocities:
      for( k=0; k<6; k++ ) {
         if( FEAT[4+k] != 0.0 ) {
            s[i]=d->cvel[k];
            i++;
         }
      }

      for (j=1;j<NBODY;j++) {
         // quaternion
         if( FEAT[10*j] != 0.0 ) {
            for( k=0; k<4; k++ ) {
               s[i]=d->xquat[4*j+k];
               i++;
            }
         }
         // position
         for( k=0; k<3; k++ ) {
            if( FEAT[10*j+1+k] != 0.0 ) {
         if ( FEAT[10*j+1+k] == 2.0 ) // Use xanchor
                  s[i]=d->xanchor[3*m->body_jnt_adr[j]+k];
         else // Use regular position
          s[i]=d->xpos[3*j+k];
               i++;
            }
         }
         // velocities
         for( k=0; k<6; k++ ) {
            if( FEAT[10*j+4+k] != 0.0 ) {
               s[i]=d->cvel[6*j+k];
               i++;
            }
         }
      }
    //test:
      //if( i!=NS ) s[0]=sqrt((double)-1.0);
      if( jac ) { //need to compute jacs as well
         total_mass=0;
         i=needCoM; //save room for the CoM's jac, which will be created last
         if( FEAT[0] != 0.0 ) i+=3; //save useless room for the CoM's nonexistent quat's jac
         for( j=1; j<NBODY; j++ ) {
            if( FEAT[10*j] != 0.0 ) {
               mj_jac(m, d, 0, ds, 0, j, 1);
               for( k=0; k<3; k++ ) {
                  mju_copy(jac+i*NV, ds+k*NV, NV);
                  i++;
               }
            }
            if( needCoM || FEAT[10*j+1] != 0.0 || FEAT[10*j+2] != 0.0 || FEAT[10*j+3] != 0.0 ) {
               mj_jac(m, d, ds, 0, 0, j, 1);
               for( k=0; k<3; k++ ) {
                  if( FEAT[10*j+1+k] != 0.0 ) {// keep this body's jac:
                     mju_copy(jac+i*NV, ds+k*NV, NV);
                     i++;
                  }
               }
               if( needCoM ) { // building the CoM's jac:
                  n=0;
                  total_mass+=m->body_mass[j];
                  mju_scl(ds, ds, m->body_mass[j], 3*NV);  //scale by body mass
                  for( k=0; k<3; k++ ){
                     if( FEAT[1+k] != 0.0 ) {
                        mju_addTo(jac+n*NV, ds+k*NV, NV);
                        n++;
                     }
                  }
               }
            }
         }
         if( needCoM ) mju_scl(jac, jac, 1/total_mass, needCoM*NV);
      }
   }

   if( f_j ) mju_copy(f_j, d->lc_J, (d->nlim+mjConDim*d->ncon)*NV);

   if( pdist )  {
    // store the smallest contact distance for each body.
      // first, set all contact distances to the maximum value.
      for (i = 0; i < NBODY; i++)
      pdist[i] = 1000.0;
    for (j = 0; j < d->ncon; j++)
    {
      // Address of bodies.
      if (d->con_pair[j] >= m->ncpair) continue;
      int b1 = m->geom_body_id[m->pair_geom1[d->con_pair[j]]];
      int b2 = m->geom_body_id[m->pair_geom2[d->con_pair[j]]];
      double pen = d->lc_dist[d->nlim+j];
      if (pdist[b1] > pen)
        pdist[b1] = pen;
      if (pdist[b2] > pen)
        pdist[b2] = pen;
      // Make sure the ground is one of the bodies -- otherwise we need some special handling.
      //if (b1 != 0 && b2 != 0)
      //  mexErrMsgTxt("Returning X distance and detected collision not with the ground -- consider adding some code to handle this!");
    }
      //min2full_dist(d, d->lc_dist, pdist);
      //min2full_xpos(d, d->con_xpos, pdist+SZ_DIST);
   }

   if( !tu ) return 0;

   if (m_SZ_F && m_flagPrehoc==0) {
      mj_solve2(m, d);
   }
   else {
      if (m_integrator>0) {
        mju_addTo(d->qfrc_ext, tu, NV);
      }
      mj_solve2(m, d);
   }
   if (m_SZ_F) {
      dbg = _ContactImpulse(contact_scratch, f_ctr, tu); //we give tu to contact because
   }

   if (m_integrator>0) {
      if (m_flagPrehoc==0) {
         mj_backsubM(m, d, d->qacc, tu);
         mju_scl(d->qacc, d->qacc, m->timestep, NV);
         mju_addTo(d->qvel_next, d->qacc, NV);
      }
      mj_solve3(m, d);
      mj_integrate(m, d);
   } else {
      mju_mulMatTVec(d->scratch, d->lc_JP, d->lc_f, m_SZ_F, NV);
      mju_sub(d->qvel_next, tu, d->scratch, NV);
      mju_sub(d->qvel_next, d->qvel_next, d->qvel, NV);
      mju_scl(d->qacc, d->qvel_next, 1/m->timestep, NV);
      mj_fullM(m, d, d->scratch, d->qM);
      mju_mulMatVec(d->qvel_next, d->scratch, d->qacc, NV, NV);
      //first i multiplied by JP, then mult by fullM;
      //alternatively, maybe mult by J and not mult that piece by M?
      mju_sub(d->qvel_next, d->qvel_next, d->qfrc_bias, NV);
      mju_sub(d->qvel_next, d->qvel_next, d->qfrc_ext, NV);
   }

   if( plam ) { // return contact forces.
    int sz = d->nlim + mjConDim*d->ncon;
    //mju_mulMatTVec(d->scratch, d->lc_JP, d->lc_f, sz, m->ndof); // This version computes accelerations.
    mju_mulMatTVec(d->scratch, d->lc_J, d->lc_f, sz, m->ndof); // This version computes forces.
    mju_copy(plam, d->scratch, m->ndof);
    //min2full_lam(d, d->lc_f, plam);
   }
   if( pdebug ) pdebug[0]=dbg;
   return 0;
}

bn::ndarray PyMJCWorld::ComputeContacts(const bn::ndarray& x) {

  CHECK_MJTNUM_ARRAY(x, 1);
  mjtNum* xdata = reinterpret_cast<mjtNum*>(x.get_data());
  _SetState(xdata);
  mj_kinematics(m_model, m_data);
  mj_global(m_model, m_data);

  bn::ndarray dists = bn::empty(bp::make_tuple(m_model->nbody), bn::dtype::get_builtin<mjtNum>());
  mjtNum* distsdata = reinterpret_cast<mjtNum*>(dists.get_data());
  for (int i=0; i < m_model->nbody; ++i) {
    // printf("body %i parent id %i\n", i, m_model->body_parent_id[i]);
    distsdata[i] = 1000;
  }

  // Code copied from mj_contactPrepare
  mjContact con[mjMaxConPair];
  mjModel* m=m_model;

  for(int i=0; i<m->ncpair; i++ )
    if( m->pair_enable[i] )
    {
      // get geom ids
      int gid1 = m->pair_geom1[i];
      int gid2 = m->pair_geom2[i];

      int num = conFunction[m->pair_type[i]](con, m->mindist[0],
        m_data->geom_xpos+3*gid1, m_data->geom_xmat+9*gid1, m->geom_size+mjNumSize*gid1,
        m_data->geom_xpos+3*gid2, m_data->geom_xmat+9*gid2, m->geom_size+mjNumSize*gid2);

      if( num > conInfo[m->pair_type[i]][0] )
        mju_error("too many contacts in geom pair");

      // process all returned contacts
      for(int j=0; j<num; j++ )
      {
        mjtNum dist = -con[j].depth;
        int bid1 = m_model->geom_body_id[gid1];
        int bid2 = m_model->geom_body_id[gid2];
        distsdata[bid1] = fmin(dist, distsdata[bid1]);
        distsdata[bid2] = fmin(dist, distsdata[bid2]);
      }
    }
  return dists;

}

int _ndarraysize(const bn::ndarray& arr) {
  int prod = 1;
  for (int i=0; i < arr.get_nd(); ++i) {
    prod *= arr.shape(i);
  }
  return prod;
}
template<class T>
void _copyscalardata(const bp::object& from, T& to) {
  to = bp::extract<T>(from);
}
template <short ndim,typename T>
void _copyarraydata(const bn::ndarray& from, T* to) {
  FAIL_IF_FALSE(from.get_dtype() == bn::dtype::get_builtin<T>() && from.get_nd() == ndim && from.get_flags() & bn::ndarray::C_CONTIGUOUS);
  memcpy(to, from.get_data(), _ndarraysize(from)*sizeof(T));
}

template<typename T>
void _csdihk(bp::dict d, const char* key, T& to) {
  // copy scalar data if has_key
  if (d.has_key(key)) _copyscalardata(d[key], to);
}
template<short ndim, typename T>
void _cadihk(bp::dict d, const char* key, T* to) {
  // copy array data if has_key
  if (d.has_key(key)) {
    bn::ndarray arr = bp::extract<bn::ndarray>(d[key]);
    _copyarraydata<ndim,T>(arr, to);
  }
}

void PyMJCWorld::SetModel(bp::dict d) {
  _cadihk<1>(d, "dof_armature", m_model->dof_armature);
  _cadihk<2>(d, "jnt_limit", m_model->jnt_limit);
}

bp::dict PyMJCWorld::GetModel() {
  bp::dict out;
  out["nqpos"] = m_model->nqpos;
  out["ndof"] = m_model->ndof;
  out["njnt"] = m_model->njnt;
  out["nbody"] = m_model->nbody;
  out["ngeom"] = m_model->ngeom;
  out["neqmax"] = m_model->neqmax;
  out["nlmax"] = m_model->nlmax;
  out["ncpair"] = m_model->ncpair;
  out["ncmax"] = m_model->ncmax;
  out["name"] = bp::str(m_model->name);
  out["timestep"] = m_model->timestep;
  out["gravity"] = toNdarray1<mjtNum>(m_model->gravity,3);
  out["viscosity"] = m_model->viscosity;
  out["mindist"] = toNdarray1<mjtNum>(m_model->mindist,2);
  out["erreduce"] = toNdarray1<mjtNum>(m_model->errreduce, 3);

  out["body_mass"] = toNdarray1<mjtNum>(m_model->body_mass, m_model->nbody);
  out["body_inertia"] = toNdarray2<mjtNum>(m_model->body_inertia, m_model->nbody,3);
  out["body_pos"] = toNdarray2<mjtNum>(m_model->body_pos, m_model->nbody,3);
  out["body_quat"] = toNdarray2<mjtNum>(m_model->body_quat, m_model->nbody,4);
  out["body_viscoef"] = toNdarray1<mjtNum>(m_model->body_viscoef, m_model->nbody);
  out["body_parent_id"] = toNdarray1<int>(m_model->body_parent_id, m_model->nbody);
  out["body_jnt_num"] = toNdarray1<int>(m_model->body_jnt_num, m_model->nbody);
  out["body_dof_num"] = toNdarray1<int>(m_model->body_dof_num, m_model->nbody);
  out["body_root_id"] = toNdarray1<int>(m_model->body_root_id, m_model->nbody);

  out["jnt_type"] = toNdarray1<int>((int*)m_model->jnt_type, m_model->njnt);
  out["jnt_pos"] = toNdarray2<mjtNum>(m_model->jnt_pos, m_model->njnt,3);
  out["jnt_axis"] = toNdarray2<mjtNum>(m_model->jnt_axis, m_model->njnt,3);
  out["jnt_spring"] = toNdarray1<mjtNum>(m_model->jnt_spring, m_model->njnt);
  out["jnt_damping"] = toNdarray1<mjtNum>(m_model->jnt_damping, m_model->njnt);
  out["jnt_vel_damping"] = toNdarray1<mjtNum>(m_model->jnt_vel_damping, m_model->njnt);
  out["jnt_body_id"] = toNdarray1<int>(m_model->jnt_body_id, m_model->njnt);
  out["jnt_limit"] = toNdarray2<mjtNum>(m_model->jnt_limit, m_model->njnt,2);
  out["jnt_islimited"] = toNdarray1<mjtByte>(m_model->jnt_islimited, m_model->njnt);

  out["dof_armature"] = toNdarray1<mjtNum>(m_model->dof_armature, m_model->ndof);
  out["dof_body_id"] = toNdarray1<int>(m_model->dof_body_id, m_model->ndof);
  out["dof_jnt_id"] = toNdarray1<int>(m_model->dof_jnt_id, m_model->ndof);
  out["dof_parent_id"] = toNdarray1<int>(m_model->dof_parent_id, m_model->ndof);

  out["geom_type"] = toNdarray1<int>((int*)m_model->geom_type, m_model->ngeom);
  out["geom_size"] = toNdarray2<mjtNum>(m_model->geom_size, m_model->ngeom, mjNumSize);
  out["geom_pos"] = toNdarray2<mjtNum>(m_model->geom_pos, m_model->ngeom,3);
  out["geom_quat"] = toNdarray2<mjtNum>(m_model->geom_quat, m_model->ngeom,4);
  out["geom_color"] = toNdarray2<mjtNum>(m_model->geom_color, m_model->ngeom,3);
  out["geom_body_id"] = toNdarray1<int>(m_model->geom_body_id, m_model->ngeom);
  out["geom_isoffset"] = toNdarray1<mjtByte>(m_model->geom_isoffset, m_model->ngeom);

  out["eq_body1"] = toNdarray1<int>(m_model->eq_body1, m_model->neqmax);
  out["eq_body2"] = toNdarray1<int>(m_model->eq_body2, m_model->neqmax);
  out["eq_pos1"] = toNdarray2<mjtNum>(m_model->eq_pos1, m_model->neqmax, 3);
  out["eq_pos2"] = toNdarray2<mjtNum>(m_model->eq_pos2, m_model->neqmax, 3);
  out["eq_enable"] = toNdarray1<mjtByte>(m_model->eq_enable, m_model->neqmax);

  out["pair_type"] = toNdarray1<int>((int*)m_model->pair_type, m_model->ncpair);
  out["pair_friction"] = toNdarray1<mjtNum>(m_model->pair_friction, m_model->ncpair);
  out["pair_geom1"] = toNdarray1<int>(m_model->pair_geom1, m_model->ncpair);
  out["pair_geom2"] = toNdarray1<int>(m_model->pair_geom2, m_model->ncpair);
  out["pair_enable"] = toNdarray1<mjtByte>(m_model->pair_enable, m_model->ncpair);

  return out;
}

bp::dict PyMJCWorld::GetData(const bn::ndarray& x) {
  NOTIMPLEMENTED;
}

BOOST_PYTHON_MODULE(mjcpy) {
    bn::initialize();

    bp::enum_<ContactType>("ContactType")
      .value("NONE",ContactType_SPRING)
      .value("SLCP",ContactType_SLCP)
      .value("SPRING",ContactType_SPRINGTAU)
      .value("CONVEX",ContactType_CONVEX)
      ;

    bp::class_<PyMJCWorld,boost::noncopyable>("MJCWorld","docstring here", bp::init<const std::string&>())

        .def("Step",&PyMJCWorld::Step)
        .def("StepMulti",&PyMJCWorld::StepMulti)
        .def("StepMulti2",&PyMJCWorld::StepMulti2)
        .def("StepJacobian", &PyMJCWorld::StepJacobian)
        .def("Plot",&PyMJCWorld::Plot)
        .def("SetActuatedDims",&PyMJCWorld::SetActuatedDims)
        .def("ComputeContacts", &PyMJCWorld::ComputeContacts)
        .def("SetTimestep",&PyMJCWorld::SetTimestep)
        .def("SetContactType",&PyMJCWorld::SetContactType)
        .def("GetModel",&PyMJCWorld::GetModel)
        .def("SetModel",&PyMJCWorld::SetModel)
        .def("GetData",&PyMJCWorld::GetData)
        .def("GetImage",&PyMJCWorld::GetImage)
        ;
}


#ifndef STAN_MCMC_HMC_HAMILTONIANS_BASE_HAMILTONIAN_HPP
#define STAN_MCMC_HMC_HAMILTONIANS_BASE_HAMILTONIAN_HPP

#include <stan/callbacks/logger.hpp>
#include <stan/math/prim/mat/fun/Eigen.hpp>
#include <stan/model/gradient.hpp>
#include <stan/model/log_prob_propto.hpp>
#include <Eigen/Dense>
#include <iostream>
#include <limits>
#include <stdexcept>
#include <vector>

namespace stan {
namespace mcmc {

template <class Model, class Point, class BaseRNG>
class base_hamiltonian {
 public:
  explicit base_hamiltonian(const Model& model) : model_(model) {}

  ~base_hamiltonian() {}

  typedef Point PointType;

  virtual double T(Point& z) = 0;

  double V(Point& z) { return z.V; }

  virtual double tau(Point& z) = 0;

  virtual double phi(Point& z) = 0;

  double H(Point& z) { return T(z) + V(z); }

  // The time derivative of the virial, G = \sum_{d = 1}^{D} q^{d} p_{d}.
  virtual double dG_dt(Point& z, callbacks::logger& logger) = 0;

  // tau = 0.5 p_{i} p_{j} Lambda^{ij} (q)
  virtual Eigen::VectorXd dtau_dq(Point& z, callbacks::logger& logger) = 0;

  virtual Eigen::VectorXd dtau_dp(Point& z) = 0;

  // phi = 0.5 * log | Lambda (q) | + V(q)
  virtual Eigen::VectorXd dphi_dq(Point& z, callbacks::logger& logger) = 0;

  virtual void sample_p(Point& z, BaseRNG& rng) = 0;

  void init(Point& z, callbacks::logger& logger) {
    this->update_potential_gradient(z, logger);
  }

  void update_potential(Point& z, callbacks::logger& logger) {
    try {
      z.V = -stan::model::log_prob_propto<true>(model_, z.q);
    } catch (const std::exception& e) {
      this->write_error_msg_(e, logger);
      z.V = std::numeric_limits<double>::infinity();
    }
  }

  void update_potential_gradient(Point& z, callbacks::logger& logger) {
    try {
      stan::model::gradient(model_, z.q, z.V, z.g, logger);
      z.V = -z.V;
    } catch (const std::exception& e) {
      this->write_error_msg_(e, logger);
      z.V = std::numeric_limits<double>::infinity();
    }
    z.g = -z.g;
  }

  void update_metric(Point& z, callbacks::logger& logger) {}

  void update_metric_gradient(Point& z, callbacks::logger& logger) {}

  void update_gradients(Point& z, callbacks::logger& logger) {
    update_potential_gradient(z, logger);
  }

 protected:
  const Model& model_;

  void write_error_msg_(const std::exception& e, callbacks::logger& logger) {
    logger.error(
        "Informational Message: The current Metropolis proposal "
        "is about to be rejected because of the following issue:");
    logger.error(e.what());
    logger.error(
        "If this warning occurs sporadically, such as for highly "
        "constrained variable types like covariance matrices, "
        "then the sampler is fine,");
    logger.error(
        "but if this warning occurs often then your model may be "
        "either severely ill-conditioned or misspecified.");
    logger.error("");
  }
};

}  // namespace mcmc
}  // namespace stan

#endif