#define BOOST_TEST_MODULE
#include <boost/test/unit_test.hpp>
#include <util/test_macros.hpp>
#include <stdlib.h>
#include <vector>
#include <string>
#include <functional>
#include <random>
#include <cfenv>
#include <cmath>

#include <ml_data/ml_data.hpp>
#include <optimization/optimization_interface.hpp>
#include <optimization/utils.hpp>
#include <toolkits/supervised_learning/logistic_regression.hpp>
#include <toolkits/supervised_learning/logistic_regression_opt_interface.hpp>
#include <sframe/testing_utils.hpp>

using namespace turi;
using namespace turi::supervised;

void run_logistic_regression_test(std::map<std::string, flexible_type> opts) {

  size_t examples = opts.at("examples");
  size_t features = opts.at("features");
  std::string target_column_name = "target";

  // Answers
  // -----------------------------------------------------------------------
  DenseVector coefs(features+1);
  coefs.randn();

  // Feature names
  std::vector<std::string> feature_names;
  std::vector<flex_type_enum> feature_types;
  for(size_t i=0; i < features; i++){
    feature_names.push_back(std::to_string(i));
    feature_types.push_back(flex_type_enum::FLOAT);
  }

  // Data
  std::vector<std::vector<flexible_type>> y_data;
  std::vector<std::vector<flexible_type>> X_data;
  for(size_t i=0; i < examples; i++){
    DenseVector x(features);
    x.randn();
    std::vector<flexible_type> x_tmp;
    for(size_t k=0; k < features; k++){
      x_tmp.push_back(x(k));
    }

    // Compute the prediction for this
    double t = dot(x, coefs.subvec(0, features-1)) + coefs(features);
    t = 1.0/(1.0+exp(-1.0*t));
    int c = turi::random::bernoulli(t);
    if (i == 0) c = 0; // Make sure category 0 is category 0 (for testing)
    std::vector<flexible_type> y_tmp;
    y_tmp.push_back(c);

    X_data.push_back(x_tmp);
    y_data.push_back(y_tmp);
  }

  // Options
  std::map<std::string, flexible_type> options = {
    {"convergence_threshold", 1e-2},
    {"step_size", 1.0},
    {"lbfgs_memory_level", 3},
    {"max_iterations", 10},
    {"l1_penalty", 0.0},
    {"l2_penalty", 1e-2}
  };

  // Make the data
  sframe X = make_testing_sframe(feature_names, feature_types, X_data);
  sframe y = make_testing_sframe({"target"}, {flex_type_enum::STRING}, y_data);
  std::shared_ptr<logistic_regression> model;
  model.reset(new logistic_regression);
  model->init(X,y);
  model->init_options(options);
  model->train();

  // Construct the ml_data
  ml_data data = model->construct_ml_data_using_current_metadata(X, y);

  // Check coefficients & options
  // ----------------------------------------------------------------------
  DenseVector _coefs(features+1);
  model->get_coefficients(_coefs);
  TS_ASSERT(_coefs.size() == features + 1);

  std::map<std::string, flexible_type> _options;
  _options = model->get_current_options();
  for (auto& kvp: options){
    TS_ASSERT(_options[kvp.first] == kvp.second);
  }
  TS_ASSERT(model->is_trained() == true);

  // Check predictions
  // ----------------------------------------------------------------------
  std::vector<flexible_type> pred_margin;
  std::shared_ptr<sarray<flexible_type>> _pred_margin
    = model->predict(data, "margin");
  std::vector<flexible_type> pred_class;
  std::shared_ptr<sarray<flexible_type>> _pred_class
    = model->predict(data, "class");
  std::vector<flexible_type> pred_prob;
  std::shared_ptr<sarray<flexible_type>> _pred_prob
    = model->predict(data, "probability");

  // Save predictions made by the model
  auto reader = _pred_margin->get_reader();
  reader->read_rows(0, examples, pred_margin);
  reader = _pred_class->get_reader();
  reader->read_rows(0, examples, pred_class);
  reader = _pred_prob->get_reader();
  reader->read_rows(0, examples, pred_prob);

  // Check that the predictions made by the model are right!
  for(size_t i=0; i < examples; i++){
    DenseVector x(features + 1);
    for(size_t k=0; k < features; k++){
      x(k) = X_data[i][k];
    }
    x(features) = 1;
    double t = arma::dot(x, _coefs);
    double p = pred_margin[i];
    TS_ASSERT(abs(p - t) < 1e-5);
    t = 1.0/(1.0+exp(-1.0*t));
    p = pred_prob[i];
    TS_ASSERT(abs(p - t) < 1e-5);
    int c = t >= 0.5;
    TS_ASSERT_EQUALS(pred_class[i], std::to_string(c));
  }

  // Check save and load
  // ----------------------------------------------------------------------
  dir_archive archive_write;
  archive_write.open_directory_for_write("regr_logistic_regression_tests");
  turi::oarchive oarc(archive_write);
  oarc << *model;
  archive_write.close();

  // Load it
  dir_archive archive_read;
  archive_read.open_directory_for_read("regr_logistic_regression_tests");
  turi::iarchive iarc(archive_read);
  iarc >> *model;

  // Check coefficients after saving and loading.
  // ----------------------------------------------------------------------
  DenseVector _coefs_after_load(features+1);
  model->get_coefficients(_coefs_after_load);
  TS_ASSERT(_coefs_after_load.size() == features + 1);
  TS_ASSERT(arma::approx_equal(_coefs_after_load, _coefs,"absdiff", 1e-5));
  _options = model->get_current_options();
  for (auto& kvp: options){
    TS_ASSERT(_options[kvp.first] == kvp.second);
  }
  TS_ASSERT(model->is_trained() == true);

  // Check coefficients after saving and loading.
  // ----------------------------------------------------------------------
  _pred_margin = model->predict(data, "margin");
  _pred_class = model->predict(data, "class");
  _pred_prob = model->predict(data, "probability");
  reader = _pred_margin->get_reader();
  reader->read_rows(0, examples, pred_margin);
  reader = _pred_class->get_reader();
  reader->read_rows(0, examples, pred_class);
  reader = _pred_prob->get_reader();
  reader->read_rows(0, examples, pred_prob);

  // Check that the predictions made by the model are right!
  for(size_t i=0; i < examples; i++){
    DenseVector x(features + 1);
    for(size_t k=0; k < features; k++){
      x(k) = X_data[i][k];
    }
    x(features) = 1;
    double t = arma::dot(x, _coefs);
    double p = pred_margin[i];
    TS_ASSERT(abs(p - t) < 1e-5);

    t = 1.0/(1.0+exp(-1.0*t));
    p = pred_prob[i];
    TS_ASSERT(abs(p - t) < 1e-5);
    int c = t > 0.5;
    TS_ASSERT_EQUALS(pred_class[i], std::to_string(c));
  }

  model->get_coefficients(_coefs);
  TS_ASSERT(_coefs.size() == features + 1);
  model.reset();

  // Check that we can train a model when providing a validation set
  model.reset(new logistic_regression);
  logprogress_stream << "Training with a validation set" << std::endl;
  model->init(X, y, X, y);
  model->init_options(options);
  model->train();

}

/**
 *  Check logistic regression
*/
struct logistic_regression_test  {

  public:

  void test_logistic_regression_basic_2d() {
    std::map<std::string, flexible_type> opts = {
      {"examples", 100},
      {"features", 1}};
    run_logistic_regression_test(opts);
  }

  void test_logistic_regression_small() {
    std::map<std::string, flexible_type> opts = {
      {"examples", 1000},
      {"features", 10}};
    run_logistic_regression_test(opts);
  }

};

void run_logistic_regression_opt_interface_test(std::map<std::string,
    flexible_type> opts) {

  size_t examples = opts.at("examples");
  size_t features = opts.at("features");
  std::string target_column_name = "target";
  std::vector<std::string> column_names = {"user", "item"};

  // Answers
  // -----------------------------------------------------------------------
  DenseVector coefs(features+1);
  coefs.randn();

  // Feature names
  std::vector<std::string> feature_names;
  std::vector<flex_type_enum> feature_types;
  for(size_t i=0; i < features; i++){
    feature_names.push_back(std::to_string(i));
    feature_types.push_back(flex_type_enum::FLOAT);
  }

  // Data
  std::vector<std::vector<flexible_type>> y_data;
  std::vector<std::vector<flexible_type>> X_data;
  for(size_t i=0; i < examples; i++){
    DenseVector x(features);
    x.randn();
    std::vector<flexible_type> x_tmp;
    for(size_t k=0; k < features; k++){
      x_tmp.push_back(x(k));
    }

    // Compute the prediction for this
    double t = dot(x, coefs.subvec(0, features-1)) + coefs(features);
    t = 1.0/(1.0+exp(-1.0*t));
    int c = turi::random::bernoulli(t);
    std::vector<flexible_type> y_tmp;
    y_tmp.push_back(c);

    X_data.push_back(x_tmp);
    y_data.push_back(y_tmp);
  }

  // Options
  std::map<std::string, flexible_type> options = {
    {"convergence_threshold", 1e-2},
    {"step_size", 1.0},
    {"lbfgs_memory_level", 3},
    {"mini_batch_size", 1},
    {"max_iterations", 10},
    {"solver", "auto"},
  };

  // Construct the ml_data
  // Make the data
  sframe X = make_testing_sframe(feature_names, feature_types, X_data);
  sframe y = make_testing_sframe({"target"}, {flex_type_enum::STRING}, y_data);
  std::shared_ptr<logistic_regression> model;
  model.reset(new logistic_regression);
  model->init(X, y);

  // Construct the ml_data
  ml_data data = model->construct_ml_data_using_current_metadata(X, y);
  ml_data valid_data;

  std::shared_ptr<logistic_regression_opt_interface> lr_interface;
  lr_interface.reset(new logistic_regression_opt_interface(data, valid_data, *model));

  // Check examples & variables.
  TS_ASSERT(lr_interface->num_variables() == features + 1);
  TS_ASSERT(lr_interface->num_examples() == examples);

  size_t variables = lr_interface->num_variables();
  for(size_t i=0; i < 10; i++){

    DenseVector point(variables);
    point.randn();

    // Check gradients, functions and hessians.
    DenseVector gradient(variables);
    double func_value;
    DenseMatrix hessian(variables, variables);

    func_value = lr_interface->compute_function_value(point);
    lr_interface->compute_gradient(point, gradient);
    lr_interface->compute_hessian(point, hessian);
    TS_ASSERT(check_gradient(*lr_interface, point, gradient));
    if( variables <= 2){
      TS_ASSERT(check_hessian(*lr_interface, point, hessian));
    }

    // Check first order & second order computations
    DenseVector _gradient(variables);
    double _func_value;
    DenseMatrix _hessian(variables, variables);

    lr_interface->compute_first_order_statistics(point, _gradient,
      _func_value);
    TS_ASSERT(abs(func_value - _func_value) < 1e-5);
    TS_ASSERT(arma::approx_equal(gradient, _gradient,"absdiff", 1e-10));
    lr_interface->compute_second_order_statistics(point, _hessian, _gradient,
      _func_value);
    TS_ASSERT(abs(func_value - _func_value) < 1e-5);
    TS_ASSERT(arma::approx_equal(gradient, _gradient,"absdiff", 1e-10));
    TS_ASSERT(arma::approx_equal(hessian, _hessian,"absdiff", 1e-10));

  }

  model.reset();
  lr_interface.reset();
}

/**
 *  Check logistic regression opt interface
*/
struct logistic_regression_opt_interface_test  {

  public:

  void test_logistic_regression_opt_interface_basic_2d() {
    std::map<std::string, flexible_type> opts = {
      {"examples", 100},
      {"features", 1}};
    run_logistic_regression_opt_interface_test(opts);
  }

  void test_logistic_regression_opt_interface_small() {
    std::map<std::string, flexible_type> opts = {
      {"examples", 1000},
      {"features", 10}};
    run_logistic_regression_opt_interface_test(opts);
  }

};

BOOST_FIXTURE_TEST_SUITE(_logistic_regression_test, logistic_regression_test)
BOOST_AUTO_TEST_CASE(test_logistic_regression_basic_2d) {
  logistic_regression_test::test_logistic_regression_basic_2d();
}
BOOST_AUTO_TEST_CASE(test_logistic_regression_small) {
  logistic_regression_test::test_logistic_regression_small();
}
BOOST_AUTO_TEST_SUITE_END()
BOOST_FIXTURE_TEST_SUITE(_logistic_regression_opt_interface_test, logistic_regression_opt_interface_test)
BOOST_AUTO_TEST_CASE(test_logistic_regression_opt_interface_basic_2d) {
  logistic_regression_opt_interface_test::test_logistic_regression_opt_interface_basic_2d();
}
BOOST_AUTO_TEST_CASE(test_logistic_regression_opt_interface_small) {
  logistic_regression_opt_interface_test::test_logistic_regression_opt_interface_small();
}
BOOST_AUTO_TEST_SUITE_END()


#include <boost/math/special_functions/ellint_rd.hpp>


#include "drake/solvers/system_identification.h"

#include <random>  // Used only with deterministic seeds!

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

#include "drake/common/polynomial.h"
#include "drake/common/trig_poly.h"

namespace drake {
namespace solvers {
namespace {

typedef SystemIdentification<double> SID;

GTEST_TEST(SystemIdentificationTest, LumpedSingle) {
  Polynomiald x = Polynomiald("x");
  Polynomiald y = Polynomiald("y");
  Polynomiald a = Polynomiald("a");
  Polynomiald b = Polynomiald("b");
  Polynomiald c = Polynomiald("c");

  /* From the SystemIdentification.h doxygen */
  Polynomiald input = (a * x) + (b * x) + (a * c * y) + (a * c * y * y);

  std::set<Polynomiald::VarType> parameters = {
    a.GetSimpleVariable(),
    b.GetSimpleVariable(),
    c.GetSimpleVariable()};
  SID::LumpingMapType lump_map =
      SID::GetLumpedParametersFromPolynomial(input, parameters);
  EXPECT_EQ(lump_map.size(), 2u);
  EXPECT_EQ(lump_map.count(a + b), 1u);
  EXPECT_EQ(lump_map.count(a * c), 1u);
}

GTEST_TEST(SystemIdentificationTest, LumpedMulti) {
  Polynomiald x = Polynomiald("x");
  Polynomiald y = Polynomiald("y");
  Polynomiald a = Polynomiald("a");
  Polynomiald b = Polynomiald("b");
  Polynomiald c = Polynomiald("c");

  std::vector<Polynomiald> input = {
    (a * x) + (b * x) + (a * c * y),
    (a * c * y * y),
    2 * a,
    a};

  std::set<Polynomiald::VarType> parameters = {
    a.GetSimpleVariable(),
    b.GetSimpleVariable(),
    c.GetSimpleVariable()};
  SID::LumpingMapType lump_map =
      SID::GetLumpedParametersFromPolynomials(input, parameters);

  // Note that we expect that 'a' and '2*a' will collapse to one lumped param.
  EXPECT_EQ(lump_map.size(), 3u);
  EXPECT_EQ(lump_map.count(a), 1u);
  EXPECT_EQ(lump_map.count(a + b), 1u);
  EXPECT_EQ(lump_map.count(a * c), 1u);

  // TODO(ggould-tri) The above code should be able to be more cleanly written
  // using gmock as something like:
  //
  // EXPECT_THAT(lump_map, ElementsAre(std::make_pair(a, _),
  //                                   std::make_pair((a + b), _),
  //                                   std::make_pair((a * c), _)));
  //
  // but the author could not get this working.
}

GTEST_TEST(SystemIdentificationTest, LumpedParameterRewrite) {
  Polynomiald x = Polynomiald("x");
  Polynomiald y = Polynomiald("y");
  Polynomiald a = Polynomiald("a");
  Polynomiald b = Polynomiald("b");
  Polynomiald c = Polynomiald("c");

  std::vector<Polynomiald> input = {
    (a * x) + (b * x) + (3 * a * c * y),
    (a * x) + (2 * b * x) + (3 * a * c * y),
    (a * c * y * y),
    2 * a,
    a};

  std::set<Polynomiald::VarType> parameters = {
    a.GetSimpleVariable(),
    b.GetSimpleVariable(),
    c.GetSimpleVariable()};
  SID::LumpingMapType lump_map =
      SID::GetLumpedParametersFromPolynomials(input, parameters);

  // A point for testing numeric stability.
  std::map<Polynomiald::VarType, double> eval_point = {
    {x.GetSimpleVariable(), 1},
    {y.GetSimpleVariable(), 2},
    {a.GetSimpleVariable(), 3},
    {b.GetSimpleVariable(), 5},
    {c.GetSimpleVariable(), 7},
  };
  // Compute the value of each lumped parameter at eval_point; store those
  // values into the eval_point (so that now it provides both lumped and
  // un-lumped values and can be used to evaluate either the original or the
  // rewritten polynomial).
  for (const auto& poly_var_pair : lump_map) {
    eval_point[poly_var_pair.second] =
        poly_var_pair.first.EvaluateMultivariate(eval_point);
  }

  for (const Polynomiald& poly : input) {
    Polynomiald rewritten =
        SID::RewritePolynomialWithLumpedParameters(poly, lump_map);

    // No non-lumped parameters should remain in rewritten.
    EXPECT_EQ(rewritten.GetVariables().count(a.GetSimpleVariable()), 0u);
    EXPECT_EQ(rewritten.GetVariables().count(b.GetSimpleVariable()), 0u);
    EXPECT_EQ(rewritten.GetVariables().count(c.GetSimpleVariable()), 0u);

    // Rewritten has the same or smaller number of variables and terms.
    EXPECT_LE(rewritten.GetVariables().size(), poly.GetVariables().size());
    EXPECT_LE(rewritten.GetMonomials().size(), poly.GetMonomials().size());

    // Rewriting in terms of lumped parameters should never change the
    // actual value of a polynomial at a particular point.
    EXPECT_EQ(poly.EvaluateMultivariate(eval_point),
              rewritten.EvaluateMultivariate(eval_point));

    // TODO(ggould-tri) The above tests do not ensure that the original and
    // rewritten polys are everywhere and always structurally identical, just
    // nearly always identical in their EvaluateMultivariate behaviour.
  }
}

GTEST_TEST(SystemIdentificationTest, BasicEstimateParameters) {
  const Polynomiald x = Polynomiald("x");
  const auto x_var = x.GetSimpleVariable();
  const Polynomiald y = Polynomiald("y");
  const auto y_var = y.GetSimpleVariable();
  const Polynomiald z = Polynomiald("z");
  const auto z_var = z.GetSimpleVariable();
  const Polynomiald a = Polynomiald("a");
  const auto a_var = a.GetSimpleVariable();
  const Polynomiald b = Polynomiald("b");
  const auto b_var = b.GetSimpleVariable();
  const Polynomiald c = Polynomiald("c");
  const auto c_var = c.GetSimpleVariable();

  /// Parameter estimation will try to make this Polynomial evaluate to zero:
  const Polynomiald poly = (a * x) + (b * x * x) + (c * y) - z;

  { // A very simple test case in which the error is zero.
    const std::vector<SID::PartialEvalType> sample_points {
      {{x_var, 1}, {y_var, 1}, {z_var, 3}},
      {{x_var, 1}, {y_var, 2}, {z_var, 4}},
      {{x_var, 2}, {y_var, 1}, {z_var, 7}},
      {{x_var, 2}, {y_var, 2}, {z_var, 8}}};

    const SID::PartialEvalType expected_params {
      {a_var, 1}, {b_var, 1}, {c_var, 1}};

    SID::PartialEvalType estimated_params;
    double error;
    std::tie(estimated_params, error) =
        SID::EstimateParameters(VectorXPoly::Constant(1, poly),
                                sample_points);

    EXPECT_LT(error, 1e-5);
    EXPECT_EQ(estimated_params.size(), 3u);
    for (const auto& var : {a_var, b_var, c_var}) {
      // `9 * error` here in case all of the RMS error was in a single term.
      EXPECT_NEAR(estimated_params[var], expected_params.at(var),
                  9 * error);
    }
  }

  { // Test with some error injected.
    const std::vector<SID::PartialEvalType> sample_points {
      {{x_var, 1}, {y_var, 1}, {z_var, 3.05}},
      {{x_var, 1}, {y_var, 2}, {z_var, 3.95}},
      {{x_var, 2}, {y_var, 1}, {z_var, 7.05}},
      {{x_var, 2}, {y_var, 2}, {z_var, 8.05}}};

    const SID::PartialEvalType expected_params {
      {a_var, 1}, {b_var, 1}, {c_var, 1}};

    SID::PartialEvalType estimated_params;
    double error;
    std::tie(estimated_params, error) =
        SID::EstimateParameters(VectorXPoly::Constant(1, poly),
                                sample_points);

    EXPECT_LT(error, 0.1);
    EXPECT_EQ(estimated_params.size(), 3u);
    for (const auto& var : {a_var, b_var, c_var}) {
      EXPECT_NEAR(estimated_params[var], expected_params.at(var), 4 * error);
    }
  }
}

/// Test to check parameter estimation for a basic spring-mass system.
///@{

struct State { double acceleration, velocity, position, force; };
static const double kMass = 1;
static const double kDamping = 0.1;
static const double kSpring = 2;
static const double kNoise = 0.01;
static const double kNoiseSeed = 1;

State AdvanceState(const State& previous, double input_force, double dt) {
  State next{};
  next.force = input_force;
  next.acceleration = (input_force -
                       previous.velocity * kDamping -
                       previous.position * kSpring) / kMass;
  next.velocity = previous.velocity +
      (dt * (previous.acceleration + next.acceleration) / 2);
  next.position = previous.position +
      (dt * (previous.velocity + next.velocity) / 2);
  return next;
}

std::vector<State> MakeTestData() {
  static const double kDt = 0.01;
  static const double kDuration1 = 1;
  static const double kInputForce1 = 1;
  static const double kDuration2 = 3;
  static const double kInputForce2 = 0;
  static const State kInitial {0, 0, 0, 0};

  std::vector<State> result { kInitial };
  State current = kInitial;
  double t = 0;
  while (t < kDuration1) {
    current = AdvanceState(current, kInputForce1, kDt);
    result.push_back(current);
    t += kDt;
  }
  while (t < kDuration1 + kDuration2) {
    current = AdvanceState(current, kInputForce2, kDt);
    result.push_back(current);
    t += kDt;
  }

  return result;
}

// TODO(ggould-tri) It is likely that much of the logic below will be
// boilerplate shared by all manipulator identification; it should eventually
// be pulled into a function of its own inside of system_identification.
GTEST_TEST(SystemIdentificationTest, SpringMassIdentification) {
  Polynomiald pos = Polynomiald("pos");
  auto pos_var = pos.GetSimpleVariable();
  Polynomiald velocity = Polynomiald("vel");
  auto velocity_var = velocity.GetSimpleVariable();
  Polynomiald acceleration = Polynomiald("acc");
  auto acceleration_var = acceleration.GetSimpleVariable();
  Polynomiald input_force = Polynomiald("f_in");
  auto input_force_var = input_force.GetSimpleVariable();
  Polynomiald mass = Polynomiald("m");
  auto mass_var = mass.GetSimpleVariable();
  Polynomiald damping = Polynomiald("b");
  auto damping_var = damping.GetSimpleVariable();
  Polynomiald spring = Polynomiald("k");
  auto spring_var = spring.GetSimpleVariable();

  // Code style violations here:
  // * Vector initializations use two statements on one line for clarity.
  // * The short names and upper/lower case here are conventional in the
  //   manipulator formulation.
  //
  // We write the manipulator as:
  //   H*vdot + C*v + g = B*u + f
  // Where f embodies any forces not appropriate to C.
  VectorXPoly v(1); v << velocity;
  VectorXPoly vdot(1); vdot << acceleration;
  VectorXPoly H(1); H << mass;
  VectorXPoly C(1); C << 0;
  VectorXPoly g(1); g << (spring * pos);
  VectorXPoly f(1); f << (velocity * damping);
  VectorXPoly B(1); B << 1;
  VectorXPoly u(1); u << input_force;

  const VectorXPoly manipulator_left = (H * vdot) + (C * v) + g;
  const VectorXPoly manipulator_right = (B * u) + f;

  std::default_random_engine noise_generator;
  noise_generator.seed(kNoiseSeed);
  std::uniform_real_distribution<double> noise_distribution(-kNoise, kNoise);
  auto noise = std::bind(noise_distribution, noise_generator);

  const std::vector<State> oracular_data = MakeTestData();
  std::vector<SID::PartialEvalType> measurements;
  for (const State& oracular_state : oracular_data) {
    SID::PartialEvalType measurement;
    measurement[pos_var] = oracular_state.position + noise();
    measurement[velocity_var] = oracular_state.velocity + noise();
    measurement[acceleration_var] = oracular_state.acceleration + noise();
    measurement[input_force_var] = oracular_state.force + noise();
    measurements.push_back(measurement);
  }

  SID::PartialEvalType estimated_params;
  double error;
  std::tie(estimated_params, error) =
      SID::EstimateParameters(manipulator_left - manipulator_right,
                              measurements);

  // Multiple layers of naive discrete-time numeric integration yields a very
  // high error value here, which almost all lands in the damping constant
  // because it is the smallest term in the equation of motion.
  //
  // The value for the error check here is an arbitrary empirical observation,
  // to catch changes that heavily regress accuracy.
  EXPECT_LT(error, 2e-2);

  EXPECT_EQ(estimated_params.size(), 3u);
  EXPECT_NEAR(estimated_params[mass_var], kMass, kNoise);
  EXPECT_NEAR(estimated_params[damping_var], kDamping,
              measurements.size() * error);
  EXPECT_NEAR(estimated_params[spring_var], kSpring, kNoise);
}

GTEST_TEST(SystemIdentificationTest, PendulaIdentification) {
  // Simulate two pendula that swing independently but are actuated with the
  // same torque.  The pendula have lengths l1 = 1, l2 = 2; their masses m1
  // and m2 are both 1.  Gravity is an earth-conventional -9.8.
  //
  // The comments about nomenclature from the previous test apply here as
  // well: Variable naming is conventional rather than style-conformant.

  const TrigPolyd theta1(Polynomiald("th", 1),
                         Polynomiald("s", 1), Polynomiald("c", 1));
  const TrigPolyd theta2(Polynomiald("th", 2),
                         Polynomiald("s", 2), Polynomiald("c", 2));
  VectorXTrigPoly q(2); q << theta1, theta2;

  const TrigPolyd theta1dot(Polynomiald("th.", 1),
                            Polynomiald("s.", 1), Polynomiald("c.", 1));
  const TrigPolyd theta2dot(Polynomiald("th.", 2),
                            Polynomiald("s.", 2), Polynomiald("c.", 2));
  VectorXTrigPoly qdot(2); qdot << theta1dot, theta2dot;

  const TrigPolyd theta1dotdot(Polynomiald("th..", 1),
                               Polynomiald("s..", 1), Polynomiald("c..", 1));
  const TrigPolyd theta2dotdot(Polynomiald("th..", 2),
                               Polynomiald("s..", 2), Polynomiald("c..", 2));
  VectorXTrigPoly qdotdot(2); qdotdot << theta1dotdot, theta2dotdot;

  const TrigPolyd l1(Polynomiald("l", 1));  //< Length of arm 1.
  const TrigPolyd l2(Polynomiald("l", 2));  //< Length of arm 2.
  const TrigPolyd m1(Polynomiald("m", 1));  //< Mass of arm 1.
  const TrigPolyd m2(Polynomiald("m", 2));  //< Mass of arm 2.

  const TrigPolyd gravity(Polynomiald("g"));  //< gravity
  const TrigPolyd tau(Polynomiald("tau"));  //< torque

  // The following matrices and vectors are the components of the Manipulator.
  Eigen::Matrix<TrigPolyd, 2, 2> H;  //< Inertia matrix.
  H << (m1 * l1 * l1), 0,
       0, (m2 * l2 * l2);
  Eigen::Matrix<TrigPolyd, 2, 2> C;  //< Coriolis matrix.
  C << 0, 0, 0, 0;
  Eigen::Matrix<TrigPolyd, 2, 1> g;  //< Field function (gravity).
  g << m1 * gravity * l1 * sin(theta1), m2 * gravity * l2 * sin(theta2);
  Eigen::Matrix<TrigPolyd, 2, 1> B;  //< Input transmission mapping.
  B << 1, 1;
  Eigen::Matrix<TrigPolyd, 2, 1> f;  //< Dissipative forces.
  f << 0, 0;

  VectorXTrigPoly u(1); u << tau;  //< Input signals.

  const VectorXTrigPoly manipulator_left = (H * qdotdot) + (C * qdot) + g;
  const VectorXTrigPoly manipulator_right = (B * u) + f;
  const VectorXTrigPoly to_estimate = manipulator_left - manipulator_right;

  // Create convenience variables for our q/qdot/u values.  Convenience vars
  // have an underscore prefix for slightly easier understandability.
  const TrigPolyd::VarType th1_var = theta1.poly().GetSimpleVariable();
  const TrigPolyd::VarType th2_var = theta2.poly().GetSimpleVariable();
  const TrigPolyd::VarType th1d_var =
      theta1dot.poly().GetSimpleVariable();
  const TrigPolyd::VarType th2d_var =
      theta2dot.poly().GetSimpleVariable();
  const TrigPolyd::VarType th1dd_var =
      theta1dotdot.poly().GetSimpleVariable();
  const TrigPolyd::VarType th2dd_var =
      theta2dotdot.poly().GetSimpleVariable();
  const TrigPolyd::VarType tau_var = tau.poly().GetSimpleVariable();

  const double kG = 9.8;
  const double kPi = 3.14159265;
  const double kPi2 = kPi / 2;

  const std::vector<typename SID::PartialEvalType> pendula_data = {
    {{tau_var, 0.},
     {th1_var, 0.}, {th1d_var, 0.}, {th1dd_var, 0.},
     {th2_var, 0.}, {th2d_var, 0.}, {th2dd_var, 0.}},
    {{tau_var, 0.},
     {th1_var, kPi2}, {th1d_var, 0.}, {th1dd_var, -kG},
     {th2_var, kPi2}, {th2d_var, 0.}, {th2dd_var, -0.25 * kG}},
    {{tau_var, 0.},
     {th1_var, -kPi}, {th1d_var, 0.}, {th1dd_var, 0.},
     {th2_var, -kPi}, {th2d_var, 0.}, {th2dd_var, 0.}},
    {{tau_var, 0.},
     {th1_var, -kPi2}, {th1d_var, 0.}, {th1dd_var, kG},
     {th2_var, -kPi2}, {th2d_var, 0.}, {th2dd_var, 0.25 * kG}},
    {{tau_var, 1.},
     {th1_var, 0.}, {th1d_var, 0.}, {th1dd_var, 1.},
     {th2_var, 0.}, {th2d_var, 0.}, {th2dd_var, 0.25}},
    {{tau_var, kG},
     {th1_var, kPi2}, {th1d_var, 0.}, {th1dd_var, 0.},
     {th2_var, kPi2}, {th2d_var, 0.}, {th2dd_var, 0.}},
    {{tau_var, 1.},
     {th1_var, -kPi}, {th1d_var, 0.}, {th1dd_var, 1.},
     {th2_var, -kPi}, {th2d_var, 0.}, {th2dd_var, 0.25}},
    {{tau_var, -kG},
     {th1_var, -kPi2}, {th1d_var, 0.}, {th1dd_var, 0.},
     {th2_var, -kPi2}, {th2d_var, 0.}, {th2dd_var, 0.}},
    };

  SID::SystemIdentificationResult result =
      SID::LumpedSystemIdentification(to_estimate, pendula_data);

  // Check result.rms_error.
  const double epsilon = 1e-5;  // Moderate, empirical epsilon for weak solvers.
  EXPECT_LT(result.rms_error, epsilon);
  const double max_per_term_error =
      result.rms_error *
      (result.lumped_parameters.size() * result.lumped_parameters.size());

  // Check result.lumped_parameters.
  Polynomiald mgl1 = (m1 * gravity * l1).poly();
  Polynomiald mgl2 = (m2 * gravity * l2).poly();
  Polynomiald mll1 = (m1 * l1 * l1).poly();
  Polynomiald mll2 = (m2 * l2 * l2).poly();
  EXPECT_EQ(result.lumped_parameters.size(), static_cast<size_t>(4));
  Polynomiald::VarType mgl1_var = result.lumped_parameters.at(mgl1);
  Polynomiald::VarType mgl2_var = result.lumped_parameters.at(mgl2);
  Polynomiald::VarType mll1_var = result.lumped_parameters.at(mll1);
  Polynomiald::VarType mll2_var = result.lumped_parameters.at(mll2);

  // Check result.lumped_polys.
  std::set<Polynomiald::VarType> expected_vars_1 = {
    mgl1_var, mll1_var, th1_var, th1dd_var, tau_var};
  EXPECT_EQ(result.lumped_polys[0].GetVariables(), expected_vars_1);
  std::set<Polynomiald::VarType> expected_vars_2 = {
    mgl2_var, mll2_var, th2_var, th2dd_var, tau_var};
  EXPECT_EQ(result.lumped_polys[1].GetVariables(), expected_vars_2);

  // Check result.lumped_parameter_values
  EXPECT_NEAR(result.lumped_parameter_values[mgl1_var], kG, max_per_term_error);
  EXPECT_NEAR(result.lumped_parameter_values[mgl2_var], kG, max_per_term_error);
  EXPECT_NEAR(result.lumped_parameter_values[mll1_var], 1, max_per_term_error);
  EXPECT_NEAR(result.lumped_parameter_values[mll2_var], 4, max_per_term_error);

  // Check result.partially_evaluated_polys.
  for (const auto& point : pendula_data) {
    EXPECT_NEAR(result.partially_evaluated_polys[0].EvaluateMultivariate(point),
                0, max_per_term_error);
    EXPECT_NEAR(result.partially_evaluated_polys[1].EvaluateMultivariate(point),
                0, max_per_term_error);
  }
}

///@}

}  // anonymous namespace
}  // namespace solvers
}  // namespace drake


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

  @copyright 2016 NumScale SAS

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

#if defined(DOXYGEN_ONLY)
namespace boost { namespace simd
{

 /*!

    @ingroup group-trigonometric
    Function object implementing tand capabilities

    tangent of the input in degrees: \f$\sin(\pi x/180)/\cos(\pi x/180) \f$.

    @par Semantic:

    For every parameter of floating type

    @code
    auto r = tand(x);
    @endcode

    As most other trigonometric function tand can be called with a second optional parameter
    which is a tag on speed and accuracy (see @ref cos for further details)

    @see tan, tanpi

  **/
  Value tand(Value const & v0);
} }
#endif

#include <boost/simd/function/scalar/tand.hpp>
#include <boost/simd/function/simd/tand.hpp>

#endif


// Boost.Geometry (aka GGL, Generic Geometry Library)

// Copyright (c) 2007-2012 Barend Gehrels, Amsterdam, the Netherlands.
// Copyright (c) 2008-2012 Bruno Lalande, Paris, France.
// Copyright (c) 2009-2012 Mateusz Loskot, London, UK.

// This file was modified by Oracle on 2020.
// Modifications copyright (c) 2020, Oracle and/or its affiliates.
// Contributed and/or modified by Adam Wulkiewicz, on behalf of Oracle

// Parts of Boost.Geometry are redesigned from Geodan's Geographic Library
// (geolib/GGL), copyright (c) 1995-2010 Geodan, Amsterdam, the Netherlands.

// Use, modification and distribution is subject to the Boost Software License,
// Version 1.0. (See accompanying file LICENSE_1_0.txt or copy at
// http://www.boost.org/LICENSE_1_0.txt)

#ifndef BOOST_GEOMETRY_EXTENSIONS_NSPHERE_ALGORITHMS_AREA_HPP
#define BOOST_GEOMETRY_EXTENSIONS_NSPHERE_ALGORITHMS_AREA_HPP

#include <type_traits>

#include <boost/math/constants/constants.hpp>

#include <boost/geometry/algorithms/area.hpp>
#include <boost/geometry/extensions/nsphere/core/radius.hpp>
#include <boost/geometry/extensions/nsphere/core/tags.hpp>

namespace boost { namespace geometry
{

#ifndef DOXYGEN_NO_DETAIL
namespace detail { namespace area
{

template<typename C>
struct circle_area
{
    typedef typename coordinate_type<C>::type coordinate_type;

    // Returning the coordinate precision, but if integer, returning a double
    typedef std::conditional_t
        <
            std::is_integral<coordinate_type>::value,
            double,
            coordinate_type
        > return_type;

    template <typename S>
    static inline return_type apply(C const& c, S const&)
    {
        // Currently only works for Cartesian circles
        // Todo: use strategy
        // Todo: use concept
        assert_dimension<C, 2>();

        return_type r = get_radius<0>(c);
        r *= r * boost::math::constants::pi<return_type>();
        return r;
    }
};

}} // namespace detail::area

#endif // DOXYGEN_NO_DETAIL

#ifndef DOXYGEN_NO_DISPATCH
namespace dispatch
{

template <typename Geometry>
struct area<Geometry, nsphere_tag>
    : detail::area::circle_area<Geometry>
{};

} // namespace dispatch
#endif // DOXYGEN_NO_DISPATCH

}} // namespace boost::geometry

#endif // BOOST_GEOMETRY_EXTENSIONS_NSPHERE_ALGORITHMS_AREA_HPP


#include <scitbx/array_family/simple_io.h>
#include <scitbx/array_family/initialiser.h>
#include <scitbx/array_family/ref_reductions.h>
#include <scitbx/array_family/tiny_algebra.h>
#include <scitbx/array_family/ref_algebra.h>
#include <scitbx/array_family/misc_functions.h>
#include <scitbx/matrix/special_matrices.h>
#include <scitbx/matrix/move.h>
#include <scitbx/error.h>
#include <scitbx/random.h>
#include <scitbx/array_family/simple_io.h>
#include <iostream>

#include <boost/random/mersenne_twister.hpp>
#include <boost/random/uniform_int.hpp>

#include <scitbx/matrix/householder.h>
#include <scitbx/matrix/tests.h>
#include <scitbx/matrix/tests/utils.h>

namespace af = scitbx::af;
using scitbx::fn::approx_equal;
using namespace scitbx::matrix;

#if defined(__GNUC__) && __GNUC__ == 4 && __GNUC_MINOR__ == 1 \
 && defined(__i386__) && defined(__linux)
// avoid internal compiler error
#undef SCITBX_ASSERT
#define SCITBX_ASSERT(cond)
#endif

double tol = 1e-12;

struct test_case
{
  matrix_t a0;
  double thresh;

  test_case(int m, int n, double ratio_threshold=10)
    : a0(dim(m,n)), thresh(ratio_threshold)
  {}

  void check_qr(bool thin_q) {
    matrix_const_ref_t a0_ = a0.const_ref();

    matrix_t a = a0.deep_copy();
    matrix_ref_t a_ = a.ref();

    int m = a_.n_rows(), n=a_.n_columns();

    householder::qr_decomposition<double> qr(a_);

    // Get R
    matrix_t r(dim(thin_q ? std::min(m,n) : m, n));
    matrix_ref_t r_ = r.ref();
    for (int i=0; i<std::min(m,n); ++i)
    for (int j=i; j<n; ++j) {
      r_(i,j) = a_(i,j);
    }

    // Accumulate Q out-of-place
    matrix_t q = qr.q(thin_q);
    matrix_const_ref_t q_ = q.const_ref();
    SCITBX_ASSERT( normality_ratio(q_) < thresh );

    // Accumulate Q in-place if it makes sense
    if (thin_q && m >= n) {
      qr.accumulate_q_in_place();
      SCITBX_ASSERT( equality_ratio(q_, a_) < thresh );
    }

    // Check A = QR
    matrix_t q_r = af::matrix_multiply(q_, r_);
    matrix_const_ref_t q_r_ = q_r.const_ref();
    SCITBX_ASSERT( equality_ratio(a0_, q_r_) < thresh );
  }

  void check_lq(bool thin_q) {
    matrix_const_ref_t a0_ = a0.const_ref();

    matrix_t a = a0.deep_copy();
    matrix_ref_t a_ = a.ref();

    int m = a_.n_rows(), n=a_.n_columns();

    householder::lq_decomposition<double> lq(a_);

    // Get L
    matrix_t l(dim(m, thin_q ? std::min(m,n) : n));
    matrix_ref_t l_ = l.ref();
    for (int j=0; j<std::min(m,n); ++j) {
    for (int i=j; i<m; ++i)
      l_(i,j) = a_(i,j);
    }

    // Accumulate Q out-of-place
    matrix_t q = lq.q(thin_q);
    matrix_const_ref_t q_ = q.const_ref();
    SCITBX_ASSERT( normality_ratio(q_) < thresh );

    // Accumulate Q in-place if it makes sense
    if (thin_q && m <= n) {
      lq.accumulate_q_in_place();
      SCITBX_ASSERT( equality_ratio(q_, a_) < thresh );
    }

    // Check A = LQ
    matrix_t l_q = af::matrix_multiply(l_, q_);
    matrix_const_ref_t l_q_ = l_q.const_ref();
    SCITBX_ASSERT( equality_ratio(a0_, l_q_) < thresh );
  }

  void check_bidiagonalisation(bool thin) {
    matrix_const_ref_t a0_ = a0.const_ref();

    matrix_t a = a0.deep_copy();
    matrix_ref_t a_ = a.ref();

    int m = a_.n_rows(), n=a_.n_columns();

    householder::bidiagonalisation<double> bidiag(a_);

    // Get B
    matrix_t b(dim(thin ? std::min(m,n) : m, thin ? std::min(m,n) : n));
    matrix_ref_t b_ = b.ref();
    if (m >=n) copy_upper_bidiagonal(b_, a_);
    else copy_lower_bidiagonal(b_, a_);

    matrix_t u = bidiag.u(thin);
    matrix_ref_t u_ = u.ref();

    matrix_t v = bidiag.v(thin);
    matrix_ref_t v_ = v.ref();

    // Check that U and V are orthogonal
    SCITBX_ASSERT( normality_ratio(u_) < thresh );
    SCITBX_ASSERT( normality_ratio(v_) < thresh );

    // Check that A = U B V^T
    matrix_t u_b_vt = product_U_M_VT(u_, b_, v_);
    matrix_const_ref_t u_b_vt_ = u_b_vt.const_ref();
    SCITBX_ASSERT( equality_ratio(a0_, u_b_vt_) < thresh );
  }
};

struct lotkin_test_case : test_case
{
  lotkin_test_case(int m, int n) : test_case(m,n)
  {
    matrix_ref_t a_ = this->a0.ref();
    for (int i=0; i<m; ++i) for (int j=0; j<n; ++j) {
      a_(i,j) = i > 0 ? 1./(i+j+1) : 1;
    }
  }
};

struct graded_test_case : test_case
{
  householder::random_normal_matrix_generator<
    double, scitbx::boost_random::mt19937> gen;

  graded_test_case(int n, double x)
    : test_case(n,n), gen(n,n)
  {
    vec_t d(n), f(n-1);
    for (int i=0; i<n; ++i) {
      d[i] = std::pow(x, i);
      if (i < n-1) f[i] = d[i];
    }
    matrix_t u = gen.normal_matrix();
    matrix_t v = gen.normal_matrix();
    matrix_ref_t vt_ = v.ref();
    vt_.transpose_in_place();
    matrix_t b = upper_bidiagonal(d.ref(), f.ref());
    a0 = af::matrix_multiply(u.ref(), b.ref());
    a0 = af::matrix_multiply(a0.ref(), vt_);
  }
};

void exercise_householder_zeroing_vector() {
  // Householder zeroing vector (1)
  {
    af::tiny<double, 4> x;
    af::init(x) = 3, 1, 5, 1;
    householder::reflection<double> p(x.ref());
    af::tiny<double, 3> expected;
    af::init(expected) = -1, -5, -1;
    expected /= 3.;
    af::ref<double> obtained(&p.v[0], 3);
    SCITBX_ASSERT(approx_equal(p.beta, 0.5, tol))(p.beta);
    SCITBX_ASSERT(expected.ref().all_approx_equal(obtained, tol));
    af::ref<double> overwritten(&x[1], 3);
    SCITBX_ASSERT(expected.ref().all_approx_equal(overwritten, tol));
  }

  // Householder zeroing vector (2)
  {
    af::tiny<double, 4> x;
    af::init(x) = -3, 1, 5, 1;
    householder::reflection<double> p(x.ref());
    af::tiny<double, 3> expected;
    af::init(expected) = -1, -5, -1;
    expected /= 9.;
    af::ref<double> obtained(&p.v[0], 3);
    SCITBX_ASSERT(approx_equal(p.beta, 3./2, tol))(p.beta);
    SCITBX_ASSERT(expected.ref().all_approx_equal(obtained, tol));
    af::ref<double> overwritten(&x[1], 3);
    SCITBX_ASSERT(expected.ref().all_approx_equal(overwritten, tol));
  }

  // Householder zeroing vector in-place
  {
    householder::reflection<double> p(4);
    af::ref<double> v(&p.v[0], 4);
    af::init(v) = 3, 1, 5, 1;
    af::tiny<double, 3> expected;
    af::init(expected) = -1, -5, -1;
    expected /= 3.;
    p.zero_vector(4);
    af::ref<double> obtained(&p.v[0], 3);
    SCITBX_ASSERT(approx_equal(p.beta, 0.5, tol))(p.beta);
    SCITBX_ASSERT(expected.ref().all_approx_equal(obtained, tol));
  }

  // Householder zeroing matrix columns or rows
  {
    int const m=6, n=7;
    matrix_t a0(dim(m,n));
    af::init(a0) = 11, 12, 13, 14, 15, 16, 17,
                    21, 22,  1, 24, 25, 26, 27,
                    31, 32, -1, 34, 35, 36, 37,
                    41, 42,  2, 44, 45, 46, 47,
                    51, 52,  1, 54, 55, 56, 57,
                    61, 62,  3, 64, 65, 66, 67;

    matrix_t a = a0.deep_copy();
    householder::reflection<double> p(m, n,
                                      householder::applied_on_left_tag(),
                                      false);
    p.zero_vector(af::column_below(a.ref(), 1, 2));
    p.apply_on_left_to_lower_right_block(a.ref(), 1, 3);
    /* Mathematica:
      a = Table[10 i + j, {i, 6}, {j, 7}];
      a[[2 ;;, 3]] = {1, -1, 2, 1, 3};
      a // MatrixForm
      v = {0, 1, 1/3, -2/3, -1/3, -1};
      beta = 3/4;
      p = IdentityMatrix[6] - beta KroneckerProduct[v, v];
      b = p.a;
      b // MatrixForm
    */
    matrix_t expected(dim(m,n));
    af::init(expected) =  11, 12, 13   ,  14,  15   ,  16   ,  17   ,
                          21, 22,  4   ,  81, 165./2,  84   , 171./2,
                          31, 32,  1./3,  53, 325./6, 166./3, 113./2,
                          41, 42, -2./3,   6,  20./3,  22./3,   8   ,
                          51, 52, -1./3,  35, 215./6, 110./3,  75./2,
                          61, 62, -1.  ,   7,  15./2,   8   ,  17./2;
    SCITBX_ASSERT(expected.all_approx_equal(a, tol));

    matrix_t a_t = af::matrix_transpose(a0.ref());

    householder::reflection<double> q(n, m,
                                      householder::applied_on_right_tag(),
                                      false);
    q.zero_vector(af::row_right_of(a_t.ref(), 2, 1));
    q.apply_on_right_to_lower_right_block(a_t.ref(), 3, 1);
    SCITBX_ASSERT(matrix_transpose(expected.ref()).all_approx_equal(a_t, tol));
  }
}

void exercise_householder_applied_to_symmetric_matrix() {
  double tol = 1e-12;
  scitbx::random::mersenne_twister rnd;
  for (int n=2; n<10; ++n) {
    vec_t x = rnd.random_double(n);
    householder::reflection<double> p(
      n, n, householder::applied_on_left_and_right_tag());
    p.zero_vector(x.ref(), false);
    matrix_t id_n = identity<double>(n);
    matrix_ref_t pp = id_n.ref();
    p.apply_on_left_to_lower_right_block(pp, 0, 0);

    symmetric_matrix_packed_u_t a(n);
    double *a_ = a.begin();
    for (int i=0; i<n; ++i) for (int j=i; j<n; ++j) *a_++ = rnd.random_double();
    af::const_ref<double> r(a.begin(), a.size());
    matrix_t a1 = packed_u_as_symmetric(r);

    matrix_t b1 = product_U_M_VT(pp, a1.const_ref(), pp);
    symmetric_matrix_packed_u_t b(symmetric_as_packed_u(b1.ref(), tol), n);

    p.apply_to_lower_right_block(a.ref(), 0);
    SCITBX_ASSERT(a.all_approx_equal(b, tol));
  }
}

void exercise_householder_accumulation() {
  boost::mt19937 urng;
  boost::uniform_int<> gen(0, 64);
  for (int m=3; m<=6; ++m) for (int n=3; n<=6; ++n) {
    // set up test case
    matrix_t a(dim(m,n), -1.);
    matrix_ref_t a_ = a.ref();
    vec_t beta(std::min(m-1, n));
    for (int j=0; j<std::min(m-1, n); ++j) {
      double v2 = 1;
      for (int i=j+1; i<m; ++i) {
        a(i,j) = gen(urng);
        v2 += a(i,j)*a(i,j);
      }
      beta[j] = 2./v2;
    }
    matrix_t at = af::matrix_transpose(a.ref());
    matrix_ref_t at_ = at.ref();

    // test accumulation of the thin Q if m >= n or the full Q otherwise

    // on the right
    {
      matrix_t q(dim(m, std::min(m,n)));
      matrix_ref_t q_ = q.ref();
      householder::reflection<double> h(m, n,
                                        householder::applied_on_left_tag(),
                                        true);
      h.accumulate_factored_form_in_columns(q_, a_, beta.ref());
      h.accumulate_in_place_factored_form_in_columns(a_, beta.ref());
      for (int i=0; i<q_.n_rows(); ++i)
      for (int j=0; j<q_.n_columns(); ++j) {
        approx_equal(q_(i,j), a_(i,j), 1e-15);
      }
    }

    // on the left
    {
      std::swap(m,n);
      int p = std::min(m,n); // <|workaround for gcc 3.3.4 compilation error
      matrix_t q(dim(p, n)); // <|
      matrix_ref_t q_ = q.ref();
      householder::reflection<double> h(m, n,
                                        householder::applied_on_right_tag(),
                                        true);
      h.accumulate_factored_form_in_rows(
        q_, at_, beta.ref(), householder::product_in_reverse_row_order);
      h.accumulate_in_place_factored_form_in_rows(at_, beta.ref());
      for (int i=0; i<q_.n_rows(); ++i)
      for (int j=0; j<q_.n_columns(); ++j) {
        approx_equal(q_(i,j), at_(i,j), 1e-15);
      }
    }
  }
}

void exercise_householder() {
  // Householder QR (Lotkin matrix: ill-conditioned)
  std::vector<lotkin_test_case> cases;

  cases.push_back( lotkin_test_case(3,2) );
  cases.push_back( lotkin_test_case(2,3) );
  cases.push_back( lotkin_test_case(3,3) );

  cases.push_back( lotkin_test_case(4,3) );
  cases.push_back( lotkin_test_case(3,4) );
  cases.push_back( lotkin_test_case(4,4) );

  cases.push_back( lotkin_test_case(5,3) );
  cases.push_back( lotkin_test_case(3,5) );

  for (int i=0; i<cases.size(); ++i) {
    lotkin_test_case &t = cases[i];
    t.check_qr(false); // full QR
    t.check_qr(true);  // thin QR
    t.check_lq(false); // full LQ
    t.check_lq(true);  // thin LQ
  }

  graded_test_case t(10, 1.e-10);
  t.check_qr(true);
}

void exercise_bidiagonalisation() {
  /* Householder bidiagonalisation (Golub and Van Loan example 5.4.2)
     The last diagonal entry in B is zero which makes it interesting.
     In the book, the authors used higher a precision than available on standard
     hardware: hence the difference between their result and those obtained
     here. But the SVD is still correct to machine precision as we assert it is.
  */
  {
    matrix_t a0(dim(4,3));
    af::init(a0) =  1,  2,  3,
                    4,  5,  6,
                    7,  8,  9,
                   10, 11, 12;
    matrix_const_ref_t a0_ = a0.const_ref();
    matrix_t a = a0.deep_copy();
    householder::bidiagonalisation<double> bidiag(a.ref());
    matrix_t b(dim(4,3));
    af::init(b) = 12.8840987267251, 21.876432827428,  0             ,
                   0              ,  2.246235240294, -0.613281332054,
                   0              ,  0             ,  0             ,
                   0              ,  0             ,  0             ;
    matrix_const_ref_t b_ = b.const_ref();
    matrix_t u = bidiag.u(false);
    matrix_ref_t u_ = u.ref();
    matrix_t v = bidiag.v(false);
    matrix_ref_t v_ = v.ref();
    SCITBX_ASSERT( normality_ratio(u_) < 10 );
    SCITBX_ASSERT( normality_ratio(v_) < 10 );
    matrix_t u_b_vt = product_U_M_VT(u_, b_, v_);
    matrix_const_ref_t u_b_vt_ = u_b_vt.const_ref();
    SCITBX_ASSERT( equality_ratio(u_b_vt_, a0_, 1e-13) < 10 );
  }
  {
    matrix_t a0(dim(3,3));
    af::init(a0) = 24.3104915623, 0.0          , 0.0                ,
                   26.4083512403, 1.26450969612, 0.0                ,
                   28.5062109182, 2.52901939223, 3.74525472711e-15 ;
    matrix_const_ref_t a0_ = a0.const_ref();
    matrix_t a = a0.deep_copy();
    matrix_ref_t a_ = a.ref();
    householder::bidiagonalisation<double> bidiag(a.ref());
    matrix_t u = bidiag.u();
    matrix_ref_t u_ = u.ref();
    matrix_t v = bidiag.v();
    matrix_ref_t v_ = v.ref();
    SCITBX_ASSERT( normality_ratio(u_) < 10 );
    SCITBX_ASSERT( normality_ratio(v_) < 10 );
    matrix_t b(dim(3, 3));
    matrix_ref_t b_ = b.ref();
    copy_upper_bidiagonal(b_, a_);
    matrix_t u_b_vt = product_U_M_VT(u_, b_, v_);
    matrix_const_ref_t u_b_vt_ = u_b_vt.const_ref();
    SCITBX_ASSERT( equality_ratio(u_b_vt_, a0_, 1e-13) < 10 );
  }

  // Householder bidiagonalisation (Lotkin matrix)
  {
    lotkin_test_case t(7,5);
    t.check_bidiagonalisation(true);
    t.check_bidiagonalisation(false);
  }
  {
    lotkin_test_case t(10,5);
    t.check_bidiagonalisation(true);
    t.check_bidiagonalisation(false);
  }
  {
    lotkin_test_case t(5,5);
    t.check_bidiagonalisation(true);
    t.check_bidiagonalisation(false);
  }
  {
    lotkin_test_case t(5,7);
    t.check_bidiagonalisation(true);
    t.check_bidiagonalisation(false);
  }
  {
    lotkin_test_case t(5,10);
    t.check_bidiagonalisation(true);
    t.check_bidiagonalisation(false);
  }
  {
    graded_test_case t(10, 1.e10);
    t.check_bidiagonalisation(false);
  }
}

int main() {
  exercise_householder_zeroing_vector();
  exercise_householder_applied_to_symmetric_matrix();
  exercise_householder_accumulation();
  exercise_householder();
  exercise_bidiagonalisation();
  std::cout << "OK\n";
  return 0;
}


#include <limits>
#include <vector>
#include <gtest/gtest.h>

#include <boost/random/mersenne_twister.hpp>
#include <boost/math/distributions.hpp>
#include <boost/math/tools/promotion.hpp>

#include <stan/math/rev.hpp>
#include <stan/math/prim/fun/sign.hpp>
#include <stan/math/prim/fun/fabs.hpp>
#include <stan/math/prim/fun/log1m.hpp>

template <typename T1, typename T2, typename T3, typename T4>
inline typename boost::math::tools::promote_args<T1, T2, T3, T4>::type
skew_de_ccdf_test(const T1& y, const T2& mu, const T3& sigma, const T4& tau) {
  using stan::math::log1m;
  using stan::math::log1m_exp;
  using std::exp;
  using std::log;

  if (y < mu) {
    return log1m_exp(log(tau) - 2 / sigma * (1 - tau) * (mu - y));
  } else {
    return log1m_exp(log1m((1 - tau) * exp(-2 / sigma * tau * (y - mu))));
  }
}

TEST(ProbDistributionsSkewedDoubleExponential,
     lccdf_computes_correct_gradients) {
  using stan::math::skew_double_exponential_lccdf;

  for (double ys : {-1.7, 0.2, 0.5, 0.9, 1.1, 3.2, 8.3}) {
    for (double mus : {-1.8, 0.1, 0.55, 0.89, 1.3, 4.2, 9.3}) {
      for (double sigmas : {0.1, 0.5, 1.1, 10.1}) {
        for (double taus : {0.01, 0.1, 0.5, 0.6, 0.9, 0.99}) {
          stan::math::var y = ys;
          stan::math::var mu = mus;
          stan::math::var sigma = sigmas;
          stan::math::var tau = taus;

          stan::math::var lp = skew_double_exponential_lccdf(y, mu, sigma, tau);
          std::vector<stan::math::var> theta;
          theta.push_back(y);
          theta.push_back(mu);
          theta.push_back(sigma);
          theta.push_back(tau);
          std::vector<double> grads;
          lp.grad(theta, grads);

          stan::math::var y_true = ys;
          stan::math::var mu_true = mus;
          stan::math::var sigma_true = sigmas;
          stan::math::var tau_true = taus;

          stan::math::var lp_test
              = skew_de_ccdf_test(y_true, mu_true, sigma_true, tau_true);
          std::vector<stan::math::var> theta_true;
          theta_true.push_back(y_true);
          theta_true.push_back(mu_true);
          theta_true.push_back(sigma_true);
          theta_true.push_back(tau_true);
          std::vector<double> grads_true;
          lp_test.grad(theta_true, grads_true);

          EXPECT_NEAR(grads_true[0], grads[0], 0.01);
          EXPECT_NEAR(grads_true[1], grads[1], 0.01);
          EXPECT_NEAR(grads_true[2], grads[2], 0.01);
          EXPECT_NEAR(grads_true[3], grads[3], 0.01);
        }
      }
    }
  }
}

TEST(ProbDistributionsSkewedDoubleExponential,
     lccdf_works_on_scalar_arguments) {
  using stan::math::skew_double_exponential_lccdf;

  for (double ys : {0.2, 0.9, 1.1, 3.2}) {
    for (double mus : {0.1, 1.3, 3.0}) {
      for (double sigmas : {0.1, 1.1, 3.2}) {
        for (double taus : {0.01, 0.1, 0.5, 0.9, 0.99}) {
          EXPECT_NEAR(skew_de_ccdf_test(ys, mus, sigmas, taus),
                      skew_double_exponential_lccdf(ys, mus, sigmas, taus),
                      0.001);
        }
      }
    }
  }
}

TEST(ProbDistributionsSkewedDoubleExponential,
     lccdf_works_on_vector_arguments) {
  using stan::math::skew_double_exponential_lccdf;

  std::vector<double> ys{0.2, 0.9, 1.1, 3.2};

  for (double mus : {0.1, 1.3, 3.0}) {
    for (double sigmas : {0.1, 1.1, 3.2}) {
      for (double taus : {0.01, 0.1, 0.5, 0.9, 0.99}) {
        double x = 0.0;
        for (double y : ys)
          x += skew_de_ccdf_test(y, mus, sigmas, taus);
        EXPECT_NEAR(x, skew_double_exponential_lccdf(ys, mus, sigmas, taus),
                    0.001);
      }
    }
  }
}

TEST(ProbDistributionsSkewedDoubleExponential,
     lccdf_works_on_vectorial_y_and_mu) {
  using stan::math::skew_double_exponential_lccdf;
  std::vector<double> ys{0.2, 0.9, 1.1};
  std::vector<double> mus{0.1, 1.3, 3.0};

  for (double sigmas : {0.1, 1.1, 3.2}) {
    for (double taus : {0.01, 0.1, 0.5, 0.9, 0.99}) {
      double x = 0.0;
      for (int i = 0; i < 3; i++)
        x += skew_de_ccdf_test(ys[i], mus[i], sigmas, taus);

      EXPECT_NEAR(x, skew_double_exponential_lccdf(ys, mus, sigmas, taus),
                  0.001);
    }
  }
}

TEST(ProbDistributionsSkewedDoubleExponential,
     lccdf_works_on_vectorial_y_and_sigma) {
  using stan::math::skew_double_exponential_lccdf;
  std::vector<double> ys{0.2, 0.9, 1.1};
  std::vector<double> sigmas{0.1, 1.1, 3.2};

  for (double mus : {0.1, 1.3, 3.0}) {
    for (double taus : {0.01, 0.1, 0.5, 0.9, 0.99}) {
      double x = 0.0;
      for (int i = 0; i < 3; i++)
        x += skew_de_ccdf_test(ys[i], mus, sigmas[i], taus);

      EXPECT_NEAR(x, skew_double_exponential_lccdf(ys, mus, sigmas, taus),
                  0.001);
    }
  }
}

TEST(ProbDistributionsSkewedDoubleExponential,
     lccdf_works_on_vectorial_y_and_tau) {
  using stan::math::skew_double_exponential_lccdf;
  std::vector<double> ys{0.2, 0.9, 1.1};
  std::vector<double> taus{0.1, 0.5, 0.9};

  for (double mus : {0.1, 1.3, 3.0}) {
    for (double sigmas : {0.1, 1.1, 3.2}) {
      double x = 0.0;
      for (int i = 0; i < 3; i++)
        x += skew_de_ccdf_test(ys[i], mus, sigmas, taus[i]);

      EXPECT_NEAR(x, skew_double_exponential_lccdf(ys, mus, sigmas, taus),
                  0.001);
    }
  }
}

TEST(ProbDistributionsSkewedDoubleExponential,
     lccdf_works_on_vectorial_mu_sigma_and_tau) {
  using stan::math::skew_double_exponential_lccdf;

  std::vector<double> mus{0.1, 1.3, 3.0};
  std::vector<double> sigmas{0.1, 1.1, 3.2};
  std::vector<double> taus{0.1, 0.5, 0.9};

  for (double ys : {0.1, 1.3, 3.0}) {
    double x = 0.0;
    for (int i = 0; i < 3; i++)
      x += skew_de_ccdf_test(ys, mus[i], sigmas[i], taus[i]);
    EXPECT_NEAR(x, skew_double_exponential_lccdf(ys, mus, sigmas, taus), 0.001);
  }
}

TEST(ProbDistributionsSkewedDoubleExponential, lccdf_check_errors) {
  using stan::math::skew_double_exponential_lccdf;
  static double inff = std::numeric_limits<double>::infinity();
  EXPECT_THROW(stan::math::skew_double_exponential_lccdf(1.0, 0.0, -1, 0.5),
               std::domain_error);
  EXPECT_THROW(stan::math::skew_double_exponential_lccdf(1.0, 0.0, 0.1, -0.5),
               std::domain_error);
  EXPECT_THROW(stan::math::skew_double_exponential_lccdf(inff, 0.0, 0.1, 1.5),
               std::domain_error);
  EXPECT_THROW(stan::math::skew_double_exponential_lccdf(1.0, inff, 0.1, 1.5),
               std::domain_error);
}

TEST(ProbDistributionsSkewedDoubleExponential, lccdf_check_inconsistent_size) {
  using stan::math::skew_double_exponential_lccdf;

  std::vector<double> mus{0.1, 1.3, 3.0};
  std::vector<double> sigmas{0.1, 1.1, 3.2, 1.0};
  std::vector<double> taus{0.1, 0.5, 0.9};
  EXPECT_THROW(
      stan::math::skew_double_exponential_lccdf(1.0, mus, sigmas, taus),
      std::invalid_argument);
}

TEST(ProbDistributionsSkewedDoubleExponential, cdf_log_matches_lccdf) {
  double y = 0.8;
  double mu = 2;
  double sigma = 2.3;
  double tau = 0.1;

  EXPECT_FLOAT_EQ(
      (stan::math::skew_double_exponential_lccdf(y, mu, sigma, tau)),
      (stan::math::skew_double_exponential_ccdf_log(y, mu, sigma, tau)));
  EXPECT_FLOAT_EQ(
      (stan::math::skew_double_exponential_lccdf<double, double, double,
                                                 double>(y, mu, sigma, tau)),
      (stan::math::skew_double_exponential_ccdf_log<double, double, double,
                                                    double>(y, mu, sigma,
                                                            tau)));
}


#ifndef INCLUDE_SWIFT_VIO_SIMPLE_IMU_ODOMETRY_HPP_
#define INCLUDE_SWIFT_VIO_SIMPLE_IMU_ODOMETRY_HPP_

#include <Eigen/Core>
#include <swift_vio/Measurements.hpp>

#include <okvis/kinematics/Transformation.hpp>
#include <okvis/kinematics/operators.hpp>
#include <okvis/FrameTypedefs.hpp>
#include <okvis/ImuMeasurements.hpp>
#include <okvis/Variables.hpp>
#include <okvis/assert_macros.hpp>

namespace swift_vio {
namespace ode {
template<typename Scalar>
int predictStates(const GenericImuMeasurementDeque<Scalar> & imuMeasurements,
            const Scalar gravityMag,
            std::pair<Eigen::Matrix<Scalar, 3, 1>, Eigen::Quaternion<Scalar>> &T_WS,
            Eigen::Matrix<Scalar, 9,1>& speedBgBa,
            const Scalar t_start, const Scalar t_end)
{
    Eigen::Matrix<Scalar, 3,1> r_old(T_WS.first), r_new(r_old), v_old(speedBgBa.template head<3>()), v_new(v_old);
    Eigen::Quaternion<Scalar> q_old(T_WS.second.conjugate()), q_new(q_old); //rotation from world to sensor
    assert(imuMeasurements.front().timeStamp <= t_start && imuMeasurements.back().timeStamp >= t_end);

    Scalar time = t_start;
    Scalar end = t_end;
    Scalar Delta_t = Scalar(0);
    bool hasStarted = false;
    int i = 0;
    Scalar nexttime;
    for(typename GenericImuMeasurementDeque<Scalar>::const_iterator it = imuMeasurements.begin();
          it != imuMeasurements.end(); ++it) {

      Eigen::Matrix<Scalar,3,1> omega_S_0 = it->template gyroscopes;
      Eigen::Matrix<Scalar,3,1> acc_S_0 = it->template accelerometers;
      Eigen::Matrix<Scalar,3,1> omega_S_1 = (it + 1)->template gyroscopes;
      Eigen::Matrix<Scalar,3,1> acc_S_1 = (it + 1)->template accelerometers;

      // time delta
      if ((it + 1) == imuMeasurements.end()) {
        nexttime = t_end;
      } else
        nexttime = (it + 1)->timeStamp;
      Scalar dt = nexttime - time;

      if (end < nexttime) {
        Scalar interval = nexttime - it->timeStamp;
        nexttime = t_end;
        dt = nexttime - time;
        if (dt == 0.0)
            break;
        const Scalar r = dt / interval;
        omega_S_1 = ((Scalar(1.0) - r) * omega_S_0 + r * omega_S_1).eval();
        acc_S_1 = ((Scalar(1.0) - r) * acc_S_0 + r * acc_S_1).eval();
      }

      if (dt <= Scalar(0.0)) {
        continue;
      }
      Delta_t += dt;

      if (!hasStarted) {
        hasStarted = true;
        const Scalar r = dt / (nexttime - it->timeStamp);
        omega_S_0 = (r * omega_S_0 + (Scalar(1.0) - r) * omega_S_1).eval();
        acc_S_0 = (r * acc_S_0 + (Scalar(1.0) - r) * acc_S_1).eval();
      }

      // actual propagation
      Eigen::Matrix<Scalar,3,1> a_est = Scalar(0.5)*(acc_S_0+acc_S_1) - speedBgBa.template tail<3>();
      Eigen::Matrix<Scalar,3,1> w_est = Scalar(0.5)*(omega_S_0+omega_S_1) - speedBgBa.template segment<3>(3);
      Eigen::Matrix<Scalar,3,1> gW(Scalar(0), Scalar(0), -gravityMag);

      Eigen::Quaternion<Scalar> qb = okvis::kinematics::rvec2quat(-w_est*dt);
      q_new=qb*q_old;

      Eigen::Matrix<Scalar,3,1> vel_inc1=(q_old.conjugate()._transformVector(a_est*dt)+q_new.conjugate()._transformVector(a_est*dt))*Scalar(0.5);
      Eigen::Matrix<Scalar,3,1> vel_inc2=gW*dt;

      v_new= v_old + vel_inc1+vel_inc2;
      r_new= r_old + (v_new+v_old)*dt*Scalar(0.5);

      time = nexttime;
      ++i;

      if (nexttime == t_end)
        break;

      r_old=r_new;
      v_old=v_new;
      q_old=q_new;
    }
    assert(nexttime == t_end);

    T_WS.second = q_new.conjugate();
    T_WS.first = r_new;
    speedBgBa. template head<3>()=v_new;
    return i;
}

/**
 * @brief predictStates accumulate the rotation vector given gyro measurements.
 * i.e., \f$ \int_{t_1}^{t_2} \hat{\omega} - \mathbf{b}_g \f$
 */
template <typename Scalar>
int predictStates(
    const std::deque<
        okvis::Measurement<Eigen::Vector3d>,
        Eigen::aligned_allocator<okvis::Measurement<Eigen::Vector3d>>>
        &gyroMeasurements,
    Eigen::Vector3d *rotVector, const Eigen::Matrix<Scalar, 3, 1> &gyroBias,
    const okvis::Time t_start, const okvis::Time t_end) {
  rotVector->setZero();
  assert(gyroMeasurements.front().timeStamp <= t_start &&
         gyroMeasurements.back().timeStamp >= t_end);

  okvis::Time time = t_start;
  okvis::Time end = t_end;
  Scalar Delta_t = Scalar(0);
  bool hasStarted = false;
  int i = 0;
  okvis::Time nexttime;
  for (typename std::deque<okvis::Measurement<Eigen::Vector3d>,
                           Eigen::aligned_allocator<okvis::Measurement<
                               Eigen::Vector3d>>>::const_iterator it =
           gyroMeasurements.begin();
       it != gyroMeasurements.end(); ++it) {
    Eigen::Matrix<Scalar, 3, 1> omega_S_0 = it->measurement;
    Eigen::Matrix<Scalar, 3, 1> omega_S_1 = (it + 1)->measurement;
    // time delta
    if ((it + 1) == gyroMeasurements.end()) {
      nexttime = t_end;
    } else
      nexttime = (it + 1)->timeStamp;

    Scalar dt = (Scalar)(nexttime - time).toSec();
    if (end < nexttime) {
      Scalar interval = (Scalar)(nexttime - it->timeStamp).toSec();
      nexttime = t_end;
      dt = (Scalar)(nexttime - time).toSec();
      if (dt == 0.0)
        break;
      const Scalar r = dt / interval;
      omega_S_1 = ((Scalar(1.0) - r) * omega_S_0 + r * omega_S_1).eval();
    }

    if (dt <= Scalar(0.0)) {
      continue;
    }
    Delta_t += dt;

    if (!hasStarted) {
      hasStarted = true;
      const Scalar r = dt / ((Scalar)(nexttime - it->timeStamp).toSec());
      omega_S_0 = (r * omega_S_0 + (Scalar(1.0) - r) * omega_S_1).eval();
    }

    // actual propagation
    Eigen::Matrix<Scalar, 3, 1> w_est =
        Scalar(0.5) * (omega_S_0 + omega_S_1) - gyroBias;
    (*rotVector) += w_est * dt;

    time = nexttime;
    ++i;

    if (nexttime == t_end)
      break;
  }
  assert(nexttime == t_end);
  return i;
}

// time_pair[0] timestamp of the provided state values. time_pair[0] >= time_pair[1],
template<typename Scalar>
int predictStatesBackward(const GenericImuMeasurementDeque<Scalar> & imuMeasurements,
                    const Scalar gravityMag,
                    std::pair<Eigen::Matrix<Scalar, 3, 1>, Eigen::Quaternion<Scalar>> &T_WS,
                    Eigen::Matrix<Scalar, 9,1>& speedBgBa,
                    const Scalar t_start, const Scalar t_end)
{
    Eigen::Matrix<Scalar, 3,1> r_old(T_WS.first), r_new(r_old), v_old(speedBgBa.template head<3>()), v_new(v_old);
    Eigen::Quaternion<Scalar> q_old(T_WS.second.conjugate()), q_new(q_old); //rotation from world to sensor
    assert(imuMeasurements.front().timeStamp <= t_end && imuMeasurements.back().timeStamp >= t_start);

    Scalar time = t_start;
    Scalar end = t_end;
    Scalar Delta_t = Scalar(0);
    bool hasStarted = false;
    int i = 0;
    Scalar nexttime;
    for(typename GenericImuMeasurementDeque<Scalar>::const_reverse_iterator it = imuMeasurements.rbegin();
          it != imuMeasurements.rend(); ++it) {

      Eigen::Matrix<Scalar,3,1> omega_S_0 = it->template gyroscopes;
      Eigen::Matrix<Scalar,3,1> acc_S_0 = it->template accelerometers;
      Eigen::Matrix<Scalar,3,1> omega_S_1 = (it + 1)->template gyroscopes;
      Eigen::Matrix<Scalar,3,1> acc_S_1 = (it + 1)->template accelerometers;

      // time delta
      if ((it + 1) == imuMeasurements.rend()) {
        nexttime = t_end;
      } else
        nexttime = (it + 1)->timeStamp;
      Scalar dt = nexttime - time;

      if (end > nexttime) {
        Scalar interval = nexttime - it->timeStamp;
        nexttime = t_end;
        dt = nexttime - time;
        if (dt == 0.0)
            break;
        const Scalar r = dt / interval;
        omega_S_1 = ((Scalar(1.0) - r) * omega_S_0 + r * omega_S_1).eval();
        acc_S_1 = ((Scalar(1.0) - r) * acc_S_0 + r * acc_S_1).eval();
      }

      if (dt >= Scalar(0.0)) {
        continue;
      }
      Delta_t += dt;

      if (!hasStarted) {
        hasStarted = true;
        const Scalar r = dt / (nexttime - it->timeStamp);
        omega_S_0 = (r * omega_S_0 + (Scalar(1.0) - r) * omega_S_1).eval();
        acc_S_0 = (r * acc_S_0 + (Scalar(1.0) - r) * acc_S_1).eval();
      }

      // actual propagation

      Eigen::Matrix<Scalar,3,1> a_est = Scalar(0.5)*(acc_S_0+acc_S_1) - speedBgBa.template tail<3>();
      Eigen::Matrix<Scalar,3,1> w_est = Scalar(0.5)*(omega_S_0+omega_S_1) - speedBgBa.template segment<3>(3);
      Eigen::Matrix<Scalar,3,1> gW(Scalar(0), Scalar(0), -gravityMag);

      Eigen::Quaternion<Scalar> qb = okvis::kinematics::rvec2quat(-w_est*dt);
      q_new=qb*q_old;

      Eigen::Matrix<Scalar,3,1> vel_inc1=(q_old.conjugate()._transformVector(a_est*dt)+q_new.conjugate()._transformVector(a_est*dt))*Scalar(0.5);
      Eigen::Matrix<Scalar,3,1> vel_inc2=gW*dt;

      v_new= v_old + vel_inc1+vel_inc2;
      r_new= r_old + (v_new+v_old)*dt*Scalar(0.5);

      time = nexttime;
      ++i;

      if (nexttime == t_end)
        break;

      r_old=r_new;
      v_old=v_new;
      q_old=q_new;
    }
    assert(nexttime == t_end);

    T_WS.second = q_new.conjugate();
    T_WS.first = r_new;
    speedBgBa.template head<3>()=v_new;
    return i;
}

template<typename Scalar>
int predictStates(const okvis::ImuMeasurementDeque & imuMeasurements,
            const Scalar gravityMag,
            std::pair<Eigen::Matrix<Scalar, 3, 1>, Eigen::Quaternion<Scalar>> &T_WS,
            Eigen::Matrix<Scalar, 9,1>& speedBgBa,
            const okvis::Time t_start, const okvis::Time t_end)
{
    Eigen::Matrix<Scalar, 3,1> r_old(T_WS.first), r_new(r_old), v_old(speedBgBa.template head<3>()), v_new(v_old);
    Eigen::Quaternion<Scalar> q_old(T_WS.second.conjugate()), q_new(q_old); //rotation from world to sensor
    assert(imuMeasurements.front().timeStamp <= t_start && imuMeasurements.back().timeStamp >= t_end);

    okvis::Time time = t_start;
    okvis::Time end = t_end;
    Scalar Delta_t = Scalar(0);
    bool hasStarted = false;
    int i = 0;
    okvis::Time nexttime;
    for(typename okvis::ImuMeasurementDeque::const_iterator it = imuMeasurements.begin();
          it != imuMeasurements.end(); ++it) {

      Eigen::Matrix<Scalar,3,1> omega_S_0 = it->measurement.gyroscopes;
      Eigen::Matrix<Scalar,3,1> acc_S_0 = it->measurement.accelerometers;
      Eigen::Matrix<Scalar,3,1> omega_S_1 = (it + 1)->measurement.gyroscopes;
      Eigen::Matrix<Scalar,3,1> acc_S_1 = (it + 1)->measurement.accelerometers;

      // time delta
      if ((it + 1) == imuMeasurements.end()) {
        nexttime = t_end;
      } else
        nexttime = (it + 1)->timeStamp;

      Scalar dt = (Scalar)(nexttime - time).toSec();

      if (end < nexttime) {
        Scalar interval = (Scalar)(nexttime - it->timeStamp).toSec();
        nexttime = t_end;
        dt =(Scalar)(nexttime - time).toSec();
        if (dt == 0.0)
            break;
        const Scalar r = dt / interval;
        omega_S_1 = ((Scalar(1.0) - r) * omega_S_0 + r * omega_S_1).eval();
        acc_S_1 = ((Scalar(1.0) - r) * acc_S_0 + r * acc_S_1).eval();
      }

      if (dt <= Scalar(0.0)) {
        continue;
      }
      Delta_t += dt;

      if (!hasStarted) {
        hasStarted = true;
        const Scalar r = dt / ((Scalar)(nexttime - it->timeStamp).toSec());
        omega_S_0 = (r * omega_S_0 + (Scalar(1.0) - r) * omega_S_1).eval();
        acc_S_0 = (r * acc_S_0 + (Scalar(1.0) - r) * acc_S_1).eval();
      }

      // actual propagation
      Eigen::Matrix<Scalar,3,1> a_est = Scalar(0.5)*(acc_S_0+acc_S_1) - speedBgBa.template tail<3>();
      Eigen::Matrix<Scalar,3,1> w_est = Scalar(0.5)*(omega_S_0+omega_S_1) - speedBgBa.template segment<3>(3);
      Eigen::Matrix<Scalar,3,1> gW(Scalar(0), Scalar(0), -gravityMag);

      Eigen::Quaternion<Scalar> qb = okvis::kinematics::rvec2quat(-w_est*dt);
      q_new=qb*q_old;

      Eigen::Matrix<Scalar,3,1> vel_inc1=(q_old.conjugate()._transformVector(a_est*dt)+q_new.conjugate()._transformVector(a_est*dt))*Scalar(0.5);
      Eigen::Matrix<Scalar,3,1> vel_inc2=gW*dt;

      v_new= v_old + vel_inc1+vel_inc2;
      r_new= r_old + (v_new+v_old)*dt*Scalar(0.5);

      time = nexttime;
      ++i;

      if (nexttime == t_end)
        break;

      r_old=r_new;
      v_old=v_new;
      q_old=q_new;
    }
    assert(nexttime == t_end);

    T_WS.second = q_new.conjugate();
    T_WS.first = r_new;
    speedBgBa. template head<3>()=v_new;
    return i;
}

// time_pair[0] timestamp of the provided state values. time_pair[0] >= time_pair[1],
template<typename Scalar>
int predictStatesBackward(const okvis::ImuMeasurementDeque & imuMeasurements,
                    const Scalar gravityMag,
                    std::pair<Eigen::Matrix<Scalar, 3, 1>, Eigen::Quaternion<Scalar>> &T_WS,
                    Eigen::Matrix<Scalar, 9,1>& speedBgBa,
                    const okvis::Time t_start, const okvis::Time t_end)
{
    Eigen::Matrix<Scalar, 3,1> r_old(T_WS.first), r_new(r_old), v_old(speedBgBa.template head<3>()), v_new(v_old);
    Eigen::Quaternion<Scalar> q_old(T_WS.second.conjugate()), q_new(q_old); //rotation from world to sensor
    assert(imuMeasurements.front().timeStamp <= t_end && imuMeasurements.back().timeStamp >= t_start);
    // if this assertion fails during optimization, it often means the time offset variables diverge. Solution:
    // either try to enlarge the range of imu reading segment, or decrease the time variables' std to effectively lock them

    okvis::Time time = t_start;
    okvis::Time end = t_end;
    Scalar Delta_t = Scalar(0);
    bool hasStarted = false;
    int i = 0;
    okvis::Time nexttime;
    for(typename okvis::ImuMeasurementDeque::const_reverse_iterator it = imuMeasurements.rbegin();
          it != imuMeasurements.rend(); ++it) {

      Eigen::Matrix<Scalar,3,1> omega_S_0 = it->measurement.gyroscopes;
      Eigen::Matrix<Scalar,3,1> acc_S_0 = it->measurement.accelerometers;
      Eigen::Matrix<Scalar,3,1> omega_S_1 = (it + 1)->measurement.gyroscopes;
      Eigen::Matrix<Scalar,3,1> acc_S_1 = (it + 1)->measurement.accelerometers;

      // time delta
      if ((it + 1) == imuMeasurements.rend()) {
        nexttime = t_end;
      } else
        nexttime = (it + 1)->timeStamp;
      Scalar dt = (Scalar)(nexttime - time).toSec();

      if (end > nexttime) {
        Scalar interval = (Scalar)(nexttime - it->timeStamp).toSec();
        nexttime = t_end;
        dt = (Scalar)(nexttime - time).toSec();
        if (dt == 0.0)
            break;
        const Scalar r = dt / interval;
        omega_S_1 = ((Scalar(1.0) - r) * omega_S_0 + r * omega_S_1).eval();
        acc_S_1 = ((Scalar(1.0) - r) * acc_S_0 + r * acc_S_1).eval();
      }

      if (dt >= Scalar(0.0)) {
        continue;
      }
      Delta_t += dt;

      if (!hasStarted) {
        hasStarted = true;
        const Scalar r = dt / ((Scalar)(nexttime - it->timeStamp).toSec());
        omega_S_0 = (r * omega_S_0 + (Scalar(1.0) - r) * omega_S_1).eval();
        acc_S_0 = (r * acc_S_0 + (Scalar(1.0) - r) * acc_S_1).eval();
      }

      // actual propagation

      Eigen::Matrix<Scalar,3,1> a_est = Scalar(0.5)*(acc_S_0+acc_S_1) - speedBgBa.template tail<3>();
      Eigen::Matrix<Scalar,3,1> w_est = Scalar(0.5)*(omega_S_0+omega_S_1) - speedBgBa.template segment<3>(3);
      Eigen::Matrix<Scalar,3,1> gW(Scalar(0), Scalar(0), -gravityMag);

      Eigen::Quaternion<Scalar> qb = okvis::kinematics::rvec2quat(-w_est*dt);
      q_new=qb*q_old;

      Eigen::Matrix<Scalar,3,1> vel_inc1=(q_old.conjugate()._transformVector(a_est*dt)+q_new.conjugate()._transformVector(a_est*dt))*Scalar(0.5);
      Eigen::Matrix<Scalar,3,1> vel_inc2=gW*dt;

      v_new= v_old + vel_inc1+vel_inc2;
      r_new= r_old + (v_new+v_old)*dt*Scalar(0.5);

      time = nexttime;
      ++i;

      if (nexttime == t_end)
        break;

      r_old=r_new;
      v_old=v_new;
      q_old=q_new;
    }
    assert(nexttime == t_end);

    T_WS.second = q_new.conjugate();
    T_WS.first = r_new;
    speedBgBa.template head<3>()=v_new;
    return i;
}

// linear interpolation
inline void interpolateInertialData(const okvis::ImuMeasurementDeque &imuMeas,
                                    const okvis::Time &queryTime,
                                    okvis::ImuMeasurement &queryValue) {
  auto iterLeft = imuMeas.begin(), iterRight = imuMeas.end();
  if (iterLeft->timeStamp > queryTime)
    throw std::runtime_error("iterLeft->timeStamp > queryTime: Imu "
                             "measurements has wrong timestamps");
  for (auto iter = imuMeas.begin(); iter != imuMeas.end(); ++iter) {
    if (iter->timeStamp < queryTime) {
      iterLeft = iter;
    } else if (iter->timeStamp == queryTime) {
      queryValue = *iter;
      return;
    } else {
      iterRight = iter;
      break;
    }
  }
  double ratio = (queryTime - iterLeft->timeStamp).toSec() /
                 (iterRight->timeStamp - iterLeft->timeStamp).toSec();
  queryValue.timeStamp = queryTime;
  Eigen::Vector3d omega_S0 =
      (iterRight->measurement.gyroscopes - iterLeft->measurement.gyroscopes) *
          ratio +
      iterLeft->measurement.gyroscopes;
  Eigen::Vector3d acc_S0 = (iterRight->measurement.accelerometers -
                            iterLeft->measurement.accelerometers) *
                               ratio +
                           iterLeft->measurement.accelerometers;
  queryValue.measurement.gyroscopes = omega_S0;
  queryValue.measurement.accelerometers = acc_S0;
}
}  // namespace ode
}  // namespace swift_vio
#endif


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

/* Test_matmul.cpp - Testing the functionality of multiplying an encrypted
 * vector by a plaintext matrix, either over the extension- or the
 * base-field/ring.
 */

#include <cassert>
#include <NTL/lzz_pXFactoring.h>
#include "FHE.h"
#include "timing.h"
#include "EncryptedArray.h"
#include "matmul.h"

static MatMulBase* buildRandomMatrix(EncryptedArray& ea);
static MatMulBase* buildRandomBlockMatrix(const EncryptedArray& ea);

void  TestIt(long m, long p, long r, long d, long L, bool verbose)
{
  cout << "*** TestIt: m=" << m
       << ", p=" << p
       << ", r=" << r
       << ", d=" << d
       << ", L=" << L
       << endl;

  FHEcontext context(m, p, r);
  buildModChain(context, L, /*c=*/3);

  FHESecKey secretKey(context);
  const FHEPubKey& publicKey = secretKey;
  secretKey.GenSecKey(/*w=*/64); // A Hamming-weight-w secret key

  ZZX G;
  if (d == 0)
    G = context.alMod.getFactorsOverZZ()[0];
  else
    G = makeIrredPoly(p, d);

  if (verbose) {
    context.zMStar.printout();
    cout << endl;
    cout << "G = " << G << "\n";
  }

  addSome1DMatrices(secretKey); // compute key-switching matrices that we need
  addFrbMatrices(secretKey); // compute key-switching matrices that we need
  EncryptedArray ea(context, G);

  // Test a "dense" matrix over the extension field
  {
    // choose a random plaintext square matrix
    unique_ptr<MatMulBase> ptr(buildRandomMatrix(ea));

    // choose a random plaintext vector
    NewPlaintextArray v(ea);
    random(ea, v);

    // encrypt the random vector
    Ctxt ctxt(publicKey);
    ea.encrypt(ctxt, publicKey, v);
    Ctxt ctxt2 = ctxt;

    cout << " Multiplying with MatMulBase... " << std::flush;
    matMul(ctxt2, *ptr, cachezzX); // multiply ciphertext and build cache
    matMul(v, *ptr);     // multiply the plaintext vector

    NewPlaintextArray v1(ea);
    ea.decrypt(ctxt2, secretKey, v1); // decrypt the ciphertext vector

    if (equals(ea, v, v1))        // check that we've got the right answer
      cout << "Nice!!\n";
    else
      cout << "Grrr@*\n";

    cout << " Multiplying with MatMulBase+dcrt cache... " << std::flush;
    ctxt2 = ctxt;
    matMul(ctxt2, *ptr, cacheDCRT); // upgrade cache and use in multiplication

    ea.decrypt(ctxt2, secretKey, v1); // decrypt the ciphertext vector

    if (equals(ea, v, v1))        // check that we've got the right answer
      cout << "Nice!!\n";
    else
      cout << "Grrr@*\n";
  }
  {
    // choose a random plaintext square matrix
    unique_ptr<MatMulBase> ptr(buildRandomMatrix(ea));

    // choose a random plaintext vector
    NewPlaintextArray v(ea);
    random(ea, v);

    // encrypt the random vector
    Ctxt ctxt(publicKey);
    ea.encrypt(ctxt, publicKey, v);
    cout << " Multiplying with MatMulBase+zzx cache... " << std::flush;
    buildCache4MatMul(*ptr, cachezzX);// build the cache
    matMul(ctxt, *ptr);               // then use it
    matMul(v, *ptr);     // multiply the plaintext vector

    NewPlaintextArray v1(ea);
    ea.decrypt(ctxt, secretKey, v1); // decrypt the ciphertext vector

    if (equals(ea, v, v1))        // check that we've got the right answer
      cout << "Nice!!\n";
    else
      cout << "Grrr@*\n";
  }
  // Test a "diagonal sparse" matrix over the extension field
  {
    // choose a random plaintext square matrix
    unique_ptr<MatMulBase> ptr(buildRandomMatrix(ea));

    // choose a random plaintext vector
    NewPlaintextArray v(ea);
    random(ea, v);

    // encrypt the random vector
    Ctxt ctxt(publicKey);
    ea.encrypt(ctxt, publicKey, v);
    Ctxt ctxt2 = ctxt;

    cout << "\n Multiplying with Sparse MatMulBase... " << std::flush;
    matMul_sparse(ctxt2, *ptr, cachezzX); // multiply ciphertext and build cache
    matMul(v, *ptr);     // multiply the plaintext vector

    NewPlaintextArray v1(ea);
    ea.decrypt(ctxt2, secretKey, v1); // decrypt the ciphertext vector

    if (equals(ea, v, v1))        // check that we've got the right answer
      cout << "Nice!!\n";
    else
      cout << "Grrr@*\n";

    cout << " Multiplying with Sparse MatMulBase+dcrt cache... " << std::flush;
    ctxt2 = ctxt;
    matMul_sparse(ctxt2, *ptr, cacheDCRT); // upgrade the cache

    ea.decrypt(ctxt2, secretKey, v1); // decrypt the ciphertext vector

    if (equals(ea, v, v1))        // check that we've got the right answer
      cout << "Nice!!\n";
    else
      cout << "Grrr@*\n";
  }
  {
    // choose a random plaintext square matrix
    unique_ptr<MatMulBase> ptr(buildRandomMatrix(ea));

    // choose a random plaintext vector
    NewPlaintextArray v(ea);
    random(ea, v);

    // encrypt the random vector
    Ctxt ctxt(publicKey);
    ea.encrypt(ctxt, publicKey, v);

    cout << " Multiplying with Sparse MatMulBase+zzx cache... " << std::flush;
    buildCache4MatMul_sparse(*ptr, cachezzX); // build the cache
    matMul_sparse(ctxt, *ptr);                // then use it
    matMul(v, *ptr);                          // multiply plaintext vector

    NewPlaintextArray v1(ea);
    ea.decrypt(ctxt, secretKey, v1); // decrypt the ciphertext vector

    if (equals(ea, v, v1))        // check that we've got the right answer
      cout << "Nice!!\n";
    else
      cout << "Grrr@*\n";
  }

  // Test a "block matrix" over the base field
  {
    // choose a random plaintext square matrix
    std::unique_ptr<MatMulBase> ptr(buildRandomBlockMatrix(ea));

    // choose a random plaintext vector
    NewPlaintextArray v(ea);
    random(ea, v);

    // encrypt the random vector
    Ctxt ctxt(publicKey);
    ea.encrypt(ctxt, publicKey, v);
    Ctxt ctxt2 = ctxt;

    cout << "\n Multiplying with BlockMatMul... "  << std::flush;
    blockMatMul(ctxt2, *ptr, cachezzX); // multiply ciphertext and build cache
    blockMatMul(v, *ptr);      // multiply the plaintext vector

    NewPlaintextArray v1(ea);
    ea.decrypt(ctxt2, secretKey, v1); // decrypt the ciphertext vector

    if (equals(ea, v, v1))        // check that we've got the right answer
      cout << "Nice!!\n";
    else
      cout << "Grrr...\n";

    cout << " Multiplying with BlockMatMul+dcrt cache... " << std::flush;
    ctxt2 = ctxt;
    blockMatMul(ctxt2, *ptr, cacheDCRT); // upgrade the cache

    ea.decrypt(ctxt2, secretKey, v1); // decrypt the ciphertext vector
    if (equals(ea, v, v1))        // check that we've got the right answer
      cout << "Nice!!\n";
    else
      cout << "Grrr@*\n";
  }
  {
    // choose a random plaintext square matrix
    std::unique_ptr<MatMulBase> ptr(buildRandomBlockMatrix(ea));

    // choose a random plaintext vector
    NewPlaintextArray v(ea);
    random(ea, v);

    // encrypt the random vector
    Ctxt ctxt(publicKey);
    ea.encrypt(ctxt, publicKey, v);

    cout << " Multiplying with BlockMatMul+zzx cache... " << std::flush;
    buildCache4BlockMatMul(*ptr, cachezzX); // build the cache
    blockMatMul(ctxt, *ptr);                // then use it
    blockMatMul(v, *ptr);     // multiply the plaintext vector

    NewPlaintextArray v1(ea);
    ea.decrypt(ctxt, secretKey, v1); // decrypt the ciphertext vector

    if (equals(ea, v, v1))        // check that we've got the right answer
      cout << "Nice!!\n";
    else
      cout << "Grrr@*\n";
  }
}

void usage(char *prog)
{
  cout << "Usage: "<<prog<<" [ optional parameters ]...\n";
  cout << "  optional parameters have the form 'attr1=val1 attr2=val2 ...'\n";
  cout << "  e.g, 'm=2047 p=2 L=4'\n\n";
  cout << "  m defines the cyclotomic polynomial Phi_m(X)\n";
  cout << "  p is the plaintext base [default=2]" << endl;
  cout << "  r is the lifting [default=1]" << endl;
  cout << "  d is the degree of the field extension [default==1]\n";
  cout << "    (d == 0 => factors[0] defined the extension)\n";
  cout << "  L is the # of primes in the modulus chain [default=4]\n";
  cout << "  verbose print timing info [default=0]\n";
  exit(0);
}

/* Testing the functionality of multiplying an encrypted vector by a plaintext
 * matrix, either over the extension- or the base-field/ring.
 */
int main(int argc, char *argv[])
{
  argmap_t argmap;
  argmap["m"] = "2047";
  argmap["p"] = "2";
  argmap["r"] = "1";
  argmap["d"] = "1";
  argmap["L"] = "4";
  argmap["verbose"] = "0";

  // get parameters from the command line
  if (!parseArgs(argc, argv, argmap)) usage(argv[0]);

  long m = atoi(argmap["m"]);
  long p = atoi(argmap["p"]);
  long r = atoi(argmap["r"]);
  long d = atoi(argmap["d"]);
  long L = atoi(argmap["L"]);
  bool v = atoi(argmap["verbose"]);

  //  setTimersOn();
  setTimersOn();
  TestIt(m, p, r, d, L, v);
  cout << endl;
  if (v) {
    printAllTimers();
    cout << endl;
  }
}

template<class type> class RandomMatrix : public MatMul<type> {
  PA_INJECT(type)
  vector< vector< RX > > data;

public:
  ~RandomMatrix() {/*cout << "destructor: random dense matrix\n";*/}
  RandomMatrix(const EncryptedArray& _ea): MatMul<type>(_ea) {
    long n = _ea.size();
    long d = _ea.getDegree();
    long bnd = 2*n; // non-zero with probability 1/bnd

    RBak bak; bak.save(); _ea.getContext().alMod.restoreContext();
    data.resize(n);
    for (long i = 0; i < n; i++) {
      data[i].resize(n);
      for (long j = 0; j < n; j++) {
        bool zEntry = (RandomBnd(bnd) > 0);
        if (zEntry) clear(data[i][j]);
        else        random(data[i][j], d);
      }
    }
  }

  virtual bool get(RX& out, long i, long j) const {
    assert(i >= 0 && i < this->getEA().size());
    assert(j >= 0 && j < this->getEA().size());
    if (IsZero(data[i][j])) return true;
    out = data[i][j];
    return false;
  }
};
static MatMulBase* buildRandomMatrix(EncryptedArray& ea)
{
  switch (ea.getTag()) {
    case PA_GF2_tag: { return new RandomMatrix<PA_GF2>(ea); }
    case PA_zz_p_tag:{ return new RandomMatrix<PA_zz_p>(ea); }
    default: return nullptr;
  }
}

template<class type> class RandomBlockMatrix : public BlockMatMul<type> {
  PA_INJECT(type)

  std::vector< std::vector< mat_R > > data;

public:
  virtual ~RandomBlockMatrix() {}
  RandomBlockMatrix(const EncryptedArray& _ea): BlockMatMul<type>(_ea)
  {
    RBak bak; bak.save(); _ea.getAlMod().restoreContext();
    long n = _ea.size();
    long d = _ea.getDegree();
    long bnd = 2*n; // non-zero with probability 1/bnd

    data.resize(n);
    for (long i = 0; i < n; i++) {
      data[i].resize(n);
      for (long j = 0; j < n; j++) {
        data[i][j].SetDims(d, d);

        bool zEntry = (RandomBnd(bnd) > 0);

        for (long u = 0; u < d; u++)
          for (long v = 0; v < d; v++)
            if (zEntry)
              clear(data[i][j][u][v]);
            else
              random(data[i][j][u][v]);
      }
    }
  }

  virtual bool get(mat_R& out, long i, long j) const {
    assert(i >= 0 && i < this->getEA().size());
    assert(j >= 0 && j < this->getEA().size());
    if (IsZero(data[i][j])) return true;
    out = data[i][j];
    return false;
  }

  const std::vector< std::vector< mat_R > >& getData() const {return data;}
};

static MatMulBase* buildRandomBlockMatrix(const EncryptedArray& ea)
{
  switch (ea.getTag()) {
    case PA_GF2_tag: {
      return new RandomBlockMatrix<PA_GF2>(ea);
    }

    case PA_zz_p_tag: {
      return new RandomBlockMatrix<PA_zz_p>(ea);
    }

    default: return 0;
  }
}


#include <catch2/catch.hpp>

#include "engine.hpp"

#include <iostream>
#include <boost/units/io.hpp>

using EvolutionaryWalker::LibUtils::Point;
using EvolutionaryWalker::Physics::Engine;
using EvolutionaryWalker::Physics::LengthQuantity;

namespace
{

}

TEST_CASE("Basic uses of Physics::Engine", "[physics]")
{
    using boost::units::si::meters;
    using boost::units::si::seconds;
    using boost::units::si::kilograms;
    using boost::units::si::meters_per_second_squared;

    SECTION("Basic spring oscillator")
    {
        Engine e{9.61 * meters_per_second_squared};
        const auto ref = e.addFixedNode({{0 * meters, 10 * meters}});
        const auto p = e.addNode({{0 * meters, 0 * meters}}, 1 * kilograms);
        e.addSpring(ref, p, {5 * meters, 1 * meters_per_second_squared});

        e.init();
        const auto step = .0001 * seconds;
        for(std::size_t i = 0; i < 100; i++)
        {
            e.step(step);
            std::cout << "y(" << ((double)i * step) << ") = " << e.node(p)[1] << std::endl;
        }
    }
}


//  ================================================================
//  Created by Gregory Kramida on 10/24/18.
//  Copyright (c) 2018 Gregory Kramida
//  Licensed under the Apache License, Version 2.0 (the "License");
//  you may not use this file except in compliance with the License.
//  You may obtain a copy of the License at

//  http://www.apache.org/licenses/LICENSE-2.0

//  Unless required by applicable law or agreed to in writing, software
//  distributed under the License is distributed on an "AS IS" BASIS,
//  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
//  See the License for the specific language governing permissions and
//  limitations under the License.
//  ================================================================
#pragma once

//libraries
#include <Eigen/Eigen>
#include <unsupported/Eigen/CXX11/Tensor>

//local
#include "vector2.hpp"
#include "vector3.hpp"
#include "matrix2.hpp"
#include "matrix3.hpp"
#include "nestable_type_metainfo.hpp"

namespace math{
typedef unsigned char uchar;
typedef unsigned short ushort;
typedef unsigned int uint;
typedef unsigned long ulong;

typedef class math::Vector2<short> Vector2s;
typedef class math::Vector2<int> Vector2i;
typedef class math::Vector2<float> Vector2f;
typedef class math::Vector2<double> Vector2d;

typedef class math::Vector3<short> Vector3s;
typedef class math::Vector3<double> Vector3d;
typedef class math::Vector3<int> Vector3i;
typedef class math::Vector3<uint> Vector3ui;
typedef class math::Vector3<uchar> Vector3u;
typedef class math::Vector3<float> Vector3f;
typedef class math::Matrix2<float> Matrix2f;
typedef class math::Matrix3<float> Matrix3f;

typedef Eigen::MatrixXf MatrixXf;
typedef Eigen::MatrixXd MatrixXd;
typedef Eigen::Matrix<math::Vector2<float>, Eigen::Dynamic, Eigen::Dynamic> MatrixXv2f;
typedef Eigen::Matrix<math::Vector2<float>, Eigen::Dynamic, Eigen::Dynamic, Eigen::RowMajor> MatrixXv2f_rm;
typedef Eigen::Matrix<math::Matrix2<float>, Eigen::Dynamic, Eigen::Dynamic> MatrixXm2f;

typedef class Eigen::Tensor<float,3> Tensor3f;
typedef class Eigen::Tensor<float,0> Tensor0f;
typedef Eigen::Tensor<math::Vector2<float>,3> Tensor3v2f;
typedef Eigen::Tensor<math::Vector3<float>,3> Tensor3v3f;
typedef Eigen::Tensor<math::Matrix3<float>,3> Tensor3m3f;

typedef class Eigen::Matrix<unsigned short, Eigen::Dynamic,Eigen::Dynamic> MatrixXus;
typedef class Eigen::Matrix<unsigned short, Eigen::Dynamic,1> MatrixX1us;
typedef class Eigen::Matrix<unsigned short, 1, Eigen::Dynamic> Matrix1Xus;
typedef class Eigen::Matrix<unsigned char, Eigen::Dynamic,Eigen::Dynamic> MatrixXuc;
typedef class Eigen::Matrix<unsigned char, Eigen::Dynamic,1> MatrixX1uc;
typedef class Eigen::Matrix<unsigned char, 1, Eigen::Dynamic> Matrix1Xuc;
typedef class Eigen::Matrix<unsigned char, Eigen::Dynamic,Eigen::Dynamic> MatrixXuc;

}//namespace math


#include <gtest/gtest.h>

#include <Eigen/Core>
#include <algorithm>
#include <csim/init_ops.hpp>
#include <csim/memory_ops.hpp>
#include <csim/stat_ops.hpp>
#include <csim/update_ops.hpp>
#include <csim/update_ops_cpp.hpp>
#include <string>

#include "../util/util.hpp"

void test_single_diagonal_matrix_gate(
    std::function<void(UINT, const CTYPE*, CTYPE*, ITYPE)> func) {
    const UINT n = 6;
    const ITYPE dim = 1ULL << n;
    const UINT max_repeat = 10;

    Eigen::MatrixXcd Identity(2, 2), Z(2, 2);
    Identity << 1, 0, 0, 1;
    Z << 1, 0, 0, -1;

    Eigen::Matrix<std::complex<double>, 2, 2, Eigen::RowMajor> U;

    UINT target;
    double icoef, zcoef, norm;

    auto state = allocate_quantum_state(dim);
    initialize_Haar_random_state(state, dim);
    Eigen::VectorXcd test_state = Eigen::VectorXcd::Zero(dim);
    for (ITYPE i = 0; i < dim; ++i) test_state[i] = state[i];

    Eigen::MatrixXcd whole_I = Eigen::MatrixXcd::Identity(dim, dim);

    for (UINT rep = 0; rep < max_repeat; ++rep) {
        // single qubit diagonal matrix gate
        target = rand_int(n);
        icoef = rand_real();
        zcoef = rand_real();
        norm = sqrt(icoef * icoef + zcoef * zcoef);
        icoef /= norm;
        zcoef /= norm;
        U = icoef * Identity + 1.i * zcoef * Z;
        Eigen::VectorXcd diag = U.diagonal();
        func(target, (CTYPE*)diag.data(), state, dim);
        test_state =
            get_expanded_eigen_matrix_with_identity(target, U, n) * test_state;
        state_equal(state, test_state, dim, "single diagonal gate");
    }
    release_quantum_state(state);
}

TEST(UpdateTest, SingleDiagonalMatrixTest) {
    test_single_diagonal_matrix_gate(single_qubit_diagonal_matrix_gate);
    test_single_diagonal_matrix_gate(
        single_qubit_diagonal_matrix_gate_single_unroll);
#ifdef _OPENMP
    test_single_diagonal_matrix_gate(
        single_qubit_diagonal_matrix_gate_parallel_unroll);
#endif
#ifdef _USE_SIMD
    test_single_diagonal_matrix_gate(
        single_qubit_diagonal_matrix_gate_single_simd);
#ifdef _OPENMP
    test_single_diagonal_matrix_gate(
        single_qubit_diagonal_matrix_gate_parallel_simd);
#endif
#endif
}

void test_single_phase_gate(
    std::function<void(UINT, CTYPE, CTYPE*, ITYPE)> func) {
    const UINT n = 6;
    const ITYPE dim = 1ULL << n;
    const UINT max_repeat = 10;

    Eigen::Matrix<std::complex<double>, 2, 2, Eigen::RowMajor> U;

    UINT target;
    double angle;

    auto state = allocate_quantum_state(dim);
    initialize_Haar_random_state(state, dim);
    Eigen::VectorXcd test_state = Eigen::VectorXcd::Zero(dim);
    for (ITYPE i = 0; i < dim; ++i) test_state[i] = state[i];

    Eigen::MatrixXcd whole_I = Eigen::MatrixXcd::Identity(dim, dim);

    for (UINT rep = 0; rep < max_repeat; ++rep) {
        // single qubit phase matrix gate
        target = rand_int(n);
        angle = rand_real();
        U << 1, 0, 0, cos(angle) + 1.i * sin(angle);
        CTYPE t = cos(angle) + 1.i * sin(angle);
        func(target, t, state, dim);
        test_state =
            get_expanded_eigen_matrix_with_identity(target, U, n) * test_state;
        state_equal(state, test_state, dim, "single phase gate");
    }
    release_quantum_state(state);
}

TEST(UpdateTest, SinglePhaseGateTest) {
    test_single_phase_gate(single_qubit_phase_gate);
    test_single_phase_gate(single_qubit_phase_gate_single_unroll);
#ifdef _OPENMP
    test_single_phase_gate(single_qubit_phase_gate_parallel_unroll);
#endif
#ifdef _USE_SIMD
    test_single_phase_gate(single_qubit_phase_gate_single_simd);
#ifdef _OPENMP
    test_single_phase_gate(single_qubit_phase_gate_parallel_simd);
#endif
#endif
}