// -*- mode: C++; c-indent-level: 4; c-basic-offset: 4; indent-tabs-mode: nil; -*-
//   Learn more about Rcpp at:
//
//   http://www.rcpp.org/
//   http://adv-r.had.co.nz/Rcpp.html
//   http://gallery.rcpp.org/
//

// we only include RcppArmadillo.h which pulls Rcpp.h in for us
#include "RcppArmadillo.h"
using namespace Rcpp;

// via the depends attribute we tell Rcpp to create hooks for
// RcppArmadillo so that the build process will know what to do
//
// [[Rcpp::depends(RcppArmadillo)]]
// [[Rcpp::depends(BH)]]

// one include file from Boost to access the digamma function
#include <boost/math/special_functions/digamma.hpp>

/*----------------------------------------------------------------------------*/
// SlalomModel class definition and module //

//' SlalomModel C++ class
//'
//' @description
//' A C++ class for SlalomModel models.
//'
//' @param Y_init matrix of expression values
//' @param pi_init G x K matrix with each entry being the prior
//' probability for a gene g being active for factor k.
//' @param X_init matrix of initial factor states (N x K)
//' @param W_init G x K matrix of initial weights
//' @param prior_alpha numeric vector of length two giving prior values
//' for the gamma hyperparameters of the precisions
//' @param prior_epsilon numeric vector of length two giving prior values
//' for the gamma hyperparameters of the residual variances
//'
//' @return
//' an object of the SlalomModel class
//'
//' @useDynLib slalom, .registration=TRUE, .fixes="Rcpp_"
//' @importFrom Rcpp evalCpp
//' @exportClass Rcpp_SlalomModel
//' @name SlalomModel
//' @aliases SlalomModel
class SlalomModel {
    // base class for slalom holding models
public:
    // declare necessary variables for alpha
    double alpha_pa;
    double alpha_pb;
    arma::vec alpha_a;
    arma::vec alpha_b;
    arma::vec alpha_E1;
    arma::vec alpha_lnE1;
    // declare necessary variables for epsilon
    double epsilon_pa;
    double epsilon_pb;
    arma::vec epsilon_a;
    arma::vec epsilon_b;
    arma::vec epsilon_E1;
    arma::vec epsilon_lnE1;
    arma::vec epsilon_diagSigmaS;
    // declare necessary variables for X
    arma::mat X_E1;
    arma::mat X_diagSigmaS;
    arma::mat X_init;
    // declare necessary variables for W
    arma::mat W_E1;
    arma::mat W_sigma2;
    arma::mat W_E2diag;
    arma::mat W_gamma0;
    arma::mat W_gamma1;
    // declare necessary variables for Z
    arma::mat Z_E1;
    // declare other variables
    arma::mat I;
    arma::mat Pi_a;
    arma::mat Pi_b;
    arma::mat Pi_pa;
    arma::mat Pi_E1;
    arma::mat Y;
    arma::mat pseudo_Y;
    arma::vec YY;
    arma::mat Known;
    int K;
    int N;
    int G;
    int nAnnotated;
    int nHidden;
    int nKnown;
    int nIterations;
    int minIterations;
    int iterationCount;
    double tolerance;
    bool forceIterations;
    bool shuffle;
    bool converged;
    double nScale;
    std::vector< std::string > noiseModel;
    double onF;
    arma::vec nOn;
    arma::vec iUnannotatedDense;
    arma::vec iUnannotatedSparse;
    arma::uvec doUpdate;
    arma::uvec pretrain_order;
    bool dropFactors;
    bool learnPi;
    // names
    Rcpp::StringVector termNames;
    Rcpp::StringVector cellNames;
    Rcpp::StringVector geneNames;
    // methods
    void train(void);
    void update(void);
    void updateEpsilon(void);
    void updateAlpha(const int);
    void updateX(const int);
    void updateW(const int);
    void updatePi(const int);
    // constructor
    SlalomModel() {}
    SlalomModel(arma::mat Y_init, arma::mat pi_init, arma::mat X_init,
                arma::mat W_init, arma::vec prior_alpha, arma::vec prior_epsilon) :
        Y(Y_init), Pi_E1(pi_init), I(pi_init), W_gamma0(pi_init), X_E1(X_init),
        W_E1(W_init) {
        K = pi_init.n_cols;
        N = Y_init.n_rows;
        G = Y_init.n_cols;
        nScale = 100.0;
        // initialise squared column sums
        YY = arma::zeros(G);
        for (int g=0; g < G; g++) {
            YY(g) = arma::accu(Y.col(g) % Y.col(g));
        }
        // initialise alpha variables
        alpha_pa = prior_alpha[0];
        alpha_pb = prior_alpha[1];
        alpha_a = arma::ones(K) * alpha_pa;
        alpha_b = arma::ones(K) * alpha_pb;
        alpha_E1 = alpha_b / alpha_a;
        alpha_lnE1 = arma::ones(K);
        for (int i = 0; i < K; i++) {
            alpha_lnE1(i) = boost::math::digamma(alpha_a(i));
        }
        alpha_lnE1 = alpha_lnE1 - arma::log(alpha_b);
        // initialise epsilon variables
        epsilon_pa = prior_epsilon[0];
        epsilon_pb = prior_epsilon[1];
        epsilon_a = arma::ones(G) * epsilon_pa;
        epsilon_b = arma::ones(G) * epsilon_pb;
        epsilon_E1 = epsilon_b / epsilon_a;
        epsilon_lnE1 = arma::ones(G);
        for (int i = 0; i < G; i++) {
            epsilon_lnE1(i) = boost::math::digamma(epsilon_a(i));
        }
        epsilon_lnE1 = epsilon_lnE1 - arma::log(epsilon_b);
        epsilon_diagSigmaS = arma::zeros(G);
        // initialise X variables
        X_diagSigmaS = arma::ones(N, K);
        // initialise W variables
        W_gamma1 = 1.0 - W_gamma0;
        W_sigma2 = arma::ones(G, K);
        W_E2diag = arma::zeros(G, K);
        // initialise iterations
        converged = false;
        iterationCount = 0;
        tolerance = 1e-08;
        learnPi = true;
        doUpdate = arma::ones<arma::uvec>(K);
        minIterations = 1;
        nIterations = 1;
    }
};

/*----------------------------------------------------------------------------*/
// SlalomModel class train method

// train method
void SlalomModel::train(void) {
    /*
     Iterate updates of weights (with spike-and-slab prior), ARD parameters, factors, annd noise parameters.

     No arguments, but utilises the following parameters defined in the object:
     nIterations          (int): Number of iterations.
     forceIterations     (bool): Force the model to update `nIteration` times.
     tolerance          (float): Tolerance to monitor convergence of reconstruction error
     minIterations        (int): Minimum number of iterations the model should perform.
     */
    arma::umat Ion = (this->W_gamma0 > .5);
    arma::mat Zr = this->X_E1 * (this->W_E1.t() % Ion.t());
    arma::mat Zd = this->Z_E1 - Zr;
    arma::rowvec error(1);
    arma::rowvec error_old(1);
    error.fill(arma::mean(arma::mean((arma::abs(Zd)))));
    double meanerr = arma::mean(error);
    // mean absolute error
    error_old.fill(100);
    this->converged = false;
    // iterate model to train it
    this->iterationCount = 0;
    for (int iter = 0; iter <= this->nIterations; iter++) {
        // t = time.time();
        this->update();
        this->iterationCount++;
        if (iter % 100 == 0) {
            Rcout << "iteration " << iter << std::endl;
        }
        if (iter % 50 == 0) {
            error_old = error;
            Zr = this->X_E1 * this->W_E1.t();
            Zd = this->Z_E1 - Zr;
            error.fill(arma::mean(arma::mean((arma::abs(Zd)))));
            // mean absolute error
            this->converged = arma::approx_equal(error_old, error, "absdiff",
                                           this->tolerance);
            double meanerr = arma::mean(error);
        }
        if ( this->converged && !(this->forceIterations) &&
             (iter > this->minIterations) ) {
            Rcout << "Model converged after " << iter << " iterations." << std::endl;
            break;
        }
        // this->Z_E1 = this->X_E1 * this->W_E1.t();
    }
    if (!converged) {
        Rcout << "Model not converged after " <<  this->nIterations << " iterations." << std::endl;
    }
}

/*----------------------------------------------------------------------------*/
// SlalomModel class update methods //

// wrapper around R's RNG such that we get a uniform distribution over
// [0,n) as required by the STL algorithm
inline int randWrapper(const int n) { return floor(unif_rand() * n); }

void SlalomModel::update(void) {
    /*
     Do one update of weights (with spike-and-slab prior), ARD parameters, factors, annd noise parameters.

     */
    this->epsilon_diagSigmaS = arma::zeros(this->K);
    // arma::umat Ion = (this->W_gamma0 > .5);
    // check above: set parameter to zero with every update?
    arma::uvec kRange;
    if (this->iterationCount == 1) {
        kRange = this->pretrain_order;
    } else {
        kRange = arma::regspace<arma::uvec>(0,  (this->K - 1));
    }
    if (this->shuffle == true && this->iterationCount > 1) {
        arma::uvec kunfix = arma::regspace<arma::uvec>(this->nKnown,
                                                       (this->K - 1));
        std::random_shuffle(kunfix.begin(), kunfix.end(), randWrapper);
        for (int i = 0; i < (this->K - this->nKnown); ++i) {
            int ii = i + this->nKnown;
            kRange(ii) = kunfix(i);
        }
    }
    // switch factors off such that they can't be turned back on
    // Regular iterator, non-C++11
    for (int m = 0; m < this->K; m++) {
        int k = kRange[m];
        if (this->doUpdate[k]) {
            if (this->dropFactors == false || this->iterationCount < 10 ||
                (this->alpha_E1[k] / arma::var(this->X_E1.col(k))) < 1e10) {
                this->updateW(k);
                if (this->learnPi) {
                    if (arma::any((this->iUnannotatedSparse - 1) == k)) {
                        this->updatePi(k);
                    }
                }
                this->updateAlpha(k);
                this->updateX(k);
            } else {
                this->doUpdate[k] = 0;
                Rcout << "Switched off factor " << k << std::endl;
            }
        }
    }
    // const char noise_gauss = "gauss";
    // if (std::strcmp(this->noiseModel, noise_gauss) == 0) {
    //     // this->updateEpsilon();
    // };
    this->updateEpsilon();
}

void SlalomModel::updateW(const int k) {
    // update the factor weights
    // define logPi values
    arma::vec logPi;
    int Muse = arma::accu(this->doUpdate);
    if (k < this->nKnown || arma::any((this->iUnannotatedSparse - 1) == k) ||
        arma::any((this->iUnannotatedDense - 1) == k)) {
        logPi = arma::log(this->Pi_E1.col(k) / (1.0 - this->Pi_E1.col(k)));
        // careful of divide-by-zero errors here
    } else if (this->nScale > 0 && this->nScale < this->N) {
        logPi = arma::log(this->Pi_E1.col(k) / (1.0 - this->Pi_E1.col(k)));
        // careful of divide-by-zero errors here
        arma::uvec isOFF_ = arma::find(this->Pi_E1.col(k) < 0.5);
        arma::uvec kvec(1);
        kvec(0) = k;
        logPi(isOFF_) = ((this->N / this->nScale) *
            arma::log(this->Pi_E1(isOFF_, kvec) /
                (1 - this->Pi_E1(isOFF_, kvec))));
        // careful of divide-by-zero errors here
        arma::uvec isON_ = arma::find(this->Pi_E1.col(k) > 0.5);
        if (this->onF > 1.0) {
            logPi(isON_) = this->onF * arma::log(
                this->Pi_E1(isON_, kvec) / (1 - this->Pi_E1(isON_, kvec)));
        }
    } else {
        logPi = arma::log(this->Pi_E1.col(k) / (1 - this->Pi_E1.col(k)));
    }
    arma::vec sigma2Sigmaw;
    sigma2Sigmaw = (1.0 / this->epsilon_E1) * this->alpha_E1(k);
    arma::uvec set1 = arma::regspace<arma::uvec>(0, 1, k - 1);    // zero-indexing
    arma::uvec set2 = arma::regspace<arma::uvec>(k + 1, 1, this->K - 1);
    arma::uvec setMinus = arma::join_cols(set1, set2);
    arma::uvec idx = arma::find(this->doUpdate(setMinus) == 1);
    setMinus = setMinus(idx);

    arma::vec SmTSk = arma::sum(
        (arma::repmat(this->X_E1.col(k), 1, Muse - 1) %
             this->X_E1.cols(setMinus)), 0).t();
    double tmp = arma::as_scalar(this->X_E1.col(k).t() * this->X_E1.col(k));
    double SmTSm = (tmp + arma::sum(this->X_diagSigmaS.col(k)));

    arma::vec b;
    arma::vec diff;
    b = (this->W_gamma0.cols(setMinus) % this->W_E1.cols(setMinus)) * SmTSk;
    diff = (this->X_E1.col(k).t() * this->Z_E1).t() - b;
    arma::vec diff2 = diff % diff;
    arma::vec SmTSmSig = SmTSm + sigma2Sigmaw; // same as Py

    // update gamma and W
    arma::vec u_qm = (logPi + 0.5 * arma::log(sigma2Sigmaw) - 0.5 *
        arma::log(SmTSmSig) + (0.5 * this->epsilon_E1) %
        (diff2 / SmTSmSig));
    this->W_gamma0.col(k) = 1.0 / (1 + arma::exp(-u_qm));
    this->W_gamma1.col(k) = 1.0 - this->W_gamma0.col(k);
    this->W_E1.col(k) = (diff / SmTSmSig);
    this->W_sigma2.col(k) = (1.0 / this->epsilon_E1) / SmTSmSig;
    this->W_E2diag.col(k) = ((this->W_E1.col(k) % this->W_E1.col(k)) +
        this->W_sigma2.col(k));
}

void SlalomModel::updatePi(const int k) {
    // update Pi in a SlalomModel class object
    this->Pi_a.col(k) = this->Pi_pa.col(k) + arma::accu(this->W_gamma0.col(k));
    this->Pi_b.col(k) = (this->Pi_pa.col(k) + this->G -
        arma::accu(this->W_gamma0.col(k)));
    this->Pi_E1.col(k) = this->Pi_a.col(k) / this->Pi_a.col(k) +
        this->Pi_b.col(k);
}

void SlalomModel::updateAlpha(const int k) {
    // pdate alpha in a SlalomModel class object - precisions
    double Ewdwd = arma::accu(this->W_gamma0.col(k) % this->W_E2diag.col(k));
    // elementwise mult.
    this->alpha_a(k) = this->alpha_pa + 0.5 * Ewdwd;
    this->alpha_b(k) = this->alpha_pb + arma::accu(this->W_gamma0.col(k)) / 2.0;
    this->alpha_E1(k) = this->alpha_b(k) / this->alpha_a(k);
}

void SlalomModel::updateX(const int k) {
    // update the factor states
    arma::uvec set1 = arma::regspace<arma::uvec>(0, 1, k - 1);
    arma::uvec set2 = arma::regspace<arma::uvec>(k + 1, 1, this->K - 1);
    arma::uvec setMinus = arma::join_cols(set1, set2);
    arma::uvec idx = arma::find(this->doUpdate(setMinus) == 1);
    setMinus = setMinus(idx);

    arma::vec SW_sigma;
    arma::vec SW2_sigma;

    SW2_sigma = ((this->W_gamma0.col(k) % this->W_E2diag.col(k)) %
                     this->epsilon_E1);
    double alphaSm = arma::sum(SW2_sigma);

    for (int i = 0; i < this->N; i++) {
        this->X_diagSigmaS(i, k) = 1.0 / (1.0 + alphaSm);
    }

    if (k >= this->nKnown) {
        SW_sigma = (this->W_gamma0.col(k) %
                        this->W_E1.col(k)) % this->epsilon_E1;     // Gx1 vec
        arma::mat b0 = (this->X_E1.cols(setMinus) *
            (this->W_gamma0.cols(setMinus) % this->W_E1.cols(setMinus)).t());  // NxG mat.
        arma::vec b = b0 * SW_sigma;    // NxG x Gx1 -> Nx1 vec
        arma::vec barmuX = (this->Z_E1 * SW_sigma) - b;
        // update X
        this->X_E1.col(k) = barmuX / (1.0 + alphaSm);
        // keep diagSigmaS
        this->epsilon_diagSigmaS(k) = arma::accu(this->X_diagSigmaS.col(k));
    }
}

void SlalomModel::updateEpsilon(void) {
    // update Epsilon (vectorised) - noise parameters
    arma::uvec update_cols = arma::find(this->doUpdate == 1);
    arma::mat SW_sigma = (this->W_gamma0.cols(update_cols) %
                              this->W_E1.cols(update_cols));  // elementwise mult.; GxK mat
    arma::mat SW2_sigma = (this->W_gamma0.cols(update_cols) %
                               this->W_E2diag.cols(update_cols));  // elementwise mult.; GxK mat
    arma::mat muSTmuS = (this->X_E1.cols(update_cols).t() *
        this->X_E1.cols(update_cols));     // KxK matrix
    arma::mat newmat = this->Z_E1.t() * this->X_E1.cols(update_cols);
    arma::vec t1 = arma::sum(SW_sigma % newmat, 1);  // K length vec
    arma::vec t2 = arma::sum(SW2_sigma % arma::repmat(muSTmuS.diag().t() +
        this->epsilon_diagSigmaS(update_cols).t(), this->G, 1), 1);
    // set diagonals to zeros in muSTmuS for next calculation
    muSTmuS.diag().zeros();
    arma::vec t3 = arma::sum((SW_sigma * muSTmuS) % SW_sigma, 1);
    this->epsilon_E1 = 1.0 / ((0.5 * (this->YY  + (-2 * t1  + t2 + t3))) /
        (0.5 * this->N));
    this->epsilon_a.fill(0.5 * this->N + this->epsilon_pa);
    this->epsilon_b = this->epsilon_pb + 0.5 * (this->YY + (-2 * t1 + t2 + t3));
    for (int i = 0; i < this->epsilon_E1.n_elem; i++) {
        if (this->epsilon_E1(i) > 1e6) {
            this->epsilon_E1(i) = 1e6;
        }
    }
}

/*----------------------------------------------------------------------------*/
// SlalomModel Rcpp module

// Define a module to make the C++ class available to R
RCPP_MODULE(SlalomModel) {
    using namespace Rcpp;

    class_<SlalomModel>("SlalomModel")
        // expose the default constructor
        .constructor()
        .constructor<arma::mat, arma::mat, arma::mat, arma::mat, arma::vec,
        arma::vec>()
        //
        // fields
        .field("K", &SlalomModel::K)
        .field("N", &SlalomModel::N)
        .field("G", &SlalomModel::G)
        .field("nScale", &SlalomModel::nScale)
        .field("nAnnotated", &SlalomModel::nAnnotated)
        .field("nHidden", &SlalomModel::nHidden)
        .field("nKnown", &SlalomModel::nKnown)
        .field("nIterations", &SlalomModel::nIterations)
        .field("minIterations", &SlalomModel::minIterations)
        .field("iterationCount", &SlalomModel::iterationCount)
        .field("forceIterations", &SlalomModel::forceIterations)
        .field("tolerance", &SlalomModel::tolerance)
        .field("shuffle", &SlalomModel::shuffle)
        .field("converged", &SlalomModel::converged)
        .field("noiseModel", &SlalomModel::noiseModel)
        .field("onF", &SlalomModel::onF)
        // alpha
        .field("alpha_pa", &SlalomModel::alpha_pa)
        .field("alpha_pb", &SlalomModel::alpha_pb)
        .field("alpha_a", &SlalomModel::alpha_a)
        .field("alpha_b", &SlalomModel::alpha_b)
        .field("alpha_E1", &SlalomModel::alpha_E1)
        .field("alpha_lnE1", &SlalomModel::alpha_lnE1)
        // epsilon
        .field("epsilon_pa", &SlalomModel::epsilon_pa)
        .field("epsilon_pb", &SlalomModel::epsilon_pb)
        .field("epsilon_a", &SlalomModel::epsilon_a)
        .field("epsilon_b", &SlalomModel::epsilon_b)
        .field("epsilon_E1", &SlalomModel::epsilon_E1)
        .field("epsilon_lnE1", &SlalomModel::epsilon_lnE1)
        .field("epsilon_diagSigmaS", &SlalomModel::epsilon_diagSigmaS)
        // X
        .field("X_E1", &SlalomModel::X_E1)
        .field("X_diagSigmaS", &SlalomModel::X_diagSigmaS)
        .field("X_init", &SlalomModel::X_init)
        // W
        .field("W_E1", &SlalomModel::W_E1)
        .field("W_sigma2", &SlalomModel::W_sigma2)
        .field("W_E2diag", &SlalomModel::W_E2diag)
        .field("W_gamma0", &SlalomModel::W_gamma0)
        .field("W_gamma1", &SlalomModel::W_gamma1)
        // Z
        .field("Z_E1", &SlalomModel::Z_E1)
        // Pi
        .field("Pi_a", &SlalomModel::Pi_a)
        .field("Pi_pa", &SlalomModel::Pi_pa)
        .field("Pi_b", &SlalomModel::Pi_b)
        .field("Pi_E1", &SlalomModel::Pi_E1)
        .field("I", &SlalomModel::I)
        // other variables
        .field("Known", &SlalomModel::Known)
        .field("Y", &SlalomModel::Y)
        .field("pseudo_Y", &SlalomModel::pseudo_Y)
        .field("YY", &SlalomModel::YY)
        .field("iUnannotatedDense", &SlalomModel::iUnannotatedDense)
        .field("iUnannotatedSparse", &SlalomModel::iUnannotatedSparse)
        .field("nOn", &SlalomModel::nOn)
        .field("doUpdate", &SlalomModel::doUpdate)
        .field("pretrain_order", &SlalomModel::pretrain_order)
        .field("learnPi", &SlalomModel::learnPi)
        .field("dropFactors", &SlalomModel::dropFactors)
        // names
        .field("termNames", &SlalomModel::termNames)
        .field("cellNames", &SlalomModel::cellNames)
        .field("geneNames", &SlalomModel::geneNames)
        // methods
        .method("train", &SlalomModel::train , "Train the SlalomModel")
        .method("update", &SlalomModel::update , "Update the SlalomModel")
        .method("updateW", &SlalomModel::updateW , "Update W")
        .method("updateX", &SlalomModel::updateX , "Update X")
        .method("updatePi", &SlalomModel::updatePi , "Update Pi")
        .method("updateEpsilon", &SlalomModel::updateEpsilon , "Update Epsilon")
        .method("updateAlpha", &SlalomModel::updateAlpha , "Update alpha")
        ;
}


//---------------------------------------------------------------------------//
//  MIT License
//
//  Copyright (c) 2020-2021 Mikhail Komarov <nemo@nil.foundation>
//  Copyright (c) 2020-2021 Nikita Kaskov <nemo@nil.foundation>

//
//  Permission is hereby granted, free of charge, to any person obtaining a copy
//  of this software and associated documentation files (the "Software"), to deal
//  in the Software without restriction, including without limitation the rights
//  to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
//  copies of the Software, and to permit persons to whom the Software is
//  furnished to do so, subject to the following conditions:
//
//  The above copyright notice and this permission notice shall be included in all
//  copies or substantial portions of the Software.
//
//  THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
//  IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
//  FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
//  AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
//  LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
//  OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
//  SOFTWARE.
//---------------------------------------------------------------------------//

#define BOOST_TEST_MODULE vanilla_challenges_test

#include <boost/test/data/monomorphic.hpp>
#include <boost/test/data/test_case.hpp>
#include <boost/test/unit_test.hpp>

#include <nil/filecoin/storage/proofs/porep/stacked/vanilla/challenges.hpp>

using namespace nil::filecoin;

BOOST_AUTO_TEST_SUITE(vanilla_challenges_test_suite)

BOOST_AUTO_TEST_CASE(challenge_derivation) {
    const std::size_t n = 200;
    const std::size_t layers = 100;

    const auto challenges = LayerChallenges(layers, n);
    const std::size_t leaves = 1 << 30;
    const auto rng = thread_rng();
    const Pedersen::digest_type replica_id = Pedersen::digest_type::random(rng);
    const auto seed : [u8; 32] = rng.gen();
    const std::size_t partitions = 5;
    const std::size_t total_challenges = partitions * n;

    auto layers_with_duplicates = 0;

    for (_layer in 1.. = layers) {
        auto histogram = HashMap();
        for (std::size_t k = 0; k < partitions; ++k) {
            const auto challenges = challenges.derive(leaves, &replica_id, &seed, k as u8);

            for (challenges::iterator challenge = challenges.begin(); challenge != challenges.end(); ++challenge) {
                const auto counter = histogram.entry(challenge).or_insert(0);
                *counter += 1;
            }
        }

        const auto unique_challenges = histogram.len();

        if (unique_challenges < total_challenges) {
                layers_with_duplicates += 1;
        }
    }

    // If we generate 100 layers with 1,000 challenges in each, at most two layers can contain
    // any duplicates for this assertion to succeed.
    //
    // This test could randomly fail (anything's possible), but if it happens regularly something is wrong.
    BOOST_ASSERT (layers_with_duplicates < 3);
}

// This test shows that partitioning (k = 0..partitions) generates the same challenges as
// generating the same number of challenges with only one partition (k = 0).
BOOST_AUTO_TEST_CASE(challenge_partition_equivalence) {
    const std::size_t n = 40;
    const std::size_t leaves = 1 << 30;
    const auto rng = thread_rng();
    const Pedersen::digest_type replica_id  = Pedersen::digest_type::random(rng);
    const auto seed : [u8; 32] = rng.gen();
    const std::size_t partitions = 5;
    const std::size_t layers = 100;
    const std::size_t total_challenges = n * partitions;

    for (_layer in 1.. = layers) {
        const auto one_partition_challenges =
            LayerChallenges(layers, total_challenges).derive(leaves, &replica_id, &seed, 0, );
        const auto many_partition_challenges =
            (0..partitions)
                .flat_map(| k | {LayerChallenges(layers, n).derive(leaves, &replica_id, &seed, k as u8)})
                .collect::<Vec<_>>();

        BOOST_ASSERT(one_partition_challenges == many_partition_challenges);
    }
}

BOOST_AUTO_TEST_SUITE_END()


#ifndef TCM_BLAS_WRAPPER_HPP
#define TCM_BLAS_WRAPPER_HPP

#include <type_traits>
#include <complex>
#include <cassert>

#include <boost/numeric/conversion/cast.hpp>

#include <detail/config.hpp>
#include <detail/utils.hpp>
#include <detail/iterator.hpp>

///////////////////////////////////////////////////////////////////////////////
/// \file blas_wrapper.hpp
/// \brief Defines template versions of some BLAS routines.
///////////////////////////////////////////////////////////////////////////////

#define REGISTER_BLAS(general_f, s_f, d_f, c_f, z_f)                           \
	namespace {                                                                \
		template<class... Args>                                                \
        __attribute__((always_inline))                                         \
		inline                                                                 \
		auto general_f##_impl( utils::Type2Type<float>                         \
		                     , Args&&... args) noexcept                        \
		{ return s_f(std::forward<Args>(args)...); }                           \
		                                                                       \
		template<class... Args>                                                \
        __attribute__((always_inline))                                         \
		inline                                                                 \
		auto general_f##_impl( utils::Type2Type<double>                        \
		                     , Args&&... args) noexcept                        \
		{ return d_f(std::forward<Args>(args)...); }                           \
		                                                                       \
		template<class... Args>                                                \
        __attribute__((always_inline))                                         \
		inline                                                                 \
		auto general_f##_impl( utils::Type2Type<std::complex<float>>           \
		                     , Args&&... args) noexcept                        \
		{ return c_f(std::forward<Args>(args)...); }                           \
		                                                                       \
		template<class... Args>                                                \
        __attribute__((always_inline))                                         \
		inline                                                                 \
		auto general_f##_impl( utils::Type2Type<std::complex<double>>          \
		                     , Args&&... args) noexcept                        \
		{ return z_f(std::forward<Args>(args)...); }                           \
	}                                                                          \
                                                                               \
	template<class T, class... Args>                                           \
	__attribute__((always_inline))                                             \
	inline                                                                     \
	auto general_f(Args&&... args) noexcept                                    \
	{ return general_f##_impl( utils::Type2Type<T>{}                           \
	                         , std::forward<Args>(args)...);                   \
	}

namespace tcm {

namespace import {

enum class Operator : char {None = 'N', T = 'T', H = 'C'};

#ifdef USING_INTEL_MKL
	using blas_int = long long int;
#else
#	ifdef USING_ATLAS
		using blas_int = int;
#	else
#		error "Need BLAS"
#	endif
#endif

// ============================================================================
//                                DOT PRODUCT
// ============================================================================

extern "C" {

float sdot_
( blas_int const* N
, float const* X, blas_int const* INCX
, float const* Y, blas_int const* INCY );

double ddot_
( blas_int const* N
, double const* X, blas_int const* INCX
, double const* Y, blas_int const* INCY );

std::complex<float> cdotc_
( blas_int const* N
, std::complex<float> const* X, blas_int const* INCX
, std::complex<float> const* Y, blas_int const* INCY );

std::complex<double> zdotc_
( blas_int const* N
, std::complex<double> const* X, blas_int const* inc_X
, std::complex<double> const* Y, blas_int const* inc_Y );

} // extern "C"

REGISTER_BLAS(dotc, sdot_, ddot_, cdotc_, zdotc_)

template< class _T
        , class = std::enable_if_t
                  <   std::is_same<_T, float>()
                   or std::is_same<_T, double>()
                   or std::is_same<_T, std::complex<float>>()
                   or std::is_same<_T, std::complex<double>>()
                  >
        >
inline
auto dot( std::size_t const n
        , _T const* X, blas_int const INCX
        , _T const* Y, blas_int const INCY
        ) -> _T
{
	if(n == 0) return _T{0.0};

	assert(X != nullptr and INCX != 0);
	assert(Y != nullptr and INCY != 0);
	/*
	utils::assert_valid( const_blas_iterator<_T>{X           , INCX}
	                   , const_blas_iterator<_T>{X + n * INCX, INCX} );
	utils::assert_valid( const_blas_iterator<_T>{Y           , INCY}
	                   , const_blas_iterator<_T>{Y + n * INCY, INCY} );
	*/

	auto const N = boost::numeric_cast<blas_int>(n);
	auto const result = dotc<_T>(&N, X, &INCX, Y, &INCY);

	utils::assert_valid(result);
	return result;
}

// ============================================================================
//                                   AXPY
// ============================================================================

extern "C" {

void saxpy_
( blas_int const* N, float const* A
, float const* X, blas_int const* INCX
, float      * Y, blas_int const* INCY );

void daxpy_
( blas_int const* N, double const* A
, double const* X, blas_int const* INCX
, double      * Y, blas_int const* INCY );

void caxpy_
( blas_int const* N, std::complex<float> const* A
, std::complex<float> const* X, blas_int const* INCX
, std::complex<float>      * Y, blas_int const* INCY );

void zaxpy_
( blas_int const* N, std::complex<double> const* A
, std::complex<double> const* X, blas_int const* INCX
, std::complex<double>      * Y, blas_int const* INCY );

} // extern "C"

REGISTER_BLAS(axpy, saxpy_, daxpy_, caxpy_, zaxpy_)

template< class _T
        , class = std::enable_if_t
                  <   std::is_same<_T, float>()
                   or std::is_same<_T, double>()
                   or std::is_same<_T, std::complex<float>>()
                   or std::is_same<_T, std::complex<double>>()
                  >
        >
inline
auto axpy( std::size_t const n, _T const a
         , _T const* X, blas_int const INCX
         , _T      * Y, blas_int const INCY ) -> void
{
	if(n == 0) return;

	assert(X != nullptr and INCX != 0);
	assert(Y != nullptr and INCY != 0);
	/*
	assert_valid( const_blas_iterator<_T>{X           , INCX}
	            , const_blas_iterator<_T>{X + n * INCX, INCX} );
	assert_valid( const_blas_iterator<_T>{Y           , INCY}
	            , const_blas_iterator<_T>{Y + n * INCY, INCY} );
	*/
	auto const N = boost::numeric_cast<blas_int>(n);

	axpy<_T>(&N, &a, X, &INCX, Y, &INCY);

	/*
	assert_valid( const_blas_iterator<_T>{Y           , INCY}
	            , const_blas_iterator<_T>{Y + n * INCY, INCY} );
	*/
}

// ============================================================================
//                         MATRIX-VECTOR MULTIPLICATION
// ============================================================================

extern "C" {

void sgemv_
( char const* TRANSA
, blas_int const* M, blas_int const* N
, float const* ALPHA, float const* A, blas_int const* LDA
, float const* X, blas_int const* INCX
, float const* BETA, float* Y, blas_int const* INCY );

void dgemv_
( char const* TRANSA
, blas_int const* M, blas_int const* N
, double const* ALPHA, double const* A, blas_int const* LDA
, double const* X, blas_int const* INCX
, double const* BETA, double* Y, blas_int const* INCY );

void cgemv_
( char const* TRANSA
, blas_int const* M, blas_int const* N
, std::complex<float> const* ALPHA
, std::complex<float> const* A, blas_int const* LDA
, std::complex<float> const* X, blas_int const* INCX
, std::complex<float> const* BETA
, std::complex<float>* Y, blas_int const* INCY );

void zgemv_
( char const* TRANSA
, blas_int const* M, blas_int const* N
, std::complex<double> const* ALPHA
, std::complex<double> const* A, blas_int const* LDA
, std::complex<double> const* X, blas_int const* INCX
, std::complex<double> const* BETA
, std::complex<double>* Y, blas_int const* INCY );

} // extern "C"

REGISTER_BLAS(gemv, sgemv_, dgemv_, cgemv_, zgemv_)

template< class _T
        , class = std::enable_if_t
                  <   std::is_same<_T, float>()
                   or std::is_same<_T, double>()
                   or std::is_same<_T, std::complex<float>>()
                   or std::is_same<_T, std::complex<double>>()
                  >
        >
auto gemv( Operator const op_A
         , std::size_t const m, std::size_t const n
         , _T const ALPHA, _T const* A, blas_int const LDA
         , _T const* X, blas_int const INCX
         , _T const BETA, _T* Y, blas_int const INCY ) -> void
{
	if(n == 0 or m == 0) return;

	assert(A != nullptr and LDA  != 0);
	assert(X != nullptr and INCX != 0);
	assert(Y != nullptr and INCY != 0);
	/*
	utils::assert_valid(A, A + LDA * n);
	utils::assert_valid
		( const_blas_iterator<_T>{X, INCX}
		, const_blas_iterator<_T>{ X + (op_A == Operator::None ? n : m) * INCX
		                         , INCX }
		);
	utils::assert_valid
		( const_blas_iterator<_T>{Y, INCY}
		, const_blas_iterator<_T>{ Y + (op_A == Operator::None ? m : n) * INCY
		                         , INCY }
		);
	*/

	auto const TRANS = static_cast<char>(op_A);
	auto const M     = boost::numeric_cast<blas_int>(m);
	auto const N     = boost::numeric_cast<blas_int>(n);

	assert(LDA >= M);

	gemv<_T>( &TRANS, &M, &N
	        , &ALPHA, A, &LDA, X, &INCX
	        , &BETA, Y, &INCY );

	/*
	utils::assert_valid
		( const_blas_iterator<_T>{Y, INCY}
		, const_blas_iterator<_T>{ Y + (op_A == Operator::None ? m : n) * INCY
		                         , INCY }
		);
	*/
}

// ============================================================================
//                           MATRIX-MATRIX MULTIPLICATION
// ============================================================================

extern "C" {

void sgemm_
( char const* TRANSA, char const* TRANSB
, blas_int const* M, blas_int const* N, blas_int const* K
, float const* ALPHA
, float const* A, blas_int const* LDA
, float const* B, blas_int const* LDB
, float const* BETA
, float* C, blas_int const* LDC );

void dgemm_
( char const* TRANSA, char const* TRANSB
, blas_int const* M, blas_int const* N, blas_int const* K
, double const* ALPHA
, double const* A, blas_int const* LDA
, double const* B, blas_int const* LDB
, double const* BETA
, double* C, blas_int const* LDC );

void cgemm_
( char const* TRANSA, char const* TRANSB
, blas_int const* M, blas_int const* N, blas_int const* K
, std::complex<float> const* ALPHA
, std::complex<float> const* A, blas_int const* LDA
, std::complex<float> const* B, blas_int const* LDB
, std::complex<float> const* BETA
, std::complex<float>* C, blas_int const* LDC );

void zgemm_
( char const* TRANSA, char const* TRANSB
, blas_int const* M, blas_int const* N, blas_int const* K
, std::complex<double> const* ALPHA
, std::complex<double> const* A, blas_int const* LDA
, std::complex<double> const* B, blas_int const* LDB
, std::complex<double> const* BETA
, std::complex<double>* C, blas_int const* LDC );

} // extern "C"

REGISTER_BLAS(gemm, sgemm_, dgemm_, cgemm_, zgemm_)

template< class _T
        , class = std::enable_if_t
                  <   std::is_same<_T, float>()
                   or std::is_same<_T, double>()
                   or std::is_same<_T, std::complex<float>>()
                   or std::is_same<_T, std::complex<double>>()
                  >
        >
inline
auto gemm( Operator const op_A, Operator const op_B
         , std::size_t const m, std::size_t const n, std::size_t const k
         , _T const ALPHA, _T const* A, blas_int const LDA
         , _T const* B, blas_int const LDB
         , _T const BETA, _T* C, blas_int const LDC ) -> void
{
	if(m == 0 or n == 0 or k == 0) return;

	assert(A != nullptr and LDA != 0);
	assert(B != nullptr and LDB != 0);
	assert(C != nullptr and LDC != 0);
	/*
	utils::assert_valid(A, A + LDA * (op_A == Operator::None ? k : m));
	utils::assert_valid(B, B + LDB * (op_B == Operator::None ? n : k));
	utils::assert_valid(C, C + LDC * n);
	*/

	auto const M      = boost::numeric_cast<blas_int>(m);
	auto const N      = boost::numeric_cast<blas_int>(n);
	auto const K      = boost::numeric_cast<blas_int>(k);
	auto const TRANSA = static_cast<char>(op_A);
	auto const TRANSB = static_cast<char>(op_B);

	assert( LDA >= (op_A == Operator::None ? M : K) );
	assert( LDB >= (op_B == Operator::None ? K : N) );
	assert( LDC >= M );

	gemm<_T>( &TRANSA, &TRANSB, &M, &N, &K
	        , &ALPHA, A, &LDA, B, &LDB
	        , &BETA, C, &LDC );

	/*
	utils::assert_valid(C, C + LDC * n);
	*/
}

} // namespace import

} // namespace tcm

#endif // TCM_BLAS_WRAPPER_HPP


/**
 * @file
 * @copyright This code is licensed under the 3-clause BSD license.\n
 *            Copyright ETH Zurich, Laboratory for Physical Chemistry, Reiher Group.\n
 *            See LICENSE.txt for details.
 */

#include <Sparrow/Implementations/Nddo/Utils/IntegralsEvaluationUtils/AtomPairOverlap.h>
#include <Utils/DataStructures/AtomicGtos.h>
#include <Utils/DataStructures/GtoExpansion.h>
#include <Utils/DataStructures/SlaterToGaussian.h>
#include <gmock/gmock.h>
#include <Eigen/Core>
#include <memory>

namespace Scine {
namespace Sparrow {

using namespace testing;
using namespace nddo;

class AAtomPairOverlapMatrix : public Test {
 public:
  AtomPairOverlap<Utils::derivOrder::zero> pairOverlap;
  Utils::GtoExpansion gs1, gs2, gp1, gp2, gd1, gd2;
  Utils::AtomicGtos aS1, aS2, aP1, aP2, aD1, aD2;
  Eigen::Vector3d arbitraryVector;

  void SetUp() override {
    arbitraryVector = Eigen::Vector3d(1.1, 2.0, 0.23);
    gs1 = Utils::SlaterToGaussian::getGTOExpansion(6, 1, 0);
    gs2 = Utils::SlaterToGaussian::getGTOExpansion(6, 4, 0, 1.42);
    gp1 = Utils::SlaterToGaussian::getGTOExpansion(6, 2, 1, 2.2);
    gp2 = Utils::SlaterToGaussian::getGTOExpansion(6, 3, 1);
    gd1 = Utils::SlaterToGaussian::getGTOExpansion(6, 5, 2, 0.92);
    gd2 = Utils::SlaterToGaussian::getGTOExpansion(6, 4, 2);

    aS1.setS(gs1);
    aS2.setS(gs2);
    aP1 = aS1;
    aP1.setP(gp1);
    aP2 = aS2;
    aP2.setP(gp2);
    aD1 = aP1;
    aD1.setD(gd1);
    aD2 = aP2;
    aD2.setD(gd2);
  }
};

TEST_F(AAtomPairOverlapMatrix, ReturnsAtomPairBlockOfCorrectSize) {
  ASSERT_THAT(pairOverlap.getMatrixBlock(aS1, aS2, arbitraryVector).size(), Eq(1));
  ASSERT_THAT(pairOverlap.getMatrixBlock(aD1, aD2, arbitraryVector).size(), Eq(81));
  ASSERT_THAT(pairOverlap.getMatrixBlock(aD1, aP2, arbitraryVector).rows(), Eq(9));
  ASSERT_THAT(pairOverlap.getMatrixBlock(aD1, aP2, arbitraryVector).cols(), Eq(4));
}

TEST_F(AAtomPairOverlapMatrix, ReturnsIdentityForOverlapOfAtomWithItself) {
  Eigen::MatrixXd blockS = pairOverlap.getMatrixBlock(aS1, aS1, Eigen::Vector3d(0, 0, 0));
  Eigen::MatrixXd blockP = pairOverlap.getMatrixBlock(aP1, aP1, Eigen::Vector3d(0, 0, 0));
  Eigen::MatrixXd blockD = pairOverlap.getMatrixBlock(aD1, aD1, Eigen::Vector3d(0, 0, 0));

  ASSERT_TRUE(blockS.isApprox(Eigen::Matrix<double, 1, 1>::Identity(), 1e-7));
  ASSERT_TRUE(blockP.isApprox(Eigen::Matrix<double, 4, 4>::Identity(), 1e-7));
  ASSERT_TRUE(blockD.isApprox(Eigen::Matrix<double, 9, 9>::Identity(), 1e-7));
}
} // namespace Sparrow
} // namespace Scine


/*
 * This file is part of the statismo library.
 *
 * Author: Marcel Luethi (marcel.luethi@unibas.ch)
 *
 * Copyright (c) 2011 University of Basel
 * 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 project's author 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
 * 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 __STATIMO_CORE_REDUCED_VARIANCE_MODEL_BUILDER_HXX_
#define __STATIMO_CORE_REDUCED_VARIANCE_MODEL_BUILDER_HXX_

#include "statismo/core/ReducedVarianceModelBuilder.h"
#include "statismo/core/CommonTypes.h"
#include "statismo/core/Exceptions.h"

#include <Eigen/SVD>

namespace statismo
{

template <typename T>
UniquePtrType<typename ReducedVarianceModelBuilder<T>::StatisticalModelType>
ReducedVarianceModelBuilder<T>::BuildNewModelWithLeadingComponents(const StatisticalModelType * inputModel,
                                                                   unsigned numberOfPrincipalComponents) const

{
  STATISMO_LOG_INFO("Building new model");
  STATISMO_LOG_INFO("Number of principal components: " + std::to_string(numberOfPrincipalComponents));

  auto reducedModel =
    StatisticalModelType::SafeCreate(inputModel->GetRepresenter(),
                                     inputModel->GetMeanVector(),
                                     inputModel->GetOrthonormalPCABasisMatrix().leftCols(numberOfPrincipalComponents),
                                     inputModel->GetPCAVarianceVector().topRows(numberOfPrincipalComponents),
                                     inputModel->GetNoiseVariance());

  // Write the parameters used to build the models into the builderInfo
  typename ModelInfo::BuilderInfoList builderInfoList = inputModel->GetModelInfo().GetBuilderInfoList();

  BuilderInfo::ParameterInfoList bi;
  bi.emplace_back("NumberOfPincipalComponents ", std::to_string(numberOfPrincipalComponents));

  BuilderInfo::DataInfoList di;
  builderInfoList.emplace_back("ReducedVarianceModelBuilder", di, bi);

  // If the scores matrix is not set, or if we have for some reasons not as many score entries as the number of
  // principal components, we simply work with what is there.
  unsigned numComponentsForScores =
    std::min(static_cast<unsigned>(inputModel->GetModelInfo().GetScoresMatrix().rows()), numberOfPrincipalComponents);

  reducedModel->SetModelInfo(
    ModelInfo(inputModel->GetModelInfo().GetScoresMatrix().topRows(numComponentsForScores), builderInfoList));
  return reducedModel;
}

template <typename T>
UniquePtrType<typename ReducedVarianceModelBuilder<T>::StatisticalModelType>
ReducedVarianceModelBuilder<T>::BuildNewModelWithVariance(const StatisticalModelType * inputModel,
                                                          double                       totalVariance) const
{
  STATISMO_LOG_INFO("Building new model");
  STATISMO_LOG_INFO("Total variance: " + std::to_string(totalVariance));

  VectorType pcaVariance = inputModel->GetPCAVarianceVector();
  double     modelVariance = pcaVariance.sum();

  // count the number of modes required for the model
  double   cumulatedVariance = 0;
  unsigned numComponentsToReachPrescribedVariance = 0;
  for (unsigned i = 0; i < pcaVariance.size(); i++)
  {
    cumulatedVariance += pcaVariance(i);
    numComponentsToReachPrescribedVariance++;
    if (cumulatedVariance / modelVariance >= totalVariance)
      break;
  }
  return BuildNewModelWithLeadingComponents(inputModel, numComponentsToReachPrescribedVariance);
}

template <typename T>
UniquePtrType<typename ReducedVarianceModelBuilder<T>::StatisticalModelType>
ReducedVarianceModelBuilder<T>::BuildNewModelFromModel(const StatisticalModelType * inputModel,
                                                       double                       totalVariance) const
{
  STATISMO_LOG_INFO("Building new model from model");
  return BuildNewModelWithVariance(inputModel, totalVariance);
}

} // namespace statismo

#endif


#include <boost/graph/metis.hpp>


#include <cstdio>
#include <cmath>
#include <iostream>
#include <fstream>
#include <sstream>
#include <vector>
#include <map>
#include <cfloat>

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

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

// HDF5
#include <H5Cpp.h>

// OpenMP
#include <omp.h>

// Point-triangle distance and ray-triangle intersection.
#include "triangle_point/poitri.h"
#include "triangle_ray/raytri.h"
#include "box_triangle/aabb_triangle_overlap.h"

/** \brief Compute triangle point distance and corresponding closest point.
 * \param[in] point point
 * \param[in] v1 first vertex
 * \param[in] v2 second vertex
 * \param[in] v3 third vertex
 * \param[out] ray corresponding closest point
 * \return distance
 */
float triangle_point_distance(const Eigen::Vector3f &point, const Eigen::Vector3f &v1, const Eigen::Vector3f &v2, const Eigen::Vector3f &v3,
    Eigen::Vector3f &closest_point) {

  Vec3f x0(point.data());
  Vec3f x1(v1.data());
  Vec3f x2(v2.data());
  Vec3f x3(v3.data());

  Vec3f r(0);
  float distance = point_triangle_distance(x0, x1, x2, x3, r);

  for (int d = 0; d < 3; d++) {
    closest_point(d) = r[d];
  }

  return distance;
}

/** \brief Test triangle ray intersection.
 * \param[in] origin origin of ray
 * \param[in] dest destination of ray
 * \param[in] v1 first vertex
 * \param[in] v2 second vertex
 * \param[in] v3 third vertex
 * \return intersects
 */
bool triangle_ray_intersection(const Eigen::Vector3f &origin, const Eigen::Vector3f &dest,
    const Eigen::Vector3f &v1, const Eigen::Vector3f &v2, const Eigen::Vector3f &v3, float &t) {

  double _origin[3] = {origin(0), origin(1), origin(2)};
  double _dir[3] = {dest(0) - origin(0), dest(1) - origin(1), dest(2) - origin(2)};
  double _v1[3] = {v1(0), v1(1), v1(2)};
  double _v2[3] = {v2(0), v2(1), v2(2)};
  double _v3[3] = {v3(0), v3(1), v3(2)};

  // t is the distance, u and v are barycentric coordinates
  // http://fileadmin.cs.lth.se/cs/personal/tomas_akenine-moller/code/raytri_tam.pdf
  double _t, u, v;
  int success = intersect_triangle(_origin, _dir, _v1, _v2, _v3, &_t, &u, &v);
  t = _t;

  if (success) {
    return true;
  }

  return false;
}

/** \brief Compute triangle box intersection.
 * \param[in] min defining voxel
 * \param[in] max defining voxel
 * \param[in] v1 first vertex
 * \param[in] v2 second vertex
 * \param[in] v3 third vertex
 * \return intersects
 */
bool triangle_box_intersection(const Eigen::Vector3f &min, Eigen::Vector3f &max, const Eigen::Vector3f &v1, const Eigen::Vector3f &v2, const Eigen::Vector3f &v3) {
  float half_size[3] = {
    (max(0) - min(0))/2.,
    (max(1) - min(1))/2.,
    (max(2) - min(2))/2.
  };

  float center[3] = {
    max(0) - half_size[0],
    max(1) - half_size[1],
    max(2) - half_size[2]
  };

  float vertices[3][3] = {{v1(0), v1(1), v1(2)}, {v2(0), v2(1), v2(2)}, {v3(0), v3(1), v3(2)}};
  return triBoxOverlap(center, half_size, vertices);
}

/** \brief Specifies the voxelization mode, i.e. which point of a voxel to use for SDF computation. */
enum VoxelizationMode {
  CENTER = 0,
  CORNER = 1
};

/** \brief Just encapsulating vertices and faces. */
class Mesh {
public:
  /** \brief Empty constructor. */
  Mesh() {

  }

  /** \brief Reading an off file and returning the vertices x, y, z coordinates and the
   * face indices.
   * \param[in] filepath path to the OFF file
   * \param[out] mesh read mesh with vertices and faces
   * \return success
   */
  static bool from_off(const std::string filepath, Mesh& mesh) {

    std::ifstream* file = new std::ifstream(filepath.c_str());
    std::string line;
    std::stringstream ss;
    int line_nb = 0;

    std::getline(*file, line);
    ++line_nb;

    if (line != "off" && line != "OFF") {
      std::cout << "[Error] Invalid header: \"" << line << "\", " << filepath << std::endl;
      return false;
    }

    size_t n_edges;
    std::getline(*file, line);
    ++line_nb;

    int n_vertices;
    int n_faces;
    ss << line;
    ss >> n_vertices;
    ss >> n_faces;
    ss >> n_edges;

    for (size_t v = 0; v < n_vertices; ++v) {
      std::getline(*file, line);
      ++line_nb;

      ss.clear();
      ss.str("");

      Eigen::Vector3f vertex;
      ss << line;
      ss >> vertex(0);
      ss >> vertex(1);
      ss >> vertex(2);

      mesh.add_vertex(vertex);
    }

    size_t n;
    for (size_t f = 0; f < n_faces; ++f) {
      std::getline(*file, line);
      ++line_nb;

      ss.clear();
      ss.str("");

      size_t n;
      ss << line;
      ss >> n;

      if(n != 3) {
        std::cout << "[Error] Not a triangle (" << n << " points) at " << (line_nb - 1) << std::endl;
        return false;
      }

      Eigen::Vector3i face;
      ss >> face(0);
      ss >> face(1);
      ss >> face(2);

      mesh.add_face(face);
    }

    if (n_vertices != mesh.num_vertices()) {
      std::cout << "[Error] Number of vertices in header differs from actual number of vertices." << std::endl;
      return false;
    }

    if (n_faces != mesh.num_faces()) {
      std::cout << "[Error] Number of faces in header differs from actual number of faces." << std::endl;
      return false;
    }

    file->close();
    delete file;

    return true;
  }

  /** \brief Write mesh to OFF file.
   * \param[in] filepath path to OFF file to write
   * \return success
   */
  bool to_off(const std::string filepath) {
    std::ofstream* out = new std::ofstream(filepath, std::ofstream::out);
    if (!static_cast<bool>(out)) {
      return false;
    }

    (*out) << "OFF" << std::endl;
    (*out) << this->num_vertices() << " " << this->num_faces() << " 0" << std::endl;

    for (unsigned int v = 0; v < this->num_vertices(); v++) {
      (*out) << this->vertices[v](0) << " " << this->vertices[v](1) << " " << this->vertices[v](2) << std::endl;
    }

    for (unsigned int f = 0; f < this->num_faces(); f++) {
      (*out) << "3 " << this->faces[f](0) << " " << this->faces[f](1) << " " << this->faces[f](2) << std::endl;
    }

    out->close();
    delete out;

    return true;
  }

  /** \brief Add a vertex.
   * \param[in] vertex vertex to add
   */
  void add_vertex(Eigen::Vector3f& vertex) {
    this->vertices.push_back(vertex);
  }

  /** \brief Get the number of vertices.
   * \return number of vertices
   */
  int num_vertices() {
    return static_cast<int>(this->vertices.size());
  }

  /** \brief Add a face.
   * \param[in] face face to add
   */
  void add_face(Eigen::Vector3i& face) {
    this->faces.push_back(face);
  }

  /** \brief Get the number of faces.
   * \return number of faces
   */
  int num_faces() {
    return static_cast<int>(this->faces.size());
  }

  /** \brief Translate the mesh.
   * \param[in] translation translation vector
   */
  void translate(const Eigen::Vector3f& translation) {
    for (int v = 0; v < this->num_vertices(); ++v) {
      for (int i = 0; i < 3; ++i) {
        this->vertices[v](i) += translation(i);
      }
    }
  }

  /** \brief Scale the mesh.
   * \param[in] scale scale vector
   */
  void scale(const Eigen::Vector3f& scale) {
    for (int v = 0; v < this->num_vertices(); ++v) {
      for (int i = 0; i < 3; ++i) {
        this->vertices[v](i) *= scale(i);
      }
    }
  }

  /** \brief Voxelize the given mesh into a SDF.
   * \param[out] sdf volume to fill with sdf values
   */
  void voxelize_sdf(Eigen::Tensor<float, 3, Eigen::RowMajor>& sdf, const VoxelizationMode &mode) {

    int height = sdf.dimension(0);
    int width = sdf.dimension(1);
    int depth = sdf.dimension(2);

    #pragma omp parallel
    {
      #pragma omp for
      for (int i = 0; i < height*width*depth; i++) {
        int d = i%depth;
        int w = (i/depth)%width;
        int h = (i/depth)/width;

        sdf(h, w, d) = FLT_MAX;

        // the box corresponding to this voxel
        Eigen::Vector3f min(w, h, d);
        Eigen::Vector3f max(w + 1, h + 1, d + 1);

        Eigen::Vector3f center(w + 0.5f, h + 0.5f, d + 0.5f);
        if (mode == VoxelizationMode::CORNER) {
          center = Eigen::Vector3f(w, h, d);
        }

        // count number of intersections.
        int num_intersect = 0;
        for (unsigned int f = 0; f < this->num_faces(); ++f) {

          Eigen::Vector3f v1 = this->vertices[this->faces[f](0)];
          Eigen::Vector3f v2 = this->vertices[this->faces[f](1)];
          Eigen::Vector3f v3 = this->vertices[this->faces[f](2)];

          Eigen::Vector3f closest_point;
          triangle_point_distance(center, v1, v2, v3, closest_point);
          float distance = (center - closest_point).norm();

          if (distance < sdf(h, w, d)) {
            sdf(h, w, d) = distance;
          }

          bool intersect = triangle_ray_intersection(center, Eigen::Vector3f(0, 0, 0), v1, v2, v3, distance);

          if (intersect && distance >= 0) {
            num_intersect++;
          }
        }

        if (num_intersect%2 == 1) {
          sdf(h, w, d) *= -1;
        }
      }
    }
  }

  /** \brief Voxelize the given mesh into an occupancy grid.
   * \param[out] occ volume to fill
   */
  void voxelize_occ(Eigen::Tensor<int, 3, Eigen::RowMajor>& occ, const VoxelizationMode &mode) {

    int height = occ.dimension(0);
    int width = occ.dimension(1);
    int depth = occ.dimension(2);

    #pragma omp parallel
    {
      #pragma omp for
      for (int i = 0; i < height*width*depth; i++) {
        int d = i%depth;
        int w = (i/depth)%width;
        int h = (i/depth)/width;

        Eigen::Vector3f min(w, h, d);
        Eigen::Vector3f max(w + 1, h + 1, d + 1);

        for (unsigned int f = 0; f < this->num_faces(); ++f) {

          Eigen::Vector3f v1 = this->vertices[this->faces[f](0)];
          Eigen::Vector3f v2 = this->vertices[this->faces[f](1)];
          Eigen::Vector3f v3 = this->vertices[this->faces[f](2)];

          bool overlap = triangle_box_intersection(min, max, v1, v2, v3);
          if (overlap) {
            occ(h, w, d) = 1;
            break;
          }
        }
      }
    }
  }

private:

  /** \brief Vertices as (x,y,z)-vectors. */
  std::vector<Eigen::Vector3f> vertices;

  /** \brief Faces as list of vertex indices. */
  std::vector<Eigen::Vector3i> faces;
};

/** \brief Write the given set of volumes to h5 file.
 * \param[in] filepath h5 file to write
 * \param[in] n number of volumes
 * \param[in] height height of volumes
 * \param[in] width width of volumes
 * \param[in] depth depth of volumes
 * \param[in] dense volume data
 */
template<int RANK>
bool write_float_hdf5(const std::string filepath, Eigen::Tensor<float, RANK, Eigen::RowMajor>& tensor) {

  try {

    /*
     * Turn off the auto-printing when failure occurs so that we can
     * handle the errors appropriately
     */
    H5::Exception::dontPrint();

    /*
     * Create a new file using H5F_ACC_TRUNC access,
     * default file creation properties, and default file
     * access properties.
     */
    H5::H5File file(filepath, H5F_ACC_TRUNC);

    /*
     * Define the size of the array and create the data space for fixed
     * size dataset.
     */
    hsize_t rank = RANK;
    hsize_t dimsf[rank];
    for (int i = 0; i < rank; i++) {
      dimsf[i] = tensor.dimension(i);

    }
    H5::DataSpace dataspace(rank, dimsf);

    /*
     * Define datatype for the data in the file.
     * We will store little endian INT numbers.
     */
    H5::IntType datatype(H5::PredType::NATIVE_FLOAT);
    datatype.setOrder(H5T_ORDER_LE);

    /*
     * Create a new dataset within the file using defined dataspace and
     * datatype and default dataset creation properties.
     */
    H5::DataSet dataset = file.createDataSet("tensor", datatype, dataspace);

    /*
     * Write the data to the dataset using default memory space, file
     * space, and transfer properties.
     */
    float* data = static_cast<float*>(tensor.data());
    dataset.write(data, H5::PredType::NATIVE_FLOAT);
  }  // end of try block

  // catch failure caused by the H5File operations
  catch(H5::FileIException error) {
    error.printError();
    return false;
  }

  // catch failure caused by the DataSet operations
  catch(H5::DataSetIException error) {
    error.printError();
    return false;
  }

  // catch failure caused by the DataSpace operations
  catch(H5::DataSpaceIException error) {
    error.printError();
    return false;
  }

  // catch failure caused by the DataSpace operations
  catch(H5::DataTypeIException error) {
    error.printError();
    return false;
  }

  return true;
}

/** \brief Write the given set of volumes to h5 file.
 * \param[in] filepath h5 file to write
 * \param[in] n number of volumes
 * \param[in] height height of volumes
 * \param[in] width width of volumes
 * \param[in] depth depth of volumes
 * \param[in] dense volume data
 */
template<int RANK>
bool write_int_hdf5(const std::string filepath, Eigen::Tensor<int, RANK, Eigen::RowMajor>& tensor) {

  try {

    /*
     * Turn off the auto-printing when failure occurs so that we can
     * handle the errors appropriately
     */
    H5::Exception::dontPrint();

    /*
     * Create a new file using H5F_ACC_TRUNC access,
     * default file creation properties, and default file
     * access properties.
     */
    H5::H5File file(filepath, H5F_ACC_TRUNC);

    /*
     * Define the size of the array and create the data space for fixed
     * size dataset.
     */
    hsize_t rank = RANK;
    hsize_t dimsf[rank];
    for (int i = 0; i < rank; i++) {
      dimsf[i] = tensor.dimension(i);

    }
    H5::DataSpace dataspace(rank, dimsf);

    /*
     * Define datatype for the data in the file.
     * We will store little endian INT numbers.
     */
    H5::IntType datatype(H5::PredType::NATIVE_INT);
    datatype.setOrder(H5T_ORDER_LE);

    /*
     * Create a new dataset within the file using defined dataspace and
     * datatype and default dataset creation properties.
     */
    H5::DataSet dataset = file.createDataSet("tensor", datatype, dataspace);

    /*
     * Write the data to the dataset using default memory space, file
     * space, and transfer properties.
     */
    int* data = static_cast<int*>(tensor.data());
    dataset.write(data, H5::PredType::NATIVE_INT);
  }  // end of try block

  // catch failure caused by the H5File operations
  catch(H5::FileIException error) {
    error.printError();
    return false;
  }

  // catch failure caused by the DataSet operations
  catch(H5::DataSetIException error) {
    error.printError();
    return false;
  }

  // catch failure caused by the DataSpace operations
  catch(H5::DataSpaceIException error) {
    error.printError();
    return false;
  }

  // catch failure caused by the DataSpace operations
  catch(H5::DataTypeIException error) {
    error.printError();
    return false;
  }

  return true;
}

/** \brief Read all files in a directory matching the given extension.
 * \param[in] directory path to directory
 * \param[out] files read file paths
 * \param[in] extension extension to filter for
 */
void read_directory(const boost::filesystem::path directory, std::map<int, boost::filesystem::path>& files, const std::string extension = ".off") {

  files.clear();
  boost::filesystem::directory_iterator end;

  for (boost::filesystem::directory_iterator it(directory); it != end; ++it) {
    if (it->path().extension().string() == extension) {
      int number = std::stoi(it->path().filename().string());
      files.insert(std::pair<int, boost::filesystem::path>(number, it->path()));
    }
  }
}

/** \brief Main entrance point of the script.
 * Expects one parameter, the path to the corresponding config file in config/.
 */
int main(int argc, char** argv) {
  boost::program_options::options_description desc("Allowed options");
  desc.add_options()
      ("help", "produce help message")
      ("mode", boost::program_options::value<std::string>()->default_value("occ"), "operation mode, 'occ' or 'sdf'")
      ("input", boost::program_options::value<std::string>(), "input, either single OFF file or directory containing OFF files where the names correspond to integers (zero padding allowed) and are consecutively numbered starting with zero")
      ("height", boost::program_options::value<int>()->default_value(32), "height of volume, corresponding to y-axis (=up)")
      ("width", boost::program_options::value<int>()->default_value(32), "width of volume, corresponding to x-axis (=right")
      ("depth", boost::program_options::value<int>()->default_value(32), "depth of volume, corresponding to z-axis (=forward)")
      ("center", boost::program_options::bool_switch()->default_value(false), "by default, the top-left-front corner is used for SDF computation; if instead the voxel centers should be used, set this flag")
      ("output", boost::program_options::value<std::string>(), "output file, will be a HDF5 file containing either a N x C x height x width x depth tensor or a C x height x width x depth tensor, where N is the number of files and C=2 the number of channels, N is discarded if only a single file is processed; should have the .h5 extension");

  boost::program_options::positional_options_description positionals;
  positionals.add("mode", 1);
  positionals.add("input", 1);
  positionals.add("output", 1);

  boost::program_options::variables_map parameters;
  boost::program_options::store(boost::program_options::command_line_parser(argc, argv).options(desc).positional(positionals).run(), parameters);
  boost::program_options::notify(parameters);

  if (parameters.find("help") != parameters.end()) {
    std::cout << desc << std::endl;
    return 0;
  }

  std::string mode = parameters["mode"].as<std::string>();
  if (mode == "occ") {
    std::cout << "Voxelizing occupancy grids." << std::endl;
  }
  else if (mode == "sdf") {
    std::cout << "Voxelizing SDFs." << std::endl;
  }
  else {
    std::cout << "Invalid mode, choose from occ or sdf." << std::endl;
    return 1;
  }

  boost::filesystem::path input(parameters["input"].as<std::string>());
  if (!boost::filesystem::is_directory(input) && !boost::filesystem::is_regular_file(input)) {
    std::cout << "Input is neither directory nor file." << std::endl;
    return 1;
  }

  boost::filesystem::path output(parameters["output"].as<std::string>());
  if (boost::filesystem::is_regular_file(output)) {
    std::cout << "Output file already exists; overwriting." << std::endl;
  }

  VoxelizationMode voxelization_mode;
  if (parameters["center"].as<bool>()) {
    voxelization_mode = VoxelizationMode::CENTER;
    std::cout << "Using the top-left-front voxel corner for voxelization." << std::endl;
  }
  else {
    voxelization_mode = VoxelizationMode::CORNER;
    std::cout << "Using the voxel center for voxelization." << std::endl;
  }

  int height = parameters["height"].as<int>();
  int width = parameters["width"].as<int>();
  int depth = parameters["depth"].as<int>();

  std::cout << "Voxelizing into " << height << " x " << width << " x " << depth << " (height x width x depth)." << std::endl;

  if (boost::filesystem::is_regular_file(input)) {
    Mesh mesh;
    bool success = Mesh::from_off(input.string(), mesh);

    if (!success) {
      std::cout << "Could not read " << input << "." << std::endl;
      return 1;
    }

    std::cout << "Read " << input << "." << std::endl;

    if (mode == "sdf") {
      Eigen::Tensor<float, 3, Eigen::RowMajor> tensor(height, width, depth);

      mesh.voxelize_sdf(tensor, voxelization_mode);
      std::cout << "Voxelized " << input << "." << std::endl;

      bool success = write_float_hdf5<3>(output.string(), tensor);

      if (!success) {
        std::cout << "Could not write " << output << "." << std::endl;
        return 1;
      }
    }
    if (mode == "occ") {
      Eigen::Tensor<int, 3, Eigen::RowMajor> tensor(height, width, depth);
      tensor.setZero();

      mesh.voxelize_occ(tensor, voxelization_mode);
      std::cout << "Voxelized " << input << "." << std::endl;

      bool success = write_int_hdf5<3>(output.string(), tensor);

      if (!success) {
        std::cout << "Could not write " << output << "." << std::endl;
        return 1;
      }
    }

    std::cout << "Wrote " << output << "." << std::endl;
    std::cout << "The output is a " << height << " x " << width << " x " << depth << " tensor." << std::endl;
  }
  else {
    std::map<int, boost::filesystem::path> input_files;
    read_directory(input, input_files);

    if (input_files.size() <= 0) {
      std::cout << "Could not find any OFF files in the input directory." << std::endl;
      return 1;
    }

    std::cout << "Read " << input_files.size() << " files." << std::endl;

    if (mode == "sdf") {
      Eigen::Tensor<float, 4, Eigen::RowMajor> tensor(input_files.size(), height, width, depth);

      int i = 0;
      for (std::map<int, boost::filesystem::path>::iterator it = input_files.begin(); it != input_files.end(); it++) {
        Mesh mesh;
        bool success = Mesh::from_off(it->second.string(), mesh);

        if (!success) {
          std::cout << "Could not read " << it->second << "." << std::endl;
          return 1;
        }

        Eigen::Tensor<float, 3, Eigen::RowMajor> slice(height, width, depth);
        mesh.voxelize_sdf(slice, voxelization_mode);
        tensor.chip(i, 0) = slice;
        std::cout << "Voxelized " << it->second << " (" << (i + 1) << " of " << input_files.size() << ")." << std::endl;

        i++;
      }

      bool success = write_float_hdf5<4>(output.string(), tensor);

      if (!success) {
        std::cout << "Could not write " << output << "." << std::endl;
        return 1;
      }
    }
    if (mode == "occ") {
      Eigen::Tensor<int, 4, Eigen::RowMajor> tensor(input_files.size(), height, width, depth);
      tensor.setZero();

      int i = 0;
      for (std::map<int, boost::filesystem::path>::iterator it = input_files.begin(); it != input_files.end(); it++) {
        Mesh mesh;
        bool success = Mesh::from_off(it->second.string(), mesh);

        if (!success) {
          std::cout << "Could not read " << it->second << "." << std::endl;
          return 1;
        }

        Eigen::Tensor<int, 3, Eigen::RowMajor> slice(height, width, depth);
        slice.setZero();

        mesh.voxelize_occ(slice, voxelization_mode);
        tensor.chip(i, 0) = slice;
        std::cout << "Voxelized " << it->second << " (" << (i + 1) << " of " << input_files.size() << ")." << std::endl;

        i++;
      }

      bool success = write_int_hdf5<4>(output.string(), tensor);

      if (!success) {
        std::cout << "Could not write " << output << "." << std::endl;
        return 1;
      }
    }

    std::cout << "Wrote " << output << "." << std::endl;
    std::cout << "The output is a " << input_files.size() << " x " << height << " x " << width << " x " << depth << " tensor." << std::endl;
  }

  return 0;
}


/*********************************************************************************
 *  OKVIS - Open Keyframe-based Visual-Inertial SLAM
 *  Copyright (c) 2015, Autonomous Systems Lab / ETH Zurich
 *
 *  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 Autonomous Systems Lab / ETH Zurich 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 OWNER 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.
 *
 *  Created on: Jan 7, 2014
 *      Author: Stefan Leutenegger (s.leutenegger@imperial.ac.uk)
 *********************************************************************************/

/**
 * @file odeHybrid.hpp
 * @brief Runge Kutta based IMU forward and backward propagation with an extended IMU model.
 * developed from ode.hpp in okvis.
 * @author Jianzhu Huai
 */

#ifndef INCLUDE_SWIFT_VIO_ODE_HYBRID_HPP_
#define INCLUDE_SWIFT_VIO_ODE_HYBRID_HPP_

#include <swift_vio/imu/ImuErrorModel.h>

#include <Eigen/Core>

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

namespace swift_vio {
namespace ode {
const int kNavErrorStateDim = 9;
// Note this function assume that the W frame has z axis along the negative gravity.
// This fact is used in computing velocity in the W frame.
__inline__ void evaluateContinuousTimeOde(
    const Eigen::Vector3d& gyr, const Eigen::Vector3d& acc, double g,
    const Eigen::Vector3d& p_WS_W, const Eigen::Quaterniond& q_WS,
    const okvis::SpeedAndBiases& sb,
    const ImuErrorModel<double>& iem, Eigen::Vector3d& p_WS_W_dot,
    Eigen::Vector4d& q_WS_dot, okvis::SpeedAndBiases& sb_dot,
    Eigen::MatrixXd* F_c_ptr = 0) {
  Eigen::Vector3d omega_S;
  Eigen::Vector3d acc_S;
  iem.estimate(gyr, acc, &omega_S, &acc_S);

  // nonlinear states
  // start with the pose
  p_WS_W_dot = sb.head<3>();

  // now the quaternion
  Eigen::Vector4d dq;
  q_WS_dot.head<3>() = 0.5 * omega_S;
  q_WS_dot[3] = 0.0;
  Eigen::Matrix3d C_WS = q_WS.toRotationMatrix();

  // the rest is straightforward
  // consider Earth's radius. Model the Earth as a sphere, since we neither
  // know the position nor yaw (except if coupled with GPS and magnetometer).
  Eigen::Vector3d G =
      -p_WS_W - Eigen::Vector3d(0, 0, 6371009);  // vector to Earth center
  sb_dot.head<3>() = (C_WS * acc_S + g * G.normalized());  // s
  // biases
  sb_dot.tail<6>().setZero();

  // linearized system:
  if (F_c_ptr) {
    F_c_ptr->setZero();
    F_c_ptr->block<3, 3>(0, 6) = Eigen::Matrix3d::Identity();
    //    F_c_ptr->block<3, 3>(3, 9) -= C_WS;
    F_c_ptr->block<3, 3>(6, 3) -= okvis::kinematics::crossMx(C_WS * acc_S);
    //    F_c_ptr->block<3, 3>(6, 12) -= C_WS;
    Eigen::Matrix<double, 9, 6> N = Eigen::Matrix<double, 9, 6>::Zero();
    N.block<3, 3>(3, 0) = C_WS;
    N.block<3, 3>(6, 3) = C_WS;
    int colsF = F_c_ptr->cols();
    Eigen::Matrix<double, 6, 6 + 27> dwaB_dbgbaSTS;
    iem.dwa_B_dbgbaSTS(omega_S, acc_S, dwaB_dbgbaSTS);
    F_c_ptr->block(0, 9, 9, colsF - 9) = N * dwaB_dbgbaSTS.rightCols(colsF - 9);
  }
}

// p_WS_W, q_WS, sb are states at k
/* * covariance error states $\Delta y = \Delta[\mathbf{p}_{WS}^W,
\Delta\mathbf{\alpha}, \Delta\mathbf{v}_S^W,
 * \Delta\mathbf{b_g}, \Delta\mathbf{b_a}, \Delta\mathbf{Tg}, \Delta\mathbf{Ts},
\Delta\mathbf{Ta}]$
 * $\tilde{R}_s^w=(I-[\delta\alpha^w]_\times)R_s^w$
 * this DEFINITION OF ROTATION ERROR is the same in Mingyang Li's later
publications, okvis, and huai ION GNSS+ 2015 $y = [\mathbf{p}_{WS}^W,
\mathbf{q}_S^W, \mathbf{v}_S^W, \mathbf{b_g}, \mathbf{b_a}]$ $u =
[\mathbf{\omega}_{WS}^S,\mathbf{a}^S]$ $h = t_{n+1}-t_n$
$\mathbf{p}_{WS}^W \oplus = \mathbf{p}_{WS}^W +$
$\mathbf{v}_{WS}^W \oplus = \mathbf{v}_{WS}^W +$
$\mathbf{q}_{S}^W \oplus \mathbf{\omega}_{WS}^S h/2 =
\mathbf{q}_{S}^W\begin{bmatrix}
cos(\mathbf{\omega}_{WS}^S h/2) \\
sin(\mathbf{\omega}_{WS}^S
h/2)\frac{\mathbf{\omega}_{WS}^S/2}{\vert\mathbf{\omega}_{WS}^S/2\vert}
\end{bmatrix}$

$k_1 = f(t_n,y_n,u_n)$
$k_2 = f(t_n+h/2,y_n\oplus k_1 h/2 ,(u_n +u_{n+1})/2)$
$k_3 = f(t_n+h/2,y_n\oplus k_2 h/2 ,(u_n +u_{n+1})/2)$
$k_4 = f(t_n+h,y_n\oplus k_3 h , u_{n+1})$
$y_{n+1}=y_n\oplus\left(h(k_1 +2k_2 +2k_3 +k_4)/6 \right )$
Caution: provide both F_tot_ptr(e.g., identity) and P_ptr(e.g., zero matrix) if
covariance is to be computed
*/
__inline__ void integrateOneStep_RungeKutta(
    const Eigen::Vector3d& gyr_0, const Eigen::Vector3d& acc_0,
    const Eigen::Vector3d& gyr_1, const Eigen::Vector3d& acc_1, double g,
    double sigma_g_c, double sigma_a_c, double sigma_gw_c, double sigma_aw_c,
    double dt, Eigen::Vector3d& p_WS_W, Eigen::Quaterniond& q_WS,
    okvis::SpeedAndBiases& sb, const ImuErrorModel<double>& iem,
    Eigen::MatrixXd* P_ptr = 0,
    Eigen::MatrixXd* F_tot_ptr = 0) {
  Eigen::Vector3d k1_p_WS_W_dot;
  Eigen::Vector4d k1_q_WS_dot;

  okvis::SpeedAndBiases k1_sb_dot;
  int covRows = 0;
  Eigen::MatrixXd k1_F_c;
  Eigen::MatrixXd k2_F_c;
  Eigen::MatrixXd k3_F_c;
  Eigen::MatrixXd k4_F_c;
  Eigen::MatrixXd* k1_F_c_ptr = nullptr;
  Eigen::MatrixXd* k2_F_c_ptr = nullptr;
  Eigen::MatrixXd* k3_F_c_ptr = nullptr;
  Eigen::MatrixXd* k4_F_c_ptr = nullptr;
  if (P_ptr || F_tot_ptr) {
      covRows = P_ptr ? P_ptr->rows() : F_tot_ptr->rows();
      k1_F_c.resize(covRows, covRows);
      k2_F_c.resize(covRows, covRows);
      k3_F_c.resize(covRows, covRows);
      k4_F_c.resize(covRows, covRows);
      k1_F_c_ptr = &k1_F_c;
      k2_F_c_ptr = &k2_F_c;
      k3_F_c_ptr = &k3_F_c;
      k4_F_c_ptr = &k4_F_c;
  }

  evaluateContinuousTimeOde(gyr_0, acc_0, g, p_WS_W, q_WS, sb, iem,
                            k1_p_WS_W_dot, k1_q_WS_dot, k1_sb_dot, k1_F_c_ptr);

  Eigen::Vector3d p_WS_W1 = p_WS_W;
  Eigen::Quaterniond q_WS1 = q_WS;
  okvis::SpeedAndBiases sb1 = sb;
  // state propagation:
  p_WS_W1 += k1_p_WS_W_dot * 0.5 * dt;
  Eigen::Quaterniond dq;
  double theta_half = k1_q_WS_dot.head<3>().norm() * 0.5 * dt;
  double sinc_theta_half = okvis::kinematics::sinc(theta_half);
  double cos_theta_half = cos(theta_half);
  dq.vec() = sinc_theta_half * k1_q_WS_dot.head<3>() * 0.5 * dt;
  dq.w() = cos_theta_half;
  q_WS1 = q_WS * dq;
  sb1 += k1_sb_dot * 0.5 * dt;

  Eigen::Vector3d k2_p_WS_W_dot;
  Eigen::Vector4d k2_q_WS_dot;
  okvis::SpeedAndBiases k2_sb_dot;
  evaluateContinuousTimeOde(0.5 * (gyr_0 + gyr_1), 0.5 * (acc_0 + acc_1), g,
                            p_WS_W1, q_WS1, sb1, iem, k2_p_WS_W_dot,
                            k2_q_WS_dot, k2_sb_dot, k2_F_c_ptr);

  Eigen::Vector3d p_WS_W2 = p_WS_W;
  Eigen::Quaterniond q_WS2 = q_WS;
  okvis::SpeedAndBiases sb2 = sb;
  // state propagation:
  p_WS_W2 += k2_p_WS_W_dot * 0.5 * dt;
  theta_half = k2_q_WS_dot.head<3>().norm() * 0.5 * dt;
  sinc_theta_half = okvis::kinematics::sinc(theta_half);
  cos_theta_half = cos(theta_half);
  dq.vec() = sinc_theta_half * k2_q_WS_dot.head<3>() * 0.5 * dt;
  dq.w() = cos_theta_half;
  // std::cout<<dq.transpose()<<std::endl;
  q_WS2 = q_WS2 * dq;
  sb2 += k1_sb_dot * 0.5 * dt;

  Eigen::Vector3d k3_p_WS_W_dot;
  Eigen::Vector4d k3_q_WS_dot;
  okvis::SpeedAndBiases k3_sb_dot;
  evaluateContinuousTimeOde(0.5 * (gyr_0 + gyr_1), 0.5 * (acc_0 + acc_1), g,
                            p_WS_W2, q_WS2, sb2, iem, k3_p_WS_W_dot,
                            k3_q_WS_dot, k3_sb_dot, k3_F_c_ptr);

  Eigen::Vector3d p_WS_W3 = p_WS_W;
  Eigen::Quaterniond q_WS3 = q_WS;
  okvis::SpeedAndBiases sb3 = sb;
  // state propagation:
  p_WS_W3 += k3_p_WS_W_dot * dt;
  theta_half = k3_q_WS_dot.head<3>().norm() * dt;
  sinc_theta_half = okvis::kinematics::sinc(theta_half);
  cos_theta_half = cos(theta_half);
  dq.vec() = sinc_theta_half * k3_q_WS_dot.head<3>() * dt;
  dq.w() = cos_theta_half;
  // std::cout<<dq.transpose()<<std::endl;
  q_WS3 = q_WS3 * dq;
  sb3 += k3_sb_dot * dt;

  Eigen::Vector3d k4_p_WS_W_dot;
  Eigen::Vector4d k4_q_WS_dot;
  okvis::SpeedAndBiases k4_sb_dot;
  evaluateContinuousTimeOde(gyr_1, acc_1, g, p_WS_W3, q_WS3, sb3, iem,
                            k4_p_WS_W_dot, k4_q_WS_dot, k4_sb_dot, k4_F_c_ptr);

  // now assemble
  p_WS_W +=
      (k1_p_WS_W_dot + 2 * (k2_p_WS_W_dot + k3_p_WS_W_dot) + k4_p_WS_W_dot) *
      dt / 6.0;
  Eigen::Vector3d theta_half_vec =
      (k1_q_WS_dot.head<3>() +
       2 * (k2_q_WS_dot.head<3>() + k3_q_WS_dot.head<3>()) +
       k4_q_WS_dot.head<3>()) *
      dt / 6.0;
  theta_half = theta_half_vec.norm();
  sinc_theta_half = okvis::kinematics::sinc(theta_half);
  cos_theta_half = cos(theta_half);
  dq.vec() = sinc_theta_half * theta_half_vec;
  dq.w() = cos_theta_half;
  q_WS = q_WS * dq;
  sb += (k1_sb_dot + 2 * (k2_sb_dot + k3_sb_dot) + k4_sb_dot) * dt / 6.0;

  q_WS.normalize();  // do not accumulate errors!

  if (F_tot_ptr) {
    // compute state transition matrix, note $\frac{d\Phi(t, t_0)}{dt}=
    // F(t)\Phi(t, t_0)$
    Eigen::MatrixXd& F_tot = *F_tot_ptr;
    const Eigen::MatrixXd& J1 = k1_F_c;
    const Eigen::MatrixXd J2 =
        k2_F_c * (Eigen::MatrixXd::Identity(covRows, covRows) + 0.5 * dt * J1);
    const Eigen::MatrixXd J3 =
        k3_F_c * (Eigen::MatrixXd::Identity(covRows, covRows) + 0.5 * dt * J2);
    const Eigen::MatrixXd J4 =
        k4_F_c * (Eigen::MatrixXd::Identity(covRows, covRows) + dt * J3);
    Eigen::MatrixXd F = Eigen::MatrixXd::Identity(covRows, covRows) +
                        dt * (J1 + 2 * (J2 + J3) + J4) / 6.0;
    F_tot =
        (F * F_tot)
            .eval();  // F is $\Phi(t_k, t_{k-1})$, F_tot is $\Phi(t_k, t_{0})$

    if (P_ptr) {
      Eigen::MatrixXd& cov = *P_ptr;
      cov = F * (cov * F.transpose()).eval();

      // add process noise
      const double Q_g = sigma_g_c * sigma_g_c * dt;
      const double Q_a = sigma_a_c * sigma_a_c * dt;
      const double Q_gw = sigma_gw_c * sigma_gw_c * dt;
      const double Q_aw = sigma_aw_c * sigma_aw_c * dt;
      cov(3, 3) += Q_g;
      cov(4, 4) += Q_g;
      cov(5, 5) += Q_g;
      cov(6, 6) += Q_a;
      cov(7, 7) += Q_a;
      cov(8, 8) += Q_a;
      cov(9, 9) += Q_gw;
      cov(10, 10) += Q_gw;
      cov(11, 11) += Q_gw;
      cov(12, 12) += Q_aw;
      cov(13, 13) += Q_aw;
      cov(14, 14) += Q_aw;

      // force symmetric
      // huai: this may help keep cov positive semi-definite after propagation
      cov = 0.5 * cov + 0.5 * cov.transpose().eval();
    }
  }
}

/* p_WS_W, q_WS, sb are states at k+1, dt= t(k+1) -t(k)
$y = [\mathbf{p}_{WS}^W, \mathbf{q}_S^W, \mathbf{v}_S^W, \mathbf{b_g},
\mathbf{b_a}]$ $u = [\mathbf{\omega}_{WS}^S,\mathbf{a}^S]$ $h = t_{n}-t_{n+1}$
$\mathbf{p}_{WS}^W \oplus = \mathbf{p}_{WS}^W +$
$\mathbf{v}_{WS}^W \oplus = \mathbf{v}_{WS}^W +$
$\mathbf{q}_{S}^W \oplus \mathbf{\omega}_{WS}^S h/2 =
\mathbf{q}_{S}^W\begin{bmatrix}
cos(\mathbf{\omega}_{WS}^S h/2) \\
sin(\mathbf{\omega}_{WS}^S
h/2)\frac{\mathbf{\omega}_{WS}^S/2}{\vert\mathbf{\omega}_{WS}^S/2\vert}
\end{bmatrix}$

$k_1 = f(t_{n+1},y_{n+1},u_{n+1})$
$k_2 = f(t_{n+1}+h/2,y_{n+1}\oplus k_1 h/2 ,(u_n +u_{n+1})/2)$
$k_3 = f(t_{n+1}+h/2,y_{n+1}\oplus k_2 h/2 ,(u_n +u_{n+1})/2)$
$k_4 = f(t_n,y_{n+1}\oplus k_3 h , u_{n})$
$y_{n}=y_{n+1}\oplus\left(h(k_1 +2k_2 +2k_3 +k_4)/6 \right )$
*/
__inline__ void integrateOneStepBackward_RungeKutta(
    const Eigen::Vector3d& gyr_0, const Eigen::Vector3d& acc_0,
    const Eigen::Vector3d& gyr_1, const Eigen::Vector3d& acc_1, double g,
    double sigma_g_c, double sigma_a_c, double sigma_gw_c, double sigma_aw_c,
    double dt, Eigen::Vector3d& p_WS_W, Eigen::Quaterniond& q_WS,
    okvis::SpeedAndBiases& sb, const ImuErrorModel<double>& iem,
    Eigen::MatrixXd* P_ptr = 0,
    Eigen::MatrixXd* F_tot_ptr = 0) {
  Eigen::Vector3d k1_p_WS_W_dot;
  Eigen::Vector4d k1_q_WS_dot;
  okvis::SpeedAndBiases k1_sb_dot;

  int covRows = 0;
  Eigen::MatrixXd k1_F_c;
  Eigen::MatrixXd k2_F_c;
  Eigen::MatrixXd k3_F_c;
  Eigen::MatrixXd k4_F_c;
  Eigen::MatrixXd* k1_F_c_ptr = nullptr;
  Eigen::MatrixXd* k2_F_c_ptr = nullptr;
  Eigen::MatrixXd* k3_F_c_ptr = nullptr;
  Eigen::MatrixXd* k4_F_c_ptr = nullptr;
  if (P_ptr || F_tot_ptr) {
      covRows = P_ptr ? P_ptr->rows() : F_tot_ptr->rows();
      k1_F_c.resize(covRows, covRows);
      k2_F_c.resize(covRows, covRows);
      k3_F_c.resize(covRows, covRows);
      k4_F_c.resize(covRows, covRows);
      k1_F_c_ptr = &k1_F_c;
      k2_F_c_ptr = &k2_F_c;
      k3_F_c_ptr = &k3_F_c;
      k4_F_c_ptr = &k4_F_c;
  }

  evaluateContinuousTimeOde(gyr_1, acc_1, g, p_WS_W, q_WS, sb, iem,
                            k1_p_WS_W_dot, k1_q_WS_dot, k1_sb_dot, k1_F_c_ptr);

  Eigen::Vector3d p_WS_W1 = p_WS_W;
  Eigen::Quaterniond q_WS1 = q_WS;
  okvis::SpeedAndBiases sb1 = sb;
  // state propagation:
  p_WS_W1 -= k1_p_WS_W_dot * 0.5 * dt;
  Eigen::Quaterniond dq;
  double theta_half = -k1_q_WS_dot.head<3>().norm() * 0.5 * dt;
  double sinc_theta_half = okvis::kinematics::sinc(theta_half);
  double cos_theta_half = cos(theta_half);
  dq.vec() = -sinc_theta_half * k1_q_WS_dot.head<3>() * 0.5 * dt;
  dq.w() = cos_theta_half;
  q_WS1 = q_WS * dq;
  sb1 -= k1_sb_dot * 0.5 * dt;

  Eigen::Vector3d k2_p_WS_W_dot;
  Eigen::Vector4d k2_q_WS_dot;
  okvis::SpeedAndBiases k2_sb_dot;
  evaluateContinuousTimeOde(0.5 * (gyr_0 + gyr_1), 0.5 * (acc_0 + acc_1), g,
                            p_WS_W1, q_WS1, sb1, iem, k2_p_WS_W_dot,
                            k2_q_WS_dot, k2_sb_dot, k2_F_c_ptr);

  Eigen::Vector3d p_WS_W2 = p_WS_W;
  Eigen::Quaterniond q_WS2 = q_WS;
  okvis::SpeedAndBiases sb2 = sb;
  // state propagation:
  p_WS_W2 -= k2_p_WS_W_dot * 0.5 * dt;
  theta_half = -k2_q_WS_dot.head<3>().norm() * 0.5 * dt;
  sinc_theta_half = okvis::kinematics::sinc(theta_half);
  cos_theta_half = cos(theta_half);
  dq.vec() = -sinc_theta_half * k2_q_WS_dot.head<3>() * 0.5 * dt;
  dq.w() = cos_theta_half;
  // std::cout<<dq.transpose()<<std::endl;
  q_WS2 = q_WS2 * dq;
  sb2 -= k1_sb_dot * 0.5 * dt;

  Eigen::Vector3d k3_p_WS_W_dot;
  Eigen::Vector4d k3_q_WS_dot;
  okvis::SpeedAndBiases k3_sb_dot;
  evaluateContinuousTimeOde(0.5 * (gyr_0 + gyr_1), 0.5 * (acc_0 + acc_1), g,
                            p_WS_W2, q_WS2, sb2, iem, k3_p_WS_W_dot,
                            k3_q_WS_dot, k3_sb_dot, k3_F_c_ptr);

  Eigen::Vector3d p_WS_W3 = p_WS_W;
  Eigen::Quaterniond q_WS3 = q_WS;
  okvis::SpeedAndBiases sb3 = sb;
  // state propagation:
  p_WS_W3 -= k3_p_WS_W_dot * dt;
  theta_half = -k3_q_WS_dot.head<3>().norm() * dt;
  sinc_theta_half = okvis::kinematics::sinc(theta_half);
  cos_theta_half = cos(theta_half);
  dq.vec() = -sinc_theta_half * k3_q_WS_dot.head<3>() * dt;
  dq.w() = cos_theta_half;
  // std::cout<<dq.transpose()<<std::endl;
  q_WS3 = q_WS3 * dq;
  sb3 -= k3_sb_dot * dt;

  Eigen::Vector3d k4_p_WS_W_dot;
  Eigen::Vector4d k4_q_WS_dot;
  okvis::SpeedAndBiases k4_sb_dot;
  evaluateContinuousTimeOde(gyr_0, acc_0, g, p_WS_W3, q_WS3, sb3, iem,
                            k4_p_WS_W_dot, k4_q_WS_dot, k4_sb_dot, k4_F_c_ptr);

  // now assemble
  p_WS_W -=
      (k1_p_WS_W_dot + 2 * (k2_p_WS_W_dot + k3_p_WS_W_dot) + k4_p_WS_W_dot) *
      dt / 6.0;
  Eigen::Vector3d theta_half_vec =
      -(k1_q_WS_dot.head<3>() +
        2 * (k2_q_WS_dot.head<3>() + k3_q_WS_dot.head<3>()) +
        k4_q_WS_dot.head<3>()) *
      dt / 6.0;
  theta_half = theta_half_vec.norm();
  sinc_theta_half = okvis::kinematics::sinc(theta_half);
  cos_theta_half = cos(theta_half);
  dq.vec() = sinc_theta_half * theta_half_vec;
  dq.w() = cos_theta_half;
  q_WS = q_WS * dq;
  sb -= (k1_sb_dot + 2 * (k2_sb_dot + k3_sb_dot) + k4_sb_dot) * dt / 6.0;

  q_WS.normalize();  // do not accumulate errors!

  if (F_tot_ptr) {
    assert(false);  // the following section is not well perused and tested
    // compute state transition matrix, note $\frac{d\Phi(t, t_0)}{dt}=
    // F(t)\Phi(t, t_0)$
    Eigen::MatrixXd& F_tot = *F_tot_ptr;
    const Eigen::MatrixXd& J1 = k1_F_c;
    const Eigen::MatrixXd J2 =
        k2_F_c * (Eigen::MatrixXd::Identity(covRows, covRows) - 0.5 * dt * J1);
    const Eigen::MatrixXd J3 =
        k3_F_c * (Eigen::MatrixXd::Identity(covRows, covRows) - 0.5 * dt * J2);
    const Eigen::MatrixXd J4 =
        k4_F_c * (Eigen::MatrixXd::Identity(covRows, covRows) - dt * J3);
    Eigen::MatrixXd F = Eigen::MatrixXd::Identity(covRows, covRows) -
                        dt * (J1 + 2 * (J2 + J3) + J4) / 6.0;
    F_tot = (F * F_tot).eval();

    if (P_ptr) {
      Eigen::MatrixXd& cov = *P_ptr;
      cov = F * (cov * F.transpose()).eval();

      // add process noise
      const double Q_g = sigma_g_c * sigma_g_c * dt;
      const double Q_a = sigma_a_c * sigma_a_c * dt;
      const double Q_gw = sigma_gw_c * sigma_gw_c * dt;
      const double Q_aw = sigma_aw_c * sigma_aw_c * dt;
      cov(3, 3) += Q_g;
      cov(4, 4) += Q_g;
      cov(5, 5) += Q_g;
      cov(6, 6) += Q_a;
      cov(7, 7) += Q_a;
      cov(8, 8) += Q_a;
      cov(9, 9) += Q_gw;
      cov(10, 10) += Q_gw;
      cov(11, 11) += Q_gw;
      cov(12, 12) += Q_aw;
      cov(13, 13) += Q_aw;
      cov(14, 14) += Q_aw;

      // force symmetric - TODO: is this really needed here?
      // cov = 0.5 * cov + 0.5 * cov.transpose().eval();
    }
  }
}
}  // namespace ode
}  // namespace swift_vio

#endif // INCLUDE_SWIFT_VIO_ODE_HYBRID_HPP_


//Link to Boost
 #define BOOST_TEST_DYN_LINK

//VERY IMPORTANT - include this last
#include <boost/test/unit_test.hpp>
#include <boost/timer.hpp>

#include <vector>
#include <iostream>
#include "test.h"
#include "../BootstrapRunner.h"
using namespace std;

BOOST_FIXTURE_TEST_SUITE(CorrelatorSet_suite, SimpleTestFixture, * utf::label("correlatorset"))

BOOST_DATA_TEST_CASE(CostDerDelta_test1, CreateCfgConfigTestData(2), cfg)
{
    CftData cftData(cfg);
    int scalarRange = min(4, cftData.MaxScalarId());

    CorrelatorSet correlatorSet(&cftData, scalarRange);
    cpx_t z = RandomComplex(0.3) + 0.5;
    int op = randomint(1, scalarRange);

    float_type actual = correlatorSet.CostDerDelta(z, op);

    float_type dim = cftData.GetPrimaryDim(op);
    cftData.SetPrimaryDim(op, dim - inc);
    correlatorSet.UpdateCftData(&cftData);
    float_type value1 = correlatorSet.Cost(z);

    cftData.SetPrimaryDim(op, dim + inc);
    correlatorSet.UpdateCftData(&cftData);
    float_type value2 = correlatorSet.Cost(z);

    float_type expect = (value2 - value1) / (2 * inc);
    MY_FLOAT_EQUAL(actual, expect, tol);
}

BOOST_DATA_TEST_CASE(CostDerDelta_test2, CreateCfgConfigTestData(4), cfg)
{
    CftData cftData(cfg);

    CorrelatorSet correlatorSet(&cftData, min(4, cftData.MaxScalarId()));
    cpx_t z = RandomComplex(0.3) + 0.5;

    int op = randomint(cftData.StressTensorId + 1, cftData.MaxPrimaryId());

    float_type actual = correlatorSet.CostDerDelta(z, op);

    float_type dim = cftData.GetPrimaryDim(op);
    cftData.SetPrimaryDim(op, dim - inc);
    correlatorSet.UpdateCftData(&cftData);
    float_type value1 = correlatorSet.Cost(z);

    cftData.SetPrimaryDim(op, dim + inc);
    correlatorSet.UpdateCftData(&cftData);
    float_type value2 = correlatorSet.Cost(z);

    float_type expect = (value2 - value1) / (2 * inc);
    MY_FLOAT_EQUAL(actual, expect, tol);
}

BOOST_DATA_TEST_CASE(CostDerOpeCoefficient_test, CreateCfgConfigTestData(4), cfg)
{
    CftData cftData(cfg);

    int scalarRange = min(4, cftData.MaxScalarId());
    CorrelatorSet correlatorSet(&cftData, scalarRange);
    cpx_t z = RandomComplex(0.3) + 0.5;

    int op1 = randomint(1, scalarRange);
    int op2 = randomint(1, scalarRange);
    int op3 = randomint(1, cftData.MaxScalarId());
    OpeCoefficientKey key(op1, op2, op3);

    float_type actual = correlatorSet.CostDerOpeCoefficient(z, key);

    float_type coef = cftData.GetOpeCoefficient(key);
    cftData.SetOpeCoefficient(key, coef - inc);
    correlatorSet.UpdateCftData(&cftData);
    float_type value1 = correlatorSet.Cost(z);

    cftData.SetOpeCoefficient(key, coef + inc);
    correlatorSet.UpdateCftData(&cftData);
    float_type value2 = correlatorSet.Cost(z);

    float_type expect = (value2 - value1) / (2 * inc);
    MY_FLOAT_EQUAL(actual, expect, tol);
}

BOOST_DATA_TEST_CASE(CostDerivatives_test, CreateCfgConfigTestData(1), cfg)
{
    CftData cftData(cfg);

    int scalarRange = min(4, cftData.MaxScalarId());
    CorrelatorSet correlatorSet(&cftData, scalarRange);
    cpx_t z = RandomComplex(0.3) + 0.5;

    int op1 = randomint(1, scalarRange);
    int op2 = randomint(1, scalarRange);
    int op3 = randomint(1, cftData.MaxScalarId());
    int op4 = randomint(1, cftData.MaxScalarId());
    OpeCoefficientKey key1(op1, op2, op3);
    OpeCoefficientKey key2(op1, op2, op4);
    vector<OpeCoefficientKey> keys = vector<OpeCoefficientKey>({key1, key2});
    vector<int> ops = vector<int>({op1, op4});
    vector<float_type> opsDerivatives;
    vector<float_type> coefsDerivatives;

    // get the derivatives results
    correlatorSet.CostDerivatives(z, ops, keys, opsDerivatives, coefsDerivatives);

    // update the cft data to be old values - inc.
    for (uint i = 0; i < ops.size(); i++) {
        float_type dim = cftData.GetPrimaryDim(ops[i]);
        cftData.SetPrimaryDim(ops[i], dim - inc);
    }

    for (uint i = 0; i < keys.size(); i++) {
        float_type coef = cftData.GetOpeCoefficient(keys[i]);
        cftData.SetOpeCoefficient(keys[i], coef - inc);
    }

    // get the cost value.
    correlatorSet.UpdateCftData(&cftData);
    float_type value1 = correlatorSet.Cost(z);

    // update the cft data to be old values - inc.
    for (uint i = 0; i < ops.size(); i++) {
        float_type dim = cftData.GetPrimaryDim(ops[i]);
        cftData.SetPrimaryDim(ops[i], dim + 2 * inc);
    }

    for (uint i = 0; i < keys.size(); i++) {
        float_type coef = cftData.GetOpeCoefficient(keys[i]);
        cftData.SetOpeCoefficient(keys[i], coef + 2 * inc);
    }

    // get the cost value.
    correlatorSet.UpdateCftData(&cftData);
    float_type value2 = correlatorSet.Cost(z);

    float_type actual = value2 - value1;
    float_type expect = .0;

    // calculate the expected value based on derivatives.
    for (uint i = 0; i < opsDerivatives.size(); i++) {
        expect += 2 * inc * opsDerivatives[i];
    }

    for (uint i = 0; i < coefsDerivatives.size(); i++) {
        expect += 2 * inc * coefsDerivatives[i];
    }

    MY_FLOAT_EQUAL(actual, expect, tol);
}

BOOST_AUTO_TEST_CASE(Performance_test, * utf::label("perf"))
{
    CftConfig cfg = CreateCfgConfigTestData(1)[0];
    CftData cftData(cfg);

    int scalarRange = min(4, cftData.MaxScalarId());
    CorrelatorSet correlatorSet(&cftData, scalarRange);
    cpx_t z = RandomComplex(0.3) + 0.5;

    int op1 = randomint(1, scalarRange);
    int op2 = randomint(1, scalarRange);
    int op3 = randomint(1, cftData.MaxScalarId());
    int op4 = randomint(1, cftData.MaxScalarId());
    OpeCoefficientKey key1(op1, op2, op3);
    OpeCoefficientKey key2(op1, op2, op4);
    OpeCoefficientKey key3(op1, op2, op1);
    OpeCoefficientKey key4(op1, op2, op2);
    vector<OpeCoefficientKey> keys = vector<OpeCoefficientKey>({key1, key2, key3, key4});
    vector<int> ops = vector<int>({op1, op2, op3, op4});
    vector<float_type> opsDerivatives;
    vector<float_type> coefsDerivatives;

    boost::timer stopwatch;

    // get the derivatives results
    correlatorSet.CostDerivatives(z, ops, keys, opsDerivatives, coefsDerivatives);

    BOOST_TEST_MESSAGE("Time to run CostDerivatives: " << stopwatch.elapsed() << " seconds.");

    stopwatch.restart();
    CorrelatorSet correlatorSet2(&cftData, scalarRange);
    for (uint i = 0; i < ops.size(); i++) {
        float_type value = correlatorSet2.CostDerDelta(z, ops[i]);
        MY_FLOAT_EQUAL(value, opsDerivatives[i], tol);
    }

    for (uint i = 0; i < keys.size(); i++) {
        float_type value = correlatorSet2.CostDerOpeCoefficient(z, keys[i]);
        MY_FLOAT_EQUAL(value, coefsDerivatives[i], tol);
    }

    BOOST_TEST_MESSAGE("Time to run CostDerDelta and CostDerOpeCoefficient seperately: " << stopwatch.elapsed() << " seconds.");
}

// test suite end
BOOST_AUTO_TEST_SUITE_END()