// Copyright (c) Microsoft Corporation.
// Licensed under the MIT license.
//
// gDLS*: Generalized Pose-and-Scale Estimation Given Scale and Gravity Priors
//
// Victor Fragoso, Joseph DeGol, Gang Hua.
// Proc. of the IEEE/CVF Conf. on Computer Vision and Pattern Recognition 2020.
//
// Please contact the author of this library if you have any questions.
// Author: Victor Fragoso (victor.fragoso@microsoft.com)

#include "gtest/gtest.h"

#include <array>
#include <random>
#include <vector>
#include <Eigen/Core>
#include <Eigen/Geometry>
#include <glog/logging.h>

#include "gdls_star/gdls_star.h"

namespace msft {
namespace {
using Eigen::AngleAxisd;
using Eigen::Quaterniond;
using Eigen::Vector3d;

using Input = GdlsStar::Input;

// Convert degrees to radians.
constexpr double DegToRad(double angle_degrees) {
  constexpr double kDegToRad = M_PI / 180.0;
  return angle_degrees * kDegToRad;
}

// Convert radiants to degrees.
constexpr double RadToDeg(double angle_radians) {
  constexpr double kRadToDeg = 180.0 / M_PI;
  return angle_radians * kRadToDeg;
}

struct TestParams {
  double max_reprojection_error;
  double max_rotation_difference;
  double max_translation_difference;
  double max_scale_difference;
  double projection_noise_std_dev;
  double down_direction_noise_std_dev;
};

struct ExpectedSolution {
  Quaterniond rotation;
  Vector3d translation;
  double scale;
};

struct TestInput {
  std::vector<Vector3d> camera_centers;
  std::vector<Vector3d> world_points;
  Vector3d query_down_direction;
  Vector3d world_down_direction;
  double scale_prior;
  double scale_penalty;
  double gravity_penalty;
};

class GdlsStarTest : public ::testing::Test {
 public:
  static void SetUpTestCase() {
    estimator = new GdlsStar;
    random_engine = new std::default_random_engine;
    prng = new std::mt19937(std::mt19937::default_seed);
    gaussian_distribution = new std::normal_distribution<double>(0.0, 1.0);
  }

  static void TearDownTestCase() {
    delete prng;
    delete random_engine;
    delete estimator;
    delete gaussian_distribution;
  }

  static Input GenerateTestingData(const ExpectedSolution& expected_solution,
                                   const TestParams& test_params,
                                   const TestInput& test_input);

  static void TestWithNoise(const ExpectedSolution& expected_solution,
                            const TestParams& test_params,
                            const TestInput& test_input);

  static void AddNoiseToRay(const double std_dev, Eigen::Vector3d* ray);

  // Estimator.
  static GdlsStar* estimator;
  static std::default_random_engine* random_engine;
  static std::mt19937* prng;
  static std::normal_distribution<double>* gaussian_distribution;
};

GdlsStar* GdlsStarTest::estimator = nullptr;
std::default_random_engine* GdlsStarTest::random_engine = nullptr;
std::mt19937* GdlsStarTest::prng = nullptr;
std::normal_distribution<double>* GdlsStarTest::gaussian_distribution = nullptr;

void GdlsStarTest::AddNoiseToRay(const double std_dev, Eigen::Vector3d* ray) {
  const double scale = CHECK_NOTNULL(ray)->norm();
  const Eigen::Quaterniond rotation =
      Eigen::Quaterniond::FromTwoVectors(Eigen::Vector3d(0, 0, 1.0), *ray);
  const Eigen::Vector3d noisy_point(std_dev * (*gaussian_distribution)(*prng),
                                    std_dev * (*gaussian_distribution)(*prng),
                                    1);
  *ray = (rotation * (scale * noisy_point)).normalized();
}

Input GdlsStarTest::GenerateTestingData(
    const ExpectedSolution& expected_solution,
    const TestParams& test_params,
    const TestInput& test_input) {
  Input input;
  const int num_points = test_input.world_points.size();
  const int num_cameras = test_input.camera_centers.size();
  std::vector<Vector3d>& camera_rays = input.ray_directions;
  std::vector<Vector3d>& ray_origins = input.ray_origins;
  camera_rays.reserve(num_points);
  ray_origins.reserve(num_points);
  input.world_points = test_input.world_points;
  const std::vector<Vector3d>& world_points = test_input.world_points;
  const std::vector<Vector3d>& camera_centers = test_input.camera_centers;
  const Quaterniond expected_rotation = expected_solution.rotation;
  const Vector3d expected_translation = expected_solution.translation;
  const double expected_scale = expected_solution.scale;
  for (int i = 0; i < num_points; ++i) {
    // Setting camera positions or ray origins.
    ray_origins.emplace_back(
        (expected_rotation * camera_centers[i % num_cameras] +
         expected_translation) / expected_scale);

    // Reproject 3D points into camera frame.
    camera_rays.emplace_back(
        (expected_rotation * world_points[i] + expected_translation -
         expected_scale * ray_origins[i]).normalized());
  }

  // Add noise to the camera rays.
  if (test_params.projection_noise_std_dev > 0.0) {
    // Adds noise to both of the rays.
    for (int i = 0; i < num_points; ++i) {
      AddNoiseToRay(test_params.projection_noise_std_dev, &camera_rays[i]);
    }
  }

  // Compute the query down direction.
  input.priors.query_down_direction =
      expected_rotation * input.priors.world_down_direction;

  if (test_params.down_direction_noise_std_dev > 0.0) {
    AddNoiseToRay(test_params.down_direction_noise_std_dev,
                  &input.priors.world_down_direction);
    AddNoiseToRay(test_params.down_direction_noise_std_dev,
                  &input.priors.query_down_direction);
  }

  // Setting penalties for the regularizers.
  input.priors.scale_penalty = test_input.scale_penalty;
  input.priors.gravity_penalty = test_input.gravity_penalty;
  input.priors.scale_prior = test_input.scale_prior;
  return input;
}

bool CheckReprojectionErrors(const Input& input,
                             const Eigen::Quaterniond& soln_rotation,
                             const Eigen::Vector3d& soln_translation,
                             const double soln_scale,
                             const double max_reprojection_error) {
  const int num_points = input.world_points.size();
  const std::vector<Vector3d>& camera_rays = input.ray_directions;
  const std::vector<Vector3d>& ray_origins = input.ray_origins;
  const std::vector<Vector3d>& world_points = input.world_points;
  double good_reprojection_errors = true;
  for (int i = 0; i < num_points; ++i) {
    const Quaterniond unrot =
        Quaterniond::FromTwoVectors(camera_rays[i], Vector3d(0, 0, 1));
    const Vector3d reprojected_point =
        (soln_rotation * world_points[i] + soln_translation) /
        soln_scale - ray_origins[i];

    const Vector3d unrot_cam_ray = unrot * camera_rays[i];
    const Vector3d unrot_reproj_pt = unrot * reprojected_point;

    const double reprojection_error =
        (unrot_cam_ray.hnormalized() - unrot_reproj_pt.hnormalized()).norm();
    good_reprojection_errors = (good_reprojection_errors &&
                                (reprojection_error < max_reprojection_error));
  }
  return good_reprojection_errors;
}

void GdlsStarTest::TestWithNoise(
    const ExpectedSolution& expected_solution,
    const TestParams& test_params,
    const TestInput& test_input) {
  // Generate testing data.
  const Input input =
      GenerateTestingData(expected_solution, test_params, test_input);

  // Estimate.
  GdlsStar::Solution solution;
  EXPECT_TRUE(estimator->EstimateSimilarityTransformation(input, &solution));

  // Check the solutions here.
  bool matched_transform = false;
  EXPECT_GT(solution.rotations.size(), 0);
  for (int i = 0; i < solution.rotations.size(); ++i) {
    const bool good_reprojection_errors =
        CheckReprojectionErrors(input,
                                solution.rotations[i],
                                solution.translations[i],
                                solution.scales[i],
                                test_params.max_reprojection_error);
    const double rotation_difference =
        expected_solution.rotation.angularDistance(solution.rotations[i]);
    const bool matched_rotation =
        rotation_difference < test_params.max_rotation_difference;
    const double translation_difference =
        (expected_solution.translation -
         solution.translations[i]).squaredNorm();
    const bool matched_translation =
        translation_difference < test_params.max_translation_difference;
    const double scale_difference =
        fabs(expected_solution.scale - solution.scales[i]);
    const bool matched_scale =
        scale_difference < test_params.max_scale_difference;
    VLOG(3) << "Matched rotation: " << matched_rotation
            << " rotation error [deg]=" << RadToDeg(rotation_difference);
    VLOG(3) << "Matched translation: " << matched_translation
            << " translation error=" << translation_difference;
    VLOG(3) << "Matched scale: " << matched_scale
            << " scale difference=" << scale_difference;
    VLOG(3) << "Good reprojection errors: " << good_reprojection_errors;
    if (matched_rotation && matched_translation && matched_scale &&
        good_reprojection_errors) {
      matched_transform = true;
    }
  }
  EXPECT_TRUE(matched_transform);
}

TEST_F(GdlsStarTest, UnconstrainedBasicTest) {
  const std::vector<Vector3d> kPoints3d = { Vector3d(-1.0, 3.0, 3.0),
                                            Vector3d(1.0, -1.0, 2.0),
                                            Vector3d(-1.0, 1.0, 2.0),
                                            Vector3d(2.0, 1.0, 3.0) };
  const std::vector<Vector3d> kImageOrigins = { Vector3d(-1.0, 0.0, 0.0),
                                                Vector3d(0.0, 0.0, 0.0),
                                                Vector3d(2.0, 0.0, 0.0),
                                                Vector3d(3.0, 0.0, 0.0) };
  // Expected solution.
  ExpectedSolution solution;
  solution.rotation =
      Quaterniond(AngleAxisd(DegToRad(13.0), Vector3d(0.0, 0.0, 1.0)));
  solution.translation = Vector3d(1.0, 1.0, 1.0);
  solution.scale = 2.5;

  // Test parameters.
  TestParams test_params;
  test_params.projection_noise_std_dev = 0.0;
  test_params.down_direction_noise_std_dev = 0.0;
  test_params.max_reprojection_error =  5.0 / 512;
  test_params.max_rotation_difference = DegToRad(1.0);
  test_params.max_translation_difference = 1e-3;
  test_params.max_scale_difference = 1e-2;

  // Test input.
  TestInput test_input;
  test_input.camera_centers = kImageOrigins;
  test_input.world_points = kPoints3d;
  test_input.query_down_direction.setZero();
  test_input.world_down_direction.setZero();
  test_input.scale_penalty = 0.0;
  test_input.gravity_penalty = 0.0;

  TestWithNoise(solution, test_params, test_input);
}

TEST_F(GdlsStarTest, UnconstrainedBasicTestWithNoise) {
  const std::vector<Vector3d> kPoints3d = { Vector3d(-1.0, 3.0, 3.0),
                                            Vector3d(1.0, -1.0, 2.0),
                                            Vector3d(-1.0, 1.0, 2.0),
                                            Vector3d(2.0, 1.0, 3.0) };
  const std::vector<Vector3d> kImageOrigins = { Vector3d(-1.0, 0.0, 0.0),
                                                Vector3d(0.0, 0.0, 0.0),
                                                Vector3d(2.0, 0.0, 0.0),
                                                Vector3d(3.0, 0.0, 0.0) };
  // Expected solution.
  ExpectedSolution solution;
  solution.rotation =
      Quaterniond(AngleAxisd(DegToRad(13.0), Vector3d(0.0, 0.0, 1.0)));
  solution.translation = Vector3d(1.0, 1.0, 1.0);
  solution.scale = 2.5;

  // Test parameters.
  TestParams test_params;
  test_params.projection_noise_std_dev = 1.0 / 512.0;
  test_params.down_direction_noise_std_dev = 0.0;
  test_params.max_reprojection_error =  5.0 / 512;
  test_params.max_rotation_difference = DegToRad(1.0);
  test_params.max_translation_difference = 1e-3;
  test_params.max_scale_difference = 1e-2;

  // Test input.
  TestInput test_input;
  test_input.camera_centers = kImageOrigins;
  test_input.world_points = kPoints3d;
  test_input.query_down_direction.setZero();
  test_input.world_down_direction.setZero();
  test_input.scale_penalty = 0.0;
  test_input.gravity_penalty = 0.0;

  TestWithNoise(solution, test_params, test_input);
}

TEST_F(GdlsStarTest, GravityRegularizedBasicTest) {
  const std::vector<Vector3d> kPoints3d = { Vector3d(-1.0, 3.0, 3.0),
                                            Vector3d(1.0, -1.0, 2.0),
                                            Vector3d(-1.0, 1.0, 2.0),
                                            Vector3d(2.0, 1.0, 3.0) };
  const std::vector<Vector3d> kImageOrigins = { Vector3d(-1.0, 0.0, 0.0),
                                                Vector3d(0.0, 0.0, 0.0),
                                                Vector3d(2.0, 0.0, 0.0),
                                                Vector3d(3.0, 0.0, 0.0) };
  // Expected solution.
  ExpectedSolution solution;
  solution.rotation = Quaterniond(
      AngleAxisd(DegToRad(13.0), Vector3d(1.0, 1.0, 1.0).normalized()));
  solution.translation = Vector3d(1.0, 1.0, 1.0);
  solution.scale = 2.5;

  // Test parameters.
  TestParams test_params;
  test_params.projection_noise_std_dev = 0.0;
  test_params.down_direction_noise_std_dev = 0.0;
  test_params.max_reprojection_error =  5.0 / 512;
  test_params.max_rotation_difference = DegToRad(1.0);
  test_params.max_translation_difference = 1e-3;
  test_params.max_scale_difference = 1e-2;

  // Test input.
  TestInput test_input;
  test_input.camera_centers = kImageOrigins;
  test_input.world_points = kPoints3d;
  test_input.query_down_direction.setZero();
  test_input.world_down_direction = Vector3d::UnitZ();
  test_input.scale_penalty = 0.0;
  test_input.gravity_penalty = 1.0;

  TestWithNoise(solution, test_params, test_input);
}

TEST_F(GdlsStarTest, GravityRegularizedBasicTestWithNoise) {
  const std::vector<Vector3d> kPoints3d = { Vector3d(-1.0, 3.0, 3.0),
                                            Vector3d(1.0, -1.0, 2.0),
                                            Vector3d(-1.0, 1.0, 2.0),
                                            Vector3d(2.0, 1.0, 3.0) };
  const std::vector<Vector3d> kImageOrigins = { Vector3d(-1.0, 0.0, 0.0),
                                                Vector3d(0.0, 0.0, 0.0),
                                                Vector3d(2.0, 0.0, 0.0),
                                                Vector3d(3.0, 0.0, 0.0) };
  // Expected solution.
  ExpectedSolution solution;
  solution.rotation =
      Quaterniond(
          AngleAxisd(DegToRad(13.0), Vector3d(1.0, 1.0, 1.0).normalized()));
  solution.translation = Vector3d(1.0, 1.0, 1.0);
  solution.scale = 2.5;

  // Test parameters.
  TestParams test_params;
  test_params.projection_noise_std_dev = 1.0 / 512.0;
  test_params.down_direction_noise_std_dev = 1.0 / 512.0;
  test_params.max_reprojection_error =  5.0 / 512;
  test_params.max_rotation_difference = DegToRad(1.0);
  test_params.max_translation_difference = 1e-2;
  test_params.max_scale_difference = 1e-2;

  // Test input.
  TestInput test_input;
  test_input.camera_centers = kImageOrigins;
  test_input.world_points = kPoints3d;
  test_input.query_down_direction.setZero();
  test_input.world_down_direction = Vector3d::UnitZ();
  test_input.scale_penalty = 0.0;
  test_input.gravity_penalty = 1.0;

  TestWithNoise(solution, test_params, test_input);
}

TEST_F(GdlsStarTest, ScaleRegularizedBasicTest) {
  const std::vector<Vector3d> kPoints3d = { Vector3d(-1.0, 3.0, 3.0),
                                            Vector3d(1.0, -1.0, 2.0),
                                            Vector3d(-1.0, 1.0, 2.0),
                                            Vector3d(2.0, 1.0, 3.0) };
  const std::vector<Vector3d> kImageOrigins = { Vector3d(-1.0, 0.0, 0.0),
                                                Vector3d(0.0, 0.0, 0.0),
                                                Vector3d(2.0, 0.0, 0.0),
                                                Vector3d(3.0, 0.0, 0.0) };
  // Expected solution.
  ExpectedSolution solution;
  solution.rotation = Quaterniond(
      AngleAxisd(DegToRad(13.0), Vector3d(1.0, 0.5, 1.0).normalized()));
  solution.translation = Vector3d(1.0, 1.0, 1.0);
  solution.scale = 2.5;

  // Test parameters.
  TestParams test_params;
  test_params.projection_noise_std_dev = 0.0;
  test_params.down_direction_noise_std_dev = 0.0;
  test_params.max_reprojection_error =  5.0 / 512;
  test_params.max_rotation_difference = DegToRad(1.0);
  test_params.max_translation_difference = 1e-3;
  test_params.max_scale_difference = 1e-2;

  // Test input.
  TestInput test_input;
  test_input.camera_centers = kImageOrigins;
  test_input.world_points = kPoints3d;
  test_input.query_down_direction.setZero();
  test_input.world_down_direction.setZero();
  test_input.scale_penalty = 1.0;
  test_input.gravity_penalty = 0.0;
  test_input.scale_prior = solution.scale;

  TestWithNoise(solution, test_params, test_input);
}

TEST_F(GdlsStarTest, ScaleRegularizedBasicTestWithNoise) {
  const std::vector<Vector3d> kPoints3d = { Vector3d(-1.0, 3.0, 3.0),
                                            Vector3d(1.0, -1.0, 2.0),
                                            Vector3d(-1.0, 1.0, 2.0),
                                            Vector3d(2.0, 1.0, 3.0) };
  const std::vector<Vector3d> kImageOrigins = { Vector3d(-1.0, 0.0, 0.0),
                                                Vector3d(0.0, 0.0, 0.0),
                                                Vector3d(2.0, 0.0, 0.0),
                                                Vector3d(3.0, 0.0, 0.0) };
  // Expected solution.
  ExpectedSolution solution;
  solution.rotation = Quaterniond(
      AngleAxisd(DegToRad(13.0), Vector3d(1.0, 0.5, 1.0).normalized()));
  solution.translation = Vector3d(1.0, 1.0, 1.0);
  solution.scale = 2.5;

  // Test parameters.
  TestParams test_params;
  test_params.projection_noise_std_dev = 1.0 / 512;
  test_params.down_direction_noise_std_dev = 0.0;
  test_params.max_reprojection_error =  5.0 / 512;
  test_params.max_rotation_difference = DegToRad(1.0);
  test_params.max_translation_difference = 1e-2;
  test_params.max_scale_difference = 1e-2;

  // Test input.
  TestInput test_input;
  test_input.camera_centers = kImageOrigins;
  test_input.world_points = kPoints3d;
  test_input.query_down_direction.setZero();
  test_input.world_down_direction.setZero();
  test_input.scale_penalty = 1.0;
  test_input.gravity_penalty = 0.0;
  test_input.scale_prior = solution.scale;

  TestWithNoise(solution, test_params, test_input);
}

TEST_F(GdlsStarTest, ScaleAndGravityRegularizedBasicTestWithNoise) {
  const std::vector<Vector3d> kPoints3d = { Vector3d(-1.0, 3.0, 3.0),
                                            Vector3d(1.0, -1.0, 2.0),
                                            Vector3d(-1.0, 1.0, 2.0),
                                            Vector3d(2.0, 1.0, 3.0) };
  const std::vector<Vector3d> kImageOrigins = { Vector3d(-1.0, 0.0, 0.0),
                                                Vector3d(0.0, 0.0, 0.0),
                                                Vector3d(2.0, 0.0, 0.0),
                                                Vector3d(3.0, 0.0, 0.0) };
  // Expected solution.
  ExpectedSolution solution;
  solution.rotation = Quaterniond(
      AngleAxisd(DegToRad(13.0), Vector3d(1.0, 0.5, 1.0).normalized()));
  solution.translation = Vector3d(1.0, 1.0, 1.0);
  solution.scale = 2.5;

  // Test parameters.
  TestParams test_params;
  test_params.projection_noise_std_dev = 1.0 / 512;
  test_params.down_direction_noise_std_dev = 1.0 / 512;
  test_params.max_reprojection_error =  5.0 / 512;
  test_params.max_rotation_difference = DegToRad(1.0);
  test_params.max_translation_difference = 1e-3;
  test_params.max_scale_difference = 1e-2;

  // Test input.
  TestInput test_input;
  test_input.camera_centers = kImageOrigins;
  test_input.world_points = kPoints3d;
  test_input.query_down_direction.setZero();
  test_input.world_down_direction = Vector3d::UnitZ();
  test_input.scale_penalty = 1.0;
  test_input.gravity_penalty = 1.0;
  test_input.scale_prior = solution.scale;

  TestWithNoise(solution, test_params, test_input);
}

TEST_F(GdlsStarTest, ScaleAndGravityRegularizedTestWithNoiseAndManyPoints) {
  // Sets some test rotations and translations.
  constexpr int kNumAxes = 8;
  constexpr int kNumPointSets = 3;
  const std::vector<Vector3d> kImageOrigins = { Vector3d(-1.0, 0.0, 0.0),
                                                Vector3d(0.0, 0.0, 0.0),
                                                Vector3d(2.0, 0.0, 0.0),
                                                Vector3d(3.0, 0.0, 0.0) };

  const std::array<int, kNumPointSets> num_points = { 100, 500, 1000 };
  std::uniform_real_distribution angle_distribution(0.0, RadToDeg(33));
  std::uniform_real_distribution scale_distribution(0.0, RadToDeg(50));
  std::uniform_real_distribution point_distribution(0.0, 1.0);
  std::vector<Vector3d> points_3d;

  int test_trial = 0;
  for (int i = 0; i < kNumAxes; ++i) {
    const Quaterniond soln_rotation(
        AngleAxisd(angle_distribution(*prng), Vector3d::Random().normalized()));
    for (int j = 0; j < kNumPointSets; ++j) {
      points_3d.clear();
      points_3d.reserve(num_points[j]);
      for (int k = 0; k < num_points[j]; ++k) {
        points_3d.emplace_back(10 * point_distribution(*prng) - 5,
                               10 * point_distribution(*prng) - 5,
                               8 * point_distribution(*prng) + 2);
      }

      // Expected solution.
      ExpectedSolution solution;
      solution.rotation = soln_rotation;
      solution.translation = Vector3d::Random();
      solution.scale = scale_distribution(*prng);

      // Test parameters.
      TestParams test_params;
      test_params.projection_noise_std_dev = 1.0 / 512;
      test_params.down_direction_noise_std_dev = 1.0 / 512;
      test_params.max_reprojection_error =  8.0 / 512;
      test_params.max_rotation_difference = DegToRad(2.0);
      test_params.max_translation_difference = 1e-2;
      test_params.max_scale_difference = 0.1;

      // Test input.
      TestInput test_input;
      test_input.camera_centers = kImageOrigins;
      test_input.world_points = std::move(points_3d);
      test_input.query_down_direction.setZero();
      test_input.world_down_direction = Vector3d::UnitZ();
      test_input.scale_penalty = 1.0;
      test_input.gravity_penalty = 1.0;
      test_input.scale_prior = solution.scale;

      VLOG(3) << ">>>> Test trial: " << ++test_trial;
      TestWithNoise(solution, test_params, test_input);
    }
  }

}

}  // namespace
}  // namespace msft


#pragma once

#include "generator/feature_builder.hpp"

#include <iterator>

#include <boost/geometry.hpp>

namespace generator
{
namespace boost_helpers
{

template <typename BoostGeometry, typename FbGeometry>
void FillBoostGeometry(BoostGeometry & geometry, FbGeometry const & fbGeometry)
{
  using BoostPoint = typename boost::geometry::point_type<BoostGeometry>::type;

  geometry.clear();
  geometry.reserve(fbGeometry.size());
  for (auto const & p : fbGeometry)
    boost::geometry::append(geometry, BoostPoint{p.x, p.y});
}

template <typename BoostPolygon>
void FillPolygon(BoostPolygon & polygon, feature::FeatureBuilder const & fb)
{
  auto const & fbGeometry = fb.GetGeometry();
  CHECK(!fbGeometry.empty(), ());

  polygon.clear();
  FillBoostGeometry(polygon.outer(), *fbGeometry.begin());
  polygon.inners().resize(fbGeometry.size() - 1);

  size_t i = 0;

  for (auto it = std::next(fbGeometry.begin()); it != fbGeometry.end(); ++it)
    FillBoostGeometry(polygon.inners()[i++], *it);

  boost::geometry::correct(polygon);
}
}  // namespace boost_helpers
}  // namespace generator


/*
 * Copyright (c) 2015 Shihao Ji and Hyokun Yun. All Rights Reserved.
 *
 * Licensed to the Apache Software Foundation (ASF) under one
 * or more contributor license agreements.  See the NOTICE file
 * distributed with this work for additional information
 * regarding copyright ownership.  The ASF licenses this file
 * to you 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.
 *
 * For more information, bug reports, fixes, contact:
 *   Shihao Ji  (shihaoji@yahoo.com)
 *   Hyokun Yun (yungilbert@gmail.com)
 */
#ifndef __WORDRANK_MODEL_HPP
#define __WORDRANK_MODEL_HPP

#include <queue>
#include <map>
#include <unordered_set>
#include <set>
#include <sstream>
#include <omp.h>
#include <boost/date_time/posix_time/posix_time.hpp>
#include <mm_malloc.h>
#include "parameter.hpp"
#include "data.hpp"

#ifdef USE_MKL
#include "mkl.h"
#endif

using namespace boost::posix_time;
using std::cout;
using std::flush;
using std::endl;
using std::pair;
using std::priority_queue;
using std::map;
using std::set;
using std::min;
using std::max;
using std::max_element;

namespace wordrank {

struct aux_info {
    int col_index;
    int row_index;
    scalar xi;
    scalar wi;
};

scalar logistic_loss(scalar x) {
    return log2f(1.f + powf(2.f, x));
}

scalar hinge_loss(scalar x) {
    return x > -1.f ? x + 1.f : 0.f;
}

scalar logistic_der(scalar x) {
    return 1.f / (1.f + powf(2.f, -x));
}

scalar hinge_der(scalar x) {
    return x > -1.f ? 1.f : 0.f;
}

class Model {

private:
    Parameter& param_;
    Data& data_;
    rng_type rng_;

public:
    Model(Parameter& param, Data& data) : param_(param), data_(data), rng_(param.random_seed_) {}

    ~Model() {}

    void run();

private:
    scalar calc_loss(scalar x) {
        if (param_.loss_type_ == LossType::LOGISTIC) {
            return logistic_loss(x);
        } else if (param_.loss_type_ == LossType::HINGE) {
            return hinge_loss(x);
        } else {
            //cerr << "unsupported loss type" << endl;
            return 1000000.f;
            //exit(1);
        }
    }

    scalar calc_der(scalar x) {
        if (param_.loss_type_ == LossType::LOGISTIC) {
            return logistic_der(x);
        } else if (param_.loss_type_ == LossType::HINGE) {
            return hinge_der(x);
        } else {
            //cerr << "unsupported loss type" << endl;
            return 1000000.f;
            //exit(1);
        }
    }

    void normalize(scalar* matrix, int m, int n) {
        for (int i = 0; i < m; i++) {
            int offset = i * n;
            scalar norm2 = 0.f;
            #pragma simd
            for (int j = 0; j < n; j++) {
                scalar v = matrix[offset + j];
                norm2 += v * v;
            }
            scalar norm2r = 1.f / sqrtf(norm2);
            if (norm2 > 1e-6f) {
                #pragma simd
                for (int j = 0; j < n; j++) matrix[offset + j] *= norm2r;
            }
        }
    }

};

}

void writelog(const char *str) {
    cout << "\n" << second_clock::local_time() << " - " << str << endl << flush;
}

void wordrank::Model::run() {

    writelog("start run()");

    // define necessary constants
    const int rank = param_.rank_;
    const int dim = param_.latent_dim_;
    const scalar alpha = param_.alpha_;
    const scalar beta = param_.beta_;
    const scalar tl = param_.t_;
    const scalar tau = param_.tau_;
    const int num_rows = data_.get_num_rows();
    const int num_cols = data_.get_num_cols();

    const int numtasks = param_.numtasks_;
    const int numthreads = param_.numthreads_;
    const int numparts = numthreads * numtasks;
    const int numrows_per_part = data_.numrows_per_part_;
    const int numcols_per_part = data_.numcols_per_part_;
    const int num_rows_upbd = numthreads * numrows_per_part * numtasks;
    const int num_cols_upbd = numthreads * numcols_per_part * numtasks;

    const int sgd_num = param_.sgd_num_;
    const TransformType trans_type = param_.trans_type_;

    auto& train_row_nnzs = data_.train_row_nnzs_;
    auto& train_col_nnzsum = data_.train_col_nnzsum_;
    auto& train_row_nnzsum = data_.train_row_nnzsum_;

    const size_t train_num_points = data_.train_total_nnz_;
    const double train_total_nnzsum = data_.train_total_nnzsum_;

    vector<vector<record_info>>& csc_indices = data_.csc_indices_;
    vector<vector<index_type>>& csc_ptrs = data_.csc_ptrs_;
    vector<index_type>& local_nnz = data_.local_nnz_;

    // create thread-specific RNGs
    rng_type rngs[numthreads];
    for (int i = 0; i < numthreads; i++) {
        rngs[i] = rng_type(param_.random_seed_ + 15791 * i + 373 * rank);
    }
    rng_type main_rng(param_.random_seed_ + 377 * rank);

    /////////////////////////////////////////////////////////////////////
    // Initialize Parameters
    /////////////////////////////////////////////////////////////////////

#ifdef USE_MKL
    mkl_set_num_threads(1);
#endif

    // indexed by thread, local_row_index, coordinate

    scalar *matrix_U = (scalar *) _mm_malloc(numthreads * numrows_per_part * dim * sizeof(scalar), 64);

    scalar *matrix_V = (scalar *) _mm_malloc(numthreads * numcols_per_part * dim * sizeof(scalar), 64);

    scalar *vector_Rec = (scalar *) _mm_malloc(numthreads * num_cols_upbd * sizeof(scalar), 64);

    // maps local_col_index to global_col_index
    int *col_indices = (int *) _mm_malloc(numthreads * numcols_per_part * sizeof(int), 64);

    {
        std::uniform_real_distribution<scalar> init_dist(0, 1.f);

        #pragma omp parallel for num_threads(numthreads)
        for (int thread_index = 0; thread_index < numthreads; thread_index++)
        {
            scalar *local_matrix_U = matrix_U + thread_index * numrows_per_part * dim;
            for (int j = 0; j < numrows_per_part * dim; j++) {
                local_matrix_U[j] = init_dist(rngs[thread_index]);
            }
            scalar *local_matrix_V = matrix_V + thread_index * numcols_per_part * dim;
            for (int j = 0; j < numcols_per_part * dim; j++) {
                local_matrix_V[j] = init_dist(rngs[thread_index]);
            }
            // normalize them
            normalize(local_matrix_U, numrows_per_part, dim);
            normalize(local_matrix_V, numcols_per_part, dim);

            int *local_col_indices = col_indices + thread_index * numcols_per_part;
            #pragma simd
            for (int j = 0; j < numcols_per_part; j++) {
                int col_index = (rank * numthreads + thread_index) * numcols_per_part + j;
                if (col_index < num_cols) {
                    local_col_indices[j] = col_index;
                } else {
                    local_col_indices[j] = -1;
                }
            }
        }
        // set the values of overflowed rows and cols to zeros
        int numRowsOverflow = (rank + 1) * numthreads * numrows_per_part - num_rows;
        if (numRowsOverflow > 0)
            memset(matrix_U + numthreads * numrows_per_part - numRowsOverflow, 0, numRowsOverflow * sizeof(scalar));

        int numColsOverflow = (rank + 1) * numthreads * numcols_per_part - num_cols;
        if (numColsOverflow > 0)
            memset(matrix_V + numthreads * numcols_per_part - numColsOverflow, 0, numColsOverflow * sizeof(scalar));

    }

    scalar *global_matrix_U = NULL;
    if (rank == 0) {
        global_matrix_U = (scalar *) _mm_malloc(num_rows_upbd * dim * sizeof(scalar), 64);
    }

    scalar *global_matrix_V = (scalar *) _mm_malloc(num_cols_upbd * dim * sizeof(scalar), 64);

    int *global_col_indices = (int *) _mm_malloc(num_cols_upbd * sizeof(int), 64);

    MPI_Allgather(matrix_V, numthreads * numcols_per_part * dim, MPI_SCALAR, global_matrix_V,
            numthreads * numcols_per_part * dim, MPI_SCALAR, MPI_COMM_WORLD);
    MPI_Allgather(col_indices, numthreads * numcols_per_part, MPI_INT, global_col_indices,
            numthreads * numcols_per_part, MPI_INT, MPI_COMM_WORLD);

    /////////////////////////////////////////////////////////////////////
    // Optimization
    /////////////////////////////////////////////////////////////////////

    int *col_locations = (int *) _mm_malloc(num_cols * sizeof(int), 64);

    // prepare auxiliary information variables
    vector<vector<aux_info>> auxs(numthreads);
    vector<vector<vector<aux_info *>>> rowwise_auxptrs(numthreads, vector<vector<aux_info *>>(numrows_per_part, vector<aux_info *>()));

    #pragma omp parallel for num_threads(numthreads)
    for (int thread_index = 0; thread_index < numthreads; thread_index++) {
        auxs[thread_index].resize(local_nnz[thread_index]);

        aux_info *ptr = &auxs[thread_index][0];
        for (int col_index = 0; col_index < num_cols; col_index++) {
            for (int j = csc_ptrs[thread_index][col_index]; j < csc_ptrs[thread_index][col_index + 1]; j++) {
                record_info& ri = csc_indices[thread_index][j];
                ptr->col_index = col_index;
                ptr->row_index = ri.row_index;
                ptr->xi = (trans_type == TransformType::RHO2) ? 0.5f : 1.0f;
                ptr->wi = ri.weight;
                rowwise_auxptrs[thread_index][ri.row_index].push_back(ptr);
                ptr++;
            }
        }
    }
    // release memory
    csc_indices.clear();

    const double stepsize = param_.learning_rate_ * train_num_points / train_total_nnzsum;
    const double reg = param_.regularization_ / ((num_cols - 1) * train_num_points);

    writelog("starts optimization");

    int *global_perm = (int *) _mm_malloc(num_cols_upbd * sizeof(int), 64);
    int *local_perm = (int *) _mm_malloc(numthreads * numcols_per_part * sizeof(int), 64);

    double cumul_computation_time = 0.0;

    /*************************************************************************/
    /* Here is the main loop                                                 */
    /*************************************************************************/
    for (int iter_num = 0; iter_num < param_.max_iteration_; iter_num++) {

        std::stringstream monitor_stream;
        monitor_stream << iter_num << ", " << param_;

        cout << "iteration: " << iter_num << endl << flush;

        if ((param_.dump_prefix_.length() > 0) && (iter_num % param_.dump_period_ == 0)) {

            cout << "dumping data" << endl << flush;

            if (rank == 0) {

                MPI_Gather(matrix_U, numthreads * numrows_per_part * dim, MPI_SCALAR, global_matrix_U,
                        numthreads * numrows_per_part * dim, MPI_SCALAR, 0, MPI_COMM_WORLD);

                std::ofstream ofile(param_.dump_prefix_ + "_word_" + boost::lexical_cast<std::string>(iter_num) + ".txt");

                scalar *user_param = global_matrix_U;
                for (int user_index = 0; user_index < num_rows; user_index++) {
                    int global_row_index = data_.row_perm_inv_[user_index];
                    if (data_.useVocab_){
                        ofile << data_.vocab_[global_row_index];
                    } else {
                        ofile << global_row_index + 1;
                    }

                    for (int i = 0; i < dim; i++) {
                        ofile << " ";
                        ofile << user_param[i];
                    }
                    ofile << endl;
                    user_param += dim;
                }
                ofile.close();
            } else {
                MPI_Gather(matrix_U, numthreads * numrows_per_part * dim, MPI_SCALAR, global_matrix_U,
                        numthreads * numrows_per_part * dim, MPI_SCALAR, 0, MPI_COMM_WORLD);
            }

            if (rank == 0) {
                std::ofstream ofile(param_.dump_prefix_ + "_context_" + boost::lexical_cast<std::string>(iter_num) + ".txt");
                scalar *item_param = global_matrix_V;
                for (int item_index = 0; item_index < num_cols_upbd; item_index++) {
                    int global_col_index = global_col_indices[item_index];

                    if (global_col_index > -1) {
                        if (data_.useVocab_) {
                            ofile << data_.vocab_[global_col_index];
                        } else {
                            ofile << global_col_index + 1;
                        }

                        for (int i = 0; i < dim; i++) {
                            ofile << " ";
                            ofile << item_param[i];
                        }
                        ofile << endl;
                    }
                    item_param += dim;
                }
                ofile.close();
            }
        }

        double compute_start_time = omp_get_wtime();

        // sample a new assignment of column indices
        if (rank == 0) {
            std::iota(global_perm, global_perm + num_cols_upbd, 0);
            std::shuffle(global_perm, global_perm + num_cols_upbd, main_rng);
        }
        MPI_Scatter(global_perm, numcols_per_part * numthreads, MPI_INT, local_perm, numcols_per_part * numthreads,
                MPI_INT, 0, MPI_COMM_WORLD);

        // copy global parameters to local space
        {
            #pragma omp parallel for num_threads(numthreads)
            for (int i = 0; i < numcols_per_part * numthreads; i++) {
                // copy the index
                col_indices[i] = global_col_indices[local_perm[i]];

                // copy the parameter value
                scalar *source_ptr = global_matrix_V + local_perm[i] * dim;
                scalar *target_ptr = matrix_V + i * dim;
                std::copy(source_ptr, source_ptr + dim, target_ptr);
            }
        }

        // inside a local machine, we use a permutation matrix
        // to determine the order of block assignment
        vector<vector<int> > perm_matrix(numthreads, vector<int>(numthreads, 0));
        {
            vector<int> col_perm(numthreads, 0);

            #pragma omp parallel for num_threads(numthreads)
            for (int i = 0; i < numthreads; i++) {
                col_perm[i] = i;
                #pragma simd
                for (int j = 0; j < numthreads; j++) {
                    perm_matrix[i][j] = (i + j) % numthreads;
                }
            }
            shuffle(perm_matrix.begin(), perm_matrix.end(), main_rng);
            shuffle(col_perm.begin(), col_perm.end(), main_rng);

            // apply column permutation
            #pragma omp parallel for num_threads(numthreads)
            for (int i = 0; i < numthreads; i++) {
                vector<int> tmp_buffer(numthreads, 0);
                std::copy(perm_matrix[i].begin(), perm_matrix[i].end(), tmp_buffer.begin());
                #pragma simd
                for (int j = 0; j < numthreads; j++) {
                    perm_matrix[i][col_perm[j]] = tmp_buffer[j];
                }
            }
        }

        // run SGD
        for (int inner_iter = 0; inner_iter < numthreads; inner_iter++) {

            #pragma omp parallel for num_threads(numthreads)
            for (int thread_index = 0; thread_index < numthreads; thread_index++)
            {
                int chunk_index = perm_matrix[inner_iter][thread_index];

                int *local_col_indices = col_indices + chunk_index * numcols_per_part;
                scalar *local_matrix_U = matrix_U + thread_index * numrows_per_part * dim;
                scalar *local_matrix_V = matrix_V + chunk_index * numcols_per_part * dim;

                std::uniform_int_distribution<int> init_dist(0, numcols_per_part - 1);
                auto sample_column = [&](int local_col_index)->int {
                    int ret = local_col_index;
                    while (true) {
                        ret = init_dist(rngs[thread_index]);
                        if ((ret != local_col_index) && (local_col_indices[ret] != -1)) {
                            return ret;
                        }
                    }
                };

                scalar backup_row_vec[dim];

                if (param_.loss_type_ == LossType::LOGISTIC) {

                    for (int sgd_iter = 0; sgd_iter < sgd_num; sgd_iter++) {

                        for (int i = 0; i < numcols_per_part; i++) {
                            int col_index = local_col_indices[i];
                            scalar *col_vec = local_matrix_V + i * dim;

                            if (col_index < 0) {
                                continue;
                            }

                            for (int j = csc_ptrs[thread_index][col_index]; j < csc_ptrs[thread_index][col_index + 1]; j++) {
                                aux_info& aux_ptr = auxs[thread_index][j];
                                int row_index = aux_ptr.row_index;
                                scalar xi = aux_ptr.xi;
                                scalar wi = aux_ptr.wi;
                                scalar *row_vec = local_matrix_U + row_index * dim;

                                int sam_local_index = sample_column(i);
                                int sam_index = local_col_indices[sam_local_index];
                                scalar *sam_vec = local_matrix_V + sam_local_index * dim;

                                std::copy(row_vec, row_vec + dim, backup_row_vec);

                                scalar org_dot = std::inner_product(row_vec, row_vec + dim, col_vec, 0.f);
                                scalar sam_dot = std::inner_product(row_vec, row_vec + dim, sam_vec, 0.f);

                                scalar coef = 0.f;

                                switch (trans_type) {
                                case TransformType::RHO1:
                                    coef = wi * xi * INV_LOG2 / (tau + powf(2.f, org_dot - sam_dot));
                                    break;
                                case TransformType::RHO2: {
                                    scalar logxi = logf(xi);
                                    scalar logxi2 = logxi * logxi;
                                    coef = wi * xi * LOG2 / (logxi2 * (tau + powf(2.f, org_dot - sam_dot)));
                                    break;
                                }
                                case TransformType::RHO3:
                                    coef = wi * powf(xi, tl) / (tau + powf(2.f, org_dot - sam_dot));
                                    break;
                                default:
                                    cerr << "unsupported transform type" << endl;
                                    exit(1);
                                }

                                #pragma simd
                                for (int t = 0; t < dim; t++) {
                                    row_vec[t] -= stepsize * (reg * train_row_nnzsum[thread_index][row_index] * row_vec[t]
                                                     + coef * (sam_vec[t] - col_vec[t]));
                                    col_vec[t] -= stepsize * (reg * train_col_nnzsum[col_index] * col_vec[t]
                                                     - coef * backup_row_vec[t]);
                                    sam_vec[t] -= stepsize * (reg * (train_total_nnzsum - train_col_nnzsum[sam_index]) * sam_vec[t]
                                                     + coef * backup_row_vec[t]);
                                }
                            }
                        }
                    } // end of sam_iter
                } // case of logistic loss (end)
                else if (param_.loss_type_ == LossType::HINGE) {

                    for (int sgd_iter = 0; sgd_iter < sgd_num; sgd_iter++) {

                        for (int i = 0; i < numcols_per_part; i++) {
                            int col_index = local_col_indices[i];
                            scalar *col_vec = local_matrix_V + i * dim;

                            if (col_index < 0) {
                                continue;
                            }

                            for (int j = csc_ptrs[thread_index][col_index]; j < csc_ptrs[thread_index][col_index + 1]; j++) {
                                aux_info& aux_ptr = auxs[thread_index][j];
                                int row_index = aux_ptr.row_index;
                                scalar xi = aux_ptr.xi;
                                scalar wi = aux_ptr.wi;
                                scalar *row_vec = local_matrix_U + row_index * dim;

                                int sam_local_index = sample_column(i);
                                int sam_index = local_col_indices[sam_local_index];
                                scalar *sam_vec = local_matrix_V + sam_local_index * dim;

                                std::copy(row_vec, row_vec + dim, backup_row_vec);

                                scalar org_dot = std::inner_product(row_vec, row_vec + dim, col_vec, 0.f);
                                scalar sam_dot = std::inner_product(row_vec, row_vec + dim, sam_vec, 0.f);

                                // if hinge loss is nonzero
                                if (sam_dot > org_dot - tau) {

                                    scalar coef = 0.f;

                                    switch (trans_type) {
                                    case TransformType::RHO1:
                                        coef = wi * xi * INV_LOG2;
                                        break;
                                    case TransformType::RHO2: {
                                        scalar logxi = logf(xi);
                                        scalar logxi2 = logxi * logxi;
                                        coef = wi * xi * LOG2 / logxi2;
                                        break;
                                    }
                                    case TransformType::RHO3:
                                        coef = wi * powf(xi, tl);
                                        break;
                                    default:
                                        cerr << "unsupported transform type" << endl;
                                        exit(1);
                                    }

                                    #pragma simd
                                    for (int t = 0; t < dim; t++) {
                                        row_vec[t] -= stepsize * (reg * train_row_nnzsum[thread_index][row_index] * row_vec[t]
                                                        + coef * (sam_vec[t] - col_vec[t]));
                                        col_vec[t] -= stepsize * (reg * train_col_nnzsum[col_index] * col_vec[t]
                                                        - coef * backup_row_vec[t]);
                                        sam_vec[t] -= stepsize * (reg * (train_total_nnzsum - train_col_nnzsum[sam_index]) * sam_vec[t]
                                                        + coef * backup_row_vec[t]);
                                    }
                                }
                            }
                        }
                    }
                } else {
                    cerr << "unsupported loss type 1" << endl << flush;
                    exit(1);
                }
            }
        }

        // at the end of the iteration, synchronize parameters
        MPI_Allgather(matrix_V, numthreads * numcols_per_part * dim, MPI_SCALAR, global_matrix_V,
                numthreads * numcols_per_part * dim, MPI_SCALAR, MPI_COMM_WORLD);
        MPI_Allgather(col_indices, numthreads * numcols_per_part, MPI_INT, global_col_indices,
                numthreads * numcols_per_part, MPI_INT, MPI_COMM_WORLD);

        #pragma omp parallel for num_threads(numthreads)
        #pragma simd
        for (int i = 0; i < num_cols_upbd; i++) {
            if (global_col_indices[i] >= 0) {
                col_locations[global_col_indices[i]] = i;
            }
        }

        // update auxiliary parameters exactly
        if (iter_num > 0 && iter_num % param_.xi_period_ == 0) {

            writelog("update auxiliary parameters");

            vector<scalar> max_xi(numthreads, -1000000.f);
            vector<scalar> min_xi(numthreads, 1000000.f);
            vector<scalar> mean_xi(numthreads, 0.f);

//            cblas_sgemm(CblasRowMajor, CblasNoTrans, CblasTrans, numthreads * numrows_per_part, num_cols_upbd,
//                    dim, 1.0f, matrix_U, dim, global_matrix_V, dim, 0.0f, matrix_Rec, num_cols_upbd);

            #pragma omp parallel for num_threads(numthreads)
            for (int thread_index = 0; thread_index < numthreads; thread_index++)
            {
                scalar* vector_Rec_ptr = vector_Rec + thread_index * num_cols_upbd;
                scalar* matrix_U_ptr = matrix_U + thread_index * numrows_per_part * dim;

                for (int row_index = 0; row_index < numrows_per_part; row_index++) {
#ifndef USE_MKL
                    scalar* row_ptr = matrix_U_ptr + row_index * dim;
                    for (int i = 0; i < num_cols_upbd; i++) {
                        scalar* col_ptr = global_matrix_V + i * dim;
                        scalar sum = 0.f;
                        #pragma simd
                        for (int j = 0; j < dim; j++) {
                            sum += row_ptr[j] * col_ptr[j];
                        }
                        vector_Rec_ptr[i] = sum;
                    }
#else
                    cblas_sgemv(CblasRowMajor, CblasNoTrans, num_cols_upbd, dim, 1.0f, global_matrix_V, dim,
                            matrix_U_ptr + row_index * dim, 1, 0.0f, vector_Rec_ptr, 1);
#endif
                    vector<aux_info *>& auxptrs = rowwise_auxptrs[thread_index][row_index];
                    int numnnzs = auxptrs.size();

                    for (int i = 0; i < numnnzs; i++) {
                        aux_info* ptr = auxptrs[i];
                        int col_index = ptr->col_index;
                        scalar wi = ptr->wi;
                        scalar val = vector_Rec_ptr[col_locations[col_index]];
                        double sum = 0;
                        #pragma simd
                        for (int j = 0; j < num_cols_upbd; j++) {
                            scalar diff = vector_Rec_ptr[j] - val;
                            sum += diff > -tau ? (tau + diff) : 0.;
                        }
                        ptr->xi = (trans_type == TransformType::RHO2) ? (0.5 * beta + 1) / (sqrt(sum) + 1 + beta) :
                                alpha / (sum + beta);
                    }

                    #pragma simd
                    for (int i = 0; i < numnnzs; i++) {
                        aux_info* ptr = auxptrs[i];
                        mean_xi[thread_index] += ptr->xi;
                        if (max_xi[thread_index] < ptr->xi)
                            max_xi[thread_index] = ptr->xi;
                        if (min_xi[thread_index] > ptr->xi)
                            min_xi[thread_index] = ptr->xi;
                    }

                }
            }

            scalar machine_max_xi = *max_element(max_xi.begin(), max_xi.end());
            scalar machine_min_xi = *min_element(min_xi.begin(), min_xi.end());
            scalar machine_mean_xi = std::accumulate(mean_xi.begin(), mean_xi.end(), 0.f) * numtasks / train_num_points;

            scalar global_max_xi = -1000000.f;
            scalar global_min_xi = 1000000.f;
            scalar global_mean_xi = 0.f;

            MPI_Allreduce(&machine_max_xi, &global_max_xi, 1, MPI_SCALAR, MPI_MAX, MPI_COMM_WORLD);
            MPI_Allreduce(&machine_min_xi, &global_min_xi, 1, MPI_SCALAR, MPI_MIN, MPI_COMM_WORLD);
            MPI_Allreduce(&machine_mean_xi, &global_mean_xi, 1, MPI_SCALAR, MPI_SUM, MPI_COMM_WORLD);

            global_mean_xi /= numtasks;

            cout << "machine min xi: " << machine_min_xi << ", machine max xi: " << machine_max_xi
                    << ", machine mean xi: " << machine_mean_xi << ", global mean xi: " << global_mean_xi << endl << flush;
        }

        double computation_time = omp_get_wtime() - compute_start_time;
        cumul_computation_time += computation_time;
        monitor_stream << ", " << cumul_computation_time;

        // calculate the objective function
        if (param_.cost_eval_)
        {
            writelog("calculate cost function");

            vector<double> local_loss_sums(numthreads, 0);
            vector<double> local_approx_loss_sums(numthreads, 0);
            vector<double> local_reg_sums(numthreads, 0);

            #pragma omp parallel for num_threads(numthreads)
            for (int thread_index = 0; thread_index < numthreads; thread_index++)
            {
                scalar* matrix_U_ptr = matrix_U + thread_index * numrows_per_part * dim;
                scalar* vector_Rec_ptr = vector_Rec + thread_index * num_cols_upbd;

                for (int n = 0; n < local_nnz[thread_index]; n++) {
                    aux_info& aux = auxs[thread_index][n];
                    int row_index = aux.row_index;
                    int col_index = aux.col_index;
                    scalar xi = aux.xi;
                    scalar wi = aux.wi;

#ifndef USE_MKL
                    scalar* row_ptr = matrix_U_ptr + row_index * dim;
                    for (int i = 0; i < num_cols_upbd; i++) {
                        scalar* col_ptr = global_matrix_V + i * dim;
                        scalar sum = 0.f;
                        #pragma simd
                        for (int j = 0; j < dim; j++) {
                            sum += row_ptr[j] * col_ptr[j];
                        }
                        vector_Rec_ptr[i] = sum;
                    }
#else
                    cblas_sgemv(CblasRowMajor, CblasNoTrans, num_cols_upbd, dim, 1.f, global_matrix_V, dim,
                            matrix_U_ptr + row_index * dim, 1, 0.f, vector_Rec_ptr, 1);
#endif

                    scalar val = vector_Rec_ptr[col_locations[col_index]];
                    scalar sum = 0.f;
                    #pragma simd
                    for (int j = 0; j < num_cols_upbd; j++) {
                        scalar diff = vector_Rec_ptr[j] - val;
                        sum += diff > -tau ? (tau + diff) : 0.f;
                    }
                    if (trans_type == TransformType::RHO1) {
                        local_loss_sums[thread_index] += wi * log2(sum);
                        local_approx_loss_sums[thread_index] += wi * ((xi * (sum + beta) - 1) * INV_LOG2 - alpha * log2(xi));
                    } else {
                        scalar logxi = logf(xi);
                        scalar logxi2 = logxi * logxi;
                        local_loss_sums[thread_index] += wi * (1 - 1 / log2(sum + 1));
                        local_approx_loss_sums[thread_index] += wi * (1 + LOG2 * (logxi - 1 + xi * (sum + 1)) / logxi2
                                + (beta * xi - 0.5 * beta * logxi) * INV_LOG2);
                    }
                }

                // calculate the regularizer term
                if (param_.regularization_ > 0) {
                    #pragma simd
                    for (int row_index = 0; row_index < numrows_per_part; row_index++) {
                        scalar *row_vec = matrix_U_ptr + row_index * dim;
                        local_reg_sums[thread_index] += train_row_nnzsum[thread_index][row_index]
                                * std::inner_product(row_vec, row_vec + dim, row_vec, 0.f);
                    }
                    int *local_col_indices = col_indices + thread_index * numcols_per_part;
                    scalar *local_matrix_V = matrix_V + thread_index * numcols_per_part * dim;
                    #pragma simd
                    for (int col_index = 0; col_index < numcols_per_part; col_index++) {
                        scalar *col_vec = local_matrix_V + col_index * dim;
                        local_reg_sums[thread_index] += train_col_nnzsum[local_col_indices[col_index]]
                                * std::inner_product(col_vec, col_vec + dim, col_vec, 0.f);
                    }
                }
            }

            double global_loss_sum = 0.0;
            double machine_loss_sum = std::accumulate(local_loss_sums.begin(), local_loss_sums.end(), 0.0);
            MPI_Allreduce(&machine_loss_sum, &global_loss_sum, 1, MPI_DOUBLE, MPI_SUM, MPI_COMM_WORLD);

            double global_approx_loss_sum = 0.0;
            double machine_approx_loss_sum = std::accumulate(local_approx_loss_sums.begin(), local_approx_loss_sums.end(), 0.0);
            MPI_Allreduce(&machine_approx_loss_sum, &global_approx_loss_sum, 1, MPI_DOUBLE, MPI_SUM, MPI_COMM_WORLD);

            double global_reg_sum = 0.0;
            if (param_.regularization_ > 0) {
                double machine_reg_sum = std::accumulate(local_reg_sums.begin(), local_reg_sums.end(), 0.0);
                MPI_Allreduce(&machine_reg_sum, &global_reg_sum, 1, MPI_DOUBLE, MPI_SUM, MPI_COMM_WORLD);
            }

            double cost = global_approx_loss_sum + param_.regularization_ * global_reg_sum;

            monitor_stream << ", cost: " << cost << " (" << global_approx_loss_sum << " + " << param_.regularization_ * global_reg_sum
                    << "), true loss: " << global_loss_sum;
        }

        if (rank == 0) {
            cout << "monitor, " << monitor_stream.str() << endl << flush;
        }

    } // end of SGD iteration

    /////////////////////////////////////////////////////////////////////
    // Deallocate Memory
    /////////////////////////////////////////////////////////////////////

    _mm_free(vector_Rec);
    _mm_free(matrix_U);
    _mm_free(matrix_V);
    if (rank == 0) _mm_free(global_matrix_U);
    _mm_free(global_matrix_V);
    _mm_free(global_col_indices);
    _mm_free(local_perm);
    _mm_free(global_perm);
    _mm_free(col_indices);
    _mm_free(col_locations);
}

#endif


#ifndef OCV_CBIR_CODE_BOOK_BUILDER_HPP
#define OCV_CBIR_CODE_BOOK_BUILDER_HPP

#include "features_indexer.hpp"

#include <opencv2/core.hpp>
#include <opencv2/hdf.hpp>

#include <armadillo>

#include <type_traits>

/*!
 *  \addtogroup ocv
 *  @{
 */
namespace ocv{

/*!
 *  \addtogroup cbir
 *  @{
 */
namespace cbir{

/**
 *build the code book by kmeans, depend on
 *armadillo, opencv, hdf5
 *
 *@tparam T type of code book
 *@tparam U type of the features
 */
template<typename T>
class code_book_builder
{
public:
    /**
     * @param fi the feature indexer which store the
     * hdf5 data
     * @param sample_size size of the sample
     * @param feature_type type of the features
     * @param seed random seed
     */
    code_book_builder(features_indexer const &fi,
                      int sample_size,
                      int feature_type,
                      unsigned int seed = 0) :
        feature_dimension_(fi.get_features_dimension()),
        fi_(fi)
    {
        static_assert(std::is_arithmetic<T>::value,
                      "T should be arithmetic");
        read_data(sample_size, feature_type, seed);
    }

    /**
     * @param fi the feature indexer which store the
     * hdf5 data
     * @param ratio ratio of the feature you want to sample
     * @param feature_type type of the features
     * @param seed random seed
     */
    code_book_builder(features_indexer const &fi,
                      double ratio,
                      int feature_type,
                      unsigned int seed = 0) :
        feature_dimension_(fi.get_features_dimension()),
        fi_(fi)
    {
        read_data(ratio, feature_type, seed);
    }

    /**
     * create code book
     * @param cluster cluster size, it should << than feature size
     * @param n_iter number of iteration, 10 is default number
     * @param print_mode true will print the else and vice versa
     */
    void create_code_book(arma::uword cluster,
                          arma::uword n_iter = 10,
                          bool print_mode = false)
    {
        arma::kmeans(code_book_, data_, cluster,
                     arma::random_subset, n_iter, print_mode);
    }

    /**
     * read features from hdf5
     * @param sample_size how many sample want to read
     * @param feature_type type of the features
     * @param seed random seed
     */
    void read_data(int sample_size, int feature_type,
                   unsigned int seed = 0)
    {
        data_.set_size(feature_dimension_[1], sample_size);
        switch(feature_type){
        case CV_8U :{
            fi_.read_random_features(sample_size,
                                     read_func<uchar>(data_), seed);
            break;
        }
        case CV_8S :{
            fi_.read_random_features(sample_size,
                                     read_func<schar>(data_), seed);
            break;
        }
        case CV_16U :{
            fi_.read_random_features(sample_size,
                                     read_func<ushort>(data_), seed);
            break;
        }
        case CV_16S :{
            fi_.read_random_features(sample_size,
                                     read_func<short>(data_), seed);
            break;
        }
        case CV_32S :{
            fi_.read_random_features(sample_size,
                                     read_func<int>(data_), seed);
            break;
        }
        case CV_32F :{
            fi_.read_random_features(sample_size,
                                     read_func<float>(data_), seed);
            break;
        }
        case CV_64F :{
            fi_.read_random_features(sample_size,
                                     read_func<double>(data_), seed);
            break;
        }
        default:{
            throw std::runtime_error("code_book_builder do not support this type");
        }
        }
    }

    /**
     * read features from hdf5
     * @param ratio ratio of the feature you want to sample
     * @param feature_type type of the features
     * @param seed random seed
     */
    void read_data(double ratio, int feature_type,
                   unsigned int seed = 0)
    {
        int const feature_size =
                static_cast<int>(feature_dimension_[0] * ratio);
        read_data(feature_size, feature_type, seed);
    }

    /**
     * get the code book
     * @return code book
     */
    arma::Mat<T> const& get_code_book() const
    {
        return code_book_;
    }

private:
    template<typename U>
    struct read_func
    {
        read_func(arma::Mat<T> &data) :
            data(data){}

        void operator()(cv::Mat const &features, int index)
        {
            auto *fptr = features.ptr<U>(0);
            //do not prefer std::copy since there are type cast warning
            //std::copy(fptr, fptr + features.cols, data.colptr(index));
            auto *cptr = data.colptr(index);
            for(int i = 0; i != features.cols; ++i){
                cptr[i] = static_cast<T>(fptr[i]);
            }
        }

        arma::Mat<T> &data;
    };

    arma::Mat<T> code_book_;
    arma::Mat<T> data_;
    std::vector<int> feature_dimension_;
    features_indexer const &fi_;
};

} /*! @} End of Doxygen Groups*/

} /*! @} End of Doxygen Groups*/

#endif // CODE_BOOK_BUILDER_H


#include <iostream>
#include <alps/alea.h>
#include <boost/random.hpp>

int main()
{
 //DEFINE RANDOM NUMBER GENERATOR
 typedef boost::minstd_rand0 random_base_type;
 typedef boost::uniform_01<random_base_type> random_type;
 random_base_type random_int;
 random_type random(random_int);

 //DEFINE OBSERVABLE
 alps::RealObservable obs_a("observable a");

 //ADD 1000 MEASUREMENTS TO THE OBSERVABLE
 for(int i = 0; i < 1000; ++i){
   obs_a << random();
 }

 //RESET OBSERVABLES (THERMALIZATION FINISHED)
 obs_a.reset(true);

 //ADD 10000 MEASUREMENTS TO THE OBSERVABLE
 for(int i = 0; i < 10000; ++i){
   obs_a << random();
 }

 //OUTPUT OBSERVABLE
 std::cout << obs_a;
}

V6 Out-Of & On Equilibrium/V6.9 Wolff Algorithm/ALPS/cumerr.cpp


/*    Copyright (c) 2010-2019, Delft University of Technology
 *    All rigths reserved
 *
 *    This file is part of the Tudat. Redistribution and use in source and
 *    binary forms, with or without modification, are permitted exclusively
 *    under the terms of the Modified BSD license. You should have received
 *    a copy of the license with this file. If not, please or visit:
 *    http://tudat.tudelft.nl/LICENSE.
 */

#define BOOST_TEST_DYN_LINK
#define BOOST_TEST_MAIN

#include <limits>
#include <boost/test/tools/floating_point_comparison.hpp>
#include <boost/test/unit_test.hpp>
#include <boost/bind/bind.hpp>
using namespace boost::placeholders;

#include <Eigen/Core>

#include "tudat/astro/gravitation/basicSolidBodyTideGravityFieldVariations.h"
#include "tudat/astro/gravitation/gravityFieldVariations.h"
#include "tudat/astro/gravitation/timeDependentSphericalHarmonicsGravityField.h"
#include "tudat/astro/gravitation/tabulatedGravityFieldVariations.h"
#include "tudat/interface/spice/spiceInterface.h"
#include "tudat/io/basicInputOutput.h"
#include "tudat/simulation/environment_setup/createGravityFieldVariations.h"

namespace tudat
{
namespace unit_tests
{

BOOST_AUTO_TEST_SUITE( test_gravity_field_variations )

using namespace tudat::gravitation;
using namespace tudat::spice_interface;

//! Function to get nominal gravity field coefficients for Jupiter
/*!
 * Function to get nominal gravity field coefficients for Jupiter (returned by reference)
 * \param cosineCoefficients Jupiter's cosine coefficients
 * \param sineCoefficients Jupiter's sine coefficients
 */
void getNominalJupiterGravityField(
        Eigen::MatrixXd& cosineCoefficients, Eigen::MatrixXd& sineCoefficients )
{
    // Define unnormalized coefficients.
    double jupiterJ2 = 14.736E-3;
    double jupiterJ3 = 1.4E-6;
    double jupiterJ4 = -587.0E-6;
    double jupiterJ6 = 31.0E-6;
    double jupiterc22 = -0.03E-6;
    double jupiters22 = -0.007E-6;

    cosineCoefficients.setZero( 6, 6 );
    sineCoefficients.setZero( 6, 6 );

    // Normalize coefficients and set in matrices.
    cosineCoefficients( 0, 0 ) = 1.0;
    cosineCoefficients( 2, 0 ) =  -jupiterJ2 / basic_mathematics::calculateLegendreGeodesyNormalizationFactor( 2, 0 );
    cosineCoefficients( 2, 2 ) =  jupiterc22 / basic_mathematics::calculateLegendreGeodesyNormalizationFactor( 2, 2 );
    cosineCoefficients( 3, 0 ) =  -jupiterJ3 / basic_mathematics::calculateLegendreGeodesyNormalizationFactor( 3, 0 );
    cosineCoefficients( 4, 0 ) =  -jupiterJ4 / basic_mathematics::calculateLegendreGeodesyNormalizationFactor( 4, 0 );
    cosineCoefficients( 5, 0 ) =  -jupiterJ6 / basic_mathematics::calculateLegendreGeodesyNormalizationFactor( 6, 0 );
    sineCoefficients( 2, 2 ) =  jupiters22/ basic_mathematics::calculateLegendreGeodesyNormalizationFactor( 2, 2 );
}

//! Function to get tabulated gravity field variations
/*!
 * Function to get tabulated gravity field variations, based on arbitrary values, purely for test
 * purposes.
 * \param cosineCoefficientCorrections Map of corrections to cosine coefficients as function of
 * time.
 * \param sineCoefficientCorrections Map of corrections to sine coefficients as function of
 * time.
 */
void getTabulatedGravityFieldVariationValues(
        std::map< double, Eigen::MatrixXd >& cosineCoefficientCorrections,
        std::map< double, Eigen::MatrixXd >& sineCoefficientCorrections )
{

    // Define settings for coefficient variations
    double angularFrequency = 2.0 * mathematical_constants::PI / 86400.0;
    double amplitude = 1.0E-7;
    double startTime = 0.99E7;
    double endTime = 1.01E7;
    double timeStep = 3600.0;

    // Iterate over time steps
    double currentTime = startTime;
    while( currentTime < endTime )
    {
        // Initialize correctiuons to zero
        cosineCoefficientCorrections[ currentTime ].setZero( 4, 5 );
        sineCoefficientCorrections[ currentTime ].setZero( 4, 5 );

        // Set corrections for current time step.
        for( unsigned int i = 1; i < 5; i++ )
        {
            for( unsigned int j = 0; j <= i; j++ )
            {
                cosineCoefficientCorrections[ currentTime ]( i - 1, j ) =
                        amplitude *
                        std::cos( currentTime * angularFrequency + static_cast< double >( i * j ) );
                if( j != 0 )
                {
                    sineCoefficientCorrections[ currentTime ]( i - 1, j ) =
                            amplitude *
                            std::cos( currentTime * angularFrequency -
                                      static_cast< double >( i * j ) );
                }
            }
        }
        currentTime += timeStep;
    }
}

//! Function to get predefined tabulated gravity field variations.
/*!
 * Function to get predefined tabulated gravity field variations (non-physical values, purelu
 * for test purposes) using getTabulatedGravityFieldVariationValues function
 * \return Predefined tabulated gravity field variations.
 */
std::shared_ptr< TabulatedGravityFieldVariations > getTabulatedGravityFieldVariations( )
{
    // Get coefficient tables.
    std::map< double, Eigen::MatrixXd > cosineCoefficientCorrections;
    std::map< double, Eigen::MatrixXd > sineCoefficientCorrections;
    getTabulatedGravityFieldVariationValues(
                cosineCoefficientCorrections, sineCoefficientCorrections );

    // Create correction object.
    return std::make_shared< TabulatedGravityFieldVariations >(
                cosineCoefficientCorrections, sineCoefficientCorrections, 1, 0 );
}

//! Function to get predefined gravity field variations object.
/*!
 *  Function to get predefined gravity field variations object, consisting of dummy tabulated
 *  variations from getTabulatedGravityFieldVariationValues and degree 2 tidal variations.
 * \return Predefined gravity field variations object.
 */
std::shared_ptr< GravityFieldVariationsSet > getTestGravityFieldVariations( )
{
    // Define bodies raising rides.
    std::vector< std::string > deformingBodies;
    deformingBodies.push_back( "Io" );
    deformingBodies.push_back( "Europa" );

    // Retrieve required data of bodies raising tides.
    std::vector< std::function< Eigen::Vector6d( const double ) > >
            deformingBodyStateFunctions;
    std::vector< std::function< double( ) > > deformingBodyMasses;
    for( unsigned int i = 0; i < deformingBodies.size( ); i++ )
    {
        deformingBodyStateFunctions.push_back(
                    std::bind( &getBodyCartesianStateAtEpoch,
                               deformingBodies.at( i ), "SSB", "J2000",
                               "None", std::placeholders::_1 ) );
        deformingBodyMasses.push_back(
                    std::bind( &getBodyGravitationalParameter,
                               deformingBodies.at( i ) ) );
    }

    // Define Love numbers ( constant for degree 2 only)
    std::complex< double > constantLoveNumber( 0.5, 0.5E-3 );
    std::vector< std::complex< double > > constantSingleDegreeLoveNumber( 3, constantLoveNumber );
    std::map< int, std::vector< std::complex< double > > > loveNumbers;
    loveNumbers[ 2 ] = constantSingleDegreeLoveNumber;

    // Set up gravity field variation of Jupiter due to Galilean moons.
    std::shared_ptr< GravityFieldVariations > solidBodyGravityFieldVariations =
            std::make_shared< BasicSolidBodyTideGravityFieldVariations >(
                std::bind( &getBodyCartesianStateAtEpoch,
                           "Jupiter", "SSB", "J2000", "None", std::placeholders::_1 ),
                std::bind( &computeRotationQuaternionBetweenFrames,
                           "J2000", "IAU_Jupiter", std::placeholders::_1 ),
                deformingBodyStateFunctions,
                getAverageRadius( "Jupiter" ),
                std::bind( &getBodyGravitationalParameter, "Jupiter" ),
                deformingBodyMasses, loveNumbers, deformingBodies );

    // Get tabulated gravity field variations.
    std::shared_ptr< GravityFieldVariations > tabulatedGravityFieldVariations =
            getTabulatedGravityFieldVariations( );

    // Create and return full gravity field variations object.
    return std::make_shared< GravityFieldVariationsSet >(
                std::vector< std::shared_ptr< GravityFieldVariations > >{ solidBodyGravityFieldVariations, tabulatedGravityFieldVariations },
                std::vector< BodyDeformationTypes >{ basic_solid_body, tabulated_variation },
                std::vector< std::string >{ "BasicTidal", "Tabulated" } );
}

BOOST_AUTO_TEST_CASE( testGravityFieldVariations )
{
    // Load spice kernels.
    spice_interface::loadStandardSpiceKernels( );

    // Define properties of nominal field
    double gravitationalParameter = getBodyGravitationalParameter( "Jupiter" );
    double referenceRadius = getAverageRadius( "Jupiter" );
    Eigen::MatrixXd nominalCosineCoefficients;
    Eigen::MatrixXd nominalSineCoefficients;
    getNominalJupiterGravityField( nominalCosineCoefficients, nominalSineCoefficients );

    // Create gravity field corrections object and retrieve corrections
    std::vector< std::shared_ptr< GravityFieldVariations > > gravityFieldVariationsList =
            getTestGravityFieldVariations( )->getVariationObjects( );

    // Define data structures for storing expected gravity field variations.
    std::pair< Eigen::MatrixXd, Eigen::MatrixXd > directGravityFieldVariations;
    Eigen::Matrix< double, 5, 5 > expectedCosineCoefficientsCorrections =
            Eigen::Matrix< double, 5, 5 >::Zero( );
    Eigen::Matrix< double, 5, 5 > expectedSineCoefficientsCorrections =
            Eigen::Matrix< double, 5, 5 >::Zero( );

    // Calculate gravity field variations directly.
    int minimumDegree, minimumOrder, numberOfDegrees, numberOfOrders;
    double testTime = 1.0E7;
    for( unsigned int i = 0; i < gravityFieldVariationsList.size( ); i++ )
    {
        minimumDegree = gravityFieldVariationsList.at( i )->getMinimumDegree( );
        numberOfDegrees = gravityFieldVariationsList.at( i )->getNumberOfDegrees( );

        minimumOrder = gravityFieldVariationsList.at( i )->getMinimumOrder( );
        numberOfOrders = gravityFieldVariationsList.at( i )->getNumberOfOrders( );

        directGravityFieldVariations =
                gravityFieldVariationsList.at( i )->
                calculateSphericalHarmonicsCorrections( testTime );
        expectedCosineCoefficientsCorrections.block(
                    minimumDegree, minimumOrder, numberOfDegrees, numberOfOrders ) +=
                directGravityFieldVariations.first;
        expectedSineCoefficientsCorrections.block(
                    minimumDegree, minimumOrder, numberOfDegrees, numberOfOrders ) +=
                directGravityFieldVariations.second;
    }

    // Create time-varying gravity field.
    std::shared_ptr< TimeDependentSphericalHarmonicsGravityField > timeDependentGravityField =
            std::make_shared< TimeDependentSphericalHarmonicsGravityField >(
                gravitationalParameter, referenceRadius, nominalCosineCoefficients,
                nominalSineCoefficients, getTestGravityFieldVariations( ) );
    timeDependentGravityField->update( 2.0 * testTime );

    // Calculate variations for current test time.
    timeDependentGravityField->update( testTime );

    // Get corrections at current time step from timeDependentGravityField
    Eigen::MatrixXd perturbedCosineCoefficients =
            timeDependentGravityField->getCosineCoefficients( );
    Eigen::MatrixXd perturbedSineCoefficients =
            timeDependentGravityField->getSineCoefficients( );
    Eigen::MatrixXd calculatedCosineCoefficientCorrections =
            ( perturbedCosineCoefficients - nominalCosineCoefficients ).block( 0, 0, 5, 5 );
    Eigen::MatrixXd calculatedSineCoefficientCorrections =
            ( perturbedSineCoefficients - nominalSineCoefficients ).block( 0, 0, 5, 5 );

    // Compare corrections against expected variations.
    for( unsigned int i = 0; i < 5; i++ )
    {
        for( unsigned int j = 0; j < 5; j++ )
        {
            BOOST_CHECK_SMALL( calculatedCosineCoefficientCorrections( i, j ) -
                               expectedCosineCoefficientsCorrections( i, j ), 1.0E-18 );
            BOOST_CHECK_SMALL( calculatedSineCoefficientCorrections( i, j ) -
                               expectedSineCoefficientsCorrections( i, j ), 1.0E-18 );
        }
    }

    // Test calculated tidal corrections against manual corrections directly from Cartesian states
    // of perturbing bodies.
    {
        // Define Love numbers
        // Define Love numbers ( constant for degree 2 only)
        std::complex< double > constantLoveNumber( 0.5, 0.5E-3 );
        std::vector< std::complex< double > > constantSingleDegreeLoveNumber( 3, constantLoveNumber );
        std::map< int, std::vector< std::complex< double > > > loveNumbers;
        loveNumbers[ 2 ] = constantSingleDegreeLoveNumber;

        // Manually calculate tidal corrections directly.
        std::pair< Eigen::MatrixXd, Eigen::MatrixXd > directIoTide =
                calculateSolidBodyTideSingleCoefficientSetCorrectionFromAmplitude(
                    loveNumbers, getBodyGravitationalParameter( "Io" )
                    / getBodyGravitationalParameter( "Jupiter" ),
                    getAverageRadius( "Jupiter" ),  spice_interface::getBodyCartesianPositionAtEpoch(
                        "Io", "Jupiter", "IAU_Jupiter", "None", testTime ), 2, 2 );
        std::pair< Eigen::MatrixXd, Eigen::MatrixXd > directEuropaTide =
                calculateSolidBodyTideSingleCoefficientSetCorrectionFromAmplitude(
                    loveNumbers, getBodyGravitationalParameter( "Europa" )
                    / getBodyGravitationalParameter( "Jupiter" ),
                    getAverageRadius( "Jupiter" ),  spice_interface::getBodyCartesianPositionAtEpoch(
                        "Europa", "Jupiter", "IAU_Jupiter", "None", testTime ), 2, 2 );

        // Calculate tidal corrections from objects
        std::pair< Eigen::MatrixXd, Eigen::MatrixXd > tidalCorrectionsFromObject =
                ( ( timeDependentGravityField->getGravityFieldVariationsSet( )
                    ->getGravityFieldVariation( basic_solid_body ) ).second
                  ->calculateSphericalHarmonicsCorrections( testTime ) );

        // Compare results.
        Eigen::MatrixXd directCosineCorrections = directIoTide.first + directEuropaTide.first;
        Eigen::MatrixXd directSineCorrections = directIoTide.second + directEuropaTide.second;

        for( unsigned int i = 0; i < 3; i++ )
        {
            BOOST_CHECK_SMALL( directCosineCorrections( 2, i )
                               - tidalCorrectionsFromObject.first( 0, i ), 1.0E-19 );
            BOOST_CHECK_SMALL( directSineCorrections( 2, i )
                               - tidalCorrectionsFromObject.second( 0, i ), 1.0E-19 );
        }
    }
}

void getPeriodicGravityFieldVariationSettings(
        std::vector< Eigen::MatrixXd >& cosineAmplitudes,
        std::vector< Eigen::MatrixXd >& sineAmplitudes,
        std::vector< double >& frequencies,
        std::vector< double >& phases,
        double& referenceEpoch,
        int& minimumDegree,
        int& minimumOrder,
        const int type = 0 )
{
    cosineAmplitudes.clear();
    sineAmplitudes.clear();
    frequencies.clear();
    phases.clear();

    cosineAmplitudes.push_back(
                ( Eigen::MatrixXd( 2, 3 )<<2.0E-8, -2.0E-8, 4.2E-9,
                  4.0E-8, -3.7E-9, 2.9E-7 ).finished( ) );
    cosineAmplitudes.push_back(
                ( Eigen::MatrixXd( 2, 3 )<<5.0E-8, -6.0E-8, 2.2E-9,
                  4.0E-7, -3.4E-7, 8.9E-7 ).finished( ) );
    cosineAmplitudes.push_back(
                ( Eigen::MatrixXd( 2, 3 )<<8.0E-8, 2.0E-7, 7.2E-8,
                  7.5E-7, -1.9E-7, 2.5E-7 ).finished( ) );

    sineAmplitudes.push_back(
                ( Eigen::MatrixXd( 2, 3 )<<0.0, 6.0E-8, 5.2E-9,
                  0.0, 7.4E-8, 5.6E-7 ).finished( ) );
    sineAmplitudes.push_back(
                ( Eigen::MatrixXd( 2, 3 )<<0.0, 5.5E-8, 9.2E-9,
                  0.0, 4.4E-8, 1.1E-7 ).finished( ) );
    sineAmplitudes.push_back(
                ( Eigen::MatrixXd( 2, 3 )<<0.0, 7.8E-7, 4.2E-8,
                  0.0, 8.5E-7, -2.5E-8 ).finished( ) );

    frequencies.resize( 3 );
    phases.resize( 3 );

    frequencies = { 2.0E-5, 7.0E-5, 4.0E-6 };
    phases = { 1.0, 2.4, 5.0 };
    referenceEpoch = 5.0E8;

    if( type == 0 )
    {
        minimumDegree = 2;
        minimumOrder = 0;
    }
    else if( type == 1 )
    {
        minimumDegree = 3;
        minimumOrder = 0;
    }
    else if( type == 2 )
    {
        minimumDegree = 3;
        minimumOrder = 1;
    }
}

BOOST_AUTO_TEST_CASE( testPeriodicGravityFieldVariations )
{
    using namespace tudat::simulation_setup;

    // Load spice kernels.
    spice_interface::loadStandardSpiceKernels( );

    // Define properties of nominal field
    double gravitationalParameter = getBodyGravitationalParameter( "Jupiter" );
    double referenceRadius = getAverageRadius( "Jupiter" );
    Eigen::MatrixXd nominalCosineCoefficients;
    Eigen::MatrixXd nominalSineCoefficients;
    getNominalJupiterGravityField( nominalCosineCoefficients, nominalSineCoefficients );
    {
        double testTime = 1.0E7;
        for( int k = 0; k < 4; k++ )
        {
            std::vector< Eigen::MatrixXd > cosineAmplitudes;
            std::vector< Eigen::MatrixXd > sineAmplitudes;
            std::vector< double > frequencies;
            std::vector< double > phases;
            double referenceEpoch;
            int minimumDegree;
            int minimumOrder;
            int testIndex = ( k < 3 ) ? k : 0;
            getPeriodicGravityFieldVariationSettings(
                        cosineAmplitudes, sineAmplitudes, frequencies, phases,
                        referenceEpoch, minimumDegree, minimumOrder, testIndex );

            std::shared_ptr< PeriodicGravityFieldVariationsSettings > variationSettings;
            if( k < 3 )
            {
                variationSettings = std::dynamic_pointer_cast< PeriodicGravityFieldVariationsSettings >(
                            periodicGravityFieldVariationsSettings(
                                cosineAmplitudes, sineAmplitudes,
                                frequencies, phases, referenceEpoch,
                                minimumDegree, minimumOrder ) );
            }
            else
            {
                variationSettings = std::dynamic_pointer_cast< PeriodicGravityFieldVariationsSettings >(
                            periodicGravityFieldVariationsSettings(
                                cosineAmplitudes.at( 0 ), sineAmplitudes.at( 0 ),
                                frequencies.at( 0 ), 0.0, referenceEpoch,
                                minimumDegree, minimumOrder ) );
            }

            // Create time-varying gravity field.
            std::shared_ptr< TimeDependentSphericalHarmonicsGravityField > timeDependentGravityField =
                    std::make_shared< TimeDependentSphericalHarmonicsGravityField >(
                        gravitationalParameter, referenceRadius, nominalCosineCoefficients,
                        nominalSineCoefficients );

            SystemOfBodies dummySystem;
            dummySystem.createEmptyBody( "Jupiter" );
            dummySystem.getBody( "Jupiter" )->setGravityFieldModel( timeDependentGravityField );
            std::vector< std::shared_ptr< GravityFieldVariationSettings > > variationSettingsList;
            variationSettingsList.push_back( variationSettings );

            std::shared_ptr< gravitation::GravityFieldVariationsSet > variations = createGravityFieldModelVariationsSet(
                        "Jupiter", dummySystem, variationSettingsList );
            timeDependentGravityField->updateCorrectionFunctions( );

            if( k < 3 )
            {
                timeDependentGravityField->update( 2.0 * testTime );

                // Calculate variations for current test time.
                timeDependentGravityField->update( testTime );

                // Get corrections at current time step from timeDependentGravityField
                Eigen::MatrixXd perturbedCosineCoefficients =
                        timeDependentGravityField->getCosineCoefficients( );
                Eigen::MatrixXd perturbedSineCoefficients =
                        timeDependentGravityField->getSineCoefficients( );
                Eigen::MatrixXd calculatedCosineCoefficientCorrections =
                        ( perturbedCosineCoefficients - nominalCosineCoefficients ).block( 0, 0, 5, 5 );
                Eigen::MatrixXd calculatedSineCoefficientCorrections =
                        ( perturbedSineCoefficients - nominalSineCoefficients ).block( 0, 0, 5, 5 );

                Eigen::MatrixXd expectedCosineCoefficientsCorrections = Eigen::MatrixXd::Zero( 5, 5 );
                Eigen::MatrixXd expectedSineCoefficientsCorrections = Eigen::MatrixXd::Zero( 5, 5 );

                for( unsigned int j = 0; j < cosineAmplitudes.size( ); j++ )
                {
                    expectedCosineCoefficientsCorrections.block( minimumDegree, minimumOrder, 2, 3 ) += cosineAmplitudes.at( j ) * std::cos(
                                frequencies.at( j ) * ( testTime - referenceEpoch ) + phases.at( j ) );
                    expectedSineCoefficientsCorrections.block( minimumDegree, minimumOrder, 2, 3 ) += sineAmplitudes.at( j ) * std::sin(
                                frequencies.at( j ) * ( testTime - referenceEpoch ) + phases.at( j ) );
                }

                // Compare corrections against expected variations.
                for( unsigned int i = 0; i < 5; i++ )
                {
                    for( unsigned int j = 0; j < 5; j++ )
                    {
                        BOOST_CHECK_SMALL( calculatedCosineCoefficientCorrections( i, j ) -
                                           expectedCosineCoefficientsCorrections( i, j ), 1.0E-18 );
                        BOOST_CHECK_SMALL( calculatedSineCoefficientCorrections( i, j ) -
                                           expectedSineCoefficientsCorrections( i, j ), 1.0E-18 );
                    }
                }
            }
            else
            {
                timeDependentGravityField->update( referenceEpoch );

                // Get corrections at current time step from timeDependentGravityField
                Eigen::MatrixXd perturbedCosineCoefficients =
                        timeDependentGravityField->getCosineCoefficients( );
                Eigen::MatrixXd perturbedSineCoefficients =
                        timeDependentGravityField->getSineCoefficients( );
                Eigen::MatrixXd calculatedCosineCoefficientCorrections =
                        ( perturbedCosineCoefficients - nominalCosineCoefficients ).block( 0, 0, 5, 5 );
                Eigen::MatrixXd calculatedSineCoefficientCorrections =
                        ( perturbedSineCoefficients - nominalSineCoefficients ).block( 0, 0, 5, 5 );

                Eigen::MatrixXd expectedCosineCoefficientsCorrections = Eigen::MatrixXd::Zero( 5, 5 );

                expectedCosineCoefficientsCorrections.block( minimumDegree, minimumOrder, 2, 3 ) =
                        cosineAmplitudes.at( 0 );

                // Compare corrections against expected variations.
                for( unsigned int i = 0; i < 5; i++ )
                {
                    for( unsigned int j = 0; j < 5; j++ )
                    {
                        BOOST_CHECK_SMALL( calculatedCosineCoefficientCorrections( i, j ) -
                                           expectedCosineCoefficientsCorrections( i, j ), 1.0E-18 );
                        BOOST_CHECK_EQUAL( calculatedSineCoefficientCorrections( i, j ), 0.0 );
                    }
                }

            }
        }
    }
}

BOOST_AUTO_TEST_SUITE_END( )

} // namespace unit_tests
} // namespace tudat


#include <cmath>

#include <string>
#include <iostream>
#include <random>
#include <algorithm>

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

#define EIGEN_INITIALIZE_MATRICES_BY_NAN

#include <cmath>
#include <cstdlib>
#include <complex>
#include <type_traits>

#include <iostream>
#include <iomanip>

// define the exception class, but don't override eigen's eigen_assert() macro itself
#include <tomographer/tools/eigen_assert_exception.h>

#include <Eigen/Core>
#include <unsupported/Eigen/KroneckerProduct>

#define BOOST_TEST_MODULE test_channelspace
#include <boost/test/unit_test.hpp>

#include <tomographer/tools/boost_test_logger.h>
#include <tomographer/mathtools/pos_semidef_util.h>
#include <tomographer/densedm/dmtypes.h>

#include "channelspace.h"

// -----------------------------------------------------------------------------
// fixture(s)

struct MyStatsCollector {
  inline void init() { }
  inline void thermalizingDone() { }
  inline void done() { }
  template<typename PointType, typename MHRandomWalk>
  inline void processSample(int k, int n, const PointType & Vpt, double fnval, MHRandomWalk & rw)
  {
    BOOST_TEST_MESSAGE("New sample: V = \n" << Vpt) ;
  }
  template<typename... Args>
  inline void rawMove(Args&&...) const { }
};

template<int JumpMode_, int LogLevel = Tomographer::Logger::LONGDEBUG>
struct qubit_to_qutrit_fixture {

  typedef ChannelTypes<> ChannelTypes;
  typedef ChannelTypes::MatrixType MatrixType;
  typedef Tomographer::DenseDM::IndepMeasLLH<ChannelTypes> DenseLLH;

  ChannelTypes cht;
  DenseLLH llh;

  std::mt19937 rng;
  Tomographer::Logger::BoostTestLogger logger;

  // the channel MHWalker
  typedef ChannelSpaceMHWalker<DenseLLH, std::mt19937, Tomographer::Logger::BoostTestLogger>
    ChannelSpaceMHWalker;

  ChannelSpaceMHWalker mhwalker;

  static inline DenseLLH _make_llh(ChannelTypes cht)
  {
    DenseLLH llh(cht);

    MatrixType Zp(2,2);
    MatrixType Zm(2,2);
    MatrixType Xp(2,2);
    MatrixType Xm(2,2);

    MatrixType out(3,3);

    Zp << 1, 0,  0, 0;
    Zm << 0, 0,  0, 1;
    Xp << 1, 1,  1, 1;
    Xm << 1,-1, -1, 1;

    out << 1, 0, 0,  0, 0, 0,  0, 0, 0; // pure |0> state on qutrit

    llh.addMeasEffect(Eigen::kroneckerProduct(Zp, out), 50);
    llh.addMeasEffect(Eigen::kroneckerProduct(Zm, out), 20);
    llh.addMeasEffect(Eigen::kroneckerProduct(Xp, out), 10);
    llh.addMeasEffect(Eigen::kroneckerProduct(Xm, out), 90);

    return llh;
  }

  qubit_to_qutrit_fixture()
    : cht(2,3),
      llh(_make_llh(cht)),
      rng(0),
      logger(LogLevel),
      mhwalker(llh, rng, logger, (JumpMode)JumpMode_)
  {
  }

};

// because BOOST macros don't like commas in their fixture argument...
template<int JumpType_>
struct qubit_to_qutrit_fixture_DEBUG : public qubit_to_qutrit_fixture<JumpType_, Tomographer::Logger::DEBUG> { };

// -----------------------------------------------------------------------------
// test suites

BOOST_AUTO_TEST_SUITE(test_channelspacemhwalker);
// =============================================================================

BOOST_FIXTURE_TEST_CASE(simple_qubit_to_qutrit_eiH, qubit_to_qutrit_fixture<RandHermExp>)
{
  MyStatsCollector stats;
  typedef Tomographer::MHRandomWalk<std::mt19937, ChannelSpaceMHWalker, MyStatsCollector,
                                    Tomographer::Logger::BoostTestLogger> MHRandomWalk;
  MHRandomWalk rw(0.1, 10, 500, 1024, mhwalker, stats, rng, logger);
  rw.run();

  BOOST_CHECK( true ) ; // dummy test case -- for now, just check that the above compiles & runs
}
BOOST_FIXTURE_TEST_CASE(simple_qubit_to_qutrit_elr, qubit_to_qutrit_fixture<ElemRotations>)
{
  MyStatsCollector stats;
  typedef Tomographer::MHRandomWalk<std::mt19937, ChannelSpaceMHWalker, MyStatsCollector,
                                    Tomographer::Logger::BoostTestLogger> MHRandomWalk;
  MHRandomWalk rw(0.1, 10, 500, 1024, mhwalker, stats, rng, logger);
  rw.run();

  BOOST_CHECK( true ) ; // dummy test case -- for now, just check that the above compiles & runs
}

BOOST_FIXTURE_TEST_CASE(jump_fn_symmetric_eiH, qubit_to_qutrit_fixture_DEBUG<RandHermExp>)
{
  // start at a fixed starting point
  typedef ChannelSpaceMHWalker::PointType PointType;
  PointType V0 = PointType::Zero(llh.dmt.dimXY2(), llh.dmt.dimX());
  V0(0,0) = V0(1,1) = 1;

  // average over many steps starting from the same V0 and normalize --> should be V0
  // again

  const int N = 5000000;
  const double step_size = 0.1;

  PointType totalV = PointType::Zero(llh.dmt.dimXY2(), llh.dmt.dimX());

  for (int k = 0; k < N; ++k) {
    PointType newV = mhwalker.jumpFn(V0, step_size) ;
    totalV += newV;
  }
  BOOST_TEST_MESSAGE("total summed newV's = \n" << totalV) ;

  totalV /= (totalV.norm() / std::sqrt((double)llh.dmt.dimX()));

  BOOST_TEST_MESSAGE("normalized = \n" << totalV) ;

  double err = (totalV - V0).cwiseAbs().maxCoeff();
  double errok = 4*step_size/std::sqrt(N); // approx. 4 sigma
  BOOST_TEST_MESSAGE("err = " << err << ", errok = " << errok) ;
  BOOST_CHECK_LT(err, errok) ;
}
BOOST_FIXTURE_TEST_CASE(jump_fn_symmetric_elr, qubit_to_qutrit_fixture_DEBUG<ElemRotations>)
{
  // start at a fixed starting point
  typedef ChannelSpaceMHWalker::PointType PointType;
  PointType V0 = PointType::Zero(llh.dmt.dimXY2(), llh.dmt.dimX());
  V0(0,0) = V0(1,1) = 1;

  // average over many steps starting from the same V0 and normalize --> should be V0
  // again

  const int N = 5000000;
  const double step_size = 0.1;

  PointType totalV = PointType::Zero(llh.dmt.dimXY2(), llh.dmt.dimX());

  for (int k = 0; k < N; ++k) {
    PointType newV = mhwalker.jumpFn(V0, step_size) ;
    totalV += newV;
  }
  BOOST_TEST_MESSAGE("total summed newV's = \n" << totalV) ;

  totalV /= (totalV.norm() / std::sqrt((double)llh.dmt.dimX()));

  BOOST_TEST_MESSAGE("normalized = \n" << totalV) ;

  double err = (totalV - V0).cwiseAbs().maxCoeff();
  double errok = 4*step_size/std::sqrt(N); // approx. 4 sigma
  BOOST_TEST_MESSAGE("err = " << err << ", errok = " << errok) ;
  BOOST_CHECK_LT(err, errok) ;
}

// =============================================================================
BOOST_AUTO_TEST_SUITE_END() ;


/*
 * graph_astar.hpp
 *
 * Graph Library for Autonomous and Dynamic Systems
 *
 * author:  Pierrick Koch <pierrick.koch@laas.fr>
 * created: 2013-07-22
 * license: BSD
 */
#ifndef GRAPH_ASTAR_HPP
#define GRAPH_ASTAR_HPP

#include <cmath>
#include <vector>
#include <algorithm>

#include <boost/graph/adjacency_list.hpp>
#include <boost/graph/astar_search.hpp>

#include "gladys/point.hpp"

namespace gladys {

struct vertex {
    point_xy_t pt;
};

struct edge {
    float weight;
    time_t t; /**< when the weight has been modified */
};
// graph
typedef boost::adjacency_list<boost::vecS, boost::vecS, boost::undirectedS,
                              vertex, edge > graph_t;
typedef graph_t::vertex_descriptor vertex_t;
typedef graph_t::edge_descriptor edge_t;
typedef std::vector<vertex_t> vertices_t;
typedef std::map<point_xy_t, vertex_t> vertex_map_t;

typedef struct {
    path_t path;
    double cost; // TMP
    double utility; // TMP
} path_cost_util_t;

typedef struct {
    path_t path;
    std::deque<double> costs ; // cumulative costs for if waypoints
} detailed_path_t;

struct found_goal {
    vertex_t g;
    found_goal(vertex_t _g) : g(_g) {};
}; // exception for termination

// visitor that terminates when we find the goal
class astar_goal_visitor : public boost::default_astar_visitor {
    vertex_t goal;
public:
    astar_goal_visitor(vertex_t _goal) : goal(_goal) {}

    /** examine_vertex is invoked when a vertex is popped from the queue
     * (i.e., it has the lowest cost on the OPEN list).
     */
    void examine_vertex(vertex_t u, const graph_t& g) {
        // search if the current vertex is in the list of goal
        if (u == goal)
            throw found_goal(u);
    }
};

class astar_goals_visitor : public boost::default_astar_visitor {
    vertices_t goals;
public:
    astar_goals_visitor(vertices_t _goals) : goals(_goals) {}

    /** examine_vertex is invoked when a vertex is popped from the queue
     * (i.e., it has the lowest cost on the OPEN list).
     */
    void examine_vertex(vertex_t u, const graph_t& g) {
        // search if the current vertex is in the list of goal
        for (auto& v : goals)
            if (u == v)
                throw found_goal(u);
    }
};

/**
* Navigation heuristics functions for graph visitor algorithms
*/
class nav_goal_heuristic : public boost::astar_heuristic<graph_t, double> {
    const graph_t& g;
    const vertex_t& goal;
public:
    nav_goal_heuristic(const graph_t& _g, const vertex_t& _goal)
        : g(_g), goal(_goal) {}

    double operator()(const vertex_t& u) {
        return distance(g[u].pt, g[goal].pt);
    }
};

/**
* Navigation heuristics functions for graph visitor algorithms
*/
class nav_goals_heuristic : public boost::astar_heuristic<graph_t, double> {
    const graph_t& g;
    const vertices_t& goals;
public:
    nav_goals_heuristic(const graph_t& _g, const vertices_t& _goals)
        : g(_g), goals(_goals) {}

    double operator()(const vertex_t& u) {
        // get closest distance to one of the goals
        auto first = goals.begin();
        double tmp, smallest = distance_sq(g[u].pt, g[*first].pt);
        ++first;
        for (; first != goals.end(); ++first) {
            tmp = distance_sq(g[u].pt, g[*first].pt);
            if (tmp < smallest) {
                smallest = tmp;
            }
        }
        return std::sqrt( smallest );
    }
};

} // namespace gladys

#endif // GRAPH_ASTAR_HPP