#pragma once
#include <type_traits>
#include <Eigen/Dense>
#include <fastad_bits/reverse/core/var_view.hpp>
#include <fastad_bits/reverse/core/eval.hpp>
#include <autoppl/util/traits/var_traits.hpp>
#include <autoppl/util/packs/ptr_pack.hpp>
#include <autoppl/util/logging.hpp>
#include <autoppl/util/time/stopwatch.hpp>
#include <autoppl/math/math.hpp>
#include <autoppl/mcmc/sampler_tools.hpp>
#include <autoppl/mcmc/result.hpp>
#include <autoppl/mcmc/base_mcmc.hpp>
#include <autoppl/mcmc/hmc/nuts/tree_utils.hpp>
#include <autoppl/mcmc/hmc/leapfrog.hpp>
#include <autoppl/mcmc/hmc/hamiltonian.hpp>
#include <autoppl/mcmc/hmc/nuts/configs.hpp>

namespace ppl {
namespace mcmc {

/**
 * Checks if state is entroping based on integrated momentum vector
 * across the path and the scaled momentum at the ends of the path.
 *
 * @param   rho             integrated momentum vector
 * @param   p_beg_scaled    scaled momentum beginning of
 *                          current path (given the direction)
 * @param   p_end_scaled    scaled momentum end of
 *                          current path (given the direction).
 */
template <class MatType1, class MatType2, class MatType3>
bool check_entropy(const MatType1& rho,
                   const MatType2& p_beg_scaled,
                   const MatType3& p_end_scaled)
{
    return rho.dot(p_beg_scaled) > 0 &&
           rho.dot(p_end_scaled) > 0;
}

/**
 * Building binary tree for sampling candidates.
 * Helper function to obtain the forward/backward-most position and momentum.
 * Accept/reject policy is based on UniformDistType parameter and GenType
 *
 * Note that the caller, i.e. nuts(), MUST have theta_adj already pre-computed
 * (theta_adj is a member of input and input will be an instance of TreeInput).
 *
 * @param   n_params            number of (continuous) parameters
 * @param   input               TreeInput-like input object
 * @param   depth               current depth of building tree
 * @param   unif_sampler        an object like std::uniform_distribution(0,1)
 *                              used for metropolis acceptance
 * @param   gen                 rng device
 * @param   momentum_handler    MomentumHandler-like object to compute
 *                              kinetic energy and momentum
 * @param   tree_cache          pointer to cache memory that will be used by build_tree.
 *                              The array of doubles must be of size n_params * 7 * max_depth.
 */
template <class InputType
        , class UniformDistType
        , class GenType
        , class MomentumHandlerType
    >
TreeOutput build_tree(size_t n_params,
                      InputType& input,
                      uint8_t depth,
                      UniformDistType& unif_sampler,
                      GenType& gen,
                      MomentumHandlerType& momentum_handler,
                      double* tree_cache)
{
    constexpr double delta_max = 1000;  // suggested by Gelman

    // base case
    if (depth == 0) {
        double new_potential = leapfrog(input.ad_expr_ref.get(),
                                        input.theta_ref.get(),
                                        input.theta_adj_ref.get(),
                                        input.tp_adj_ref.get(),
                                        input.p_most_ref.get(),
                                        momentum_handler,
                                        input.v * input.epsilon,
                                        true // always reuse previous adjoint
                                        );
        double new_kinetic = momentum_handler.kinetic(input.p_most_ref.get());
        double new_ham = hamiltonian(new_potential, new_kinetic);

        // update number of leapfrogs
        ++(input.n_leapfrog_ref.get());

        // update LSE of weights
        if (std::isnan(new_ham)) { new_ham = math::inf<double>; }
        input.log_sum_weight_ref.get() = math::lse(
                input.log_sum_weight_ref.get(),
                input.ham - new_ham);

        // update sum of probabilities
        input.sum_metro_prob_ref.get() += (input.ham - new_ham > 0) ?
                1 : std::exp(input.ham - new_ham);

        // always copy into theta_prime
        input.theta_prime_ref.get() = input.theta_ref.get();

        // update momenta of beginning of subtree (moving in the direction of input.v)
        input.p_beg_ref.get() = input.p_most_ref.get();
        input.p_beg_scaled_ref.get() =
            momentum_handler.dkinetic_dr(input.p_most_ref.get());

        // update momenta of end of subtree (moving in the direction of input.v)
        input.p_end_ref.get() = input.p_beg_ref.get();
        input.p_end_scaled_ref.get() = input.p_beg_scaled_ref.get();

        // update integrated momentum
        input.rho_ref.get() += input.p_most_ref.get();

        // return validity and new potential
        return TreeOutput(
                (new_ham - input.ham <= delta_max),
                new_potential
            );
    }

    // recursion
    Eigen::Map<Eigen::VectorXd> p_end_inner(tree_cache, n_params);
    Eigen::Map<Eigen::VectorXd> p_end_scaled_inner(tree_cache + n_params, n_params);
    Eigen::Map<Eigen::VectorXd> rho_first(tree_cache + 2*n_params, n_params);
    rho_first.setZero();
    double log_sum_weight_first = math::neg_inf<double>;

    tree_cache += 3 * n_params; // update position of tree cache

    // create a new input for first recursion
    // some references have to rebound
    InputType first_input = input;
    first_input.p_end_ref = p_end_inner;
    first_input.p_end_scaled_ref = p_end_scaled_inner;
    first_input.rho_ref = rho_first;
    first_input.log_sum_weight_ref = log_sum_weight_first;

    // build first subtree
    TreeOutput first_output =
        build_tree(n_params, first_input, depth - 1,
                   unif_sampler, gen, momentum_handler,
                   tree_cache);

    // if first subtree is already invalid, early exit
    // note that caller will break out of doubling process now,
    // so we do not have to worry about updating the other momentum vectors
    if (!first_output.valid) { return first_output; }

    // second recursion
    Eigen::Map<Eigen::VectorXd> theta_double_prime(tree_cache, n_params);
    Eigen::Map<Eigen::VectorXd> p_beg_inner(tree_cache + n_params, n_params);
    Eigen::Map<Eigen::VectorXd> p_beg_scaled_inner(tree_cache + 2*n_params, n_params);
    Eigen::Map<Eigen::VectorXd> rho_second(tree_cache + 3*n_params, n_params);
    rho_second.setZero();
    double log_sum_weight_second = math::neg_inf<double>;

    tree_cache += 4 * n_params;

    // create a new input for second recursion
    InputType second_input = input;
    second_input.theta_prime_ref = theta_double_prime;
    second_input.p_beg_ref = p_beg_inner;
    second_input.p_beg_scaled_ref = p_beg_scaled_inner;
    second_input.rho_ref = rho_second;
    second_input.log_sum_weight_ref = log_sum_weight_second;

    // build second subtree
    TreeOutput second_output =
        build_tree(n_params, second_input, depth - 1,
                   unif_sampler, gen, momentum_handler,
                   tree_cache);

    // if second subtree is invalid, early exit
    // note that we must return first output since it has the potential
    // of the first proposal and we ignore the second proposal
    if (!second_output.valid) {
        first_output.valid = false;
        return first_output;
    }

    // create output to return at the end
    TreeOutput output;

    // sample proposal and update corresponding potential
    double log_sum_weight_curr = math::lse(
            log_sum_weight_first, log_sum_weight_second
            );
    input.log_sum_weight_ref.get() = math::lse(
            input.log_sum_weight_ref.get(), log_sum_weight_curr
            );

    // note: accept_prob is mathematically guaranteed to be <= 1
    double accept_prob = std::exp(log_sum_weight_second - log_sum_weight_curr);
    bool accept = accept_or_reject(accept_prob, unif_sampler, gen);
    if (accept) {
        input.theta_prime_ref.get() =
            second_input.theta_prime_ref.get();
        output.potential = second_output.potential;
    } else {
        output.potential = first_output.potential;
    }

    // check if current subtree is still valid based
    // on entropy condition
    auto rho_curr = rho_first + rho_second;
    input.rho_ref.get() += rho_curr;
    output.valid =
        check_entropy(rho_curr,
                      input.p_beg_scaled_ref.get(),
                      input.p_end_scaled_ref.get()) &&
        check_entropy(rho_first + p_beg_inner,
                      input.p_beg_scaled_ref.get(),
                      p_beg_scaled_inner) &&
        check_entropy(p_end_inner + rho_second,
                      p_end_scaled_inner,
                      input.p_end_scaled_ref.get());

    return output;
}

/**
 * Finds a reasonable epsilon for NUTS algorithm.
 *
 * @param   eps                 initial epsilon (see Gelman's paper)
 * @param   ad_expr             AD expression bound to theta and theta_adj
 * @param   theta               vector of theta values
 * @param   theta_adj           vector of theta adjoints
 * @param   gen                 rng device
 * @param   momentum_handler    MomentumHandler-like object
 */
template <class ADExprType
        , class MatType
        , class GenType
        , class MomentumHandlerType>
double find_reasonable_epsilon(double eps,
                               ADExprType& ad_expr,
                               MatType& theta,
                               MatType& theta_adj,
                               MatType& tp_adj,
                               GenType& gen,
                               MomentumHandlerType& momentum_handler)
{
    // See (STAN) for reference: if epsilon is way out of bounds, just return eps
    if (eps <= 0 || eps > 1e7) return eps;

    const double diff_bound = std::log(0.8);

    size_t n_params = theta.rows(); // theta is expected to be vector-like

    Eigen::MatrixXd mat(n_params, 3);
    Eigen::Map<Eigen::VectorXd> r(mat.col(0).data(), n_params);
    Eigen::Map<Eigen::VectorXd> theta_orig(mat.col(1).data(), n_params);
    Eigen::Map<Eigen::VectorXd> theta_adj_orig(mat.col(2).data(), n_params);

    // sample momentum vector based on handler
    momentum_handler.sample(r, gen);

    // differentiate first to get adjoints and hamiltonian
    const double potential_orig = -ad::autodiff(ad_expr);
    double kinetic_orig = momentum_handler.kinetic(r);
    double ham_orig = hamiltonian(potential_orig, kinetic_orig);

    // save original value and adjoint
    theta_orig = theta;
    theta_adj_orig = theta_adj;

    // get current hamiltonian after leapfrog
    double potential_curr = leapfrog(
            ad_expr, theta, theta_adj, tp_adj,
            r, momentum_handler, eps, true);
    double kinetic_curr = momentum_handler.kinetic(r);
    double ham_curr = hamiltonian(potential_curr, kinetic_curr);

    int a = (ham_orig - ham_curr > diff_bound) ? 1 : -1;

    while (1) {

        // check if break condition holds
        if ( ((a == 1) && !(ham_orig - ham_curr > diff_bound)) ||
             ((a == -1) && !(ham_orig - ham_curr < diff_bound)) ) {
            break;
        }

        // update epsilon
        eps *= (a == -1) ? 0.5 : 2;

        // copy back original value and adjoint
        theta = theta_orig;
        theta_adj = theta_adj_orig;

        // recompute original hamiltonian with new momentum
        momentum_handler.sample(r, gen);
        kinetic_orig = momentum_handler.kinetic(r);
        ham_orig = hamiltonian(potential_orig, kinetic_orig);

        // leapfrog and compute current hamiltonian
        potential_curr = leapfrog(
                ad_expr, theta, theta_adj, tp_adj,
                r, momentum_handler, eps, true);
        kinetic_curr = momentum_handler.kinetic(r);
        ham_curr = hamiltonian(potential_curr, kinetic_curr);

    }

    // copy back original value and adjoint
    theta = theta_orig;
    theta_adj = theta_adj_orig;

    return eps;
}

/**
 * No-U-Turn Sampler (NUTS)
 *
 * User must ensure that the program does not have any discrete parameters.
 * Discrete data is allowed.
 *
 * @param   program     program expression used to determine log-pdf
 * @param   config      NUTS configuration object
 * @param   pack        offset pack result of activating program.
 *                      It will likely be util::OffsetPack where each offset
 *                      value is equivalent to the total number of values needed,
 *                      i.e. if pack.uc_offset is 10, there is exactly 10 unconstrained values
 *                      for the program.
 * @param   res         result object of calling NUTS that will be populated with samples and other information.
 */

template <class ProgramType
        , class OffsetPackType
        , class MCMCResultType
        , class NUTSConfigType = NUTSConfig<>>
void nuts_(ProgramType& program,
           const NUTSConfigType& config,
           const OffsetPackType& pack,
           MCMCResultType& res)
{
    assert(std::get<1>(pack).uc_offset == 0);
    assert(std::get<1>(pack).tp_offset == 0);
    assert(std::get<1>(pack).c_offset == 0);
    assert(std::get<1>(pack).v_offset == 0);

    auto& offset_pack = std::get<0>(pack);
    size_t n_params = offset_pack.uc_offset;

    // initialization of meta-variables
    std::mt19937 gen(config.seed);
    std::uniform_int_distribution direction_sampler(0, 1);
    std::uniform_real_distribution unif_sampler(0., 1.);

    // Transformed parameters, constrained parameter, visit count cache
    // This can be shared across all AD expressions since only one expression
    // will be evaluated at a time.
    Eigen::MatrixXd tp_mat(offset_pack.tp_offset, 2);
    Eigen::Map<Eigen::VectorXd> tp_val(tp_mat.col(0).data(), offset_pack.tp_offset);
    Eigen::Map<Eigen::VectorXd> tp_adj(tp_mat.col(1).data(), offset_pack.tp_offset);
    Eigen::VectorXd constrained(offset_pack.c_offset);
    Eigen::Matrix<size_t, Eigen::Dynamic, 1> visit(offset_pack.v_offset);
    tp_mat.setZero();
    constrained.setZero();
    visit.setZero();

    // momentum matrix (for stability reasons we require knowing 4 momentum)
    // left-subtree backwardmost momentum => bb
    // left-subtree forwardmost momentum => bf
    // right-subtree backwardmost momentum => fb
    // right-subtree forwardmost momentum => ff
    // scaled versions are based on hamiltonian adjusted covariance matrix
    Eigen::MatrixXd cache_mat(n_params, 18);
    cache_mat.setZero();
    Eigen::Map<Eigen::VectorXd> p_bb(cache_mat.col(0).data(), n_params);
    Eigen::Map<Eigen::VectorXd> p_bb_scaled(cache_mat.col(1).data(), n_params);
    Eigen::Map<Eigen::VectorXd> p_bf(cache_mat.col(2).data(), n_params);
    Eigen::Map<Eigen::VectorXd> p_bf_scaled(cache_mat.col(3).data(), n_params);
    Eigen::Map<Eigen::VectorXd> p_fb(cache_mat.col(4).data(), n_params);
    Eigen::Map<Eigen::VectorXd> p_fb_scaled(cache_mat.col(5).data(), n_params);
    Eigen::Map<Eigen::VectorXd> p_ff(cache_mat.col(6).data(), n_params);
    Eigen::Map<Eigen::VectorXd> p_ff_scaled(cache_mat.col(7).data(), n_params);

    // position matrix for thetas and adjoints
    Eigen::Map<Eigen::VectorXd> theta_bb(cache_mat.col(8).data(), n_params);
    Eigen::Map<Eigen::VectorXd> theta_bb_adj(cache_mat.col(9).data(), n_params);
    Eigen::Map<Eigen::VectorXd> theta_ff(cache_mat.col(10).data(), n_params);
    Eigen::Map<Eigen::VectorXd> theta_ff_adj(cache_mat.col(11).data(), n_params);
    Eigen::Map<Eigen::VectorXd> theta_curr(cache_mat.col(12).data(), n_params);
    Eigen::Map<Eigen::VectorXd> theta_curr_adj(cache_mat.col(13).data(), n_params);
    Eigen::Map<Eigen::VectorXd> theta_prime(cache_mat.col(14).data(), n_params);

    // integrated momentum vectors (more stable than checking entropy with theta_ff - theta_bb)
    // forward-subtree => rho_f
    // backward-subtree => rho_b
    // combined subtrees => rho
    Eigen::Map<Eigen::VectorXd> rho_f(cache_mat.col(15).data(), n_params);
    Eigen::Map<Eigen::VectorXd> rho_b(cache_mat.col(16).data(), n_params);
    Eigen::Map<Eigen::VectorXd> rho(cache_mat.col(17).data(), n_params);

    // build-tree helper function cache line
    Eigen::VectorXd tree_cache(n_params * 7 * config.max_depth);
    tree_cache.setZero();

    // AD Expressions for L(theta) (log-pdf up to constant at theta)
    // Note that these expressions are the only ones used ever.
    auto theta_bb_ad_expr = program.ad_log_pdf(util::make_ptr_pack(
            theta_bb.data(), theta_bb_adj.data(),
            tp_val.data(), tp_adj.data(),
            constrained.data(), visit.data() ));
    auto theta_ff_ad_expr = program.ad_log_pdf(util::make_ptr_pack(
            theta_ff.data(), theta_ff_adj.data(),
            tp_val.data(), tp_adj.data(),
            constrained.data(), visit.data() ));
    auto theta_curr_ad_expr = program.ad_log_pdf(util::make_ptr_pack(
            theta_curr.data(), theta_curr_adj.data(),
            tp_val.data(), tp_adj.data(),
            constrained.data(), visit.data() ));

    // bind every AD expression to the same cache line
    auto size_pack = theta_bb_ad_expr.bind_cache_size();
    Eigen::VectorXd ad_val_buf(size_pack(0));
    Eigen::VectorXd ad_adj_buf(size_pack(1));
    theta_bb_ad_expr.bind_cache({ad_val_buf.data(), ad_adj_buf.data()});
    theta_ff_ad_expr.bind_cache({ad_val_buf.data(), ad_adj_buf.data()});
    theta_curr_ad_expr.bind_cache({ad_val_buf.data(), ad_adj_buf.data()});

    // initializes first sample into theta_curr
    // TODO: allow users to choose how to initialize first point?
    program.bind(util::make_ptr_pack(
                theta_curr.data(), nullptr,
                tp_val.data(), nullptr,
                constrained.data(), visit.data()));
    program.init_params(gen, config.prune);

    // initialize current potential (will be "previous" starting in for-loop)
    double potential_prev = -ad::evaluate(theta_curr_ad_expr);

    // initialize momentum handler
    using var_adapter_policy_t = typename
        nuts_config_traits<NUTSConfigType>::var_adapter_policy_t;
    mcmc::MomentumHandler<var_adapter_policy_t> momentum_handler(n_params);

    // initialize step adapter
    const double log_eps = std::log(
        mcmc::find_reasonable_epsilon(
            1., // initial epsilon
            theta_curr_ad_expr, theta_curr,
            theta_curr_adj, tp_adj,
            gen, momentum_handler));
    mcmc::StepAdapter step_adapter(log_eps);        // initialize step adapter with initial log-epsilon
    step_adapter.step_config = config.step_config;  // copy step configs from user

    // initialize variance adapter
    mcmc::VarAdapter<var_adapter_policy_t> var_adapter(
            n_params, config.warmup, config.var_config.init_buffer,
            config.var_config.term_buffer, config.var_config.window_base
            );

    // construct miscellaneous objects
    auto logger = util::ProgressLogger(config.samples + config.warmup, "NUTS");
    util::StopWatch<> stopwatch_warmup;
    util::StopWatch<> stopwatch_sampling;

    // start timing warmup
    stopwatch_warmup.start();

    for (size_t i = 0; i < config.samples + config.warmup; ++i) {

        // if warmup is finished, stop timing warmup and start timing sampling
        if (i == config.warmup) {
            stopwatch_warmup.stop();
            stopwatch_sampling.start();
        }

        logger.printProgress(i);

        // re-initialize vectors to current theta as the "root" of tree
        theta_bb = theta_curr;
        theta_ff = theta_bb;
        mcmc::reset_autodiff(theta_bb_ad_expr, theta_bb_adj, tp_adj);
        theta_ff_adj = theta_bb_adj;   // no need to differentiate again

        // initialize values for multinomial sampling
        // this is the total log sum weight over full tree
        double log_sum_weight = 0.;

        // initialize values used to adapt stepsize
        size_t n_leapfrog = 0;
        double sum_metro_prob = 0.;

        // p ~ N(0, M) (depending on momentum handler)
        momentum_handler.sample(p_bb, gen);
        p_bf = p_bb;
        p_fb = p_bb;
        p_ff = p_bb;

        // scaled p by hamiltonian dkinetic_dr
        p_bb_scaled = momentum_handler.dkinetic_dr(p_bb);
        p_bf_scaled = p_bb_scaled;
        p_fb_scaled = p_bb_scaled;
        p_ff_scaled = p_bb_scaled;

        // re-initialize integrated momentum vectors
        rho = p_bb;

        const double kinetic = momentum_handler.kinetic(p_bb);
        const double ham_prev = mcmc::hamiltonian(potential_prev, kinetic);

        // Note that this object can be reused since all members
        // are guaranteed to overwritten by build_tree.
        mcmc::TreeOutput output;

        for (size_t depth = 0; depth < config.max_depth; ++depth) {

            // zero-out subtree integrated momentum vectors
            rho_b.setZero();
            rho_f.setZero();

            double log_sum_weight_subtree = math::neg_inf<double>;

            int8_t v = 2 * direction_sampler(gen) - 1; // -1 or 1
            if (v == -1) {
                auto input = mcmc::TreeInput(
                    // position information to update
                    theta_bb_ad_expr, theta_bb, theta_bb_adj, tp_adj,
                    theta_prime, p_bb,
                    // momentum vectors to update
                    p_bf, p_bb, p_bf_scaled, p_bb_scaled, rho_b,
                    // stats to update to adapt step size at the end
                    n_leapfrog, log_sum_weight_subtree, sum_metro_prob,
                    // other miscellaneous variables
                    v, std::exp(step_adapter.log_eps), ham_prev
                );
                rho_f = rho;
                p_fb = p_bb;
                p_fb_scaled = p_bb_scaled;

                output = mcmc::build_tree(n_params, input, depth,
                                          unif_sampler, gen, momentum_handler,
                                          tree_cache.data());
            } else {
                auto input = mcmc::TreeInput(
                    // correct position information to update
                    theta_ff_ad_expr, theta_ff, theta_ff_adj, tp_adj,
                    theta_prime, p_ff,
                    // correct momentum vectors to update
                    p_fb, p_ff, p_fb_scaled, p_ff_scaled, rho_f,
                    // stats to update to adapt step size at the end
                    n_leapfrog, log_sum_weight_subtree, sum_metro_prob,
                    // other miscellaneous variables
                    v, std::exp(step_adapter.log_eps), ham_prev
                );
                rho_b = rho;
                p_bf = p_ff;
                p_bf_scaled = p_ff_scaled;

                output = mcmc::build_tree(n_params, input, depth,
                                          unif_sampler, gen, momentum_handler,
                                          tree_cache.data());
            }

            // early break if starting to U-Turn
            if (!output.valid) break;

            // if new subtree's weight is greater than previous subtree's weight
            // always accept!
            if (log_sum_weight_subtree > log_sum_weight) {
                theta_curr = theta_prime;
                potential_prev = output.potential;
            } else {
                double p = std::exp(log_sum_weight_subtree - log_sum_weight);
                if (mcmc::accept_or_reject(p, unif_sampler, gen)) {
                    theta_curr = theta_prime;
                    potential_prev = output.potential;
                }
            }

            // update total log_sum_weight
            log_sum_weight = math::lse(log_sum_weight, log_sum_weight_subtree);

            // check if proposals are still
            // - entroping in the full tree
            // - entroping from backwards-subtree to forwards-subtree
            // - entroping from forwards-subtree to backwards-subtree
            // This is a much stronger than the original paper's entropy condition.
            // This most likely reduces the depth to avoid unnecessary computation.

            rho = rho_b + rho_f;

            bool valid =
                mcmc::check_entropy(rho, p_bb_scaled, p_ff_scaled) &&
                mcmc::check_entropy(rho_b + p_fb, p_bb_scaled, p_fb_scaled) &&
                mcmc::check_entropy(p_bf + rho_f, p_bf_scaled, p_ff_scaled)
                ;

            if (!valid) break;

        } // end tree doubling for-loop

        // Warmup Adapt!
        if (i < config.warmup) {

            // epsilon dual averaging
            step_adapter.adapt(sum_metro_prob / static_cast<double>(n_leapfrog));

            // adapt variance only if adapting policy is diag_var or dense_var
            if constexpr (std::is_same_v<var_adapter_policy_t, diag_var> ||
                          std::is_same_v<var_adapter_policy_t, dense_var>) {
                const bool update = var_adapter.adapt(theta_curr, momentum_handler.get_m_inverse());
                if (update) {
                    double log_eps = std::log( mcmc::find_reasonable_epsilon(
                                        std::exp(step_adapter.log_eps),
                                        theta_curr_ad_expr, theta_curr,
                                        theta_curr_adj, tp_adj,
                                        gen, momentum_handler) );
                    step_adapter.reset();
                    step_adapter.init(log_eps);
                }
            }

            // if last warmup iteration
            if (i == config.warmup - 1) {
                step_adapter.log_eps = step_adapter.log_eps_bar;
            }
        }

        // store sample theta_curr only after burning
        if (i >= config.warmup) {
            res.cont_samples.row(i-config.warmup) = theta_curr;
        }

    } // end for-loop to sample 1 point

    // stop timing sampling
    stopwatch_sampling.stop();

    // save output results
    res.warmup_time = stopwatch_warmup.elapsed();
    res.sampling_time = stopwatch_sampling.elapsed();
}

} // namespace mcmc

template <class ExprType
        , class NUTSConfigType = NUTSConfig<>>
inline auto nuts(const ExprType& expr,
                 const NUTSConfigType& config = NUTSConfigType())
{
    return mcmc::base_mcmc(expr, config,
            [](auto& program, const auto& config,
               const auto& pack, auto& res) {
                res.name = "nuts";
                mcmc::nuts_(program, config, pack, res);
            });
}

} // namespace ppl


#ifndef WITHIN_BIGGLES_KALMAN_FILTER_HPP__
#error "This file should only be included by kalman_filter.hpp"
#endif // WITHIN_BIGGLES_KALMAN_FILTER_HPP__

#include <boost/assert.hpp>
#include <boost/foreach.hpp>
#include <cmath>
#include <deque>
#include <Eigen/Dense>

#include "detail/physics.hpp"
#include "tools/debug.hpp"

namespace biggles
{

template<typename OutputIterator>
void rts_smooth(const kalman_filter& filter,
                OutputIterator output_reversed_states_and_covariances)
{
    // Rauch-Tung-Striebel smoother:
    //
    // L_t = Sigma_{t|t} A^T P_{t+1|t}^{-1}
    //
    // mu_{t|T} = mu_{t|t} + L_t(mu_{t+1|T} - mu_{t+1|t})
    //
    // P_{t|T} = P_{t|t} + L_t(P_{t+1|T} - P_{t+1|t})L^T_t

    // the state evolution matrix: X_{k+1} = A X_k + <noise>
    Eigen::Matrix4f A(Eigen::Matrix4f::Identity());
    BOOST_ASSERT(filter.dynamic_drag() == 1);
    A(0,1) = A(2,3) = filter.dynamic_drag(); // <- off-diagonal velocity integration

    const kalman_filter::states_and_cov_deque& correction_states_and_covs(filter.corrections());
    const kalman_filter::states_and_cov_deque& prediction_states_and_covs(filter.predictions());

    // final interpolated state is output of forward filter, this is mu_{T|T} and Sigma_{T|T}
    if (not is_symmetric(correction_states_and_covs.back().second) and
            measure_asymmetry(correction_states_and_covs.back().second) > 0.f) {
        OK("foo");
        OK1(correction_states_and_covs.back().second);
        OK1(correction_states_and_covs.back().first);
        BOOST_ASSERT(measure_asymmetry(correction_states_and_covs.back().second)==0.f);
    }
    *output_reversed_states_and_covariances = correction_states_and_covs.back();
    ++output_reversed_states_and_covariances;

    // iterator pointing to mu_{t+1|t}, starting at t = T-1
    kalman_filter::states_and_cov_deque::const_reverse_iterator pred_it(prediction_states_and_covs.rbegin());

    // iterator pointing to mu_{t|t}, starting at t = T-1
    kalman_filter::states_and_cov_deque::const_reverse_iterator corr_it(++correction_states_and_covs.rbegin());

    for(kalman_filter::state_covariance_pair prior_pair(correction_states_and_covs.back());
        corr_it != correction_states_and_covs.rend();
        ++pred_it, ++corr_it, ++output_reversed_states_and_covariances)
    {
        // mu_{t+1|t} and Sigma_{t+1|t}
        const Eigen::Vector4f& pred_state(pred_it->first);
        const Eigen::Matrix4f& pred_cov(pred_it->second);

        // mu_{t|t} and Sigma_{t|t}
        const Eigen::Vector4f& corr_state(corr_it->first);
        const Eigen::Matrix4f& corr_cov(corr_it->second);

        // mu_{t+1|T} and Sigma_{t+1|T}
        const Eigen::Vector4f& last_smoothed_state(prior_pair.first);
        const Eigen::Matrix4f& last_smoothed_cov(prior_pair.second);

        // calculate the RTS smoothed values
        Eigen::Matrix4f L = corr_cov * A.transpose() * pred_cov.inverse();
        Eigen::Vector4f smoothed_state = corr_state + L * (last_smoothed_state - pred_state);
        Eigen::Matrix4f smoothed_cov = corr_cov + L * (last_smoothed_cov - pred_cov) * L.transpose();
        smoothed_cov = enforce_symmetry(smoothed_cov);

        if (not is_symmetric(corr_cov) and measure_asymmetry(corr_cov) > 0.f) {
            OK("rts_smooth");
            OK1(corr_cov);
            OK(measure_asymmetry(corr_cov));
            BOOST_ASSERT(measure_asymmetry(corr_cov)==0.f);
        }

        if (not is_symmetric(smoothed_cov) and measure_asymmetry(smoothed_cov) > 0.f) {
            OK("rts_smooth");
            OK1(smoothed_cov);
            OK(measure_asymmetry(corr_cov));
            OK(measure_asymmetry(smoothed_cov));
            BOOST_ASSERT(measure_asymmetry(smoothed_cov)==0.f);
        }

        // record smoothed state and covariance
        prior_pair = kalman_filter::state_covariance_pair(smoothed_state, smoothed_cov);
        *output_reversed_states_and_covariances = prior_pair;
    }
}

}


// This file is part of SVO - Semi-direct Visual Odometry.
//
// Copyright (C) 2014 Christian Forster <forster at ifi dot uzh dot ch>
// (Robotics and Perception Group, University of Zurich, Switzerland).
//
// This file is subject to the terms and conditions defined in the file
// 'LICENSE', which is part of this source code package.

#include <svo/direct/feature_alignment.h>

#ifdef __SSE2__
# include <emmintrin.h>
#endif
#ifdef __ARM_NEON__
# include <arm_neon.h>
#endif

#include <glog/logging.h>
#include <Eigen/Dense>
#include <opencv2/highgui/highgui.hpp>
#include <opencv2/imgproc/imgproc.hpp>
#include <svo/direct/patch_utils.h>

namespace svo {
namespace feature_alignment {

#define SUBPIX_VERBOSE 0
#define SVO_DISPLAY_ALIGN_1D 0

//------------------------------------------------------------------------------
bool align1D(
    const cv::Mat& cur_img,
    const Eigen::Ref<GradientVector>& dir,                  // direction in which the patch is allowed to move
    uint8_t* ref_patch_with_border,
    uint8_t* ref_patch,
    const int n_iter,
    const bool affine_est_offset,
    const bool affine_est_gain,
    Keypoint* cur_px_estimate,
    double* h_inv)
{
  CHECK_NOTNULL(cur_px_estimate);

  constexpr int kHalfPatchSize = 4;
  constexpr int kPatchSize = 2 * kHalfPatchSize;
  constexpr int kPatchArea = kPatchSize * kPatchSize;
  bool converged = false;

  // We optimize feature position and two affine parameters.
  // Compute derivative of template and prepare inverse compositional.
  float __attribute__((__aligned__(16))) ref_patch_dv[kPatchArea];
  Eigen::Matrix3f H = Eigen::Matrix3f::Zero(3, 3);

  // Compute gradient and hessian.
  constexpr int ref_step = kPatchSize + 2;
  float* it_dv = ref_patch_dv;
  for(int y=0; y<kPatchSize; ++y)
  {
    uint8_t* it = ref_patch_with_border + (y+1)*ref_step + 1;
    for(int x=0; x<kPatchSize; ++x, ++it, ++it_dv)
    {
      Eigen::Vector3f J;
      const float dx = static_cast<float>(it[1]) - static_cast<float>(it[-1]);
      const float dy = static_cast<float>(it[ref_step]) - static_cast<float>(it[-ref_step]);
      J[0] = 0.5f * (dir(0) * dx + dir(1) * dy);

      // If not using the affine compensation, set the jacobian be zero.
      // In this way, all the blocks related to affine parameters will be zero.
      J[1] = affine_est_offset? 1.0f : 0.0f;
      J[2] = affine_est_gain? -1.0f*it[0] : 0.0f;

      *it_dv = J[0];
      H += J*J.transpose();
    }
  }
  // If not use affine compensation, force update to be zero by
  // * setting the affine parameter block in H to identity
  // * setting the residual block to zero (see below)
  if(!affine_est_offset)
  {
    H(1, 1) = 1.0;
  }
  if(!affine_est_gain)
  {
    H(2, 2) = 1.0;
  }

  if(h_inv)
    *h_inv = 1.0/H(0,0)*kPatchSize*kPatchSize;
  Eigen::Matrix3f Hinv = H.inverse();
  float mean_diff = 0;
  float alpha = 1.0;

  // Compute pixel location in new image:
  float u = cur_px_estimate->x();
  float v = cur_px_estimate->y();

  // termination condition
  const float min_update_squared = 0.03*0.03;
  const int cur_step = cur_img.step.p[0];
  #if SVO_DISPLAY_ALIGN_1D
  cv::Mat res_patch(kPatchSize, kPatchSize, CV_32FC1);
  cv::Mat cur_patch(kPatchSize, kPatchSize, CV_32FC1);
  #endif
  for(int iter = 0; iter<n_iter; ++iter)
  {
    int u_r = std::floor(u);
    int v_r = std::floor(v);
    if(u_r < kHalfPatchSize
       || v_r < kHalfPatchSize
       || u_r >= cur_img.cols - kHalfPatchSize
       || v_r >= cur_img.rows - kHalfPatchSize)
      break;

    if(std::isnan(u) || std::isnan(v)) // TODO very rarely this can happen, maybe H is singular? should not be at corner.. check
      return false;

    // compute interpolation weights
    float subpix_x = u-u_r;
    float subpix_y = v-v_r;
    float wTL = (1.0-subpix_x)*(1.0-subpix_y);
    float wTR = subpix_x * (1.0-subpix_y);
    float wBL = (1.0-subpix_x)*subpix_y;
    float wBR = subpix_x * subpix_y;

    // loop through search_patch, interpolate
    uint8_t* it_ref = ref_patch;
    float* it_ref_dv = ref_patch_dv;
    float new_chi2 = 0.0;
    Eigen::Vector3f Jres = Eigen::Vector3f::Zero();
    for(int y=0; y<kPatchSize; ++y)
    {
      uint8_t* it = (uint8_t*) cur_img.data +
          (v_r+y-kHalfPatchSize)*cur_step + u_r-kHalfPatchSize;
      for(int x=0; x<kPatchSize; ++x, ++it, ++it_ref, ++it_ref_dv)
      {
        float cur_intensity =
            wTL*it[0] + wTR*it[1] + wBL*it[cur_step] + wBR*it[cur_step+1];
        float res = cur_intensity - alpha*(*it_ref) + mean_diff;
        Jres[0] -= res*(*it_ref_dv);

        // If affine compensation is used,
        // set Jres with respect to affine parameters.
        if(affine_est_offset)
        {
          Jres[1] -= res;
        }
        if(affine_est_gain)
        {
          Jres[2] -= (-1)*res*(*it_ref);
        }
        new_chi2 += res*res;
        #if SVO_DISPLAY_ALIGN_1D
        res_patch.at<float>(y,x) = res;
        cur_patch.at<float>(y,x) = cur_intensity;
        #endif
      }
    }
    // If not using affine compensation, force update to be zero.
    if(!affine_est_offset)
    {
      Jres[1] = 0.0;
    }
    if(!affine_est_gain)
    {
      Jres[2] = 0.0;
    }

    Eigen::Vector3f update = Hinv * Jres;
    u += update[0]*dir[0];
    v += update[0]*dir[1];
    mean_diff += update[1];
    alpha += update[2];

    VLOG(300) << "It. " << iter << ": \t"
              << "\t u=" << u << ", v=" << v
              << "\t update = "
              << update[0] << ", " << update[1] << ", " << update[2]
              << "\t new chi2 = " << new_chi2;

    #if SVO_DISPLAY_ALIGN_1D
    cv::Mat res_patch_normalized, cur_patch_normalized, ref_patch_img, ref_patch_normalized;
    patch_utils::normalizeAndUpsamplePatch(res_patch, 8, &res_patch_normalized);
    patch_utils::normalizeAndUpsamplePatch(cur_patch, 8, &cur_patch_normalized);
    patch_utils::patchToMat(ref_patch, kPatchSize, &ref_patch_img);
    patch_utils::normalizeAndUpsamplePatch(ref_patch_img, 8, &ref_patch_normalized);
    cv::Mat concatenated;
    patch_utils::concatenatePatches({ref_patch_normalized,
                                     cur_patch_normalized,
                                     res_patch_normalized}, &concatenated);
    cv::line(concatenated,
             cv::Point2f(concatenated.rows/2, concatenated.rows/2),
             cv::Point2f(concatenated.rows/2+15*dir(0), concatenated.rows/2 +15*dir(1)),
             cv::Scalar(0,0,255), 3);
    cv::imshow("concatenated", concatenated);
    cv::waitKey(0);
    #endif

    if(update[0]*update[0] < min_update_squared)
    {
      VLOG(300) << "converged.";
      converged = true;
      break;
    }
  }

  *cur_px_estimate << u, v;
  return converged;
}

//------------------------------------------------------------------------------
bool align2D(
    const cv::Mat& cur_img,
    uint8_t* ref_patch_with_border,
    uint8_t* ref_patch,
    const int n_iter,
    const bool affine_est_offset,
    const bool affine_est_gain,
    Keypoint& cur_px_estimate,
    bool no_simd,
    std::vector<Eigen::Vector2f> *each_step)
{
#ifdef __ARM_NEON__
  if(!no_simd)
    return align2D_NEON(cur_img, ref_patch_with_border,
                        ref_patch, n_iter, cur_px_estimate);
#endif

  if(each_step) each_step->clear();

  const int halfpatch_size_ = 4;
  const int patch_size_ = 8;
  const int patch_area_ = 64;
  bool converged=false;

  // We optimize feature position and two affine parameters.
  // compute derivative of template and prepare inverse compositional
  float __attribute__((__aligned__(16))) ref_patch_dx[patch_area_];
  float __attribute__((__aligned__(16))) ref_patch_dy[patch_area_];
  Eigen::Matrix4f H; H.setZero();

  // compute gradient and hessian
  const int ref_step = patch_size_+2;
  float* it_dx = ref_patch_dx;
  float* it_dy = ref_patch_dy;
  for(int y=0; y<patch_size_; ++y)
  {
    uint8_t* it = ref_patch_with_border + (y+1)*ref_step + 1;
    for(int x=0; x<patch_size_; ++x, ++it, ++it_dx, ++it_dy)
    {
      Eigen::Vector4f J;
      J[0] = 0.5 * (it[1] - it[-1]);
      J[1] = 0.5 * (it[ref_step] - it[-ref_step]);

      // If not using the affine compensation, force the jacobian to be zero.
      J[2] = affine_est_offset? 1.0 : 0.0;
      J[3] = affine_est_gain? -1.0*it[0]: 0.0;

      *it_dx = J[0];
      *it_dy = J[1];
      H += J*J.transpose();
    }
  }
  // If not use affine compensation, force update to be zero by
  // * setting the affine parameter block in H to identity
  // * setting the residual block to zero (see below)
  if(!affine_est_offset)
  {
    H(2, 2) = 1.0;
  }
  if(!affine_est_gain)
  {
    H(3, 3) = 1.0;
  }
  Eigen::Matrix4f Hinv = H.inverse();
  float mean_diff = 0;
  float alpha = 1.0;

  // Compute pixel location in new image:
  float u = cur_px_estimate.x();
  float v = cur_px_estimate.y();

  if(each_step) each_step->push_back(Eigen::Vector2f(u, v));

  // termination condition
  const float min_update_squared = 0.03*0.03; // TODO I suppose this depends on the size of the image (ate)
  const int cur_step = cur_img.step.p[0];
  //float chi2 = 0;
  Eigen::Vector4f update; update.setZero();
  for(int iter = 0; iter<n_iter; ++iter)
  {
    int u_r = std::floor(u);
    int v_r = std::floor(v);
    if(u_r < halfpatch_size_
       || v_r < halfpatch_size_
       || u_r >= cur_img.cols-halfpatch_size_
       || v_r >= cur_img.rows-halfpatch_size_)
      break;

    if(std::isnan(u) || std::isnan(v)) // TODO very rarely this can happen, maybe H is singular? should not be at corner.. check
      return false;

    // compute interpolation weights
    float subpix_x = u-u_r;
    float subpix_y = v-v_r;
    float wTL = (1.0-subpix_x)*(1.0-subpix_y);
    float wTR = subpix_x * (1.0-subpix_y);
    float wBL = (1.0-subpix_x)*subpix_y;
    float wBR = subpix_x * subpix_y;

    // loop through search_patch, interpolate
    uint8_t* it_ref = ref_patch;
    float* it_ref_dx = ref_patch_dx;
    float* it_ref_dy = ref_patch_dy;
    //float new_chi2 = 0.0;
    Eigen::Vector4f Jres; Jres.setZero();
    for(int y=0; y<patch_size_; ++y)
    {
      uint8_t* it = (uint8_t*) cur_img.data + (v_r+y-halfpatch_size_)*cur_step + u_r-halfpatch_size_;
      for(int x=0; x<patch_size_; ++x, ++it, ++it_ref, ++it_ref_dx, ++it_ref_dy)
      {
        float search_pixel = wTL*it[0] + wTR*it[1] + wBL*it[cur_step] + wBR*it[cur_step+1];
        float res = search_pixel - alpha*(*it_ref) + mean_diff;
        Jres[0] -= res*(*it_ref_dx);
        Jres[1] -= res*(*it_ref_dy);

        // If affine compensation is used,
        // set Jres with respect to affine parameters.
        if(affine_est_offset)
        {
          Jres[2] -= res;
        }

        if(affine_est_gain)
        {
          Jres[3] -= (-1)*res*(*it_ref);
        }
        //new_chi2 += res*res;
      }
    }
    // If not use affine compensation, force update to be zero.
    if(!affine_est_offset)
    {
      Jres[2] = 0.0;
    }
    if(!affine_est_gain)
    {
      Jres[3] = 0.0;
    }
    /*
    if(iter > 0 && new_chi2 > chi2)
    {
#if SUBPIX_VERBOSE
      cout << "error increased." << endl;
#endif
      u -= update[0];
      v -= update[1];
      break;
    }
    chi2 = new_chi2;
*/
    update = Hinv * Jres;
    u += update[0];
    v += update[1];
    mean_diff += update[2];
    alpha += update[3];

    if(each_step) each_step->push_back(Eigen::Vector2f(u, v));

#if SUBPIX_VERBOSE
    cout << "Iter " << iter << ":"
         << "\t u=" << u << ", v=" << v
         << "\t update = " << update[0] << ", " << update[1]
         << "\t new chi2 = " << new_chi2 << endl;
#endif

    if(update[0]*update[0]+update[1]*update[1] < min_update_squared)
    {
#if SUBPIX_VERBOSE
      cout << "converged." << endl;
#endif
      converged=true;
      break;
    }
  }

  cur_px_estimate << u, v;
  (void)no_simd;

  return converged;
}

#define  DESCALE(x,n)     (((x) + (1 << ((n)-1))) >> (n)) // rounds to closest integer and descales

//------------------------------------------------------------------------------
bool align2D_SSE2(
    const cv::Mat& cur_img,
    uint8_t* ref_patch_with_border,
    uint8_t* ref_patch,
    const int n_iter,
    Keypoint& cur_px_estimate)
{
  // TODO: This function should not be used as the alignment is not robust to illumination changes!
  const int halfpatch_size = 4;
  const int patch_size = 8;
  const int patch_area = 64;
  bool converged=false;
  const int W_BITS = 14;

  // compute derivative of template and prepare inverse compositional
  int16_t __attribute__((__aligned__(16))) ref_patch_dx[patch_area];
  int16_t __attribute__((__aligned__(16))) ref_patch_dy[patch_area];

  // compute gradient and hessian
  const int ref_step = patch_size+2;
  int16_t* it_dx = ref_patch_dx;
  int16_t* it_dy = ref_patch_dy;
  float A11=0, A12=0, A22=0;
  for(int y=0; y<patch_size; ++y)
  {
    uint8_t* it = ref_patch_with_border + (y+1)*ref_step + 1;
    for(int x=0; x<patch_size; ++x, ++it, ++it_dx, ++it_dy)
    {
      int16_t dx = static_cast<int16_t>(it[1]) - it[-1];
      int16_t dy = static_cast<int16_t>(it[ref_step]) - it[-ref_step];
      *it_dx = dx;
      *it_dy = dy;  // we are missing a factor 1/2
      A11 += static_cast<float>(dx*dx); // we are missing a factor 1/4
      A12 += static_cast<float>(dx*dy);
      A22 += static_cast<float>(dy*dy);
    }
  }

  // Compute pixel location in new image:
  float u = cur_px_estimate.x();
  float v = cur_px_estimate.y();

  // termination condition
  const float min_update_squared = 0.03*0.03;
#ifdef __SSE2__
  const int cur_step = cur_img.step.p[0];
#endif
  const float Dinv = 1.0f/(A11*A22 - A12*A12); // we are missing an extra factor 16
  float chi2 = 0;
  float update_u = 0, update_v = 0;

  for(int iter = 0; iter<n_iter; ++iter)
  {
    int u_r = std::floor(u);
    int v_r = std::floor(v);
    if(u_r < halfpatch_size || v_r < halfpatch_size || u_r >= cur_img.cols-halfpatch_size || v_r >= cur_img.rows-halfpatch_size)
      break;

    if(std::isnan(u) || std::isnan(v)) // TODO very rarely this can happen, maybe H is singular? should not be at corner.. check
      return false;

    float subpix_x = u-u_r;
    float subpix_y = v-v_r;
    float b1=0, b2=0;
    float new_chi2 = 0.0;

#ifdef __SSE2__
    // compute bilinear interpolation weights
    int wTL = static_cast<int>((1.0f-subpix_x)*(1.0f-subpix_y)*(1 << W_BITS));
    int wTR = static_cast<int>(subpix_x * (1.0f-subpix_y)*(1 << W_BITS));
    int wBL = static_cast<int>((1.0f-subpix_x)*subpix_y*(1 << W_BITS));
    int wBR = (1 << W_BITS) - wTL - wTR - wBL;

    __m128i qw0 = _mm_set1_epi32(wTL + (wTR << 16)); // Sets the 4 signed 32-bit integer values to [wTL, wTR].
    __m128i qw1 = _mm_set1_epi32(wBL + (wBR << 16));
    __m128i z = _mm_setzero_si128();
    __m128 qb0 = _mm_setzero_ps(); // 4 floats
    __m128 qb1 = _mm_setzero_ps(); // 4 floats
    __m128i qdelta = _mm_set1_epi32(1 << (W_BITS-1));
    for(int y=0; y<patch_size; ++y)
    {
      const uint8_t* it  = (const uint8_t*) cur_img.data + (v_r+y-halfpatch_size)*cur_step + u_r-halfpatch_size;

      // Iptr is aligned!
      //__m128i diff0 = _mm_load_si128((const __m128i*)(ref_patch + y*8));
      __m128i diff = _mm_unpacklo_epi8(_mm_loadl_epi64((const __m128i*)(ref_patch + y*8)), z);

      // load the lower 64 bits and unpack [8u 0 8u 0..]
      __m128i v00 = _mm_unpacklo_epi8(_mm_loadl_epi64((const __m128i*)(it)), z);
      __m128i v01 = _mm_unpacklo_epi8(_mm_loadl_epi64((const __m128i*)(it + 1)), z);
      __m128i v10 = _mm_unpacklo_epi8(_mm_loadl_epi64((const __m128i*)(it + cur_step)), z);
      __m128i v11 = _mm_unpacklo_epi8(_mm_loadl_epi64((const __m128i*)(it + cur_step + 1)), z);

      // interpolate top row and bottom row
      // _mm_unpacklo_epi16: Interleaves the lower 4 signed or unsigned 16-bit integers in a with the lower 4 signed or unsigned 16-bit integers in b.
      // _mm_madd_epi16: Multiplies the 8 signed 16-bit integers from a by the 8 signed 16-bit integers from b.
      __m128i t0 = _mm_add_epi32(_mm_madd_epi16(_mm_unpacklo_epi16(v00, v01), qw0),
                                 _mm_madd_epi16(_mm_unpacklo_epi16(v10, v11), qw1));
      __m128i t1 = _mm_add_epi32(_mm_madd_epi16(_mm_unpackhi_epi16(v00, v01), qw0),
                                 _mm_madd_epi16(_mm_unpackhi_epi16(v10, v11), qw1));

      // _mm_srai_epi32: Shifts the 4 signed 32-bit integers in A right by count bits while shifting in the sign bit.
      t0 = _mm_srai_epi32(_mm_add_epi32(t0, qdelta), W_BITS); // adding qdelta is for rounding closest int
      t1 = _mm_srai_epi32(_mm_add_epi32(t1, qdelta), W_BITS);

      // compute 8xres:
      // _mm_packs_epi32: Packs the 8 signed 32-bit integers from a and b into signed 16-bit integers and saturates.
      diff = _mm_subs_epi16(_mm_packs_epi32(t0, t1), diff);

      // load gradient dX and dY, both are aligned!
      v00 = _mm_load_si128((const __m128i*)(ref_patch_dx + y*patch_size)); // [dx1, dx2, dx3, dx4 ...]
      v01 = _mm_load_si128((const __m128i*)(ref_patch_dy + y*patch_size)); // [dy1, dy2, dy3, dy4 ...]

      // _mm_mulhi_epi16: Multiplies the 8 signed 16-bit integers from a by the 8 signed 16-bit integers from b.
      v10 = _mm_mullo_epi16(v00, diff); // Packs the lower 16 bits of the 8 signed 32-bit results. [15:0]
      v11 = _mm_mulhi_epi16(v00, diff); // Packs the upper 16 bits of the 8 signed 32-bit results. [31:16]

      // _mm_unpacklo_epi16: Interleaves the lower 4 signed or unsigned 16-bit integers with the lower 4 signed or unsigned 16-bit integers in b.
      v00 = _mm_unpacklo_epi16(v10, v11);
      v10 = _mm_unpackhi_epi16(v10, v11);

      // convert to float and add to dx
      qb0 = _mm_add_ps(qb0, _mm_cvtepi32_ps(v00));
      qb0 = _mm_add_ps(qb0, _mm_cvtepi32_ps(v10));

      // same with dY
      v10 = _mm_mullo_epi16(v01, diff);
      v11 = _mm_mulhi_epi16(v01, diff);
      v00 = _mm_unpacklo_epi16(v10, v11);
      v10 = _mm_unpackhi_epi16(v10, v11);
      qb1 = _mm_add_ps(qb1, _mm_cvtepi32_ps(v00));
      qb1 = _mm_add_ps(qb1, _mm_cvtepi32_ps(v10));
    }

    float __attribute__((__aligned__(16))) buf[4];
    _mm_store_ps(buf, qb0);
    b1 += buf[0]+buf[1]+buf[2]+buf[3];
    _mm_store_ps(buf, qb1);
    b2 += buf[0]+buf[1]+buf[2]+buf[3];
#endif

    // compute -A^-1*b
    update_u = ((A12*b2 - A22*b1) * Dinv)*2; // * 2 to compensate because above, we did not compute the derivative correctly
    update_v = ((A12*b1 - A11*b2) * Dinv)*2;
    u += update_u;
    v += update_v;

#if SUBPIX_VERBOSE
    cout << "Iter " << iter << ":"
         << "\t u=" << u << ", v=" << v
         << "\t update = " << update_u << ", " << update_v
         << "\t new chi2 = " << new_chi2 << endl;
#endif

    if(update_u*update_u+update_v*update_v < min_update_squared)
    {
#if SUBPIX_VERBOSE
      cout << "converged." << endl;
#endif
      converged=true;
      break;
    }
    chi2 = new_chi2;
    (void)chi2;
    (void)subpix_x;
    (void)subpix_y;
    (void)ref_patch;
  }

  cur_px_estimate << u, v;
  (void)W_BITS;
  return converged;
}

//------------------------------------------------------------------------------
bool align2D_NEON (const cv::Mat& cur_img,
                   uint8_t* ref_patch_with_border,
                   uint8_t* ref_patch,
                   const int n_iter,
                   Keypoint& cur_px_estimate)
{
  const int halfpatch_size = 4;
  const int patch_size = 8;
  const int patch_area = 64;
  bool converged=false;
  const int W_BITS = 14;

  // compute derivative of template and prepare inverse compositional
  int16_t __attribute__((__aligned__(16))) ref_patch_dx[patch_area];
  int16_t __attribute__((__aligned__(16))) ref_patch_dy[patch_area];

  // compute gradient and hessian
  const int ref_step = patch_size+2;
  int16_t* it_dx = ref_patch_dx;
  int16_t* it_dy = ref_patch_dy;
  Eigen::Matrix3f H; H.setZero();
  for(int y=0; y<patch_size; ++y)
  {
    uint8_t* it = ref_patch_with_border + (y+1)*ref_step + 1;
    for(int x=0; x<patch_size; ++x, ++it, ++it_dx, ++it_dy)
    {
      *it_dx = static_cast<int16_t>(it[1]) - it[-1];
      *it_dy = static_cast<int16_t>(it[ref_step]) - it[-ref_step]; // divide by 2 missing
      Eigen::Vector3f J(*it_dx, *it_dy, 1.0f);
      H += J*J.transpose();
    }
  }
  Eigen::Matrix3f Hinv = H.inverse();
  float mean_diff = 0.0;

  // Compute pixel location in new image:
  float u = cur_px_estimate.x();
  float v = cur_px_estimate.y();

  // termination condition
  const float min_update_squared = 0.03*0.03;
  const int cur_step = cur_img.step.p[0];
  Eigen::Vector3f update;
  Eigen::Vector3f Jres;
  for(int iter = 0; iter<n_iter; ++iter)
  {
    const int u_r = std::floor(u);
    const int v_r = std::floor(v);
    if(u_r < halfpatch_size || v_r < halfpatch_size || u_r >= cur_img.cols-halfpatch_size || v_r >= cur_img.rows-halfpatch_size)
      break;

    if(std::isnan(u) || std::isnan(v)) // TODO very rarely this can happen, maybe H is singular? should not be at corner.. check
      return false;

    const float subpix_x = u-u_r;
    const float subpix_y = v-v_r;
    float b1=0, b2=0;

#ifdef __ARM_NEON__
    const int SHIFT_BITS = 7;
    const uint16_t wTL = static_cast<uint16_t>((1.0f-subpix_x)*(1.0f-subpix_y)*(1<<SHIFT_BITS));
    const uint16_t wTR = static_cast<uint16_t>(subpix_x*(1.0f-subpix_y)*(1<<SHIFT_BITS));
    const uint16_t wBL = static_cast<uint16_t>((1.0f-subpix_x)*subpix_y*(1<<SHIFT_BITS));
    const uint16_t wBR = static_cast<uint16_t>((1 << SHIFT_BITS) - wTL - wTR - wBL);

    // initialize result to zero
    int32x4_t vb1 = vdupq_n_s32(0);
    int32x4_t vb2 = vdupq_n_s32(0);
    int16x8_t vmean_diff = vdupq_n_s16( (int16_t) (mean_diff+0.5) );
    int16x8_t vres_sum = vdupq_n_s16(0);
    for(int y=0; y<patch_size; ++y)
    {
      const uint8_t* it  = (const uint8_t*) cur_img.data + (v_r+y-halfpatch_size)*cur_step + u_r-halfpatch_size;

      // load and convert from uint8 to uint16
      uint16x8_t v00 = vmovl_u8( vld1_u8( it ) );
      uint16x8_t v01 = vmovl_u8( vld1_u8( it + 1 ) );
      uint16x8_t v10 = vmovl_u8( vld1_u8( it + cur_step ) );
      uint16x8_t v11 = vmovl_u8( vld1_u8( it + cur_step + 1 ) );

      // vector multiply by scalar
      v00 = vmulq_n_u16( v00, wTL );
      v01 = vmulq_n_u16( v01, wTR );
      v10 = vmulq_n_u16( v10, wBL );
      v11 = vmulq_n_u16( v11, wBR );

      // add all results together
      v00 = vaddq_u16( v00, vaddq_u16( v01, vaddq_u16( v10, v11 ) ) );

      // descale: shift right by constant
      v00 = vrshrq_n_u16(v00, SHIFT_BITS);

      // compute difference between reference and interpolated patch,
      // use reinterpet-cast to make signed [-255,255]
      int16x8_t res = vsubq_s16(vreinterpretq_s16_u16(v00), vreinterpretq_s16_u16(vmovl_u8(vld1_u8( ref_patch + y*8 ))));

      // correct res with mean difference
      res = vaddq_s16(res, vmean_diff);

      // compute sum of the residual
      vres_sum = vaddq_s16(vres_sum, res);

      // Eigen::Vector multiply accumulate long: vmla -> Vr[i] := Va[i] + Vb[i] * Vc[i]
      // int32x4_t  vmlal_s16(int32x4_t a, int16x4_t b, int16x4_t c);    // VMLAL.S16 q0,d0,d0
      int16x8_t grad = vld1q_s16(ref_patch_dx + y*patch_size);
      vb1 = vmlal_s16(vb1, vget_low_s16(grad), vget_low_s16(res));
      vb1 = vmlal_s16(vb1, vget_high_s16(grad), vget_high_s16(res));

      grad = vld1q_s16(ref_patch_dy + y*patch_size);
      vb2 = vmlal_s16(vb2, vget_low_s16(grad), vget_low_s16(res));
      vb2 = vmlal_s16(vb2, vget_high_s16(grad), vget_high_s16(res));
    }

    // finally, sum results of vb1, vb2 and vres_sum
    int32x2_t tmp;
    tmp = vpadd_s32(vget_low_s32(vb1), vget_high_s32(vb1));
    Jres[0] = -vget_lane_s32(tmp, 0) - vget_lane_s32(tmp, 1);

    tmp = vpadd_s32(vget_low_s32(vb2), vget_high_s32(vb2));
    Jres[1] = -vget_lane_s32(tmp, 0) - vget_lane_s32(tmp, 1);

    int32x4_t vres_sum1 = vpaddlq_s16(vres_sum);
    tmp = vpadd_s32(vget_low_s32(vres_sum1), vget_high_s32(vres_sum1));
    Jres[2] = -vget_lane_s32(tmp, 0) - vget_lane_s32(tmp, 1);
#endif

    update = Hinv * Jres * 2; // * 2 to compensate because above, we did not compute the derivative correctly
    u += update[0];
    v += update[1];
    mean_diff += update[2];

#if SUBPIX_VERBOSE
    cout << "Iter " << iter << ":"
         << "\t u=" << u << ", v=" << v
         << "\t update = " << update[0] << ", " << update[1] << endl;
#endif

    if(update[0]*update[0]+update[1]*update[1] < min_update_squared)
    {
#if SUBPIX_VERBOSE
      cout << "converged." << endl;
#endif
      converged=true;
      break;
    }

    (void)subpix_x;
    (void)subpix_y;
    (void)b1;
    (void)b2;
    (void)cur_step;
  }

  cur_px_estimate << u, v;

  (void)W_BITS;
  (void)ref_patch;
  return converged;
}

//------------------------------------------------------------------------------
void alignPyr2DVec(
    const std::vector<cv::Mat>& img_pyr_ref,
    const std::vector<cv::Mat>& img_pyr_cur,
    const int max_level,
    const int min_level,
    const std::vector<int>& patch_sizes,
    const int n_iter,
    const float min_update_squared,
    const std::vector<cv::Point2f>& px_ref,
    std::vector<cv::Point2f>& px_cur,
    std::vector<uint8_t>& status)
{
  for(size_t i=0; i<px_ref.size(); ++i)
  {
    Eigen::Vector2i px_ref_level_0(px_ref[i].x, px_ref[i].y);
    Keypoint px_cur_level_0(px_cur[i].x, px_cur[i].y);
    bool res = alignPyr2D(
          img_pyr_ref, img_pyr_cur, max_level, min_level, patch_sizes,
          n_iter, min_update_squared, px_ref_level_0, px_cur_level_0);
    status[i] = res ? 1 : 0;
    px_cur[i] = cv::Point2f(px_cur_level_0[0], px_cur_level_0[1]);

    // TODO: add subpixel of rounding px_ref_level_0.
    //       typically not a problem because ref comes directly from feature
    //       extraction and is int.
  }
}

//------------------------------------------------------------------------------
bool alignPyr2D(
    const std::vector<cv::Mat>& img_pyr_ref,
    const std::vector<cv::Mat>& img_pyr_cur,
    const int max_level,
    const int min_level,
    const std::vector<int>& patch_sizes,
    const int n_iter,
    const float min_update_squared,
    const Eigen::Vector2i& px_ref_level_0,
    Keypoint& px_cur_level_0)
{
#define SVO_PYRAMIDAL_VERBOSE 0
#define SVO_DESCALE(x,n)     (((x) + (1 << ((n)-1))) >> (n))

#if SVO_PYRAMIDAL_VERBOSE
  std::cout << "--------" << std::endl;
#endif

  // compute derivative of template and prepare inverse compositional
  const int max_patch_area = patch_sizes[0]*patch_sizes[0];
  uint8_t __attribute__((__aligned__(16))) ref_patch[max_patch_area];
  int16_t __attribute__((__aligned__(16))) ref_patch_dx[max_patch_area];
  int16_t __attribute__((__aligned__(16))) ref_patch_dy[max_patch_area];
  bool converged = false;

  for(int level = max_level; level >= min_level; --level)
  {
    const int patch_size = patch_sizes[level];
    CHECK(patch_size % 8 == 0)
        << ": alignPyr2D is only made for patch sizes multiples of 8!";
    const int halfpatch_size = patch_size/2;
    const int scale = (1<<level);
    const cv::Mat& img_ref = img_pyr_ref[level];
    const cv::Mat& img_cur = img_pyr_cur[level];
    const int width = img_ref.cols;
    const int height = img_ref.rows;
    const int step = img_ref.step; // should be number of bytes of each row
    const Eigen::Vector2f px_ref_flt = px_ref_level_0.cast<float>() / scale - Eigen::Vector2f(halfpatch_size, halfpatch_size);
    const Eigen::Vector2i px_ref = px_ref_flt.cast<int>();
    const Eigen::Vector2f px_ref_offset = px_ref_flt - px_ref.cast<float>();

    if(px_ref[0] < 1 || px_ref[1] < 1
       || px_ref[0] >= width-patch_size-1 || px_ref[1] >= height-patch_size-1)
    {
#if SVO_PYRAMIDAL_VERBOSE
      std::cout << "reference pixel is too close to the border" << std::endl;
#endif
      continue;
    }

    // compute gradient and hessian
    uint8_t* it_patch = ref_patch;
    int16_t* it_dx = ref_patch_dx;
    int16_t* it_dy = ref_patch_dy;
    Eigen::Matrix2f H; H.setZero();
    for(int y=0; y<patch_size; ++y)
    {
      uint8_t* it = img_ref.data + (px_ref[1]+y)*step + (px_ref[0]);
      for(int x=0; x<patch_size; ++x, ++it, ++it_patch, ++it_dx, ++it_dy)
      {
        *it_patch = *it;
        *it_dx = static_cast<int16_t>(it[1]) - it[-1];
        *it_dy = static_cast<int16_t>(it[step]) - it[-step]; // divide by 2 missing
        Eigen::Vector2f J(*it_dx, *it_dy);
        H += J*J.transpose();
      }
    }
    Eigen::Matrix2f Hinv = H.inverse();

    // Compute pixel location in new image:
    float u = px_cur_level_0[0] / scale - halfpatch_size - px_ref_offset[0];
    float v = px_cur_level_0[1] / scale - halfpatch_size - px_ref_offset[1];
    Eigen::Vector2f update; update.setZero();
    bool go_to_next_level = false;
    const int SHIFT_BITS = 7;
    converged = false; // reset at every level
    for(int iter = 0; iter<n_iter; ++iter)
    {
      if(std::isnan(u) || std::isnan(v)) // TODO very rarely this can happen, maybe H is singular? should not be at corner.. check
      {
#if SVO_PYRAMIDAL_VERBOSE
        std::cout << "update is NaN" << std::endl;
#endif
        return false;
      }

      go_to_next_level = false;
      const int u_r = std::floor(u);
      const int v_r = std::floor(v);
      if(u_r < 0 || v_r < 0 || u_r >= width-patch_size || v_r >= height-patch_size)
      {
#if SVO_PYRAMIDAL_VERBOSE
        std::cout << "current patch is too close to the border "
                  << "u,v: " << u_r << ", " << v_r
                  << "\t w,h:" << width << ", " << height
                  << "\t patch_size = " << patch_size << std::endl;
#endif
        go_to_next_level = true;
        break;
      }

      // compute interpolation weights
      const float subpix_x = u-u_r;
      const float subpix_y = v-v_r;
      const uint16_t wTL = static_cast<uint16_t>((1.0f-subpix_x)*(1.0f-subpix_y)*(1<<SHIFT_BITS));
      const uint16_t wTR = static_cast<uint16_t>(subpix_x*(1.0f-subpix_y)*(1<<SHIFT_BITS));
      const uint16_t wBL = static_cast<uint16_t>((1.0f-subpix_x)*subpix_y*(1<<SHIFT_BITS));
      const uint16_t wBR = (1 << SHIFT_BITS) - wTL - wTR - wBL;

      // loop through search_patch, interpolate
      uint8_t* it_ref = ref_patch;
      Eigen::Vector2f Jres; Jres.setZero();

#ifdef __ARM_NEON__
      // initialize result to zero
      int32x4_t vb1 = vdupq_n_s32(0);
      int32x4_t vb2 = vdupq_n_s32(0);
      for(int y=0; y<patch_size; ++y)
      {
        uint8_t* it = img_cur.data + (v_r+y)*step + (u_r);
        for(int x=0; x<patch_size; x+=8, it+=8, it_ref+=8)
        {
          // load and convert from uint8 to uint16
          uint16x8_t v00 = vmovl_u8( vld1_u8( it ) );
          uint16x8_t v01 = vmovl_u8( vld1_u8( it + 1 ) );
          uint16x8_t v10 = vmovl_u8( vld1_u8( it + step ) );
          uint16x8_t v11 = vmovl_u8( vld1_u8( it + step + 1 ) );

          // vector multiply by scalar
          v00 = vmulq_n_u16( v00, wTL );
          v01 = vmulq_n_u16( v01, wTR );
          v10 = vmulq_n_u16( v10, wBL );
          v11 = vmulq_n_u16( v11, wBR );

          // add all results together
          v00 = vaddq_u16( v00, vaddq_u16( v01, vaddq_u16( v10, v11 ) ) );

          // descale: shift right by constant
          v00 = vrshrq_n_u16(v00, SHIFT_BITS);

          // compute difference between reference and interpolated patch,
          // use reinterpet-cast to make signed [-255,255]
          int16x8_t res = vsubq_s16(vreinterpretq_s16_u16(v00),
                                    vreinterpretq_s16_u16(vmovl_u8(vld1_u8( it_ref ))));

          // Vector multiply accumulate long: vmla -> Vr[i] := Va[i] + Vb[i] * Vc[i]
          // int32x4_t  vmlal_s16(int32x4_t a, int16x4_t b, int16x4_t c);    // VMLAL.S16 q0,d0,d0
          int16x8_t grad = vld1q_s16(ref_patch_dx + y*patch_size + x);
          vb1 = vmlal_s16(vb1, vget_low_s16(grad), vget_low_s16(res));
          vb1 = vmlal_s16(vb1, vget_high_s16(grad), vget_high_s16(res));

          grad = vld1q_s16(ref_patch_dy + y*patch_size + x);
          vb2 = vmlal_s16(vb2, vget_low_s16(grad), vget_low_s16(res));
          vb2 = vmlal_s16(vb2, vget_high_s16(grad), vget_high_s16(res));
        }
      }

      // finally, sum results of vb1, vb2
      int32x2_t tmp;
      tmp = vpadd_s32(vget_low_s32(vb1), vget_high_s32(vb1));
      Jres[0] = -vget_lane_s32(tmp, 0) - vget_lane_s32(tmp, 1);

      tmp = vpadd_s32(vget_low_s32(vb2), vget_high_s32(vb2));
      Jres[1] = -vget_lane_s32(tmp, 0) - vget_lane_s32(tmp, 1);
#else
      int16_t* it_ref_dx = ref_patch_dx;
      int16_t* it_ref_dy = ref_patch_dy;

      for(int y=0; y<patch_size; ++y)
      {
        uint8_t* it = img_cur.data + (v_r+y)*step + (u_r);
        for(int x=0; x<patch_size; ++x, ++it, ++it_ref, ++it_ref_dx, ++it_ref_dy)
        {
          uint16_t cur = SVO_DESCALE(wTL*it[0] + wTR*it[1] + wBL*it[step] + wBR*it[step+1], SHIFT_BITS);
          //float cur = wTL*it[0] + wTR*it[1] + wBL*it[step] + wBR*it[step+1];
          float res = static_cast<float>(cur) - *it_ref;
          Jres[0] -= res*(*it_ref_dx);
          Jres[1] -= res*(*it_ref_dy);
        }
      }
#endif

      update = Hinv * Jres * 2.0; // * 2 to compensate because above, we did not compute the derivative correctly
      u += update[0];
      v += update[1];

#if SVO_PYRAMIDAL_VERBOSE
      std::cout << "Level: " << level
                << "\t Iter: " << iter << ":"
                << "\t u=" << u << ", v=" << v
                << "\t update = " << update[0] << ", " << update[1] << std::endl;
#endif

      if(update[0]*update[0]+update[1]*update[1] < min_update_squared)
      {
#if SVO_PYRAMIDAL_VERBOSE
        std::cout << "level = " << level << ": converged." << std::endl;
#endif
        converged=true;
        break;
      }
    } // end iterations

    px_cur_level_0 = Keypoint((u+halfpatch_size+px_ref_offset[0])*scale,
        (v+halfpatch_size+px_ref_offset[1])*scale);

    if(!converged && !go_to_next_level)
      return false; // no need to go to next level when we didn't converge at this one.

  } // end pyramid

  return converged;
}

} // namespace feature_alignment
} // namespace svo


#pragma once
#include <armadillo>
#include <iostream>

namespace arm_simu{

	typedef void (*input_handle)(double t,arma::vec*);

	class System{
		arma::mat A;
		arma::mat B;
		arma::mat C;
		arma::mat D;
		arma::rowvec x0;
		arma::rowvec x0_copy;
		input_handle handle;
		int input_size;
		public:
			System(arma::mat A,arma::mat B,arma::mat C,arma::mat D,arma::rowvec x0,input_handle handle,int input_size);
			~System();
			arma::mat& getA();
			arma::mat& getB();
			arma::mat& getC();
			arma::mat& getD();
			arma::rowvec& getX0();
			arma::mat compute_state(arma::vec t_frame);
			arma::mat compute_output(arma::mat state_result);//to be implement
		private:
	};

}


#include <iostream>
#include <sstream>
#include <map>
#include <experimental/filesystem>
#include <opencv2/core.hpp>
#include <opencv2/imgcodecs.hpp>
#include <opencv2/imgproc.hpp>
#include <opencv2/calib3d/calib3d.hpp>
#include <Eigen/Geometry>
#include "mpe/pose_estimation/pose_estimation.h"
#include "mpe/pose_estimation/marker_pose.h"
#include "mpe/utils/marker_detector.h"
#include "mpe/utils/path_loader.h"
#include "mpe/utils/draw_cube.h"

mpe::PoseEstimation::PoseEstimation()
 : fs{"/output/estimation.yml", cv::FileStorage::WRITE} {
  read_config();
}

void mpe::PoseEstimation::help(const std::string images_dir) {
  std::cout << "\nPose Estimation\n" << std::endl;
  std::cout << "This is a marker pose estimation program." << std::endl;
  std::cout << "Usage: pose_estimation [images_dir]" << std::endl
            << "(default: ./images --> The program will process every file in input/images/)\n" << std::endl;
}

void mpe::PoseEstimation::read_config() {
  std::cout << "\nReading parameters from input/config.yml" << std::endl;
  cv::FileStorage fs("/input/config.yml", cv::FileStorage::READ);
  marker_size_ = fs["marker_size_meter"];
  intrinsic_matrix_ = (cv::Mat_<double>(3, 3) <<
    fs["camera_intrinsic"]["fx"],                            0, fs["camera_intrinsic"]["cx"],
                               0, fs["camera_intrinsic"]["fy"], fs["camera_intrinsic"]["cy"],
                               0,                            0,                            1);
  distortion_coeff_ = (cv::Mat_<double>(5, 1) << fs["distortion_coeff"]["k1"],
                                             fs["distortion_coeff"]["k2"],
                                             fs["distortion_coeff"]["p1"],
                                             fs["distortion_coeff"]["p2"],
                                             fs["distortion_coeff"]["k3"]);
  std::cout << "\nMarker size in meter: " << marker_size_ << std::endl;
  std::cout << "\nCamera intrinsic matrix: \n" << intrinsic_matrix_ << std::endl;
  std::cout << "\nCamera distortion coefficient \n" << distortion_coeff_ << std::endl;
}

void mpe::PoseEstimation::draw() {
  for (int i = 0; i < ids_.size(); i++) {
    draw_cube(image_out_, ids_[i], intrinsic_matrix_, distortion_coeff_,
              rotation_[i], translation_[i], marker_size_);
  }
}

void mpe::PoseEstimation::output(const std::string output_path) {
  // Output image
  imwrite(output_path, image_out_);

  // Output pose
  std::map<int, MarkerPose> results;
  for (int i = 0; i < ids_.size(); ++i) {
    Eigen::Vector3d rotation(rotation_[i][0], rotation_[i][1], rotation_[i][2]);
    double theta = rotation.norm();
    Eigen::Quaterniond q(Eigen::AngleAxisd(theta, rotation / theta));
    std::stringstream ss;
    ss << "  Marker " << ids_[i] << '\n';
    ss << "    translation (x, y, z) = (";
    ss << translation_[i][0] << ", ";
    ss << translation_[i][1] << ", ";
    ss << translation_[i][2] << ")\n";
    ss << "    rotation (w, x, y, z) = (";
    ss << q.w() << ", ";
    ss << q.vec()[0] << ", ";
    ss << q.vec()[1] << ", ";
    ss << q.vec()[2] << ")" << std::endl;
    results[ids_[i]] = MarkerPose(translation_[i], q, ss.str());
  }
  // Output result (sorted according to the index of the marker)
  for (auto result : results) {
    std::stringstream marker_name;
    marker_name << "marker_" << result.first;
    std::cout << result.second.text;
    fs << marker_name.str() << "{";
    fs <<   "translation" << "{";
    fs <<     "x" << result.second.x;
    fs <<     "y" << result.second.y;
    fs <<     "z" << result.second.z;
    fs <<   "}";
    fs <<   "rotation" << "{";
    fs <<     "w" << result.second.qw;
    fs <<     "x" << result.second.qx;
    fs <<     "y" << result.second.qy;
    fs <<     "z" << result.second.qz;
    fs <<   "}";
    fs << "}";
  }

  std::cout << std::endl;
}

void mpe::PoseEstimation::estimate_pose(const std::string image_path, const std::string output_path) {
  // Read the image
  image_ = cv::imread(image_path);
  image_.copyTo(image_out_);

  // Detect markers
  MarkerDetector().detect(image_, corners_, ids_);

  if (ids_.size()) {
    // Estimate pose
    cv::aruco::estimatePoseSingleMarkers(corners_, marker_size_, intrinsic_matrix_,
                                         distortion_coeff_, rotation_, translation_);

    // Draw axes for each detected marker
    draw();
  }

  // Output
  output(output_path);

}

void mpe::PoseEstimation::estimate(const std::string images_dir) {
  help(images_dir);

  PathLoader path_loader(images_dir);

  for (auto image_path : path_loader.load()) {
    std::string filename = std::experimental::filesystem::path(image_path).filename();
    std::string output_path = "/output/" + filename;
    std::cout << "Processing " << filename
              << " --> " << output_path.substr(1) << std::endl;

    fs << filename.substr(0, filename.find(".")) << "{";

    try {
      estimate_pose(image_path, output_path);
    } catch (const std::exception& e) {
      std::cerr << e.what() << std::endl;
      continue;
    }

    fs << "}";
  }

  fs.release();

  std::cout << "Done!" << std::endl;
}

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

  // Path to the directory containing input images
  const std::string images_dir = argc > 1 ? argv[1] : "/input/images";
  mpe::PoseEstimation().estimate(images_dir);

  return 0;
}


#ifndef STAN_MATH_PRIM_SCAL_FUN_TRUNC_HPP
#define STAN_MATH_PRIM_SCAL_FUN_TRUNC_HPP

#include <stan/math/prim/scal/fun/boost_policy.hpp>
#include <stan/math/prim/scal/fun/is_nan.hpp>
#include <boost/math/special_functions/trunc.hpp>
#include <limits>

namespace stan {
  namespace math {

    /**
     * Return the nearest integral value that is not larger in
     * magnitude than the specified argument.
     *
     * @param[in] x Argument.
     * @return The truncated argument.
     */
    inline double trunc(double x) {
      if (is_nan(x))
        return std::numeric_limits<double>::quiet_NaN();
      return boost::math::trunc(x, boost_policy_t());
    }

    /**
     * Return the nearest integral value that is not larger in
     * magnitude than the specified argument.
     *
     * @param[in] x Argument.
     * @return The truncated argument.
     */
    inline double trunc(int x) {
      return trunc(static_cast<double>(x));
    }

  }
}
#endif


/******************************************************************************

  This source file is part of the Avogadro project.

  Copyright 2008-2009 Marcus D. Hanwell
  Copyright 2010-2013 Kitware, Inc.

  This source code is released under the New BSD License, (the "License").

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

******************************************************************************/

#include "slaterset.h"

#include <cmath>
#include <iostream>

#include <Eigen/LU>

using std::vector;
using std::cout;
using std::endl;
using Eigen::SelfAdjointEigenSolver;

namespace Avogadro {
namespace Core {

SlaterSet::SlaterSet() : m_initialized(false)
{
}

SlaterSet::~SlaterSet()
{
}

bool SlaterSet::addSlaterIndices(const std::vector<int>& i)
{
  m_slaterIndices = i;
  return true;
}

bool SlaterSet::addSlaterTypes(const std::vector<int>& t)
{
  m_initialized = false;
  m_slaterTypes = t;
  return true;
}

bool SlaterSet::addZetas(const std::vector<double>& z)
{
  m_initialized = false;
  m_zetas = z;
  return true;
}

bool SlaterSet::addPQNs(const std::vector<int>& pqns)
{
  m_initialized = false;
  m_pqns = pqns;
  return true;
}

bool SlaterSet::addOverlapMatrix(const Eigen::MatrixXd& m)
{
  m_initialized = false;
  m_overlap.resize(m.rows(), m.cols());
  m_overlap = m;
  return true;
}

bool SlaterSet::addEigenVectors(const Eigen::MatrixXd& e)
{
  m_eigenVectors.resize(e.rows(), e.cols());
  m_eigenVectors = e;
  return true;
}

bool SlaterSet::addDensityMatrix(const Eigen::MatrixXd& d)
{
  m_density.resize(d.rows(), d.cols());
  m_density = d;
  return true;
}

unsigned int SlaterSet::molecularOrbitalCount(ElectronType)
{
  return static_cast<unsigned int>(m_overlap.cols());
}

void SlaterSet::outputAll()
{
}

void SlaterSet::initCalculation()
{
  if (m_initialized)
    return;

  m_normalized.resize(m_overlap.cols(), m_overlap.rows());

  SelfAdjointEigenSolver<MatrixX> s(m_overlap);
  MatrixX p = s.eigenvectors();
  MatrixX m =
    p * s.eigenvalues().array().inverse().array().sqrt().matrix().asDiagonal() *
    p.inverse();
  m_normalized = m * m_eigenVectors;

  if (!(m_overlap * m * m).eval().isIdentity())
    cout << "Identity test FAILED - do you need a newer version of Eigen?\n";

  m_factors.resize(m_zetas.size());
  m_PQNs = m_pqns;
  // Calculate the normalizations of the orbitals.
  for (size_t i = 0; i < m_zetas.size(); ++i) {
    switch (m_slaterTypes[i]) {
      case S:
        m_factors[i] = pow(2.0 * m_zetas[i], m_pqns[i] + 0.5) *
                       sqrt(1.0 / (4.0 * M_PI) / factorial(2 * m_pqns[i]));
        m_PQNs[i] -= 1;
        break;
      case PX:
      case PY:
      case PZ:
        m_factors[i] = pow(2.0 * m_zetas[i], m_pqns[i] + 0.5) *
                       sqrt(3.0 / (4.0 * M_PI) / factorial(2 * m_pqns[i]));
        m_PQNs[i] -= 2;
        break;
      case X2:
        m_factors[i] = 0.5 * pow(2.0 * m_zetas[i], m_pqns[i] + 0.5) *
                       sqrt(15.0 / (4.0 * M_PI) / factorial(2 * m_pqns[i]));
        m_PQNs[i] -= 3;
        break;
      case XZ:
        m_factors[i] = pow(2.0 * m_zetas[i], m_pqns[i] + 0.5) *
                       sqrt(15.0 / (4.0 * M_PI) / factorial(2 * m_pqns[i]));
        m_PQNs[i] -= 3;
        break;
      case Z2:
        m_factors[i] = (0.5 / sqrt(3.0)) *
                       pow(2.0 * m_zetas[i], m_pqns[i] + 0.5) *
                       sqrt(15.0 / (4.0 * M_PI) / factorial(2 * m_pqns[i]));
        m_PQNs[i] -= 3;
        break;
      case YZ:
      case XY:
        m_factors[i] = pow(2.0 * m_zetas[i], m_pqns[i] + 0.5) *
                       sqrt(15.0 / (4.0 * M_PI) / factorial(2 * m_pqns[i]));
        m_PQNs[i] -= 3;
        break;
      default:
        cout << "Orbital " << i << " not handled, type " << m_slaterTypes[i]
             << endl;
    }
  }
  // Convert the exponents into Angstroms
  for (size_t i = 0; i < m_zetas.size(); ++i)
    m_zetas[i] = m_zetas[i] / BOHR_TO_ANGSTROM_D;

  m_initialized = true;
}

inline unsigned int SlaterSet::factorial(unsigned int n)
{
  if (n <= 1)
    return n;
  return (n * factorial(n - 1));
}

} // End namespace Core
} // End namespace Avogadro


/*
 * distributions_boost.cpp
 *
 *  Created on: 7 Sep 2018
 *      Author: admin
 */
#include <Eigen/Eigen>
#include <math.h>
#include "distributions_boost.hpp"
#include <boost/random/gamma_distribution.hpp>
#include "boost/random.hpp"
#include "boost/generator_iterator.hpp"
#include <boost/random/normal_distribution.hpp>
#include <boost/random/beta_distribution.hpp>
#include <boost/random/uniform_01.hpp>
#include <iostream>
#include <fstream>

using namespace std;

/*
Distributions_boost::Distributions_boost(unsigned int pseed):seed(pseed){
    std::cout << "Initiating dist.rng with pseed = " << pseed << std::endl;
    rng = boost::mt19937(seed);
}
*/

Distributions_boost::Distributions_boost() {
    rng = boost::mt19937(0);
}

Distributions_boost::~Distributions_boost() {}

void Distributions_boost::reset_rng(unsigned int seed) {
    rng = boost::mt19937(seed);
}

void Distributions_boost::write_rng_state_to_file(const std::string rngfp) {
    //cout << "Dumping rng state to file " << rngfp << endl;
    ofstream file;
	file.open(rngfp.c_str(), ios::out | ios::trunc | ios::binary);
	file << rng;
	file.close();
}

void Distributions_boost::read_rng_state_from_file(const std::string rngfp) {
    //cout << "Reading rng state from file " << rngfp << endl;
    ifstream file(rngfp.c_str(), ios::in | ios::binary);
    if (file.is_open()) {
        file >> rng;
    } else {
        cout << "*FATAL*: Unable to read from file " << rngfp << endl;
        exit(1);
    }
}

double Distributions_boost::rgamma(double shape, double scale){
    boost::random::gamma_distribution<double> myGamma(shape, scale);
    boost::random::variate_generator<boost::mt19937&, boost::random::gamma_distribution<> > rand_gamma(rng, myGamma);
    return rand_gamma();
}

double Distributions_boost::unif_rng(){
    boost::random::uniform_real_distribution<double> myU(0,1);
    boost::random::variate_generator<boost::mt19937&, boost::random::uniform_real_distribution<> > real_variate_generator(rng, myU);
    return real_variate_generator();
}

Eigen::VectorXd Distributions_boost::dirichlet_rng(Eigen::VectorXi alpha) {
    int len;
    len=alpha.size();
    Eigen::VectorXd result(len);
    for(int i=0;i<len;i++)
        result[i]=rgamma((double)alpha[i], 1.0);
    result/=result.sum();
    return result;
}

Eigen::VectorXd Distributions_boost::dirichlet_rng(Eigen::VectorXd alpha) {
    int len;
    len=alpha.size();
    Eigen::VectorXd result(len);
    for(int i=0; i<len; i++)
        result[i]=rgamma(alpha[i],(double)1.0);
    result/=result.sum();
    return result;
}

double Distributions_boost::inv_gamma_rng(double shape,double scale){
    return ((double)1.0 / rgamma(shape, 1.0/scale));
}

double Distributions_boost::gamma_rng(double shape,double scale){
    return rgamma(shape, scale);
}

double Distributions_boost::inv_gamma_rate_rng(double shape,double rate){
    return (double)1.0 / gamma_rate_rng(shape, rate);
}

double Distributions_boost::gamma_rate_rng(double shape,double rate){
    return rgamma(shape,(double)1.0/rate);
}

double Distributions_boost::inv_scaled_chisq_rng(double dof,double scale){
    return inv_gamma_rng((double)0.5*dof, (double)0.5*dof*scale);
}

double Distributions_boost::norm_rng(double mean,double sigma2){
    boost::normal_distribution<double> nd(mean, std::sqrt((double)sigma2));
    boost::variate_generator<boost::mt19937&,boost::normal_distribution<> > var_nor(rng, nd);
    return var_nor();
}

inline double runif(unsigned int seed){
    boost::mt19937 rng(seed);
    static boost::uniform_01<boost::mt19937> zeroone(rng);
    return zeroone();
}

double Distributions_boost::component_probs(double b2,Eigen::VectorXd pi){
    double sum;
    double p;
    p=runif(seed);
    sum= pi[0]*exp((-0.5*b2)/(5e-2))/sqrt(5e-2)+pi[1]*exp((-0.5*b2));
    if(p<=(pi[0]*exp((-0.5*b2)/(5e-2))/sqrt(5e-2))/sum)
        return 5e-2;
    else
        return 1;
}

double Distributions_boost::beta_rng(double a,double b){
    boost::random::beta_distribution<double> mybeta(a, b);
    boost::random::variate_generator<boost::mt19937&, boost::random::beta_distribution<> > rand_beta(rng, mybeta);
    return rand_beta();
}


// Copyright (c) 2018 Graphcore Ltd. All rights reserved.
#include "Constraint.hpp"
#include "Scheduler.hpp"
#include <memory>
#define BOOST_TEST_MODULE Sum
#include <boost/test/unit_test.hpp>

using namespace popsolver;

BOOST_AUTO_TEST_CASE(PropagateNoChange) {
  Variable a(0), b(1), c(2);
  auto sum = std::unique_ptr<Sum>(new Sum(c, {a, b}));
  Domains domains;
  domains.push_back({DataType{7}, DataType{8}});  // a
  domains.push_back({DataType{2}, DataType{5}});  // b
  domains.push_back({DataType{9}, DataType{13}}); // c
  Scheduler scheduler(domains, {sum.get()});
  bool success = sum->propagate(scheduler);
  BOOST_CHECK(success);
  BOOST_CHECK_EQUAL(scheduler.getDomains()[a].min(), DataType{7});
  BOOST_CHECK_EQUAL(scheduler.getDomains()[a].max(), DataType{8});
  BOOST_CHECK_EQUAL(scheduler.getDomains()[b].min(), DataType{2});
  BOOST_CHECK_EQUAL(scheduler.getDomains()[b].max(), DataType{5});
  BOOST_CHECK_EQUAL(scheduler.getDomains()[c].min(), DataType{9});
  BOOST_CHECK_EQUAL(scheduler.getDomains()[c].max(), DataType{13});
}

BOOST_AUTO_TEST_CASE(PropagateResult) {
  Variable a(0), b(1), c(2);
  auto sum = std::unique_ptr<Sum>(new Sum(c, {a, b}));
  Domains domains;
  domains.push_back({DataType{7}, DataType{8}});  // a
  domains.push_back({DataType{2}, DataType{5}});  // b
  domains.push_back({DataType{1}, DataType{20}}); // c
  Scheduler scheduler(domains, {sum.get()});
  bool success = sum->propagate(scheduler);
  BOOST_CHECK(success);
  BOOST_CHECK_EQUAL(scheduler.getDomains()[a].min(), DataType{7});
  BOOST_CHECK_EQUAL(scheduler.getDomains()[a].max(), DataType{8});
  BOOST_CHECK_EQUAL(scheduler.getDomains()[b].min(), DataType{2});
  BOOST_CHECK_EQUAL(scheduler.getDomains()[b].max(), DataType{5});
  BOOST_CHECK_EQUAL(scheduler.getDomains()[c].min(), DataType{9});
  BOOST_CHECK_EQUAL(scheduler.getDomains()[c].max(), DataType{13});
}

BOOST_AUTO_TEST_CASE(PropagateOperands) {
  Variable a(0), b(1), c(2);
  auto sum = std::unique_ptr<Sum>(new Sum(c, {a, b}));
  Domains domains;
  domains.push_back({DataType{7}, DataType{1000}}); // a
  domains.push_back({DataType{0}, DataType{9}});    // b
  domains.push_back({DataType{9}, DataType{13}});   // c
  Scheduler scheduler(domains, {sum.get()});
  bool success = sum->propagate(scheduler);
  BOOST_CHECK(success);
  BOOST_CHECK_EQUAL(scheduler.getDomains()[a].min(), DataType{7});
  BOOST_CHECK_EQUAL(scheduler.getDomains()[a].max(), DataType{13});
  BOOST_CHECK_EQUAL(scheduler.getDomains()[b].min(), DataType{0});
  BOOST_CHECK_EQUAL(scheduler.getDomains()[b].max(), DataType{6});
  BOOST_CHECK_EQUAL(scheduler.getDomains()[c].min(), DataType{9});
  BOOST_CHECK_EQUAL(scheduler.getDomains()[c].max(), DataType{13});
}

BOOST_AUTO_TEST_CASE(PropagateBoth) {
  Variable a(0), b(1), c(2);
  auto sum = std::unique_ptr<Sum>(new Sum(c, {a, b}));
  Domains domains;
  domains.push_back({DataType{7}, DataType{8}});  // a
  domains.push_back({DataType{0}, DataType{5}});  // b
  domains.push_back({DataType{9}, DataType{20}}); // c
  Scheduler scheduler(domains, {sum.get()});
  bool success = sum->propagate(scheduler);
  BOOST_CHECK(success);
  BOOST_CHECK_EQUAL(scheduler.getDomains()[a].min(), DataType{7});
  BOOST_CHECK_EQUAL(scheduler.getDomains()[a].max(), DataType{8});
  BOOST_CHECK_EQUAL(scheduler.getDomains()[b].min(), DataType{1});
  BOOST_CHECK_EQUAL(scheduler.getDomains()[b].max(), DataType{5});
  BOOST_CHECK_EQUAL(scheduler.getDomains()[c].min(), DataType{9});
  BOOST_CHECK_EQUAL(scheduler.getDomains()[c].max(), DataType{13});
}

BOOST_AUTO_TEST_CASE(AvoidOverflow) {
  Variable a(0), b(1), c(2);
  auto sum = std::unique_ptr<Sum>(new Sum(c, {a, b}));
  Domains domains;
  domains.push_back({DataType{1}, DataType::max() - DataType{1}}); // a
  domains.push_back({DataType{2}, DataType::max() - DataType{1}}); // b
  domains.push_back({DataType{20}, DataType{21}});                 // c
  Scheduler scheduler(domains, {sum.get()});
  bool success = sum->propagate(scheduler);
  BOOST_CHECK(success);
  BOOST_CHECK_EQUAL(scheduler.getDomains()[a].min(), DataType{1});
  BOOST_CHECK_EQUAL(scheduler.getDomains()[a].max(), DataType{19});
  BOOST_CHECK_EQUAL(scheduler.getDomains()[b].min(), DataType{2});
  BOOST_CHECK_EQUAL(scheduler.getDomains()[b].max(), DataType{20});
  BOOST_CHECK_EQUAL(scheduler.getDomains()[c].min(), DataType{20});
  BOOST_CHECK_EQUAL(scheduler.getDomains()[c].max(), DataType{21});
}


////////////////////////////////////////////////////////////////////////////////
// TensionFieldTheory.hh
////////////////////////////////////////////////////////////////////////////////
/*! @file
//  Implements a field theory relaxed energy density for a fully generic,
//  potentially anisotropic, 2D "C-based" energy density. Here "C-based" means
//  the energy density is expressed in terms of the Cauchy-Green deformation
//  tensor.
//  Our implementation is based on the applied math paper
//  [Pipkin1994:"Relaxed energy densities for large deformations of membranes"]
//  whose optimality criterion for the wrinkling strain we use to obtain a
//  slightly less expensive optimization formulation.
//
//  We implement the second derivatives needed for a Newton-based equilibrium
//  solver; these are nontrivial since they must account for the dependence of
//  the wrinkling strain on C.
*/
//  Author:  Julian Panetta (jpanetta), julian.panetta@gmail.com
//  Created:  06/28/2020 19:01:21
////////////////////////////////////////////////////////////////////////////////
#ifndef TENSIONFIELDTHEORY_HH
#define TENSIONFIELDTHEORY_HH
#include <Eigen/Dense>
#include <MeshFEM/Types.hh>
#include <stdexcept>
#include <iostream>
#include <MeshFEM/newton_optimizer/dense_newton.hh>

template<class _Psi>
struct IsotropicWrinkleStrainProblem {
    using Psi      = _Psi;
    using Real     = typename Psi::Real;
    using VarType  = Eigen::Matrix<Real, 1, 1>;
    using HessType = Eigen::Matrix<Real, 1, 1>;
    using M2d      = Mat2_T<Real>;

    IsotropicWrinkleStrainProblem(Psi &psi, const M2d &C, const Vec2_T<Real> &n)
        : m_psi(psi), m_C(C), m_nn(n * n.transpose()) { }

    void setC(const M2d &C) { m_C = C; }
    const M2d &getC() const { return m_C; }

    size_t numVars() const { return 1; }
    void setVars(const VarType &vars) {
        m_a = vars;
        m_psi.setC(m_C + m_a[0] * m_nn);
    }
    const VarType &getVars() const { return m_a; }

    Real energy()      const { return m_psi.energy(); }
    VarType gradient() const { return  VarType(0.5 * doubleContract(m_nn, m_psi.PK2Stress())          ); }
    HessType hessian() const { return HessType(0.5 * doubleContract(m_psi.delta_PK2Stress(m_nn), m_nn)); }

    void solve() { dense_newton(*this, /* maxIter = */ 100, /*gradTol = */1e-14, /* verbose = */ false); }

    EIGEN_MAKE_ALIGNED_OPERATOR_NEW
private:
    VarType m_a = VarType::Zero();
    Psi &m_psi;
    M2d m_C, m_nn;
};

template<class _Psi>
struct AnisotropicWrinkleStrainProblem {
    using Psi      = _Psi;
    using Real     = typename Psi::Real;
    using VarType  = Vec2_T<Real>;
    using HessType = Mat2_T<Real>;
    using M2d      = Mat2_T<Real>;

    AnisotropicWrinkleStrainProblem(Psi &psi, const M2d &C, const VarType &n)
        : m_psi(psi), m_C(C), m_ntilde(n) { }

    void setC(const M2d &C) { m_C = C; }
    const M2d &getC() const { return m_C; }

    size_t numVars() const { return 2; }
    void setVars(const VarType &vars) {
        m_ntilde = vars;
        m_psi.setC(m_C + m_ntilde * m_ntilde.transpose());
    }
    const VarType &getVars() const { return m_ntilde; }
    Real energy() const { return m_psi.energy(); }

    // S : (0.5 * (n delta_n^T + delta_n n^T))
    //  = S : n delta_n^T = (S n) . delta_n
    VarType gradient() const { return m_psi.PK2Stress() * m_ntilde; }

    //   psi(n * n^T)
    //  dpsi = n^T psi'(n * n^T) . dn
    // d2psi = dn_a^T psi'(n * n^T) . dn_b + n^T (psi'' : (n * dn_a^T)) . dn_b
    //       = psi' : (dn_a dn_b^T) + ...
    HessType hessian() const {
        HessType h = m_psi.PK2Stress(); // psi' : (dn_a dn_b^T)
        M2d dnn(M2d::Zero());
        dnn.col(0) = m_ntilde;
        h.row(0) += m_ntilde.transpose() * m_psi.delta_PK2Stress(symmetrized_x2(dnn));
        dnn.col(0).setZero();
        dnn.col(1) = m_ntilde;
        h.row(1) += m_ntilde.transpose() * m_psi.delta_PK2Stress(symmetrized_x2(dnn));
        if (h.array().isNaN().any()) throw std::runtime_error("NaN Hessian");
        if (std::abs(h(0, 1) - h(1, 0)) > 1e-10 * std::abs(h(1, 0)) + 1e-10)
            throw std::runtime_error("Asymmetric Hessian");
        return h;
    }

    void solve() { dense_newton(*this, /* maxIter = */ 100, /*gradTol = */1e-14, /* verbose = */ false); }

    EIGEN_MAKE_ALIGNED_OPERATOR_NEW
private:
    Psi &m_psi;
    M2d m_C;
    VarType m_ntilde = VarType::Zero();
};

// Define a relaxed 2D C-based energy based on a given 2D C-based energy `Psi_C`
template<class Psi_C>
struct RelaxedEnergyDensity {
    static_assert(Psi_C::EDType == EDensityType::CBased,
                  "Tension field theory only works on C-based energy densities");
    static constexpr size_t N = Psi_C::N;
    static constexpr size_t Dimension = N;
    static constexpr EDensityType EDType = EDensityType::CBased;

    using Matrix = typename Psi_C::Matrix;
    using Real   = typename Psi_C::Real;
    using ES     = Eigen::SelfAdjointEigenSolver<Matrix>;
    using V2d    = Vec2_T<Real>;
    static_assert(N == 2, "Tension field theory relaxation only defined for 2D energies");

    static std::string name() {
        return std::string("Relaxed") + Psi_C::name();
    }

    // Forward all constructor arguments to m_psi
    template<class... Args>
    RelaxedEnergyDensity(Args&&... args)
         : m_psi(std::forward<Args>(args)...),
           m_anisoProb(m_psi, Matrix::Identity(), V2d::Zero()) { }

    // We need a custom copy constructor since m_anisoProb contains a
    // reference to this->m_psi
    RelaxedEnergyDensity(const RelaxedEnergyDensity &b)
        : m_psi(b.m_psi),
          m_anisoProb(m_psi, Matrix::Identity(), V2d::Zero()) {
        setC(b.m_C);
    }

    // Note: UninitializedDeformationTag argument must be a rvalue reference so it exactly
    // matches the type passed by the constructor call
    // RelaxedEnergyDensity(b, UninitializedDeformationTag()); otherwise the
    // perfect forwarding constructor above will be preferred for this call,
    // incorrectly forwarding b to Psi's constructor.
    RelaxedEnergyDensity(const RelaxedEnergyDensity &b, UninitializedDeformationTag &&)
        : m_psi(b.m_psi, UninitializedDeformationTag()),
          m_anisoProb(m_psi, Matrix::Identity(), V2d::Zero()) {
        setC(b.m_C);
    }

    void setC(const Matrix &C) {
        m_C = C;
        if (!relaxationEnabled()) {
            m_psi.setC(C);
            return;
        }

        // Note: Eigen guarantees eigenvalues are sorted in ascending order.
        ES C_eigs(C);
        // std::cout << "C eigenvalues: " << C_eigs.eigenvalues().transpose() << std::endl;
        // Detect full compression
        if (C_eigs.eigenvalues()[1] < 1) {
            m_tensionState = 0;
            m_wrinkleStrain = -C;
            return;
        }

        m_psi.setC(C);
        ES S_eigs(m_psi.PK2Stress());
        // Detect full tension
        // std::cout << "S eigenvalues: " << S_eigs.eigenvalues().transpose() << std::endl;
        if (S_eigs.eigenvalues()[0] >= -1e-12) {
            m_tensionState = 2;
            m_wrinkleStrain.setZero();
            return;
        }

        // Handle partial tension
        m_tensionState = 1;
        // In the isotropic case, principal stress and strain directions
        // coincide, and the wrinkling strain must be in the form
        // a n n^T, where n is the eigenvector corresponding to the smallest
        // stress eigenvalue and a > 0 is an unknown.
        // This simplifies the determination of wrinkling strain to a convex
        // 1D optimization problem that we solve with Newton's method.
        V2d n = S_eigs.eigenvectors().col(0);
        using IWSP = IsotropicWrinkleStrainProblem<Psi_C>;
        IWSP isoProb(m_psi, C, n);
        // std::cout << "Solving isotropic wrinkle strain problem" << std::endl;
        isoProb.setVars(typename IWSP::VarType{0.0});
        isoProb.solve();
        Real a = isoProb.getVars()[0];
        if (a < 0) throw std::runtime_error("Invalid wrinkle strain");

        // We use this isotropic assumption to obtain initial guess for the
        // anisotropic case, where the wrinkling strain is in the form
        //      n_tilde n_tilde^T
        // with a 2D vector n_tilde as the unknown.
        m_anisoProb.setC(C);
        // std::cout << "Solving anisotropic wrinkle strain problem" << std::endl;
        m_anisoProb.setVars(std::sqrt(a) * n);
        m_anisoProb.solve();
        auto ntilde = m_anisoProb.getVars();
        m_wrinkleStrain = -ntilde * ntilde.transpose();

        // {
        //     ES S_eigs_new(m_psi.PK2Stress());
        //     std::cout << "new S eigenvalues: " << S_eigs_new.eigenvalues().transpose() << std::endl;
        // }
    }

    Real energy() const {
        if (relaxationEnabled() && fullCompression()) return 0.0;
        return m_psi.energy();
    }

    Matrix PK2Stress() const {
        if (relaxationEnabled() && fullCompression()) return Matrix::Zero();
        // By envelope theorem, the wrinkling strain's perturbation can be
        // neglected, and the stress is simply the stress of the underlying
        // material model evaluated on the "elastic strain".
        return m_psi.PK2Stress();
    }

    template<class Mat_>
    Matrix delta_PK2Stress(const Mat_ &dC) const {
        if (!relaxationEnabled() || fullTension()) return m_psi.delta_PK2Stress(dC);
        if (fullCompression()) return Matrix::Zero();

        // n solves m_anisoProb:
        //     (psi') n = 0
        // delta_n solves:
        //     (psi'' : dC + n delta n) n + (psi') delta_n = 0
        //     (psi'' : n delta_n) n + (psi') delta_n = -(psi'' : dC) n
        //     H delta_n = -(psi'' : dC) n
        auto ntilde = m_anisoProb.getVars();
        V2d delta_n = -m_anisoProb.hessian().inverse() * (m_psi.delta_PK2Stress(dC.matrix()) * ntilde);

        // In the partial tension case, we need to account for the wrinkling
        // strain perturbation.
        return m_psi.delta_PK2Stress(dC.matrix() + ntilde * delta_n.transpose()
                                                 + delta_n * ntilde.transpose());
    }

    template<class Mat_, class Mat2_>
    Matrix delta2_PK2Stress(const Mat_ &/* dF_a */, const Mat2_ &/* dF_b */) const {
        throw std::runtime_error("Unimplemented");
    }

    V2d principalBiotStrains() const {
        ES es(m_C);
        return es.eigenvalues().array().sqrt() - 1.0;
    }

    const Matrix &wrinkleStrain() const { return m_wrinkleStrain; }

    bool fullCompression() const { return m_tensionState == 0; }
    bool partialTension()  const { return m_tensionState == 1; }
    bool fullTension()     const { return m_tensionState == 2; }
    int tensionState()     const { return m_tensionState; }

    // Copying is super dangerous since we m_anisoProb uses a reference to m_psi...
    RelaxedEnergyDensity &operator=(const RelaxedEnergyDensity &) = delete;

    // Direct access to the energy density for debugging or
    // to change the material properties
    Psi_C &psi() { return m_psi; }
    const Psi_C &psi() const { return m_psi; }

    // Turn on or off the tension field theory approximation
    void setRelaxationEnabled(bool enable) {
        m_relaxationEnabled = enable;
        setC(m_C);
    }

    bool relaxationEnabled() const { return m_relaxationEnabled; }

    void copyMaterialProperties(const RelaxedEnergyDensity &other) { psi().copyMaterialProperties(other.psi()); }

    EIGEN_MAKE_ALIGNED_OPERATOR_NEW
private:
    // Tension state:
    //      0: compression in all directions
    //      1: partial tension
    //      2: tension in all directions
    int m_tensionState = 2;
    Psi_C m_psi;
    Matrix m_wrinkleStrain = Matrix::Zero(),
           m_C = Matrix::Identity(); // full Cauchy-Green deformation tensor.
    AnisotropicWrinkleStrainProblem<Psi_C> m_anisoProb;
    bool m_relaxationEnabled = true;
};

#endif /* end of include guard: TENSIONFIELDTHEORY_HH */