#include <iostream>
#include <vector>
#include <cmath>

#include <boost/timer/timer.hpp>

#include <poplar/DeviceManager.hpp>
#include <poplar/IPUModel.hpp>

#include <poplar/Engine.hpp>
#include <poplar/Program.hpp>
#include <poplin/codelets.hpp>
#include <popops/ElementWise.hpp>
#include <popops/codelets.hpp>
#include <popops/ExprOp.hpp>
#include <popops/Pad.hpp>
#include <poputil/TileMapping.hpp>

#include "KalmanFilter.h"

using namespace poplar;
using namespace poplar::program;
using namespace popops;

float d = 1.0;
float sigma = 10E-2;
int N = 5;
float z = 0.1;
float x0 = 0.01;
float theta0 = 10E-3;

int main(int argc, char const *argv[]) {

  Device dev = KalmanFilter::connectToIPU();

  // IPUModel ipuModel;
  // ipuModel.compileIPUCode = true;

  // Device dev = ipuModel.createDevice();

  Graph graph(dev.getTarget());

  popops::addCodelets(graph);
  poplin::addCodelets(graph);
  graph.addCodelets("matrixInverseVertex.cpp");
  graph.addCodelets("matrixProductVertex.cpp");
  graph.addCodelets("scaledAddVertex.cpp");
  graph.addCodelets("packHitsVertex.cpp");

  int n_inputs = 1;
  int batch_size = 1;

  std::vector<Tensor> inputs(n_inputs);
  std::vector<Tensor> inputs_batch(n_inputs);

  for (uint i = 0; i < inputs.size(); i++) {

    std::string iStr = std::to_string(i);

    inputs[i] = graph.addVariable(FLOAT, {5, 2}, "x_in" + iStr);
    inputs_batch[i] = graph.addVariable(FLOAT, {uint(batch_size), 5 * 2}, "x_in_batch" + iStr); // Check dims!
    graph.setTileMapping(inputs[i], i);
    graph.setTileMapping(inputs_batch[i], i);
  }

  Sequence preProg;

  std::vector<DataStream> inStreams(n_inputs);
  std::vector<Tensor> covs(n_inputs);
  std::vector<Tensor> qs(n_inputs);
  std::vector<Tensor> hs(n_inputs);
  std::vector<Tensor> gs(n_inputs);
  std::vector<Tensor> fs(n_inputs);
  std::vector<Tensor> d(n_inputs);
  std::vector<Tensor> dInit(n_inputs);
  std::vector<Tensor> dSkip(n_inputs);
  std::vector<Tensor> scatterInto(n_inputs);
  std::vector<Tensor> loop(n_inputs);
  std::vector<Tensor> zero(n_inputs);
  std::vector<Tensor> one(n_inputs);
  std::vector<Tensor> loop_batch(n_inputs);
  std::vector<Tensor> hitThisLoop(n_inputs);

  std::vector<Tensor> p_proj_all(n_inputs);
  std::vector<Tensor> C_proj_all(n_inputs);

  std::vector<Tensor> p_filt_all(n_inputs);
  std::vector<Tensor> C_filt_all(n_inputs);

  std::vector<Tensor> p_smooth(n_inputs);
  std::vector<Tensor> C_smooth(n_inputs);

  for (uint i = 0; i < covs.size(); i++) {

    std::string iStr = std::to_string(i);

    inStreams[i] = graph.addHostToDeviceFIFO("inStream" + iStr, FLOAT, 5 * 2 * batch_size);
    preProg.add(Copy(inStreams[i], inputs_batch[i]));

  }

  Sequence prog;

  std::vector<Tensor> covFlat(covs.size());

  for (uint i = 0; i < covs.size(); i++) {

    std::string iStr = std::to_string(i);

    loop[i] = graph.addVariable(INT, {1}, "loop");
    graph.setTileMapping(loop[i], i);

    scatterInto[i] = graph.addConstant<int>(INT, {1}, {0});
    graph.setTileMapping(scatterInto[i], i);

    zero[i] = graph.addConstant<int>(INT, {1}, {0});
    graph.setTileMapping(zero[i], i);

    one[i] = graph.addConstant<int>(INT, {1}, {1});
    graph.setTileMapping(one[i], i);

    prog.add(Copy(inputs_batch[i].slice(0, 1, 0).reshape({5, 2}), inputs[i]));

    covFlat[i] = graph.addConstant<float>(FLOAT, {16, 1}, {sigma * sigma, 0., 0., 0.,
                                                          0., M_PI, 0., 0.,
                                                          0., 0., sigma * sigma, 0.,
                                                          0., 0., 0., M_PI
                                                          });

    Tensor qFlat = graph.addConstant<float>(FLOAT, {16, 1}, {0.});

    Tensor hFlat = graph.addConstant<float>(FLOAT, {16, 1}, {1., 0., 0., 0.,
                                                             0., 0., 0., 0.,
                                                             0., 0., 1., 0.,
                                                             0., 0., 0., 0.});
    Tensor fFlat = graph.addConstant<float>(FLOAT, {16, 1}, {1., 1., 0., 0.,
                                                             0., 1., 0., 0.,
                                                             0., 0., 1., 1.,
                                                             0., 0., 0., 1.});
    Tensor gFlat = graph.addConstant<float>(FLOAT, {16, 1}, {float(1.0)/(sigma * sigma), 0., 0., 0.,
                                                             0, 0., 0., 0.,
                                                             0, 0., float(1.0)/(sigma * sigma), 0.,
                                                             0, 0., 0., 0.,
                                                             });

    d[i] = graph.addVariable(FLOAT, {1, 1}, "d" + iStr);
    dInit[i] = graph.addConstant<float>(FLOAT, {1, 1}, {1.});
    dSkip[i] = graph.addConstant<float>(FLOAT, {1, 1}, {2.});

    prog.add(Copy(dInit[i], d[i]));

    covs[i] = graph.addVariable(FLOAT, {4, 4}, "cov" + iStr);
    prog.add(Copy(covFlat[i].reshape({4, 4}), covs[i]));

    p_proj_all[i] = graph.addVariable(FLOAT, {5, 4, 1}, "p_proj_all" + iStr);
    C_proj_all[i] = graph.addVariable(FLOAT, {5, 4, 4}, "C_proj_all" + iStr);

    p_filt_all[i] = graph.addVariable(FLOAT, {5, 4, 1}, "p_filt_all" + iStr);
    C_filt_all[i] = graph.addVariable(FLOAT, {5, 4, 4}, "C_filt_all" + iStr);

    p_smooth[i] = graph.addVariable(FLOAT, {4, 1}, "p_smooth" + iStr);
    C_smooth[i] = graph.addVariable(FLOAT, {4, 4}, "C_smooth" + iStr);

    graph.setTileMapping(p_proj_all[i], i);
    graph.setTileMapping(C_proj_all[i], i);
    graph.setTileMapping(p_filt_all[i], i);
    graph.setTileMapping(C_filt_all[i], i);

    graph.setTileMapping(p_smooth[i], i);
    graph.setTileMapping(C_smooth[i], i);

    qs[i] = qFlat.reshape({4, 4});

    hs[i] = hFlat.reshape({4, 4});
    gs[i] = gFlat.reshape({4, 4});
    fs[i] = fFlat.reshape({4, 4});

    graph.setTileMapping(covFlat[i], i);
    graph.setTileMapping(qFlat, i);
    graph.setTileMapping(d[i], i);
    graph.setTileMapping(dInit[i], i);
    graph.setTileMapping(dSkip[i], i);

    graph.setTileMapping(covs[i], i);
    graph.setTileMapping(qs[i], i);

    graph.setTileMapping(gFlat, i);
    graph.setTileMapping(hFlat, i);
    graph.setTileMapping(fFlat, i);
    graph.setTileMapping(hs[i], i);
    graph.setTileMapping(gs[i], i);
    graph.setTileMapping(fs[i], i);
  }

  // Init p with hits
  auto [packIterationTensorsInit, ps] = KalmanFilter::packIterationTensors(graph, loop, scatterInto, inputs);

  prog.add(packIterationTensorsInit);

  // Prepare hits for each loop
  auto [packIterationTensorsProg, hits] = KalmanFilter::packIterationTensors(graph, loop, scatterInto, inputs);

  auto [projProg, p_proj, C_proj] = KalmanFilter::project(graph, ps, covs, fs, qs);

  auto [filterSeq, outP, C] = KalmanFilter::filter(graph, hs, gs, hits, p_proj, C_proj);

  std::vector<Tensor> p_proj_chi2(ps.size());
  std::vector<Tensor> p_filt_chi2(ps.size());
  std::vector<Tensor> C_proj_chi2(ps.size());

  std::vector<Tensor> chiSqThreshold(ps.size());

  // For chi2 calc, test
  for (uint i = 0; i < ps.size(); i++) {

    std::string iStr = std::to_string(i);

    p_proj_chi2[i] = graph.addVariable(FLOAT, {4, 1}, "p_proj_chi2" + iStr);
    graph.setTileMapping(p_proj_chi2[i], i);

    p_filt_chi2[i] = graph.addVariable(FLOAT, {4, 1}, "p_filt_chi2" + iStr);
    graph.setTileMapping(p_filt_chi2[i], i);

    C_proj_chi2[i] = graph.addVariable(FLOAT, {4, 4}, "C_proj_chi2" + iStr);
    graph.setTileMapping(C_proj_chi2[i], i);

    chiSqThreshold[i] = graph.addConstant<float>(FLOAT, {1}, {0.4});
    graph.setTileMapping(chiSqThreshold[i], i);

  }

  // ChiSq computation

  auto [resSeq, res] = KalmanFilter::calcResidual(graph, hits, p_filt_chi2, hs);

  auto [chiSqSeq, chiSq] = KalmanFilter::calcChiSq(graph, res, gs, C_proj_chi2, p_proj_chi2, p_filt_chi2);

  auto [chiSqTestSeq, chiSqTestPred] = KalmanFilter::chiSqTest(graph, chiSq, chiSqThreshold);

  // Update loop index
  Sequence updateIterator;

  for (uint i = 0; i < inputs.size(); i++) {

    auto [computeIterate, itr] = KalmanFilter::iterate(graph, loop[i], i);
    updateIterator.add(Execute(computeIterate));
    updateIterator.add(Copy(itr, loop[i]));

  }

  Sequence planeLoop;

  planeLoop.add(packIterationTensorsProg);

  // For EKF
  // planeLoop.add(stateProg);
  // planeLoop.add(jacProg);

  planeLoop.add(projProg);
  planeLoop.add(filterSeq);

  // Save proj and filt states for smoothing step

  for (uint i = 0; i < ps.size(); i++) {

    auto [append_p_proj_Seq, p_proj_new] = KalmanFilter::appendTo(graph, p_proj[i], loop[i], p_proj_all[i], i);
    auto [append_C_proj_Seq, C_proj_new] = KalmanFilter::appendTo(graph, C_proj[i], loop[i], C_proj_all[i], i);
    auto [append_p_filt_Seq, p_filt_new] = KalmanFilter::appendTo(graph, outP[i], loop[i], p_filt_all[i], i);
    auto [append_C_filt_Seq, C_filt_new] = KalmanFilter::appendTo(graph, C[i], loop[i], C_filt_all[i], i);

    planeLoop.add(Execute(append_p_proj_Seq));
    planeLoop.add(Execute(append_C_proj_Seq));
    planeLoop.add(Execute(append_p_filt_Seq));
    planeLoop.add(Execute(append_C_filt_Seq));

    planeLoop.add(Copy(p_proj_new, p_proj_all[i]));
    planeLoop.add(Copy(C_proj_new, C_proj_all[i]));
    planeLoop.add(Copy(p_filt_new, p_filt_all[i]));
    planeLoop.add(Copy(C_filt_new, C_filt_all[i]));

  }

  // For chi2 calc, test
  // for (uint i = 0; i < ps.size(); i++) {
  //
  //   planeLoop.add(Copy(p_proj[i], p_proj_chi2[i]));
  //   planeLoop.add(Copy(outP[i], p_filt_chi2[i]));
  //   planeLoop.add(Copy(C_proj[i], C_proj_chi2[i]));
  //
  // }

  // planeLoop.add(resSeq);
  // planeLoop.add(chiSqSeq);
  // planeLoop.add(chiSqTestSeq);

  planeLoop.add(updateIterator);

  // Set up p for next projection, using p_filt
  // Set up covs for next projection, using C(_filt)
  for (uint i = 0; i < inputs.size(); i++) {
    planeLoop.add(Copy(outP[i], ps[i]));
    planeLoop.add(Copy(C[i], covs[i]));
  }

  // for (uint i = 0; i < inputs.size(); i++) {
  //
  //   auto [skipSwitchSeq, switchP, switchC, outD] = skipSwitch(graph, ps[i], covs[i],
  //                                                     outP[i], C[i],
  //                                                     chiSq[i],
  //                                                     i);
  //   planeLoop.add(Execute(skipSwitchSeq));
  //   planeLoop.add(Copy(switchP, ps[i]));
  //   planeLoop.add(Copy(switchC, covs[i]));
  //   planeLoop.add(Copy(outD, d[i]));
  // }

  prog.add(poplar::program::Repeat(5, planeLoop));

  for (uint i = 0; i < inputs.size(); i++) {
    // Reset iterator for smoothing step
    prog.add(Copy(zero[i], loop[i]));

    // Copy last filtered state to initial smoothing state
    prog.add(Copy(outP[i], p_smooth[i]));
    prog.add(Copy(C[i], C_smooth[i]));
  }

  std::vector<Tensor> p_proj_smooth(inputs.size());
  std::vector<Tensor> C_proj_smooth(inputs.size());
  std::vector<Tensor> p_filt_smooth(inputs.size());
  std::vector<Tensor> C_filt_smooth(inputs.size());

  Sequence smoothLoop;

  for (uint i = 0; i < inputs.size(); i++) {

    // Offset iterator for filtered states (when starting from idx == 0)
    auto [loopFiltSeq, loopFilt] = KalmanFilter::iterate(graph, loop[i], i);
    smoothLoop.add(Execute(loopFiltSeq));

    auto [sm_p_proj_seq, p_proj] = KalmanFilter::smoothingState(graph, p_proj_all[i], loop[i], i);
    auto [sm_C_proj_seq, C_proj] = KalmanFilter::smoothingState(graph, C_proj_all[i], loop[i], i);
    auto [sm_p_filt_seq, p_filt] = KalmanFilter::smoothingState(graph, p_filt_all[i], loopFilt, i);
    auto [sm_C_filt_seq, C_filt] = KalmanFilter::smoothingState(graph, C_filt_all[i], loopFilt, i);

    smoothLoop.add(Execute(sm_p_proj_seq));
    smoothLoop.add(Execute(sm_C_proj_seq));
    smoothLoop.add(Execute(sm_p_filt_seq));
    smoothLoop.add(Execute(sm_C_filt_seq));

    p_proj_smooth[i] = p_proj;
    C_proj_smooth[i] = C_proj;
    p_filt_smooth[i] = p_filt;
    C_filt_smooth[i] = C_filt;

  }

  auto [smoothSeq, p_smooth_new, C_smooth_new] = KalmanFilter::smooth(graph, p_smooth, C_smooth, p_filt_smooth, C_filt_smooth, p_proj_smooth, C_proj_smooth, fs);

  smoothLoop.add(smoothSeq);

  smoothLoop.add(PrintTensor("p_smooth_new", p_smooth_new[0]));

  for (uint i = 0; i < inputs.size(); i++) {

    // Propagate last smoothed state
    smoothLoop.add(Copy(p_smooth_new[i], p_smooth[i]));
    smoothLoop.add(Copy(C_smooth_new[i], C_smooth[i]));

  }

  smoothLoop.add(updateIterator);

  prog.add(poplar::program::Repeat(4, smoothLoop));

  // After loop over planes, for the next in the batch:

  for (uint i = 0; i < inputs.size(); i++) {
    // Reset plane iterator for next batch
    prog.add(Copy(zero[i], loop[i]));
  }

  Sequence progRepeat;
  progRepeat.add(poplar::program::Repeat(batch_size, prog));

  Sequence progMain;
  progMain.add(preProg);
  progMain.add(progRepeat);

  // Engine engine(graph, prog);
  Engine engine(graph, progMain);
  engine.load(dev);

  // Test input
  std::vector<float> v1 = {
    -0.02062073, -0.12062073,
    -0.02062073, -0.22062073,
    -0.12062073, -0.42062073,
    -0.12062073, -0.52062073,
    -0.22062073, -0.62062073,
  };

  std::vector<std::vector<float>> vs;
  for (uint i = 0; i < inputs.size(); i++) {
    // Instead of 1 * 5 inputs, N * 5 inputs
    std::vector<float> v10;
    for (uint j = 0; j < batch_size; j++) {
      v10.insert(std::end(v10), std::begin(v1), std::end(v1));
    }
    vs.push_back(v10);
  }

  for (uint i = 0; i < inputs.size(); i++) {
    engine.connectStream(inStreams[i], &vs[i][0], &vs[i][5 * 2 * batch_size]);
  }

  {
  boost::timer::auto_cpu_timer t;
  engine.run(0);
  }

  return 0;
}


/*
 * Copyright (c) 2019 Opticks Team. All Rights Reserved.
 *
 * This file is part of Opticks
 * (see https://bitbucket.org/simoncblyth/opticks).
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *   http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

#include <cmath>
#include <cassert>
#include <cmath>
#include <cstring>

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

#include "OpticksCSG.h"

#include "NPrism.hpp"
#include "NBBox.hpp"
#include "NPlane.hpp"
#include "NPart.hpp"

nprism::nprism(float apex_angle_degrees, float height_mm, float depth_mm, float fallback_mm)
{
    param.f.x = apex_angle_degrees  ;
    param.f.y = height_mm  ;
    param.f.z = depth_mm  ;
    param.f.w = fallback_mm  ;
}

nprism::nprism(const nquad& param_)
{
    param = param_ ;
}

float nprism::height()
{
    return param.f.y > 0.f ? param.f.y : param.f.w ;
}
float nprism::depth()
{
    return param.f.z > 0.f ? param.f.z : param.f.w ;
}
float nprism::hwidth()
{
    float pi = boost::math::constants::pi<float>() ;
    return height()*tan((pi/180.f)*param.f.x/2.0f) ;
}

void nprism::dump(const char* msg)
{
    param.dump(msg);
}

nbbox nprism::bbox()
{
    float h  = height();
    float hw = hwidth();
    float d  = depth();

    nbbox bb ;
    bb.min = {-hw,0.f,-d/2.f } ;
    bb.max = { hw,  h, d/2.f } ;

    return bb ;
}

npart nprism::part()
{
    // hmm more dupe of hemi-pmt.cu/make_prism
    // but if could somehow make vector types appear
    // the same could use same code with CUDA ?

    nbbox bb = bbox();

    npart p ;
    p.zero();
    p.setParam(param) ;
    p.setTypeCode(CSG_PRISM);
    p.setBBox(bb);

    return p ;
}


#include "matrix_product_dense.h"
#include <optional>
#include <Eigen/Core>
#include <Eigen/LU>

template<typename T> using MatrixType  = Eigen::Matrix<T, Eigen::Dynamic, Eigen::Dynamic, Eigen::ColMajor>;
template<typename T> using VectorType  = Eigen::Matrix<T, Eigen::Dynamic, 1, Eigen::ColMajor>;
template<typename T> using VectorTypeT = Eigen::Matrix<T, 1, Eigen::Dynamic, Eigen::RowMajor>;

namespace dense_lu {
    std::optional<Eigen::PartialPivLU<MatrixType<double>>>               lu_real       = std::nullopt;
    std::optional<Eigen::PartialPivLU<MatrixType<std::complex<double>>>> lu_cplx       = std::nullopt;

    void reset(){
        lu_real.reset();
        lu_cplx.reset();
    }
    template<typename Scalar>
    void init(){
        if constexpr (std::is_same_v<Scalar,double>)
            dense_lu::lu_real = Eigen::PartialPivLU<MatrixType<Scalar>>();
        if constexpr (std::is_same_v<Scalar,std::complex<double>>)
            dense_lu::lu_cplx = Eigen::PartialPivLU<MatrixType<Scalar>>();
    }
}

template<typename Scalar>
DenseMatrixProduct<Scalar>::~DenseMatrixProduct(){
    dense_lu::reset();
}

// Pointer to data constructor, copies the matrix into an internal Eigen matrix.
template<typename Scalar>
DenseMatrixProduct<Scalar>::DenseMatrixProduct(const Scalar *const A_, const int L_,const bool copy_data, const eigutils::eigSetting::Form form_, const eigutils::eigSetting::Side side_)
    : A_ptr(A_), L(L_), form(form_), side(side_) {
    if (copy_data){
        A_stl.resize(L*L);
        std::copy(A_ptr,A_ptr + L*L, A_stl.begin());
        A_ptr = A_stl.data();
    }
    dense_lu::init<Scalar>();
    init_profiling();
}

template<typename Scalar>
void DenseMatrixProduct<Scalar>::print() const {
    Eigen::Map<const MatrixType<Scalar>> A_matrix (A_ptr, L, L);
    std::cout << "A_matrix: \n" << A_matrix << std::endl;
}

// Function definitions

template<typename Scalar>
void DenseMatrixProduct<Scalar>::FactorOP()

/*  Partial pivot LU decomposition
 *  Factors P(A-sigma*I) = LU
 */
{
    if(readyFactorOp) { return; }
    std::cout << "Starting LU \n";
    Eigen::Map<const MatrixType<Scalar>> A_matrix (A_ptr, L, L);

    t_factorOp.tic();
    assert(readyShift and "Shift value sigma has not been set.");

    if constexpr(std::is_same_v<Scalar, double>) {
        dense_lu::lu_real = Eigen::PartialPivLU<MatrixType<Scalar>>();
        dense_lu::lu_real.value().compute(A_matrix - sigmaR * Eigen::MatrixXd::Identity(L, L));
    }
    if constexpr(std::is_same_v<Scalar, std::complex<double>>) {
        Scalar sigma = std::complex<double>(sigmaR, sigmaI);
        dense_lu::lu_cplx = Eigen::PartialPivLU<MatrixType<Scalar>>();
        dense_lu::lu_cplx.value().compute(A_matrix - sigma * Eigen::MatrixXd::Identity(L, L));
    }

    readyFactorOp = true;
    t_factorOp.toc();
    std::cout << "Finished LU \n";
    std::cout << "Time LU Op [ms]: " << std::fixed << std::setprecision(3) << t_factorOp.get_last_time_interval() * 1000 << '\n';
}

template<typename Scalar>
void DenseMatrixProduct<Scalar>::MultOPv(Scalar *x_in_ptr, Scalar *x_out_ptr) {
    using namespace eigutils::eigSetting;
    assert(readyFactorOp and "FactorOp() has not been run yet.");
    switch(side) {
        case Side::R: {
            Eigen::Map<VectorType<Scalar>> x_in(x_in_ptr, L);
            Eigen::Map<VectorType<Scalar>> x_out(x_out_ptr, L);
            if constexpr(std::is_same_v<Scalar, double>)
                x_out.noalias() = dense_lu::lu_real.value().solve(x_in);
            if constexpr(std::is_same_v<Scalar, std::complex<double>>)
                x_out.noalias() = dense_lu::lu_cplx.value().solve(x_in);
            break;
        }
        case Side::L: {
            Eigen::Map<VectorTypeT<Scalar>> x_in(x_in_ptr, L);
            Eigen::Map<VectorTypeT<Scalar>> x_out(x_out_ptr, L);
            if constexpr(std::is_same_v<Scalar, double>)
                x_out.noalias() = x_in * dense_lu::lu_real.value().inverse();
            if constexpr(std::is_same_v<Scalar, std::complex<double>>)
                x_out.noalias() = x_in * dense_lu::lu_cplx.value().inverse();
            break;
        }
    }
    counter++;
}

template<typename Scalar>
void DenseMatrixProduct<Scalar>::MultAx(Scalar *x_in, Scalar *x_out) {
    using namespace eigutils::eigSetting;
    Eigen::Map<const MatrixType<Scalar>> A_matrix (A_ptr, L, L);
    switch(form) {
        case Form::NONSYMMETRIC:
            switch(side) {
                case Side::R: {
                    Eigen::Map<VectorType<Scalar>> x_vec_in(x_in, L);
                    Eigen::Map<VectorType<Scalar>> x_vec_out(x_out, L);
                    x_vec_out.noalias() = A_matrix * x_vec_in;
                    break;
                }
                case Side::L: {
                    Eigen::Map<VectorTypeT<Scalar>> x_vec_in(x_in, L);
                    Eigen::Map<VectorTypeT<Scalar>> x_vec_out(x_out, L);
                    x_vec_out.noalias() = x_vec_in * A_matrix;
                    break;
                }
            }
            break;
        case Form::SYMMETRIC: {
            Eigen::Map<VectorType<Scalar>> x_vec_in(x_in, L);
            Eigen::Map<VectorType<Scalar>> x_vec_out(x_out, L);
            x_vec_out.noalias() = A_matrix.template selfadjointView<Eigen::Upper>() * x_vec_in;
            break;
        }
    }
    counter++;
}

// Explicit instantiations
template class DenseMatrixProduct<double>;
template class DenseMatrixProduct<std::complex<double>>;


/**
 * @file
 * @brief Implementation of linear Lagrangian finite elements for the Dirichlet
 *        problem for the Laplacian
 * @author Ralf Hiptmair
 * @date   October 2018
 * @copyright MIT License
 */

#include <cmath>

#include <boost/filesystem.hpp>
#include <boost/program_options.hpp>

#include <lf/assemble/assemble.h>
#include <lf/geometry/geometry.h>
#include <lf/io/io.h>
#include <lf/mesh/hybrid2d/hybrid2d.h>
#include <lf/refinement/refinement.h>
#include "lf/fe/fe.h"
#include "lf/mesh/test_utils/test_meshes.h"
#include "lf/mesh/utils/utils.h"

static unsigned int dbg_ctrl = 0;
const unsigned int dbg_dofh = 1;
const unsigned int dbg_mesh = 2;
const unsigned int dbg_mat = 4;
const unsigned int dbg_vec = 8;
const unsigned int dbg_bdf = 16;
const unsigned int dbg_elim = 32;
const unsigned int dbg_basic = 64;
const unsigned int dbg_trp = 128;

/**
 * @brief Build a vector of boundary flags for degrees of freedom
 *
 * @param dofh Local-to-global mapper managing indexing for shape
 * functions
 * @return boolean vector of boundary flags, whose length agrees with the
 *         number of global shape functions managed by `dofh`.
 *
 * Every global shape functions belongs to a unique mesh entity. If that
 * entity is contained in the boundary of the mesh, then the flag array
 * entry with the index of the global shape function is set to `true`,
 * otherwise to `false`.
 */
std::vector<bool> flagBoundaryDOFs(const lf::assemble::DofHandler &dofh) {
  const lf::base::size_type N{dofh.NoDofs()};
  // Flag all entities on the boundary
  auto bd_flags{lf::mesh::utils::flagEntitiesOnBoundary(dofh.Mesh())};
  // Run through all global shape functions and check whether
  // they are associated with an entity on the boundary. Store
  // this information in a boolean vector
  std::vector<bool> tmp_bd_flags(N, false);
  for (lf::assemble::gdof_idx_t dofnum = 0; dofnum < N; ++dofnum) {
    const lf::mesh::Entity &dof_entity{dofh.Entity(dofnum)};
    tmp_bd_flags[dofnum] = bd_flags(dof_entity);
    SWITCHEDSTATEMENT(
        dbg_ctrl, dbg_bdf,
        std::cout << "FBD: dof " << dofnum << "@ " << dof_entity << " ["
                  << dofh.Mesh()->Index(dof_entity) << "] ";
        if (tmp_bd_flags[dofnum]) { std::cout << "ON BOUNDARY"; } std::cout
        << std::endl);
  }
  return tmp_bd_flags;
}

/** @brief Eliminate degrees of freedom located on the boundary
 *
 * No longer in use. Replaced with lf::assemble::fix_flagged_solution_components
 */
void eliminateBoundaryDofs(const std::vector<bool> &tmp_bd_flags,
                           lf::assemble::COOMatrix<double> *A) {
  const lf::base::size_type N = tmp_bd_flags.size();
  LF_ASSERT_MSG((A->cols() == N) && (A->rows() == N),
                "Matrix dimension mismath");

  // Remove rows associated with dofs on the boundary
  auto new_last = std::remove_if(
      A->triplets().begin(), A->triplets().end(),
      [&tmp_bd_flags](lf::assemble::COOMatrix<double>::Triplet &triplet) {
        SWITCHEDSTATEMENT(dbg_ctrl, dbg_elim, if (tmp_bd_flags[triplet.row()]) {
          std::cout << "EBD: removing " << triplet.row() << ',' << triplet.col()
                    << "[" << triplet.value() << "]" << std::endl;
        });
        return tmp_bd_flags[triplet.row()];
      });
  A->triplets().erase(new_last, A->triplets().end());
  // Add unit diagonal entries to rows belonging to dofs on the boundary
  for (lf::assemble::gdof_idx_t dofnum = 0; dofnum < N; ++dofnum) {
    if (tmp_bd_flags[dofnum]) {
      A->AddToEntry(dofnum, dofnum, 1.0);
    }
  }
}

/** @brief Insert sampled values from Dirichlet data into right hand side
 * vector
 *
 */
void insertDirichletDataRHS(const std::vector<bool> &tmp_bd_flags,
                            Eigen::VectorXd *rhs,
                            const Eigen::VectorXd &dirichlet_values) {
  const lf::base::size_type N = tmp_bd_flags.size();
  LF_VERIFY_MSG((rhs->size() == N), "rhs vector size mismatch");
  LF_VERIFY_MSG((dirichlet_values.size() == N), "data vector size mismatch");

  for (lf::assemble::gdof_idx_t dofnum = 0; dofnum < N; ++dofnum) {
    if (tmp_bd_flags[dofnum]) {
      // Shape function associated with the boundary
      (*rhs)[dofnum] = dirichlet_values[dofnum];
    }
  }
}

/** @brief Solves Dirichlet problem for the Laplacian
 *
 * @tparam SOLFUNC functor providing boundary values/exact solution
 * @tparam RHSFUNC functor for right hand side source function
 * @param mesh reference to the mesh on which the FE solution is to be
 * computed
 * @param u functor object for solution, also used to sample Dirichlet
 * data
 * @param f object supplying right-hand-side source function
 *
 * @return L2 norm of the nodal error, which is just the L2 norm of the
 *         discretization error approximated by means of the the 2D
 * trapezoidal rule.
 */
template <typename SOLFUNC, typename RHSFUNC>
double L2ErrorLinearFEDirichletLaplacian(
    const std::shared_ptr<const lf::mesh::Mesh> &mesh_p, SOLFUNC &&u,
    RHSFUNC &&f) {
  LF_ASSERT_MSG(mesh_p != nullptr, "Invalid mesh pointer");
  LF_ASSERT_MSG((mesh_p->DimMesh() == 2) && (mesh_p->DimWorld() == 2),
                "For 2D planar meshes only!");
  // Debugging output
  SWITCHEDSTATEMENT(
      dbg_ctrl, dbg_basic,
      std::cout << "Dirichlet Laplacian: Linear FE L2 error on mesh with "
                << mesh_p->Size(0) << " cells, " << mesh_p->Size(1)
                << " edges, " << mesh_p->Size(2) << " nodes" << std::endl;)
  SWITCHEDSTATEMENT(
      dbg_ctrl, dbg_mesh,
      const int tmp_mesh_ctrl = lf::mesh::hybrid2d::Mesh::output_ctrl_;
      lf::mesh::hybrid2d::Mesh::output_ctrl_ = 100;
      lf::mesh::utils::PrintInfo(*mesh_p, std::cout);
      lf::mesh::hybrid2d::Mesh::output_ctrl_ = tmp_mesh_ctrl);
  // Initialize objects for local computations
  lf::fe::LinearFELaplaceElementMatrix loc_mat_laplace{};
  lf::fe::LinearFELocalLoadVector<double, decltype(f)> loc_vec_sample(f);
  // Initialization of index mapping for linear finite elements
  lf::assemble::UniformFEDofHandler loc_glob_map(
      mesh_p, {{lf::base::RefEl::kPoint(), 1}});
  SWITCHEDSTATEMENT(dbg_ctrl, dbg_dofh,
                    std::cout << loc_glob_map << std::endl;);
  // Dimension of finite element space
  const lf::assemble::size_type N_dofs(loc_glob_map.NoDofs());
  // Matrix in triplet format holding Galerkin matrix
  lf::assemble::COOMatrix<double> mat(N_dofs, N_dofs);
  // Building the Galerkin matrix (trial space = test space)
  // This Galerkin matrix is oblivious of Dirichlet boundary conditions
  mat = lf::assemble::AssembleMatrixLocally<lf::assemble::COOMatrix<double>>(
      0, loc_glob_map, loc_mat_laplace);
  // Debugging output
  SWITCHEDSTATEMENT(dbg_ctrl, dbg_trp, mat.PrintInfo(std::cout));
  SWITCHEDSTATEMENT(dbg_ctrl, dbg_mat,
                    std::cout << "Full " << mat.rows() << 'x' << mat.cols()
                              << " stiffness matrix, " << mat.triplets().size()
                              << " tripets:\n"
                              << mat.makeDense() << std::endl);

  // Filling the right-hand-side vector
  auto rhsvec = lf::assemble::AssembleVectorLocally<Eigen::VectorXd>(
      0, loc_glob_map, loc_vec_sample);
  // Sample Dirichlet date from the exact solution
  Eigen::VectorXd dirichlet_data(loc_glob_map.NoDofs());
  const Eigen::Matrix<double, 0, 1> ref_coord{
      Eigen::Matrix<double, 0, 1>::Zero(0, 1)};
  for (const lf::mesh::Entity &node : mesh_p->Entities(mesh_p->DimMesh())) {
    LF_ASSERT_MSG(node.RefEl() == lf::base::RefEl::kPoint(),
                  "Wrong topological type for a node");
    const Eigen::Vector2d point = node.Geometry()->Global(ref_coord);
    const lf::assemble::size_type num_int_dof =
        loc_glob_map.NoInteriorDofs(node);
    LF_ASSERT_MSG(num_int_dof == 1, "Node with " << num_int_dof << " dof");
    const lf::base::RandomAccessRange<const lf::assemble::gdof_idx_t> gsf_idx(
        loc_glob_map.InteriorGlobalDofIndices(node));
    const lf::assemble::gdof_idx_t node_dof_idx = gsf_idx[0];
    dirichlet_data[node_dof_idx] = u(point);
  }

  // modify linear system in order to take into account boundary data
  std::vector<bool> tmp_bd_flags{flagBoundaryDOFs(loc_glob_map)};
  // >> Old version
  // eliminateBoundaryDofs(tmp_bd_flags, mat);
  // insertDirichletDataRHS(tmp_bd_flags, rhsvec, dirichlet_data);
  // Identify dof indices associated with the boundary
  // >>
  // >> Equivalent new versions
  lf::assemble::fix_flagged_solution_comp_alt<double>(
      [&tmp_bd_flags,
       &dirichlet_data](lf::assemble::gdof_idx_t i) -> std::pair<bool, double> {
        LF_ASSERT_MSG((i < tmp_bd_flags.size()) && (i < dirichlet_data.size()),
                      "Illegal index " << i);
        return std::make_pair(tmp_bd_flags[i], dirichlet_data[i]);
      },
      mat, rhsvec);
  // lf::assemble::fix_flagged_solution_components<double>(
  //   tmp_bd_flags, dirichlet_data, mat, rhsvec);
  // >>
  // Debugging output
  SWITCHEDSTATEMENT(dbg_ctrl, dbg_mat,
                    std::cout << "Reduced " << mat.rows() << 'x' << mat.cols()
                              << " stiffness matrix, " << mat.triplets().size()
                              << " triplets:\n"
                              << mat.makeDense() << std::endl);

  // Initialize sparse matrix
  Eigen::SparseMatrix<double> stiffness_matrix(mat.makeSparse());
  // Solve linear system
  Eigen::SparseLU<Eigen::SparseMatrix<double>> solver;
  solver.compute(stiffness_matrix);
  Eigen::VectorXd sol_vec = solver.solve(rhsvec);
  if (solver.info() != Eigen::Success) {
    std::cout << "solver failed!" << std::endl;
  }

  // Compute the norm of nodal error cell by cell
  double nodal_err = 0.0;
  for (const lf::mesh::Entity &cell : mesh_p->Entities(0)) {
    const lf::base::RandomAccessRange<const lf::assemble::gdof_idx_t>
        cell_dof_idx(loc_glob_map.GlobalDofIndices(cell));
    LF_ASSERT_MSG(loc_glob_map.NoLocalDofs(cell) == cell.RefEl().NumNodes(),
                  "Inconsistent node number");
    const lf::base::size_type num_nodes = cell.RefEl().NumNodes();
    double sum = 0.0;
    for (int k = 0; k < num_nodes; ++k) {
      sum += std::pow(
          sol_vec[cell_dof_idx[k]] - dirichlet_data[cell_dof_idx[k]], 2);
    }
    nodal_err += lf::geometry::Volume(*cell.Geometry()) * (sum / num_nodes);
  }
  return std::sqrt(nodal_err);
}

/** @brief Solves Dirichlet problem for the Laplacian on a sequence of
 *         regularly refined meshes
 * @param coarse_mesh_p pointer to coarsest mesh
 * @param reflevels number of refinements
 * @param u exact solution of the boundary value problem
 * @param f right hand side source function
 * @return L2 norms of discretization errors on each refinement level
 *
 */
template <typename SOLFUNCTOR, typename RHSFUNCTOR>
std::vector<double> SolveDirLaplSeqMesh(
    std::shared_ptr<lf::mesh::Mesh> coarse_mesh_p, unsigned int reflevels,
    SOLFUNCTOR &&u, RHSFUNCTOR &&f) {
  // Prepare for creating a hierarchy of meshes
  std::shared_ptr<lf::mesh::hybrid2d::MeshFactory> mesh_factory_ptr =
      std::make_shared<lf::mesh::hybrid2d::MeshFactory>(2);
  lf::refinement::MeshHierarchy multi_mesh(std::move(coarse_mesh_p),
                                           mesh_factory_ptr);

  // Perform several steps of regular refinement of the given mesh
  for (int refstep = 0; refstep < reflevels; ++refstep) {
    multi_mesh.RefineRegular(/*lf::refinement::RefPat::rp_barycentric*/);
  }
  // Solve Dirichlet boundary value problem on every level
  lf::assemble::size_type L = multi_mesh.NumLevels();
  std::vector<double> errors(L);
  for (int level = 0; level < L; level++) {
    errors.push_back(
        L2ErrorLinearFEDirichletLaplacian(multi_mesh.getMesh(level), u, f));
  }
  return errors;
}

int main(int argc, const char **argv) {
  // Pointer to the current mesh
  std::shared_ptr<lf::mesh::Mesh> mesh_p;

  // Processing command line arguments
  bool verbose = false;
  namespace po = boost::program_options;
  po::options_description desc("Allowed options");
  // clang-format off
  desc.add_options()
  ("help,h", "-h -v -f <filename> -s <selection>")
  ("filename,f", "File to load coarse mesh from ")
  ("selector,s", po::value<int>()->default_value(0), "Selection of test mesh")
  ("reflevels,r", po::value<int>()->default_value(2), "Number of refinement levels")
  ("bvpsel,b", po::value<int>()->default_value(0),
   "Selector for Dirichlet data and rhs function")
  ("verbose,v", po::bool_switch(&verbose),"Enable verbose mode");
  // clang-format on
  po::variables_map vm;
  po::store(po::parse_command_line(argc, argv, desc), vm);
  if (vm.count("help") > 0) {
    std::cout << desc << std::endl;
    std::cout << "Internal variables that can be set by name=value args"
              << std::endl;
    lf::base::ListCtrlVars(std::cout);
  } else {
    lf::base::ReadCtrVarsCmdArgs(argc, argv);
    std::cout << "*** Solving Dirichlet problems for the Laplacian ***"
              << std::endl;
    // Retrieve number of degrees of freedom for each entity type from
    // command line arguments
    if (vm.count("filename") > 0) {
      // A filename was specified
      std::string filename{vm["filename"].as<std::string>()};
      if (filename.length() > 0) {
        std::cout << "Reading mesh from file " << filename << std::endl;
        boost::filesystem::path here = __FILE__;
        auto mesh_file_path = here.parent_path() / filename.c_str();
        auto mesh_factory =
            std::make_unique<lf::mesh::hybrid2d::MeshFactory>(2);
        lf::io::GmshReader reader(std::move(mesh_factory),
                                  mesh_file_path.string());
        mesh_p = reader.mesh();
      }
    } else {
      std::cout << "No mesh file supplied, using GenerateHybrid2DTestMesh()"
                << std::endl;
      if (vm.count("selector") > 0) {
        const int selector = vm["selector"].as<int>();
        std::cout << "Using test mesh no " << selector << std::endl;
        if ((selector >= 0) && (selector <= 4)) {
          mesh_p = lf::mesh::test_utils::GenerateHybrid2DTestMesh(selector);
        }
      }
    }
    // Set number of refinement levels
    unsigned int reflevels = 2;
    if (vm.count("reflevels") > 0) {
      reflevels = vm["reflevels"].as<int>();
    }
    unsigned int bvpsel = 0;
    if (vm.count("bvpsel") > 0) {
      bvpsel = vm["bvpsel"].as<int>();
    }
    if (mesh_p == nullptr) {
      // Default mesh
      std::cout << "Using default mesh; test mesh 0" << std::endl;
      mesh_p = lf::mesh::test_utils::GenerateHybrid2DTestMesh(0);
    }
    // At this point a pointer to the mesh is stored in mesh_p
    // Output summary information about the coarsest mesh
    std::cout << "Coarse mesh: " << mesh_p->Size(0) << " cells, "
              << mesh_p->Size(1) << " edges, " << mesh_p->Size(2) << " vertices"
              << std::endl;
    std::cout << reflevels << " refinement levels requested" << std::endl;

    // Problem data provided by function pointers
    std::function<double(const Eigen::Vector2d &)> u, f;

    // Initialize the problem data
    std::cout << "Problem setting " << bvpsel << " selected" << std::endl;
    switch (bvpsel) {
      case 0: {
        // A linear solution, no error, if contained in FE space
        f = [](const Eigen::Vector2d &) { return 0.0; };
        u = [](const Eigen::Vector2d &x) { return (x[0] + 2.0 * x[1]); };
        break;
      }
      case 1: {
        // Quadratic polynomial solution
        f = [](const Eigen::Vector2d &) { return -4.0; };
        u = [](const Eigen::Vector2d &x) {
          return (std::pow(x[0], 2) + std::pow(x[1], 2));
        };
        break;
      }
      default: {
        LF_VERIFY_MSG(false, "Illegal problem number");
        break;
      }
    }

    // Set debugging switches
    lf::fe::LinearFELaplaceElementMatrix::dbg_ctrl = 0;
    // LinearFELaplaceElementMatrix::dbg_geo |
    // LinearFELaplaceElementMatrix::dbg_locmat;
    lf::fe::LinearFELocalLoadVector<double, decltype(f)>::dbg_ctrl = 0;
    lf::assemble::DofHandler::output_ctrl_ = 6;
    dbg_ctrl = dbg_basic;  // | dbg_mat | dbg_mesh | dbg_dofh | dbg_trp;
    // lf::assemble::ass_mat_dbg_ctrl = 255;

    // Compute finite element solution and error
    auto L2errs = SolveDirLaplSeqMesh(mesh_p, reflevels, u, f);
    int level = 0;
    for (auto &err : L2errs) {
      std::cout << "L2 rrror on level " << level << " = " << err << std::endl;
      level++;
    }
  }
  return 0;
}


// Distributed under the MIT License.
// See LICENSE.txt for details.

#pragma once

#include <array>
#include <boost/optional.hpp>
#include <cstddef>
#include <limits>

#include "DataStructures/Tensor/TypeAliases.hpp"
#include "Utilities/TypeTraits.hpp"

namespace PUP {
class er;
}  // namespace PUP

namespace CoordinateMaps {

/*!
 * \ingroup CoordinateMapsGroup
 *
 * \brief Redistributes gridpoints on the sphere.
 * \image html EquatorialCompression.png "A sphere with an `aspect_ratio` of 3."
 *
 * \details A mapping from the sphere to itself which depends on a single
 * parameter, the `aspect_ratio` \f$\alpha\f$, which is the ratio of the
 * horizontal length to the vertical height for a given point. This parameter
 * name was chosen because points with \f$\tan \theta = 1\f$ get mapped to
 * points with \f$\tan \theta' = \alpha\f$. In general, gridpoints located
 * at an angle \f$\theta\f$ from the pole are mapped to a new angle
 * \f$\theta'\f$ satisfying \f$\tan \theta' = \alpha \tan \theta\f$.
 *
 * For an `aspect_ratio` greater than one, the gridpoints are mapped towards
 * the equator, leading to an equatorially compressed grid. For an
 * `aspect_ratio` less than one, the gridpoints are mapped towards the poles.
 * Note that the aspect ratio must be positive.
 *
 * We define the auxiliary variables \f$ r := \sqrt{x^2 + y^2 +z^2}\f$
 * and \f$ \rho := \sqrt{x^2 + y^2 + \alpha^{-2} z^2}\f$.
 *
 * The map corresponding to this transformation in cartesian coordinates
 * is then given by:
 *
 * \f[\vec{x}'(x,y,z) =
 * \frac{r}{\rho}\begin{bmatrix}
 * x\\
 * y\\
 * \alpha^{-1} z\\
 * \end{bmatrix}\f]
 *
 */
class EquatorialCompression {
 public:
  static constexpr size_t dim = 3;
  explicit EquatorialCompression(double aspect_ratio) noexcept;
  EquatorialCompression() = default;
  ~EquatorialCompression() = default;
  EquatorialCompression(EquatorialCompression&&) = default;
  EquatorialCompression(const EquatorialCompression&) = default;
  EquatorialCompression& operator=(const EquatorialCompression&) = default;
  EquatorialCompression& operator=(EquatorialCompression&&) = default;

  template <typename T>
  std::array<tt::remove_cvref_wrap_t<T>, 3> operator()(
      const std::array<T, 3>& source_coords) const noexcept;

  boost::optional<std::array<double, 3>> inverse(
      const std::array<double, 3>& target_coords) const noexcept;

  template <typename T>
  tnsr::Ij<tt::remove_cvref_wrap_t<T>, 3, Frame::NoFrame> jacobian(
      const std::array<T, 3>& source_coords) const noexcept;

  template <typename T>
  tnsr::Ij<tt::remove_cvref_wrap_t<T>, 3, Frame::NoFrame> inv_jacobian(
      const std::array<T, 3>& source_coords) const noexcept;

  // clang-tidy: google runtime references
  void pup(PUP::er& p) noexcept;  // NOLINT

 private:
  template <typename T>
  std::array<tt::remove_cvref_wrap_t<T>, 3> angular_distortion(
      const std::array<T, 3>& coords, double inverse_alpha) const noexcept;
  template <typename T>
  tnsr::Ij<tt::remove_cvref_wrap_t<T>, 3, Frame::NoFrame>
  angular_distortion_jacobian(const std::array<T, 3>& coords,
                              double inverse_alpha) const noexcept;
  friend bool operator==(const EquatorialCompression& lhs,
                         const EquatorialCompression& rhs) noexcept;

  double aspect_ratio_{std::numeric_limits<double>::signaling_NaN()};
  double inverse_aspect_ratio_{std::numeric_limits<double>::signaling_NaN()};
};
bool operator!=(const EquatorialCompression& lhs,
                const EquatorialCompression& rhs) noexcept;
}  // namespace CoordinateMaps


/* ScaFES
 * Copyright (c) 2011-2015, ZIH, TU Dresden, Federal Republic of Germany.
 * For details, see the files COPYING and LICENSE in the base directory
 * of the package.
 */

#ifndef MODULATEDGAUSSIANSOURCE_HPP_
#define    MODULATEDGAUSSIANSOURCE_HPP_

#include <boost/property_tree/ptree.hpp>

template <typename T>
class ModulatedGaussianSource {
public:
    ModulatedGaussianSource(const ModulatedGaussianSource&)=delete;

    ~ModulatedGaussianSource() { };

    ModulatedGaussianSource(const T& dt, const boost::property_tree::ptree& p)
    : dt_(dt)
    , f_mod_(p.get<T>("source.f_mod"))
    , t0_(p.get<T>("source.t0"))
    , spread_(p.get<T>("source.spread")) { };

    template <typename TT> TT operator()(const TT & n) const
    {
        TT v=(t0_-n)/spread_;
        v*=v;
        v*=-.5;
        v=exp(v);
        TT mod=cos(2.*M_PI*f_mod_*n*dt_);
        TT ret=v*mod;
        return ret;
    }

private:
    const T dt_;
    const T f_mod_;
    const T t0_;
    const T spread_;
};

#endif    /* MODULATEDGAUSSIANSOURCE_HPP_ */


/* Copyright (c) 2018, Skolkovo Institute of Science and Technology (Skoltech)
 * All rights reserved.
 *
 * See LICENSE file in the root of the mrob library.
 *
 *
 * planeRegistration.cpp
 *
 *  Created on: Jan 28, 2019
 *      Author: Gonzalo Ferrer
 *              g.ferrer@skoltech.ru
 *              Mobile Robotics Lab, Skoltech
 */

#include "mrob/pc_registration.hpp"
#include "mrob/plane_registration.hpp"
#include <Eigen/LU> // for inverse and determinant
#include <Eigen/Eigenvalues>
#include <iostream>

#include <chrono>

using namespace mrob;

PlaneRegistration::PlaneRegistration():
        numberPlanes_(0), numberPoses_(0),isSolved_(0), trajectory_(new std::vector<SE3>(8,SE3())),
        solveMode_(SolveMode::GRADIENT),
        c1_(1e-4), c2_(0.9), alpha_(0.75), beta_(0.1)
{
    // Optmizer does not establish the size for this matrices and thus it is required
    gradient_.resize(6);
    hessian_.resize(6,6);
}

PlaneRegistration::~PlaneRegistration()
{

}

void PlaneRegistration::set_number_planes_and_poses(uint_t numberPlanes, uint_t numberPoses)
{
    planes_.clear();
    planes_.reserve(numberPlanes);
    trajectory_->clear();
    trajectory_->resize(numberPoses, SE3());
    numberPlanes_ = numberPlanes;
    numberPoses_ = numberPoses;

    previousState_.clear();
    previousState_.resize(numberPoses, Mat61::Zero());
}

void PlaneRegistration::reset_solution()
{
    // trajectory is reset
    trajectory_->clear();
    trajectory_->resize(numberPoses_, SE3());
    previousState_.clear();
    previousState_.resize(numberPoses_, Mat61::Zero());
}

uint_t PlaneRegistration::solve(SolveMode mode, bool singleIteration)
{
    // just in case some methods, such as gradient, has several modes
    solveMode_ = mode;
    switch(mode)
    {
        case SolveMode::INITIALIZE:
            return solve_initialize();
        case SolveMode::GRADIENT:
        case SolveMode::GRADIENT_BENGIOS_NAG:
            return solve_interpolate_gradient(singleIteration);
        case SolveMode::GN_HESSIAN:
            return optimize(NEWTON_RAPHSON);
        case SolveMode::GN_CLAMPED_HESSIAN:
            return solve_interpolate_hessian(singleIteration);
        case SolveMode::LM_SPHER:
            return optimize(LEVENBERG_MARQUARDT_SPHER);
        case SolveMode::LM_ELLIP:
            return optimize(LEVENBERG_MARQUARDT_ELLIP);
        default:
            return 0;
    }
}

uint_t PlaneRegistration::solve_interpolate_gradient(bool singleIteration)
{
    // This function is from the first implementation, is just kept for comparisons (but should not be used to solve the problem)
    // iterative process, on convergence basis | error_k - error_k-1| < tol
    solveIters_ = 0;
    double previousError = 1e20, diffError = 10;
    do
    {
        // 1) calculate plane estimation given the current trajectory
        double  initialError = 0.0;
        for (auto it = planes_.cbegin();  it != planes_.cend(); ++it)
        {
            initialError += it->second->estimate_plane();
        }
        diffError = previousError - initialError;
        previousError = initialError;

        std::cout << "current error iteration " << solveIters_ << " = "<< initialError << std::endl;

        // 2) calculate Gradient = Jacobian^T. We maintain the nomenclature Jacobian for coherence on the project,
        //    but actually this Jacobian should be transposed.
        Mat61 jacobian = Mat61::Zero(), accumulatedJacobian = Mat61::Zero();
        double  numberPoints, tau = 1.0 / (double)(numberPoses_-1);
        for (uint_t t = 1 ; t < numberPoses_; ++t)
        {
            numberPoints = 0.0;
            jacobian.setZero();
            for (auto it = planes_.cbegin();  it != planes_.cend(); ++it)
            {
                jacobian += it->second->calculate_gradient(t);
                numberPoints += it->second->get_number_points(t);
            }
            // XXX this could be changed to time stamps later
            accumulatedJacobian +=  (tau *  t  / numberPoints / numberPoses_) * jacobian;

        }
        // 3) update results Tf = exp(-dxi) * Tf (our convention, we expanded from the left)
        // 3-1)
        Mat61 dxi, xiFinal;
        if (solveMode_ == SolveMode::GRADIENT)
        {
            double alpha = alpha_;
            dxi = -alpha * accumulatedJacobian;
            //std::cout << "\nINterpolate jacobian : = " << accumulatedJacobian.transpose() << ", and increment update = " << dxi << std::endl;
        }
        // 3.4-B) Bengio's NAG: a modification to NAG as proposed in Bengio-2013. Fixed parameters
        //          1) momentum or velocity  v_k = beta_k-1 v _k-1 - alpha_k-1 Grad f (x_k-1)
        //          2) x_k+1 = x_k + beta_k+1 beta_k * v_k - (1 + beta_k+1)*alpha_k * Grad f(x_k)
        if (solveMode_ == SolveMode::GRADIENT_BENGIOS_NAG)
        {
            double alpha = alpha_;
            double beta = beta_;
            // x update
            dxi = beta * beta * previousState_.back() - (1 + beta) * alpha * accumulatedJacobian;

            // momentum
            previousState_.back() = beta * previousState_.back() - alpha * accumulatedJacobian;
        }
        trajectory_->back().update_lhs(dxi);
        xiFinal = trajectory_->back().ln_vee();
        for (uint_t t = 1 ; t < numberPoses_-1; ++t)
        {
            dxi = tau * t * xiFinal;// SE3 does not like all derived classes TODO
            trajectory_->at(t) = SE3(dxi);
        }
        ++solveIters_;
    }while(fabs(diffError) > 1e-4 && !singleIteration && solveIters_ < 1e4);
    return solveIters_;
}

// TO BE DEPRECATED. Only used for clamped Hessian, and that is shown NOT to work.
uint_t PlaneRegistration::solve_interpolate_hessian(bool singleIteration)
{
    // iterative process, on convergence basis | error_k - error_k-1| < tol
    // For now, only 1 iteration
    solveIters_ = 0;
    double previousError = 1e20, diffError = 10;

    do
    {

        // 1) calculate plane estimation given the current trajectory. Same as solve_interpolate
        double  initialError = get_current_error();

        diffError = previousError - initialError;
        previousError = initialError;

        solveIters_++;
        std::cout << "current error iteration " << solveIters_ << " = "<< initialError << std::endl;

        // 2) calculate Gradient and Hessian
        Mat61 gradient = Mat61::Zero();
        Mat6 hessian = Mat6::Zero();
        gradient__.setZero();
        hessian__.setZero();
        double  tau = 1.0 / (double)(numberPoses_-1);
        for (uint_t t = 1 ; t < numberPoses_; ++t)
        {
            gradient.setZero();
            hessian.setZero();
            for (auto it = planes_.cbegin();  it != planes_.cend(); ++it)
            {
                gradient += it->second->calculate_gradient(t);
                hessian += it->second->calculate_hessian(t);
            }
            // TODO this should be changed to time stamps later
            gradient__ +=  (tau *  t)  * gradient;
            hessian__ += (tau *  t) * hessian.selfadjointView<Eigen::Upper>();
        }
        // 3) calculate update Tf = exp(-dxi) * Tf (our convention, we expanded from the left)
        Mat61 dxi;
        if (solveMode_ == SolveMode::GN_CLAMPED_HESSIAN)
        {
            // we clamp the vector spaces corresponding to negative eigenvals
            Mat6 pseudoInv = Mat6::Zero();
            Eigen::SelfAdjointEigenSolver<Mat6> eigs(hessian__);
            for (uint_t i = 0; i < 6 ; ++i)
            {
                if(eigs.eigenvalues()[i] > 1e-4 || true ) //TODO set tolerance
                {
                    std::cout << "POSITIVE. cos distance to grad = " << eigs.eigenvectors().col(i).dot(gradient)/gradient.norm()
                              << ", eigs = " << eigs.eigenvalues()[i] << std::endl;
                    pseudoInv += (1.0/eigs.eigenvalues()(i)) * eigs.eigenvectors().col(i) * eigs.eigenvectors().col(i).transpose();
                    // XXX why is this function not monotonically decreasing? this is annoying, but makes clamping a bad idea: LM!
                }
                else
                {
                    std::cout << "NEGATIVE. cos distance to grad = " << eigs.eigenvectors().col(i).dot(gradient)/gradient.norm()
                              << ", eigs = " << eigs.eigenvalues()[i] << std::endl;

                }
            }
            dxi = - pseudoInv * gradient__;
        }
        else
            dxi = - hessian__.inverse() * gradient__;
        trajectory_->back().update_lhs(dxi);

        // 4) update full trajectory. Here we assume a full rank matrix TODO check for degenerate cases
        Mat61 xiFinal = trajectory_->back().ln_vee();
        for (uint_t t = 1 ; t < numberPoses_-1; ++t)
        {
            dxi = tau * t * xiFinal;
            trajectory_->at(t) = SE3(dxi);
        }
    }while(fabs(diffError) > 1e-4 && !singleIteration && solveIters_ < 1e4);
    return solveIters_;
}

uint_t PlaneRegistration::solve_quaternion_plane()
{
    solveIters_ = 0;
    double previousError = 1e20, diffError = 10;

    // TODO create factor graph or call dense solver?

    return solveIters_;
}

uint_t PlaneRegistration::solve_initialize()
{
    // TODO Maybe solve this as a plane-to-point alignment wrt T0
    // Initialize matrices of points
    MatX X(numberPlanes_,3), Y(numberPlanes_,3);

    // create points Y (from frame t =0), minimum 3 planes per pose
    uint_t t = 0;
    for (auto it = planes_.cbegin();  it != planes_.cend(); ++it)
    {
        it->second->calculate_all_matrices_S();
        Y.row(t) = it->second->get_mean_point(0);
        ++t;
    }

    // create points X, for t = 1, ... T, minimum 3 planes per pose
    for (t = 1; t < numberPoses_; ++t)
    {
        uint_t cont = 0;
        X.setZero();
        for (auto it = planes_.cbegin();  it != planes_.cend(); ++it)
        {
            X.row(cont) = it->second->get_mean_point(t);
            ++cont;
        }
        // Arun solver
        SE3 estimatedPose;
        if (!PCRegistration::arun(X,Y,estimatedPose))
            return 0;
        trajectory_->at(t) = estimatedPose;
    }

    // update current trajectory
    return 1;
}

double PlaneRegistration::get_current_error() const
{
    double  currentError = 0.0;
    for (auto it = planes_.cbegin();  it != planes_.cend(); ++it)
        currentError += it->second->estimate_plane();
    return currentError;
}

void PlaneRegistration::add_plane(uint_t id, std::shared_ptr<Plane> &plane)
{
    plane->set_trajectory(trajectory_);
    planes_.emplace(id, plane);
}

double PlaneRegistration::calculate_poses_rmse(std::vector<SE3> & groundTruth) const
{
    assert(groundTruth.size() >= numberPoses_ && "PlaneRegistration::calculate_poses_rmse: number of poses from GT is incorrect\n");
    double rmse= 0.0;
    uint_t t = 0;
    for (auto &pose: groundTruth)
    {
        Mat61 dxi = (trajectory_->at(t) * pose.inv()).ln_vee();
        //std::cout << pose.ln_vee().transpose() << " and solution inv first pose \n" << (invFirstPose * trajectory_->at(t) * pose.inv()).ln_vee().transpose() <<std::endl;
        rmse += dxi.dot(dxi)/(double)numberPoses_;
        ++t;
    }
    return std::sqrt(rmse);
}

//XXX this can be a reference, who does this interface with pybinds?
std::vector<Mat31> PlaneRegistration::get_point_cloud(uint_t time)
{
	std::vector<Mat31> aggregated_pc;
	for (auto it = planes_.cbegin();  it != planes_.cend(); ++it)
    {
        std::vector<Mat31>& plane_pc = it->second->get_points(time);
        // aggregating elements
        aggregated_pc.insert(aggregated_pc.end(), plane_pc.begin(), plane_pc.end());
	}
	return aggregated_pc;
}

Mat4 PlaneRegistration::get_trajectory(uint_t time)
{
    assert(time < numberPoses_ && "CreatePoints::getPointCloud: temporal index larger than number of calculated poses\n");
    if (time < numberPoses_ )
        return trajectory_->at(time).T();
    return Mat4::Identity();
}

void PlaneRegistration::print(bool plotPlanes) const
{
    std::cout << "Printing plane registration data :"<< std::endl;
    for (SE3 &transf : *trajectory_)
        transf.print();
    if (plotPlanes)
    {
        for (auto it = planes_.cbegin();  it != planes_.cend(); ++it)
            it->second->print();
    }
}

// resturns: [0]error, [1]iters, hessdet[2], conditioningNumber[3]
std::vector<double> PlaneRegistration::print_evaluate()
{
    std::vector<double> result(6,0.0);

    result[0] = get_current_error();
    result[1] = solveIters_;

    MatX allPlanes(numberPlanes_,4);
    MatX allNormals(numberPlanes_,3);
    uint_t i = 0;

    switch(solveMode_)
    {
        case SolveMode::GRADIENT:
        case SolveMode::GRADIENT_BENGIOS_NAG:
            hessian__.setZero();
            break;
        case SolveMode::GN_HESSIAN:
        case SolveMode::LM_SPHER:
        case SolveMode::LM_ELLIP:
            gradient__ = gradient_;
            hessian__ = hessian_;
    }

    // Normals on planes, check for rank
    for (auto plane : planes_)
    {
        Mat41 pi = plane.second->get_plane();
        //std::cout << "plane : \n" << pi << std::endl;
        allPlanes.row(i) = pi;
        allNormals.row(i) = pi.head(3)/(pi.head(3).norm());
        ++i;
    }
    // Orthogonality between planes (4 dim)
    std::cout << "current gradient \n" << gradient__ << std::endl;
    // XXX : NO, Orthogonolaity was not an issue
    //std::cout << "solution\n" << allPlanes << "\nOrthogonality between planes: \n" << allPlanes * allPlanes.transpose() <<
    //              "\n and det  = \n" << allPlanes.determinant() << std::endl;

    // XXX No, Orthogonality between normals
    //std::cout << "Orthogonality between normals: \n" << allNormals * allNormals.transpose() <<
    //             "\n and det = \n" << allNormals.determinant() << std::endl;

    // Hessian rank and eigen, look for negative vaps. Lasta hessina calculateds
    Eigen::EigenSolver<MatX> eigs(hessian__);
    std::cout << "eigen values are: \n" << eigs.eigenvalues() << std::endl;
    // Determinant of stacked normals
    std::cout << "det(Hessian) = \n" << hessian__ << std::endl;
    // Hessian conditioning number
    Eigen::JacobiSVD<Mat6> svd(hessian__, Eigen::ComputeThinU | Eigen::ComputeThinV);
    // TODO remove
    //std::cout << "SVD decomposition : \n" << svd.singularValues() <<
    //              "\n vectors :\n" << svd.matrixU() <<
    //             "\n and conditioning number = " << svd.singularValues()(0)/svd.singularValues()(5) << std::endl;

    result[2] = hessian__.determinant();
    //TODO count this and optimze
    result[3] = svd.singularValues()(0)/svd.singularValues()(5);

    return result;
}

// From parent class Optimizer:
// TOOD for now replicated, later we will substitute
matData_t PlaneRegistration::calculate_error()
{
    return get_current_error();
}

void PlaneRegistration::calculate_gradient_hessian()
{
    Mat61 gradient = Mat61::Zero();
    Mat6 hessian = Mat6::Zero();
    gradient_.setZero();
    hessian_.setZero();
    double  tau = 1.0 / (double)(numberPoses_-1);
    for (uint_t t = 1 ; t < numberPoses_; ++t)
    {
        gradient.setZero();
        hessian.setZero();
        for (auto it = planes_.cbegin();  it != planes_.cend(); ++it)
        {
            gradient += it->second->calculate_gradient(t);
            hessian += it->second->calculate_hessian(t);
        }
        // TODO this should be changed to time stamps later
        gradient_ +=  (tau *  t)  * gradient;
        hessian_ += (tau *  t) * hessian.selfadjointView<Eigen::Upper>();
    }
}

void PlaneRegistration::update_state(const MatX1 &dx)
{
    trajectory_->back().update_lhs(dx);
    Mat61 xiFinal = trajectory_->back().ln_vee();
    double  tau = 1.0 / (double)(numberPoses_-1);
    Mat61 dxi;
    for (uint_t t = 1 ; t < numberPoses_-1; ++t)
    {
        dxi = tau * t * xiFinal;
        trajectory_->at(t) = SE3(dxi);
    }
}

void PlaneRegistration::bookkeep_state()
{
    bookept_trajectory_ = trajectory_->back();
}

void PlaneRegistration::update_state_from_bookkeep()
{
    trajectory_->back() = bookept_trajectory_;
    Mat61 xiFinal = bookept_trajectory_.ln_vee();
    double  tau = 1.0 / (double)(numberPoses_-1);
    Mat61 dxi;
    for (uint_t t = 1 ; t < numberPoses_-1; ++t)
    {
        dxi = tau * t * xiFinal;
        trajectory_->at(t) = SE3(dxi);
    }
    calculate_error();// planes get recalculated, which is a requisite for later
    //(this class construction, in general grad should be self-contained...)
}


#include <anie/details/kernels.hpp>

#include <boost/compute.hpp>

namespace anie::details
{
	const std::string kernel_matrix_add = BOOST_COMPUTE_STRINGIZE_SOURCE(
		__kernel void matrix_add(__global double* lhs, __global const double* rhs,
								 const uint width)
		{
			const uint y = get_global_id(0);

			for (uint i = 0; i < width; ++i)
			{
				lhs[width * y + i] += rhs[width * y + i];
			}
		}
	);
	const std::string kernel_matrix_sub = BOOST_COMPUTE_STRINGIZE_SOURCE(
		__kernel void matrix_sub(__global double* lhs, __global const double* rhs,
								 const uint width)
		{
			const uint y = get_global_id(0);

			for (uint i = 0; i < width; ++i)
			{
				lhs[width * y + i] -= rhs[width * y + i];
			}
		}
	);
	const std::string kernel_matrix_multiply = BOOST_COMPUTE_STRINGIZE_SOURCE(
		__kernel void matrix_multiply(__global double* dest, __global const double* src_lhs, __global const double* src_rhs,
									  const uint src_lhs_width, const uint src_rhs_width)
		{
			const uint width = get_global_size(0);
			const uint x = get_global_id(0);
			const uint y = get_global_id(1);
			const uint index = y * width + x;

			double sum = 0.;

			for (uint i = 0; i < src_lhs_width; ++i)
			{
				sum += src_lhs[src_lhs_width * y + i] * src_rhs[src_rhs_width * i + x];
			}

			dest[index] = sum;
		}
	);
	const std::string kernel_matrix_transpose = BOOST_COMPUTE_STRINGIZE_SOURCE(
		__kernel void matrix_transpose(__global double* dest, __global const double* src,
								 const uint src_width, const uint src_height)
		{
			const uint y = get_global_id(0);

			for (uint i = 0; i < src_width; ++i)
			{
				dest[i * src_height + y] = src[y * src_width + i];
			}
		}
	);
}


/* ----------------------------------------------------------------------------

 * GTSAM Copyright 2010, Georgia Tech Research Corporation,
 * Atlanta, Georgia 30332-0415
 * All Rights Reserved
 * Authors: Frank Dellaert, et al. (see THANKS for the full author list)

 * See LICENSE for the license information

 * -------------------------------------------------------------------------- */

/**
 * @file  Pose3.cpp
 * @brief 3D Pose
 */

#include <gtsam/geometry/Pose3.h>
#include <gtsam/geometry/Pose2.h>
#include <gtsam/geometry/concepts.h>
#include <gtsam/base/Lie-inl.h>
#include <boost/foreach.hpp>
#include <iostream>
#include <cmath>

using namespace std;

namespace gtsam {

  /** Explicit instantiation of base class to export members */
  INSTANTIATE_LIE(Pose3);

  /** instantiate concept checks */
  GTSAM_CONCEPT_POSE_INST(Pose3);

  static const Matrix3 I3 = eye(3), Z3 = zeros(3, 3), _I3=-I3;
  static const Matrix6 I6 = eye(6);

  /* ************************************************************************* */
  Pose3::Pose3(const Pose2& pose2) :
      R_(Rot3::rodriguez(0, 0, pose2.theta())),
      t_(Point3(pose2.x(), pose2.y(), 0)) {
  }

  /* ************************************************************************* */
  // Calculate Adjoint map
  // Ad_pose is 6*6 matrix that when applied to twist xi, returns Ad_pose(xi)
  // Experimental - unit tests of derivatives based on it do not check out yet
  Matrix6 Pose3::AdjointMap() const {
    const Matrix3 R = R_.matrix();
    const Vector3 t = t_.vector();
    Matrix3 A = skewSymmetric(t)*R;
    Matrix6 adj;
    adj << R, Z3, A, R;
    return adj;
  }

  /* ************************************************************************* */
  Matrix6 Pose3::adjointMap(const Vector& xi) {
    Matrix3 w_hat = skewSymmetric(xi(0), xi(1), xi(2));
    Matrix3 v_hat = skewSymmetric(xi(3), xi(4), xi(5));
    Matrix6 adj;
    adj << w_hat, Z3, v_hat, w_hat;

    return adj;
  }

  /* ************************************************************************* */
  Vector Pose3::adjoint(const Vector& xi, const Vector& y, boost::optional<Matrix&> H) {
    if (H) {
      *H = zeros(6,6);
      for (int i = 0; i<6; ++i) {
        Vector dxi = zero(6); dxi(i) = 1.0;
        Matrix Gi = adjointMap(dxi);
        (*H).col(i) = Gi*y;
      }
    }
    return adjointMap(xi)*y;
  }

  /* ************************************************************************* */
  Vector Pose3::adjointTranspose(const Vector& xi, const Vector& y, boost::optional<Matrix&> H) {
    if (H) {
      *H = zeros(6,6);
      for (int i = 0; i<6; ++i) {
        Vector dxi = zero(6); dxi(i) = 1.0;
        Matrix GTi = adjointMap(dxi).transpose();
        (*H).col(i) = GTi*y;
      }
    }
    Matrix adjT = adjointMap(xi).transpose();
    return adjointMap(xi).transpose() * y;
  }

  /* ************************************************************************* */
  Matrix6 Pose3::dExpInv_exp(const Vector& xi) {
    // Bernoulli numbers, from Wikipedia
    static const Vector B = Vector_(9, 1.0, -1.0/2.0, 1./6., 0.0, -1.0/30.0, 0.0, 1.0/42.0, 0.0, -1.0/30);
    static const int N = 5; // order of approximation
    Matrix res = I6;
    Matrix6 ad_i = I6;
    Matrix6 ad_xi = adjointMap(xi);
    double fac = 1.0;
    for (int i = 1 ; i<N; ++i) {
      ad_i = ad_xi * ad_i;
      fac = fac*i;
      res = res + B(i)/fac*ad_i;
    }
    return res;
  }

  /* ************************************************************************* */
  void Pose3::print(const string& s) const {
    cout << s;
    R_.print("R:\n");
    t_.print("t: ");
  }

  /* ************************************************************************* */
  bool Pose3::equals(const Pose3& pose, double tol) const {
    return R_.equals(pose.R_,tol) && t_.equals(pose.t_,tol);
  }

  /* ************************************************************************* */
  /** Modified from Murray94book version (which assumes w and v normalized?) */
  Pose3 Pose3::Expmap(const Vector& xi) {

    // get angular velocity omega and translational velocity v from twist xi
    Point3 w(xi(0),xi(1),xi(2)), v(xi(3),xi(4),xi(5));

    double theta = w.norm();
    if (theta < 1e-10) {
      static const Rot3 I;
      return Pose3(I, v);
    }
    else {
      Point3 n(w/theta); // axis unit vector
      Rot3 R = Rot3::rodriguez(n.vector(),theta);
      double vn = n.dot(v); // translation parallel to n
      Point3 n_cross_v = n.cross(v); // points towards axis
      Point3 t = (n_cross_v - R*n_cross_v)/theta + vn*n;
      return Pose3(R, t);
    }
  }

  /* ************************************************************************* */
  Vector6 Pose3::Logmap(const Pose3& p) {
    Vector3 w = Rot3::Logmap(p.rotation()), T = p.translation().vector();
    double t = w.norm();
    if (t < 1e-10) {
      Vector6 log;
      log << w, T;
      return log;
    }
    else {
      Matrix3 W = skewSymmetric(w/t);
      // Formula from Agrawal06iros, equation (14)
      // simplified with Mathematica, and multiplying in T to avoid matrix math
      double Tan = tan(0.5*t);
      Vector3 WT = W*T;
      Vector3 u = T - (0.5*t)*WT + (1 - t/(2.*Tan)) * (W * WT);
      Vector6 log;
      log << w, u;
      return log;
    }
  }

  /* ************************************************************************* */
  Pose3 Pose3::retractFirstOrder(const Vector& xi) const {
      Vector3 omega(sub(xi, 0, 3));
      Point3 v(sub(xi, 3, 6));
      Rot3 R = R_.retract(omega);  // R is done exactly
      Point3 t = t_ + R_ * v; // First order t approximation
      return Pose3(R, t);
  }

  /* ************************************************************************* */
  // Different versions of retract
  Pose3 Pose3::retract(const Vector& xi, Pose3::CoordinatesMode mode) const {
    if(mode == Pose3::EXPMAP) {
      // Lie group exponential map, traces out geodesic
      return compose(Expmap(xi));
    } else if(mode == Pose3::FIRST_ORDER) {
      // First order
      return retractFirstOrder(xi);
    } else {
      // Point3 t = t_.retract(v.vector()); // Incorrect version retracts t independently
      // Point3 t = t_ + R_ * (v+Point3(omega).cross(v)/2); // Second order t approximation
      assert(false);
      exit(1);
    }
  }

  /* ************************************************************************* */
  // different versions of localCoordinates
  Vector6 Pose3::localCoordinates(const Pose3& T, Pose3::CoordinatesMode mode) const {
    if(mode == Pose3::EXPMAP) {
      // Lie group logarithm map, exact inverse of exponential map
      return Logmap(between(T));
    } else if(mode == Pose3::FIRST_ORDER) {
      // R is always done exactly in all three retract versions below
      Vector3 omega = R_.localCoordinates(T.rotation());

      // Incorrect version
      // Independently computes the logmap of the translation and rotation
      // Vector v = t_.localCoordinates(T.translation());

      // Correct first order t inverse
      Point3 d = R_.unrotate(T.translation() - t_);

      // TODO: correct second order t inverse
      Vector6 local;
      local << omega(0),omega(1),omega(2),d.x(),d.y(),d.z();
      return local;
    } else {
      assert(false);
      exit(1);
    }
  }

  /* ************************************************************************* */
  Matrix4 Pose3::matrix() const {
    const Matrix3 R = R_.matrix();
    const Vector3 T = t_.vector();
    Eigen::Matrix<double,1,4> A14;
    A14 << 0.0, 0.0, 0.0, 1.0;
    Matrix4 mat;
    mat << R, T, A14;
    return mat;
  }

  /* ************************************************************************* */
  Pose3 Pose3::transform_to(const Pose3& pose) const {
    Rot3 cRv = R_ * Rot3(pose.R_.inverse());
    Point3 t = pose.transform_to(t_);
    return Pose3(cRv, t);
  }

  /* ************************************************************************* */
  Point3 Pose3::transform_from(const Point3& p,
      boost::optional<Matrix&> H1, boost::optional<Matrix&> H2) const {
    if (H1) {
      const Matrix R = R_.matrix();
      Matrix DR = R*skewSymmetric(-p.x(), -p.y(), -p.z());
      H1->resize(3,6);
      (*H1) << DR, R;
    }
    if (H2) *H2 = R_.matrix();
    return R_ * p + t_;
  }

  /* ************************************************************************* */
  Point3 Pose3::transform_to(const Point3& p,
            boost::optional<Matrix&> H1, boost::optional<Matrix&> H2) const {
    const Point3 result = R_.unrotate(p - t_);
    if (H1) {
      const Point3& q = result;
      Matrix DR = skewSymmetric(q.x(), q.y(), q.z());
      H1->resize(3,6);
      (*H1) << DR, _I3;
    }
    if (H2) *H2 = R_.transpose();
    return result;
  }

  /* ************************************************************************* */
  Pose3 Pose3::compose(const Pose3& p2,
        boost::optional<Matrix&> H1, boost::optional<Matrix&> H2) const {
    if (H1) *H1 = p2.inverse().AdjointMap();
    if (H2) *H2 = I6;
    return (*this) * p2;
  }

  /* ************************************************************************* */
  Pose3 Pose3::inverse(boost::optional<Matrix&> H1) const {
    if (H1) *H1 = -AdjointMap();
    Rot3 Rt = R_.inverse();
    return Pose3(Rt, Rt*(-t_));
  }

  /* ************************************************************************* */
  // between = compose(p2,inverse(p1));
  Pose3 Pose3::between(const Pose3& p2, boost::optional<Matrix&> H1,
      boost::optional<Matrix&> H2) const {
    Pose3 result = inverse()*p2;
    if (H1) *H1 = -result.inverse().AdjointMap();
    if (H2) *H2 = I6;
    return result;
  }

  /* ************************************************************************* */
  double Pose3::range(const Point3& point,
      boost::optional<Matrix&> H1,
      boost::optional<Matrix&> H2) const {
    if (!H1 && !H2) return transform_to(point).norm();
    Point3 d = transform_to(point, H1, H2);
    double x = d.x(), y = d.y(), z = d.z(),
       d2 = x * x + y * y + z * z, n = sqrt(d2);
    Matrix D_result_d = Matrix_(1, 3, x / n, y / n, z / n);
    if (H1) *H1 = D_result_d * (*H1);
    if (H2) *H2 = D_result_d * (*H2);
    return n;
  }

  /* ************************************************************************* */
  double Pose3::range(const Pose3& point,
        boost::optional<Matrix&> H1, boost::optional<Matrix&> H2) const {
     double r = range(point.translation(), H1, H2);
     if (H2) {
       Matrix H2_ = *H2 * point.rotation().matrix();
       *H2 = zeros(1, 6);
       insertSub(*H2, H2_, 0, 3);
     }
     return r;
  }

  /* ************************************************************************* */
  boost::optional<Pose3> align(const vector<Point3Pair>& pairs) {
    const size_t n = pairs.size();
    if (n<3) return boost::none; // we need at least three pairs

    // calculate centroids
    Vector cp = zero(3),cq = zero(3);
    BOOST_FOREACH(const Point3Pair& pair, pairs) {
      cp += pair.first.vector();
      cq += pair.second.vector();
    }
    double f = 1.0/n;
    cp *= f; cq *= f;

    // Add to form H matrix
    Matrix H = zeros(3,3);
    BOOST_FOREACH(const Point3Pair& pair, pairs) {
      Vector dp = pair.first.vector()  - cp;
      Vector dq = pair.second.vector() - cq;
      H += dp * dq.transpose();
    }

    // Compute SVD
    Matrix U,V;
    Vector S;
    svd(H,U,S,V);

    // Recover transform with correction from Eggert97machinevisionandapplications
    Matrix UVtranspose = U * V.transpose();
    Matrix detWeighting = eye(3,3);
    detWeighting(2,2) = UVtranspose.determinant();
    Rot3 R(Matrix(V * detWeighting * U.transpose()));
    Point3 t = Point3(cq) - R * Point3(cp);
    return Pose3(R, t);
  }

  /* ************************************************************************* */
  std::ostream &operator<<(std::ostream &os, const Pose3& pose) {
    os << pose.rotation() << "\n" << pose.translation() << endl;
    return os;
  }

} // namespace gtsam


#ifndef GEOMETRY_UTILS_HPP
#define GEOMETRY_UTILS_HPP
/**
 * @file geometry_utils.hpp
 * @author Rafael Rey (rreyarc@upo.es)
 * @brief A set of geometry utilities functions
 * @version 0.1
 * @date 2021-06-29
 *
 * @copyright Copyright (c) 2021
 *
 */
#include <iostream>
#include <vector>
#include <math.h>
#include <Eigen/Dense>

#include "utils/utils.hpp"
#include "utils/LineOfSight.hpp"
#include "utils/world.hpp"

namespace Planners
{
    namespace utils
    {
        namespace geometry
        {
            /**
             * @brief Return the integrated distance alongside the _path object with the _resolution given
             *
             * @param _path CoordinateList object of points (discrete)
             * @param _resolution
             * @return float continous path length
             */
            float calculatePathLength(const CoordinateList &_path, const double &_resolution);

            /**
             * @brief Get the Adjacent Path object
             *
             * @param _path
             * @param _algorithm
             * @return utils::CoordinateList
             */
            utils::CoordinateList getAdjacentPath(const utils::CoordinateList &_path, const utils::DiscreteWorld &_world);
            /**
             * @brief Discrete distance multiplied by dist_scale_factor_
             *
             * @param n1
             * @param n2
             * @return unsigned int
             */
            unsigned int distanceBetween2Nodes(const Node &_n1, const Node &_n2);
            /**
             * @brief  Discrete distance multiplied by dist_scale_factor_
             *
             * @param n1
             * @param n2
             * @return unsigned int
             */
            unsigned int distanceBetween2Nodes(const Node *_n1, const Node *_n2);

            /**
             * @brief
             *
             * @param _v1
             * @param _v2
             * @return unsigned int
             */
            unsigned int distanceBetween2Nodes(const Vec3i &_v1, const Vec3i &_v2);
            /**
             * @brief Discrete distance
             *
             * @param n1
             * @param n2
             * @return unsigned int
             */
            unsigned int NodesBetween2Nodes(const Node &_n1, const Node &_n2);
            /**
             * @brief  Discrete distance             *
             * @param n1
             * @param n2
             * @return unsigned int
             */
            unsigned int NodesBetween2Nodes(const Node *_n1, const Node *_n2);

            /**
             * @brief
             *
             * @param _v1
             * @param _v2
             * @return unsigned int
             */
            unsigned int NodesBetween2Nodes(const Vec3i &_v1, const Vec3i &_v2);

            /**
             * @brief Returns the absolute value vector
             *
             * @param _vec
             * @return Vec3i
             */
            Vec3i abs(const Vec3i &_vec);

            int dotProduct(const Vec3i &_v1, const Vec3i &_v2);

            double moduleVector(const Vec3i &_v);

            double angleBetweenThreePoints(const Vec3i &_v1, const Vec3i &_v2, const Vec3i &_v3);

            double angleBetweenThreePoints(const Eigen::Vector3d &_v1, const Eigen::Vector3d &_v2, const Eigen::Vector3d &_v3);

            double getCircunferenceRadius(const Vec3i &_v1, const Vec3i &_v2, const Vec3i &_v3);

            double getCircunferenceRadius(const Eigen::Vector3d &_v1, const Eigen::Vector3d &_v2, const Eigen::Vector3d &_v3);

        }//namespace geometry
    }//namespace utils
}

#endif


#include <NTL/vec_ZZ_pE.h>

NTL_START_IMPL

void InnerProduct(ZZ_pE& x, const vec_ZZ_pE& a, const vec_ZZ_pE& b)
{
   long n = min(a.length(), b.length());
   long i;
   ZZ_pX accum, t;

   clear(accum);
   for (i = 0; i < n; i++) {
      mul(t, rep(a[i]), rep(b[i]));
      add(accum, accum, t);
   }

   conv(x, accum);
}

void InnerProduct(ZZ_pE& x, const vec_ZZ_pE& a, const vec_ZZ_pE& b,
                  long offset)
{
   if (offset < 0) LogicError("InnerProduct: negative offset");
   if (NTL_OVERFLOW(offset, 1, 0)) ResourceError("InnerProduct: offset too big");

   long n = min(a.length(), b.length()+offset);
   long i;
   ZZ_pX accum, t;

   clear(accum);
   for (i = offset; i < n; i++) {
      mul(t, rep(a[i]), rep(b[i-offset]));
      add(accum, accum, t);
   }

   conv(x, accum);
}

void mul(vec_ZZ_pE& x, const vec_ZZ_pE& a, const ZZ_pE& b_in)
{
   ZZ_pE b = b_in;
   long n = a.length();
   x.SetLength(n);
   long i;
   for (i = 0; i < n; i++)
      mul(x[i], a[i], b);
}

void mul(vec_ZZ_pE& x, const vec_ZZ_pE& a, const ZZ_p& b_in)
{
   NTL_ZZ_pRegister(b);
   b = b_in;
   long n = a.length();
   x.SetLength(n);
   long i;
   for (i = 0; i < n; i++)
      mul(x[i], a[i], b);
}

void mul(vec_ZZ_pE& x, const vec_ZZ_pE& a, long b_in)
{
   NTL_ZZ_pRegister(b);
   b = b_in;
   long n = a.length();
   x.SetLength(n);
   long i;
   for (i = 0; i < n; i++)
      mul(x[i], a[i], b);
}

void add(vec_ZZ_pE& x, const vec_ZZ_pE& a, const vec_ZZ_pE& b)
{
   long n = a.length();
   if (b.length() != n) LogicError("vector add: dimension mismatch");

   x.SetLength(n);
   long i;
   for (i = 0; i < n; i++)
      add(x[i], a[i], b[i]);
}

void sub(vec_ZZ_pE& x, const vec_ZZ_pE& a, const vec_ZZ_pE& b)
{
   long n = a.length();
   if (b.length() != n) LogicError("vector sub: dimension mismatch");

   x.SetLength(n);
   long i;
   for (i = 0; i < n; i++)
      sub(x[i], a[i], b[i]);
}

void negate(vec_ZZ_pE& x, const vec_ZZ_pE& a)
{
   long n = a.length();

   x.SetLength(n);
   long i;
   for (i = 0; i < n; i++)
      negate(x[i], a[i]);
}

void clear(vec_ZZ_pE& x)
{
   long n = x.length();
   long i;
   for (i = 0; i < n; i++)
      clear(x[i]);
}

long IsZero(const vec_ZZ_pE& a)
{
   long n = a.length();
   long i;

   for (i = 0; i < n; i++)
      if (!IsZero(a[i]))
         return 0;

   return 1;
}

vec_ZZ_pE operator+(const vec_ZZ_pE& a, const vec_ZZ_pE& b)
{
   vec_ZZ_pE res;
   add(res, a, b);
   NTL_OPT_RETURN(vec_ZZ_pE, res);
}

vec_ZZ_pE operator-(const vec_ZZ_pE& a, const vec_ZZ_pE& b)
{
   vec_ZZ_pE res;
   sub(res, a, b);
   NTL_OPT_RETURN(vec_ZZ_pE, res);
}

vec_ZZ_pE operator-(const vec_ZZ_pE& a)
{
   vec_ZZ_pE res;
   negate(res, a);
   NTL_OPT_RETURN(vec_ZZ_pE, res);
}

ZZ_pE operator*(const vec_ZZ_pE& a, const vec_ZZ_pE& b)
{
   ZZ_pE res;
   InnerProduct(res, a, b);
   return res;
}

void VectorCopy(vec_ZZ_pE& x, const vec_ZZ_pE& a, long n)
{
   if (n < 0) LogicError("VectorCopy: negative length");
   if (NTL_OVERFLOW(n, 1, 0)) ResourceError("overflow in VectorCopy");

   long m = min(n, a.length());

   x.SetLength(n);

   long i;

   for (i = 0; i < m; i++)
      x[i] = a[i];

   for (i = m; i < n; i++)
      clear(x[i]);
}

void random(vec_ZZ_pE& x, long n)
{
   x.SetLength(n);
   for (long i = 0; i < n; i++) random(x[i]);
}

NTL_END_IMPL