// Copyright © 2016-2021 Thomas Nagler and Thibault Vatter
//
// This file is part of the vinecopulib library and licensed under the terms of
// the MIT license. For a copy, see the LICENSE file in the root directory of
// vinecopulib or https://vinecopulib.github.io/vinecopulib/.

#include <vinecopulib/misc/tools_stats.hpp>
#include <vinecopulib/misc/tools_stl.hpp>

#include <boost/graph/prim_minimum_spanning_tree.hpp>
#include <cmath>
#include <iostream>
#include <wdm/eigen.hpp>

namespace vinecopulib {

namespace tools_select {

using namespace tools_stl;

//! @brief Calculates criterion for tree selection.
//! @param data Observations.
//! @param tree_criterion The criterion.
//! @param weights Vector of weights for each observation (can be empty).
inline double
calculate_criterion(const Eigen::MatrixXd& data,
                    std::string tree_criterion,
                    Eigen::VectorXd weights)
{
  double w = 0.0;
  Eigen::MatrixXd data_no_nan = data;
  tools_eigen::remove_nans(data_no_nan, weights);
  double freq =
    static_cast<double>(data_no_nan.rows()) / static_cast<double>(data.rows());
  if (data_no_nan.rows() > 10) {
    if (tree_criterion == "mcor") {
      w = tools_stats::pairwise_mcor(data_no_nan, weights);
    } else if (tree_criterion == "joe") {
      // mutual information for Gaussian copula
      w = wdm::wdm(tools_stats::qnorm(data_no_nan), "pearson", weights)(0, 1);
      w = -0.5 * std::log(1 - w * w);
    } else {
      w = wdm::wdm(data_no_nan, tree_criterion, weights)(0, 1);
    }

    if (std::isnan(w)) {
      w = 0.0;
    }
  }
  return std::fabs(w) * std::sqrt(freq);
}

//! @brief Evaluates maximal criterion for tree selection.
//! @param data Observations.
//! @param tree_criterion The criterion.
//! @param weights Vector of weights for each observation (can be empty).
inline Eigen::MatrixXd
calculate_criterion_matrix(const Eigen::MatrixXd& data,
                           const std::string& tree_criterion,
                           const Eigen::VectorXd& weights)
{
  size_t n = data.rows();
  size_t d = data.cols();
  Eigen::MatrixXd mat(d, d);
  mat.diagonal() = Eigen::VectorXd::Constant(d, 1.0);
  Eigen::MatrixXd pair_data(n, 2);
  for (size_t i = 1; i < d; ++i) {
    for (size_t j = 0; j < i; ++j) {
      pair_data.col(0) = data.col(i);
      pair_data.col(1) = data.col(j);
      Eigen::VectorXd pair_w = weights;
      mat(i, j) = calculate_criterion(pair_data, tree_criterion, pair_w);
      mat(j, i) = mat(i, j);
    }
  }
  return mat;
}

//! computes
inline std::vector<size_t>
get_disc_cols(std::vector<std::string> var_types)
{
  size_t d = var_types.size();
  std::vector<size_t> disc_cols(d);
  size_t disc_count = 0;
  for (size_t i = 0; i < d; ++i) {
    if (var_types[i] == "d") {
      disc_cols[i] = disc_count++;
    } else {
      disc_cols[i] = 0;
    }
  }
  return disc_cols;
}

inline VinecopSelector::VinecopSelector(const Eigen::MatrixXd& data,
                                        const FitControlsVinecop& controls,
                                        std::vector<std::string> var_types)
  : n_(data.rows())
  , d_(var_types.size())
  , var_types_(var_types)
  , controls_(controls)
  , pool_(controls_.get_num_threads())
  , trees_(std::vector<VineTree>(1))
  , threshold_(controls.get_threshold())
  , psi0_(controls.get_psi0())
{
  vine_struct_ = RVineStructure(tools_stl::seq_int(1, d_), 1, false);
}

inline VinecopSelector::VinecopSelector(const Eigen::MatrixXd& data,
                                        const RVineStructure& vine_struct,
                                        const FitControlsVinecop& controls,
                                        std::vector<std::string> var_types)
  : VinecopSelector(data, controls, var_types)
{
  vine_struct_ = vine_struct;
  structure_known_ = false;
}

inline std::vector<std::vector<Bicop>>
VinecopSelector::get_pair_copulas() const
{
  return pair_copulas_;
}

inline RVineStructure
VinecopSelector::get_rvine_structure() const
{
  return vine_struct_;
}

//! @brief Instantiates the object for storing pair copulas.
//! @param d Dimension of the vine copula.
//! @param trunc_lvl A truncation level (optional).
//! @return A nested vector such that `pc_store[t][e]` contains a Bicop.
//!     object for the pair copula corresponding to tree `t` and edge `e`.
inline std::vector<std::vector<Bicop>>
VinecopSelector::make_pair_copula_store(size_t d, size_t trunc_lvl)
{
  if (d < 1) {
    throw std::runtime_error("dimension must be be > 0.");
  }

  size_t n_trees = std::min(d - 1, trunc_lvl);
  std::vector<std::vector<Bicop>> pc_store(n_trees);
  for (size_t t = 0; t < n_trees; ++t) {
    pc_store[t].resize(d - 1 - t);
  }

  return pc_store;
}

inline void
VinecopSelector::select_all_trees(const Eigen::MatrixXd& data)
{
  loglik_ = 0.0;
  initialize_new_fit(data);
  for (size_t t = 0; t < d_ - 1; ++t) {
    select_tree(t); // select pair copulas (+ structure) of tree t
    loglik_ += get_loglik_of_tree(t);

    if (controls_.get_show_trace()) {
      std::stringstream tree_heading;
      std::cout << "** Tree: " << t << std::endl;
      print_pair_copulas_of_tree(t);
    }

    if (controls_.get_trunc_lvl() == t + 1) {
      // don't need to fit the remaining trees
      break;
    }
  }
  finalize(controls_.get_trunc_lvl());
}

inline void
VinecopSelector::sparse_select_all_trees(const Eigen::MatrixXd& data)
{
  // family set must be reset after each iteration of the threshold search
  auto family_set = controls_.get_family_set();
  double d = static_cast<double>(d_);

  std::vector<double> thresholded_crits;
  double mbicv_opt = std::numeric_limits<double>::max();
  bool needs_break = false;
  while (!needs_break) {
    // restore family set in case previous threshold iteration also
    // truncated the model
    controls_.set_family_set(family_set);
    controls_.set_trunc_lvl(std::numeric_limits<size_t>::max());
    initialize_new_fit(data);

    // decrease the threshold
    // (in the first iteration thresholded_crits is empty and the threshold is
    // set to 1.0, which fits an independence model)
    if (controls_.get_select_threshold()) {
      controls_.set_threshold(get_next_threshold(thresholded_crits));
      if (controls_.get_show_trace()) {
        std::cout << "***** threshold: " << controls_.get_threshold()
                  << std::endl;
      }
    }

    // helper variables for checking whether an optimum was found
    double mbicv = 0.0;
    double mbicv_trunc = 0.0;
    double loglik = 0.0;
    bool select_trunc_lvl = controls_.get_select_trunc_lvl();
    bool select_threshold = controls_.get_select_threshold();
    double num_changed = 0.0;
    double num_total = d * (d - 1.0) / 2.0;

    for (size_t t = 0; t < d_ - 1; ++t) {
      if (controls_.get_trunc_lvl() < t) {
        break; // don't need to fit the remaining trees
      }

      // select pair copulas (and possibly tree structure)
      select_tree(t);
      num_changed += d - 1 - static_cast<double>(t);

      // update fit statistic
      double loglik_tree = get_loglik_of_tree(t);
      loglik += loglik_tree;
      double mbicv_tree = get_mbicv_of_tree(t, loglik_tree);
      mbicv_trunc += mbicv_tree;

      // print trace for this tree level
      if (controls_.get_show_trace()) {
        std::cout << "** Tree: " << t;
        if (select_trunc_lvl) {
          std::cout << ", mbicv: " << mbicv_tree << ", loglik: " << loglik_tree;
        }
        std::cout << std::endl;
        print_pair_copulas_of_tree(t);
      }

      // mbicv comparison for truncation level (check only after 10% of
      // copulas change to avoid getting stuck in local minimum)
      if (num_changed / num_total > 0.1) {
        num_changed = 0.0;
        if (select_trunc_lvl & (mbicv_trunc >= mbicv) & (t > 0)) {
          // mbicv did not improve
          // check if it can be improved by removing trees
          loglik -= loglik_tree;
          mbicv_trunc -= mbicv_tree;
          while (t > 1) {
            loglik_tree = get_loglik_of_tree(t - 1);
            mbicv_tree = get_mbicv_of_tree(t - 1, loglik_tree);
            if (mbicv_tree <= 0)
              break;
            loglik -= loglik_tree;
            mbicv_trunc -= mbicv_tree;
            t--;
          }
          set_current_fit_as_opt(loglik);
          controls_.set_trunc_lvl(t);
          if (!select_threshold) {
            // fixed threshold, no need to continue
            needs_break = true;
          }
        } else {
          mbicv = mbicv_trunc;
        }
      }
    }

    if (controls_.get_show_trace()) {
      std::cout << "--> mbicv = " << mbicv << ", loglik = " << loglik
                << std::endl
                << std::endl;
    }

    // check whether mbicv-optimal model has been found
    if (mbicv == 0.0) {
      //// CASE: 0-truncated model is best for this threshold
      set_current_fit_as_opt(loglik);
      if (!select_threshold) {
        // threshold is fixed and trunc_lvl has been found -> stop
        needs_break = true;
      }
    } else if (mbicv >= mbicv_opt) {
      //// CASE: old model is optimal
      needs_break = true;
    } else {
      //// CASE: optimum hasn't been found
      set_current_fit_as_opt(loglik);
      mbicv_opt = mbicv;
      // while loop is only for threshold selection
      needs_break = needs_break || !select_threshold;
      // threshold is too close to 0
      needs_break = needs_break || (controls_.get_threshold() < 0.01);
      // prepare for possible next iteration
      thresholded_crits = get_thresholded_crits();
    }
  }

  // set final model
  trees_ = trees_opt_;
  finalize(controls_.get_trunc_lvl());
}

inline void
VinecopSelector::set_tree_to_indep(size_t t)
{
  // trees_[0] is base tree, see make_base_tree()
  for (auto e : boost::edges(trees_[t + 1])) {
    trees_[t + 1][e].pair_copula = Bicop();
  }
}

// extracts the current loglik value
inline double
VinecopSelector::get_loglik() const
{
  return loglik_;
}

// extracts the current threshold value
inline double
VinecopSelector::get_threshold() const
{
  return threshold_;
}

// extracts the number of observations
inline size_t
VinecopSelector::get_nobs() const
{
  return n_;
}

//! chooses threshold for next iteration such that at a proportion of at
//! least 2.5% of the previously thresholded pairs become non-thresholded.
inline double
VinecopSelector::get_next_threshold(std::vector<double>& thresholded_crits)
{
  if (thresholded_crits.size() == 0) {
    return 1.0;
  }
  // sort in descending order
  std::sort(thresholded_crits.begin(), thresholded_crits.end());
  std::reverse(thresholded_crits.begin(), thresholded_crits.end());
  // pick threshold that changes at least alpha*100 % of the pair-copulas
  double alpha = 0.05;
  size_t m = thresholded_crits.size();
  double new_index = std::ceil(static_cast<double>(m) * alpha) - 1;
  return thresholded_crits[static_cast<size_t>(new_index)];
}

//! @brief Adds edges allowed by either the proximity condition or the vine
//  structure.
//!
//! If all the edges allowed by the proximity condition are included, then
//! the function also calculates the edge weight
//! (e.g., 1-|tau| for tree_criterion = "itau").
//!
//! @param vine_tree Tree of a vine.
inline void
VinecopSelector::add_allowed_edges(VineTree& vine_tree)
{
  std::string tree_criterion = controls_.get_tree_criterion();
  if (structure_known_) {
    double threshold = controls_.get_threshold();
    std::mutex m;
    auto add_edge = [&](size_t v0) {
      tools_interface::check_user_interrupt(v0 % 50 == 0);
      for (size_t v1 = 0; v1 < v0; ++v1) {
        // check proximity condition: common neighbor in previous tree
        // (-1 means 'no common neighbor')
        if (find_common_neighbor(v0, v1, vine_tree) > -1) {
          auto pc_data = get_pc_data(v0, v1, vine_tree);
          double crit = calculate_criterion(
            pc_data, tree_criterion, controls_.get_weights());
          double w = 1.0 - static_cast<double>(crit >= threshold) * crit;
          {
            std::lock_guard<std::mutex> lk(m);
            auto e = boost::add_edge(v0, v1, w, vine_tree).first;
            vine_tree[e].weight = w;
            vine_tree[e].crit = crit;
          }
        }
      }
    };

    pool_.map(add_edge, boost::vertices(vine_tree));
    pool_.wait();
  } else {
    size_t tree = d_ - boost::num_vertices(vine_tree);
    size_t edges = boost::num_vertices(vine_tree) - 1;
    size_t trunc_lvl = vine_struct_.get_trunc_lvl();
    if (tree < trunc_lvl) {
      for (size_t v0 = 0; v0 < edges; ++v0) {
        tools_interface::check_user_interrupt(v0 % 10000 == 0);
        size_t v1 = vine_struct_.min_array(tree, v0) - 1;
        Eigen::MatrixXd pc_data = get_pc_data(v0, v1, vine_tree);
        EdgeIterator e = boost::add_edge(v0, v1, 1.0, vine_tree).first;
        double crit = calculate_criterion(
          pc_data.leftCols(2), tree_criterion, controls_.get_weights());
        vine_tree[e].weight = 1.0;
        vine_tree[e].crit = crit;
      }
    }
  }
}

//! @brief Selects the edges using the minimum spanning tree.
//!
//! See, e.g., Czado (2010), "Pair-copula constructions of multivariate
//! copulas", url: https://mediatum.ub.tum.de/doc/1079253/file.pdf
//! @param vine_tree Tree of a vine.
inline void
VinecopSelector::select_edges(VineTree& vine_tree)
{
  // has no effect if the input is already a tree
  min_spanning_tree(vine_tree);
}

inline void
VinecopSelector::finalize(size_t trunc_lvl)
{
  pair_copulas_ = make_pair_copula_store(d_, trunc_lvl);
  trunc_lvl = pair_copulas_.size(); // trunc_lvl may be <size_t>::max()

  if (structure_known_) {
    using namespace tools_stl;
    trees_opt_ = trees_;
    TriangularArray<size_t> mat(d_, trunc_lvl);
    std::vector<size_t> order(d_);

    if (trunc_lvl > 0) {
      std::vector<size_t> ning_set;

      // fill matrix column by column
      for (size_t col = 0; col < d_ - 1; ++col) {
        tools_interface::check_user_interrupt();
        // matrix above trunc_lvl is left empty
        size_t t =
          std::max(std::min(trunc_lvl, d_ - 1 - col), static_cast<size_t>(1));
        // start with highest tree in this column
        for (auto e : boost::edges(trees_[t])) {

          // find an edge that contains a leaf
          size_t v0 = boost::source(e, trees_[t]);
          size_t v1 = boost::target(e, trees_[t]);
          size_t min_deg = std::min(boost::out_degree(v0, trees_[t]),
                                    boost::out_degree(v1, trees_[t]));
          if (min_deg > 1) {
            continue; // not a leaf
          }
          // find position of leaf in the edge
          ptrdiff_t pos = (boost::out_degree(v1, trees_[t]) == 1);
          if (pos == 1) {
            trees_[t][e].pair_copula.flip();
          }

          // fill diagonal entry with leaf index
          order[col] = trees_[t][e].conditioned[pos];

          // entry in row t-1 is other index of the edge
          mat(t - 1, col) = trees_[t][e].conditioned[std::abs(1 - pos)];

          // assign fitted pair copula to appropriate entry, see
          // `Vinecop::get_pair_copula()`.
          if (trunc_lvl > 0) {
            pair_copulas_[t - 1][col] = trees_[t][e].pair_copula;
          }

          // initialize running set with full conditioning set of this edge
          ning_set = trees_[t][e].conditioning;

          // remove edge (must not be reused in another column!)
          boost::remove_edge(v0, v1, trees_[t]);
          break;
        }

        // fill column bottom to top
        for (size_t k = 1; k < t; ++k) {
          auto check_set = cat(order[col], ning_set);
          for (auto e : boost::edges(trees_[t - k])) {
            // search for an edge in lower tree that shares all
            // indices in the conditioning set + diagonal entry
            if (!is_same_set(trees_[t - k][e].all_indices, check_set)) {
              continue;
            }
            // found suitable edge ->
            // next matrix entry is conditioned variable of new edge
            // that's not equal to the diagonal entry of this column
            auto e_new = trees_[t - k][e];
            ptrdiff_t pos = (order[col] == e_new.conditioned[1]);
            if (pos == 1) {
              e_new.pair_copula.flip();
            }
            mat(t - k - 1, col) = e_new.conditioned[std::abs(1 - pos)];

            // assign fitted pair copula to appropriate entry, see
            // Vinecop::get_pair_copula().
            pair_copulas_[t - 1 - k][col] = e_new.pair_copula;

            // start over with conditioned set of next edge
            ning_set = e_new.conditioning;

            // remove edge (must not be reused in another column!)
            size_t v0 = boost::source(e, trees_[t - k]);
            size_t v1 = boost::target(e, trees_[t - k]);
            boost::remove_edge(v0, v1, trees_[t - k]);
            break;
          }
        }
      }

      // The last column contains a single element which must be different
      // from all other diagonal elements. Based on the properties of an
      // R-vine matrix, this must be the element next to it.
      order[d_ - 1] = mat(0, d_ - 2);

      // change to user-facing format
      // (variable index starting at 1 instead of 0)
      for (size_t i = 0; i < std::min(d_ - 1, trunc_lvl); ++i) {
        for (size_t j = 0; j < d_ - i - 1; ++j) {
          mat(i, j) += 1;
        }
      }
      for (size_t i = 0; i < d_; i++)
        order[i] += 1;
    } else {
      // order doesn't matter for truncated
      order = tools_stl::seq_int(1, d_);
    }

    // return as RVineStructure
    vine_struct_ = RVineStructure(order, mat);
  } else {

    for (size_t tree = 0; tree < pair_copulas_.size(); tree++) {
      size_t edge = 0;
      for (auto e : boost::edges(trees_[tree + 1])) {
        // trees_[0] is base tree, vine copula starts at trees_[1]
        pair_copulas_[tree][edge] = trees_[tree + 1][e].pair_copula;
        edge++;
      }
    }

    vine_struct_.truncate(trunc_lvl);
  }
}

//! @brief Gets pair copula pseudo-observations from h-functions.
//! @param v0,v1 vertex indices.
//! @param tree A vine tree.
//! @return The pseudo-observations for the pair coula, extracted from
//!     the h-functions calculated in the previous tree.
inline void
VinecopSelector::add_pc_info(const EdgeIterator& e, VineTree& tree)
{
  auto v0 = boost::source(e, tree);
  auto v1 = boost::target(e, tree);
  size_t n = tree[v0].hfunc1.size();
  tree[e].pc_data = Eigen::MatrixXd(n, 2);

  // find positions of common vertex in pair indices
  size_t ei_common = find_common_neighbor(v0, v1, tree);
  ptrdiff_t pos0 = find_position(ei_common, tree[v0].prev_edge_indices);
  ptrdiff_t pos1 = find_position(ei_common, tree[v1].prev_edge_indices);

  tree[e].var_types[0] = tree[v0].var_types[std::abs(1 - pos0)];
  tree[e].var_types[1] = tree[v1].var_types[std::abs(1 - pos1)];

  // collect pseudo observations for next tree
  tree[e].pc_data.col(0) = get_hfunc(tree[v0], pos0 == 0);
  tree[e].pc_data.col(1) = get_hfunc(tree[v1], pos1 == 0);
  if ((tree[e].var_types[0] == "d") | (tree[e].var_types[1] == "d")) {
    tree[e].pc_data.conservativeResize(n, 4);
    tree[e].pc_data.col(2) = get_hfunc_sub(tree[v0], pos0 == 0);
    tree[e].pc_data.col(3) = get_hfunc_sub(tree[v1], pos1 == 0);
  }

  tree[e].conditioned =
    set_sym_diff(tree[v0].all_indices, tree[v1].all_indices);
  tree[e].conditioning = intersect(tree[v0].all_indices, tree[v1].all_indices);
  tree[e].all_indices = cat(tree[e].conditioned, tree[e].conditioning);
}

inline Eigen::VectorXd
VinecopSelector::get_hfunc(const VertexProperties& vertex_data, bool is_first)
{
  if (is_first) {
    return vertex_data.hfunc1;
  } else {
    return vertex_data.hfunc2;
  }
}

inline Eigen::VectorXd
VinecopSelector::get_hfunc_sub(const VertexProperties& vertex_data,
                               bool is_first)
{
  if (is_first) {
    if (vertex_data.hfunc1_sub.size()) {
      return vertex_data.hfunc1_sub;
    } else {
      return vertex_data.hfunc1;
    }
  } else {
    if (vertex_data.hfunc2_sub.size()) {
      return vertex_data.hfunc2_sub;
    } else {
      return vertex_data.hfunc2;
    }
  }
}

inline Eigen::MatrixXd
VinecopSelector::get_pc_data(size_t v0, size_t v1, const VineTree& tree)
{
  size_t ei_common = find_common_neighbor(v0, v1, tree);
  auto pos0 = find_position(ei_common, tree[v0].prev_edge_indices);
  auto pos1 = find_position(ei_common, tree[v1].prev_edge_indices);

  Eigen::MatrixXd pc_data(tree[v0].hfunc1.size(), 2);
  pc_data.col(0) = get_hfunc(tree[v0], pos0 == 0);
  pc_data.col(1) = get_hfunc(tree[v1], pos1 == 0);
  return pc_data;
}

//! @brief Selects and fits next tree of the vine.
//!
//! The next tree is found the following way:
//!     1. Edges of the previous tree become edges in the new tree.
//!     2. All edges allowed by the proximity condition are added to the
//!        new graph.
//!     3. Collapse the new graph to a maximum spanning tree for edge
//!        weight.
//!     4. Populate edges with conditioned/conditioning sets and pseudo-
//!        observations.
//!     5. Fit and select a copula model for each edge.
//!
//! @param prev_tree Tree T_{k}.
//! @param controls The controls for fitting a vine copula
//!     (see FitControlsVinecop).
//! @param tree_opt The current optimal tree (used only for sparse
//!     selection).
inline void
VinecopSelector::select_tree(size_t t)
{
  auto new_tree = edges_as_vertices(trees_[t]);
  remove_edge_data(trees_[t]); // no longer needed

  if (t >= vine_struct_.get_trunc_lvl()) {
    // only important if proximity_ was previously false (partial selection)
    structure_known_ = true;
  }
  add_allowed_edges(new_tree);
  if (boost::num_vertices(new_tree) > 2) {
    select_edges(new_tree);
  }
  if (boost::num_vertices(new_tree) > 0) {
    add_edge_info(new_tree);      // for pc estimation and next tree
    remove_vertex_data(new_tree); // no longer needed
    if (controls_.get_selection_criterion() == "mbicv") {
      // adjust prior probability to tree level
      controls_.set_psi0(std::pow(psi0_, t + 1));
    }
    if (trees_opt_.size() > t + 1) {
      select_pair_copulas(new_tree, trees_opt_[t + 1]);
    } else {
      select_pair_copulas(new_tree);
    }
  }
  // make sure there is space for new tree
  trees_.resize(t + 2);
  trees_[t + 1] = new_tree;
}

inline double
VinecopSelector::get_mbicv_of_tree(size_t t, double loglik)
{
  double npars = get_npars_of_tree(t);
  size_t non_indeps = get_num_non_indeps_of_tree(t);
  size_t indeps = d_ - t - 1 - non_indeps;
  double psi0 = std::pow(psi0_, t + 1);
  double log_prior = static_cast<double>(non_indeps) * std::log(psi0) +
                     static_cast<double>(indeps) * std::log(1.0 - psi0);
  double n_eff = static_cast<double>(n_);
  if (controls_.get_weights().size() > 0) {
    n_eff = std::pow(controls_.get_weights().sum(), 2);
    n_eff /= controls_.get_weights().array().pow(2).sum();
  }

  return -2 * loglik + std::log(n_eff) * npars - 2 * log_prior;
}

//! @brief Calculates the log-likelihood of a tree.
inline double
VinecopSelector::get_loglik_of_tree(size_t t)
{
  double ll = 0.0;
  // trees_[0] is base tree, see make_base_tree()
  for (const auto& e : boost::edges(trees_[t + 1])) {
    ll += trees_[t + 1][e].pair_copula.get_loglik();
  }
  return ll;
}

//! @brief Calculates the numbers of parameters of a tree.
inline double
VinecopSelector::get_npars_of_tree(size_t t)
{
  double npars = 0.0;
  // trees_[0] is base tree, see make_base_tree()
  for (const auto& e : boost::edges(trees_[t + 1])) {
    npars += trees_[t + 1][e].pair_copula.get_npars();
  }
  return npars;
}

//! @brief Calculates the numbers of independence copulas in a tree.
inline size_t
VinecopSelector::get_num_non_indeps_of_tree(size_t t)
{
  size_t num_non_indeps = 0;
  // trees_[0] is base tree, see make_base_tree()
  for (const auto& e : boost::edges(trees_[t + 1])) {
    num_non_indeps += static_cast<size_t>(
      trees_[t + 1][e].pair_copula.get_family() == BicopFamily::indep);
  }
  return num_non_indeps;
}

//! @brief Prints indices, family, and parameters for each pair-copula
//! @param tree A vine tree.
inline void
VinecopSelector::print_pair_copulas_of_tree(size_t t)
{
  // trees_[0] is the base tree, see make_base_tree()
  for (auto e : boost::edges(trees_[t + 1])) {
    std::cout << get_pc_index(e, trees_[t + 1]) << " <-> "
              << trees_[t + 1][e].pair_copula.str() << std::endl;
  }
}

//! @brief Gets all criterion values that got thresholded to zero.
inline std::vector<double>
VinecopSelector::get_thresholded_crits()
{
  std::vector<double> crits;
  for (size_t t = 1; t < trees_.size(); ++t) {
    for (auto e : boost::edges(trees_[t])) {
      if (trees_[t][e].crit < controls_.get_threshold()) {
        crits.push_back(trees_[t][e].crit);
      }
    }
  }

  return crits;
}

inline void
VinecopSelector::initialize_new_fit(const Eigen::MatrixXd& data)
{
  trees_[0] = make_base_tree(data);
}

inline void
VinecopSelector::set_current_fit_as_opt(const double& loglik)
{
  threshold_ = controls_.get_threshold();
  trees_opt_ = trees_;
  loglik_ = loglik;
}

//! @brief Instantiates base tree of the vine.
//!
//! The base tree is a star on d + 1 variables, where the conditioned
//! set of each edge consists of a single number. When building the next
//! tree, the edges become vertices. Because the base graph was a star
//! all edges are allowed by the proximity condition, and the edges will
//! have a conditioned set consisting of the two vertex indices. This
//! will be the first actual tree of the vine.
//!
//! @param data nxd matrix of copula data.
//! @return A VineTree object containing the base graph.
inline VineTree
VinecopSelector::make_base_tree(const Eigen::MatrixXd& data)
{
  VineTree base_tree(d_);
  auto order = vine_struct_.get_order();
  auto disc_cols = get_disc_cols(var_types_);

  // a star connects the root node (d) with all other nodes
  for (size_t target = 0; target < d_; ++target) {
    tools_interface::check_user_interrupt(target % 10000 == 0);
    // add edge and extract edge iterator
    auto e = add_edge(d_, target, base_tree).first;
    // inititialize hfunc1 with actual data for variable "target"
    // data need are reordered to correspond to natural order (neccessary
    // when structure is fixed)
    base_tree[e].hfunc1 = data.col(order[target] - 1);
    if (var_types_[order[target] - 1] == "d") {
      base_tree[e].hfunc1_sub = data.col(d_ + disc_cols[order[target] - 1]);
      base_tree[e].var_types = { "d", "d" };
    }

    // identify edge with variable "target" and initialize sets
    base_tree[e].conditioned.reserve(2);
    base_tree[e].conditioned.push_back(order[target] - 1);
    base_tree[e].conditioning.reserve(d_ - 2);
    base_tree[e].all_indices = base_tree[e].conditioned;
  }

  return base_tree;
}

//! @brief Converts the edge set into vertex set of a new graph.
//!
//! Further information about the structure is passed along:
//!     - conditioned/conditioning set,
//!     - indices of vertices connected by the edge in the previous tree.
//!
//! @param tree T_{k}.
//! @return A edge-less graph of vertices, each representing one edge of the
//!     previous tree.
inline VineTree
VinecopSelector::edges_as_vertices(const VineTree& prev_tree)
{
  // start with full graph
  size_t d = num_edges(prev_tree);
  VineTree new_tree(d);

  // copy & paste information from previous tree
  int i = 0;
  for (auto e : boost::edges(prev_tree)) {
    new_tree[i].hfunc1 = prev_tree[e].hfunc1;
    new_tree[i].hfunc2 = prev_tree[e].hfunc2;
    new_tree[i].hfunc1_sub = prev_tree[e].hfunc1_sub;
    new_tree[i].hfunc2_sub = prev_tree[e].hfunc2_sub;
    new_tree[i].conditioned = prev_tree[e].conditioned;
    new_tree[i].conditioning = prev_tree[e].conditioning;
    new_tree[i].all_indices = prev_tree[e].all_indices;
    new_tree[i].prev_edge_indices.reserve(2);
    new_tree[i].prev_edge_indices.push_back(boost::source(e, prev_tree));
    new_tree[i].prev_edge_indices.push_back(boost::target(e, prev_tree));
    new_tree[i].var_types = prev_tree[e].var_types;
    ++i;
  }

  return new_tree;
}

//! @brief Finds common neighbor in previous tree.
//! @param v0,v1 vertices in the tree.
//! @param tree the current tree.
//! @return Gives the index of the vertex in the previous tree that was
//!     shared by e0, e1, the edge representations of v0, v1.
inline ptrdiff_t
VinecopSelector::find_common_neighbor(size_t v0,
                                      size_t v1,
                                      const VineTree& tree)
{
  auto ei0 = tree[v0].prev_edge_indices;
  auto ei1 = tree[v1].prev_edge_indices;
  auto ei_common = intersect(ei0, ei1);

  if (ei_common.size() == 0) {
    return -1;
  } else {
    return ei_common[0];
  }
}

//! @brief Computes a fit id; can be used to re-use already fitted pair-copulas.
//! @param edge.
inline double
VinecopSelector::compute_fit_id(const EdgeProperties& e)
{
  double id = 0.0;
  if (controls_.needs_sparse_select()) {
    // the formula is quite arbitrary, but sufficient for
    // identifying situations where fits can be re-used
    id = (e.pc_data.col(0) - 2 * e.pc_data.col(1)).sum();
    id += 5.0 * static_cast<double>(e.crit < controls_.get_threshold());
  }

  return id;
}

//! @brief Collapses a graph to the minimum spanning tree.
//! @param graph The input graph.
//! @return the input graph with all non-MST edges removed.
inline void
VinecopSelector::min_spanning_tree(VineTree& graph)
{
  size_t d = num_vertices(graph);
  std::vector<size_t> targets(d);
  prim_minimum_spanning_tree(graph, targets.data());
  for (size_t v1 = 0; v1 < d; ++v1) {
    for (size_t v2 = 0; v2 < v1; ++v2) {
      if ((v2 != targets[v1]) & (v1 != targets[v2])) {
        boost::remove_edge(v1, v2, graph);
      }
    }
  }
}

//! @brief Adds conditioned info and data for each edge.
//!
//! See, e.g., Czado (2010), "Pair-copula constructions of multivariate
//! copulas", url: https://mediatum.ub.tum.de/doc/1079253/file.pdf
//! @param tree A vine tree.
inline void
VinecopSelector::add_edge_info(VineTree& tree)
{
  for (auto e : boost::edges(tree)) {
    add_pc_info(e, tree);
  }
}

//! @brief Removes data (hfunc1/hfunc2/pc_data) from all edges of a vine tree.
//! @param tree A vine tree.
inline void
VinecopSelector::remove_edge_data(VineTree& tree)
{
  for (auto e : boost::edges(tree)) {
    tree[e].hfunc1 = Eigen::VectorXd();
    tree[e].hfunc2 = Eigen::VectorXd();
    tree[e].hfunc1_sub = Eigen::VectorXd();
    tree[e].hfunc2_sub = Eigen::VectorXd();
    tree[e].pc_data = Eigen::MatrixXd(0, 2);
  }
}

//! @brief Removes data (hfunc1/hfunc2/pc_data) from all vertices of a vine
//! tree.
//! @param tree A vine tree.
inline void
VinecopSelector::remove_vertex_data(VineTree& tree)
{
  for (auto v : boost::vertices(tree)) {
    tree[v].hfunc1 = Eigen::VectorXd();
    tree[v].hfunc2 = Eigen::VectorXd();
    tree[v].hfunc1_sub = Eigen::VectorXd();
    tree[v].hfunc2_sub = Eigen::VectorXd();
  }
}

//! @brief Fits and selects a pair copula for each edges.
//! @param tree A vine tree preprocessed with `add_edge_info()`.
//! @param tree_opt The current optimal tree (used only for sparse
//!     selection).
inline void
VinecopSelector::select_pair_copulas(VineTree& tree, const VineTree& tree_opt)
{
  auto select_pc = [&](EdgeIterator e) -> void {
    tools_interface::check_user_interrupt();
    bool is_thresholded = (tree[e].crit < controls_.get_threshold());
    bool used_old_fit = false;

    tree[e].fit_id = compute_fit_id(tree[e]);
    if (boost::num_edges(tree_opt) > 0) {
      auto old_fit = find_old_fit(tree[e].fit_id, tree_opt);
      if (old_fit.second) { // indicates if match was found
        // data and thresholding status haven't changed,
        // we can use old fit
        used_old_fit = true;
        tree[e].pair_copula = tree_opt[old_fit.first].pair_copula;
      }
    }

    if (!used_old_fit) {
      tree[e].pair_copula = vinecopulib::Bicop();
      tree[e].pair_copula.set_var_types(tree[e].var_types);
      if (!is_thresholded) {
        tree[e].pair_copula.select(tree[e].pc_data, controls_);
      }
    }

    tree[e].hfunc1 = tree[e].pair_copula.hfunc1(tree[e].pc_data);
    tree[e].hfunc2 = tree[e].pair_copula.hfunc2(tree[e].pc_data);
    if (tree[e].var_types[1] == "d") {
      auto sub_data = tree[e].pc_data;
      sub_data.col(1) = sub_data.col(3);
      tree[e].hfunc1_sub = tree[e].pair_copula.hfunc1(sub_data);
    }
    if (tree[e].var_types[0] == "d") {
      auto sub_data = tree[e].pc_data;
      sub_data.col(0) = sub_data.col(2);
      tree[e].hfunc2_sub = tree[e].pair_copula.hfunc2(sub_data);
    }
  };

  // make sure that Bicop.select() doesn't spawn new threads
  size_t num_threads = controls_.get_num_threads();
  controls_.set_num_threads(0);
  pool_.map(select_pc, boost::edges(tree));
  pool_.wait();
  controls_.set_num_threads(num_threads);
}

//! @brief Finds the fitted pair-copula from the previous iteration.
inline FoundEdge
VinecopSelector::find_old_fit(double fit_id, const VineTree& old_graph)
{
  auto edge = boost::edge(0, 1, old_graph).first;
  bool fit_with_same_id = false;
  for (auto e : boost::edges(old_graph)) {
    if (fit_id == old_graph[e].fit_id) {
      fit_with_same_id = true;
      edge = e;
    }
  }
  return std::make_pair(edge, fit_with_same_id);
}

//! @brief Gets edge index for the vine (like 1, 2; 3).
//! @param e A descriptor for the edge.
//! @param tree A vine tree.
inline std::string
VinecopSelector::get_pc_index(const EdgeIterator& e, const VineTree& tree)
{
  std::stringstream index;
  // add 1 everywhere for user-facing representation (boost::graph
  // starts at 0)
  index << tree[e].conditioned[0] + 1 << "," << tree[e].conditioned[1] + 1;
  if (tree[e].conditioning.size() > 0) {
    index << " | ";
    for (unsigned int i = 0; i < tree[e].conditioning.size(); ++i) {
      index << tree[e].conditioning[i] + 1;
      if (i < tree[e].conditioning.size() - 1)
        index << ",";
    }
  }

  return index.str().c_str();
}
}
}


/***************************************************************************
 *  @file       longest_common_subsequence.hpp
 *  @author     alan.w
 *  @date       09  August 2014
 *  @remark     CLRS Algorithms implementation, using C++ templates.
 ***************************************************************************/

#ifndef LONGEST_COMMON_SUBSEQUENCE_HPP
#define LONGEST_COMMON_SUBSEQUENCE_HPP

#include <functional>
#include "matrix.hpp"
#include "color.hpp"

namespace ch15 {
/**
 * @brief  as a part of an entry
 */
enum class Arrow
{
    NA, LEFT, UP, DIAGONAL
};

/**
 * @brief The Entry struct
 *
 * as an entry of Table
 */
template<typename T, typename A>
struct Entry
{
    //! Ctors
    Entry() = default;
    Entry(const T& len, const A& arw = Arrow::NA):
        length(len), arrow(arw)
    {}

    //! update this entry
    Entry operator ()(const T& len, const A& arw = Arrow::NA)
    {
        length  =   len;
        arrow   =   arw;

        return *this;
    }

    T length = 0;
    A arrow = Arrow::NA;
};

/**
 * @brief operator <<
 *
 * called when printing a boost matrix
 */
template<typename T, typename A>
inline std::ostream&
operator <<(std::ostream& os , const Entry<T,A>& etr)
{
    os << color::green("[")
       << etr.length << ",";

    switch (etr.arrow)
    {
    case Arrow::NA      :   os << " ";      break;
    case Arrow::LEFT    :   os << "<";      break;
    case Arrow::DIAGONAL:   os << "\\";     break;
    case Arrow::UP      :   os << "^";      break;
    }

    os << color::green("]");

    return os;
}

/**
 * @brief operator >=
 *
 * overloaded for comparison in build_lcs_table
 */
template<typename T, typename A>
inline bool
operator >=(const Entry<T,A>& lhs, const Entry<T,A>& rhs)
{
    return lhs.length >= rhs.length;
}

//! @alias for Table
template<typename T, typename A>
using Table = ch15::Matrix<Entry<T,A>>;

/**
 * @brief build lcs table
 * @param lhs   X
 * @param rhs   Y
 *
 * @pseudocode  LCS-LENGTH
 * @page    394, CLRS
 * @complx  O(mn)
 *
 * @note    the tables b and c in the peudocode have been merged in to one.
 *          so only one table is returned
 */
template<typename Range, typename A = Arrow>
Table<typename Range::size_type, A>
build_lcs_table(const Range& lhs, const Range& rhs)
{
    //! types def
    using SizeType  =   typename Range::size_type;
    using TableType =   Table<SizeType, A>;

    //! init
    SizeType rows = lhs.size() + 1;
    SizeType cols = rhs.size() + 1;
    TableType lcs(rows, cols);
    //! ^^^^^^^^^^^^^^^^^^^^^^@attention  :
    //!     the Entry's default ctor grantee that its member length has
    //!     default value 0.so no need to implementlines 4 - 7 explicitly.

    //! build lcs table
    for(SizeType r = 1; r != rows; ++r)
        for(SizeType c = 1; c != cols; ++c)
        {
            if(lhs[r - 1]   ==  rhs[c - 1])
                lcs(r,c) (lcs(r - 1, c - 1).length + 1   ,  Arrow::DIAGONAL);
            else if(lcs(r - 1, c) >= lcs(r, c - 1))
                lcs(r,c) (lcs(r - 1, c).length           ,  Arrow::UP);
            else
                lcs(r,c) (lcs(r, c - 1).length           ,  Arrow::LEFT);
        }

    return lcs;
}

/**
 * @brief The LongestCommonSubsequence class
 */
template<typename Range, typename A = ch15::Arrow>
class LongestCommonSubsequence
{
public:
    //! types def
    using SizeType  =   typename Range::size_type;
    using TableType =   ch15::Table<SizeType, A>;
    using Pointer   =   const Range*;

    //! Ctor
    LongestCommonSubsequence(const Range& l, const Range& r):
        lhs(&l),
        rhs(&r),
        maze(ch15::build_lcs_table(l,r))
    {}

    /**
     * @brief print_maze
     *
     * print the maze built by build_lcs_table
     */
    void print_maze()const
    {
        ch15::print(maze);
    }

    /**
     * @brief return the generated longest common sequence
     *
     * @pseudocode PRINT-LCS
     * @page    395, CLRS
     * @complx  O(m + n)
     */
    Range generate() const
    {
        assert(lhs && rhs);
        using Lambda = std::function<void(SizeType, SizeType)>;
        Range lcs;

        //! a recursive lamda that performs the real work
        Lambda build_lcs = [&lcs, &build_lcs, this](SizeType row, SizeType col)
        {
            //! stop condition
            if(row == 0 || col == 0)    return;

            //! build the longest common sequence
            if(maze(row, col).arrow ==  Arrow::DIAGONAL)
            {
                build_lcs(row - 1, col - 1);
                lcs.push_back((*lhs)[row - 1]);
            }
            else if(maze(row, col).arrow ==  Arrow::UP)
                build_lcs(row - 1, col);
            else
                build_lcs(row, col - 1);
        };

        //! call the lambda
        build_lcs(lhs->size(), rhs->size());
        return lcs;
    }

private:
    Pointer lhs;
    Pointer rhs;
    TableType maze;
};

}//namespace
#endif // LONGEST_COMMON_SUBSEQUENCE_HPP

//! @test   print ch15::table
//#include <iostream>
//#include <boost/numeric/ublas/io.hpp>
//#include "longest_common_subsequence.hpp"
//#include "color.hpp"

//int main()
//{
//    using TableType = ch15::Table<int, ch15::Arrow>;

//    TableType tbl(3,3, ch15::Entry<int, ch15::Arrow>());
//    std::cout << tbl << std::endl;

//    std::cout << color::red("\nend\n");
//    return 0;
//}

//! @test   build_lcs_table
//!         i.e. LCS-LENGTH
//!
//#include <iostream>
//#include <boost/numeric/ublas/io.hpp>
//#include "longest_common_subsequence.hpp"
//#include "color.hpp"

//int main()
//{
//    //! the same strings as that on page 394, CLRS.
//    std::vector<std::string> lhs = {"A", "B", "C", "B", "D", "A", "B"};
//    std::vector<std::string> rhs = {"B", "D", "C", "A", "B", "A"};

//    auto ret = ch15::build_lcs_table(lhs, rhs);
//    ch15::print(ret);

//    std::cout << color::red("\nend\n");
//    return 0;
//}

//! @test   The LongestCommonSubsequence class
//!         i.e. LCS-LENGTH and PRINT-LCS
//!
//#include <iostream>
//#include <boost/numeric/ublas/io.hpp>
//#include "longest_common_subsequence.hpp"
//#include "color.hpp"

//int main()
//{
//    //! strings used on page 394, CLRS.
//    std::string lhs = "ABCBDAB";
//    std::string rhs = "BDCABA";

//    using LCS   =   ch15::LongestCommonSubsequence<std::string>;
//    LCS lcs(lhs, rhs);
//    lcs.print_maze();

//    auto sequence = lcs.generate();
//    std::cout << "The longest common sequence = ";
//    std::cout << color::yellow(sequence) << std::endl;

//    std::cout << color::red("\nend\n");
//    return 0;
//}
//! @output:
//[0, ] [0, ] [0, ] [0, ] [0, ] [0, ] [0, ]

//[0, ] [0,^] [0,^] [0,^] [1,\] [1,<] [1,\]

//[0, ] [1,\] [1,<] [1,<] [1,^] [2,\] [2,<]

//[0, ] [1,^] [1,^] [2,\] [2,<] [2,^] [2,^]

//[0, ] [1,\] [1,^] [2,^] [2,^] [3,\] [3,<]

//[0, ] [1,^] [2,\] [2,^] [2,^] [3,^] [3,^]

//[0, ] [1,^] [2,^] [2,^] [3,\] [3,^] [4,\]

//[0, ] [1,\] [2,^] [2,^] [3,^] [4,\] [4,^]

//The longest common sequence = BCBA

//end


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

// Note: this code has been downloaded from the homepage of the "Computer
// Vision Laboratory" at EPFL Lausanne, and was originally developped by the
// authors of [4]. I only adapted it to Eigen.

#ifndef OPENGV_ABSOLUTE_POSE_MODULES_EPNP_HPP_
#define OPENGV_ABSOLUTE_POSE_MODULES_EPNP_HPP_

#include <stdlib.h>
#include <Eigen/Eigen>

namespace opengv
{
namespace absolute_pose
{
namespace modules
{

class Epnp
{
public:
  Epnp(void);
  ~Epnp();

  void set_maximum_number_of_correspondences(const int n);
  void reset_correspondences(void);
  void add_correspondence(
      const double X,
      const double Y,
      const double Z,
      const double x,
      const double y,
      const double z);

  double compute_pose(double R[3][3], double T[3]);

  void relative_error(
      double & rot_err,
      double & transl_err,
      const double Rtrue[3][3],
      const double ttrue[3],
      const double Rest[3][3],
      const double test[3]);

  void print_pose(const double R[3][3], const double t[3]);
  double reprojection_error(const double R[3][3], const double t[3]);

private:
  void choose_control_points(void);
  void compute_barycentric_coordinates(void);
  void fill_M(
      Eigen::MatrixXd & M,
      const int row,
      const double * alphas,
      const double u,
      const double v);
  void compute_ccs(const double * betas, const Eigen::MatrixXd & ut);
  void compute_pcs(void);

  void solve_for_sign(void);

  void find_betas_approx_1(
      const Eigen::Matrix<double,6,10> & L_6x10,
      const Eigen::Matrix<double,6,1> & Rho,
      double * betas);
  void find_betas_approx_2(
      const Eigen::Matrix<double,6,10> & L_6x10,
      const Eigen::Matrix<double,6,1> & Rho,
      double * betas);
  void find_betas_approx_3(
      const Eigen::Matrix<double,6,10> & L_6x10,
      const Eigen::Matrix<double,6,1> & Rho,
      double * betas);
  void qr_solve(
      Eigen::Matrix<double,6,4> & A,
      Eigen::Matrix<double,6,1> & b,
      Eigen::Matrix<double,4,1> & X);

  double dot(const double * v1, const double * v2);
  double dist2(const double * p1, const double * p2);

  void compute_rho(Eigen::Matrix<double,6,1> & Rho);
  void compute_L_6x10(
      const Eigen::MatrixXd & Ut,
      Eigen::Matrix<double,6,10> & L_6x10 );

  void gauss_newton(
      const Eigen::Matrix<double,6,10> & L_6x10,
      const Eigen::Matrix<double,6,1> & Rho,
      double current_betas[4]);
  void compute_A_and_b_gauss_newton(
      const Eigen::Matrix<double,6,10> & L_6x10,
      const Eigen::Matrix<double,6,1> & Rho,
      double cb[4],
      Eigen::Matrix<double,6,4> & A,
      Eigen::Matrix<double,6,1> & b);

  double compute_R_and_t(
      const Eigen::MatrixXd & Ut,
      const double * betas,
      double R[3][3],
      double t[3]);

  void estimate_R_and_t(double R[3][3], double t[3]);

  void copy_R_and_t(
      const double R_dst[3][3],
      const double t_dst[3],
      double R_src[3][3],
      double t_src[3]);

  void mat_to_quat(const double R[3][3], double q[4]);

  double uc, vc, fu, fv;

  double * pws, * us, * alphas, * pcs;
  int * signs; //added!
  int maximum_number_of_correspondences;
  int number_of_correspondences;

  double cws[4][3], ccs[4][3];
  double cws_determinant;
};

}
}
}

#endif /* OPENGV_ABSOLUTE_POSE_MODULES_EPNP_HPP_ */


#pragma once

#define EIGEN_NO_AUTOMATIC_RESIZING

#include "Tableau.hpp"
#include "SM_utils.hpp"
#include "Solids.hpp"
#include "ReplacementDictionary.hpp"
#include "Equilibrium.hpp"

#include <Eigen/Dense>

#include <stdexcept>
#include <algorithm>
#include <cstddef>
#include <tuple>
#include <utility>

namespace AQSystemSolver {
    [[nodiscard]] inline auto possiblyAddSolid(const Eigen::VectorXd& solubilityProducts, ReplacementDict& replacementDict, SolidSystem& solidSystem){
        const auto addAttempt=solidSystem.getSolidToAdd(solubilityProducts);
        if(addAttempt.success) [[likely]] {
            replacementDict.addSolid(addAttempt.solid);
        }
        return addAttempt;
    }
    [[nodiscard]] inline auto possiblyRemoveSolid(const Eigen::VectorXd& solidAmts, ReplacementDict& replacementDict, SolidSystem& solidSystem){
        const auto removeAttempt=solidSystem.getSolidToRemove(solidAmts);
        if(removeAttempt.success) [[unlikely]] {
            replacementDict.removeSolid(removeAttempt.solid);
        }
        return removeAttempt;
    }

    inline constexpr double GUESS_TOTALS_FACTOR=1e-5;

    inline constexpr double MAX_PERCENT_ERROR_ADDEND=1e-5;

    inline constexpr std::size_t MAX_ITER=30;

    inline constexpr double STEP_WHEN_NEGATIVE=0.9;

    [[nodiscard]] std::pair<Eigen::RowVectorXd, Eigen::VectorXd> inline solveWithReplacement(const TableauWithTotals<>& replacedTableau) {

        Eigen::RowVectorXd currentSolution=Eigen::RowVectorXd::Constant(replacedTableau.cols(), GUESS_TOTALS_FACTOR);
        for(std::size_t iter=0; iter<MAX_ITER; ++iter){
            Eigen::VectorXd speciesConcentrations=replacedTableau.evalTerms(currentSolution);

            const Eigen::MatrixXd addends=replacedTableau.evalAddends(speciesConcentrations);

            const Eigen::RowVectorXd maxAddend=abs(addends.array()).colwise().maxCoeff();

            //x^c*c/x=(x^c)'
            const Eigen::MatrixXd jacobian=addends.transpose()*(replacedTableau.getCoefficients().array().rowwise()/currentSolution.array()).matrix();

            const Eigen::RowVectorXd yResult=replacedTableau.eval(addends);

            if(abs(yResult.array()/maxAddend.array()).maxCoeff()<MAX_PERCENT_ERROR_ADDEND) [[unlikely]] {
                return std::make_pair(std::move(currentSolution), std::move(speciesConcentrations));
            }

            Eigen::RowVectorXd delta=jacobian.partialPivLu().solve(yResult.transpose()).transpose();

            //don't go into the negatives, instead divide by 10
            delta=(delta.array()<currentSolution.array()).select(delta, currentSolution*STEP_WHEN_NEGATIVE);

            currentSolution-=delta;
        }
        throw std::runtime_error("Failed to converge");
    }

    inline auto solveForEquilibrium(const TableauWithTotals<>& tableau, const SolidSystem& initialSolidSystem, const ReplacementDict& origReplacementDict){
        Eigen::RowVectorXd currentSolution;
        Eigen::VectorXd speciesConcentrations;
        Eigen::VectorXd solubilityProducts;
        Eigen::VectorXd solidAmts;

        ReplacementDict replacementDict{origReplacementDict};
        SolidSystem solidSystem=initialSolidSystem.createNewWithInitialConditions();
        replacementDict.addSolidSystem(solidSystem);

        for(;;){
            const auto currentReplacedTableau=replacementDict.createReplacedTableau(tableau);
            if(currentReplacedTableau.cols()!=0){
                std::tie(currentSolution, speciesConcentrations)=solveWithReplacement(currentReplacedTableau);
            } else {
                currentSolution=Eigen::RowVectorXd(0);
                speciesConcentrations=currentReplacedTableau.getConstants();
            }
            solubilityProducts=solidSystem.calculateSolubilityProducts(currentSolution, replacementDict);

            solidAmts=solidSystem.calculateSolidAmts(-tableau.eval(speciesConcentrations));

            SolidSystem::SolidChangeAttempt removalAttempt=possiblyRemoveSolid(solidAmts, replacementDict, solidSystem);
            if(removalAttempt.success) {
                continue;
            }
            SolidSystem::SolidChangeAttempt addAttempt=possiblyAddSolid(solubilityProducts, replacementDict, solidSystem);
            if(addAttempt.success) {
                continue;
            }

            //there used to be a huge piece of code to try and switch two solids (remove one and add another within a single iteration). It was a mess, huge, and didn't really help recovery
            if(removalAttempt.solid!=nullptr){
                //we've looped, and the adding didn't help
                throw std::runtime_error("failed to recover from loop. You may have to provide an initial guess for which solids are present.");
            }
            if(addAttempt.solid!=nullptr){
                //we've looped while trying to add.
                throw std::runtime_error("failed to recover from loop. You may have to provide an initial guess for which solids are present.");
            }
            break;
        }
        return Equilibrium{
            tableau,
            std::move(replacementDict), std::move(currentSolution), std::move(speciesConcentrations),
            std::move(solidSystem), solidAmts, solubilityProducts
        };
    }
} // namespace AQSystemSolver


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

#ifndef SMOOTH__FEEDBACK__MPC_FUNC_HPP_
#define SMOOTH__FEEDBACK__MPC_FUNC_HPP_

/**
 * @file
 * @brief Functions for Model-Predictive Control (MPC) on Lie groups.
 */

#include <Eigen/Core>

#include <cassert>
#include <chrono>
#include <smooth/diff.hpp>
#include <smooth/lie_group.hpp>

#include "common.hpp"
#include "qp.hpp"

namespace smooth::feedback {

/**
 * @brief Optimal control problem defintiion.
 *
 * @tparam G State space Lie group.
 * @tparam U Input space Lie group.
 *
 * The optimal control problem is
 * \f[
 *   \begin{cases}
 *    \min_{u(\cdot)} & \int_{0}^T \left( (g(t) \ominus g_{des}(t))^T Q (g(t) \ominus g_{des}(t))^T
 *     + (u(t) \ominus u_{des}(t))^T R (u(t) \ominus u_{des}(t))^T \right)
 *     + (g(T) \ominus g_{des}(T))^T Q_T (g(T) \ominus g_{des}(T))               \\
 *    \text{s.t.}    & g(0) = x_0,                                               \\
 *                    & l_G \leq A_G (g(t) \ominus c_G) \leq u_G,                \\
 *                    & l_U \leq A_U (u(t) \ominus c_U) \leq u_U,
 *   \end{cases}
 * \f]
 * where the cost matrices must be positive semi-definite.
 */
template<LieGroup G, Manifold U>
struct OptimalControlProblem
{
  /// State tangent dimension
  static constexpr Eigen::Index Nx = Dof<G>;
  /// Input tangent dimension
  static constexpr Eigen::Index Nu = Dof<U>;

  /// Time horizon
  double T{1};
  /// Initial state
  G x0{Default<G>()};

  /// Desired input trajectory
  std::function<U(double)> udes = [](double) -> U { return Default<U>(); };
  /// Desired state trajectory
  std::function<G(double)> gdes = [](double) -> G { return Default<G>(); };

  /// Input bounds
  ManifoldBounds<U> ulim{};
  /// State bounds
  ManifoldBounds<G> glim{};

  /// MPC weights struct
  struct Weights
  {
    /// Running state cost
    Eigen::Matrix<double, Nx, Nx> Q = Eigen::Matrix<double, Nx, Nx>::Identity();
    /// Final state cost
    Eigen::Matrix<double, Nx, Nx> QT = Eigen::Matrix<double, Nx, Nx>::Identity();
    /// Running input cost
    Eigen::Matrix<double, Nu, Nu> R = Eigen::Matrix<double, Nu, Nu>::Identity();
  };

  /// MPC weights values
  Weights weights{};
};

/**
 * @brief Struct to define a linearization point.
 */
template<LieGroup G, Manifold U>
struct LinearizationInfo
{
  /**
   * @brief state linearization trajectory with first derivative
   * \f$ g_{lin}: \mathbb{R} \rightarrow (G, \mathbb{R}^{\dim \mathfrak{g}}) \f$
   */
  std::function<std::pair<G, Tangent<G>>(double)> g = [](double) -> std::pair<G, Tangent<G>> {
    return {Default<G>(), Tangent<G>::Zero()};
  };

  /**
   * @brief input linearization trajectory \f$ u_{lin}(t) :  \mathbb{R} \rightarrow U \f$
   */
  std::function<U(double)> u = [](double) -> U { return Default<U>(); };

  /**
   * @brief Domain of validity of state linearization
   *
   *  Defines an upper bound \f$ \bar a \f$ s.t. the linearization is valid for \f$ g \f$ s.t.
   * \f[
   *   -\bar a \leq g \ominus_r g_{lin} \leq \bar a
   * \f]
   * holds component-wise.
   */
  Eigen::Matrix<double, Dof<G>, 1> g_domain =
    Eigen::Matrix<double, Dof<G>, 1>::Constant(std::numeric_limits<double>::infinity());
};

/**
 * @brief Allocate QP sparsity pattern (part 1 of ocp_to_qp()).
 *
 * Variables: [x_1, ..., x_K, u_0, ..., u_{K-1}]
 *
 * Constraints:
 *  - Dynamics constraints            (K * Nx)
 *  - Input constraints               (Nu_ineq * Nu)
 *  - State constraints               (Nx_ineq * Nx)
 *  - State linearization constraints (K * Nu)        [optional]
 *
 * @param[in] pbm OptimalControlProblem definition.
 * @param[in] K number of time discretization steps
 * @param[out] qp quadratic program definition with allocated matrices.
 * @param[in] lin_con set to true to allocate K * Nu state linearization constraints
 */
template<LieGroup G, Manifold U>
  requires(Dof<G> > 0 && Dof<U> > 0)
void ocp_to_qp_allocate(
  const OptimalControlProblem<G, U> & pbm,
  std::size_t K,
  QuadraticProgramSparse<double> & qp,
  bool lin_con = false)
{
  // problem info
  static constexpr int Nx = Dof<G>;
  static constexpr int Nu = Dof<U>;

  static_assert(Nx > 0, "State space dimension must be static");
  static_assert(Nu > 0, "Input space dimension must be static");

  const uint32_t n_eq     = K * Nx;
  const uint32_t NU_iq    = K * pbm.ulim.A.rows();
  const uint32_t NX_iq    = K * pbm.glim.A.rows();
  const uint32_t NXLIN_iq = lin_con ? K * Nu : 0;

  const uint32_t nvar = K * Nx + K * Nu;
  const uint32_t ncon = n_eq + NU_iq + NX_iq + NXLIN_iq;

  // Matrix sizes
  qp.P.resize(nvar, nvar);
  qp.q.resize(nvar);

  qp.A.resize(ncon, nvar);
  qp.l.resize(ncon);
  qp.u.resize(ncon);

  // SPARSITY PATTERN FOR P

  Eigen::Matrix<int, -1, 1> Pp(nvar);
  for (auto k = 0u; k != K; ++k) {
    for (auto i = 0u; i != Nu; ++i) { Pp[k * Nu + i] = 1 + i; }
  }

  for (auto k = 0u; k != K; ++k) {
    for (auto i = 0u; i != Nx; ++i) { Pp[K * Nu + k * Nx + i] = 1 + i; }
  }

  // SPARSITY PATTERN FOR A

  Eigen::Matrix<int, -1, 1> Ap(ncon);
  int Arow = 0;

  // state constraint k = 0
  Ap.segment(Arow, Nx).setConstant(1 + Nu);
  Arow += Nx;

  // state constraints k = 1, ... K
  Ap.segment(Arow, (K - 1) * Nx).setConstant(1 + Nu + Nx);
  Arow += (K - 1) * Nx;

  // input constraints
  if (NU_iq > 0) { Ap.segment(Arow, NU_iq).setConstant(Nu); }
  Arow += NU_iq;

  // state constraints
  if (NX_iq > 0) { Ap.segment(Arow, NX_iq).setConstant(Nx); }
  Arow += NX_iq;

  // state linearization constraints
  if (NXLIN_iq > 0) { Ap.segment(Arow, NXLIN_iq).setConstant(1); }

  qp.P.reserve(Pp);
  qp.A.reserve(Ap);
}

/**
 * @brief Fill QP matrices (part 2 of ocp_to_qp()).
 */
template<LieGroup G, Manifold U, typename Dyn, diff::Type DT = diff::Type::Default>
  requires(Dof<G> > 0 && Dof<U> > 0)
void ocp_to_qp_fill(
  const OptimalControlProblem<G, U> & pbm,
  std::size_t K,
  const Dyn & f,
  const LinearizationInfo<G, U> & lin,
  QuadraticProgramSparse<double> & qp)
{
  using std::placeholders::_1;

  static constexpr int Nx = Dof<G>;
  static constexpr int Nu = Dof<U>;
  const int NU            = K * Nu;

  static_assert(Nx > 0, "State space dimension must be static");
  static_assert(Nu > 0, "Input space dimension must be static");

  const bool lin_con = lin.g_domain.minCoeff() < std::numeric_limits<double>::infinity();

  const uint32_t Nu_iq    = pbm.ulim.A.rows();
  const uint32_t Nx_iq    = pbm.glim.A.rows();
  const uint32_t Nxlin_iq = lin_con ? Nx : 0;

  const double dt = pbm.T / static_cast<double>(K);

  ////////////////////
  /// FILL A, l, u ///
  ////////////////////

  assert(static_cast<uint32_t>(qp.A.rows()) == K * (Nx + Nu_iq + Nx_iq + Nxlin_iq));
  assert(static_cast<uint32_t>(qp.A.cols()) == K * (Nx + Nu));
  assert(static_cast<uint32_t>(qp.l.rows()) == K * (Nx + Nu_iq + Nx_iq + Nxlin_iq));
  assert(static_cast<uint32_t>(qp.u.rows()) == K * (Nx + Nu_iq + Nx_iq + Nxlin_iq));

  int Arow = 0;

  for (auto k = 0u; k != K; ++k) {
    using AT = Eigen::Matrix<double, Nx, Nx>;
    using BT = Eigen::Matrix<double, Nx, Nu>;
    using ET = Eigen::Matrix<double, Nx, 1>;

    const double t = k * dt;

    // LINEARIZATION

    const auto [xl, dxl] = lin.g(t);
    const auto ul        = lin.u(t);

    // clang-format off
    const auto [flin, df_xu] = diff::dr<1, DT>(
      [&f, &t]<typename T>(const CastT<T, G> & vx, const CastT<T, U> & vu) -> Tangent<CastT<T, G>> {
        return f(t, vx, vu);
      },
      wrt(xl, ul)
    );
    // clang-format on

    // cltv system \dot x = At x(t) + Bt u(t) + Et
    const AT At = -0.5 * ad<G>(flin) - 0.5 * ad<G>(dxl) + df_xu.template leftCols<Nx>();
    const BT Bt = df_xu.template rightCols<Nu>();
    const ET Et = flin - dxl;

    // TIME DISCRETIZATION

    const AT At2     = At * At;
    const double dt2 = dt * dt;
    const double dt3 = dt2 * dt;

    // dltv system x^+ = Ak x + Bk u + Ek by truncated taylor expansion of the matrix exponential
    const AT Ak = AT::Identity() + At * dt + At2 * dt2 / 2. + At2 * At * dt3 / 6.;
    const BT Bk = Bt * dt + At * Bt * dt2 / 2. + At2 * Bt * dt3 / 6.;
    const ET Ek = Et * dt + At * Et * dt2 / 2. + At2 * Et * dt3 / 6.;

    // DYNAMICS CONSTRANTS

    if (k == 0) {
      // x(1) - B u(0) = A x0 + E

      // identity matrix on x(1)
      for (auto i = 0u; i != Nx; ++i) {
        const Eigen::Index Ari = Nx * k + i;
        assert(qp.A.outerIndexPtr()[Ari + 1] - qp.A.outerIndexPtr()[Ari] == Nu + 1);

        // B matrix on u0
        for (auto j = 0u; j != Nu; ++j) { qp.A.coeffRef(Ari, Nu * k + j) = -Bk(i, j); }

        // Identity matrix on x1
        qp.A.coeffRef(Ari, NU + Nx * k + i) = 1;
      }

      // B matrix on u(0)
      for (auto i = 0u; i != Nx; ++i) {}

      qp.u.template segment<Nx>(Nx * k) = Ak * rminus(pbm.x0, xl) + Ek;
      qp.l.template segment<Nx>(Nx * k) = qp.u.template segment<Nx>(Nx * k);
    } else {
      // x(k+1) - A x(k) - B u(k) = E
      for (auto i = 0u; i != Nx; ++i) {
        const Eigen::Index Ari = Nx * k + i;
        assert(qp.A.outerIndexPtr()[Ari + 1] - qp.A.outerIndexPtr()[Ari] == Nu + Nx + 1);

        // B matrix on u(k)
        for (auto j = 0u; j != Nu; ++j) { qp.A.coeffRef(Ari, Nu * k + j) = -Bk(i, j); }

        // A matrix on x(k)
        for (auto j = 0u; j != Nx; ++j) { qp.A.coeffRef(Ari, NU + Nx * (k - 1) + j) = -Ak(i, j); }

        // identity matrix on x(k+1)
        qp.A.coeffRef(Ari, NU + Nx * k + i) = 1;
      }

      qp.u.template segment<Nx>(Nx * k) = Ek;
      qp.l.template segment<Nx>(Nx * k) = Ek;
    }
  }
  Arow += K * Nx;

  // INPUT CONSTRAINTS

  if (Nu_iq > 0) {
    for (auto k = 0u; k < K; ++k) {
      for (auto i = 0u; i != Nu_iq; ++i) {
        const Eigen::Index Ari = Arow + k * Nu_iq + i;
        assert(qp.A.outerIndexPtr()[Ari + 1] - qp.A.outerIndexPtr()[Ari] == Nu);

        for (auto j = 0u; j != Nu; ++j) { qp.A.coeffRef(Ari, k * Nu + j) = pbm.ulim.A(i, j); }
      }
      // clang-format off
      qp.l.segment(Arow + k * Nu_iq, Nu_iq) = pbm.ulim.l - pbm.ulim.A * rminus(lin.u(k * dt), pbm.ulim.c);
      qp.u.segment(Arow + k * Nu_iq, Nu_iq) = pbm.ulim.u - pbm.ulim.A * rminus(lin.u(k * dt), pbm.ulim.c);
      // clang-format on
    }
  }
  Arow += K * Nu_iq;

  // STATE CONSTRAINTS

  if (Nx_iq > 0) {
    for (auto k = 1u; k != K + 1; ++k) {
      for (auto i = 0u; i != Nx_iq; ++i) {
        const Eigen::Index Ari = Arow + (k - 1) * Nx_iq + i;
        assert(qp.A.outerIndexPtr()[Ari + 1] - qp.A.outerIndexPtr()[Ari] == Nx);

        for (auto j = 0u; j != Nx; ++j) {
          qp.A.coeffRef(Ari, NU + (k - 1) * Nx + j) = pbm.glim.A(i, j);
        }
      }
      // clang-format off
      qp.l.segment(Arow + (k - 1) * Nx_iq, Nx_iq) = pbm.glim.l - pbm.glim.A * rminus(lin.g(k * dt).first, pbm.glim.c);
      qp.u.segment(Arow + (k - 1) * Nx_iq, Nx_iq) = pbm.glim.u - pbm.glim.A * rminus(lin.g(k * dt).first, pbm.glim.c);
      // clang-format on
    }
  }
  Arow += K * Nx_iq;

  // STATE LINEARIZATION BOUNDS

  if (Nxlin_iq > 0) {
    for (auto k = 1u; k < K + 1; ++k) {
      for (auto i = 0u; i != Nxlin_iq; ++i) {
        const Eigen::Index Ari = Arow + (k - 1) * Nxlin_iq + i;
        assert(qp.A.outerIndexPtr()[Ari + 1] - qp.A.outerIndexPtr()[Ari] == 1);

        qp.A.coeffRef(Ari, NU + (k - 1) * Nx + i) = 1.;
      }
      qp.l.segment(Arow + (k - 1) * Nxlin_iq, Nxlin_iq) = -lin.g_domain;
      qp.u.segment(Arow + (k - 1) * Nxlin_iq, Nxlin_iq) = lin.g_domain;
    }
  }
  Arow += K * Nxlin_iq;

  ////////////////
  /// FILL P,q ///
  ////////////////

  assert(static_cast<uint32_t>(qp.P.cols()) == K * (Nx + Nu));
  assert(static_cast<uint32_t>(qp.P.cols()) == K * (Nx + Nu));
  assert(static_cast<uint32_t>(qp.q.rows()) == K * (Nx + Nu));

  // INPUT COSTS

  for (auto k = 0u; k < K; ++k) {
    for (auto i = 0u; i != Nu; ++i) {
      const Eigen::Index Pci = k * Nu + i;
      assert(
        qp.P.outerIndexPtr()[Pci + 1] - qp.P.outerIndexPtr()[Pci]
        == static_cast<decltype(qp.P)::StorageIndex>(i) + 1);

      for (auto j = 0u; j != i + 1; ++j) {
        qp.P.coeffRef(k * Nu + j, Pci) = pbm.weights.R(j, i) * dt;
      }
    }
    qp.q.segment(k * Nu, Nu) = pbm.weights.R * rminus(lin.u(k * dt), pbm.udes(k * dt)) * dt;
  }

  // STATE COSTS

  // intermediate states x(1) ... x(K-1)
  for (auto k = 1u; k < K; ++k) {
    for (auto i = 0u; i != Nx; ++i) {
      const Eigen::Index Pci = NU + (k - 1) * Nx + i;
      assert(
        qp.P.outerIndexPtr()[Pci + 1] - qp.P.outerIndexPtr()[Pci]
        == static_cast<decltype(qp.P)::StorageIndex>(i) + 1);

      for (auto j = 0u; j != i + 1; ++j) {
        qp.P.coeffRef(NU + (k - 1) * Nx + j, Pci) = pbm.weights.Q(j, i) * dt;
      }
    }
    qp.q.segment(NU + (k - 1) * Nx, Nx) =
      pbm.weights.Q * rminus(lin.g(k * dt).first, pbm.gdes(k * dt)) * dt;
  }

  // last state x(K) ~ x(T)
  for (auto i = 0u; i != Nx; ++i) {
    const Eigen::Index Pci = NU + (K - 1) * Nx + i;
    assert(
      qp.P.outerIndexPtr()[Pci + 1] - qp.P.outerIndexPtr()[Pci]
      == static_cast<decltype(qp.P)::StorageIndex>(i) + 1);

    for (auto j = 0u; j != i + 1; ++j) {
      qp.P.coeffRef(NU + (K - 1) * Nx + j, Pci) = pbm.weights.QT(j, i);
    }
  }
  qp.q.segment(NU + (K - 1) * Nx, Nx) =
    pbm.weights.QT * rminus(lin.g(pbm.T).first, pbm.gdes(pbm.T));

  qp.A.makeCompressed();
  qp.P.makeCompressed();
}

/**
 * @brief Convert OptimalControlProblem on \f$ (\mathbb{G}, \mathbb{U}) \f$ into a tangent space
 * QuadraticProgramSparse on \f$ (\mathbb{R}^{\dim \mathfrak{g}}, \mathbb{R}^{\dim \mathfrak{u}})
 * \f$.
 *
 * The OptimalControlProblem is encoded into a QuadraticProgram via linearization around
 * \f$(g_{lin}(t), u_{lin}(t))\f$ followed by time discretization. The variables of the QP are \f[
 * \begin{bmatrix} \mu_0 & \mu_1 & \ldots & \mu_{K - 1} & x_1 & x_2 & \ldots & x_K
 * \end{bmatrix}, \f] where the discrete time index \f$k\f$ corresponds to time \f$t_k = k
 * \frac{T}{K} \f$ for \f$ k = 0, 1, \ldots, K \f$.
 *
 * performance.
 * @tparam G problem state group type \f$ \mathbb{G} \f$
 * @tparam U problem input group type \f$ \mathbb{U} \f$
 * @tparam Dyn dynamics functor type
 * @tparam DT differentiation method to utilize
 *
 * @param pbm optimal control problem
 * @param K number of discretization points. More points create a larger QP, but the distance \f$
 * T / K \f$ between points should be smaller than the smallest system time constant for adequate
 * @param f dynamics \f$ f : \mathbb{R} \times \mathbb{G} \times \mathbb{U} \rightarrow
 * \mathbb{R}^{\dim \mathfrak g}\f$ s.t. \f$ \mathrm{d}^r g_t = f(t, g, u) \f$
 * @param lin linearization point
 *
 * @return QuadraticProgramSparse modeling the input optimal control problem.
 *
 * @note Given a solution \f$(x^*, \mu^*)\f$ to the QuadraticProgramSparse, the corresponding
 * solution to the OptimalControlProblem is \f$ u^*(t) = u_{lin}(t) \oplus \mu^*(t) \f$ and the
 * optimal trajectory is \f$ g^*(t) = g_{lin}(t) \oplus x^*(t) \f$.
 *
 * @note Constraints are added as \f$ A x \leq b - A (g_{lin} - c) \f$ and similarly for the
 * input. Beware of using constraints on non-Euclidean spaces.
 *
 * @note \p f must be differentiable w.r.t. \f$ g \f$ and \f$ u \f$  with the default \p
 * smooth::diff method (check \p smooth::diff::DefaultType). If using an automatic differentiation
 * method this means that it must be templated on the scalar type.
 */
template<LieGroup G, Manifold U, typename Dyn, diff::Type DT = diff::Type::Default>
QuadraticProgramSparse<double> ocp_to_qp(
  const OptimalControlProblem<G, U> & pbm,
  std::size_t K,
  const Dyn & f,
  const LinearizationInfo<G, U> & lin)
{
  bool lin_con = lin.g_domain.minCoeff() < std::numeric_limits<double>::infinity();

  QuadraticProgramSparse<double> qp;

  ocp_to_qp_allocate<G, U>(pbm, K, qp, lin_con);
  ocp_to_qp_fill<G, U>(pbm, K, f, lin, qp);

  return qp;
}

}  // namespace smooth::feedback

#endif  // SMOOTH__FEEDBACK__MPC_FUNC_HPP_


#include "FloatingObject.hpp"
#include <algorithm>
#include <vector>
#include <mutex>
#include <shared_mutex>
#include <Eigen\Geometry>

#define NOMINMAX
#include <Windows.h>

namespace {
	const float pi = 3.14159265359;
	constexpr float AIR_DENSITY = 1.293e-9; // [kg/mm3]
}

namespace dynaman {
	FloatingObject::FloatingObject(
		const Eigen::Vector3f& positionTarget,
		const Eigen::Vector3f& lowerbound,
		const Eigen::Vector3f& upperbound,
		float radius,
		float weight
	):
		m_position(positionTarget),
		m_velocity(Eigen::Vector3f::Zero()),
		m_integral(Eigen::Vector3f::Zero()),
		m_lowerbound(lowerbound),
		m_upperbound(upperbound),
		m_weight(weight),
		m_radius(radius),
		m_is_tracked(false),
		m_lastDeterminationTime(0),
		velocityBufferSize(3),
		positionBuffer(velocityBufferSize, positionTarget),
		velocityBuffer(velocityBufferSize, Eigen::Vector3f::Zero()),
		dTBuffer(velocityBufferSize, 1),
		m_pTrajectory(std::make_shared<TrajectoryConstantState>(positionTarget)
		)
	{}

	FloatingObjectPtr FloatingObject::Create(
		const Eigen::Vector3f& posTgt,
		const Eigen::Vector3f& lowerbound,
		const Eigen::Vector3f& upperbound,
		float radius,
		float weight
	){
		return std::make_shared<FloatingObject>(posTgt, lowerbound, upperbound, radius, weight);
	}

	float FloatingObject::sphereMass()
	{
		return AIR_DENSITY * 4.0f * pi * m_radius * m_radius * m_radius / 3.0f;
	}

	float FloatingObject::radius() {
		return m_radius;
	}

	float FloatingObject::weight()
	{
		return m_weight;
	}

	float FloatingObject::totalMass()
	{
		return sphereMass() + weight();
	}

	DWORD FloatingObject::lastDeterminationTime() {
		std::lock_guard<std::mutex> lock(mtxState);
		return m_lastDeterminationTime;
	}

	Eigen::Vector3f FloatingObject::position()
	{
		std::lock_guard<std::mutex> lock(mtxState);
		return m_position;
	}

	Eigen::Vector3f FloatingObject::velocity()
	{
		std::lock_guard<std::mutex> lock(mtxState);
		return averageVelocity(this->velocityBuffer, this->dTBuffer);
	}

	Eigen::Vector3f FloatingObject::integral()
	{
		std::lock_guard<std::mutex> lock(mtxState);
		return m_integral;
	}

	Eigen::Vector3f FloatingObject::positionTarget(DWORD systime_ms)
	{
		std::shared_lock<std::shared_mutex> lock(mtxTrajectory);
		return m_pTrajectory->pos(systime_ms);
	}

	Eigen::Vector3f FloatingObject::velocityTarget(DWORD systime_ms)
	{
		std::shared_lock<std::shared_mutex> lock(mtxTrajectory);
		return m_pTrajectory->vel(systime_ms);
	}

	Eigen::Vector3f FloatingObject::accelTarget(DWORD systime_ms)
	{
		std::shared_lock<std::shared_mutex> lock(mtxTrajectory);
		return m_pTrajectory->accel(systime_ms);
	}

	void FloatingObject::getStates(
		Eigen::Vector3f& pos,
		Eigen::Vector3f& vel,
		Eigen::Vector3f& integ
	) {
		std::lock_guard<std::mutex> lock(mtxState);
		pos = m_position;
		vel = averageVelocity(this->velocityBuffer, this->dTBuffer);
		integ = m_integral;
	}

	void FloatingObject::updateStates(DWORD determinationTime, const Eigen::Vector3f& positionNew)
	{
		std::lock_guard<std::mutex> lock(mtxState);
		float dt = (float)(determinationTime - m_lastDeterminationTime) / 1000.f; // [sec]
		m_velocity = (positionNew - m_position) / dt;
		dTBuffer.push_back(dt);
		dTBuffer.pop_front();
		positionBuffer.push_back(positionNew);
		positionBuffer.pop_front();
		velocityBuffer.push_back(m_velocity);
		velocityBuffer.pop_front();
		if (isTracked())
		{
			m_integral += (0.5f * (positionNew + m_position) - positionTarget()) * dt;
		}
		m_position = positionNew;
		m_lastDeterminationTime = determinationTime;
	}

	void FloatingObject::resetIntegral() {
		std::lock_guard<std::mutex> lock(mtxState);
		m_integral.setZero();
	}

	void FloatingObject::getStatesTarget(
		Eigen::Vector3f& posTgt,
		Eigen::Vector3f& velTgt,
		Eigen::Vector3f& accelTgt,
		DWORD time
	) {
		std::lock_guard<std::shared_mutex> lock(mtxTrajectory);
		posTgt = m_pTrajectory->pos(time);
		velTgt = m_pTrajectory->vel(time);
		accelTgt = m_pTrajectory->accel(time);
	}

	void FloatingObject::updateStatesTarget(
		const Eigen::Vector3f& _positionTarget,
		const Eigen::Vector3f& _velocityTarget,
		const Eigen::Vector3f& _accelTarget)
	{
		auto constTrajPtr = TrajectoryConstantState::Create(_positionTarget, _velocityTarget, _accelTarget);
		setTrajectory(constTrajPtr);
	}

	bool FloatingObject::isTracked() {
		std::lock_guard<std::mutex> lock(mtxTrack);
		return m_is_tracked;
	}

	void FloatingObject::setTrackingStatus(bool is_tracked) {
		std::lock_guard<std::mutex> lock(mtxTrack);
		this->m_is_tracked = is_tracked;
	}

	Eigen::Vector3f FloatingObject::averageVelocity(
		std::deque<Eigen::Vector3f> velocityBuffer,
		std::deque<float> intervalBuffer
	) {
		auto itrVel = velocityBuffer.begin();
		auto itrDT = intervalBuffer.begin();
		Eigen::Vector3f distSum(0.f, 0.f, 0.f);
		float period = 0;
		while (itrVel != velocityBuffer.end())
		{
			distSum += (*itrDT) * (*itrVel);
			period += *itrDT;
			itrDT++;
			itrVel++;
		}
		return distSum /= period;
	}

	Eigen::Vector3f FloatingObject::averagePosition() {
		Eigen::Vector3f posAverage(0, 0, 0);
		std::lock_guard<std::mutex> lock(mtxState);
		for (auto itr = positionBuffer.begin(); itr != positionBuffer.end(); itr++) {
			posAverage += *itr;
		}
		return posAverage / positionBuffer.size();
	}

	bool FloatingObject::isInsideWorkspace()
	{
		auto pos = position();
		Eigen::Vector3f v0 = pos - lowerbound();
		Eigen::Vector3f v1 = pos - upperbound();
		return (v0.x() * v1.x() <= 0) && (v0.y() * v1.y() <= 0) && (v0.z() * v1.z() <= 0);
	}

	void FloatingObject::setTrajectory(std::shared_ptr<Trajectory> newTrajectoryPtr)
	{
		std::lock_guard<std::shared_mutex> lock(mtxTrajectory);
		m_pTrajectory = newTrajectoryPtr;
	}

	Eigen::Vector3f FloatingObject::lowerbound() {
		return m_lowerbound;
	}

	Eigen::Vector3f FloatingObject::upperbound() {
		return m_upperbound;
	}

	std::string FloatingObject::logCtrlHeader() const {
		return "sys_time,x,y,z,vx,vy,vz,ix,iy,iz,xTgt,yTgt,zTgt,vxTgt,vyTgt,vzTgt,axTgt,ayTgt,azTgt";
	}

	std::ofstream& operator<<(std::ostream& ofs, FloatingObjectPtr pObject) {
		DWORD time = pObject->lastDeterminationTime();
		Eigen::Vector3f pos, vel, integ, posTgt, velTgt, accelTgt;
		pObject->getStates(pos, vel, integ);
		pObject->getStatesTarget(posTgt, velTgt, accelTgt, time);
		ofs
			<< time << ","
			<< pos.x() << "," << pos.y() << "," << pos.z() << ","
			<< vel.x() << "," << vel.y() << "," << vel.z() << ","
			<< integ.x() << "," << integ.y() << "," << integ.z() << ","
			<< posTgt.x() << "," << posTgt.y() << "," << posTgt.z() << ","
			<< velTgt.x() << "," << velTgt.y() << "," << velTgt.z() << ","
			<< accelTgt.x() << "," << accelTgt.y() << "," << accelTgt.z();
	}
}


#ifndef _CHEBSPEC_CPP_
#define _CHEBSPEC_CPP_

#include <armadillo>

#include "solvspec.cpp"
#include "../include/chebspec.h"

using namespace std;
using namespace arma;

const double PI = 3.141592653589793;

int ChebSpec::N = -1;
int ChebSpec::L = 0;
vector<double> ChebSpec::x = vector<double>();
vector<vector<double> > ChebSpec::Tm = vector<vector<double> >();
vector<vector<double> > ChebSpec::TmInt = vector<vector<double> >();
vector<vector<double> > ChebSpec::D = vector<vector<double> >();
vector<vector<double> > ChebSpec::D2 = vector<vector<double> >();

ChebSpec::ChebSpec(){}

void ChebSpec::setPotential(const vector<double> &XX, const vector<double> &VV)
{
  SolvSpec::setPotential(XX, VV);
  a = XX[0];
  b = XX[XX.size() - 1];
  scal = pow((b - a) / 2, 2.0);
}

void ChebSpec::findSpectrum(int nEigen)
{
  if(PotVals.size() < 2)
  {
    cout << "Please use the setPotential() function before using this one." << endl;
    return;
  }
  // compute V
  V.resize(L);
  for(int i = 0; i < L; i++)
  {
    V[i] = scal * potFunc.interp(0.5 * ((b - a) * x[i] + b + a));
  }
  // compute Ehat
  Ehat.resize(L);
  for(int i = 0; i < L; i++)
  {
    Ehat[i].resize(L);
    for(int j = 0; j < L; j++)
    {
      Ehat[i][j] = D2[i][j] * V[j];
    }
  }
  // compute UE and US
  UE.resize(L);
  US.resize(L);
  for(int i = 0; i < L; i++)
  {
    UE[i].resize(L);
    US[i].resize(L);
    for(int j = 0; j < L; j++)
    {
      UE[i][j] = 0.5 * (x[i] + 1) * D2[N][j];
      US[i][j] = 0.5 * (x[i] + 1) * Ehat[N][j];
    }
  }
  // compute the A and B matrices, remove the first/last row/column
  int Nr = N - 1;
  mat A(Nr, Nr);
  mat B(Nr, Nr);

  //lapack routines use the column-major order!
  for(int i = 0; i < Nr; i++)
  {
    for(int j = 0; j < Nr; j++)
    {
      // define A
      A(i, j) = -US[i + 1][j + 1] + Ehat[i + 1][j + 1];
      if(i == j)   // add the identity matrix
      {
        A(i, j) = A[i + j * Nr] - 1;
      }
      // now B
      B(i, j) = D2[i + 1][j + 1] - UE[i + 1][j + 1];
    }
  }

  mat B1A = inv(B) * A;
  cx_vec cxeigval;
  cx_mat cxeigvec;

  // diagonalize
  eig_gen(cxeigval, cxeigvec, B1A);
  mat eigvec = real(cxeigvec);
  vec eigval = real(cxeigval);
  // we need to sort this
  for(int i = 0; i < Nr; i++)
  {
    for(int j = 0; j < Nr; j++)
    {
      if(eigval(i) < eigval(j))
      {
        double temp = eigval(i);
        eigval(i) = eigval(j);
        eigval(j) = temp;
        eigvec.swap_cols(i, j);
      }
    }
  }
  // now we just need to put everything in our internal format
  // finally safelly add all the modes found to the spectrum (already in a nice way)
  spectrum.clear();
  spectrum.potential.push_back(X); spectrum.potential.push_back(PotVals);
  for (int i = 0; i < nEigen; i++)
  {
    // build the wavefunction
    vector<Point> wf;
    for(int j = 0; j < Nr; j++)
    {
      wf.push_back(Point(0.5 * ((b - a) * x[j] + b + a), eigvec(j,i)));
    }
    // normalization loop
    double c = 0;
    for(int j = 0; j < Nr; j++)
    {
      c += D[N][j + 1] * wf[j].y * wf[j].y * 0.5 * (b - a);
    }
    // set all the wavefunctions start growing positive from the left
    double s = 1;
    // get the sign of the derivative, this is important since it may be the case
    // that the routine return the same eigenvector with a different sign, in that case
    // the overall coefficient we would like to fit would have change the sign
    for(int j = 4; j < Nr; j++)
    {
      // filter the data and compute the derivative
      double der = (25/12)*int(1e3 * wf[j].y)-4*int(1e3 * wf[j-1].y)+3*int(1e3 * wf[j-2].y)-(4/3) * int(1e3 * wf[j-3].y)+(1/4)* int(1e3 * wf[j-4].y);
      der /= 1e3;
      if(abs(der) > 1e-2)
      {
        s = der / abs(der);
        break;
      }
    }
    for(int j = 0; j < Nr; j++)
    {
      wf[j].y = s * wf[j].y / sqrt(c);
    }
    Mode m(eigval(i) / scal, wf);
    spectrum.addMode(m);
  }
}

void ChebSpec::showMatrix(double* A, int Nr)
{
  cout << endl;
  for(int i = 0; i < Nr; i++)
  {
    for(int j = 0; j < Nr; j++)
    {
      cout << A[i + j * Nr] << " ";
    }
    cout << endl;
  }
}

ChebSpec::~ChebSpec(){}

// [[Rcpp::export]]
void chebSetN(int n)
{
  if(ChebSpec::N == n)
    return;
  ChebSpec::N = n;
  ChebSpec::L = n + 1;
  cout << "computing chebyshev matrices, N = " << ChebSpec::N << endl;
  // initialize x
  ChebSpec::x.clear();
  ChebSpec::x.resize(ChebSpec::L);
  for(int i = 0; i < ChebSpec::L; i++)
  {
    ChebSpec::x[i] = -cos(PI * i / ChebSpec::N);
  }
  // Tm, the matrix of the Chebyshev polynomial values in the grid. Checked.
  ChebSpec::Tm.clear();
  ChebSpec::Tm.resize(ChebSpec::L);
  for(int i = 0; i < ChebSpec::L; i++)
  {
    ChebSpec::Tm[i].resize(ChebSpec::L);
    for(int j = 0; j < ChebSpec::L; j++)
    {
      if(i == 0)      // first
      {
        ChebSpec::Tm[i][j] = 1;
      }
      else if(i == 1) // second
      {
        ChebSpec::Tm[i][j] = ChebSpec::x[j];
      }
      else
      {
        ChebSpec::Tm[i][j] = 2 * ChebSpec::x[j] * ChebSpec::Tm[i - 1][j] - ChebSpec::Tm[i - 2][j];
      }
    }
  }

  // TmIntInt, the matrix of the integrated from -1 Chebyshev polynomial values in the grid. Checked.
  ChebSpec::TmInt.clear();
  ChebSpec::TmInt.resize(ChebSpec::L);
  for(int i = 0; i < ChebSpec::L; i++)
  {
    ChebSpec::TmInt[i].resize(ChebSpec::L);
    for(int j = 0; j < ChebSpec::L; j++)
    {
      if(i == 0)      // first
      {
        ChebSpec::TmInt[i][j] = ChebSpec::x[j] + 1;
      }
      else if(i == 1) // second
      {
        ChebSpec::TmInt[i][j] =  0.5 * (ChebSpec::x[j] * ChebSpec::x[j] - 1);
      }
      else if(i < ChebSpec::L - 1)
      {
        ChebSpec::TmInt[i][j] = ChebSpec::Tm[i + 1][j] / (2. * (i + 1)) - ChebSpec::Tm[i - 1][j] / (2. * (i - 1)) + pow(-1, i + 1.) / (pow(i, 2.) - 1);
      }
      else   // case i = L - 1, right extreme
      {
        ChebSpec::TmInt[i][j] = (2 * ChebSpec::x[j] * ChebSpec::Tm[i][j] - ChebSpec::Tm[i - 1][j]) / (2.0 * (i + 1)) - ChebSpec::Tm[i - 1][j] / (2.0 * (i - 1)) + pow(-1, i + 1) / (pow(i, 2) - 1);
      }
    }
  }

  // D matrix: the integral from -1 to x. Checked.
  ChebSpec::D.clear();
  ChebSpec::D.resize(ChebSpec::L);
  for(int i = 0; i < ChebSpec::L; i++)
  {
    ChebSpec::D[i].resize(ChebSpec::L);
    for(int j = 0; j < ChebSpec::L; j++)
    {
      double s = 0;
      for(int k = 0; k < ChebSpec::L; k++)
      {
        double f = 1;
        if(k == 0 || k == ChebSpec::L - 1) // prime '' sum
        {
          f = 0.5;
        }
        s += f * ChebSpec::Tm[k][j] * ChebSpec::TmInt[k][i];
      }

      ChebSpec::D[i][j] = (2. / ChebSpec::N) * s;
      // prime '' sum
      if(j == 0 || j == ChebSpec::L - 1)
        ChebSpec::D[i][j] = 0.5 * ChebSpec::D[i][j];
    }
  }
  // finally the D2 = D*D matrix
  // check, the sum of the last row has to be 2 and 0 the sum of the first
  ChebSpec::D2.clear();
  ChebSpec::D2.resize(ChebSpec::L);
  for(int i = 0; i < ChebSpec::L; i++)
  {
    ChebSpec::D2[i].resize(ChebSpec::L);
    for(int j = 0; j < ChebSpec::L; j++)
    {
      double s = 0;
      for(int k = 0; k < ChebSpec::L; k++)
      {
        s += ChebSpec::D[i][k] * ChebSpec::D[k][j];
      }
      ChebSpec::D2[i][j] = s;
    }
  }

  // check, the sum of the last row has to be 2
  double s = 0;
  for(int k = 0; k < ChebSpec::L; k++)
  {
    s += ChebSpec::D2[0][k];
  }
}

;
#endif


/* Copyright (c) 2016 - 2019, the adamantine authors.
 *
 * This file is subject to the Modified BSD License and may not be distributed
 * without copyright and license information. Please refer to the file LICENSE
 * for the text and further information on this license.
 */

#ifndef MATERIAL_PROPERTY_HH
#define MATERIAL_PROPERTY_HH

#include <types.hh>
#include <utils.hh>

#include <deal.II/base/function_parser.h>
#include <deal.II/distributed/tria.h>
#include <deal.II/dofs/dof_accessor.h>
#include <deal.II/dofs/dof_handler.h>
#include <deal.II/fe/fe_dgq.h>
#include <deal.II/grid/filtered_iterator.h>
#include <deal.II/lac/la_parallel_vector.h>
#include <deal.II/lac/la_vector.h>

#include <boost/property_tree/ptree.hpp>

#include <array>
#include <limits>
#include <unordered_map>

namespace adamantine
{
/**
 * This class stores the material properties for all the materials
 */
template <int dim>
class MaterialProperty
{
public:
  /**
   * Constructor.
   * \param[in] database requires the following entries:
   *   - n_materials: unsigned int in \f$(0,\infty)\f$
   *   - material_X: property tree associated with material_X
   *   where X is a number
   *   - material_X.Y: property_tree where Y is either liquid, powder, or
   *   solid [optional]
   *   - material_X.Y.Z: string where Z is either density, specific_heat,
   *   or thermal_conductivity, describe the behavior of the property as a
   *   function of the temperatur (e.g. "2.*T") [optional]
   *   - material.X.A: A is either solidus, liquidus, or latent_heat
   *   [optional]
   */
  MaterialProperty(
      MPI_Comm const &communicator,
      dealii::parallel::distributed::Triangulation<dim> const &tria,
      boost::property_tree::ptree const &database);

  /**
   * Return the value of the given property, for a given cell and a given field
   * state.
   */
  template <typename NumberType>
  double
  get(typename dealii::Triangulation<dim>::active_cell_iterator const &cell,
      Property prop,
      dealii::LA::distributed::Vector<NumberType> const &field_state) const;

  /**
   * Return the average temperature on every cell given the enthalpy.
   */
  template <typename NumberType>
  dealii::LA::distributed::Vector<NumberType> enthalpy_to_temperature(
      dealii::DoFHandler<dim> const &enthalpy_dof_handler,
      dealii::LA::distributed::Vector<NumberType> const &enthalpy);

  /**
   * Reinitialize the DoFHandler associated with MaterialProperty and resize the
   * state vectors.
   */
  void reinit_dofs();

  /**
   * Update the material state, i.e, the ratio of liquid, powder, and solid.
   */
  template <typename NumberType>
  void
  update_state(dealii::DoFHandler<dim> const &enthalpy_dof_handler,
               dealii::LA::distributed::Vector<NumberType, dealii::MemorySpace::CUDA> const &enthalpy);

  /**
   * Get the array of material state vectors. The order of the different state
   * vectos is given by the MaterialState enum. Each entry in the vector
   * correspond to a cell in the mesh and has a value between 0 and 1. The sum
   * of the states for a given cell is equal to 1.
   */
  std::array<dealii::LA::distributed::Vector<double>,
             static_cast<unsigned int>(MaterialState::SIZE)> &
  get_state();

  double get_state_ratio(
      typename dealii::Triangulation<dim>::active_cell_iterator const &cell,
      MaterialState material_state) const;

  /**
   * Return \f$ -\frac{H_{liquidus}}{\rho C_P} + T_{liquidus} \f$
   */
  double get_liquid_beta(
      typename dealii::Triangulation<dim>::active_cell_iterator const &cell)
      const;

  /**
   * Return \f$ \frac{T_{liquidus}-T_{solidus}}{\mathcal{L}} \f$
   */
  double get_mushy_alpha(
      typename dealii::Triangulation<dim>::active_cell_iterator const &cell)
      const;

  /**
   * Return \f$ -H_{solidus} \frac{T_{liquidus}-T_{solidus}}{\mathcal{L}} +
   * T_{solidus} \f$
   */
  double get_mushy_beta(
      typename dealii::Triangulation<dim>::active_cell_iterator const &cell)
      const;

  /**
   * Return the underlying the DoFHandler.
   */
  dealii::DoFHandler<dim> const &get_dof_handler() const;

private:
  /**
   * Maximum different number of states a given material can be.
   */
  static unsigned int constexpr _n_material_states =
      static_cast<unsigned int>(MaterialState::SIZE);

  /**
   * Number of properties defined.
   */
  static unsigned int constexpr _n_properties =
      static_cast<unsigned int>(Property::SIZE);

  /**
   * Set the values in _state from the values of the user index of the
   * Triangulation.
   */
  void set_state();

  /**
   * Fill the _properties map.
   */
  void fill_properties(boost::property_tree::ptree const &database);

  /**
   * If the density and the specific heat do not depend on the temperature, the
   * relationship between the temperature and the enthalpy can be written by
   * piecewise function \f$ T = \alpha H + \beta \f$. This function computes the
   * constants \f$ \alpha \f$ and \f$ \beta \f$.
   */
  void compute_constants();

  /**
   * Return the index of the dof associated to the cell.
   */
  double get_dof_index(
      typename dealii::Triangulation<dim>::active_cell_iterator const &cell)
      const;

  /**
   * Compute the average of the enthalpy on every cell.
   */
  template <typename NumberType>
  dealii::LA::distributed::Vector<NumberType, dealii::MemorySpace::CUDA> compute_average_enthalpy(
      dealii::DoFHandler<dim> const &enthalpy_dof_handler,
      dealii::LA::distributed::Vector<NumberType, dealii::MemorySpace::CUDA> const &enthalpy) const;

  /**
   * MPI communicator.
   */
  MPI_Comm _communicator;
  /**
   * Map of \f$ -\frac{H_{liquidus}}{\rho C_P} + T_{liquidus} \f$ for each
   * material.
   */
  std::unordered_map<dealii::types::material_id, double> _liquid_beta;
  /**
   * Map of \f$ \frac{T_{liquidus}-T_{solidus}}{\mathcal{L}} \f$ for each
   * material.
   */
  std::unordered_map<dealii::types::material_id, double> _mushy_alpha;
  /**
   * Map of \f$ -H_{solidus} \frac{T_{liquidus}-T_{solidus}}{\mathcal{L}} +
   * T_{solidus} \f$ for each material.
   */
  std::unordered_map<dealii::types::material_id, double> _mushy_beta;
  /**
   * Map that stores functions describing the properties of the material.
   */
  std::unordered_map<
      dealii::types::material_id,
      std::array<
          std::array<std::unique_ptr<dealii::FunctionParser<1>>, _n_properties>,
          _n_material_states>>
      _properties;
  /**
   * Array of vector describing the ratio of each state in each cell. Each
   * vector corresponds to a state defined in the MaterialState enum.
   */
  std::array<dealii::LA::distributed::Vector<double>, _n_material_states>
      _state;
  /**
   * Discontinuous piecewise constant finite element.
   */
  dealii::FE_DGQ<dim> _fe;
  /**
   * DoFHandler associated to the _state array.
   */
  dealii::DoFHandler<dim> _mp_dof_handler;
};

template <int dim>
inline std::array<dealii::LA::distributed::Vector<double>,
                  static_cast<unsigned int>(MaterialState::SIZE)> &
MaterialProperty<dim>::get_state()
{
  return _state;
}

template <int dim>
inline double MaterialProperty<dim>::get_state_ratio(
    typename dealii::Triangulation<dim>::active_cell_iterator const &cell,
    MaterialState material_state) const
{
  double const mp_dof_index = get_dof_index(cell);
  unsigned int const mat_state = static_cast<unsigned int>(material_state);

  return _state[mat_state][mp_dof_index];
}

template <int dim>
inline double MaterialProperty<dim>::get_liquid_beta(
    typename dealii::Triangulation<dim>::active_cell_iterator const &cell) const
{
  dealii::types::material_id material_id = cell->material_id();
  auto const liquid_beta = _liquid_beta.find(material_id);
  ASSERT(liquid_beta != _liquid_beta.end(), "Material not found.");

  return liquid_beta->second;
}

template <int dim>
inline double MaterialProperty<dim>::get_mushy_alpha(
    typename dealii::Triangulation<dim>::active_cell_iterator const &cell) const
{
  dealii::types::material_id material_id = cell->material_id();
  auto const mushy_alpha = _mushy_alpha.find(material_id);
  ASSERT(mushy_alpha != _mushy_alpha.end(), "Material not found.");

  return mushy_alpha->second;
}

template <int dim>
inline double MaterialProperty<dim>::get_mushy_beta(
    typename dealii::Triangulation<dim>::active_cell_iterator const &cell) const
{
  dealii::types::material_id material_id = cell->material_id();
  auto const mushy_beta = _mushy_beta.find(material_id);
  ASSERT(mushy_beta != _mushy_beta.end(), "Material not found.");

  return mushy_beta->second;
}

template <int dim>
inline double MaterialProperty<dim>::get_dof_index(
    typename dealii::Triangulation<dim>::active_cell_iterator const &cell) const
{
  // Get a DoFCellAccessor from a Triangulation::active_cell_iterator.
  dealii::DoFAccessor<dim, dealii::DoFHandler<dim>, false> dof_accessor(
      &_mp_dof_handler.get_triangulation(), cell->level(), cell->index(),
      &_mp_dof_handler);
  std::vector<dealii::types::global_dof_index> mp_dof(1.);
  dof_accessor.get_dof_indices(mp_dof);

  return mp_dof[0];
}

template <int dim>
inline dealii::DoFHandler<dim> const &
MaterialProperty<dim>::get_dof_handler() const
{
  return _mp_dof_handler;
}
} // namespace adamantine

#endif