#include "formulation/scalar/diffusion.hpp"

#include <array>
#include <memory>
#include <cstdlib>

#include <deal.II/dofs/dof_handler.h>
#include <deal.II/fe/fe_q.h>
#include <deal.II/grid/grid_generator.h>
#include <deal.II/grid/tria.h>

#include "data/cross_sections/material_cross_sections.hpp"
#include "domain/finite_element/tests/finite_element_mock.hpp"
#include "data/material/tests/material_mock.hpp"
#include "test_helpers/gmock_wrapper.h"
#include "test_helpers/test_helper_functions.h"
#include "test_helpers/test_assertions.hpp"

namespace {

using ::testing::Le;
using ::testing::Ge;
using ::testing::DoDefault;
using ::testing::NiceMock;
using ::testing::Return;
using ::testing::_;
using ::testing::Ref;
using ::testing::AtLeast;
using namespace bart;

namespace test_helpers = bart::test_helpers;
using test_helpers::AreEqual;

class FormulationCFEMDiffusionTest : public ::testing::Test {
 protected:
  FormulationCFEMDiffusionTest()
      : dof_handler_(triangulation_),
        fe_(1) {};
  using Matrix = dealii::FullMatrix<double>;
  std::shared_ptr<domain::finite_element::FiniteElementMock<2>> fe_mock_ptr;
  std::shared_ptr<data::cross_sections::MaterialCrossSections> cross_sections_ptr;

  dealii::DoFHandler<2>::active_cell_iterator cell_ptr_;
  dealii::Triangulation<2> triangulation_;
  dealii::DoFHandler<2> dof_handler_;
  dealii::FE_Q<2> fe_;

  const int fissile_material_id_ = 0, non_fissile_material_id_ = 1;

  void SetUp() override;
  void SetUpDealii();
};

void FormulationCFEMDiffusionTest::SetUp() {
  SetUpDealii();
  // Make mock objects. Cross-sections is a struct that cannot be mocked, but
  // we can mock the material object it is based on.
  NiceMock<data::material::MaterialMock> mock_material;
  fe_mock_ptr = std::make_shared<NiceMock<domain::finite_element::FiniteElementMock<2>>>();

  ON_CALL(*fe_mock_ptr, dofs_per_cell()).WillByDefault(Return(2));
  ON_CALL(*fe_mock_ptr, n_cell_quad_pts()).WillByDefault(Return(2));
  ON_CALL(*fe_mock_ptr, n_face_quad_pts()).WillByDefault(Return(2));

  for (int q = 0; q < 2; ++q) {
    ON_CALL(*fe_mock_ptr, Jacobian(q)).WillByDefault(Return((1 + q)*3));
    ON_CALL(*fe_mock_ptr, FaceJacobian(q)).WillByDefault(Return((1 + q)*3));
    for (int i = 0; i < 2; ++i) {
      ON_CALL(*fe_mock_ptr, ShapeValue(i, q)).WillByDefault(Return(i + q));
      ON_CALL(*fe_mock_ptr, FaceShapeValue(i,q)).WillByDefault(Return(i + q));

      dealii::Tensor<1, 2> gradient_tensor;
      gradient_tensor[0] = i;
      gradient_tensor[1] = q;

      ON_CALL(*fe_mock_ptr, ShapeGradient(i, q)).WillByDefault(Return(gradient_tensor));
    }
  }

  // Cross-section data for fake two-group
  std::array<double, 4> sigma_s_values{0.25, 0.5, 0.75, 1.0};
  dealii::FullMatrix<double> sigma_s_matrix{2,2, sigma_s_values.begin()};
  std::unordered_map<int, dealii::FullMatrix<double>> sigma_s{{this->fissile_material_id_, sigma_s_matrix}};

  std::unordered_map<int, std::vector<double>> sigma_t{{this->fissile_material_id_, {1.0, 2.0}}}; // i + 1

  std::unordered_map<int, std::vector<double>> q{{this->non_fissile_material_id_, {1.0, 2.0}}};

  std::unordered_map<int, bool> fissile_id{{this->fissile_material_id_, true},
                                           {this->non_fissile_material_id_, false}};

  ON_CALL(mock_material, GetSigT()).WillByDefault(Return(sigma_t));
  ON_CALL(mock_material, GetSigS()).WillByDefault(Return(sigma_s));
  ON_CALL(mock_material, GetDiffusionCoef()).WillByDefault(Return(sigma_t));
  ON_CALL(mock_material, GetQ()).WillByDefault(Return(q));
  ON_CALL(mock_material, GetFissileIDMap()).WillByDefault(Return(fissile_id));
  ON_CALL(mock_material, GetChiNuSigF()).WillByDefault(Return(sigma_s));

  cross_sections_ptr = std::make_shared<data::cross_sections::MaterialCrossSections>(mock_material);
}

// Set up a simple deal.ii problem so that the cell points to something, this
// is for material ID
void FormulationCFEMDiffusionTest::SetUpDealii() {
  dealii::GridGenerator::hyper_cube(triangulation_, 0, 1);
  dof_handler_.distribute_dofs(fe_);
  for (auto cell = dof_handler_.begin_active(); cell != dof_handler_.end(); ++cell) {
    if (cell->is_locally_owned()) {
      cell_ptr_ = cell;
      cell_ptr_->set_material_id(this->fissile_material_id_);
    }
  }
}

TEST_F(FormulationCFEMDiffusionTest, ConstructorTest) {
  EXPECT_CALL(*fe_mock_ptr, dofs_per_cell()).WillOnce(DoDefault());
  EXPECT_CALL(*fe_mock_ptr, n_cell_quad_pts()).WillOnce(DoDefault());
  EXPECT_CALL(*fe_mock_ptr, n_face_quad_pts()).WillOnce(DoDefault());

  formulation::scalar::Diffusion<2> test_diffusion(fe_mock_ptr,cross_sections_ptr);

  EXPECT_EQ(fe_mock_ptr.use_count(), 2);
  EXPECT_EQ(cross_sections_ptr.use_count(), 2);
  EXPECT_FALSE(test_diffusion.is_initialized());
}

TEST_F(FormulationCFEMDiffusionTest, PrecalculateTest) {
  // Shape call, by default, returns (quadrature point index + degree of freedom)
  // Gradient call will return a tensor of [i, q]
  // Expected results, we need to make arrays to put them into dealii matrices
  std::array<double, 4> shape_matrix_q_0_values = {0, 0,
                                                   0, 1};
  std::array<double, 4> shape_matrix_q_1_values = {1, 2,
                                                   2, 4};
  std::array<double, 4> gradient_matrix_q_0_values = {0, 0,
                                                      0, 1};
  std::array<double, 4> gradient_matrix_q_1_values = {1, 1,
                                                      1, 2};

  dealii::FullMatrix<double> shape_matrix_q_0{2,2,shape_matrix_q_0_values.begin()};
  dealii::FullMatrix<double> shape_matrix_q_1{2,2,shape_matrix_q_1_values.begin()};
  dealii::FullMatrix<double> gradient_matrix_q_0{2,2,gradient_matrix_q_0_values.begin()};
  dealii::FullMatrix<double> gradient_matrix_q_1{2,2,gradient_matrix_q_1_values.begin()};

  formulation::scalar::Diffusion<2> test_diffusion(fe_mock_ptr,cross_sections_ptr);

  // Set call expectations
  EXPECT_CALL(*fe_mock_ptr, SetCell(_)).Times(1);
  // These are called at least 4 times which is the minimum required for 2 dofs and 2 quadrature points
  EXPECT_CALL(*fe_mock_ptr, ShapeValue(_,_)).Times(::testing::AtLeast(4)).WillRepeatedly(DoDefault());
  EXPECT_CALL(*fe_mock_ptr, ShapeGradient(_,_)).Times(::testing::AtLeast(4)).WillRepeatedly(DoDefault());

  test_diffusion.Precalculate(cell_ptr_);
  auto shape_squared = test_diffusion.GetShapeSquared();
  auto gradient_squared = test_diffusion.GetGradientSquared();

  EXPECT_TRUE(AreEqual(shape_matrix_q_0, shape_squared.at(0)));
  EXPECT_TRUE(AreEqual(shape_matrix_q_1, shape_squared.at(1)));
  EXPECT_TRUE(AreEqual(gradient_matrix_q_0, gradient_squared.at(0)));
  EXPECT_TRUE(AreEqual(gradient_matrix_q_1, gradient_squared.at(1)));
  EXPECT_TRUE(test_diffusion.is_initialized());
}

TEST_F(FormulationCFEMDiffusionTest, FillCellConstantTermTest) {
  dealii::Vector<double> test_vector(2);
  dealii::Vector<double> expected_values{ 54, 129 };
  dealii::Vector<double> constant_vector_at_dofs(2);
  std::vector<double> constant_vector_at_quadrature{ 7, 9 };

  formulation::scalar::Diffusion<2> test_diffusion(fe_mock_ptr, cross_sections_ptr);
  EXPECT_CALL(*fe_mock_ptr, SetCell(cell_ptr_)).Times(1);
  EXPECT_CALL(*fe_mock_ptr, ValueAtQuadrature(Ref(constant_vector_at_dofs)))
      .WillOnce(Return(constant_vector_at_quadrature));

  for (int q = 0; q < 2; ++q) {
    EXPECT_CALL(*fe_mock_ptr, Jacobian(q)).Times(AtLeast(1)).WillRepeatedly(DoDefault());
    for (int i = 0; i < 2; ++i) {
      EXPECT_CALL(*fe_mock_ptr, ShapeValue(i,q)).Times(AtLeast(1)).WillRepeatedly(DoDefault());
    }
  }

  test_diffusion.FillCellConstantTerm(test_vector, cell_ptr_, constant_vector_at_dofs);
  EXPECT_TRUE(AreEqual(expected_values, test_vector));
}

TEST_F(FormulationCFEMDiffusionTest, FillCellStreamingTermTest) {
  dealii::FullMatrix<double> test_matrix(2,2);

  std::array<double, 4> expected_values{6, 6,
                                        6, 15};
  dealii::FullMatrix<double> expected_matrix(2, 2, expected_values.begin());

  formulation::scalar::Diffusion<2> test_diffusion(fe_mock_ptr, cross_sections_ptr);

  EXPECT_CALL(*fe_mock_ptr, Jacobian(_)).Times(2).WillRepeatedly(DoDefault());
  EXPECT_CALL(*fe_mock_ptr, SetCell(cell_ptr_)).Times(2);

  EXPECT_ANY_THROW({
    test_diffusion.FillCellStreamingTerm(test_matrix, cell_ptr_, 0);
  });
  EXPECT_NO_THROW({
    test_diffusion.Precalculate(cell_ptr_);
    test_diffusion.FillCellStreamingTerm(test_matrix, cell_ptr_, 0);
  });

  EXPECT_TRUE(AreEqual(expected_matrix, test_matrix));
}

TEST_F(FormulationCFEMDiffusionTest, FillCellCollisionTermTest) {
  dealii::FullMatrix<double> test_matrix(2,2);

  std::array<double, 4> expected_values{4.5, 9.0,
                                        9.0, 20.25};
  dealii::FullMatrix<double> expected_matrix(2, 2, expected_values.begin());

  formulation::scalar::Diffusion<2> test_diffusion(fe_mock_ptr, cross_sections_ptr);

  EXPECT_CALL(*fe_mock_ptr, Jacobian(_)).Times(2).WillRepeatedly(DoDefault());
  EXPECT_CALL(*fe_mock_ptr, SetCell(cell_ptr_)).Times(2);

  EXPECT_ANY_THROW({
    test_diffusion.FillCellCollisionTerm(test_matrix, cell_ptr_, 0);
                   });
  EXPECT_NO_THROW({
    test_diffusion.Precalculate(cell_ptr_);
    test_diffusion.FillCellCollisionTerm(test_matrix, cell_ptr_, 0);
                  });

  EXPECT_TRUE(AreEqual(expected_matrix, test_matrix));
}

TEST_F(FormulationCFEMDiffusionTest, FillBoundaryTermTestReflective) {
  dealii::FullMatrix<double> test_matrix(2,2);
  dealii::FullMatrix<double> expected_matrix(2,2);

  auto boundary = formulation::scalar::Diffusion<2>::BoundaryType::kReflective;

  formulation::scalar::Diffusion<2> test_diffusion(fe_mock_ptr, cross_sections_ptr);

  EXPECT_ANY_THROW({
    test_diffusion.FillBoundaryTerm(test_matrix, cell_ptr_, 0, boundary);
                   });
  EXPECT_NO_THROW({
    test_diffusion.Precalculate(cell_ptr_);
    test_diffusion.FillBoundaryTerm(test_matrix, cell_ptr_, 0, boundary);
                  });

  EXPECT_TRUE(AreEqual(expected_matrix, test_matrix));
}

TEST_F(FormulationCFEMDiffusionTest, FillBoundaryTermTestVacuum) {
  dealii::FullMatrix<double> test_matrix(2,2);

  auto boundary = formulation::scalar::Diffusion<2>::BoundaryType::kVacuum;

  std::array<double, 4> expected_values{3, 6,
                                        6, 13.5};

  dealii::FullMatrix<double> expected_matrix(2, 2, expected_values.begin());

  formulation::scalar::Diffusion<2> test_diffusion(fe_mock_ptr, cross_sections_ptr);

  EXPECT_CALL(*fe_mock_ptr, SetFace(cell_ptr_, domain::FaceIndex(0))).Times(1);
  EXPECT_CALL(*fe_mock_ptr, FaceJacobian(_)).Times(2).WillRepeatedly(DoDefault());

  test_diffusion.Precalculate(cell_ptr_);
  test_diffusion.FillBoundaryTerm(test_matrix, cell_ptr_, 0, boundary);

  EXPECT_TRUE(AreEqual(expected_matrix, test_matrix));

}

TEST_F(FormulationCFEMDiffusionTest, FillCellFixedSourceNonFissile) {
  dealii::Vector<double> expected_vector{6, 15};
  dealii::Vector<double> test_vector(2);

  formulation::scalar::Diffusion<2> test_diffusion(fe_mock_ptr,cross_sections_ptr);

  cell_ptr_->set_material_id(this->non_fissile_material_id_);
  EXPECT_CALL(*fe_mock_ptr, SetCell(_)).Times(1);
  EXPECT_CALL(*fe_mock_ptr, Jacobian(_)).Times(2).WillRepeatedly(DoDefault());
  EXPECT_CALL(*fe_mock_ptr, ShapeValue(_,_)).Times(4).WillRepeatedly(DoDefault());

  EXPECT_NO_THROW(test_diffusion.FillCellFixedSource(test_vector, cell_ptr_, 0));

  EXPECT_TRUE(AreEqual(expected_vector, test_vector));
}

TEST_F(FormulationCFEMDiffusionTest, FillCellFixedSourceFissile) {
  dealii::Vector<double> expected_vector{0, 0};
  dealii::Vector<double> test_vector(2);

  formulation::scalar::Diffusion<2> test_diffusion(fe_mock_ptr, cross_sections_ptr);

  EXPECT_NO_THROW(test_diffusion.FillCellFixedSource(test_vector, cell_ptr_, 0));

  EXPECT_TRUE(AreEqual(expected_vector, test_vector));
}

TEST_F(FormulationCFEMDiffusionTest, FillFissionSourceTest) {

  formulation::scalar::Diffusion<2> test_diffusion(fe_mock_ptr, cross_sections_ptr);

  dealii::Vector<double> test_vector(2);
  int group = 0;
  double k_effective = 1.05;
  // Make in-group moment
  std::vector<double> in_group_moment_values{0.5, 0.5};
  system::moments::MomentVector in_group_moment(in_group_moment_values.begin(), in_group_moment_values.end());

  /* Make out-group moments (specifically in-group values in this are different
   * This is a somewhat tortured process due to the way dealii::Vectors are
   * defined (system::moments::MomentVector is an alias) */
  std::vector<double> group_1_moment_values{1.0, 1.0};
  system::moments::MomentVector group_0_moment{0.75, 0.75};
  system::moments::MomentVector group_1_moment{1.0, 1.0};
  system::moments::MomentsMap out_group_moments;
  out_group_moments[{0,0,0}] = group_0_moment;
  out_group_moments[{1,0,0}] = group_1_moment;

  // Expected answer
  dealii::Vector<double> expected_vector{5.0, 12.5};

  EXPECT_CALL(*fe_mock_ptr, SetCell(_)).Times(1);
  EXPECT_CALL(*fe_mock_ptr, ValueAtQuadrature(group_1_moment)).WillOnce(Return(group_1_moment_values));
  EXPECT_CALL(*fe_mock_ptr, ValueAtQuadrature(in_group_moment)).WillOnce(Return(in_group_moment_values));

  test_diffusion.FillCellFissionSource(test_vector, cell_ptr_, group, k_effective, in_group_moment, out_group_moments);

  EXPECT_TRUE(AreEqual(expected_vector, test_vector));

}

TEST_F(FormulationCFEMDiffusionTest, FillScatteringSourceTest) {

  formulation::scalar::Diffusion<2> test_diffusion(fe_mock_ptr, cross_sections_ptr);

  dealii::Vector<double> test_vector(2);
  // Make in-group moment
  std::vector<double> in_group_moment_values{0.5, 0.5};
  system::moments::MomentVector in_group_moment(in_group_moment_values.begin(), in_group_moment_values.end());

  /* Make out-group moments (specifically in-group values in this are different
   * This is a somewhat tortured process due to the way dealii::Vectors are
   * defined (system::moments::MomentVector is an alias) */
  std::vector<double> group_0_moment_values{0.75, 0.75};
  std::vector<double> group_1_moment_values{1.0, 1.0};
  system::moments::MomentVector group_0_moment(group_0_moment_values.begin(), group_0_moment_values.end());
  system::moments::MomentVector group_1_moment(group_1_moment_values.begin(), group_1_moment_values.end());
  system::moments::MomentsMap out_group_moments;
  out_group_moments[{0,0,0}] = group_0_moment;
  out_group_moments[{1,0,0}] = group_1_moment;

  // Expected solutions
  std::vector<double> expected_group_0_values{3.0, 7.5};
  dealii::Vector<double> expected_group_0_vector(expected_group_0_values.begin(), expected_group_0_values.end());
  std::vector<double> expected_group_1_values{3.375, 8.4375};
  dealii::Vector<double> expected_group_1_vector(expected_group_1_values.begin(), expected_group_1_values.end());
  std::array<dealii::Vector<double>, 2> expected_vectors{ expected_group_0_vector, expected_group_1_vector };

  EXPECT_CALL(*fe_mock_ptr, SetCell(_)).Times(2);
  EXPECT_CALL(*fe_mock_ptr, ValueAtQuadrature(group_1_moment)).WillOnce(Return(group_1_moment_values));
  EXPECT_CALL(*fe_mock_ptr, ValueAtQuadrature(group_0_moment)).WillOnce(Return(group_0_moment_values));

  for (int group = 0; group < 2; ++group) {
    test_diffusion.FillCellScatteringSource(test_vector, cell_ptr_, group, out_group_moments);

    EXPECT_TRUE(AreEqual(expected_vectors.at(group), test_vector));
    test_vector = 0;
  }
}

} // namespace


// File: train_adagrad.cc
// Author: Karl Moritz Hermann (mail@karlmoritz.com)
// Created: 01-01-2013
// Last Update: Thu 03 Oct 2013 11:48:28 AM BST

// STL
#include <iostream>
#include <cmath>
#include <algorithm>
#include <random>

// Boost
#include <boost/program_options/variables_map.hpp>
#include <boost/program_options/parsers.hpp>

// L-BFGS
#include <lbfgs.h>

// Local
#include "train_adagrad.h"
#include "train_update.h"
#include "recursive_autoencoder.h"
#include "utils.h"
#include "fast_math.h"

using namespace std;
namespace bpo = boost::program_options;

int train_adagrad(Model &model, int iterations, float eta, int batches, Real lambda)
{
  Real* vars = nullptr;
  int number_vars = 0;

  setVarsAndNumber(vars,number_vars,model);
  WeightArrayType theta(vars,number_vars);

  Real* Gt_d = new Real[number_vars]();
  // Real* Ginv_d = new Real[number_vars]();
  WeightArrayType Gt(Gt_d,number_vars);
  // WeightArrayType Ginv(Ginv_d,number_vars);
  // Gt.setOnes(); // initialize to ones..

  Real* gradient = new Real[number_vars]();
  lbfgsfloatval_t error = 0;

  // Remove L2 regularization as AdaGrad uses L1 instead
  model.calc_L2 = false;

  int size = model.corpus.size();
  int num_batches = min(batches,size/2);
  int batchsize = max((size/num_batches),2);
  // eta = eta / num_batches;
  cout << "Batch size: " << batchsize << "  eta " << eta << endl;

  Real update;
  Real l1_reg;

  for (auto iteration = 0; iteration < iterations; ++iteration)
  {
    cout << "Iteration " << iteration << endl;
    std::random_shuffle ( model.indexes.begin(), model.indexes.end() );
    for (auto batch = 0; batch < num_batches; ++batch)
    {
      model.from = batch*batchsize;
      model.to = min((batch+1)*batchsize,size);
      error = computeCostAndGrad(model,nullptr,gradient,number_vars);

      // seeing that I need to iterate ...
      for (int i = 0; i < number_vars; ++i) {
        Gt_d[i] += gradient[i]*gradient[i];
        // Update weight: ( eta / \sqrt(Sum square gradients) ) * gradient
        if (Gt_d[i] != 0) {
        update = vars[i] - ((eta / (sqrt(Gt_d[i]))) * gradient[i]);
        l1_reg = (eta / (sqrt(Gt_d[i]))) * lambda;
        vars[i] = signum(update) * max(0.0, abs(update) - l1_reg);
        }
      }
    }

    if (iteration % model.rae.config.dump_freq == 0)
    {
      printf("Dumping model ...\n");
      dumpModel(model,iteration);
    }
  }

  delete [] gradient;
  // delete [] Ginv_d;
  delete [] Gt_d;
  return 0;
}


//
// Copyright (c) 2016-2017 CNRS
//

#include <boost/fusion/container/generation/make_list.hpp>
#include <pinocchio/multibody/model.hpp>
#include "pinocchio/multibody/data.hpp"
#include "pinocchio/parsers/sample-models.hpp"
#include <pinocchio/algorithm/crba.hpp>
#include <pinocchio/algorithm/aba.hpp>
#include <pinocchio/algorithm/check.hpp>
#include <pinocchio/algorithm/default-check.hpp>
#include <iostream>

using namespace pinocchio;

#include <boost/test/unit_test.hpp>
#include <boost/utility/binary.hpp>

// Dummy checker.
struct Check1 : public AlgorithmCheckerBase<Check1>
{
  bool checkModel_impl( const Model& ) const { return true; }
};

BOOST_AUTO_TEST_SUITE ( BOOST_TEST_MODULE )

BOOST_AUTO_TEST_CASE ( test_check )
{
  using namespace boost::fusion;

  pinocchio::Model model; buildModels::humanoidRandom(model);

  BOOST_CHECK(model.check (Check1()));
  BOOST_CHECK(model.check (CRBAChecker()));
  BOOST_CHECK(! model.check (ABAChecker())); // some inertias are negative ... check fail.
  BOOST_FOREACH(Inertia& Y,model.inertias)
      Y.mass() = Y.inertia().data()[0] = Y.inertia().data()[3] = Y.inertia().data()[5] = 1.0 ;
  BOOST_CHECK(model.check (ABAChecker())); // some inertias are negative ... check fail.

  BOOST_CHECK(model.check(makeAlgoCheckerList(Check1(),ParentChecker(),CRBAChecker()) ));
  BOOST_CHECK(model.check(DEFAULT_CHECKERS));

  pinocchio::Data data(model);
  BOOST_CHECK(checkData(model,data));
  BOOST_CHECK(model.check(data));

}

BOOST_AUTO_TEST_SUITE_END ()


// Copyright (C) 2015 The Regents of the University of California (Regents).
// All rights reserved.
//
// Redistribution and use in source and binary forms, with or without
// modification, are permitted provided that the following conditions are
// met:
//
//     * Redistributions of source code must retain the above copyright
//       notice, this list of conditions and the following disclaimer.
//
//     * Redistributions in binary form must reproduce the above
//       copyright notice, this list of conditions and the following
//       disclaimer in the documentation and/or other materials provided
//       with the distribution.
//
//     * Neither the name of The Regents or University of California nor the
//       names of its contributors may be used to endorse or promote products
//       derived from this software without specific prior written permission.
//
// THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
// AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
// IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
// ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDERS OR CONTRIBUTORS BE
// LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR
// CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF
// SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS
// INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN
// CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE)
// ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
// POSSIBILITY OF SUCH DAMAGE.
//
// Please contact the author of this library if you have any questions.
// Author: Chris Sweeney (cmsweeney@cs.ucsb.edu)

#include "theia/sfm/global_pose_estimation/linear_rotation_estimator.h"

#include <Eigen/Core>
#include <Eigen/SparseCore>
#include <Eigen/SparseQR>
#include <ceres/rotation.h>
#include <unordered_map>
#include <vector>

#include "spectra/include/SymEigsShiftSolver.h"

#include "theia/math/matrix/spectra_linear_operator.h"
#include "theia/sfm/pose/util.h"
#include "theia/sfm/twoview_info.h"
#include "theia/sfm/types.h"
#include "theia/util/hash.h"
#include "theia/util/map_util.h"

namespace theia {
using Eigen::Matrix;
using Eigen::Matrix3d;
using Eigen::MatrixXd;

typedef Eigen::Triplet<double> TripletEntry;

namespace {

void Fill3x3SparseMatrix(const Eigen::Matrix3d& mat,
                         const int row,
                         const int col,
                         std::vector<TripletEntry>* triplet_list) {
  for (int i = 0; i < 3; i++) {
    for (int j = 0; j < 3; j++) {
      // Only add the value if it is not zero.
      triplet_list->emplace_back(row + i, col + j, mat(i, j));
    }
  }
}

}  // namespace

// Estimates the global orientations of all views based on an initial
// guess. Returns true on successful estimation and false otherwise.
bool LinearRotationEstimator::EstimateRotations(
    const std::unordered_map<ViewIdPair, TwoViewInfo>& view_pairs,
    std::unordered_map<ViewId, Eigen::Vector3d>* global_orientations) {
  for (const auto& view_pair : view_pairs) {
    AddRelativeRotationConstraint(view_pair.first, view_pair.second.rotation_2);
  }
  return EstimateRotations(global_orientations);
}

// Add the relative rotation matrix constraint that minimizes the Frobenius norm
// of the matrices. That is, set up the linear system such that
// ||R_j - R_ij * R_i|| is minimized
void LinearRotationEstimator::AddRelativeRotationConstraint(
    const ViewIdPair& view_id_pair, const Eigen::Vector3d& relative_rotation) {
  static const int kNumRotationMatrixDimensions = 3;

  // Add a new view-id to sparse matrix index mapping. This is a no-op if the
  // view has already been assigned to a matrix index.
  InsertIfNotPresent(&view_id_map_, view_id_pair.first, view_id_map_.size());
  InsertIfNotPresent(&view_id_map_, view_id_pair.second, view_id_map_.size());
  const int view1_index = view_id_map_[view_id_pair.first];
  const int view2_index = view_id_map_[view_id_pair.second];

  // Add the corresponding entries to A^t * A. Note that for each row, we may
  // represent the row as a block matrix: [B | C]. Thus, the corresponding
  // entries for A^t * A are:
  //
  //   [B | C]^t * [B | C] = [B^t * B | B^t * C]
  //                         [C^t * B | C^t * C]
  //
  // Note that for rotation matrices, B^t * B = C^t * B = I. Thus we have:
  //
  //   [B | C]^t * [B | C] = [   I    | B^t * C]
  //                         [C^t * B |    I   ]
  //
  // We only store the upper triangular portion of the matrix since it is
  // symmetric. In our case, one of B or C will be the identity, which further
  // simplifies the matrix entries.
  //
  // First, add the identity entries along the diagonal.
  for (int i = 0; i < 3; i++) {
    constraint_entries_.emplace_back(
        kNumRotationMatrixDimensions * view1_index + i,
        kNumRotationMatrixDimensions * view1_index + i,
        1.0);
    constraint_entries_.emplace_back(
        kNumRotationMatrixDimensions * view2_index + i,
        kNumRotationMatrixDimensions * view2_index + i,
        1.0);
  }

  // Add the 3x3 matrix B^t * C. This corresponds either to -R_ij or -R_ij^t
  // depending on the order of the view indices.
  Eigen::Matrix3d relative_rotation_matrix;
  ceres::AngleAxisToRotationMatrix(
      relative_rotation.data(),
      ceres::ColumnMajorAdapter3x3(relative_rotation_matrix.data()));

  // We seek to insert B^t * C into the linear system as noted above. If the
  // view1 index comes before the view2 index, then B = -R_ij and C =
  // I. Otherwise, B = I and C = -R_ij.
  if (view1_index < view2_index) {
    Fill3x3SparseMatrix(-relative_rotation_matrix.transpose(),
                        kNumRotationMatrixDimensions * view1_index,
                        kNumRotationMatrixDimensions * view2_index,
                        &constraint_entries_);
  } else {
    Fill3x3SparseMatrix(-relative_rotation_matrix,
                        kNumRotationMatrixDimensions * view2_index,
                        kNumRotationMatrixDimensions * view1_index,
                        &constraint_entries_);
  }
}

// Given the relative rotation constraints added with
// AddRelativeRotationConstraint, this method returns the robust estimation of
// global camera orientations. Like the method above, this requires an initial
// estimate of the global orientations.
bool LinearRotationEstimator::EstimateRotations(
    std::unordered_map<ViewId, Eigen::Vector3d>* global_orientations) {
  CHECK_GT(constraint_entries_.size(), 0);
  CHECK_NOTNULL(global_orientations);
  static const int kNumRotationMatrixDimensions = 3;

  // Setup the sparse linear system.
  Eigen::SparseMatrix<double> constraint_matrix(
      view_id_map_.size() * kNumRotationMatrixDimensions,
      view_id_map_.size() * kNumRotationMatrixDimensions);
  constraint_matrix.setFromTriplets(constraint_entries_.begin(),
                                    constraint_entries_.end());

  // Compute the 3 eigenvectors corresponding to the smallest eigenvalues. These
  // orthogonal vectors will contain the solution rotation matrices.
  SparseSymShiftSolveLLT op(constraint_matrix);
  Spectra::
      SymEigsShiftSolver<double, Spectra::LARGEST_MAGN, SparseSymShiftSolveLLT>
          eigs(&op, 3, 6, 0.0);
  eigs.init();
  eigs.compute();

  // The solution appears in the first three eigenvectors.
  const Eigen::MatrixXd solution =
      eigs.eigenvectors().leftCols<kNumRotationMatrixDimensions>();

  // Project all solutions into a valid SO3 rotation space. The linear system
  // above makes no constraint on the space of the solutions, so the final
  // solutions are not guaranteed to be valid rotations (e.g., det(R) may not be
  // +1).
  global_orientations->reserve(view_id_map_.size());
  for (const auto& view_id_map : view_id_map_) {
    const Matrix3d non_so3_rotation =
        solution
            .block<kNumRotationMatrixDimensions, kNumRotationMatrixDimensions>(
                kNumRotationMatrixDimensions * view_id_map.second, 0);
    const Matrix3d rotation = ProjectToRotationMatrix(non_so3_rotation);

    // Convert to angle axis.
    Eigen::Vector3d rotation_aa;
    ceres::RotationMatrixToAngleAxis(
        ceres::ColumnMajorAdapter3x3(rotation.data()), rotation_aa.data());
    global_orientations->emplace(view_id_map.first, rotation_aa);
  }

  return true;
}

bool LinearRotationEstimator::EstimateRotationsWrapper(
    const std::unordered_map<ViewIdPair, TwoViewInfo>& view_pairs,
    std::unordered_map<ViewId, Eigen::Vector3d>& global_orientations) {
  return EstimateRotations(view_pairs, &global_orientations);
}

}  // namespace theia


#define BOOST_TEST_MODULE "test_bond_length_interaction"

#ifdef BOOST_TEST_DYN_LINK
#include <boost/test/unit_test.hpp>
#else
#include <boost/test/included/unit_test.hpp>
#endif

#include <test/util/utility.hpp>

#include <mjolnir/core/BoundaryCondition.hpp>
#include <mjolnir/core/SimulatorTraits.hpp>
#include <mjolnir/forcefield/local/BondLengthInteraction.hpp>
#include <mjolnir/forcefield/local/HarmonicPotential.hpp>
#include <mjolnir/util/make_unique.hpp>

#include <random>

BOOST_AUTO_TEST_CASE(BondLength_numerical_difference)
{
    namespace test = mjolnir::test;

    mjolnir::LoggerManager::set_default_logger("test_bondlength_interaction.log");

    using traits_type      = mjolnir::SimulatorTraits<double, mjolnir::UnlimitedBoundary>;
    using real_type        = traits_type::real_type;
    using coord_type       = traits_type::coordinate_type;
    using boundary_type    = traits_type::boundary_type;
    using system_type      = mjolnir::System<traits_type>;
    using potential_type   = mjolnir::HarmonicPotential<real_type>;
    using interaction_type = mjolnir::BondLengthInteraction<traits_type, potential_type>;

    const real_type k(100.0);
    const real_type native(std::sqrt(3.0));

    potential_type   potential(k, native);
    interaction_type interaction("none", {{ {{0,1}}, potential}});

    std::mt19937 mt(123456789);

    for(std::size_t i = 0; i < 1000; ++i)
    {
        system_type sys(2, boundary_type{});
        test::clear_everything(sys);

        sys.position(0) = coord_type( 0.0, 0.0, 0.0);
        sys.position(1) = coord_type( 1.0, 1.0, 1.0);

        test::apply_random_rotation(sys, mt);
        test::apply_random_perturbation(sys, mt, 0.01);

        constexpr real_type tol = 1e-6;
        constexpr real_type dr  = 1e-6;

        test::check_force(sys, interaction, tol, dr);
        test::check_virial(sys, interaction, tol);
        test::check_force_and_virial(sys, interaction, tol);
        test::check_force_and_energy(sys, interaction, tol);
        test::check_force_energy_virial(sys, interaction, tol);
    }
}

BOOST_AUTO_TEST_CASE(BondLength_calc_force)
{
    namespace test = mjolnir::test;

    mjolnir::LoggerManager::set_default_logger("test_bondlength_interaction.log");

    using traits_type      = mjolnir::SimulatorTraits<double, mjolnir::UnlimitedBoundary>;
    using real_type        = traits_type::real_type;
    using coord_type       = traits_type::coordinate_type;
    using boundary_type    = traits_type::boundary_type;
    using system_type      = mjolnir::System<traits_type>;
    using potential_type    = mjolnir::HarmonicPotential<real_type>;
    using interaction_type = mjolnir::BondLengthInteraction<traits_type, potential_type>;

    constexpr real_type tol = 1e-8;

    auto normalize = [](const coord_type& v){return v / mjolnir::math::length(v);};

    const real_type k(100.);
    const real_type native(2.0);

    potential_type    potential(k, native);
    interaction_type interaction("none", {{ {{0,1}}, potential}});

    system_type sys(2, boundary_type{});
    test::clear_everything(sys);

    const real_type dr = 1e-3;
    real_type dist = 1e0;
    for(int i = 0; i < 2000; ++i)
    {
        sys.position(0) = coord_type(0,0,0);
        sys.position(1) = coord_type(0,0,0);
        sys.force(0)    = coord_type(0,0,0);
        sys.force(1)    = coord_type(0,0,0);
        sys.position(1)[0] = dist;

        const real_type deriv = potential.derivative(dist);
        const real_type coef  = std::abs(deriv);

        interaction.calc_force(sys);

        const real_type force_strength1 = mjolnir::math::length(sys.force(0));
        const real_type force_strength2 = mjolnir::math::length(sys.force(1));

        // direction
        if(i == 1000) // most stable point
        {
            BOOST_TEST(force_strength1 == 0.0, boost::test_tools::tolerance(tol));
            BOOST_TEST(force_strength2 == 0.0, boost::test_tools::tolerance(tol));
        }
        else if(i < 1000) // repulsive
        {
            BOOST_TEST(coef == force_strength1, boost::test_tools::tolerance(tol));
            BOOST_TEST(coef == force_strength2, boost::test_tools::tolerance(tol));

            const real_type dir1 = mjolnir::math::dot_product(
                normalize(sys.force(0)), normalize(sys.position(0) - sys.position(1)));
            const real_type dir2 = mjolnir::math::dot_product(
                normalize(sys.force(1)), normalize(sys.position(1) - sys.position(0)));

            BOOST_TEST(dir1 == 1.0, boost::test_tools::tolerance(tol));
            BOOST_TEST(dir2 == 1.0, boost::test_tools::tolerance(tol));
        }
        else if(i > 1000) // attractive
        {
            BOOST_TEST(coef == force_strength1, boost::test_tools::tolerance(tol));
            BOOST_TEST(coef == force_strength2, boost::test_tools::tolerance(tol));

            const real_type dir1 = mjolnir::math::dot_product(
                normalize(sys.force(0)), normalize(sys.position(1) - sys.position(0)));
            const real_type dir2 = mjolnir::math::dot_product(
                normalize(sys.force(1)), normalize(sys.position(0) - sys.position(1)));

            BOOST_TEST(dir1 == 1e0, boost::test_tools::tolerance(tol));
            BOOST_TEST(dir2 == 1e0, boost::test_tools::tolerance(tol));
        }
        BOOST_TEST(mjolnir::math::length(sys.force(0) + sys.force(1)) == 0.0,
                   boost::test_tools::tolerance(tol));

        dist += dr;
    }
}


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

// Copyright (c) 2015, Oracle and/or its affiliates.

// Contributed and/or modified by Menelaos Karavelas, on behalf of Oracle

// Distributed under the Boost Software License, Version 1.0.
// (See accompanying file LICENSE_1_0.txt or copy at
// http://www.boost.org/LICENSE_1_0.txt)

#ifndef BOOST_GEOMETRY_ALGORITHMS_DETAIL_ENVELOPE_INTERSECTS_ANTIMERIDIAN_HPP
#define BOOST_GEOMETRY_ALGORITHMS_DETAIL_ENVELOPE_INTERSECTS_ANTIMERIDIAN_HPP

#include <boost/geometry/core/access.hpp>
#include <boost/geometry/core/coordinate_system.hpp>

#include <boost/geometry/util/math.hpp>

#include <boost/geometry/algorithms/detail/normalize.hpp>

namespace boost { namespace geometry
{

namespace detail { namespace envelope
{

struct intersects_antimeridian
{
    template <typename Units, typename CoordinateType>
    static inline bool apply(CoordinateType const& lon1,
                             CoordinateType const& lat1,
                             CoordinateType const& lon2,
                             CoordinateType const& lat2)
    {
        typedef math::detail::constants_on_spheroid
            <
                CoordinateType, Units
            > constants;

        return
            math::equals(math::abs(lat1), constants::max_latitude())
            ||
            math::equals(math::abs(lat2), constants::max_latitude())
            ||
            math::larger(math::abs(lon1 - lon2), constants::half_period());
    }

    template <typename Segment>
    static inline bool apply(Segment const& segment)
    {
        return apply(detail::indexed_point_view<Segment, 0>(segment),
                     detail::indexed_point_view<Segment, 1>(segment));
    }

    template <typename Point>
    static inline bool apply(Point const& p1, Point const& p2)
    {
        Point p1_normalized = detail::return_normalized<Point>(p1);
        Point p2_normalized = detail::return_normalized<Point>(p2);

        return apply
            <
                typename coordinate_system<Point>::type::units
            >(geometry::get<0>(p1_normalized),
              geometry::get<1>(p1_normalized),
              geometry::get<0>(p2_normalized),
              geometry::get<1>(p2_normalized));
    }
};

}} // namespace detail::envelope

}} // namespace boost::geometry

#endif // BOOST_GEOMETRY_ALGORITHMS_DETAIL_ENVELOPE_INTERSECTS_ANTIMERIDIAN_HPP


//==============================================================================
//         Copyright 2003 - 2011 LASMEA UMR 6602 CNRS/Univ. Clermont II
//         Copyright 2009 - 2011 LRI    UMR 8623 CNRS/Univ Paris Sud XI
//
//          Distributed under the Boost Software License, Version 1.0.
//                 See accompanying file LICENSE.txt or copy at
//                     http://www.boost.org/LICENSE_1_0.txt
//==============================================================================
#ifndef BOOST_SIMD_TOOLBOX_ARITHMETIC_FUNCTIONS_SCALAR_MULS_HPP_INCLUDED
#define BOOST_SIMD_TOOLBOX_ARITHMETIC_FUNCTIONS_SCALAR_MULS_HPP_INCLUDED

#include <boost/simd/toolbox/arithmetic/functions/muls.hpp>
#include <boost/simd/include/functions/scalar/genmask.hpp>
#include <boost/simd/include/functions/scalar/abs.hpp>
#include <boost/simd/include/functions/scalar/bitofsign.hpp>
#include <boost/simd/include/functions/scalar/saturate.hpp>
#include <boost/simd/include/functions/scalar/bitwise_xor.hpp>
#include <boost/simd/include/constants/valmin.hpp>
#include <boost/simd/include/constants/valmax.hpp>
#include <boost/simd/include/constants/zero.hpp>
#include <boost/dispatch/meta/upgrade.hpp>
#include <boost/dispatch/meta/as_integer.hpp>
#
namespace boost { namespace simd { namespace ext
{
  BOOST_SIMD_FUNCTOR_IMPLEMENTATION( boost::simd::tag::muls_, tag::cpu_
                            , (A0)
                            , (scalar_< unsigned_<A0> >)(scalar_< unsigned_<A0> >)
                            )
  {
    typedef A0 result_type;
    BOOST_SIMD_FUNCTOR_CALL_REPEAT(2)
    {
      typedef typename dispatch::meta::upgrade<A0>::type utype;
      utype res = utype(a0)*utype(a1);
      return static_cast<A0>(res) | genmask(static_cast<A0>(res >> sizeof(A0)*8));
    }
  };

  BOOST_SIMD_FUNCTOR_IMPLEMENTATION( boost::simd::tag::muls_, tag::cpu_
                            , (A0)
                            , (scalar_< integer_<A0> >)(scalar_< integer_<A0> >)
                            )
  {

    typedef A0 result_type;
    BOOST_SIMD_FUNCTOR_CALL_REPEAT(2)
    {
      typedef typename dispatch::meta::upgrade<A0>::type uptype;
      return A0(saturate<A0>(static_cast<uptype>(a0)*static_cast<uptype>(a1)));
    }
  };

  BOOST_SIMD_FUNCTOR_IMPLEMENTATION( boost::simd::tag::muls_, tag::cpu_
                            , (A0)
                            , (scalar_< uint64_<A0> >)(scalar_< uint64_<A0> >)
                            )
  {
    typedef A0 result_type;
    BOOST_SIMD_FUNCTOR_CALL_REPEAT(2)
    {
      if (a1 == 0 || a0 == 0) return Zero<A0>();
      if (a1 >= a0)
      {
        A0 z = Valmax<A0>()/a1;
        if (z < a0)
          return  Valmax<A0>();
        else
          return a0*a1;
      }
      else
        return muls(a1, a0);
    }
  };

  BOOST_SIMD_FUNCTOR_IMPLEMENTATION( boost::simd::tag::muls_, tag::cpu_
                            , (A0)
                            , (scalar_< int64_<A0> >)(scalar_< int64_<A0> >)
                            )
  {
    typedef A0 result_type;
    BOOST_SIMD_FUNCTOR_CALL_REPEAT(2)
    {
      if (a1 == 0 || a0 == 0) return Zero<result_type>();
      typedef typename dispatch::meta::as_integer<result_type, unsigned>::type untype;
      result_type sign =  b_xor(bitofsign(a0), bitofsign(a1));
      untype aa0 = boost::simd::abs(a0);
      untype aa1 = boost::simd::abs(a1);
      if (aa1 >= aa0)
      {
        untype z = Valmax<result_type>()/aa1;
        if (z < aa0)
          return sign ? Valmin<result_type>():Valmax<result_type>();
        else
          return a0*a1;
      }
      else
        return muls(a1, a0);
    }
  };
} } }

#endif


#include "astar.h"
#include <Eigen/Dense>
#include <math.h>
#include <vector>
#include <map>
#include <stdio.h>
#include <iostream>
#include <queue>
#include "forward_kinematics.h"
#include <iomanip>
#include <unordered_map>
#include <bitset>
#include <string>
#include <utility>

Astar::Astar(){
  target_threshold = .025;
  value_comp_time = 0;
  get_visdata_time = 0;
  frontier_insert_time = 0;
  mem_time = 0;
}

std::vector<configState*> Astar::run(configState* s, wsState* t){

  std::cout << "Initializing Start and Target\n";

  target = t;
  start = s;
  std::cout << "Target: ";
  wsstate_tostring(target);

  std::cout << "Initialize start data structures\n";

  /* Runtime checks */
  float avg_frontier_pop = 0;
  int pop_count = 0;

  float avg_vis_add = 0;
  int vis_add_count = 0;

  float avg_expansion = 0;
  int expansion_count = 0;

  /* Init */
  visData* startv = create_visdata();
  startv->current = start;
  startv->prev = 0;
  startv->value = heuristic(start);
  frontier.push(startv);
  visData* next = 0;
  wsState* wscurrent = fk(start);
  while(distance(target, wscurrent) > target_threshold){
    clock_t pqstart = clock();
    next = frontier.top();
    frontier.pop();
    float secsElapsed = (float)(clock() - pqstart)/CLOCKS_PER_SEC;
    avg_frontier_pop += secsElapsed;
    pop_count++;
    delete(wscurrent);
    wscurrent = fk(next->current);

    if(next != 0){
      clock_t visaddstart = clock();
      add_visited(next->current, next);
      secsElapsed= (float)(clock() - visaddstart)/CLOCKS_PER_SEC;
      avg_vis_add += secsElapsed;
      vis_add_count ++;

      clock_t expstart = clock();
      expand_frontier(next->current);
      secsElapsed = (float)(clock() - expstart)/CLOCKS_PER_SEC;
      avg_expansion += secsElapsed;
      expansion_count++;
    } else {
      std::vector<configState*> path;
      return path;
    }

  }

  std::cout << "Visited Set: "<< visited_set.size() << '\n';
  std::cout << "Queue: "<< frontier.size() << '\n';

  std::cout << "Completed Search\n";

  std::cout << "Runtimes: \n\n\n\n";
  std::cout << "Frontier Popping :" << avg_frontier_pop << '\n';
  std::cout << "Expansion: " << avg_expansion << '\n';
  std::cout << "  Vis Gets: " << get_visdata_time << '\n';
  std::cout << "  Value Comp Time: " << value_comp_time << '\n';
  std::cout << "  Frontier Push Time: " << frontier_insert_time << '\n';
  std::cout << "  Mem_Time: " << mem_time << '\n';
  std::cout << "  VisAdd :" << avg_vis_add*3 << '\n';
  std::cout << "Counts: " << pop_count << '\n';

  visData* current = next;

  std::vector<configState*> path;

  std::cout << "Backtracing path\n";
  while(current != 0){
    path.push_back(current->current);
    if(current -> prev == 0){
      current = 0;
      continue;
    }
    current = get_visdata(current->prev);
  }
  std::cout << "Back trace completed\n";
  return path;
}

//Determines how close a configuration is to the target by measuring Euclidean
//  distance in the workspace
double Astar::heuristic(configState* c)
{
  wsState* state = fk(c);
  double d = distance(state,target);
  delete(state);
  return d;
}//end heurisitic

double Astar::cost(configState* c1,configState* c2)
{
  wsState* state1 = fk(c1);
  wsState* state2 = fk(c2);
  double workdist = distance(state1,state2);
  delete(state1);
  delete(state2);
  //TODO experiment with various angle penalties to determine the optimum path
  return workdist;
}

void Astar::expand_frontier(configState* c){
  if(c == 0){
    printf("not configed\n");
  }
  for(int i = 1; i < 4; i++){
    configState* pos = 0;
    configState* neg = 0;

    // clock_t memstart = clock();

    pos = clone_configstate(c);

    pos->theta[i] = c->theta[i]+1;

    neg = clone_configstate(c);
    neg->theta[i] = c->theta[i]-1;

    // mem_time += (float)(clock() - memstart)/CLOCKS_PER_SEC;

    // clock_t getvisstart = clock();
    visData* cv = get_visdata(c);
    // get_visdata_time += (float)(clock() - getvisstart)/CLOCKS_PER_SEC;
    if(!has_visited(pos)){
      visData* posv;
      // memstart = clock();

      posv = create_visdata();

      posv->current = pos;
      posv->prev = c;
      // mem_time += (float)(clock() - memstart)/CLOCKS_PER_SEC;

      clock_t getvaluestart = clock();
      posv->value = heuristic(pos)+cost(c, pos)+cv->value;
      // std::cout << "Value: " <<std::setprecision(30) << posv->value << '\n';
      value_comp_time += (float)(clock() - getvaluestart)/CLOCKS_PER_SEC;
      // clock_t fpstart = clock();
      frontier.push(posv);
      // frontier_insert_time +=(float)(clock() - fpstart)/CLOCKS_PER_SEC;

    } else {
      // memstart = clock();
      deallocate_configstate(pos);
      // mem_time += (float)(clock() - memstart)/CLOCKS_PER_SEC;

    }

    if(!has_visited(neg)){
      visData* negv;
      // memstart = clock();

      negv = create_visdata();
      negv->current = neg;
      negv->prev = c;
      // mem_time += (float)(clock() - memstart)/CLOCKS_PER_SEC;
       clock_t getvaluestart = clock();

      negv->value = heuristic(neg)+cost(c, neg)+cv->value;
       value_comp_time += (float)(clock() - getvaluestart)/CLOCKS_PER_SEC;
      // clock_t fpstart = clock();

      frontier.push(negv);
      // frontier_insert_time +=(float)(clock() - fpstart)/CLOCKS_PER_SEC;

    } else {
      // memstart = clock();
      deallocate_configstate(neg);
      // mem_time += (float)(clock() - memstart)/CLOCKS_PER_SEC;
      }

  }
  // std::cout << '\n';
}

bool Astar::has_visited(configState* c){
  bool found = true;
  try {
    visited_set.at(c);
  }
  catch(...) {
    found = false;
  }
  return found;
}

void Astar::add_visited(configState* c, visData* v){
  visited_set.insert(std::pair<configState*, visData*>(c,v));
}

visData*  Astar::get_visdata(configState* c){
  return visited_set.at(c);
}


#pragma once
#include <unordered_set>
#include <numeric>
#include "TopicModel.hpp"
#include <Eigen/Dense>
#include "../Utils/Utils.hpp"
#include "../Utils/math.h"
#include "../Utils/sample.hpp"

/*
Implementation of LDA using Collapsed Variational Bayes zero-order estimation by bab2min

* Blei, D. M., Ng, A. Y., & Jordan, M. I. (2003). Latent dirichlet allocation. Journal of machine Learning research, 3(Jan), 993-1022.

Term Weighting Scheme is based on following paper:
* Wilson, A. T., & Chew, P. A. (2010, June). Term weighting schemes for latent dirichlet allocation. In human language technologies: The 2010 annual conference of the North American Chapter of the Association for Computational Linguistics (pp. 465-473). Association for Computational Linguistics.

*/

#define GETTER(name, type, field) type get##name() const override { return field; }
namespace tomoto
{
	struct DocumentLDACVB0 : public DocumentBase
	{
	public:
		using DocumentBase::DocumentBase;

		Eigen::MatrixXf Zs;
		Eigen::VectorXf numByTopic;

		DEFINE_SERIALIZER_AFTER_BASE(DocumentBase, Zs);

		template<typename _TopicModel> void update(Float* ptr, const _TopicModel& mdl);

		int32_t getSumWordWeight() const
		{
			return this->words.size();
		}
	};

	struct ModelStateLDACVB0
	{
		Eigen::VectorXf zLikelihood;
		Eigen::VectorXf numByTopic;
		Eigen::MatrixXf numByTopicWord;

		DEFINE_SERIALIZER(numByTopic, numByTopicWord);
	};

	class ILDACVB0Model : public ITopicModel
	{
	public:
		using DefaultDocType = DocumentLDACVB0;
		static ILDACVB0Model* create(size_t _K = 1, Float _alpha = 0.1, Float _eta = 0.01, size_t _rg = std::random_device{}());

		virtual size_t addDoc(const std::vector<std::string>& words) = 0;
		virtual std::unique_ptr<DocumentBase> makeDoc(const std::vector<std::string>& words) const = 0;
		TermWeight getTermWeight() const { return TermWeight::one; };
		virtual size_t getOptimInterval() const = 0;
		virtual void setOptimInterval(size_t) = 0;
		virtual void setBurnInIteration(size_t) {}
		virtual std::vector<size_t> getCountByTopic() const = 0;
		virtual size_t getK() const = 0;
		virtual Float getAlpha() const = 0;
		virtual Float getEta() const = 0;

		virtual std::vector<Float> getWordPrior(const std::string& word) const { return {}; }
		virtual void setWordPrior(const std::string& word, const std::vector<Float>& priors) {}
	};

	template<typename _Interface = ILDACVB0Model,
		typename _Derived = void,
		typename _DocType = DocumentLDACVB0,
		typename _ModelState = ModelStateLDACVB0>
	class LDACVB0Model : public TopicModel<0, _Interface,
		typename std::conditional<std::is_same<_Derived, void>::value, LDACVB0Model<>, _Derived>::type,
		_DocType, _ModelState>
	{
	protected:
		using DerivedClass = typename std::conditional<std::is_same<_Derived, void>::value, LDACVB0Model, _Derived>::type;
		using BaseClass = TopicModel<0, _Interface, DerivedClass, _DocType, _ModelState>;
		friend BaseClass;

		static constexpr const char TWID[] = "one\0";
		static constexpr static constexpr char TMID[] = "LDA\0";

		Float alpha;
		Vector alphas;
		Float eta;
		Tid K;
		size_t optimInterval = 50;

		template<typename _List>
		static Float calcDigammaSum(_List list, size_t len, Float alpha)
		{
			auto listExpr = Vector::NullaryExpr(len, list);
			auto dAlpha = math::digammaT(alpha);
			return (math::digammaApprox(listExpr.array() + alpha) - dAlpha).sum();
		}

		void optimizeParameters(ThreadPool& pool, _ModelState* localData)
		{
			const auto K = this->K;
			for (size_t i = 0; i < 5; ++i)
			{
				Float denom = calcDigammaSum([&](size_t i) { return this->docs[i].getSumWordWeight(); }, this->docs.size(), alphas.sum());
				for (size_t k = 0; k < K; ++k)
				{
					Float nom = calcDigammaSum([&](size_t i) { return this->docs[i].numByTopic[k]; }, this->docs.size(), alphas(k));
					alphas(k) = std::max(nom / denom * alphas(k), 1e-5f);
				}
			}
		}

		const Eigen::VectorXf& getZLikelihoods(_ModelState& ld, const _DocType& doc, size_t docId, size_t vid) const
		{
			const size_t V = this->realV;
			assert(vid < V);
			auto& zLikelihood = ld.zLikelihood;
			zLikelihood = (doc.numByTopic.array().template cast<Float>() + alphas.array())
				* (ld.numByTopicWord.col(vid).array().template cast<Float>() + eta)
				/ (ld.numByTopic.array().template cast<Float>() + V * eta);
			zLikelihood /= zLikelihood.sum() + 1e-10;
			return zLikelihood;
		}

		template<int _Inc, typename _Vec>
		inline void addWordTo(_ModelState& ld, _DocType& doc, uint32_t pid, Vid vid, _Vec tDist) const
		{
			assert(vid < this->realV);
			constexpr bool _dec = _Inc < 0;
			doc.numByTopic += _Inc * tDist;
			if (_dec) doc.numByTopic = doc.numByTopic.cwiseMax(0);
			ld.numByTopic += _Inc * tDist;
			if (_dec) ld.numByTopic = ld.numByTopic.cwiseMax(0);
			ld.numByTopicWord.col(vid) += _Inc * tDist;
			if (_dec) ld.numByTopicWord.col(vid) = ld.numByTopicWord.col(vid).cwiseMax(0);
		}

		template<ParallelScheme _ps, bool _infer, typename _ExtraDocData>
		void sampleDocument(_DocType& doc, const _ExtraDocData& edd, size_t docId, _ModelState& ld, _RandGen& rgs, size_t iterationCnt, size_t partitionId = 0) const
		{
			for (size_t w = 0; w < doc.words.size(); ++w)
			{
				if (doc.words[w] >= this->realV) continue;
				addWordTo<-1>(ld, doc, w, doc.words[w], doc.Zs.col(w));
				doc.Zs.col(w) = static_cast<const DerivedClass*>(this)->getZLikelihoods(ld, doc, docId, doc.words[w]);
				addWordTo<1>(ld, doc, w, doc.words[w], doc.Zs.col(w));
			}
		}

		template<typename _DocIter, typename _ExtraDocData>
		void updatePartition(ThreadPool& pool, _ModelState* localData, _DocIter first, _DocIter last, _ExtraDocData& edd)
		{
		}

		template<ParallelScheme _ps>
		void trainOne(ThreadPool& pool, _ModelState* localData, _RandGen* rgs)
		{
			std::vector<std::future<void>> res;
			const size_t chStride = std::min(pool.getNumWorkers() * 8, this->docs.size());
			for (size_t ch = 0; ch < chStride; ++ch)
			{
				res.emplace_back(pool.enqueue([&, this, ch, chStride](size_t threadId)
				{
					forShuffled((this->docs.size() - 1 - ch) / chStride + 1, rgs[threadId](), [&, this](size_t id)
					{
						static_cast<DerivedClass*>(this)->template sampleDocument<ParallelScheme::copy_merge>(
							this->docs[id * chStride + ch], 0, id * chStride + ch,
							localData[threadId], rgs[threadId], this->globalStep);
					});
				}));
			}
			for (auto& r : res) r.get();
			static_cast<DerivedClass*>(this)->updateGlobalInfo(pool, localData);
			static_cast<DerivedClass*>(this)->mergeState(pool, this->globalState, this->tState, localData);
			if (this->globalStep >= 250 && optimInterval && (this->globalStep + 1) % optimInterval == 0)
			{
				static_cast<DerivedClass*>(this)->optimizeParameters(pool, localData);
			}
		}

		void updateGlobalInfo(ThreadPool& pool, _ModelState* localData)
		{
			std::vector<std::future<void>> res;

			this->globalState.numByTopic.setZero();
			this->globalState.numByTopicWord.setZero();
			for (auto& doc : this->docs)
			{
				doc.numByTopic = doc.Zs.rowwise().sum();
				this->globalState.numByTopic += doc.numByTopic;
				for (size_t i = 0; i < doc.words.size(); ++i)
				{
					this->globalState.numByTopicWord.col(doc.words[i]) += doc.Zs.col(i);
				}
			}

			for (size_t i = 0; i < pool.getNumWorkers(); ++i)
			{
				res.emplace_back(pool.enqueue([&, i](size_t threadId)
				{
					localData[i] = this->globalState;
				}));
			}
			for (auto& r : res) r.get();
		}

		void mergeState(ThreadPool& pool, _ModelState& globalState, _ModelState& tState, _ModelState* localData) const
		{
		}

		template<typename _DocIter>
		double getLLDocs(_DocIter _first, _DocIter _last) const
		{
			double ll = 0;
			// doc-topic distribution
			ll += (math::lgammaT(K*alpha) - math::lgammaT(alpha)*K) * std::distance(_first, _last);
			for (; _first != _last; ++_first)
			{
				auto& doc = *_first;
				ll -= math::lgammaT(doc.getSumWordWeight() + K * alpha);
				for (Tid k = 0; k < K; ++k)
				{
					ll += math::lgammaT(doc.numByTopic[k] + alpha);
				}
			}
			return ll;
		}

		double getLLRest(const _ModelState& ld) const
		{
			double ll = 0;
			const size_t V = this->realV;
			// topic-word distribution
			// it has the very-small-value problem
			ll += (math::lgammaT(V*eta) - math::lgammaT(eta)*V) * K;
			for (Tid k = 0; k < K; ++k)
			{
				ll -= math::lgammaT(ld.numByTopic[k] + V * eta);
				for (Vid v = 0; v < V; ++v)
				{
					ll += math::lgammaT(ld.numByTopicWord(k, v) + eta);
				}
			}
			return ll;
		}

		double getLL() const
		{
			return static_cast<const DerivedClass*>(this)->template getLLDocs<>(this->docs.begin(), this->docs.end())
				+ static_cast<const DerivedClass*>(this)->getLLRest(this->globalState);
		}

		void prepareShared()
		{
		}

		void prepareDoc(_DocType& doc, Float* topicDocPtr, size_t wordSize) const
		{
			doc.numByTopic = Eigen::VectorXf::Zero(K);
			doc.Zs = Eigen::MatrixXf::Zero(K, wordSize);
		}

		void initGlobalState(bool initDocs)
		{
			const size_t V = this->realV;
			this->globalState.zLikelihood = Vector::Zero(K);
			if (initDocs)
			{
				this->globalState.numByTopic = Vector::Zero(K);
				this->globalState.numByTopicWord = Matrix::Zero(K, V);
			}
		}

		struct Generator
		{
			std::uniform_int_distribution<Tid> theta;
		};

		Generator makeGeneratorForInit(const _DocType*) const
		{
			return Generator{ std::uniform_int_distribution<Tid>{0, (Tid)(K - 1)} };
		}

		template<bool _Infer>
		void updateStateWithDoc(Generator& g, _ModelState& ld, _RandGen& rgs, _DocType& doc, size_t i) const
		{
			doc.Zs.col(i).setZero();
			doc.Zs(g.theta(rgs), i) = 1;
			addWordTo<1>(ld, doc, i, doc.words[i], doc.Zs.col(i));
		}

		template<bool _Infer, typename _Generator>
		void initializeDocState(_DocType& doc, Float* topicDocPtr, _Generator& g, _ModelState& ld, _RandGen& rgs) const
		{
			std::vector<uint32_t> tf(this->realV);
			static_cast<const DerivedClass*>(this)->prepareDoc(doc, topicDocPtr, doc.words.size());

			for (size_t i = 0; i < doc.words.size(); ++i)
			{
				if (doc.words[i] >= this->realV) continue;
				static_cast<const DerivedClass*>(this)->template updateStateWithDoc<_Infer>(g, ld, rgs, doc, i);
			}
		}

		std::vector<uint64_t> _getTopicsCount() const
		{
			Eigen::VectorXf cnt = Eigen::VectorXf::Zero(K);
			for (auto& doc : this->docs)
			{
				cnt += doc.Zs.rowwise().sum();
			}

			return { cnt.data(), cnt.data() + K };
		}

		template<ParallelScheme _ps>
		size_t estimateMaxThreads() const
		{
			if (_ps == ParallelScheme::partition)
			{
				return this->realV / 4;
			}
			if (_ps == ParallelScheme::copy_merge)
			{
				return this->docs.size() / 2;
			}
			return (size_t)-1;
		}

		DEFINE_SERIALIZER(alpha, eta, K);

	public:
		LDACVB0Model(size_t _K = 1, Float _alpha = 0.1, Float _eta = 0.01, size_t _rg = std::random_device{}())
			: BaseClass(_rg), K(_K), alpha(_alpha), eta(_eta)
		{
			alphas = Vector::Constant(K, alpha);
		}
		GETTER(K, size_t, K);
		GETTER(Alpha, Float, alpha);
		GETTER(Eta, Float, eta);
		GETTER(OptimInterval, size_t, optimInterval);

		void setOptimInterval(size_t _optimInterval) override
		{
			optimInterval = _optimInterval;
		}

		size_t addDoc(const std::vector<std::string>& words) override
		{
			return this->_addDoc(this->_makeDoc(words));
		}

		std::unique_ptr<DocumentBase> makeDoc(const std::vector<std::string>& words) const override
		{
			return std::make_unique<_DocType>(as_mutable(this)->template _makeDoc<true>(words));
		}

		void updateDocs()
		{
			for (auto& doc : this->docs)
			{
				doc.template update<>(nullptr, *static_cast<DerivedClass*>(this));
			}
		}

		void prepare(bool initDocs = true, size_t minWordCnt = 0, size_t minWordDf = 0, size_t removeTopN = 0) override
		{
			if (initDocs) this->removeStopwords(minWordCnt, minWordDf, removeTopN);
			static_cast<DerivedClass*>(this)->updateWeakArray();
			static_cast<DerivedClass*>(this)->initGlobalState(initDocs);

			if (initDocs)
			{
				auto generator = static_cast<DerivedClass*>(this)->makeGeneratorForInit(nullptr);
				for (auto& doc : this->docs)
				{
					initializeDocState<false>(doc, nullptr, generator, this->globalState, this->rg);
				}
			}
			else
			{
				static_cast<DerivedClass*>(this)->updateDocs();
			}
			static_cast<DerivedClass*>(this)->prepareShared();
		}

		std::vector<size_t> getCountByTopic() const override
		{
			return static_cast<const DerivedClass*>(this)->_getTopicsCount();
		}

		std::vector<Float> getTopicsByDoc(const _DocType& doc) const
		{
			std::vector<Float> ret(K);
			Float sum = doc.getSumWordWeight() + K * alpha;
			transform(doc.numByTopic.data(), doc.numByTopic.data() + K, ret.begin(), [sum, this](size_t n)
			{
				return (n + alpha) / sum;
			});
			return ret;
		}

		std::vector<Float> _getWidsByTopic(Tid tid, bool normalize = true) const
		{
			assert(tid < K);
			const size_t V = this->realV;
			std::vector<Float> ret(V);
			Float sum = this->globalState.numByTopic[tid] + V * eta;
			auto r = this->globalState.numByTopicWord.row(tid);
			for (size_t v = 0; v < V; ++v)
			{
				ret[v] = (r[v] + eta) / sum;
			}
			return ret;
		}

		template<bool _Together, ParallelScheme _ps, typename _Iter>
		std::vector<double> _infer(_Iter docFirst, _Iter docLast, size_t maxIter, Float tolerance, size_t numWorkers) const
		{
			return {};
		}
	};

	template<typename _TopicModel>
	void DocumentLDACVB0::update(Float * ptr, const _TopicModel & mdl)
	{
		numByTopic = Eigen::VectorXf::Zero(mdl.getK());
		for (size_t i = 0; i < Zs.cols(); ++i)
		{
			numByTopic += Zs.col(i);
		}
	}

	inline ILDACVB0Model* ILDACVB0Model::create(size_t _K, Float _alpha, Float _eta, const _RandGen& _rg)
	{
		return new LDACVB0Model<>(_K, _alpha, _eta, _rg);
	}

}


// Copyright 2019, Collabora, Ltd.
// SPDX-License-Identifier: BSL-1.0
/*!
 * @file
 * @brief  Base implementations for math library.
 * @author Jakob Bornecrantz <jakob@collabora.com>
 * @author Ryan Pavlik <ryan.pavlik@collabora.com>
 * @ingroup aux_math
 */

#include <Eigen/Core>
#include <Eigen/Geometry>

#include <assert.h>

#include "math/m_api.h"
#include "math/m_eigen_interop.h"

/*
 *
 * Copy helpers.
 *
 */

static inline Eigen::Quaternionf
copy(const struct xrt_quat& q)
{
	// Eigen constructor order is different from XRT, OpenHMD and OpenXR!
	//  Eigen: `float w, x, y, z`.
	// OpenXR: `float x, y, z, w`.
	return Eigen::Quaternionf(q.w, q.x, q.y, q.z);
}

static inline Eigen::Quaternionf
copy(const struct xrt_quat* q)
{
	return copy(*q);
}

static inline Eigen::Vector3f
copy(const struct xrt_vec3& v)
{
	return Eigen::Vector3f(v.x, v.y, v.z);
}

static inline Eigen::Vector3f
copy(const struct xrt_vec3* v)
{
	return copy(*v);
}

/*
 *
 * Exported vector functions.
 *
 */
void
math_vec3_accum(const struct xrt_vec3* additional, struct xrt_vec3* inAndOut)
{
	assert(additional != NULL);
	assert(inAndOut != NULL);

	map_vec3(*inAndOut) += map_vec3(*additional);
}

/*
 *
 * Exported quaternion functions.
 *
 */

void
math_quat_rotate(const struct xrt_quat* left,
                 const struct xrt_quat* right,
                 struct xrt_quat* result)
{
	assert(left != NULL);
	assert(right != NULL);
	assert(result != NULL);

	auto l = copy(left);
	auto r = copy(right);

	auto q = l * r;

	map_quat(*result) = q;
}

void
math_quat_rotate_vec3(const struct xrt_quat* left,
                      const struct xrt_vec3* right,
                      struct xrt_vec3* result)
{
	assert(left != NULL);
	assert(right != NULL);
	assert(result != NULL);

	auto l = copy(left);
	auto r = copy(right);

	auto v = l * r;

	map_vec3(*result) = v;
}

/*
 *
 * Exported pose functions.
 *
 */

bool
math_pose_validate(const struct xrt_pose* pose)
{
	assert(pose != NULL);

	const float FLOAT_EPSILON = Eigen::NumTraits<float>::epsilon();
	auto norm = orientation(*pose).squaredNorm();
	if (norm > 1.0f + FLOAT_EPSILON || norm < 1.0f - FLOAT_EPSILON) {
		return false;
	}

	// Technically not yet a required check, but easier to stop problems
	// now than once denormalized numbers pollute the rest of our state.
	// see https://gitlab.khronos.org/openxr/openxr/issues/922
	if (!orientation(*pose).coeffs().allFinite()) {
		return false;
	}
	if (!position(*pose).allFinite()) {
		return false;
	}
	return true;
}

void
math_pose_invert(const struct xrt_pose* pose, struct xrt_pose* outPose)
{
	assert(pose != NULL);
	assert(outPose != NULL);

	// store results to temporary locals so we can do this "in-place"
	// (pose == outPose) if desired.
	Eigen::Vector3f newPosition = -position(*pose);
	// Conjugate legal here since pose must be normalized/unit length.
	Eigen::Quaternionf newOrientation = orientation(*pose).conjugate();

	position(*outPose) = newPosition;
	orientation(*outPose) = newOrientation;
}

/*!
 * Return the result of transforming a point by a pose/transform.
 */
static inline Eigen::Vector3f
transform_point(const xrt_pose& transform, const xrt_vec3& point)
{
	return orientation(transform) * map_vec3(point) + position(transform);
}

/*!
 * Return the result of transforming a pose by a pose/transform.
 */
static inline xrt_pose
transform_pose(const xrt_pose& transform, const xrt_pose& pose)
{
	xrt_pose ret;
	position(ret) = transform_point(transform, pose.position);
	orientation(ret) = orientation(transform) * orientation(pose);
	return ret;
}

void
math_pose_transform(const struct xrt_pose* transform,
                    const struct xrt_pose* pose,
                    struct xrt_pose* outPose)
{
	assert(pose != NULL);
	assert(transform != NULL);
	assert(outPose != NULL);

	xrt_pose newPose = transform_pose(*transform, *pose);
	memcpy(outPose, &newPose, sizeof(xrt_pose));
}

void
math_pose_openxr_locate(const struct xrt_pose* space_pose,
                        const struct xrt_pose* relative_pose,
                        const struct xrt_pose* base_space_pose,
                        struct xrt_pose* result)
{
	assert(space_pose != NULL);
	assert(relative_pose != NULL);
	assert(base_space_pose != NULL);
	assert(result != NULL);

	// Compilers are slighty better optimizing
	// if we copy the arguments in one go.
	const auto bsp = *base_space_pose;
	const auto rel = *relative_pose;
	const auto spc = *space_pose;
	struct xrt_pose pose;

	// Apply the invert of the base space to identity.
	math_pose_invert(&bsp, &pose);

	// Apply the pure pose from the space relation.
	math_pose_transform(&pose, &rel, &pose);

	// Apply the space pose.
	math_pose_transform(&pose, &spc, &pose);

	*result = pose;
}

/*!
 * Return the result of rotating a derivative vector by a matrix.
 *
 * This is a differential transform.
 */
static inline Eigen::Vector3f
rotate_deriv(Eigen::Matrix3f const& rotation,
             const xrt_vec3& derivativeVector,
             Eigen::Matrix3f const& rotationInverse)
{
	return ((rotation * map_vec3(derivativeVector)).transpose() *
	        rotationInverse)
	    .transpose();
}

#ifndef XRT_DOXYGEN

#define MAKE_REL_FLAG_CHECK(NAME, MASK)                                        \
	static inline bool NAME(xrt_space_relation_flags flags)                \
	{                                                                      \
		return ((flags & (MASK)) != 0);                                \
	}

MAKE_REL_FLAG_CHECK(has_some_pose_component,
                    XRT_SPACE_RELATION_POSITION_VALID_BIT |
                        XRT_SPACE_RELATION_ORIENTATION_VALID_BIT)
MAKE_REL_FLAG_CHECK(has_position, XRT_SPACE_RELATION_POSITION_VALID_BIT)
MAKE_REL_FLAG_CHECK(has_orientation, XRT_SPACE_RELATION_ORIENTATION_VALID_BIT)
MAKE_REL_FLAG_CHECK(has_lin_vel, XRT_SPACE_RELATION_LINEAR_VELOCITY_VALID_BIT)
MAKE_REL_FLAG_CHECK(has_ang_vel, XRT_SPACE_RELATION_ANGULAR_VELOCITY_VALID_BIT)
MAKE_REL_FLAG_CHECK(has_lin_acc,
                    XRT_SPACE_RELATION_LINEAR_ACCELERATION_VALID_BIT)
MAKE_REL_FLAG_CHECK(has_ang_acc,
                    XRT_SPACE_RELATION_ANGULAR_ACCELERATION_VALID_BIT)
MAKE_REL_FLAG_CHECK(has_some_derivative,
                    XRT_SPACE_RELATION_LINEAR_VELOCITY_VALID_BIT |
                        XRT_SPACE_RELATION_ANGULAR_VELOCITY_VALID_BIT |
                        XRT_SPACE_RELATION_LINEAR_ACCELERATION_VALID_BIT |
                        XRT_SPACE_RELATION_ANGULAR_ACCELERATION_VALID_BIT)

#undef MAKE_REL_FLAG_CHECK

#endif // !XRT_DOXYGEN

/*!
 * Apply a transform to a space relation.
 */
static inline void
transform_accumulate_pose(const xrt_pose& transform,
                          xrt_space_relation& relation,
                          bool do_translation = true,
                          bool do_rotation = true)
{
	assert(do_translation || do_rotation);

	// Save the quat in case we are self-transforming.
	Eigen::Quaternionf quat = orientation(transform);

	auto flags = relation.relation_flags;
	// so code looks similar
	auto in_out_relation = &relation;

	// transform (rotate and translate) the pose, if applicable.
	if (has_some_pose_component(flags)) {
		// Zero out transform parts we don't want to use,
		// because math_pose_transform doesn't take flags.
		xrt_pose transform_copy = transform;
		if (!do_translation) {
			position(transform_copy) = Eigen::Vector3f::Zero();
		}
		if (!do_rotation) {
			orientation(transform_copy) =
			    Eigen::Quaternionf::Identity();
		}

		math_pose_transform(&in_out_relation->pose, &transform,
		                    &in_out_relation->pose);
	}

	if (do_rotation && has_some_derivative(flags)) {

		// prepare matrices required for rotating derivatives from the
		// saved quat.
		Eigen::Matrix3f rot = quat.toRotationMatrix();
		Eigen::Matrix3f rotInverse = rot.inverse();

		// Rotate derivatives, if applicable.
		if (has_lin_vel(flags)) {
			map_vec3(in_out_relation->linear_velocity) =
			    rotate_deriv(rot, in_out_relation->linear_velocity,
			                 rotInverse);
		}

		if (has_ang_vel(flags)) {
			map_vec3(in_out_relation->angular_velocity) =
			    rotate_deriv(rot, in_out_relation->angular_velocity,
			                 rotInverse);
		}

		if (has_lin_acc(flags)) {
			map_vec3(in_out_relation->linear_acceleration) =
			    rotate_deriv(rot,
			                 in_out_relation->linear_acceleration,
			                 rotInverse);
		}

		if (has_ang_acc(flags)) {
			map_vec3(in_out_relation->angular_acceleration) =
			    rotate_deriv(rot,
			                 in_out_relation->angular_acceleration,
			                 rotInverse);
		}
	}
}

static const struct xrt_space_relation BLANK_RELATION = {
    XRT_SPACE_RELATION_BITMASK_ALL,
    {{0.0f, 0.0f, 0.0f, 1.0f}, {0.0f, 0.0f, 0.0f}},
    {0, 0, 0},
    {0, 0, 0},
    {0, 0, 0},
    {0, 0, 0},
};

void
math_relation_reset(struct xrt_space_relation* out)
{
	*out = BLANK_RELATION;
}

void
math_relation_accumulate_transform(const struct xrt_pose* transform,
                                   struct xrt_space_relation* in_out_relation)
{
	assert(transform != nullptr);
	assert(in_out_relation != nullptr);

	// No modifying the validity flags here.
	transform_accumulate_pose(*transform, *in_out_relation);
}

void
math_relation_accumulate_relation(
    const struct xrt_space_relation* additional_relation,
    struct xrt_space_relation* in_out_relation)
{
	assert(additional_relation != NULL);
	assert(in_out_relation != NULL);

	// Update the flags.
	xrt_space_relation_flags flags = (enum xrt_space_relation_flags)(
	    in_out_relation->relation_flags &
	    additional_relation->relation_flags);
	in_out_relation->relation_flags = flags;

	if (has_some_pose_component(flags)) {
		// First, just do the pose part (including rotating
		// derivatives, if applicable).
		transform_accumulate_pose(additional_relation->pose,
		                          *in_out_relation, has_position(flags),
		                          has_orientation(flags));
	}

	// Then, accumulate the derivatives, if required.
	if (has_lin_vel(flags)) {
		map_vec3(in_out_relation->linear_velocity) +=
		    map_vec3(additional_relation->linear_velocity);
	}

	if (has_ang_vel(flags)) {
		map_vec3(in_out_relation->angular_velocity) +=
		    map_vec3(additional_relation->angular_velocity);
	}

	if (has_lin_acc(flags)) {
		map_vec3(in_out_relation->linear_acceleration) +=
		    map_vec3(additional_relation->linear_acceleration);
	}

	if (has_ang_acc(flags)) {
		map_vec3(in_out_relation->angular_acceleration) +=
		    map_vec3(additional_relation->angular_acceleration);
	}
}

void
math_relation_openxr_locate(const struct xrt_pose* space_pose,
                            const struct xrt_space_relation* relative_relation,
                            const struct xrt_pose* base_space_pose,
                            struct xrt_space_relation* result)
{
	assert(space_pose != NULL);
	assert(relative_relation != NULL);
	assert(base_space_pose != NULL);
	assert(result != NULL);

	// Compilers are slighty better optimizing
	// if we copy the arguments in one go.
	const auto bsp = *base_space_pose;
	const auto spc = *space_pose;
	struct xrt_space_relation accumulating_relation = BLANK_RELATION;

	// Apply the invert of the base space to identity.
	math_pose_invert(&bsp, &accumulating_relation.pose);

	// Apply the pure relation between spaces.
	math_relation_accumulate_relation(relative_relation,
	                                  &accumulating_relation);

	// Apply the space pose.
	math_relation_accumulate_transform(&spc, &accumulating_relation);

	*result = accumulating_relation;
}


#include <vector>
#include <string>
#include <map>

#include "rate_model.h"

#include <armadillo>
using namespace arma;

inline int signof(double d) {return d >= 0 ? 1 : -1;}
inline double roundto(double in) {return floor(in*(1000)+0.5)/(1000);}

RateModel::RateModel(int _nstates):Q(_nstates,_nstates),labels(),Q_mask(),
    lasteigval(_nstates,_nstates),lasteigvec(_nstates,_nstates),
    eigval(_nstates,_nstates),eigvec(_nstates,_nstates),
    lasteigval_simple(_nstates,_nstates),lasteigvec_simple(_nstates,_nstates),
    eigval_simple(_nstates,_nstates),eigvec_simple(_nstates,_nstates),nstates(_nstates) {

    setup_Q();
    sameQ = false;
    lasteigval.fill(0);
    lasteigvec.fill(0);
    eigval.fill(0);
    eigvec.fill(0);
    lasteigval_simple.fill(0);
    lasteigvec_simple.fill(0);
    eigval_simple.fill(0);
    eigvec_simple.fill(0);
}

void RateModel::set_Q_cell(int from, int to, double num) {
    Q(from,to) = num;
    sameQ = false;
}

void RateModel::set_Q_diag() {
    for (unsigned int i=0; i < Q.n_rows; i++) {
        double su = 0;
        for (unsigned int j=0; j < Q.n_cols; j++) {
            if (i != j) {
                su += Q(i, j);
            }
        }
        Q(i, i) = 0-su;
    }
    sameQ = false;
}

void RateModel::setup_Q() {
    Q.fill(0);
    for (unsigned int i=0; i < Q.n_rows; i++) {
        for (unsigned int j=0; j < Q.n_cols; j++) {
            if (i != j) {
                Q(i, j) = 1./nstates;
            } else {
                Q(i, j) = -(1./nstates * (nstates-1));
            }
        }
    }
    sameQ = false;
}

void RateModel::setup_Q(vector<vector<double> > & inQ) {
    for (unsigned int i=0; i < Q.n_rows; i++) {
        for (unsigned int j=0; j < Q.n_cols; j++) {
            Q(i, j) = inQ[i][j];
        }
    }
    for (unsigned int i=0; i < Q.n_rows; i++) {
        colvec a = (sum(Q,1));
        Q(i, i) = -(a(i)-Q(i, i));
    }
    sameQ = false;
}

void RateModel::setup_Q(mat & inQ) {
    for (unsigned int i=0; i < Q.n_rows; i++) {
        for (unsigned int j=0; j < Q.n_cols; j++) {
            Q(i, j) = inQ(i, j);
        }
    }
    set_Q_diag();
    sameQ = false;
}

void RateModel::set_n_qs(int number) {
    for (int i=0; i < number; i++) {
        mat tm(nstates,nstates);
        Qs.push_back(tm);
    }
}

void RateModel::set_Q_which(mat & inQ,int which) {
    for (unsigned int i=0; i < Qs[which].n_rows; i++) {
        for (unsigned int j=0; j < Qs[which].n_cols; j++) {
            Qs[which](i, j) = inQ(i, j);
        }
    }
    for (unsigned int i=0; i < Qs[which].n_rows; i++) {
        double su = 0;
        for (unsigned int j=0; j < Qs[which].n_cols; j++) {
            if (i != j) {
                su += Qs[which](i, j);
            }
        }
        Qs[which](i, i) = 0-su;
    }
    sameQ = false;
}

void RateModel::set_Q(mat & inQ) {
    for (unsigned int i=0; i < Q.n_rows; i++) {
        for (unsigned int j=0; j < Q.n_cols; j++) {
            Q(i, j) = inQ(i, j);
        }
    }
    sameQ = false;
}

mat & RateModel::get_Q() {
    return Q;
}

cx_mat RateModel::setup_P(double bl,bool store_p_matrices) {
    //sameQ = false;
    eigvec.fill(0);
    eigval.fill(0);
    bool isImag = get_eigenvec_eigenval_from_Q(&eigval, &eigvec); // isImag is not used?
    //cout << eigval << endl;
    //cout << eigvec << endl;
    for (int i=0; i < nstates; i++) {
        eigval(i, i) = exp(eigval(i, i) * bl);
    }
    cx_mat C_inv = inv(eigvec);
    cx_mat P = eigvec * eigval * C_inv;
    neg_p = false;
    for (unsigned int i=0; i < P.n_rows; i++) {
        for (unsigned int j=0; j < P.n_cols; j++) {
            if (real(P(i, j))<0)
            neg_p = true;
        }
    }
    if (store_p_matrices == true) {
        stored_p_matrices[bl] = P;
    }
    return P;
}

void RateModel::setup_P_simple(mat & p,double bl,bool store_p_matrices) {
//    sameQ = false;
    eigvec_simple.fill(0);
    eigval_simple.fill(0);
    get_eigenvec_eigenval_from_Q_simple(&eigval_simple, &eigvec_simple);
    //cout << eigval << endl;
    //cout << eigvec << endl;
    for (int i=0; i < nstates; i++) {
        eigval_simple(i, i) = exp(eigval_simple(i, i) * bl);
    }
    mat C_inv = inv(eigvec_simple);
    p = eigvec_simple * eigval_simple * C_inv;
/*    if (store_p_matrices == true) {
    stored_p_matrices[bl] = P;
    }*/
    //return P;
}

void RateModel::get_eigenvec_eigenval_from_Q_simple(mat * eigval, mat * eigvec) {
    if (sameQ == true) {
        for (unsigned int i=0; i < Q.n_rows; i++) {
            for (unsigned int j=0; j < Q.n_cols; j++) {
                (*eigval)(i, j) = lasteigval_simple(i, j);
                (*eigvec)(i, j) = lasteigvec_simple(i, j);
            }
        }
        return;
    }
    mat tQ(nstates,nstates); tQ.fill(0);
    for (unsigned int i=0; i < Q.n_rows; i++) {
        for (unsigned int j=0; j < Q.n_cols; j++) {
            tQ(i, j) = Q(i, j);
        }
    }
    colvec eigva;
    mat eigve;
    eig_sym(eigva,eigve,tQ);
    //bool isImag = false; // not used
    for (unsigned int i=0; i < Q.n_rows; i++) {
        for (unsigned int j=0; j < Q.n_cols; j++) {
            if (i == j) {
                (*eigval)(i, j) = eigva(i);
                lasteigval_simple(i, j) = eigva(i);
            } else {
                (*eigval)(i, j) = 0;
                lasteigval_simple(i, j) = 0;
            }
            (*eigvec)(i, j) = eigve(i, j);
            lasteigvec_simple(i, j) = eigve(i, j);
        }
    }
    return;
}

/*
 * this should be used to caluculate the eigenvalues and eigenvectors
 * as U * Q * U-1 -- eigen decomposition
 *
 * this should use the armadillo library
 */
bool RateModel::get_eigenvec_eigenval_from_Q(cx_mat * eigval, cx_mat * eigvec) {
    if (sameQ == true) {
        for (unsigned int i=0; i < Q.n_rows; i++) {
            for (unsigned int j=0; j < Q.n_cols; j++) {
                (*eigval)(i, j) = lasteigval(i, j);
                (*eigvec)(i, j) = lasteigvec(i, j);
            }
        }
        return lastImag;
    }
    mat tQ(nstates,nstates); tQ.fill(0);
    for (unsigned int i=0; i < Q.n_rows; i++) {
        for (unsigned int j=0; j < Q.n_cols; j++) {
            tQ(i, j) = Q(i, j);
        }
    }
    cx_colvec eigva;
    cx_mat eigve;
    eig_gen(eigva,eigve,tQ);
    bool isImag = false;
    for (unsigned int i=0; i < Q.n_rows; i++) {
        for (unsigned int j=0; j < Q.n_cols; j++) {
            if (i == j) {
                (*eigval)(i, j) = eigva(i);
                lasteigval(i, j) = eigva(i);
            } else {
                (*eigval)(i, j) = 0;
                lasteigval(i, j) = 0;
            }
            (*eigvec)(i, j) = eigve(i, j);
            lasteigvec(i, j) = eigve(i, j);
            if (imag((*eigvec)(i, j)) > 0 || imag((*eigval)(i, j))) {
                isImag = true;
            }
        }
    }
    //lasteigval = eigval;
    //lasteigvec = eigvec;
    lastImag = isImag;
    return isImag;
}

//taking out fortran
//
/*
extern"C" {
    void wrapalldmexpv_(int * n,int* m,double * t,double* v,double * w,double* tol,double* anorm,double* wsp,int * lwsp,int* iwsp,int *liwsp, int * itrace,int *iflag,int *ia, int *ja, double *a, int *nz, double * res);
    void wrapsingledmexpv_(int * n,int* m,double * t,double* v,double * w,double* tol,double* anorm,double* wsp,int * lwsp,int* iwsp,int *liwsp, int * itrace,int *iflag,int *ia, int *ja, double *a, int *nz, double * res);
    void wrapdgpadm_(int * ideg,int * m,double * t,double * H,int * ldh,double * wsp,int * lwsp,int * ipiv,int * iexph,int *ns,int *iflag );
}

void RateModel::setup_fortran_P_whichQ(int which, mat & P, double t) {
    Q = Qs[which];
    setup_fortran_P(P,t,false);
}
*/
//taking out fortran
/*
 * runs the basic padm fortran expokit full matrix exp
 *
void RateModel::setup_fortran_P(mat & P, double t, bool store_p_matrices) {
    //
    //  return P, the matrix of dist-to-dist transition probabilities,
    //  from the model's rate matrix (Q) over a time duration (t)
    //
    int ideg = 6;
    int m = Q.n_rows; // square so you only need the rows
    int ldh = m;
    double tol = 1;
    int iflag = 0;
    int lwsp = 4*m*m+6+1;
    double * wsp = new double[lwsp];
    int * ipiv = new int[m];
    int iexph = 0;
    int ns = 0;
    double * H = new double [m*m];
    convert_matrix_to_single_row_for_fortran(Q, t, H);
    wrapdgpadm_(&ideg, &m, &tol, H, &ldh, wsp, &lwsp, ipiv, &iexph, &ns, &iflag);

    for (int i=0; i < m; i++) {
        for (int j=0; j < m; j++) {
            P(i, j) = wsp[iexph + (j-1) * m + (i-1) + m];
        }
    }
    delete [] wsp;
    delete [] ipiv;
    delete [] H;
    for (int i=0; i < nstates; i++) {
        double sum = 0.0;
        for (int j=0; j < nstates; j++) {
            sum += P(i, j);
        }
        for (int j=0; j < nstates; j++) {
            P(i, j) = (P(i, j)/sum);
        }
    }

}
*/
void RateModel::set_sameQ(bool s) {
    sameQ = s;
}

void update_simple_goldman_yang_q(mat * inm, double K, double w, mat & bigpibf,mat &bigpiK, mat & bigpiw) {
    double s = 0;
    for (int i=0; i < 61; i++) {
        for (int j=0; j < 61; j++) {
            if (bigpibf(i, j) == 0) {
                (*inm)(i, j) = 0;
            } else {
                (*inm)(i, j) = 1/61.;
                if (bigpiK(i, j) != 0) {
                    (*inm)(i, j) *= K;
                }
                if (bigpiw(i, j) != 0) {
                    (*inm)(i, j) *= w;
                }
            }
            if (i == j) {
                (*inm)(i, j) = 0;
            }
            s += (*inm)(i, j);
        }
    }
    (*inm) = (*inm)/s;
    for (unsigned int i=0; i < (*inm).n_rows; i++) {
        double su = 0;
        for (unsigned int j=0; j < (*inm).n_cols; j++) {
            if (i != j) {
            su += (*inm)(i, j);
            }
        }
        (*inm)(i, i) = 0-su;
    }
    (*inm) = (*inm)/(1/61.);
    //(*inm) = trans((*inm));
}

bool test_transition(char a, char b) {
    bool ret = false;
    if ((a == 'A' &&  b == 'G') || (a == 'C' &&  b == 'T') || (a == 'G' &&  b == 'A') || (a == 'T' &&  b == 'C')) {
        ret = true;
    }
    return ret;
}

void generate_bigpibf_K_w(mat * bf, mat * K, mat * w,map<string, string> & codon_dict,
    map<string, vector<int> > & codon_index, vector<string> & codon_list) {
    for (int i=0; i < 61; i++) {
        for (int j=0; j < 61; j++) {
            int diff = 0;
            bool transit = false;
            bool nonsyn = false;
            for (int m=0; m < 3; m++) {
                if (codon_list[i][m] != codon_list[j][m]) {
                    diff += 1;
                }
                transit = test_transition(codon_list[i][m],codon_list[j][m]);
            }
            if (diff > 1) {
                (*bf)(i, j) = 0;
            } else {
                if (codon_dict[codon_list[i]] != codon_dict[codon_list[j]]) {
                    nonsyn = true;
                }
                (*bf)(i, j) = 1;
                if (transit) {
                    (*K)(i, j) = 1;
                }
                if (nonsyn) {
                    (*w)(i, j) = 1;
                }
            }
        }
    }
}

//take out fortran
/*
void convert_matrix_to_single_row_for_fortran(mat & inmatrix, double t, double * H) {
    int count = 0;
    for (unsigned int i=0; i < inmatrix.n_cols; i++) {
        for (unsigned int j=0; j < inmatrix.n_cols; j++) {
            H[i+(j*inmatrix.n_cols)] = inmatrix(i, j)*t;
            count += 1;
        }
    }
}
*/


#ifndef TEST_UNIT_TORSTEN_COUPLED_TWOCPT_MODEL_TEST_FIXTURE
#define TEST_UNIT_TORSTEN_COUPLED_TWOCPT_MODEL_TEST_FIXTURE

#include <stan/math/rev/core/recover_memory.hpp>
#include <gtest/gtest.h>
#include <boost/numeric/odeint.hpp>
#include <stan/math/torsten/test/unit/test_fixture_model.hpp>
#include <stan/math/torsten/pmx_onecpt_model.hpp>
#include <stan/math/torsten/pmx_solve_onecpt_rk45.hpp>
#include <stan/math/torsten/pmx_solve_onecpt_bdf.hpp>
#include <stan/math/torsten/pmx_ode_model.hpp>
#include <iostream>
#include <sstream>
#include <vector>
#include <string>

struct CoupledTwoCptODE {
  /*
   * Coupled model functor
   */
  template <typename T0, typename T1, typename T2, typename T3>
  Eigen::Matrix<typename stan::return_type_t<T0, T1, T2, T3>, -1, 1>
  operator()(const T0& t,
             const Eigen::Matrix<T1, -1, 1>& y,
             const Eigen::Matrix<T2, -1, 1>& y_pk,
             const std::vector<T3>& theta,
             const std::vector<double>& x_r,
             const std::vector<int>& x_i,
             std::ostream* pstream_) const {
    typedef typename boost::math::tools::promote_args<T0, T1, T2, T3>::type
      scalar;

    scalar VC = theta[2],
      Mtt = theta[5],
      circ0 = theta[6],
      alpha = theta[7],
      gamma = theta[8],
      ktr = 4 / Mtt,
      prol = y[0] + circ0,
      transit = y[1] + circ0,
      circ = y[2] + circ0,
      conc = y_pk[1] / VC,
      Edrug = alpha * conc;

    Eigen::Matrix<scalar, -1, 1> dxdt(3);
    dxdt << ktr * prol * ((1 - Edrug) * pow(circ0 / circ, gamma) - 1),
      ktr * (prol - transit),
      ktr * (transit - circ);

    return dxdt;
  }

  /*
   * full ODE model functor
   */
  template <typename T0, typename T1, typename T2>
  Eigen::Matrix<typename stan::return_type_t<T0, T1, T2>, -1, 1>
  operator()(const T0& t,
             const Eigen::Matrix<T1, -1, 1> & x,
             std::ostream* pstream_,
             const std::vector<T2>& theta,
             const std::vector<double>& x_r,
             const std::vector<int>& x_i) const {
    typedef typename boost::math::tools::promote_args<T0, T1, T2>::type
    scalar;

    T2
      CL = theta[0],
      Q = theta[1],
      VC = theta[2],
      VP = theta[3],
      ka = theta[4],
      k10 = CL / VC,
      k12 = Q / VC,
      k21 = Q / VP,
      Mtt = theta[5],
      circ0 = theta[6],
      alpha = theta[7],
      gamma = theta[8],
      ktr = 4 / Mtt;

    scalar
      prol = x[3] + circ0,
      transit = x[4] + circ0,
      circ = x[5] + circ0,
      Edrug = alpha * x[1] / VC;

    Eigen::Matrix<scalar, -1, 1> dxdt(6);
    dxdt[0] = -ka * x[0];
    dxdt[1] = ka * x[0] - (k10 + k12) * x[1] + k21 * x[2];
    dxdt[2] = k12 * x[1] - k21 * x[2];
    dxdt[3] = ktr * prol * ((1 - Edrug) * pow(circ0 / circ, gamma) - 1);
    dxdt[4] = ktr * (prol - transit);
    dxdt[5] = ktr * (transit - circ);

    return dxdt;
  }
};

template<typename T>
struct test_coupled_twocpt : public TorstenPMXTest<test_coupled_twocpt<T> > {
  const size_t nOde;
  test_coupled_twocpt() : nOde(6) {
    this -> ncmt = nOde;           // # of PD compartments
    this -> reset_events(10);
    this -> theta.resize(1);
    this -> theta[0] = {10, 15, 35, 105, 2.0, 125, 5, 3e-4, 0.17};
    this -> biovar.resize(1);
    this -> biovar[0] = {1, 1, 1, 1, 1, 1};
    this -> tlag.resize(1);
    this -> tlag[0] = {0, 0, 0, 0, 0, 0};

    for(int i = 0; i < this -> nt; i++) {
      this -> time[i] = i * 0.25;
    }
    this -> time[this -> nt - 1] = 4.0;

    this -> amt[0] = 10000.0;
    this -> cmt[0] = 1;
    this -> evid[0] = 1;
  }
};

#endif


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

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

// This file was modified by Oracle on 2014, 2016, 2017.
// Modifications copyright (c) 2014-2017 Oracle and/or its affiliates.

// Contributed and/or modified by Adam Wulkiewicz, on behalf of Oracle

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

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

#ifndef BOOST_GEOMETRY_CORE_SRS_HPP
#define BOOST_GEOMETRY_CORE_SRS_HPP

#include <cstddef>

#include <boost/static_assert.hpp>

#include <boost/geometry/core/radius.hpp>
#include <boost/geometry/core/tag.hpp>
#include <boost/geometry/core/tags.hpp>

namespace boost { namespace geometry
{

namespace srs
{

/*!
    \brief Defines spheroid radius values for use in geographical CS calculations
    \note See http://en.wikipedia.org/wiki/Figure_of_the_Earth
          and http://en.wikipedia.org/wiki/World_Geodetic_System#A_new_World_Geodetic_System:_WGS84
*/
template <typename RadiusType>
class spheroid
{
public:
    spheroid(RadiusType const& a, RadiusType const& b)
        : m_a(a)
        , m_b(b)
    {}

    spheroid()
        : m_a(RadiusType(6378137.0))
        , m_b(RadiusType(6356752.3142451793))
    {}

    template <std::size_t I>
    RadiusType get_radius() const
    {
        BOOST_STATIC_ASSERT(I < 3);

        return I < 2 ? m_a : m_b;
    }

    template <std::size_t I>
    void set_radius(RadiusType const& radius)
    {
        BOOST_STATIC_ASSERT(I < 3);

        (I < 2 ? m_a : m_b) = radius;
    }

private:
    RadiusType m_a, m_b; // equatorial radius, polar radius
};

} // namespace srs

// Traits specializations for spheroid
#ifndef DOXYGEN_NO_TRAITS_SPECIALIZATIONS
namespace traits
{

template <typename RadiusType>
struct tag< srs::spheroid<RadiusType> >
{
    typedef srs_spheroid_tag type;
};

template <typename RadiusType>
struct radius_type< srs::spheroid<RadiusType> >
{
    typedef RadiusType type;
};

template <typename RadiusType, std::size_t Dimension>
struct radius_access<srs::spheroid<RadiusType>, Dimension>
{
    typedef srs::spheroid<RadiusType> spheroid_type;

    static inline RadiusType get(spheroid_type const& s)
    {
        return s.template get_radius<Dimension>();
    }

    static inline void set(spheroid_type& s, RadiusType const& value)
    {
        s.template set_radius<Dimension>(value);
    }
};

} // namespace traits
#endif // DOXYGEN_NO_TRAITS_SPECIALIZATIONS

namespace srs
{

/*!
    \brief Defines sphere radius value for use in spherical CS calculations
*/
template <typename RadiusType>
class sphere
{
public:
    explicit sphere(RadiusType const& r)
        : m_r(r)
    {}

    sphere()
        : m_r(RadiusType((2.0 * 6378137.0 + 6356752.3142451793) / 3.0))
    {}

    template <std::size_t I>
    RadiusType get_radius() const
    {
        BOOST_STATIC_ASSERT(I < 3);

        return m_r;
    }

    template <std::size_t I>
    void set_radius(RadiusType const& radius)
    {
        BOOST_STATIC_ASSERT(I < 3);

        m_r = radius;
    }

private:
    RadiusType m_r; // radius
};

} // namespace srs

// Traits specializations for sphere
#ifndef DOXYGEN_NO_TRAITS_SPECIALIZATIONS
namespace traits
{

template <typename RadiusType>
struct tag< srs::sphere<RadiusType> >
{
    typedef srs_sphere_tag type;
};

template <typename RadiusType>
struct radius_type< srs::sphere<RadiusType> >
{
    typedef RadiusType type;
};

template <typename RadiusType, std::size_t Dimension>
struct radius_access<srs::sphere<RadiusType>, Dimension>
{
    typedef srs::sphere<RadiusType> sphere_type;

    static inline RadiusType get(sphere_type const& s)
    {
        return s.template get_radius<Dimension>();
    }

    static inline void set(sphere_type& s, RadiusType const& value)
    {
        s.template set_radius<Dimension>(value);
    }
};

} // namespace traits
#endif // DOXYGEN_NO_TRAITS_SPECIALIZATIONS

}} // namespace boost::geometry

#endif // BOOST_GEOMETRY_CORE_SRS_HPP


/***********************************************************************
Copyright (c) 2017, Shushman Choudhury
All rights reserved.
Authors: Shushman Choudhury <shushmanchoudhury@gmail.com>

Redistribution and use in source and binary forms, with or without
modification, are permitted provided that the following conditions are
met:
  Redistributions of source code must retain the above copyright
  notice, this list of conditions and the following disclaimer.
  Redistributions in binary form must reproduce the above copyright
  notice, this list of conditions and the following disclaimer in the
  documentation and/or other materials provided with the distribution.
THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS
"AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR
A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT
HOLDER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL,
SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT
LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE,
DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY
THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
(INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
*************************************************************************/

#ifndef BATCHING_POMP_CSPACEBELIEF_KNNMODEL_HPP_
#define BATCHING_POMP_CSPACEBELIEF_KNNMODEL_HPP_

#include <functional>
#include <exception>
#include <memory>
#include <cmath>
#include <ompl/datastructures/NearestNeighbors.h>
#include <ompl/datastructures/NearestNeighborsLinear.h>
#include <ompl/datastructures/NearestNeighborsGNAT.h>
#include <Eigen/Dense>
#include "Model.hpp"
#include "BeliefPoint.hpp"

namespace batching_pomp {
namespace cspacebelief {

//! Derived class of Model that uses kNN to represent the C-space belief.

/// Implements a C-Space Belief Model using k-Nearest Neighbour lookup
/// and weighted averaging over the results of neighbouring belief points.
/// The weights are inversely proportional to the distance between
/// the queried point and the neighbouring belief points.
class KNNModel : public virtual Model<BeliefPoint>
{
public:

  /// \param[in] _KNN The value of k for KNN lookup
  /// \param[in] _distanceFunction The ompl::NearestNeighbours::Distance Function to set
  KNNModel(size_t _KNN,
           double _prior, double _priorWeight,
           const std::function<double(const BeliefPoint&, const BeliefPoint&)>& _distanceFunction)
  : mNumPoints{0}
  , mKNN{_KNN}
  , mPrior{_prior}
  , mPriorWeight{_priorWeight}
  , mDistanceFunction{_distanceFunction}
  {
    mBeliefPointNN.reset(new ompl::NearestNeighborsGNAT<BeliefPoint>());
    mBeliefPointNN->setDistanceFunction(mDistanceFunction);
  }

  //////////////////////////////////////////////////
  // Setters
  void setPrior(double _prior)
  {
    mPrior = _prior;
  }

  void setPriorWeight(double _priorWeight)
  {
    mPriorWeight = _priorWeight;
  }

  void setKNN(size_t _knn)
  {
    mKNN = _knn;
  }

  //////////////////////////////////////////////////
  // Overriden methods
  void addPoint(const BeliefPoint& data) override
  {
    mBeliefPointNN->add(data);
    mNumPoints++;
  }

  void removePoint(const BeliefPoint& data) override
  {
    mBeliefPointNN->remove(data);
    mNumPoints--;
  }

  double estimate(const BeliefPoint& query) const override
  {
    // This is important for cases where edges need to be initialized before any checks.
    if(mNumPoints == 0) {
      return mPrior;
    }

    size_t knn = std::min(mKNN , mNumPoints);

    std::vector<BeliefPoint> neighboursVect;
    mBeliefPointNN->nearestK(query,knn,neighboursVect);

    if(knn != neighboursVect.size()) {
      throw std::runtime_error("Model could not return the expected number of neighbours");
    }

    Eigen::VectorXd weights(knn);
    Eigen::VectorXd values(knn);

    for(size_t i = 0; i < knn; i++) {
      double distance = std::max(0.0001,mDistanceFunction(query, neighboursVect[i]));
      weights[i] = 1.0/distance;
      values[i] = neighboursVect[i].getValue(); // Second element of pair is value
    }
    double result{mPrior};

    result = weights.dot(values) / weights.sum();

    // Do (result + pw*p) / (1 + pw) for smoothing by prior
    result = (result + mPriorWeight*mPrior) / (1 + mPriorWeight);

    return result;
  }

private:

  // The current number of belief points in the model
  size_t mNumPoints;

  // The value of k for the kNN lookup
  size_t mKNN;

  // The prior probability in (0,1) to use for smoothing
  double mPrior;

  // The weight in (0,1) to assign to the prior for smoothing
  double mPriorWeight;

  // The kind of nearest neighbour structure for the belief model, GNAT in this case.
  std::unique_ptr<ompl::NearestNeighbors<BeliefPoint>> mBeliefPointNN;

  // The distance metric for the nearest neighbour structure
  std::function<double(const BeliefPoint&, const BeliefPoint&)> mDistanceFunction;
};

} //namespace cspacebelief
} //namespace batching_pomp

#endif //BATCHING_POMP_CSPACEBELIEF_KNNMODEL_HPP_