/*
	This file is part of solidity.

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

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

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

#include <test/yulPhaser/TestHelpers.h>

#include <tools/yulPhaser/SimulationRNG.h>

#include <boost/test/unit_test.hpp>

#include <cassert>

using namespace std;

namespace solidity::phaser::test
{

BOOST_AUTO_TEST_SUITE(Phaser)
BOOST_AUTO_TEST_SUITE(RandomTest)

BOOST_AUTO_TEST_CASE(bernoulliTrial_should_produce_samples_with_right_expected_value_and_variance)
{
	SimulationRNG::reset(1);
	constexpr size_t numSamples = 10000;
	constexpr double successProbability = 0.4;
	constexpr double relativeTolerance = 0.05;

	// For bernoulli distribution with success probability p: EX = p, VarX = p(1 - p)
	constexpr double expectedValue = successProbability;
	constexpr double variance = successProbability * (1 - successProbability);

	vector<uint32_t> samples;
	for (uint32_t i = 0; i < numSamples; ++i)
		samples.push_back(static_cast<uint32_t>(SimulationRNG::bernoulliTrial(successProbability)));

	BOOST_TEST(abs(mean(samples) - expectedValue) < expectedValue * relativeTolerance);
	BOOST_TEST(abs(meanSquaredError(samples, expectedValue) - variance) < variance * relativeTolerance);
}

BOOST_AUTO_TEST_CASE(bernoulliTrial_can_be_reset)
{
	constexpr size_t numSamples = 10;
	constexpr double successProbability = 0.4;

	SimulationRNG::reset(1);
	vector<uint32_t> samples1;
	for (uint32_t i = 0; i < numSamples; ++i)
		samples1.push_back(static_cast<uint32_t>(SimulationRNG::bernoulliTrial(successProbability)));

	vector<uint32_t> samples2;
	for (uint32_t i = 0; i < numSamples; ++i)
		samples2.push_back(static_cast<uint32_t>(SimulationRNG::bernoulliTrial(successProbability)));

	SimulationRNG::reset(1);
	vector<uint32_t> samples3;
	for (uint32_t i = 0; i < numSamples; ++i)
		samples3.push_back(static_cast<uint32_t>(SimulationRNG::bernoulliTrial(successProbability)));

	SimulationRNG::reset(2);
	vector<uint32_t> samples4;
	for (uint32_t i = 0; i < numSamples; ++i)
		samples4.push_back(static_cast<uint32_t>(SimulationRNG::bernoulliTrial(successProbability)));

	BOOST_TEST(samples1 != samples2);
	BOOST_TEST(samples1 == samples3);
	BOOST_TEST(samples1 != samples4);
	BOOST_TEST(samples2 != samples3);
	BOOST_TEST(samples2 != samples4);
	BOOST_TEST(samples3 != samples4);
}

BOOST_AUTO_TEST_CASE(uniformInt_returns_different_values_when_called_multiple_times)
{
	SimulationRNG::reset(1);
	constexpr size_t numSamples = 1000;
	constexpr uint32_t minValue = 50;
	constexpr uint32_t maxValue = 80;
	constexpr double relativeTolerance = 0.05;

	// For uniform distribution from range a..b: EX = (a + b) / 2, VarX = ((b - a + 1)^2 - 1) / 12
	constexpr double expectedValue = (minValue + maxValue) / 2.0;
	constexpr double variance = ((maxValue - minValue + 1) * (maxValue - minValue + 1) - 1) / 12.0;

	vector<uint32_t> samples;
	for (uint32_t i = 0; i < numSamples; ++i)
		samples.push_back(SimulationRNG::uniformInt(minValue, maxValue));

	BOOST_TEST(abs(mean(samples) - expectedValue) < expectedValue * relativeTolerance);
	BOOST_TEST(abs(meanSquaredError(samples, expectedValue) - variance) < variance * relativeTolerance);
}

BOOST_AUTO_TEST_CASE(uniformInt_can_be_reset)
{
	constexpr size_t numSamples = 10;
	constexpr uint32_t minValue = 50;
	constexpr uint32_t maxValue = 80;

	SimulationRNG::reset(1);
	vector<uint32_t> samples1;
	for (uint32_t i = 0; i < numSamples; ++i)
		samples1.push_back(SimulationRNG::uniformInt(minValue, maxValue));

	vector<uint32_t> samples2;
	for (uint32_t i = 0; i < numSamples; ++i)
		samples2.push_back(SimulationRNG::uniformInt(minValue, maxValue));

	SimulationRNG::reset(1);
	vector<uint32_t> samples3;
	for (uint32_t i = 0; i < numSamples; ++i)
		samples3.push_back(SimulationRNG::uniformInt(minValue, maxValue));

	SimulationRNG::reset(2);
	vector<uint32_t> samples4;
	for (uint32_t i = 0; i < numSamples; ++i)
		samples4.push_back(SimulationRNG::uniformInt(minValue, maxValue));

	BOOST_TEST(samples1 != samples2);
	BOOST_TEST(samples1 == samples3);
	BOOST_TEST(samples1 != samples4);
	BOOST_TEST(samples2 != samples3);
	BOOST_TEST(samples2 != samples4);
	BOOST_TEST(samples3 != samples4);
}

BOOST_AUTO_TEST_CASE(binomialInt_should_produce_samples_with_right_expected_value_and_variance)
{
	SimulationRNG::reset(1);
	constexpr size_t numSamples = 1000;
	constexpr uint32_t numTrials = 100;
	constexpr double successProbability = 0.2;
	constexpr double relativeTolerance = 0.05;

	// For binomial distribution with n trials and success probability p: EX = np, VarX = np(1 - p)
	constexpr double expectedValue = numTrials * successProbability;
	constexpr double variance = numTrials * successProbability * (1 - successProbability);

	vector<uint32_t> samples;
	for (uint32_t i = 0; i < numSamples; ++i)
		samples.push_back(SimulationRNG::binomialInt(numTrials, successProbability));

	BOOST_TEST(abs(mean(samples) - expectedValue) < expectedValue * relativeTolerance);
	BOOST_TEST(abs(meanSquaredError(samples, expectedValue) - variance) < variance * relativeTolerance);
}

BOOST_AUTO_TEST_CASE(binomialInt_can_be_reset)
{
	constexpr size_t numSamples = 10;
	constexpr uint32_t numTrials = 10;
	constexpr double successProbability = 0.6;

	SimulationRNG::reset(1);
	vector<uint32_t> samples1;
	for (uint32_t i = 0; i < numSamples; ++i)
		samples1.push_back(SimulationRNG::binomialInt(numTrials, successProbability));

	vector<uint32_t> samples2;
	for (uint32_t i = 0; i < numSamples; ++i)
		samples2.push_back(SimulationRNG::binomialInt(numTrials, successProbability));

	SimulationRNG::reset(1);
	vector<uint32_t> samples3;
	for (uint32_t i = 0; i < numSamples; ++i)
		samples3.push_back(SimulationRNG::binomialInt(numTrials, successProbability));

	SimulationRNG::reset(2);
	vector<uint32_t> samples4;
	for (uint32_t i = 0; i < numSamples; ++i)
		samples4.push_back(SimulationRNG::binomialInt(numTrials, successProbability));

	BOOST_TEST(samples1 != samples2);
	BOOST_TEST(samples1 == samples3);
	BOOST_TEST(samples1 != samples4);
	BOOST_TEST(samples2 != samples3);
	BOOST_TEST(samples2 != samples4);
	BOOST_TEST(samples3 != samples4);
}

BOOST_AUTO_TEST_SUITE_END()
BOOST_AUTO_TEST_SUITE_END()

}


/**
 * @file    testGaussianISAM2.cpp
 * @brief   Unit tests for GaussianISAM2
 * @author  Michael Kaess
 */

#include <gtsam/base/debug.h>
#include <gtsam/base/TestableAssertions.h>
#include <gtsam/base/LieVector.h>
#include <gtsam/geometry/Point2.h>
#include <gtsam/geometry/Pose2.h>
#include <gtsam/inference/SymbolicFactorGraph.h>
#include <gtsam/linear/GaussianBayesNet.h>
#include <gtsam/linear/GaussianSequentialSolver.h>
#include <gtsam/linear/GaussianBayesTree.h>
#include <gtsam/linear/GaussianFactorGraph.h>
#include <gtsam/nonlinear/Ordering.h>
#include <gtsam/nonlinear/Values.h>
#include <gtsam/nonlinear/NonlinearFactorGraph.h>
#include <gtsam/nonlinear/ISAM2.h>
#include <gtsam/nonlinear/Marginals.h>
#include <gtsam/slam/PriorFactor.h>
#include <gtsam/slam/BetweenFactor.h>
#include <gtsam/slam/BearingRangeFactor.h>
#include <tests/smallExample.h>

#include <CppUnitLite/TestHarness.h>

#include <boost/foreach.hpp>
#include <boost/assign/std/list.hpp> // for operator +=
#include <boost/assign.hpp>
using namespace boost::assign;

using namespace std;
using namespace gtsam;
using boost::shared_ptr;

const double tol = 1e-4;

//  SETDEBUG("ISAM2 update", true);
//  SETDEBUG("ISAM2 update verbose", true);
//  SETDEBUG("ISAM2 recalculate", true);

// Set up parameters
SharedDiagonal odoNoise = noiseModel::Diagonal::Sigmas(Vector_(3, 0.1, 0.1, M_PI/100.0));
SharedDiagonal brNoise = noiseModel::Diagonal::Sigmas(Vector_(2, M_PI/100.0, 0.1));

ISAM2 createSlamlikeISAM2(
    boost::optional<Values&> init_values = boost::none,
    boost::optional<NonlinearFactorGraph&> full_graph = boost::none,
    const ISAM2Params& params = ISAM2Params(ISAM2GaussNewtonParams(0.001), 0.0, 0, false, true)) {

  // These variables will be reused and accumulate factors and values
  ISAM2 isam(params);
  Values fullinit;
  NonlinearFactorGraph fullgraph;

  // i keeps track of the time step
  size_t i = 0;

  // Add a prior at time 0 and update isam
  {
    NonlinearFactorGraph newfactors;
    newfactors.add(PriorFactor<Pose2>(0, Pose2(0.0, 0.0, 0.0), odoNoise));
    fullgraph.push_back(newfactors);

    Values init;
    init.insert((0), Pose2(0.01, 0.01, 0.01));
    fullinit.insert((0), Pose2(0.01, 0.01, 0.01));

    isam.update(newfactors, init);
  }

  // Add odometry from time 0 to time 5
  for( ; i<5; ++i) {
    NonlinearFactorGraph newfactors;
    newfactors.add(BetweenFactor<Pose2>(i, i+1, Pose2(1.0, 0.0, 0.0), odoNoise));
    fullgraph.push_back(newfactors);

    Values init;
    init.insert((i+1), Pose2(double(i+1)+0.1, -0.1, 0.01));
    fullinit.insert((i+1), Pose2(double(i+1)+0.1, -0.1, 0.01));

    isam.update(newfactors, init);
  }

  // Add odometry from time 5 to 6 and landmark measurement at time 5
  {
    NonlinearFactorGraph newfactors;
    newfactors.add(BetweenFactor<Pose2>(i, i+1, Pose2(1.0, 0.0, 0.0), odoNoise));
    newfactors.add(BearingRangeFactor<Pose2,Point2>(i, 100, Rot2::fromAngle(M_PI/4.0), 5.0, brNoise));
    newfactors.add(BearingRangeFactor<Pose2,Point2>(i, 101, Rot2::fromAngle(-M_PI/4.0), 5.0, brNoise));
    fullgraph.push_back(newfactors);

    Values init;
    init.insert((i+1), Pose2(1.01, 0.01, 0.01));
    init.insert(100, Point2(5.0/sqrt(2.0), 5.0/sqrt(2.0)));
    init.insert(101, Point2(5.0/sqrt(2.0), -5.0/sqrt(2.0)));
    fullinit.insert((i+1), Pose2(1.01, 0.01, 0.01));
    fullinit.insert(100, Point2(5.0/sqrt(2.0), 5.0/sqrt(2.0)));
    fullinit.insert(101, Point2(5.0/sqrt(2.0), -5.0/sqrt(2.0)));

    isam.update(newfactors, init);
    ++ i;
  }

  // Add odometry from time 6 to time 10
  for( ; i<10; ++i) {
    NonlinearFactorGraph newfactors;
    newfactors.add(BetweenFactor<Pose2>(i, i+1, Pose2(1.0, 0.0, 0.0), odoNoise));
    fullgraph.push_back(newfactors);

    Values init;
    init.insert((i+1), Pose2(double(i+1)+0.1, -0.1, 0.01));
    fullinit.insert((i+1), Pose2(double(i+1)+0.1, -0.1, 0.01));

    isam.update(newfactors, init);
  }

  // Add odometry from time 10 to 11 and landmark measurement at time 10
  {
    NonlinearFactorGraph newfactors;
    newfactors.add(BetweenFactor<Pose2>(i, i+1, Pose2(1.0, 0.0, 0.0), odoNoise));
    newfactors.add(BearingRangeFactor<Pose2,Point2>(i, 100, Rot2::fromAngle(M_PI/4.0 + M_PI/16.0), 4.5, brNoise));
    newfactors.add(BearingRangeFactor<Pose2,Point2>(i, 101, Rot2::fromAngle(-M_PI/4.0 + M_PI/16.0), 4.5, brNoise));
    fullgraph.push_back(newfactors);

    Values init;
    init.insert((i+1), Pose2(6.9, 0.1, 0.01));
    fullinit.insert((i+1), Pose2(6.9, 0.1, 0.01));

    isam.update(newfactors, init);
    ++ i;
  }

  if (full_graph)
    *full_graph = fullgraph;

  if (init_values)
    *init_values = fullinit;

  return isam;
}

/* ************************************************************************* */
TEST_UNSAFE(ISAM2, ImplAddVariables) {

  // Create initial state
  Values theta;
  theta.insert(0, Pose2(.1, .2, .3));
  theta.insert(100, Point2(.4, .5));
  Values newTheta;
  newTheta.insert(1, Pose2(.6, .7, .8));

  VectorValues delta;
  delta.insert(0, Vector_(3, .1, .2, .3));
  delta.insert(1, Vector_(2, .4, .5));

  VectorValues deltaNewton;
  deltaNewton.insert(0, Vector_(3, .1, .2, .3));
  deltaNewton.insert(1, Vector_(2, .4, .5));

  VectorValues deltaRg;
  deltaRg.insert(0, Vector_(3, .1, .2, .3));
  deltaRg.insert(1, Vector_(2, .4, .5));

  vector<bool> replacedKeys(2, false);

  Ordering ordering; ordering += 100, 0;

  // Verify initial state
  LONGS_EQUAL(0, ordering[100]);
  LONGS_EQUAL(1, ordering[0]);
  EXPECT(assert_equal(delta[0], delta[ordering[100]]));
  EXPECT(assert_equal(delta[1], delta[ordering[0]]));

  // Create expected state
  Values thetaExpected;
  thetaExpected.insert(0, Pose2(.1, .2, .3));
  thetaExpected.insert(100, Point2(.4, .5));
  thetaExpected.insert(1, Pose2(.6, .7, .8));

  VectorValues deltaExpected;
  deltaExpected.insert(0, Vector_(3, .1, .2, .3));
  deltaExpected.insert(1, Vector_(2, .4, .5));
  deltaExpected.insert(2, Vector_(3, 0.0, 0.0, 0.0));

  VectorValues deltaNewtonExpected;
  deltaNewtonExpected.insert(0, Vector_(3, .1, .2, .3));
  deltaNewtonExpected.insert(1, Vector_(2, .4, .5));
  deltaNewtonExpected.insert(2, Vector_(3, 0.0, 0.0, 0.0));

  VectorValues deltaRgExpected;
  deltaRgExpected.insert(0, Vector_(3, .1, .2, .3));
  deltaRgExpected.insert(1, Vector_(2, .4, .5));
  deltaRgExpected.insert(2, Vector_(3, 0.0, 0.0, 0.0));

  vector<bool> replacedKeysExpected(3, false);

  Ordering orderingExpected; orderingExpected += 100, 0, 1;

  // Expand initial state
  ISAM2::Impl::AddVariables(newTheta, theta, delta, deltaNewton, deltaRg, replacedKeys, ordering);

  EXPECT(assert_equal(thetaExpected, theta));
  EXPECT(assert_equal(deltaExpected, delta));
  EXPECT(assert_equal(deltaNewtonExpected, deltaNewton));
  EXPECT(assert_equal(deltaRgExpected, deltaRg));
  EXPECT(assert_container_equality(replacedKeysExpected, replacedKeys));
  EXPECT(assert_equal(orderingExpected, ordering));
}

/* ************************************************************************* */
TEST_UNSAFE(ISAM2, ImplRemoveVariables) {

  // Create initial state
  Values theta;
  theta.insert(0, Pose2(.1, .2, .3));
  theta.insert(1, Pose2(.6, .7, .8));
  theta.insert(100, Point2(.4, .5));

  SymbolicFactorGraph sfg;
  sfg.push_back(boost::make_shared<IndexFactor>(Index(0), Index(2)));
  sfg.push_back(boost::make_shared<IndexFactor>(Index(0), Index(1)));
  VariableIndex variableIndex(sfg);

  VectorValues delta;
  delta.insert(0, Vector_(3, .1, .2, .3));
  delta.insert(1, Vector_(3, .4, .5, .6));
  delta.insert(2, Vector_(2, .7, .8));

  VectorValues deltaNewton;
  deltaNewton.insert(0, Vector_(3, .1, .2, .3));
  deltaNewton.insert(1, Vector_(3, .4, .5, .6));
  deltaNewton.insert(2, Vector_(2, .7, .8));

  VectorValues deltaRg;
  deltaRg.insert(0, Vector_(3, .1, .2, .3));
  deltaRg.insert(1, Vector_(3, .4, .5, .6));
  deltaRg.insert(2, Vector_(2, .7, .8));

  vector<bool> replacedKeys(3, false);
  replacedKeys[0] = true;
  replacedKeys[1] = false;
  replacedKeys[2] = true;

  Ordering ordering; ordering += 100, 1, 0;

  ISAM2::Nodes nodes(3);

  // Verify initial state
  LONGS_EQUAL(0, ordering[100]);
  LONGS_EQUAL(1, ordering[1]);
  LONGS_EQUAL(2, ordering[0]);

  // Create expected state
  Values thetaExpected;
  thetaExpected.insert(0, Pose2(.1, .2, .3));
  thetaExpected.insert(100, Point2(.4, .5));

  SymbolicFactorGraph sfgRemoved;
  sfgRemoved.push_back(boost::make_shared<IndexFactor>(Index(0), Index(1)));
  sfgRemoved.push_back(SymbolicFactorGraph::sharedFactor()); // Add empty factor to keep factor indices consistent
  VariableIndex variableIndexExpected(sfgRemoved);

  VectorValues deltaExpected;
  deltaExpected.insert(0, Vector_(3, .1, .2, .3));
  deltaExpected.insert(1, Vector_(2, .7, .8));

  VectorValues deltaNewtonExpected;
  deltaNewtonExpected.insert(0, Vector_(3, .1, .2, .3));
  deltaNewtonExpected.insert(1, Vector_(2, .7, .8));

  VectorValues deltaRgExpected;
  deltaRgExpected.insert(0, Vector_(3, .1, .2, .3));
  deltaRgExpected.insert(1, Vector_(2, .7, .8));

  vector<bool> replacedKeysExpected(2, true);

  Ordering orderingExpected; orderingExpected += 100, 0;

  ISAM2::Nodes nodesExpected(2);

  // Reduce initial state
  FastSet<Key> unusedKeys;
  unusedKeys.insert(1);
  vector<size_t> removedFactorsI; removedFactorsI.push_back(1);
  SymbolicFactorGraph removedFactors; removedFactors.push_back(sfg[1]);
  variableIndex.remove(removedFactorsI, removedFactors);
  GaussianFactorGraph linearFactors;
  FastSet<Key> fixedVariables;
  ISAM2::Impl::RemoveVariables(unusedKeys, ISAM2::sharedClique(), theta, variableIndex, delta, deltaNewton, deltaRg,
    replacedKeys, ordering, nodes, linearFactors, fixedVariables);

  EXPECT(assert_equal(thetaExpected, theta));
  EXPECT(assert_equal(variableIndexExpected, variableIndex));
  EXPECT(assert_equal(deltaExpected, delta));
  EXPECT(assert_equal(deltaNewtonExpected, deltaNewton));
  EXPECT(assert_equal(deltaRgExpected, deltaRg));
  EXPECT(assert_container_equality(replacedKeysExpected, replacedKeys));
  EXPECT(assert_equal(orderingExpected, ordering));
}

/* ************************************************************************* */
//TEST(ISAM2, IndicesFromFactors) {
//
//  using namespace gtsam::planarSLAM;
//  typedef GaussianISAM2<Values>::Impl Impl;
//
//  Ordering ordering; ordering += (0), (0), (1);
//  NonlinearFactorGraph graph;
//  graph.add(PriorFactor<Pose2>((0), Pose2(), noiseModel::Unit::Create(Pose2::dimension));
//  graph.addRange((0), (0), 1.0, noiseModel::Unit::Create(1));
//
//  FastSet<Index> expected;
//  expected.insert(0);
//  expected.insert(1);
//
//  FastSet<Index> actual = Impl::IndicesFromFactors(ordering, graph);
//
//  EXPECT(assert_equal(expected, actual));
//}

/* ************************************************************************* */
//TEST(ISAM2, CheckRelinearization) {
//
//  typedef GaussianISAM2<Values>::Impl Impl;
//
//  // Create values where indices 1 and 3 are above the threshold of 0.1
//  VectorValues values;
//  values.reserve(4, 10);
//  values.push_back_preallocated(Vector_(2, 0.09, 0.09));
//  values.push_back_preallocated(Vector_(3, 0.11, 0.11, 0.09));
//  values.push_back_preallocated(Vector_(3, 0.09, 0.09, 0.09));
//  values.push_back_preallocated(Vector_(2, 0.11, 0.11));
//
//  // Create a permutation
//  Permutation permutation(4);
//  permutation[0] = 2;
//  permutation[1] = 0;
//  permutation[2] = 1;
//  permutation[3] = 3;
//
//  Permuted<VectorValues> permuted(permutation, values);
//
//  // After permutation, the indices above the threshold are 2 and 2
//  FastSet<Index> expected;
//  expected.insert(2);
//  expected.insert(3);
//
//  // Indices checked by CheckRelinearization
//  FastSet<Index> actual = Impl::CheckRelinearization(permuted, 0.1);
//
//  EXPECT(assert_equal(expected, actual));
//}

/* ************************************************************************* */
TEST(ISAM2, optimize2) {

  // Create initialization
  Values theta;
  theta.insert((0), Pose2(0.01, 0.01, 0.01));

  // Create conditional
  Vector d(3); d << -0.1, -0.1, -0.31831;
  Matrix R(3,3); R <<
      10,          0.0,          0.0,
      0.0,           10,          0.0,
      0.0,          0.0,   31.8309886;
  GaussianConditional::shared_ptr conditional(new GaussianConditional(0, d, R, Vector::Ones(3)));

  // Create ordering
  Ordering ordering; ordering += (0);

  // Expected vector
  VectorValues expected(1, 3);
  conditional->solveInPlace(expected);

  // Clique
  ISAM2::sharedClique clique(
      ISAM2::Clique::Create(make_pair(conditional,GaussianFactor::shared_ptr())));
  VectorValues actual(theta.dims(ordering));
  internal::optimizeInPlace<ISAM2::Base>(clique, actual);

//  expected.print("expected: ");
//  actual.print("actual: ");
  EXPECT(assert_equal(expected, actual));
}

/* ************************************************************************* */
bool isam_check(const NonlinearFactorGraph& fullgraph, const Values& fullinit, const ISAM2& isam, Test& test, TestResult& result) {

  TestResult& result_ = result;
  const std::string name_ = test.getName();

  Values actual = isam.calculateEstimate();
  Ordering ordering = isam.getOrdering(); // *fullgraph.orderingCOLAMD(fullinit).first;
  GaussianFactorGraph linearized = *fullgraph.linearize(fullinit, ordering);
//  linearized.print("Expected linearized: ");
  GaussianBayesNet gbn = *GaussianSequentialSolver(linearized).eliminate();
//  gbn.print("Expected bayesnet: ");
  VectorValues delta = optimize(gbn);
  Values expected = fullinit.retract(delta, ordering);

  bool isamEqual = assert_equal(expected, actual);

  // The following two checks make sure that the cached gradients are maintained and used correctly

  // Check gradient at each node
  bool nodeGradientsOk = true;
  typedef ISAM2::sharedClique sharedClique;
  BOOST_FOREACH(const sharedClique& clique, isam.nodes()) {
    // Compute expected gradient
    FactorGraph<JacobianFactor> jfg;
    jfg.push_back(JacobianFactor::shared_ptr(new JacobianFactor(*clique->conditional())));
    VectorValues expectedGradient(*allocateVectorValues(isam));
    gradientAtZero(jfg, expectedGradient);
    // Compare with actual gradients
    int variablePosition = 0;
    for(GaussianConditional::const_iterator jit = clique->conditional()->begin(); jit != clique->conditional()->end(); ++jit) {
      const int dim = clique->conditional()->dim(jit);
      Vector actual = clique->gradientContribution().segment(variablePosition, dim);
      bool gradOk = assert_equal(expectedGradient[*jit], actual);
      EXPECT(gradOk);
      nodeGradientsOk = nodeGradientsOk && gradOk;
      variablePosition += dim;
    }
    bool dimOk = clique->gradientContribution().rows() == variablePosition;
    EXPECT(dimOk);
    nodeGradientsOk = nodeGradientsOk && dimOk;
  }

  // Check gradient
  VectorValues expectedGradient(*allocateVectorValues(isam));
  gradientAtZero(FactorGraph<JacobianFactor>(isam), expectedGradient);
  VectorValues expectedGradient2(gradient(FactorGraph<JacobianFactor>(isam), VectorValues::Zero(expectedGradient)));
  VectorValues actualGradient(*allocateVectorValues(isam));
  gradientAtZero(isam, actualGradient);
  bool expectedGradOk = assert_equal(expectedGradient2, expectedGradient);
  EXPECT(expectedGradOk);
  bool totalGradOk = assert_equal(expectedGradient, actualGradient);
  EXPECT(totalGradOk);

  return nodeGradientsOk && expectedGradOk && totalGradOk && isamEqual;
}

/* ************************************************************************* */
TEST(ISAM2, slamlike_solution_gaussnewton)
{
  // These variables will be reused and accumulate factors and values
  Values fullinit;
  NonlinearFactorGraph fullgraph;
  ISAM2 isam = createSlamlikeISAM2(fullinit, fullgraph, ISAM2Params(ISAM2GaussNewtonParams(0.001), 0.0, 0, false));

  // Compare solutions
  CHECK(isam_check(fullgraph, fullinit, isam, *this, result_));
}

/* ************************************************************************* */
TEST(ISAM2, slamlike_solution_dogleg)
{
  // These variables will be reused and accumulate factors and values
  Values fullinit;
  NonlinearFactorGraph fullgraph;
  ISAM2 isam = createSlamlikeISAM2(fullinit, fullgraph, ISAM2Params(ISAM2DoglegParams(1.0), 0.0, 0, false));

  // Compare solutions
  CHECK(isam_check(fullgraph, fullinit, isam, *this, result_));
}

/* ************************************************************************* */
TEST(ISAM2, slamlike_solution_gaussnewton_qr)
{
  // These variables will be reused and accumulate factors and values
  Values fullinit;
  NonlinearFactorGraph fullgraph;
  ISAM2 isam = createSlamlikeISAM2(fullinit, fullgraph, ISAM2Params(ISAM2GaussNewtonParams(0.001), 0.0, 0, false, false, ISAM2Params::QR));

  // Compare solutions
  CHECK(isam_check(fullgraph, fullinit, isam, *this, result_));
}

/* ************************************************************************* */
TEST(ISAM2, slamlike_solution_dogleg_qr)
{
  // These variables will be reused and accumulate factors and values
  Values fullinit;
  NonlinearFactorGraph fullgraph;
  ISAM2 isam = createSlamlikeISAM2(fullinit, fullgraph, ISAM2Params(ISAM2DoglegParams(1.0), 0.0, 0, false, false, ISAM2Params::QR));

  // Compare solutions
  CHECK(isam_check(fullgraph, fullinit, isam, *this, result_));
}

/* ************************************************************************* */
TEST(ISAM2, clone) {

  ISAM2 clone1;

  {
    ISAM2 isam = createSlamlikeISAM2();
    clone1 = isam;

    ISAM2 clone2(isam);

    // Modify original isam
    NonlinearFactorGraph factors;
    factors.add(BetweenFactor<Pose2>(0, 10,
        isam.calculateEstimate<Pose2>(0).between(isam.calculateEstimate<Pose2>(10)), noiseModel::Unit::Create(3)));
    isam.update(factors);

    CHECK(assert_equal(createSlamlikeISAM2(), clone2));
  }

  // This is to (perhaps unsuccessfully) try to currupt unallocated memory referenced
  // if the references in the iSAM2 copy point to the old instance which deleted at
  // the end of the {...} section above.
  ISAM2 temp = createSlamlikeISAM2();

  CHECK(assert_equal(createSlamlikeISAM2(), clone1));
  CHECK(assert_equal(clone1, temp));

  // Check clone empty
  ISAM2 isam;
  clone1 = isam;
  CHECK(assert_equal(ISAM2(), clone1));
}

/* ************************************************************************* */
TEST(ISAM2, permute_cached) {
  typedef boost::shared_ptr<ISAM2Clique> sharedISAM2Clique;

  // Construct expected permuted BayesTree (variable 2 has been changed to 1)
  BayesTree<GaussianConditional, ISAM2Clique> expected;
  expected.insert(sharedISAM2Clique(new ISAM2Clique(make_pair(
      boost::make_shared<GaussianConditional>(pair_list_of
          (3, Matrix_(1,1,1.0))
          (4, Matrix_(1,1,2.0)),
          2, Vector_(1,1.0), Vector_(1,1.0)),   // p(3,4)
      HessianFactor::shared_ptr()))));          // Cached: empty
  expected.insert(sharedISAM2Clique(new ISAM2Clique(make_pair(
      boost::make_shared<GaussianConditional>(pair_list_of
          (2, Matrix_(1,1,1.0))
          (3, Matrix_(1,1,2.0)),
          1, Vector_(1,1.0), Vector_(1,1.0)),     // p(2|3)
      boost::make_shared<HessianFactor>(3, Matrix_(1,1,1.0), Vector_(1,1.0), 0.0))))); // Cached: p(3)
  expected.insert(sharedISAM2Clique(new ISAM2Clique(make_pair(
      boost::make_shared<GaussianConditional>(pair_list_of
          (0, Matrix_(1,1,1.0))
          (2, Matrix_(1,1,2.0)),
          1, Vector_(1,1.0), Vector_(1,1.0)),     // p(0|2)
      boost::make_shared<HessianFactor>(1, Matrix_(1,1,1.0), Vector_(1,1.0), 0.0))))); // Cached: p(1)
  // Change variable 2 to 1
  expected.root()->children().front()->conditional()->keys()[0] = 1;
  expected.root()->children().front()->children().front()->conditional()->keys()[1] = 1;

  // Construct unpermuted BayesTree
  BayesTree<GaussianConditional, ISAM2Clique> actual;
  actual.insert(sharedISAM2Clique(new ISAM2Clique(make_pair(
      boost::make_shared<GaussianConditional>(pair_list_of
          (3, Matrix_(1,1,1.0))
          (4, Matrix_(1,1,2.0)),
          2, Vector_(1,1.0), Vector_(1,1.0)),   // p(3,4)
      HessianFactor::shared_ptr()))));          // Cached: empty
  actual.insert(sharedISAM2Clique(new ISAM2Clique(make_pair(
      boost::make_shared<GaussianConditional>(pair_list_of
          (2, Matrix_(1,1,1.0))
          (3, Matrix_(1,1,2.0)),
          1, Vector_(1,1.0), Vector_(1,1.0)),     // p(2|3)
      boost::make_shared<HessianFactor>(3, Matrix_(1,1,1.0), Vector_(1,1.0), 0.0))))); // Cached: p(3)
  actual.insert(sharedISAM2Clique(new ISAM2Clique(make_pair(
      boost::make_shared<GaussianConditional>(pair_list_of
          (0, Matrix_(1,1,1.0))
          (2, Matrix_(1,1,2.0)),
          1, Vector_(1,1.0), Vector_(1,1.0)),     // p(0|2)
      boost::make_shared<HessianFactor>(2, Matrix_(1,1,1.0), Vector_(1,1.0), 0.0))))); // Cached: p(2)

  // Create permutation that changes variable 2 -> 0
  Permutation permutation = Permutation::Identity(5);
  permutation[2] = 1;

  // Permute BayesTree
  actual.root()->permuteWithInverse(permutation);

  // Check
  EXPECT(assert_equal(expected, actual));
}

/* ************************************************************************* */
TEST(ISAM2, removeFactors)
{
  // This test builds a graph in the same way as the "slamlike" test above, but
  // then removes the 2nd-to-last landmark measurement

  // These variables will be reused and accumulate factors and values
  Values fullinit;
  NonlinearFactorGraph fullgraph;
  ISAM2 isam = createSlamlikeISAM2(fullinit, fullgraph, ISAM2Params(ISAM2GaussNewtonParams(0.001), 0.0, 0, false));

  // Remove the 2nd measurement on landmark 0 (Key 100)
  FastVector<size_t> toRemove;
  toRemove.push_back(12);
  isam.update(NonlinearFactorGraph(), Values(), toRemove);

  // Remove the factor from the full system
  fullgraph.remove(12);

  // Compare solutions
  CHECK(isam_check(fullgraph, fullinit, isam, *this, result_));
}

/* ************************************************************************* */
TEST_UNSAFE(ISAM2, removeVariables)
{
  // These variables will be reused and accumulate factors and values
  Values fullinit;
  NonlinearFactorGraph fullgraph;
  ISAM2 isam = createSlamlikeISAM2(fullinit, fullgraph, ISAM2Params(ISAM2GaussNewtonParams(0.001), 0.0, 0, false));

  // Remove the measurement on landmark 0 (Key 100)
  FastVector<size_t> toRemove;
  toRemove.push_back(7);
  toRemove.push_back(14);
  isam.update(NonlinearFactorGraph(), Values(), toRemove);

  // Remove the factors and variable from the full system
  fullgraph.remove(7);
  fullgraph.remove(14);
  fullinit.erase(100);

  // Compare solutions
  CHECK(isam_check(fullgraph, fullinit, isam, *this, result_));
}

/* ************************************************************************* */
TEST_UNSAFE(ISAM2, swapFactors)
{
  // This test builds a graph in the same way as the "slamlike" test above, but
  // then swaps the 2nd-to-last landmark measurement with a different one

  Values fullinit;
  NonlinearFactorGraph fullgraph;
  ISAM2 isam = createSlamlikeISAM2(fullinit, fullgraph);

  // Remove the measurement on landmark 0 and replace with a different one
  {
    size_t swap_idx = isam.getFactorsUnsafe().size()-2;
    FastVector<size_t> toRemove;
    toRemove.push_back(swap_idx);
    fullgraph.remove(swap_idx);

    NonlinearFactorGraph swapfactors;
//    swapfactors.add(BearingRange<Pose2,Point2>(10, 100, Rot2::fromAngle(M_PI/4.0 + M_PI/16.0), 4.5, brNoise); // original factor
    swapfactors.add(BearingRangeFactor<Pose2,Point2>(10, 100, Rot2::fromAngle(M_PI/4.0 + M_PI/16.0), 5.0, brNoise));
    fullgraph.push_back(swapfactors);
    isam.update(swapfactors, Values(), toRemove);
  }

  // Compare solutions
  EXPECT(assert_equal(fullgraph, NonlinearFactorGraph(isam.getFactorsUnsafe())));
  EXPECT(isam_check(fullgraph, fullinit, isam, *this, result_));

  // Check gradient at each node
  typedef ISAM2::sharedClique sharedClique;
  BOOST_FOREACH(const sharedClique& clique, isam.nodes()) {
    // Compute expected gradient
    FactorGraph<JacobianFactor> jfg;
    jfg.push_back(JacobianFactor::shared_ptr(new JacobianFactor(*clique->conditional())));
    VectorValues expectedGradient(*allocateVectorValues(isam));
    gradientAtZero(jfg, expectedGradient);
    // Compare with actual gradients
    int variablePosition = 0;
    for(GaussianConditional::const_iterator jit = clique->conditional()->begin(); jit != clique->conditional()->end(); ++jit) {
      const int dim = clique->conditional()->dim(jit);
      Vector actual = clique->gradientContribution().segment(variablePosition, dim);
      EXPECT(assert_equal(expectedGradient[*jit], actual));
      variablePosition += dim;
    }
    EXPECT_LONGS_EQUAL(clique->gradientContribution().rows(), variablePosition);
  }

  // Check gradient
  VectorValues expectedGradient(*allocateVectorValues(isam));
  gradientAtZero(FactorGraph<JacobianFactor>(isam), expectedGradient);
  VectorValues expectedGradient2(gradient(FactorGraph<JacobianFactor>(isam), VectorValues::Zero(expectedGradient)));
  VectorValues actualGradient(*allocateVectorValues(isam));
  gradientAtZero(isam, actualGradient);
  EXPECT(assert_equal(expectedGradient2, expectedGradient));
  EXPECT(assert_equal(expectedGradient, actualGradient));
}

/* ************************************************************************* */
TEST(ISAM2, constrained_ordering)
{
  // These variables will be reused and accumulate factors and values
  ISAM2 isam(ISAM2Params(ISAM2GaussNewtonParams(0.001), 0.0, 0, false));
  Values fullinit;
  NonlinearFactorGraph fullgraph;

  // We will constrain x3 and x4 to the end
  FastMap<Key, int> constrained;
  constrained.insert(make_pair((3), 1));
  constrained.insert(make_pair((4), 2));

  // i keeps track of the time step
  size_t i = 0;

  // Add a prior at time 0 and update isam
  {
    NonlinearFactorGraph newfactors;
    newfactors.add(PriorFactor<Pose2>(0, Pose2(0.0, 0.0, 0.0), odoNoise));
    fullgraph.push_back(newfactors);

    Values init;
    init.insert((0), Pose2(0.01, 0.01, 0.01));
    fullinit.insert((0), Pose2(0.01, 0.01, 0.01));

    isam.update(newfactors, init);
  }

  CHECK(isam_check(fullgraph, fullinit, isam, *this, result_));

  // Add odometry from time 0 to time 5
  for( ; i<5; ++i) {
    NonlinearFactorGraph newfactors;
    newfactors.add(BetweenFactor<Pose2>(i, i+1, Pose2(1.0, 0.0, 0.0), odoNoise));
    fullgraph.push_back(newfactors);

    Values init;
    init.insert((i+1), Pose2(double(i+1)+0.1, -0.1, 0.01));
    fullinit.insert((i+1), Pose2(double(i+1)+0.1, -0.1, 0.01));

    if(i >= 3)
      isam.update(newfactors, init, FastVector<size_t>(), constrained);
    else
      isam.update(newfactors, init);
  }

  // Add odometry from time 5 to 6 and landmark measurement at time 5
  {
    NonlinearFactorGraph newfactors;
    newfactors.add(BetweenFactor<Pose2>(i, i+1, Pose2(1.0, 0.0, 0.0), odoNoise));
    newfactors.add(BearingRangeFactor<Pose2,Point2>(i, 100, Rot2::fromAngle(M_PI/4.0), 5.0, brNoise));
    newfactors.add(BearingRangeFactor<Pose2,Point2>(i, 101, Rot2::fromAngle(-M_PI/4.0), 5.0, brNoise));
    fullgraph.push_back(newfactors);

    Values init;
    init.insert((i+1), Pose2(1.01, 0.01, 0.01));
    init.insert(100, Point2(5.0/sqrt(2.0), 5.0/sqrt(2.0)));
    init.insert(101, Point2(5.0/sqrt(2.0), -5.0/sqrt(2.0)));
    fullinit.insert((i+1), Pose2(1.01, 0.01, 0.01));
    fullinit.insert(100, Point2(5.0/sqrt(2.0), 5.0/sqrt(2.0)));
    fullinit.insert(101, Point2(5.0/sqrt(2.0), -5.0/sqrt(2.0)));

    isam.update(newfactors, init, FastVector<size_t>(), constrained);
    ++ i;
  }

  // Add odometry from time 6 to time 10
  for( ; i<10; ++i) {
    NonlinearFactorGraph newfactors;
    newfactors.add(BetweenFactor<Pose2>(i, i+1, Pose2(1.0, 0.0, 0.0), odoNoise));
    fullgraph.push_back(newfactors);

    Values init;
    init.insert((i+1), Pose2(double(i+1)+0.1, -0.1, 0.01));
    fullinit.insert((i+1), Pose2(double(i+1)+0.1, -0.1, 0.01));

    isam.update(newfactors, init, FastVector<size_t>(), constrained);
  }

  // Add odometry from time 10 to 11 and landmark measurement at time 10
  {
    NonlinearFactorGraph newfactors;
    newfactors.add(BetweenFactor<Pose2>(i, i+1, Pose2(1.0, 0.0, 0.0), odoNoise));
    newfactors.add(BearingRangeFactor<Pose2,Point2>(i, 100, Rot2::fromAngle(M_PI/4.0 + M_PI/16.0), 4.5, brNoise));
    newfactors.add(BearingRangeFactor<Pose2,Point2>(i, 101, Rot2::fromAngle(-M_PI/4.0 + M_PI/16.0), 4.5, brNoise));
    fullgraph.push_back(newfactors);

    Values init;
    init.insert((i+1), Pose2(6.9, 0.1, 0.01));
    fullinit.insert((i+1), Pose2(6.9, 0.1, 0.01));

    isam.update(newfactors, init, FastVector<size_t>(), constrained);
    ++ i;
  }

  // Compare solutions
  EXPECT(isam_check(fullgraph, fullinit, isam, *this, result_));

  // Check that x3 and x4 are last, but either can come before the other
  EXPECT(isam.getOrdering()[(3)] == 12 && isam.getOrdering()[(4)] == 13);

  // Check gradient at each node
  typedef ISAM2::sharedClique sharedClique;
  BOOST_FOREACH(const sharedClique& clique, isam.nodes()) {
    // Compute expected gradient
    FactorGraph<JacobianFactor> jfg;
    jfg.push_back(JacobianFactor::shared_ptr(new JacobianFactor(*clique->conditional())));
    VectorValues expectedGradient(*allocateVectorValues(isam));
    gradientAtZero(jfg, expectedGradient);
    // Compare with actual gradients
    int variablePosition = 0;
    for(GaussianConditional::const_iterator jit = clique->conditional()->begin(); jit != clique->conditional()->end(); ++jit) {
      const int dim = clique->conditional()->dim(jit);
      Vector actual = clique->gradientContribution().segment(variablePosition, dim);
      EXPECT(assert_equal(expectedGradient[*jit], actual));
      variablePosition += dim;
    }
    LONGS_EQUAL(clique->gradientContribution().rows(), variablePosition);
  }

  // Check gradient
  VectorValues expectedGradient(*allocateVectorValues(isam));
  gradientAtZero(FactorGraph<JacobianFactor>(isam), expectedGradient);
  VectorValues expectedGradient2(gradient(FactorGraph<JacobianFactor>(isam), VectorValues::Zero(expectedGradient)));
  VectorValues actualGradient(*allocateVectorValues(isam));
  gradientAtZero(isam, actualGradient);
  EXPECT(assert_equal(expectedGradient2, expectedGradient));
  EXPECT(assert_equal(expectedGradient, actualGradient));
}

/* ************************************************************************* */
TEST(ISAM2, slamlike_solution_partial_relinearization_check)
{
  // These variables will be reused and accumulate factors and values
  Values fullinit;
  NonlinearFactorGraph fullgraph;
  ISAM2Params params(ISAM2GaussNewtonParams(0.001), 0.0, 0, false);
  params.enablePartialRelinearizationCheck = true;
  ISAM2 isam = createSlamlikeISAM2(fullinit, fullgraph, params);

  // Compare solutions
  CHECK(isam_check(fullgraph, fullinit, isam, *this, result_));
}

namespace {
  bool checkMarginalizeLeaves(ISAM2& isam, const FastList<Key>& leafKeys) {
    Matrix expectedAugmentedHessian, expected3AugmentedHessian;
    vector<Index> toKeep;
    const Index lastVar = isam.getOrdering().size() - 1;
    for(Index i=0; i<=lastVar; ++i)
      if(find(leafKeys.begin(), leafKeys.end(), isam.getOrdering().key(i)) == leafKeys.end())
        toKeep.push_back(i);

    // Calculate expected marginal from iSAM2 tree
    GaussianFactorGraph isamAsGraph(isam);
    GaussianSequentialSolver solver(isamAsGraph);
    GaussianFactorGraph marginalgfg = *solver.jointFactorGraph(toKeep);
    expectedAugmentedHessian = marginalgfg.augmentedHessian();

    //// Calculate expected marginal from cached linear factors
    //assert(isam.params().cacheLinearizedFactors);
    //GaussianSequentialSolver solver2(isam.linearFactors_, isam.params().factorization == ISAM2Params::QR);
    //expected2AugmentedHessian = solver2.jointFactorGraph(toKeep)->augmentedHessian();

    // Calculate expected marginal from original nonlinear factors
    GaussianSequentialSolver solver3(
      *isam.getFactorsUnsafe().linearize(isam.getLinearizationPoint(), isam.getOrdering()),
      isam.params().factorization == ISAM2Params::QR);
    expected3AugmentedHessian = solver3.jointFactorGraph(toKeep)->augmentedHessian();

    // Do marginalization
    isam.marginalizeLeaves(leafKeys);

    // Check
    GaussianFactorGraph actualMarginalGraph(isam);
    Matrix actualAugmentedHessian = actualMarginalGraph.augmentedHessian();
    //Matrix actual2AugmentedHessian = linearFactors_.augmentedHessian();
    Matrix actual3AugmentedHessian = isam.getFactorsUnsafe().linearize(
      isam.getLinearizationPoint(), isam.getOrdering())->augmentedHessian();
    assert(actualAugmentedHessian.unaryExpr(std::ptr_fun(&std::isfinite<double>)).all());

    // Check full marginalization
    //cout << "treeEqual" << endl;
    bool treeEqual = assert_equal(expectedAugmentedHessian, actualAugmentedHessian, 1e-6);
    //bool linEqual = assert_equal(expected2AugmentedHessian, actualAugmentedHessian, 1e-6);
    //cout << "nonlinEqual" << endl;
    bool nonlinEqual = assert_equal(expected3AugmentedHessian, actualAugmentedHessian, 1e-6);
    //bool linCorrect = assert_equal(expected3AugmentedHessian, expected2AugmentedHessian, 1e-6);
    //actual2AugmentedHessian.bottomRightCorner(1,1) = expected3AugmentedHessian.bottomRightCorner(1,1); bool afterLinCorrect = assert_equal(expected3AugmentedHessian, actual2AugmentedHessian, 1e-6);
    //cout << "nonlinCorrect" << endl;
    actual3AugmentedHessian.bottomRightCorner(1,1) = expected3AugmentedHessian.bottomRightCorner(1,1); bool afterNonlinCorrect = assert_equal(expected3AugmentedHessian, actual3AugmentedHessian, 1e-6);

    bool ok = treeEqual && /*linEqual &&*/ nonlinEqual && /*linCorrect &&*/ /*afterLinCorrect &&*/ afterNonlinCorrect;
    return ok;
  }
}

/* ************************************************************************* */
TEST_UNSAFE(ISAM2, marginalizeLeaves1)
{
  ISAM2 isam;

  NonlinearFactorGraph factors;
  factors.add(PriorFactor<LieVector>(0, LieVector(0.0), noiseModel::Unit::Create(1)));

  factors.add(BetweenFactor<LieVector>(0, 1, LieVector(0.0), noiseModel::Unit::Create(1)));
  factors.add(BetweenFactor<LieVector>(1, 2, LieVector(0.0), noiseModel::Unit::Create(1)));
  factors.add(BetweenFactor<LieVector>(0, 2, LieVector(0.0), noiseModel::Unit::Create(1)));

  Values values;
  values.insert(0, LieVector(0.0));
  values.insert(1, LieVector(0.0));
  values.insert(2, LieVector(0.0));

  FastMap<Key,int> constrainedKeys;
  constrainedKeys.insert(make_pair(0,0));
  constrainedKeys.insert(make_pair(1,1));
  constrainedKeys.insert(make_pair(2,2));

  isam.update(factors, values, FastVector<size_t>(), constrainedKeys);

  FastList<Key> leafKeys;
  leafKeys.push_back(isam.getOrdering().key(0));
  EXPECT(checkMarginalizeLeaves(isam, leafKeys));
}

/* ************************************************************************* */
TEST_UNSAFE(ISAM2, marginalizeLeaves2)
{
  ISAM2 isam;

  NonlinearFactorGraph factors;
  factors.add(PriorFactor<LieVector>(0, LieVector(0.0), noiseModel::Unit::Create(1)));

  factors.add(BetweenFactor<LieVector>(0, 1, LieVector(0.0), noiseModel::Unit::Create(1)));
  factors.add(BetweenFactor<LieVector>(1, 2, LieVector(0.0), noiseModel::Unit::Create(1)));
  factors.add(BetweenFactor<LieVector>(0, 2, LieVector(0.0), noiseModel::Unit::Create(1)));
  factors.add(BetweenFactor<LieVector>(2, 3, LieVector(0.0), noiseModel::Unit::Create(1)));

  Values values;
  values.insert(0, LieVector(0.0));
  values.insert(1, LieVector(0.0));
  values.insert(2, LieVector(0.0));
  values.insert(3, LieVector(0.0));

  FastMap<Key,int> constrainedKeys;
  constrainedKeys.insert(make_pair(0,0));
  constrainedKeys.insert(make_pair(1,1));
  constrainedKeys.insert(make_pair(2,2));
  constrainedKeys.insert(make_pair(3,3));

  isam.update(factors, values, FastVector<size_t>(), constrainedKeys);

  FastList<Key> leafKeys;
  leafKeys.push_back(isam.getOrdering().key(0));
  EXPECT(checkMarginalizeLeaves(isam, leafKeys));
}

/* ************************************************************************* */
TEST_UNSAFE(ISAM2, marginalizeLeaves3)
{
  ISAM2 isam;

  NonlinearFactorGraph factors;
  factors.add(PriorFactor<LieVector>(0, LieVector(0.0), noiseModel::Unit::Create(1)));

  factors.add(BetweenFactor<LieVector>(0, 1, LieVector(0.0), noiseModel::Unit::Create(1)));
  factors.add(BetweenFactor<LieVector>(1, 2, LieVector(0.0), noiseModel::Unit::Create(1)));
  factors.add(BetweenFactor<LieVector>(0, 2, LieVector(0.0), noiseModel::Unit::Create(1)));

  factors.add(BetweenFactor<LieVector>(2, 3, LieVector(0.0), noiseModel::Unit::Create(1)));

  factors.add(BetweenFactor<LieVector>(3, 4, LieVector(0.0), noiseModel::Unit::Create(1)));
  factors.add(BetweenFactor<LieVector>(4, 5, LieVector(0.0), noiseModel::Unit::Create(1)));
  factors.add(BetweenFactor<LieVector>(3, 5, LieVector(0.0), noiseModel::Unit::Create(1)));

  Values values;
  values.insert(0, LieVector(0.0));
  values.insert(1, LieVector(0.0));
  values.insert(2, LieVector(0.0));
  values.insert(3, LieVector(0.0));
  values.insert(4, LieVector(0.0));
  values.insert(5, LieVector(0.0));

  FastMap<Key,int> constrainedKeys;
  constrainedKeys.insert(make_pair(0,0));
  constrainedKeys.insert(make_pair(1,1));
  constrainedKeys.insert(make_pair(2,2));
  constrainedKeys.insert(make_pair(3,3));
  constrainedKeys.insert(make_pair(4,4));
  constrainedKeys.insert(make_pair(5,5));

  isam.update(factors, values, FastVector<size_t>(), constrainedKeys);

  FastList<Key> leafKeys;
  leafKeys.push_back(isam.getOrdering().key(0));
  EXPECT(checkMarginalizeLeaves(isam, leafKeys));
}

/* ************************************************************************* */
TEST_UNSAFE(ISAM2, marginalizeLeaves4)
{
  ISAM2 isam;

  NonlinearFactorGraph factors;
  factors.add(PriorFactor<LieVector>(0, LieVector(0.0), noiseModel::Unit::Create(1)));
  factors.add(BetweenFactor<LieVector>(0, 2, LieVector(0.0), noiseModel::Unit::Create(1)));
  factors.add(BetweenFactor<LieVector>(1, 2, LieVector(0.0), noiseModel::Unit::Create(1)));

  Values values;
  values.insert(0, LieVector(0.0));
  values.insert(1, LieVector(0.0));
  values.insert(2, LieVector(0.0));

  FastMap<Key,int> constrainedKeys;
  constrainedKeys.insert(make_pair(0,0));
  constrainedKeys.insert(make_pair(1,1));
  constrainedKeys.insert(make_pair(2,2));

  isam.update(factors, values, FastVector<size_t>(), constrainedKeys);

  FastList<Key> leafKeys;
  leafKeys.push_back(isam.getOrdering().key(1));
  EXPECT(checkMarginalizeLeaves(isam, leafKeys));
}

/* ************************************************************************* */
TEST_UNSAFE(ISAM2, marginalizeLeaves5)
{
  // Create isam2
  ISAM2 isam = createSlamlikeISAM2();

  // Marginalize
  FastList<Key> marginalizeKeys;
  marginalizeKeys.push_back(isam.getOrdering().key(0));
  EXPECT(checkMarginalizeLeaves(isam, marginalizeKeys));
}

/* ************************************************************************* */
TEST(ISAM2, marginalCovariance)
{
  // Create isam2
  ISAM2 isam = createSlamlikeISAM2();

  // Check marginal
  Matrix expected = Marginals(isam.getFactorsUnsafe(), isam.getLinearizationPoint()).marginalCovariance(5);
  Matrix actual = isam.marginalCovariance(5);
  EXPECT(assert_equal(expected, actual));
}

/* ************************************************************************* */
int main() { TestResult tr; return TestRegistry::runAllTests(tr);}
/* ************************************************************************* */


#include <ros/ros.h>
#include <pcl/io/pcd_io.h>
#include <pcl/point_types.h>
#include <pcl_conversions/pcl_conversions.h>
#include <pcl/PCLPointCloud2.h>
#include <pcl/common/common.h>
#include <pcl/common/centroid.h>
#include <pcl/common/pca.h>

#include <sensor_msgs/PointCloud2.h>
#include <cv_bridge/cv_bridge.h>
#include <sensor_msgs/image_encodings.h>
#include <sensor_msgs/CameraInfo.h>
#include <sensor_msgs/Image.h>
#include <tf/transform_listener.h>
#include <tf_conversions/tf_eigen.h>
#include <message_filters/subscriber.h>
#include <message_filters/synchronizer.h>
#include <message_filters/sync_policies/approximate_time.h>

#include <image_geometry/pinhole_camera_model.h>

#include <opencv2/highgui/highgui.hpp>

#include <apc_msgs/ImageToWorld.h>
#include <apc_msgs/AlignedDimensions.h>

#include <apc_msgs/Camera_Info.h>

#include <boost/foreach.hpp>

#include <sstream>

#include <cmath>

/**
*Node for detecting 3D clusters and align them with corresponding 2D images
*/

using namespace message_filters;
using namespace std;
using namespace pcl;
using namespace cv;
using namespace ros;

typedef struct Proposal{
    cv::Rect bounding_box;
    std::vector<cv::Point2d> imagePoints;
    uint32_t label;
} proposal_t;

class ImageToWorld {

public:

    ImageToWorld(ros::NodeHandle nh):
        nh_(nh),
        camera_initialized_(false)
    {
        // Just do this once.
        std::string camera_info_topic_name = "/realsense_wrist/rgb/camera_info";
        sensor_msgs::CameraInfoConstPtr camera_info_ptr = ros::topic::waitForMessage<sensor_msgs::CameraInfo>(camera_info_topic_name);
        onCameraInfo(camera_info_ptr);

        image_to_world_service_handle = nh_.advertiseService("image_to_world", &ImageToWorld::image_to_world_service, this);

    }

    /**
    * @brief onCameraInfo gets the intrinsic camara info
    * @param camara_info the info
    */
    void onCameraInfo(const sensor_msgs::CameraInfoConstPtr &camara_info){

        if (camera_initialized_) return;
        else{
            ROS_INFO_STREAM("Camera Initialised!");
            model_.fromCameraInfo(camara_info);
            camera_initialized_ = true;
        }
    }

    bool image_to_world_service(apc_msgs::ImageToWorld::Request &req,
                                     apc_msgs::ImageToWorld::Response &res)
    {
        pcl::PointCloud<pcl::PointXYZRGB> input_cloud;
        pcl::fromROSMsg(req.cloud, input_cloud);
        auto image_x = req.image_x.data;
        auto image_y = req.image_y.data;

        auto world_x = 0.0;
        auto world_y = 0.0;
        auto world_z = 0.0;
        auto best_distance = 10000000.0;

        auto num_points = 0;

        for(auto &point: input_cloud) {
            num_points++;
            if(point.r == 0) {
                // Sometimes the pointcloud has 100k points at the origin.
                // This remedies it
                continue;
            }
            cv::Point3d xyz = cv::Point3d(point.x, point.y, point.z);
            cv::Point2d pt = model_.project3dToPixel(xyz);
            if (!std::isnan(pt.x) && !std::isnan(pt.y)){
                auto x_dist = std::abs(pt.x - image_x);
                auto y_dist = std::abs(pt.y - image_y);

                auto distance = std::sqrt(x_dist*x_dist + y_dist*y_dist);

                if (distance < best_distance) {
                    world_x = xyz.x;
                    world_y = xyz.y;
                    world_z = xyz.z;
                    best_distance = distance;
                }
                if (distance < 1.0) {
                    break;
                }
            }
        }

        cv::Point2d pixel_point(image_x, image_y);
        cv::Point3d xyz_unit_vector = model_.projectPixelTo3dRay(pixel_point);

        cv::Point3d xyz;

        if (num_points > 1) {
            xyz.x = xyz_unit_vector.x * world_z/xyz_unit_vector.z;
            xyz.y = xyz_unit_vector.y * world_z/xyz_unit_vector.z;
            xyz.z = xyz_unit_vector.z * world_z/xyz_unit_vector.z;
            ROS_INFO_STREAM("World depth from rgb is: " << xyz.z);
        } else {
            xyz.x = xyz_unit_vector.x * 0.5/xyz_unit_vector.z;
            xyz.y = xyz_unit_vector.y * 0.5/xyz_unit_vector.z;
            xyz.z = xyz_unit_vector.z * 0.5/xyz_unit_vector.z;
            ROS_INFO_STREAM("No points in provided segment. Choosing world depth of 0.5m (camera to item)");
        }

        res.world_x.data = xyz.x;
        res.world_y.data = xyz.y;
        res.world_z.data = xyz.z;

        return true;
    }

private:

    ros::NodeHandle nh_;
    image_geometry::PinholeCameraModel model_;
    bool camera_initialized_;

    ros::ServiceServer image_to_world_service_handle;
};

int main(int argc, char** argv) {

    ros::init(argc,argv,"image_to_world_node");
    ros::NodeHandle nh("~");
    ImageToWorld proposal(nh);

    ros::spin();

    ROS_INFO("Terminating image_to_world_node ...");

    return 0;
}


#ifndef RADIUMENGINE_RAY_HPP
#define RADIUMENGINE_RAY_HPP

#include <Core/Math/LinearAlgebra.hpp>
#include <Core/RaCore.hpp>
#include <Eigen/Geometry>

namespace Ra {
namespace Core {
namespace Math {

/// A structure representing a ray in space with an origin and a direction.
using Ray = Eigen::ParametrizedLine<Scalar, 3>;
inline Ray transformRay( const Ray& r, const Transform& t ) {
    return Ray( t * r.origin(), t.linear() * r.direction() );
}
} // namespace Math
} // namespace Core
} // namespace Ra

#endif // RADIUMENGINE_RAY_HPP


#include "../../include/waveO1/observer.hpp"
#include "../../include/waveO1/utilities.hpp"
#include "../../include/mymfem/utilities.hpp"

#include <fstream>
#include <iostream>
#include <Eigen/Core>

double WaveO1Observer
:: eval_dgpFtimeError (std::shared_ptr<GridFunction>& p,
                      std::shared_ptr<GridFunction>& v,
                      double t) const
{
    double errp=0;
    double errv=0;

    // assumes that all elements have the same geometry
    auto faceElGeomType = p->FESpace()->GetMesh()
                ->GetFaceGeometryType(0);

    // L2 error in pressure
    // set integration rules
    IntegrationRules rule1{};
    const IntegrationRule *ir1;
    int order1 = 2*p->FESpace()->GetFE(0)->GetOrder()+2;
    ir1 = &rule1.Get(faceElGeomType, order1);
    m_pE_coeff->SetTime(t);
    m_assembleDgpSmoothError->SetIntegrationRule(ir1);
    errp = m_assembleDgpSmoothError->Ftime(p.get(), m_pE_coeff.get());

    // L2 error in velocity
    // set integration rules
    IntegrationRules rule2{};
    const IntegrationRule *ir2;
    int order2 = 2*v->FESpace()->GetFE(0)->GetOrder()+2;
    ir2 = &rule2.Get(faceElGeomType, order2);
    m_vE_coeff->SetTime(t);
    m_assembleDgpSmoothError->SetIntegrationRule(ir2);
    errv = m_assembleDgpSmoothError->Ftime(v.get(), m_vE_coeff.get());

    return (errp + errv);
}

//! Computes DG error for a space-like face
//! at a specified time
double WaveO1Observer
:: eval_dgpFspaceError
(std::shared_ptr<GridFunction>& p,
 std::shared_ptr<GridFunction>& v,
 double t,
 std::shared_ptr<WaveO1InvSqMediumCoeff>&) const
{
    double errp=0;
    double errv=0;

    // assumes that all elements have the same geometry
    auto elGeomType = p->FESpace()
                ->GetFE(0)->GetGeomType();

    IntegrationRules rule1{}, rule2{};
    const IntegrationRule *ir1, *ir2;

    m_pE_coeff->SetTime(t);
    m_vE_coeff->SetTime(t);

    // pressure
    int order1 = 2*p->FESpace()->GetFE(0)->GetOrder()+2;
    ir1 = &rule1.Get(elGeomType, order1);
    m_assembleDgpSmoothError->SetIntegrationRule(ir1);
    errp = m_assembleDgpSmoothError->Fspace
            (p.get(), m_pE_coeff.get());
            //(p.get(), m_pE_coeff.get(), invSqMed.get());

    // velocity
    int order2 = 2*v->FESpace()->GetFE(0)->GetOrder()+2;
    ir2 = &rule2.Get(elGeomType, order2);
    m_assembleDgpSmoothError->SetIntegrationRule(ir2);
    m_vE->ProjectCoefficient(*m_vE_coeff);
    errv = m_assembleDgpSmoothError->Fspace(v.get(), m_vE_coeff.get());

    return (errp + errv);
}

//! Computes DG^{+} error
std::tuple <double, double> WaveO1Observer
:: eval_xtDgpError (BlockVector& W) const
{
    double errDgFtime=0;
    double errDgFspace=0;

    int Nt = m_tWspace->GetNE();
    int xdimW1 = m_xW1space->GetTrueVSize();
    int xdimW2 = m_xW2space->GetTrueVSize();

    Vector& p = W.GetBlock(0);
    Vector& v = W.GetBlock(1);

    Vector pSol(xdimW1);
    std::shared_ptr<GridFunction> pGSol
            = std::make_shared<GridFunction>
            (m_xW1space, pSol);

    Vector vSol(xdimW2);
    std::shared_ptr<GridFunction> vGSol
            = std::make_shared<GridFunction>
            (m_xW2space, vSol);

    // medium
    auto invSqMed = std::make_shared<WaveO1InvSqMediumCoeff>
            (m_testCase);

    ElementTransformation *tTrans = nullptr;
    const FiniteElement *tFe = nullptr;
    Vector tShape;
    Array<int> tVdofs;

    // DG^{+} error for time-like faces
    Eigen::VectorXd bufErrFtime(Nt);
    bufErrFtime.setZero();

    double errFtime;
    for (int n=0; n<Nt; n++)
    {
        m_tWspace->GetElementVDofs(n, tVdofs);
        tTrans = m_tWspace->GetElementTransformation(n);

        tFe = m_tWspace->GetFE(n);
        int tNdofs = tFe->GetDof();
        tShape.SetSize(tNdofs);

        int order = 2*tFe->GetOrder()+2;
        const IntegrationRule *ir
                = &IntRules.Get(tFe->GetGeomType(), order);

        for (int i = 0; i < ir->GetNPoints(); i++)
        {
            const IntegrationPoint &ip = ir->IntPoint(i);
            tTrans->SetIntPoint(&ip);
            tFe->CalcShape(ip, tShape);

            // build solution at time t
            Vector t;
            tTrans->Transform(ip, t);
            build_xSol_FG(p, tShape, tVdofs, pSol);
            build_xSol_FG(v, tShape, tVdofs, vSol);

            errFtime = eval_dgpFtimeError(pGSol, vGSol, t(0));

            double w = ip.weight*tTrans->Weight();
            bufErrFtime(n) += w*errFtime;
        }
    }
    errDgFtime = std::sqrt(bufErrFtime.sum());

    // DG^{+} error for space-like faces
    Eigen::VectorXd bufErrFspace(Nt);
    bufErrFspace.setZero();

    Vector t;
    IntegrationPoint ip0;
    ip0.Set1w(0, 1);

    for (int n=0; n<Nt; n++)
    {
        int np = n;

        // build solution at tn^{+}
        m_tWspace->GetElementVDofs(np, tVdofs);
        tTrans = m_tWspace->GetElementTransformation(np);
        tFe = m_tWspace->GetFE(np);
        int tNdofs = tFe->GetDof();
        tShape.SetSize(tNdofs);

        tTrans->SetIntPoint(&ip0);
        tFe->CalcShape(ip0, tShape);
        tTrans->Transform(ip0, t);
        build_xSol_FG(p, tShape, tVdofs, pSol);
        build_xSol_FG(v, tShape, tVdofs, vSol);

        // compute error
        bufErrFspace(n) = eval_dgpFspaceError
                (pGSol, vGSol, t(0), invSqMed);
    }
    errDgFspace = std::sqrt(2*bufErrFspace.sum());

    return {errDgFtime, errDgFspace};
}

// End of file


//Backend for the instanton finder in the AdS/Flat spacetime case.
#define DLL_EXPORT
#define INSTANTON_SOLVER_DLL_API //Comment this out if we wish to import these
//from a DLL instead.

#include "project_specific.h"

#include "multi_precision_definitions.h"
#include <boost/array.hpp>//For boost arrays (needed for ode_types)
#include "AdS.h"
#include <initializer_list>//Allow us to use initialiser lists for vector.
    //This simplifies the code somewhat.

//------------------------------------------------------------------------------
//Main function to export:
template< class value_type , class time_type >
void odeSolveAdSFlat(value_type false_vacuum , value_type true_vacuum ,
                     value_type barrier , potential< value_type >& V ,
                     time_type chimax , int odeSolverToUse,
                     value_type RelTol, value_type AbsTol, value_type stepError,
                     value_type xi,
                     value_type lowerBound , value_type upperBound ,
                     solution_grid< time_type , std::vector< value_type > ,
                     value_type >& solOut,value_type& DSout,
                     value_type precision,std::ostream& outStream,
                     bool track_scale_factor,bool compute_H0_linearisation)
{
//INPUTS:
    //Ok. we proceed by bisecting until we have a solution which overshoots
    //beyond
    //chimax, which stands in for infinity.

    //Overshoots occur when the solution crosses zero before chimax.
    //Undershoots when
    //its derivative crosses zero before then.

    //Instantiate the events function, state-type, dynamic BCs etc...
    typedef std::vector< value_type > AdSFlatSolType;//solution without action,
        //for speed.
    typedef std::vector< value_type > AdSFlatSolTypeAction;
    //typedef odeAdSFlat< AdSFlatSolType , time_type , value_type>odeRHSAdSFlat;
    typedef odeAdSFlat_delta_a< AdSFlatSolType , time_type , value_type>
            odeRHSAdSFlat;
    typedef eventsAdSFlat< AdSFlatSolType , time_type , value_type > evAdSFlat;
    /*typedef dynamicBCsAdSFlat_taylor< value_type , AdSFlatSolType , time_type>
            dynBCAdSFlat;*/
    typedef dynamicBCs_flatfv_delta< value_type , AdSFlatSolType , time_type>
            dynBCAdSFlat;
    /*typedef odeAdSFlatAction<AdSFlatSolTypeAction,time_type,value_type>
            odeRHSAdSFlatAction;*/
    typedef odeAdSFlatAction_delta_a<AdSFlatSolTypeAction,time_type,value_type>
            odeRHSAdSFlatAction;
    typedef eventsAdSFlat< AdSFlatSolTypeAction , time_type , value_type >
            evAdSFlatAction;
    typedef dynamicBCsAdSFlat_taylor< value_type , AdSFlatSolTypeAction ,
                                      time_type > dynBCAdSFlatAction;
    typedef eventsGravFriction<AdSFlatSolType,time_type,value_type> eventsGrav;
    typedef odeGravFriction< AdSFlatSolType , time_type , value_type > odeGrav;
    typedef dynamicBCsGravFriction_taylor< value_type , AdSFlatSolType ,
                                      time_type > dynBCGrav;
    typedef odeGravFrictionAction<AdSFlatSolType,time_type,value_type>
                                    odeGravAction;

    //Instantiation for overshoot/undershoot:

    const value_type _1p0 = value_type(1.0);
    const value_type _0p0 = value_type(0.0);
    const value_type _2p0 = value_type(2.0);
    const value_type _6p0 = value_type(6.0);
    std::cout.precision(50);

    //Potential
    value_type W0 = value_type(0.0);
    value_type h;
    if(xi >= _0p0 && xi < _1p0/_6p0)
    {
        h = _1p0;//h = phi_scale/M_p so h = 1 is Planck units.
    }
    else
    {
        h = _1p0/sqrt(abs(xi));
        //h = value_type(1.0);//h = phi_scale/M_p so h = 1 is Planck units.
    }
	value_type V0 = _0p0;//V at phi = 0.
    odeRHSAdSFlat odeToSolveTrack(V,xi, h, V0); //ode to solve (gravitational
                                                           //instanton equation)
    odeGrav odeToSolveNoTrack(V,xi, h, V0);
    odeRHS<std::vector<value_type>,time_type,value_type>* odeToSolve;
    if(track_scale_factor)
    {
        odeToSolve = &odeToSolveTrack;
    }
    else
    {
        odeToSolve = &odeToSolveNoTrack;
    }

    evAdSFlat eventsTrack(false_vacuum,true_vacuum,false_vacuum,outStream);
    //events function.
    eventsGrav eventsNoTrack(false_vacuum,true_vacuum);
    events_function< std::vector<value_type> , time_type, value_type >* events;
    if(track_scale_factor)
    {
        events = &eventsTrack;
    }
    else
    {
        events = &eventsNoTrack;
    }
        //Checks for events labelled IE:
        //IE = 1 -> y crosses zero (terminal event, => overshoot)
        //IE = 2 -> y' crosses zero (terminal event, => undershoot)
        //IE = 3 -> a crosses zero (non-terminal event, for reference only)
        //IE = 4 -> y crosses lowerBound (usually 0, but could be different -
                                            //terminal event)
        //IE = 5 -> y crosses upperBound (terminal event, indicates bad initial
                                          //condition or perhaps a->0)
    dynBCAdSFlat bcAdjusterTrack(xi,h,V0,V,stepError,1);
    dynBCGrav bcAdjusterNoTrack(xi,h,V0,V,stepError,1);
    //Performs analytic
    // first
    //step using a taylor series approximation.
        //stepError - determines size of initial step.
        //version (last argument) - specifies which version of the ode we are
            //using.
            //version = 0 -> normal ode, no action.
            //version = 1 -> ode with action
            //version = 2 -> linearised ode.
    dynamicBCs< value_type , std::vector<value_type> , time_type >* bcAdjuster;
    if(track_scale_factor)
    {
        bcAdjuster = &bcAdjusterTrack;
    }
    else
    {
        bcAdjuster = &bcAdjusterNoTrack;
    }

    //Search loop:
    int nMax = 200;
    int counter = 0;
    //Current bounds on location of bounce:
    value_type lower = lowerBound;//Barrier
    value_type upper = upperBound;//True vacuum
    outStream << "Lower initial = " << lower << std::endl;
    outStream << "Upper initial = " << upper << std::endl;
    value_type guess;
    multi diff = abs(upper - lower);
    outStream << "diff = " << diff << std::endl;
    while(diff > precision)
    {
        //outStream << "Loop stage 1" <<std::endl;
        //Bisect to get initial guess:
        guess = (lower + upper)/_2p0;
        outStream << "\nn = " << counter << ", guess = " << guess << "\n";
        outStream << "\ndiff = " << upper - lower;
        //outStream << "Loop stage 2" <<std::endl;

        //Search for FWHMif(track_scale_factor)
        if(track_scale_factor)
        {
            eventsTrack.value_to_check = (guess + false_vacuum)/_2p0;
        }

        //Formulate initial conditions vector:
        AdSFlatSolType y0;
        //value_type y0List[] = { guess , _0p0 , _0p0 , _1p0 };
        value_type y0List[] = { guess , _0p0 , _0p0 , _1p0 , _0p0, _0p0 };

        y0.assign(y0List,y0List + 6);
        //AdSFlatSolType y0NoTrack;
        //value_type y0NoTrackList[] = { guess , _0p0 , _0p0 , _1p0};
        //y0NoTrack.assign(y0TrackList,y0TrackList + 4);

        //AdSFlatSolType y0 = track_scale_factor ? y0Track : y0NoTrack;
        //Step away to avoid co-ordinate singularity at chi = 0:
        //outStream << "Loop stage 3" <<std::endl;
        AdSFlatSolType y01step = y0;
        time_type epsilon = (*bcAdjuster)(y0,y01step);//Copies solution after
            //step of epsilon into y01step. Last argument tells us whether
            //we step forwards (0) or backwards (1) in time.
        //Set up time range:
        //outStream << "Loop stage 4" <<std::endl;
        //time_type tspanArray[2] = { epsilon , chimax };
        std::vector< time_type > tspan;
        time_type tspan_list[] = { epsilon , chimax };
        tspan.assign(tspan_list,tspan_list + 2);
        //tspan.assign(tspanArray,tspanArray + 2);
        //Solve the ode:
        //outStream << "Loop stage 5" <<std::endl;
        solution_grid< time_type , AdSFlatSolType , value_type > solTemp;
        time_type initStepMax = time_type(0.01)/time_type(chimax);
        //outStream << "Integrating..." <<std::endl;
        int nSuccess = ode_solve<time_type,AdSFlatSolType,value_type>
            (*odeToSolve,y01step,tspan,(*events),odeSolverToUse,solTemp,RelTol,
             AbsTol,initStepMax);
        if(nSuccess != 1)
        {
            outStream << "nSuccess = " << nSuccess
                      << " when solving for AdS-Flat instantons." << std::endl;
            throw "Integration failure.";
        }
        //Now check whether we found an overshoot or an undershoot:
        bool overshoots = false;
        bool undershoots = false;
        for(int j = 0; j < int(solTemp.IE.size());j++)
        {
            switch(solTemp.IE[j])
            {
            case 1:
                //overshoot.
                //outStream << "Overshoot." << std::endl;
                //reportToCaller("Overshoot.\n");
                if(track_scale_factor)
                {
                    overshoots = true;
                    upper = guess;
                }
                //else we just found a' = 0, so not terminal.
                break;
            case 2:
                //undershoot, IF it occurs on the opposite side of the barrier
                //to the true vacuum (theoretically impossible to occur on the
                //other side, so must be a numerical artefact if it does).
                if((true_vacuum - barrier)*(solTemp.YE[j][0] - barrier)
                   < value_type(0.0))
                {
                    //outStream << "Undershoot." << std::endl;
                    //reportToCaller("Undershoot.\n");
                    undershoots = true;
                    lower = guess;
                }
                break;
            case 3:
                //Found a = 0 point.
                outStream << "Found a = 0. Error, or possibly background is "
                          << "de-Sitter like." << std::endl;
                //reportToCaller("Found a = 0. Error, or possibly background is
                                 //de-Sitter like.");
                throw "Unexpected ode-event encountered during integration.";
            case 4:
                //crossed upper bound, so overshoot for tracking a'/a version,
                //FWHM for the a',a separately tracked version.
                if(!track_scale_factor)
                {
                    overshoots = true;
                    upper = guess;
                    eventsNoTrack.y0upperBound = upper;
                }
                break;
            case 5:
                //Crossed lower bound (in both cases).
                //Usually indicates an overshoot:
                //outStream << "Overshoot." << std::endl;
                //reportToCaller("Overshoot.\n");
                overshoots = true;
                upper = guess;
                if(!track_scale_factor)
                {
                    eventsNoTrack.y0upperBound = upper;
                }
                break;
            case 6:
                //Crossed upper bound. Only
                //occurs for the a'/a version.
                //Indicates a bad initial condition.
                outStream << "Upper bound crossed during ode integration. "
                          << "Possibly a bad initial condition (wrong side of "
                          << "true vacuum) or solution undershot and event "
                          << "detection failed to detect this." << std::endl;
                //reportToCaller("Upper bound crossed during ode integration.
                    //Possibly a bad initial condition (wrong side
                    //of true vacuum) or solution undershot and event detection
                    //failed to detect this.\n");
                throw "Unexpected ode-event encountered during integration.";
            default:
                outStream << "Unrecognised ode-event." << std::endl;
                //reportToCaller("Unrecognised ode-event.\n");
                throw "Unexpected ode-event encountered during integration.";
            }
            if(overshoots || undershoots)
            {
                break;
            }
        }
        counter++;
        diff = abs(lower - upper);
        //outStream << "y0Upper - y0lower = " << diff << std::endl;
        if(overshoots)
        {
            outStream << "Overshoot.\n" << std::endl;
        }
        else if(undershoots)
        {
            outStream << "Undershoot.\n" << std::endl;
        }

        //Dump all available data if we get neither an overshoot nor an
        //undershoot, as this is
        //theoretically impossible so we want to understand what went wrong:
        else if(!(overshoots||undershoots))
        {
            outStream << "Neither undershoot nor overshoot detected. "
                      << "(Perhaps integration range should be extended?) "
                      << "Relevant data:\n";
            outStream << "no. of events = " << int(solTemp.IE.size())
                      << std::endl;
            for(int i = 0; i < int(solTemp.IE.size());i++)
            {
                outStream << "IE[" << i << "] = " << solTemp.IE[i] << std::endl;
                outStream << "YE[" << i << "] = " << std::endl;
                for(int j = 0;j < y01step.size();j++)
                {
                    outStream << solTemp.YE[i][j] << std::endl;
                }
                outStream << "TE[" << i << "] = " << solTemp.TE[i] << std::endl;
            }
            outStream << "y0 = \n";
            for(int i = 0;i < y01step.size();i++)
            {
                outStream << y01step[i] << std::endl;
            }
            outStream << "yend = \n";
            for(int i = 0;i < y01step.size();i++)
            {
                outStream << solTemp.Y[int(solTemp.Y.size() - 1)][i]
                          << std::endl;
            }
            outStream << "Tend = " << solTemp.T[int(solTemp.T.size()) - 1]
                      << std::endl;
            //Now exit, throwing an error:
            throw "Error - Neither overshoot nor undershoot detected.\n";
        }

    }
    //We found the bounce. Are we within precision?
    if(diff > precision)
    {
        outStream << "Loop exited before finding bounce." << std::endl;
        //reportToCaller("Loop exited before finding bounce.\n");
    }

    //Compute action:
    odeRHSAdSFlatAction odeToSolveActionTrack(V,xi,h,V0);
    evAdSFlatAction eventsActionTrack(false_vacuum,true_vacuum,false_vacuum,
                                 outStream);
    dynBCAdSFlatAction bcAdjusterActionTrack(xi,h,V0,V,stepError,2);
    odeGravAction odeToSolveActionNoTrack(V,xi,h,V0);
    dynBCGrav bcAdjusterActionNoTrack(xi,h,V0,V,stepError,2);
    odeRHS<std::vector<value_type>,time_type,value_type>* odeToSolveAction;
    dynamicBCs< value_type , std::vector<value_type> , time_type >*
        bcAdjusterAction;
    eventsGrav eventsActionNoTrack(false_vacuum,true_vacuum);
    events_function< std::vector<value_type> , time_type, value_type >*
        eventsAction;
    if(track_scale_factor)
    {
        odeToSolveAction = &odeToSolveActionTrack;
        bcAdjusterAction = &bcAdjusterActionTrack;
        eventsAction = &eventsActionTrack;
    }
    else
    {
        odeToSolveAction = &odeToSolveActionNoTrack;
        bcAdjusterAction = &bcAdjusterActionNoTrack;
        eventsAction = &eventsActionNoTrack;
    }
    //outStream << "Computing final bounce solution... " << std::endl;

    //Check for when the solution passes through its FWHM:
    //This feature isn't yet implemented for the NoTrack version, so we don't
    //use it in that case.
    if(track_scale_factor)
    {
        eventsActionTrack.value_to_check = (guess + false_vacuum)/_2p0;
    }

    //Boundary conditions at x = 0
    AdSFlatSolTypeAction y0Action;
    //value_type y0ActionList[] = {guess,_0p0,_0p0,_1p0,_0p0};
    value_type y0ActionList[] = {guess,_0p0,_0p0,_1p0,_0p0,_0p0,_0p0};
    //y0Action.assign(y0ActionList,y0ActionList + 5);
    y0Action.assign(y0ActionList,y0ActionList + 7);
    AdSFlatSolTypeAction y0Action1step = y0Action;
    //Same BCs work for NoTrack, where we use theta = log(a + 1)

    //Step to boundary conditions slightly offset from x = 0:
    time_type epsAction = (*bcAdjusterAction)(y0Action,y0Action1step);

    //Integration range:
    time_type tspanArray[2] = { epsAction , chimax };
    std::vector< time_type > tspan;
    tspan.assign(tspanArray,tspanArray + 2);
    time_type initStepMax = time_type(0.01)/time_type(chimax);

    //Integrate:
    int nSuccess = ode_solve<time_type,AdSFlatSolTypeAction,value_type>
    (*odeToSolveAction,y0Action1step,tspan,*eventsAction,odeSolverToUse,solOut,
     RelTol,AbsTol,initStepMax);
    if(nSuccess != 1)
    {
        outStream << "nSuccess = " << nSuccess << " when solving for AdS-Flat"
                  << "instantons." << std::endl;
        throw "Integration failure.";
    }
    DSout = solOut.Y[int(solOut.T.size()) - 1][4];
    outStream << "Done." << std::endl;
    //Done.
}
//------------------------------------------------------------------------------
template< class value_type , class time_type >
void odeSolveAdSFlatSingle(value_type y0, value_type false_vacuum ,
                           value_type true_vacuum , value_type barrier ,
                           potential< value_type >& V , time_type chimax ,
                           int odeSolverToUse,
                           value_type RelTol, value_type AbsTol,
                           value_type stepError, value_type xi,
                           value_type lowerBound , value_type upperBound ,
                           solution_grid< time_type , std::vector< value_type >,
                           value_type >& solOut,value_type& DSout,
                           value_type precision,std::ostream& outStream,
                           bool track_scale_factor)
{
//INPUTS:

     outStream << "Starting Integration for y0 = \n" << y0 << std::endl;
     outStream << "\nV(y0) = " << V(y0) << std::endl;

     const value_type _1p0 = value_type(1.0);
     const value_type _0p0 = value_type(0.0);
     const value_type _2p0 = value_type(2.0);
     const value_type _6p0 = value_type(6.0);

    //Ok. we proceed by bisecting until we have a solution which overshoots
    //beyond
    //chimax, which stands in for infinity.

    //Overshoots occur when the solution crosses zero before chimax.
    //Undershoots when
    //its derivative crosses zero before then.

    //Instantiate the events function, state-type, dynamic BCs etc...
    typedef std::vector< value_type > AdSFlatSolType;//solution without action,
        //for speed.
    typedef std::vector< value_type > AdSFlatSolTypeAction;
    typedef odeAdSFlat< AdSFlatSolType , time_type , value_type > odeRHSAdSFlat;
    typedef eventsAdSFlat< AdSFlatSolType , time_type , value_type > evAdSFlat;
    typedef dynamicBCsAdSFlat_taylor< value_type , AdSFlatSolType , time_type >
            dynBCAdSFlat;
    typedef odeAdSFlatAction< AdSFlatSolTypeAction , time_type , value_type >
            odeRHSAdSFlatAction;
    typedef eventsAdSFlat< AdSFlatSolTypeAction , time_type , value_type >
            evAdSFlatAction;
    typedef dynamicBCsAdSFlat_taylor< value_type , AdSFlatSolTypeAction ,
                                      time_type > dynBCAdSFlatAction;
    typedef eventsGravFriction<AdSFlatSolType,time_type,value_type> eventsGrav;
    typedef odeGravFriction< AdSFlatSolType , time_type , value_type > odeGrav;
    typedef dynamicBCsGravFriction_taylor< value_type , AdSFlatSolType ,
                                      time_type > dynBCGrav;
    typedef odeGravFrictionAction<AdSFlatSolType,time_type,value_type>
                                    odeGravAction;

    //Instantiation for overshoot/undershoot:

    //Compute action:
    value_type W0 = _0p0;
    value_type h;
    if(xi >= _0p0 && xi < _1p0/_6p0)
    {
        h = _1p0;//h = phi_scale/M_p so h = 1 is Planck units.
    }
    else
    {
        h = _1p0/sqrt(abs(xi));
        //h = value_type(1.0);//h = phi_scale/M_p so h = 1 is Planck units.
    }
    //Setup odeRHS, event-detector, and dynamicBC objects needed:
    odeRHSAdSFlatAction odeToSolveActionTrack(V,xi,h,W0);
    evAdSFlatAction eventsActionTrack(false_vacuum,true_vacuum,false_vacuum,
                                 outStream);
    dynBCAdSFlatAction bcAdjusterActionTrack(xi,h,W0,V,stepError,2);
    odeGravAction odeToSolveActionNoTrack(V,xi,h,W0);
    dynBCGrav bcAdjusterActionNoTrack(xi,h,W0,V,stepError,2);
    odeRHS<std::vector<value_type>,time_type,value_type>* odeToSolveAction;
    dynamicBCs< value_type , std::vector<value_type> , time_type >*
        bcAdjusterAction;
    eventsGrav eventsActionNoTrack(false_vacuum,true_vacuum);
    events_function< std::vector<value_type> , time_type, value_type >*
        eventsAction;
    if(track_scale_factor)
    {
        odeToSolveAction = &odeToSolveActionTrack;
        bcAdjusterAction = &bcAdjusterActionTrack;
        eventsAction = &eventsActionTrack;
    }
    else
    {
        odeToSolveAction = &odeToSolveActionNoTrack;
        bcAdjusterAction = &bcAdjusterActionNoTrack;
        eventsAction = &eventsActionNoTrack;
    }

    AdSFlatSolTypeAction y0Action;
    value_type y0ActionList[] = { y0 , _0p0 , _0p0 , _1p0 , _0p0 };
    y0Action.assign(y0ActionList,y0ActionList + 5);
    AdSFlatSolTypeAction y0Action1step = y0Action;
    outStream << "Compute boundary conditions.\n";
    time_type epsAction = (*bcAdjusterAction)(y0Action,y0Action1step);
    outStream << "Done.";
    outStream << "\ny0 = " << y0Action1step[0];
    outStream << "\ny0p = " << y0Action1step[1];
    outStream << "\na0 = " << y0Action1step[2];
    outStream << "\na0p = " << y0Action1step[3];
    outStream << "\nS = " << y0Action1step[4];
    time_type tspanArray[2] = { epsAction , chimax };
    std::vector< time_type > tspan;
    tspan.assign(tspanArray,tspanArray + 2);
    time_type initStepMax = time_type(0.01)/time_type(chimax);
    outStream << "Starting integration.";
    int nSuccess = ode_solve<time_type,AdSFlatSolTypeAction,value_type>
    (*odeToSolveAction,y0Action1step,tspan,*eventsAction,odeSolverToUse,solOut,
     RelTol,AbsTol,initStepMax);
    if(nSuccess != 1)
    {
        outStream << "nSuccess = " << nSuccess << " when solving for AdS-Flat"
                  << " instantons." << std::endl;
        throw "Integration failure.";
    }
    DSout = solOut.Y[int(solOut.T.size()) - 1][4];
    //Done.
}
//------------------------------------------------------------------------------

//------------------------------------------------------------------------------
//Explicit instantiation and export of this function:
template DLL_EXPORT void odeSolveAdSFlat
    < multi , multi >
    (multi , multi , multi , potential< multi >& , multi , int,multi, multi ,
     multi, multi, multi , multi ,
     solution_grid< multi , std::vector< multi > , multi >&,
     multi&, multi, std::ostream&,bool,bool);
//------------------------------------------------------------------------------
template DLL_EXPORT void odeSolveAdSFlatSingle
    < multi , multi >
    (multi , multi , multi , multi , potential< multi >& , multi , int, multi,
     multi , multi, multi, multi , multi ,
     solution_grid< multi , std::vector< multi > , multi >&,
     multi&, multi, std::ostream&,bool);
//------------------------------------------------------------------------------


//
// Created by KH on 3/7/2020.
//

#include "CNode.h"
#include "CTree.h"

#include <klib/math.h>
#include <klib/random.h>
#include <klib/vecutils.h>

#include <memory>
#include <boost/container/small_vector.hpp>
#include <boost/container/flat_set.hpp>
#include <boost/iterator/counting_iterator.hpp>

#include <tsl/robin_set.h>

namespace igen {

static const double Epsilon = 1E-4;

static const int V_LOG_BUILD = 100;

CNode::CNode(CTree *tree, CNode *parent, int id, std::array<boost::sub_range<vec<PConfig>>, 2> configs) :
        tree(tree), parent(parent), id_(id), configs_(std::move(configs)) {
    depth_ = (parent ? parent->depth() + 1 : 0);
    if (parent) {
        tested_vars_ = parent->tested_vars_;
    } else {
        tested_vars_ = dom()->create_vec_vars<bool>();
    }
}

const PDomain &CNode::dom() const {
    return tree->ctx()->dom();
}

const PVarDomain &CNode::dom(int var_id) const {
    return dom()->vars()[var_id];
}

void CNode::calc_freq() {
    int n_vars = (int) dom()->n_vars();
    auto &freq = tree->t_freq;
    for (auto &a : freq) for (auto &b : a) for (auto &c : b) c = 0;
    for (const auto &c : configs_[0]) {
        const auto &vals = c->values();
        for (int var_id = 0; var_id < n_vars; ++var_id) {
            freq[0][var_id][vals[var_id]]++;
        }
    }
    for (const auto &c : configs_[1]) {
        const auto &vals = c->values();
        for (int var_id = 0; var_id < n_vars; ++var_id) {
            freq[1][var_id][vals[var_id]]++;
        }
    }
}

void CNode::calc_inf_gain() {
    splitvar = -1;
    double log2ntotal = log2(n_total());
    auto &freq = tree->t_freq;
    auto &info = tree->t_info, &gain = tree->t_gain;

    // Calc info
    for (int var_id = 0; var_id < dom()->n_vars(); ++var_id) {
        if (tested_vars_[var_id]) continue;
        double sum = 0;
        for (int val = 0; val < dom()->n_values(var_id); ++val) {
            int n = freq[0][var_id][val] + freq[1][var_id][val];
            sum += n * log2(n);
        }
        info[var_id] = log2ntotal - sum / n_total();
    }

    // Calc base_info
    double base_info = 0;
    for (int hit = 0; hit <= 1; ++hit) {
        int n = int(configs_[hit].size());
        base_info += n * log2(n);
    }
    base_info = log2ntotal - base_info / n_total();

    // Calc gain
    double total_gain = 0;
    for (int var_id = 0; var_id < dom()->n_vars(); ++var_id) {
        if (tested_vars_[var_id]) continue;
        double &g = (gain[var_id] = 0);
        for (int val = 0; val < dom()->n_values(var_id); ++val) {
            double sum = 0;
            int total = 0;
            for (int hit = 0; hit <= 1; ++hit) {
                int n = freq[hit][var_id][val];
                sum += n * log2(n);
                total += n;
            }
            sum = total * log2(total) - sum;
            g += sum;
        }
        g = base_info - g / n_total();
        total_gain += g;
    }

    // Calc avgain & mdl
    avgain = 0;
    possible = 0;
    for (int var = 0; var < dom()->n_vars(); ++var) {
        if (tested_vars_[var]) continue;
        if (gain[var] >= Epsilon &&
            (tree->multi_val_ || dom(var)->n_values() < 0.3 * (tree->n_cases_ + 1))) {
            possible++;
            avgain += gain[var];
        } else {
            //Gain[Att] = None;
        }
    }

    avgain /= possible;
    mdl = log2(possible) / int(hit_configs().size() + miss_configs().size());
    mingain = avgain * tree->avgain_wt_ + mdl * tree->mdl_wt_;
}

int CNode::select_best_var(bool first_pass) {
    auto &info = tree->t_info, &gain = tree->t_gain;
    splitvar = -1;
    find_pass = first_pass ? 1 : 2;
    double bestratio = -1000;
    int bestnbr = dom()->n_all_values();

    for (int var = 0; var < dom()->n_vars(); ++var) {
        if (tested_vars_[var]) continue;

        double inf = info[var];

        if (first_pass) {
            if (gain[var] < 0.999 * mingain || inf <= 0) continue;
        } else {
            if (inf <= 0) inf = Epsilon;
        }

        double val = gain[var] / inf;
        int nbr = dom()->n_values(var);

        if (val > bestratio
            || (val > 0.999 * bestratio && (nbr < bestnbr || (nbr == bestnbr && gain[var] > gain[var])))) {
            splitvar = var;
            bestratio = val;
            bestnbr = nbr;
        }
    }

    return splitvar;
}

void CNode::create_childs() {
    CHECK(0 <= splitvar && splitvar < dom()->n_vars());
    int nvalues = dom()->n_values(splitvar);

    tested_vars_[splitvar] = true;
    childs.resize(nvalues);

    for (int hit = 0; hit <= 1; ++hit) {
        auto &c = configs_[hit];
        const auto &freq = tree->t_freq[hit][splitvar];
        auto &tmp_conf = tree->t_conf;
        tmp_conf.assign(std::make_move_iterator(c.begin()), std::make_move_iterator(c.end()));
        auto &curpos = tree->t_curpos[hit];
        curpos.resize(nvalues);
        curpos[0] = c.begin();
        for (int i = 1; i < nvalues; ++i) {
            curpos[i] = curpos[i - 1] + freq[i - 1];
        }
        for (auto &x : tmp_conf) {
            int v = x->values()[splitvar];
            *curpos[v]++ = move(x);
        }
        CHECK(curpos[nvalues - 1] == c.end());
        //DCHECK(std::is_sorted(c.begin(), c.end(), [splitvar = splitvar](const auto &a, const auto &b) {
        //    return a->get(splitvar) < b->get(splitvar);
        //}));
        //std::stable_sort(c.begin(), c.end(), [splitvar = splitvar](const auto &a, const auto &b) {
        //    return a->get(splitvar) < b->get(splitvar);
        //});
    }

    for (int i = 0; i < nvalues; ++i) {
        auto &pos = tree->t_curpos;
        auto &c = configs_;
        boost::sub_range<vec<PConfig>> miss_conf = {i == 0 ? c[0].begin() : pos[0][i - 1], pos[0][i]};
        boost::sub_range<vec<PConfig>> hit_conf = {i == 0 ? c[1].begin() : pos[1][i - 1], pos[1][i]};
        childs[i] = new CNode(tree, this, i, {miss_conf, hit_conf});
    }
}

bool CNode::evaluate_split() {
    if (is_leaf()) {
        CHECK(n_hits() == 0 || n_misses() == 0); // Same config lead to 2 different result?
        GVLOG(V_LOG_BUILD) << "\n<" << depth() << ">: " << n_total() << " cases\n    "
                           << (leaf_value() ? "HIT" : "MISS");
        min_cases_in_one_leaf_ = n_total();
        return false;
    }

    calc_freq();
    calc_inf_gain();
    if (select_best_var(true) == -1) select_best_var(false);
    CHECK_NE(splitvar, -1);
    VLOG_BLOCK(V_LOG_BUILD, print_tmp_state(log << "\n<" << depth() << ">: " << n_total() << " cases\n"));

    split_by = dom()->vars().at(splitvar);
    create_childs();

    min_cases_in_one_leaf_ = std::numeric_limits<int>::max();
    for (auto &c : childs) {
        c->evaluate_split();
        min_cases_in_one_leaf_ = std::min(min_cases_in_one_leaf_, c->min_cases_in_one_leaf_);
    }

    return split_by != nullptr;
}

std::ostream &CNode::print_tmp_state(std::ostream &output, const str &indent) const {
    auto &freq = tree->t_freq;
    auto &info = tree->t_info, &gain = tree->t_gain;
    for (const auto &var : dom()->vars()) {
        // if (info[var->id()] < 0 || gain[var->id()] < 0)
        //     continue;
        if (tested_vars_[var->id()]) continue;
        output << indent << "Var " << var->name() << ": \n";
        fmt::print(output, "{}{}[{:>4}{:>8}{:>8}]\n", indent, indent, "Val", "MISS", "HIT");
        for (int i = 0; i < var->n_values(); ++i) {
            fmt::print(output, "{}{}[{:>4}{:>8}{:>8}]\n", indent, indent,
                       i, freq[0][var->id()][i], freq[1][var->id()][i]);
        }
        fmt::print(output, "{}{}info {:.3f}, gain {:.3f}, val {:.3f}\n", indent, indent,
                   info[var->id()], gain[var->id()], gain[var->id()] / info[var->id()]);
    }
    fmt::print(output, "{}av gain={:.3f}, MDL ({}) = {:.3f}, min={:.3f}\n",
               indent, avgain, possible, mdl, mingain);
    fmt::print(output, "{}best var ", indent);
    if (splitvar == -1) {
        output << "N/A\n";
    } else {
        fmt::print(output, "{}: info {:.3f}, gain {:.3f}, val {:.3f}{}",
                   dom()->name(splitvar), info[splitvar], gain[splitvar], gain[splitvar] / info[splitvar],
                   find_pass == 1 ? "\n" : ", second-pass\n");
    }
    return output;
}

bool CNode::leaf_value() const {
    CHECK(is_leaf());
    if (hit_configs().empty() && miss_configs().empty())
        return tree->default_hit_;
    return !hit_configs().empty();
}

z3::expr CNode::build_zexpr_mixed() const {
    if (is_leaf()) {
        return tree->ctx()->zbool(leaf_value());
    }
    CHECK(splitvar != -1 && int(childs.size()) == dom(splitvar)->n_values());

    z3::expr res(tree->ctx_mut()->zctx()), e = res;
    bool empty_res = true;

    for (int val = 0; val < int(childs.size()); ++val) {
        bool need_or = false;
        if (childs[val]->is_leaf()) {
            if (childs[val]->leaf_value()) {
                e = dom(splitvar)->eq(val);
                need_or = true;
            }
        } else {
            e = dom(splitvar)->eq(val) && childs[val]->build_zexpr_mixed();
            need_or = true;
        }
        if (need_or) {
            if (empty_res) res = e, empty_res = false;
            else res = res || e;
        }
    }
    CHECK(!empty_res);
    return res;
}

void CNode::build_zexpr_disj_conj(z3::expr_vector &vec_res, const expr &cur_expr) const {
    if (is_leaf() && leaf_value()) {
        vec_res.push_back(cur_expr);
        return;
    }
    for (int val = 0; val < int(childs.size()); ++val) {
        expr next_expr = depth_ == 0 ?
                         dom(splitvar)->eq(val) :
                         cur_expr && dom(splitvar)->eq(val);
        childs[val]->build_zexpr_disj_conj(vec_res, next_expr);
    }
}

void CNode::print_node(std::ostream &output, str &prefix) const {
    if (is_leaf()) {
        output << (leaf_value() ? "HIT" : "MISS") << " (" << n_total() << ")\n";
        return;
    }
    CHECK_NE(splitvar, -1);
    CHECK_GE(n_childs(), 2);
    const str &var_name = dom(splitvar)->name();

    prefix.append(":   ");
    for (int v = 0; v < n_childs(); ++v) {
        if (v == n_childs() - 1) prefix.at(prefix.size() - 4) = ' ';

        const PMutCNode n = childs[v];
        if (v == 0) {
            if (prefix.size() >= 8)
                output << std::string_view(prefix.data(), prefix.size() - 8) << ":...";
        } else {
            output << std::string_view(prefix.data(), prefix.size() - 4);
        }
        output << var_name << " = " << dom(splitvar)->label(v) << (n->is_leaf() ? ": " : ":\n");
        n->print_node(output, prefix);
    }
    prefix.resize(prefix.size() - 4);
}

std::pair<bool, int> CNode::test_config(const PConfig &conf) const {
    if (is_leaf()) return {leaf_value(), n_total()};
    CHECK_NE(splitvar, -1);
    return childs.at(size_t(conf->get(splitvar)))->test_config(conf);
}

std::pair<bool, int> CNode::test_add_config(const PConfig &conf, bool val) {
    if (is_leaf()) {
        bool leaf_val = leaf_value();
        if (leaf_val == val) {
            min_cases_in_one_leaf_++;
            new_configs_.push_back(conf);
            CHECK_EQ(min_cases_in_one_leaf_, n_total() + new_configs_.size());
        }
        return {leaf_val, n_total()};
    }
    CHECK_NE(splitvar, -1);
    auto res = childs.at(size_t(conf->get(splitvar)))->test_add_config(conf, val);
    min_cases_in_one_leaf_ = std::numeric_limits<int>::max();
    for (auto &c : childs) {
        min_cases_in_one_leaf_ = std::min(min_cases_in_one_leaf_, c->min_cases_in_one_leaf_);
    }
    return res;
}

// =====================================================================================================================

void CNode::gather_small_leaves(vec<PConfig> &res, int min_confs, int max_confs, const PMutConfig &curtpl) const {
    if (is_leaf()) {
        if (min_confs <= min_cases_in_one_leaf_ && min_cases_in_one_leaf_ <= max_confs)
            res.emplace_back(new Config(curtpl));
        return;
    }
    CHECK_NE(splitvar, -1);
    for (int v = 0; v < n_childs(); ++v) {
        CHECK(curtpl->get(splitvar) == -1);
        curtpl->set(splitvar, v);
        childs[v]->gather_small_leaves(res, min_confs, max_confs, curtpl);
        curtpl->set(splitvar, -1);
    }
}

void CNode::gather_leaves_nodes(vec<ptr<const CNode>> &res, int min_confs, int max_confs) const {
    if (is_leaf()) {
        if (min_confs <= min_cases_in_one_leaf_ && min_cases_in_one_leaf_ <= max_confs)
            res.emplace_back(this);
        return;
    }
    CHECK_NE(splitvar, -1);
    for (int v = 0; v < n_childs(); ++v) {
        childs[v]->gather_leaves_nodes(res, min_confs, max_confs);
    }
}

void CNode::gen_tpl(Config &conf) const {
    if (is_leaf()) {
        CHECK_EQ(splitvar, -1);
    }
    if (parent != nullptr) {
        CHECK_NE(parent->splitvar, -1);
        conf.set(parent->splitvar, id_);
        parent->gen_tpl(conf);
    }
}

vec<ptr<Config>> CNode::gen_one_convering_configs(int lim) const {
    const PMutContext ctx = tree->ctx();
    const PDomain dom = tree->ctx()->dom();
    //====
    Config templ(ctx);
    templ.set_all(-1);
    gen_tpl(templ);
    //====
    tsl::robin_set<hash_t> shash;
    shash.reserve((size_t) n_total());
    // Note: either hit_configs() or miss_configs() is empty
    CHECK(hit_configs().empty() || miss_configs().empty());
    bool use_solver = ((n_total() + sz(new_configs_)) <= 1000);
    if (use_solver) {
        for (const auto &c : hit_configs()) shash.insert(c->hash());
        for (const auto &c : miss_configs()) shash.insert(c->hash());
        for (const auto &c : new_configs_) shash.insert(c->hash());
    }
    //====
    vec<sm_vec<int>> SetVAL;
    SetVAL.reserve((size_t) dom->n_vars());
    int n_finished = 0;
    for (int i = 0; i < dom->n_vars(); i++) {
        if (templ.get(i) == -1) {
            SetVAL.emplace_back(boost::counting_iterator<int>(0), boost::counting_iterator<int>(dom->n_values(i)));
        } else {
            SetVAL.emplace_back();
            n_finished++;
        }
    }

    VLOG(100, "Tpl: \n") << templ;

    vec<PMutConfig> ret;
    std::unique_ptr<Z3Scope> z3;
    while (n_finished < dom->n_vars() && int(ret.size()) < lim) {
        PMutConfig conf = new Config(ctx);
        sm_vec<int, 128> vrand;
        for (int i = 0; i < dom->n_vars(); i++) {
            sm_vec<int> &st = SetVAL[i];
            int tmplval = templ.values()[i];
            if (tmplval != -1) {
                conf->set(i, tmplval);
            } else if (!st.empty()) {
                auto it = Rand.get(st);
                conf->set(i, *it);
                unordered_erase(st, it);
                if (st.empty())
                    n_finished++;
            } else {
                conf->set(i, Rand.get(dom->n_values(i)));
                vrand.push_back(i);
            }
        }
        if (shash.insert(conf->hash()).second) {
            if (z3 != nullptr) (*z3)->add(!conf->to_expr());
            ret.emplace_back(move(conf));
        } else if (use_solver) {
            for (int id : vrand) conf->set(id, -1);

            if (z3 == nullptr) {
                z3 = std::make_unique<Z3Scope>(ctx->zscope());
                for (const auto &c : hit_configs()) (*z3)->add(!c->to_expr());
                for (const auto &c : miss_configs()) (*z3)->add(!c->to_expr());
                for (const auto &c : new_configs_) (*z3)->add(!c->to_expr());
                for (const auto &c : ret) (*z3)->add(!c->to_expr());
                //(*z3)->set("timeout", 600u * 1000u); // 600s
            }

            expr e = conf->to_expr();
            if ((*z3)->check(1, &e) == z3::sat) {
                z3::model m = (*z3)->get_model();
                for (int id : vrand)
                    conf->set(id, dom->var(id)->val_id_of(m.get_const_interp(dom->var(id)->zvar().decl())));

                VLOG(100, "Found cex for tpl using solver ({} conds): \n", shash.size()) << templ << '\n' << *conf;
                bool insert_res = shash.insert(conf->hash(true /*recalc hash*/)).second;
                CHECK(insert_res) << *conf;
                (*z3)->add(!conf->to_expr());
                ret.emplace_back(move(conf));
            } else {
                VLOG(100, "Can't find cex for tpl: ") << templ;
            }
        }
    }

    VLOG_BLOCK(100, {
        log << "CEX:\n";
        for (const auto &c : ret) log << *c << '\n';
    });

    return ret;
}

std::ostream &CNode::serialize(std::ostream &out, bool &last_is_char) const {
    if (is_leaf()) {
        out << "MH"[leaf_value()];
        last_is_char = true;
    } else {
        if (last_is_char) out << ' ', last_is_char = false;
        out << splitvar << ' ';
        for (const auto &c : childs) c->serialize(out, last_is_char);
    }
    return out;
}

std::istream &CNode::deserialize(std::istream &inp) {
    if (inp >> splitvar) {
        CHECK(0 <= splitvar && splitvar < dom()->n_vars());
        childs.resize(dom()->n_values(splitvar));
        for (int i = 0; i < sz(childs); ++i) {
            childs[i] = new CNode(tree, this, i, {hit_configs(), miss_configs()});
            childs[i]->deserialize(inp);
        }
    } else if (char c; inp.clear(), inp >> c) {
        CHECK(c == 'H' || c == 'M');
        boost::sub_range<vec<PConfig>> empty = {tree->configs_[0].begin(), tree->configs_[0].begin()};
        configs_[0] = configs_[1] = empty;
        configs_[c == 'H'] = tree->configs_[0];
    } else {
        CHECK(0);
    }
    return inp;
}

void CNode::build_interpreter(vec<int> &dat) const {
    if (is_leaf()) {
        dat.push_back(leaf_value() ? CTree::kInterpretHit : CTree::kInterpretMiss);
    } else {
        dat.push_back(splitvar);
        CHECK_EQ(childs.size(), dom(splitvar)->n_values());
        int beg = sz(dat), nchilds = sz(childs);
        dat.resize(beg + nchilds - 1);
        for (int i = 0; i < nchilds; ++i) {
            childs[i]->build_interpreter(dat);
            if (i < nchilds - 1)
                dat[beg + i] = sz(dat);
        }
    }
}

}