// 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, John B. Mains
//
// 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__OCP_HPP_
#define SMOOTH__FEEDBACK__OCP_HPP_

/**
 * @file
 * @brief Optimal control problem definition and conversion to QP.
 */

#include <Eigen/Core>
#include <smooth/lie_group.hpp>

#include "collocation.hpp"
#include "nlp.hpp"
#include "traits.hpp"

namespace smooth::feedback {

/**
 * @brief Optimal control problem definition
 * @tparam _X state space
 * @tparam _U input space
 *
 * Problem is defined on the interval \f$ t \in [0, t_f] \f$.
 * \f[
 * \begin{cases}
 *  \min              & \theta(t_f, x_0, x_f, q)                                         \\
 *  \text{s.t.}       & x(0) = x_0                                                       \\
 *                    & x(t_f) = x_f                                                     \\
 *                    & \dot x(t) = f(t, x(t), u(t))                                     \\
 *                    & q = \int_{0}^{t_f} g(t, x(t), u(t)) \mathrm{d}t                  \\
 *                    & c_{rl} \leq c_r(t, x(t), u(t)) \leq c_{ru} \quad t \in [0, t_f]  \\
 *                    & c_{el} \leq c_e(t_f, x_0, x_f, q) \leq c_{eu}
 * \end{cases}
 * \f]
 *
 * The optimal control problem depends on arbitrary functions \f$ \theta, f, g, c_r, c_e \f$.
 * The type of those functions are template pararamters in this structure.
 *
 * @note To enable automatic differentiation \f$ \theta, f, g, c_r, c_e \f$ must be templated over the
 * scalar type.
 */
template<LieGroup _X, Manifold _U, typename Theta, typename F, typename G, typename CR, typename CE>
struct OCP
{
  using X = _X;
  using U = _U;

  /// @brief State dimension \f$ n_{x} \f$
  std::size_t nx;
  /// @brief Input dimension \f$ n_{u} \f$
  std::size_t nu;
  /// @brief Number of integrals \f$ n_{q} \f$
  std::size_t nq;
  /// @brief Number of running constraints \f$ n_{cr} \f$
  std::size_t ncr;
  /// @brief Number of end constraints \f$ n_{ce} \f$
  std::size_t nce;

  /// @brief Objective function \f$ \theta : R \times X \times X \times R^{n_q} \rightarrow R \f$
  Theta theta;

  /// @brief System dynamics \f$ f : R \times X \times U \rightarrow Tangent<X> \f$
  F f;
  /// @brief Integrals \f$ g : R \times X \times U \rightarrow R^{n_q} \f$
  G g;

  /// @brief Running constraint \f$ c_r : R \times X \times U \rightarrow R^{n_{cr}} \f$
  CR cr;
  /// @brief Running constraint lower bound \f$ c_{rl} \in R^{n_{cr}} \f$
  Eigen::VectorXd crl;
  /// @brief Running constraint upper bound \f$ c_{ru} \in R^{n_{cr}} \f$
  Eigen::VectorXd cru;

  /// @brief End constraint \f$ c_e : R \times X \times X \times R^{n_q} \rightarrow R^{n_{ce}} \f$
  CE ce;
  /// @brief End constraint lower bound \f$ c_{el} \in R^{n_{ce}} \f$
  Eigen::VectorXd cel;
  /// @brief End constraint upper bound \f$ c_{eu} \in R^{n_{ce}} \f$
  Eigen::VectorXd ceu;
};

/// @brief Concept that is true for OCP specializations
template<typename T>
concept OCPType = traits::is_specialization_of_v<T, OCP>;

/// @brief OCP defined on flat spaces
template<typename Theta, typename F, typename G, typename CR, typename CE>
using FlatOCP = OCP<Eigen::VectorXd, Eigen::VectorXd, Theta, F, G, CR, CE>;

/// @brief Concept that is true for FlatOCP specializations
template<typename T>
concept FlatOCPType = OCPType<T> &&(
  std::is_same_v<typename T::X, Eigen::VectorXd> && std::is_same_v<typename T::U, Eigen::VectorXd>);

/**
 * @brief Check if an OCP is properly defined.
 */
inline bool check_ocp(const OCPType auto & ocp)
{
  using X = typename std::decay_t<decltype(ocp)>::X;
  using U = typename std::decay_t<decltype(ocp)>::U;

  const double t = 0;
  const X x      = Default<X>(ocp.nx);
  const U u      = Default<U>(ocp.nu);

  if (!(static_cast<std::size_t>(dof(x)) == ocp.nx)) { return false; }
  if (!(static_cast<std::size_t>(dof(u)) == ocp.nu)) { return false; }

  const auto dx = ocp.f.template operator()<double>(t, x, u);
  if (!(static_cast<std::size_t>(dx.size()) == ocp.nx)) { return false; }

  const auto g = ocp.g.template operator()<double>(t, x, u);
  if (!(static_cast<std::size_t>(g.size()) == ocp.nq)) { return false; }

  [[maybe_unused]] const double obj = ocp.theta.template operator()<double>(t, x, x, g);

  const auto cr = ocp.cr.template operator()<double>(t, x, u);
  if (!(static_cast<std::size_t>(cr.size()) == ocp.ncr)) { return false; }
  if (!(static_cast<std::size_t>(ocp.crl.size()) == ocp.ncr)) { return false; }
  if (!(static_cast<std::size_t>(ocp.cru.size()) == ocp.ncr)) { return false; }

  const auto ce = ocp.ce.template operator()<double>(t, x, x, g);
  if (!(static_cast<std::size_t>(ce.size()) == ocp.nce)) { return false; }
  if (!(static_cast<std::size_t>(ocp.cel.size()) == ocp.nce)) { return false; }
  if (!(static_cast<std::size_t>(ocp.ceu.size()) == ocp.nce)) { return false; }

  return true;
}

/**
 * @brief Flatten a LieGroup OCP by defining it in the tangent space around a trajectory.
 *
 * @param ocp OCPType defined on a LieGroup
 * @param xl nominal state trajectory
 * @param ul nominal state trajectory
 *
 * @return FlatOCPType in variables (xe, ue) obtained via variables change x = xl ⊕ xe, u = ul ⊕ ue,
 */
inline auto flatten_ocp(const OCPType auto & ocp, auto && xl_fun, auto && ul_fun)
{
  using Eigen::VectorX;
  using X = typename std::decay_t<decltype(ocp)>::X;
  using U = typename std::decay_t<decltype(ocp)>::U;

  assert(Dof<X> == -1 || ocp.nx == Dof<X>);
  assert(Dof<U> == -1 || ocp.nu == Dof<U>);

  auto f_new = [f = ocp.f, xl_fun = xl_fun, ul_fun = ul_fun]<typename T>(
                 const T & t, const VectorX<T> & xe, const VectorX<T> & ue) -> VectorX<T> {
    using X_T = smooth::CastT<T, X>;
    using U_T = smooth::CastT<T, U>;

    // can not double-differentiate, so we neglect derivative of linearization w.r.t. t
    const double tdbl    = static_cast<double>(t);
    const auto [xl, dxl] = diff::dr(xl_fun, wrt(tdbl));
    const auto ul        = ul_fun(tdbl);

    const X_T x = rplus(xl.template cast<T>(), xe);
    const U_T u = rplus(ul.template cast<T>(), ue);

    return dr_expinv<X_T>(xe) * f.template operator()<T>(t, x, u)
         - dl_expinv<X_T>(xe) * dxl.template cast<T>();
  };

  auto g_new = [g = ocp.g, xl_fun = xl_fun, ul_fun = ul_fun]<typename T>(
                 const T & t, const VectorX<T> & xe, const VectorX<T> & ue) -> VectorX<T> {
    return g.template operator()<T>(t, rplus(xl_fun(t), xe), rplus(ul_fun(t), ue));
  };

  auto cr_new = [cr = ocp.cr, xl_fun = xl_fun, ul_fun = ul_fun]<typename T>(
                  const T & t, const VectorX<T> & xe, const VectorX<T> & ue) -> VectorX<T> {
    return cr.template operator()<T>(t, rplus(xl_fun(t), xe), rplus(ul_fun(t), ue));
  };

  auto theta_new =
    [theta = ocp.theta, xl_fun = xl_fun]<typename T>(
      const T & tf, const VectorX<T> & xe0, const VectorX<T> & xef, const VectorX<T> & q) -> T {
    return theta.template operator()<T>(tf, rplus(xl_fun(T(0.)), xe0), rplus(xl_fun(tf), xef), q);
  };

  auto ce_new = [ce = ocp.ce, xl_fun = xl_fun]<typename T>(
                  const T & tf,
                  const VectorX<T> & xe0,
                  const VectorX<T> & xef,
                  const VectorX<T> & q) -> VectorX<T> {
    return ce.template operator()<T>(tf, rplus(xl_fun(T(0.)), xe0), rplus(xl_fun(tf), xef), q);
  };

  return FlatOCP<
    decltype(theta_new),
    decltype(f_new),
    decltype(g_new),
    decltype(cr_new),
    decltype(ce_new)>{
    .nx    = ocp.nx,
    .nu    = ocp.nu,
    .nq    = ocp.nq,
    .ncr   = ocp.ncr,
    .nce   = ocp.nce,
    .theta = std::move(theta_new),
    .f     = std::move(f_new),
    .g     = std::move(g_new),
    .cr    = std::move(cr_new),
    .crl   = ocp.crl,
    .cru   = ocp.cru,
    .ce    = std::move(ce_new),
    .cel   = ocp.cel,
    .ceu   = ocp.ceu,
  };
}

/**
 * @brief Solution to OCP problem.
 */
template<LieGroup X, Manifold U>
struct OCPSolution
{
  double t0;
  double tf;

  /// @brief Integral values
  Eigen::VectorXd Q;

  /// @brief Callable functions for state and input
  std::function<U(double)> u;
  std::function<X(double)> x;

  /// @brief Multipliers for integral constraints
  Eigen::VectorXd lambda_q;

  /// @brief Multipliers for endpoint constraints
  Eigen::VectorXd lambda_ce;

  /// @brief Multipliers for dynamics equality constraint
  std::function<Eigen::VectorXd(double)> lambda_dyn;

  /// @brief Multipliers for active running constraints
  std::function<Eigen::VectorXd(double)> lambda_cr;
};

/// @brief Solution to OCP problem defined on flat spaces
using FlatOCPSolution = OCPSolution<Eigen::VectorXd, Eigen::VectorXd>;

/**
 * @brief Unflatten a FlatOCPSolution
 *
 * If flat_sol is a solution to flat_ocp = flatten_ocp(ocp, xl_fun, ul_fun),
 * then unflatten_ocpsol(flat_sol, xl_fun, ul_fun) is a solution to ocp.
 */
template<LieGroup X, Manifold U>
OCPSolution<X, U> unflatten_ocpsol(const FlatOCPSolution & flatsol, auto && xl_fun, auto && ul_fun)
{
  auto u_unflat = [ul_fun = std::forward<decltype(ul_fun)>(ul_fun),
                   usol   = flatsol.u](double t) -> U { return rplus(ul_fun(t), usol(t)); };

  auto x_unflat = [xl_fun = std::forward<decltype(xl_fun)>(xl_fun),
                   xsol   = flatsol.x](double t) -> X { return rplus(xl_fun(t), xsol(t)); };

  return {
    .t0         = flatsol.t0,
    .tf         = flatsol.tf,
    .Q          = flatsol.Q,
    .u          = std::move(u_unflat),
    .x          = std::move(x_unflat),
    .lambda_q   = flatsol.lambda_q,
    .lambda_ce  = flatsol.lambda_ce,
    .lambda_dyn = flatsol.lambda_dyn,
    .lambda_cr  = flatsol.lambda_cr,
  };
}

namespace detail {

/// @brief Variable and constraint structure of an OCP NLP
auto ocp_nlp_structure(const FlatOCPType auto & ocp, const MeshType auto & mesh)
{
  std::size_t N = mesh.N_colloc();

  // variable layout
  std::array<std::size_t, 4> var_len{
    1,                 // tf
    ocp.nq,            // integrals
    ocp.nx * (N + 1),  // states
    ocp.nu * N,        // inputs
  };

  // constraint layout
  std::array<std::size_t, 4> con_len{
    ocp.nx * N,   // derivatives
    ocp.nq,       // other integrals
    ocp.ncr * N,  // running constraints
    ocp.nce,      // end constraints
  };

  std::array<std::size_t, 5> var_beg{0};
  std::partial_sum(var_len.begin(), var_len.end(), var_beg.begin() + 1);

  std::array<std::size_t, 5> con_beg{0};
  std::partial_sum(con_len.begin(), con_len.end(), con_beg.begin() + 1);

  return std::make_tuple(var_beg, var_len, con_beg, con_len);
}

}  // namespace detail

/**
 * @brief Formulate an OCP as a NLP using collocation on a Mesh.
 *
 * @param ocp Optimal control problem definition
 * @param mesh collocation point structure
 * @return encoding of ocp as a nonlinear program
 *
 * @see ocpsol_to_nlpsol(), nlpsol_to_ocpsol()
 */
NLP ocp_to_nlp(const FlatOCPType auto & ocp, const MeshType auto & mesh)
{
  const auto [var_beg, var_len, con_beg, con_len] = detail::ocp_nlp_structure(ocp, mesh);

  // OBJECTIVE FUNCTION

  auto f = [var_beg = var_beg, var_len = var_len, ocp = ocp](const Eigen::VectorXd & x) -> double {
    const auto [tfvar_B, qvar_B, xvar_B, uvar_B, n] = var_beg;
    const auto [tfvar_L, qvar_L, xvar_L, uvar_L]    = var_len;

    assert(std::size_t(x.size()) == n);

    const double t0         = 0;
    const double tf         = x(tfvar_B);
    const Eigen::VectorXd Q = x.segment(qvar_B, qvar_L);
    const Eigen::MatrixXd X = x.segment(xvar_B, xvar_L).reshaped(ocp.nx, xvar_L / ocp.nx);

    return colloc_eval_endpt<false>(1, ocp.nx, ocp.theta, t0, tf, X, Q);
  };

  // OBJECTIVE JACOBIAN

  auto df_dx = [var_beg = var_beg, var_len = var_len, ocp = ocp](
                 const Eigen::VectorXd & x) -> Eigen::SparseMatrix<double> {
    const auto [tfvar_B, qvar_B, xvar_B, uvar_B, n] = var_beg;
    const auto [tfvar_L, qvar_L, xvar_L, uvar_L]    = var_len;

    const double t0         = 0;
    const double tf         = x(tfvar_B);
    const Eigen::VectorXd Q = x.segment(qvar_B, qvar_L);
    const Eigen::MatrixXd X = x.segment(xvar_B, xvar_L).reshaped(ocp.nx, xvar_L / ocp.nx);

    const auto [fval, df_dt0, df_dtf, df_dvecX, df_dQ] =
      colloc_eval_endpt<true>(1, ocp.nx, ocp.theta, t0, tf, X, Q);

    return sparse_block_matrix({
      {df_dtf, df_dQ, df_dvecX, Eigen::SparseMatrix<double>(1, uvar_L)},
    });
  };

  // CONSTRAINT FUNCTION

  auto g = [var_beg = var_beg,
            var_len = var_len,
            con_beg = con_beg,
            con_len = con_len,
            mesh    = mesh,
            ocp     = ocp](const Eigen::VectorXd & x) -> Eigen::VectorXd {
    const auto [tfvar_B, qvar_B, xvar_B, uvar_B, n] = var_beg;
    const auto [tfvar_L, qvar_L, xvar_L, uvar_L]    = var_len;

    const auto [dcon_B, qcon_B, crcon_B, cecon_B, m] = con_beg;
    const auto [dcon_L, qcon_L, crcon_L, cecon_L]    = con_len;

    assert(std::size_t(x.size()) == n);

    const double t0         = 0;
    const double tf         = x(tfvar_B);
    const Eigen::VectorXd Q = x.segment(qvar_B, qvar_L);
    const Eigen::MatrixXd X = x.segment(xvar_B, xvar_L).reshaped(ocp.nx, xvar_L / ocp.nx);
    const Eigen::MatrixXd U = x.segment(uvar_B, uvar_L).reshaped(ocp.nu, uvar_L / ocp.nu);

    Eigen::VectorXd ret(m);
    // clang-format off
    ret.segment(dcon_B, dcon_L)   = colloc_dyn<false>(ocp.nx, ocp.f, mesh, t0, tf, X, U);
    ret.segment(qcon_B, qcon_L)   = colloc_int<false>(ocp.nq, ocp.g, mesh, t0, tf, Q, X, U).reshaped();
    ret.segment(crcon_B, crcon_L) = colloc_eval<false>(ocp.ncr, ocp.cr, mesh, t0, tf, X, U).reshaped();
    ret.segment(cecon_B, cecon_L) = colloc_eval_endpt<false>(ocp.nce, ocp.nx, ocp.ce, t0, tf, X, Q).reshaped();
    // clang-format on
    return ret;
  };

  // CONSTRAINT JACOBIAN
  auto dg_dx =
    [var_beg = var_beg, var_len = var_len, mesh = mesh, ocp = ocp](const Eigen::VectorXd & x) {
      const auto [tfvar_B, qvar_B, xvar_B, uvar_B, n] = var_beg;
      const auto [tfvar_L, qvar_L, xvar_L, uvar_L]    = var_len;

      assert(std::size_t(x.size()) == n);

      const double t0          = 0;
      const double tf          = x(tfvar_B);
      const Eigen::VectorXd Qm = x.segment(qvar_B, qvar_L);
      const Eigen::MatrixXd Xm = x.segment(xvar_B, xvar_L).reshaped(ocp.nx, xvar_L / ocp.nx);
      const Eigen::MatrixXd Um = x.segment(uvar_B, uvar_L).reshaped(ocp.nu, uvar_L / ocp.nu);

      // clang-format off
      const auto [Fval, dF_dt0, dF_dtf, dF_dX, dF_dU]        = colloc_dyn<true>(ocp.nx, ocp.f, mesh, t0, tf, Xm, Um);
      const auto [Gval, dG_dt0, dG_dtf, dG_dQ, dG_dX, dG_dU] = colloc_int<true>(ocp.nq, ocp.g, mesh, t0, tf, Qm, Xm, Um);
      const auto [CRval, dCR_dt0, dCR_dtf, dCR_dX, dCR_dU]   = colloc_eval<true>(ocp.ncr, ocp.cr, mesh, t0, tf, Xm, Um);
      const auto [CEval, dCE_dt0, dCE_dtf, dCE_dX, dCE_dQ]   = colloc_eval_endpt<true>(ocp.nce, ocp.nx, ocp.ce, t0, tf, Xm, Qm);
      // clang-format on

      return sparse_block_matrix({
        // clang-format off
        { dF_dtf,     {},  dF_dX,  dF_dU},
        { dG_dtf,  dG_dQ,  dG_dX,  dG_dU},
        {dCR_dtf,     {}, dCR_dX, dCR_dU},
        {dCE_dtf, dCE_dQ, dCE_dX,     {}},
        // clang-format on
      });
    };

  const auto [tfvar_B, qvar_B, xvar_B, uvar_B, n] = var_beg;
  const auto [tfvar_L, qvar_L, xvar_L, uvar_L]    = var_len;

  const auto [dcon_B, qcon_B, crcon_B, cecon_B, m] = con_beg;
  const auto [dcon_L, qcon_L, crcon_L, cecon_L]    = con_len;

  // VARIABLE BOUNDS

  Eigen::VectorXd xl = Eigen::VectorXd::Constant(n, -std::numeric_limits<double>::infinity());
  Eigen::VectorXd xu = Eigen::VectorXd::Constant(n, std::numeric_limits<double>::infinity());

  xl.segment(tfvar_B, tfvar_L).setZero();  // tf lower bounded by zero

  // CONSTRAINT BOUNDS

  Eigen::VectorXd gl(m);
  Eigen::VectorXd gu(m);

  // derivative constraints are equalities
  gl.segment(dcon_B, dcon_L).setZero();
  gu.segment(dcon_B, dcon_L).setZero();

  // integral constraints are equalities
  gl.segment(qcon_B, qcon_L).setZero();
  gu.segment(qcon_B, qcon_L).setZero();

  // running constraints
  gl.segment(crcon_B, crcon_L) = ocp.crl.replicate(mesh.N_colloc(), 1);
  gu.segment(crcon_B, crcon_L) = ocp.cru.replicate(mesh.N_colloc(), 1);

  // end constraints
  gl.segment(cecon_B, cecon_L) = ocp.cel;
  gu.segment(cecon_B, cecon_L) = ocp.ceu;

  return {
    .n       = n,
    .m       = m,
    .f       = std::move(f),
    .xl      = std::move(xl),
    .xu      = std::move(xu),
    .g       = std::move(g),
    .gl      = std::move(gl),
    .gu      = std::move(gu),
    .df_dx   = std::move(df_dx),
    .dg_dx   = std::move(dg_dx),
    .d2f_dx2 = {},
    .d2g_dx2 = {},
  };
}

/**
 * @brief Convert nonlinear program solution to ocp solution
 */
FlatOCPSolution nlpsol_to_ocpsol(
  const FlatOCPType auto & ocp, const MeshType auto & mesh, const NLPSolution & nlp_sol)
{
  const std::size_t N                             = mesh.N_colloc();
  const auto [var_beg, var_len, con_beg, con_len] = detail::ocp_nlp_structure(ocp, mesh);

  const auto [tfvar_B, qvar_B, xvar_B, uvar_B, n] = var_beg;
  const auto [tfvar_L, qvar_L, xvar_L, uvar_L]    = var_len;

  const auto [dcon_B, qcon_B, crcon_B, cecon_B, m] = con_beg;
  const auto [dcon_L, qcon_L, crcon_L, cecon_L]    = con_len;

  const double t0 = 0;
  const double tf = nlp_sol.x(tfvar_B);

  const Eigen::VectorXd Q = nlp_sol.x.segment(qvar_B, qvar_L);

  // state vector has a value at the endpoint

  Eigen::MatrixXd X(ocp.nx, N + 1);
  X = nlp_sol.x.segment(xvar_B, xvar_L).reshaped(ocp.nx, xvar_L / ocp.nx);

  auto xfun = [t0 = t0, tf = tf, mesh = mesh, X = std::move(X)](double t) -> Eigen::VectorXd {
    return mesh.template eval<Eigen::VectorXd>((t - t0) / (tf - t0), X.colwise(), 0, true);
  };

  // for these we repeat last point since there are no values for endpoint

  Eigen::MatrixXd U(ocp.nu, N);
  U = nlp_sol.x.segment(uvar_B, uvar_L).reshaped(ocp.nu, uvar_L / ocp.nu);

  auto ufun = [t0 = t0, tf = tf, mesh = mesh, U = std::move(U)](double t) -> Eigen::VectorXd {
    return mesh.template eval<Eigen::VectorXd>((t - t0) / (tf - t0), U.colwise(), 0, false);
  };

  Eigen::MatrixXd Ldyn(ocp.nx, N);
  Ldyn = nlp_sol.lambda.segment(dcon_B, dcon_L).reshaped(ocp.nx, dcon_L / ocp.nx);

  auto ldfun =
    [t0 = t0, tf = tf, mesh = mesh, Ldyn = std::move(Ldyn)](double t) -> Eigen::VectorXd {
    return mesh.template eval<Eigen::VectorXd>((t - t0) / (tf - t0), Ldyn.colwise(), 0, false);
  };

  Eigen::MatrixXd Lcr(ocp.ncr, N);
  Lcr = nlp_sol.lambda.segment(crcon_B, crcon_L).reshaped(ocp.ncr, crcon_L / ocp.ncr);

  auto lcrfun = [t0 = t0, tf = tf, mesh = mesh, Lcr = std::move(Lcr)](double t) -> Eigen::VectorXd {
    return mesh.template eval<Eigen::VectorXd>((t - t0) / (tf - t0), Lcr.colwise(), 0, false);
  };

  return {
    .t0         = t0,
    .tf         = tf,
    .Q          = std::move(Q),
    .u          = std::move(ufun),
    .x          = std::move(xfun),
    .lambda_q   = nlp_sol.lambda.segment(qcon_B, qcon_L),
    .lambda_ce  = nlp_sol.lambda.segment(cecon_B, cecon_L),
    .lambda_dyn = std::move(ldfun),
    .lambda_cr  = std::move(lcrfun),
  };
}

/**
 * @brief Convert ocp  solution to nonlinear program
 */
NLPSolution ocpsol_to_nlpsol(
  const FlatOCPType auto & ocp, const MeshType auto & mesh, const FlatOCPSolution & ocpsol)
{
  const std::size_t N                             = mesh.N_colloc();
  const auto [var_beg, var_len, con_beg, con_len] = detail::ocp_nlp_structure(ocp, mesh);

  const auto [tfvar_B, qvar_B, xvar_B, uvar_B, n] = var_beg;
  const auto [tfvar_L, qvar_L, xvar_L, uvar_L]    = var_len;

  const auto [dcon_B, qcon_B, crcon_B, cecon_B, m] = con_beg;
  const auto [dcon_L, qcon_L, crcon_L, cecon_L]    = con_len;

  const auto [nodes, weights] = mesh.all_nodes_and_weights();

  const double t0 = 0;
  const double tf = ocpsol.tf;

  Eigen::VectorXd x(n);
  Eigen::VectorXd zl(n), zu(n);
  Eigen::VectorXd lambda(m);

  zl.setZero();
  zu.setZero();

  x(tfvar_B)                = ocpsol.tf;
  x.segment(qvar_B, qvar_L) = ocpsol.Q;

  lambda.segment(qcon_B, qcon_L)   = ocpsol.lambda_q;
  lambda.segment(cecon_B, cecon_L) = ocpsol.lambda_ce;

  for (auto i = 0u; const auto tau : nodes) {
    x.segment(xvar_B + i * ocp.nx, ocp.nx) = ocpsol.x(t0 + tau * (tf - t0));
    if (i < N) {
      x.segment(uvar_B + i * ocp.nu, ocp.nu)         = ocpsol.u(t0 + tau * (tf - t0));
      lambda.segment(dcon_B + i * ocp.nx, ocp.nx)    = ocpsol.lambda_dyn(t0 + tau * (tf - t0));
      lambda.segment(crcon_B + i * ocp.ncr, ocp.ncr) = ocpsol.lambda_cr(t0 + tau * (tf - t0));
    }
    ++i;
  }

  return {
    .status = NLPSolution::Status::Unknown,
    .x      = std::move(x),
    .zl     = Eigen::VectorXd::Zero(n),
    .zu     = Eigen::VectorXd::Zero(n),
    .lambda = std::move(lambda),
  };
}

}  // namespace smooth::feedback

#endif  // SMOOTH__FEEDBACK__OCP_HPP_


#include <cassert>
#include <cmath>
#include <boost/multiprecision/gmp.hpp>

#include <array>
#include <algorithm>
#include <bitset>
#include <numeric>
#include <sstream>
#include <utility>
#include <vector>

#include <functional>

#include <QtDebug>
#include <QTime>

#include "lrucache.hpp"

using namespace std;
using boost::multiprecision::mpq_rational;

QDebug operator<<(QDebug d, const mpq_rational& r) {
    d.nospace();
    d.noquote();
    stringstream s;
    s << r;
    d << QString::fromStdString(s.str());
    return d.resetFormat();
}

template<class T, std::size_t N>
QDebug operator<<(QDebug d, const array<T, N>& a) {
    d.nospace();
    d.noquote();

    d << "std::array(";
    for(auto i = begin(a); i != end(a); ++i) {
        d << *i;
        if (distance(i, end(a)) > 1) {
            d << ", ";
        }
    }
    d << ")";
    return d.resetFormat();
}

template<class T>
QDebug operator<<(QDebug d, const vector<T>& v) {
    d.nospace();
    d.noquote();
    d << "std::vector(";
    for(auto i = begin(v); i != end(v); ++i) {
        d << *i;
        if (distance(i, end(v)) > 1) {
            d << ", ";
        }
    }
    d << ")";
    return d.resetFormat();
}

template<class S, class T>
QDebug operator<<(QDebug d, const pair<S, T>& p) {
    d.nospace();
    d.noquote();
    d << "std::pair(" << p.first << ", " << p.second << ")";
    return d.resetFormat();
}

// warning: this is only supposed to work for our narrow task of 10 balls for 7 colors. this could go wrong in other cases.
namespace std {
template<size_t L>
struct hash<array<int8_t, L>>
{
    size_t operator()(const array<int8_t, L>& a) const {
        int shift = 0;
        size_t ret = 0;

        foreach (int8_t i, a) {
            ret += i << shift;
            shift += 4;
        }

        return ret;
    }
};
}

template<int NumberOfColors, int NumberOfBallsPerColor>
class Urn {
public:
    Urn() {
        _urn[NumberOfBallsPerColor] = NumberOfColors;
        for (auto i = 0; i < NumberOfBallsPerColor; i += 1) {
            _urn[i] = 0;
        }
    }

    //bool operator==(const Urn<NumberOfColors, NumberOfBallsPerColor>& u) const {
    //    return _urn == u._urn;
    //}

    //bool operator<(const Urn<NumberOfColors, NumberOfBallsPerColor>& u) const {
    //    for (auto i = begin(_urn), j = begin(u._urn); i != end(_urn); ++i, ++j) {
    //        if (*i < *j) return true;
    //    }
    //    return false;
    //}

    bool sumOfElementsIsCorrect() const {
        int8_t sum = 0;
        for (auto i = begin(_urn); i != end(_urn); ++i) {
            sum += *i;
        }
        return sum == NumberOfColors;
    }

    int distinctDrawnColors() const {
        return NumberOfColors - _urn[NumberOfBallsPerColor];
    }

    vector<pair<int /*index*/, mpq_rational /*prob*/>> pickOptions(int remainingBalls) const {
        vector<pair<int, mpq_rational>> options;

        int balls = remainingBalls;

        for (int i = 1; i <= NumberOfBallsPerColor; i += 1) {
            if (_urn[i] > 0) {
                options.emplace_back(make_pair(i, mpq_rational((int)_urn[i]*i, balls)));
            }
        }
        return options;
    }

    void pick(int i) {
        assert(i > 0 && i < (int) _urn.size()); // remove for more speed
        assert(_urn[i] > 0);                    // remove for more speed
        _urn[i] -= 1;
        _urn[i-1] += 1;
    }

    mpq_rational expectedNumberOfDistinctColors(int picks_) const {

        assert(picks_ <= NumberOfColors*NumberOfBallsPerColor);

        function<mpq_rational(Urn<NumberOfColors, NumberOfBallsPerColor>, mpq_rational, int, lru_cache<array<int8_t, NumberOfBallsPerColor + 1>, mpq_rational>*)> recurse = [&recurse, &picks_](Urn<NumberOfColors, NumberOfBallsPerColor> u, mpq_rational p, int picks, lru_cache<array<int8_t, NumberOfBallsPerColor + 1>, mpq_rational>* cache) -> mpq_rational {

            if (picks == 0) {
                return u.distinctDrawnColors();
            }

            if (cache->exists(u._urn)) {
                return cache->get(u._urn);
            }

            mpq_rational e = 0;

            for (const auto& pick : u.pickOptions(NumberOfBallsPerColor*NumberOfColors - (picks_ - picks))) {
                auto uCopy = u;
                uCopy.pick(pick.first);

                e += recurse(uCopy, p*pick.second, picks-1, cache)*pick.second;
            }

            cache->put(u._urn, e);

            return e;
        };

        lru_cache<array<int8_t, NumberOfBallsPerColor + 1>, mpq_rational> cache(100000);

        auto ret = recurse(*this, 1, picks_, &cache);

        return ret;
    }

//private:
public: // for testing
    array<int8_t, NumberOfBallsPerColor + 1> _urn; // _urn[numberofballs] = numberofcolors that have as many balls
};

void test() {
    Urn<7, 10> u710; assert(u710._urn.size() == 11);
    assert(u710._urn[10] == 7);
    for(auto i = 0; i < 10; i += 1) assert(u710._urn[i] == 0);
    assert(u710.sumOfElementsIsCorrect());

    Urn<1, 1> u11;
    assert(u11.expectedNumberOfDistinctColors(1) == 1);

    Urn<1, 2> u12;
    assert(u12.expectedNumberOfDistinctColors(1) == 1);
    assert(u12.expectedNumberOfDistinctColors(2) == 1);

    Urn<2, 3> u22;
    assert(u22.expectedNumberOfDistinctColors(2) == 1*mpq_rational(4, 10) + 2*mpq_rational( 6, 10));
    assert(u22.expectedNumberOfDistinctColors(3) == 1*mpq_rational(4, 40) + 2*mpq_rational(36, 40));
    assert(u22.expectedNumberOfDistinctColors(4) == 2);
    assert(u22.expectedNumberOfDistinctColors(5) == 2);
    assert(u22.expectedNumberOfDistinctColors(6) == 2);

    Urn<3, 2> u32_;
    assert(u32_.expectedNumberOfDistinctColors(2) == 1*mpq_rational( 3, 15) + 2*mpq_rational(12, 15));
    assert(u32_.expectedNumberOfDistinctColors(3) == 2*mpq_rational(18, 30) + 3*mpq_rational(12, 30));
    //assert(u32_.expectedNumberOfDistinctColors(4) == );
    //assert(u32_.expectedNumberOfDistinctColors(5) == );
    assert(u32_.expectedNumberOfDistinctColors(6) == 3);
}

void doit() {
    Urn<7, 10> u;
    qDebug() << u.expectedNumberOfDistinctColors(20);
}

int main() {

    test();

    doit();

    return 0;
}


//Author: Dr. Shantanu Shahane
#include <stdio.h>
#include <math.h>
#include <stdlib.h>
#include <time.h>
#include <float.h>
#include <string.h>
#include <iostream>
#include <vector>
#include <algorithm>
#include "class.hpp"
#include <unistd.h>
#include <limits.h>
#include <Eigen/Dense>
#include <unsupported/Eigen/SparseExtra>
#include <Eigen/SparseLU>
#include <Eigen/OrderingMethods>
#include <Eigen/Eigenvalues>
#include <Eigen/Core>
#include <Eigen/SparseCore>
#include <Spectra/GenEigsSolver.h>
#include <Spectra/MatOp/SparseGenMatProd.h>
#include <Spectra/GenEigsRealShiftSolver.h>
#include <Spectra/MatOp/SparseGenRealShiftSolve.h>
using namespace std;

POINTS::POINTS(PARAMETERS &parameters)
{
    cout << "\n";
    clock_t clock_t1 = clock();
    read_points_xyz_msh(parameters);
    read_points_flag_msh(parameters);
    calc_vert_normal(parameters);
    calc_boundary_face_area(parameters);
    delete_corner_edge_vertices(parameters);
    cout << "POINTS::read_points_xyz_msh original mesh nv: " << nv_original << ", nelem: " << nelem_original << endl;
    cout << "POINTS::POINTS after deleting corners nv: " << nv << endl;
    cout << "\n";
    for (int iv = 0; iv < nv; iv++)
        bc_tag.push_back(INTERIOR_TAG); //initialize to interior; should be set to correct value in main file if required
    for (int icv = 0; icv < nelem_original; icv++)
        elem_bc_tag_original.push_back(INTERIOR_TAG); //initialize to interior; should be set to correct value by calling points.calc_elem_bc_tag(parameters) in main file if required
    parameters.points_timer = ((double)(clock() - clock_t1)) / CLOCKS_PER_SEC;
}

void POINTS::set_periodic_bc(PARAMETERS &parameters, vector<string> periodic_axis)
{
    int dim = parameters.dimension;
    if (periodic_axis.size() > dim)
    {
        cout << "\n\nPOINTS::set_periodic_bc number of periodic axes: " << periodic_axis.size() << " should not be greater than problem dimension: " << dim << "\n\n";
        throw bad_exception();
    }

    for (int ip = 0; ip < periodic_axis.size(); ip++)
        if (periodic_axis[ip] == "x" || periodic_axis[ip] == "X")
            parameters.periodic_bc_index.push_back(0);
        else if (periodic_axis[ip] == "y" || periodic_axis[ip] == "Y")
            parameters.periodic_bc_index.push_back(1);
        else if ((periodic_axis[ip] == "z" || periodic_axis[ip] == "Z") && dim == 3)
            parameters.periodic_bc_index.push_back(2);
        else
        {
            cout << "\n\nPOINTS::set_periodic_bc undefined periodic_axis: " << periodic_axis[ip] << "\n\n";
            throw bad_exception();
        }
    cout << "POINTS::set_periodic_bc periodic axes:";
    for (int ip = 0; ip < periodic_axis.size(); ip++)
        cout << " " << periodic_axis[ip];
    cout << "\n\n";
    vector<int> ind_p = parameters.periodic_bc_index;
    double x_min = xyz[0], x_max = xyz[0], y_min = xyz[1], y_max = xyz[1], z_min = xyz[2], z_max = xyz[2];
    for (int iv = 0; iv < nv; iv++)
    {
        if (x_min > xyz[dim * iv])
            x_min = xyz[dim * iv];
        if (x_max < xyz[dim * iv])
            x_max = xyz[dim * iv];
        if (y_min > xyz[dim * iv + 1])
            y_min = xyz[dim * iv + 1];
        if (y_max < xyz[dim * iv + 1])
            y_max = xyz[dim * iv + 1];
        if (dim == 3)
        {
            if (z_min > xyz[dim * iv + 2])
                z_min = xyz[dim * iv + 2];
            if (z_max < xyz[dim * iv + 2])
                z_max = xyz[dim * iv + 2];
        }
    }
    xyz_min.push_back(x_min), xyz_min.push_back(y_min);
    xyz_max.push_back(x_max), xyz_max.push_back(y_max);
    if (dim == 3)
        xyz_min.push_back(z_min), xyz_max.push_back(z_max);
    for (int i1 = 0; i1 < xyz_min.size(); i1++)
        xyz_length.push_back(xyz_max[i1] - xyz_min[i1]);

    delete_periodic_bc_vertices(parameters);

    vector<bool> empty_bool;
    for (int iv = 0; iv < nv; iv++)
    {
        periodic_bc_flag.push_back(empty_bool);
        for (int ip = 0; ip < ind_p.size(); ip++)
            periodic_bc_flag[iv].push_back(false);
    }
    for (int iv = 0; iv < nv; iv++)
        if (boundary_flag[iv])
            for (int ip = 0; ip < ind_p.size(); ip++)
                if ((fabs(xyz[dim * iv + ind_p[ip]] - xyz_min[ind_p[ip]]) < 1E-5) || (fabs(xyz[dim * iv + ind_p[ip]] - xyz_max[ind_p[ip]]) < 1E-5)) //periodic point is not a real boundary
                    boundary_flag[iv] = false, periodic_bc_flag[iv][ip] = true;

    vector<int> empty_int;
    for (int iv = 0; iv < nv; iv++)
    {
        periodic_bc_section.push_back(empty_int);
        for (int ip = 0; ip < ind_p.size(); ip++)
            periodic_bc_section[iv].push_back(100);
    }
    double xyz_1, xyz_2;
    for (int iv = 0; iv < nv; iv++) //takes values [-1,0,1] for [near_min,middle,near_max] sections respectively
        for (int ip = 0; ip < ind_p.size(); ip++)
        {
            xyz_1 = xyz_min[ind_p[ip]] + (xyz_length[ind_p[ip]] / 3.0);
            xyz_2 = xyz_min[ind_p[ip]] + (2.0 * xyz_length[ind_p[ip]] / 3.0);
            if (xyz[dim * iv + ind_p[ip]] < xyz_1)
                periodic_bc_section[iv][ip] = -1; //section in range [min, xyz_1)
            else if ((xyz_1 <= xyz[dim * iv + ind_p[ip]]) && (xyz[dim * iv + ind_p[ip]] <= xyz_2))
                periodic_bc_section[iv][ip] = 0; //section in range [xyz_1 -> xyz_2]
            else
                periodic_bc_section[iv][ip] = 1; //section in range (xyz_2, max]}
        }
}

void POINTS::delete_periodic_bc_vertices(PARAMETERS &parameters)
{ //delete corner vertices for 2D problems and both corner and edge vertices for 3D problems
    int iv_offset = 0, iv1, dim = parameters.dimension, nv1 = boundary_flag.size();
    vector<int> ind_p = parameters.periodic_bc_index;
    vector<bool> delete_vertices_temp;
    for (int iv0 = 0; iv0 < nv_original; iv0++)
    {
        if (corner_edge_vertices[iv0]) //already deleted vertex
            iv_offset++;
        else
        {
            iv1 = iv0 - iv_offset;
            delete_vertices_temp.push_back(false);
            if (boundary_flag[iv1])
                for (int ip = 0; ip < ind_p.size(); ip++)
                    if (fabs(xyz[dim * iv1 + ind_p[ip]] - xyz_max[ind_p[ip]]) < 1E-5)
                    { //delete higher end of periodic bc
                        delete_vertices_temp[iv1] = true;
                        corner_edge_vertices[iv0] = true; //this vertex is deleted (used in CLOUD::calc_iv_original_nearest_vert)
                    }
        }
    }

    vector<double> xyz_temp, normal_temp;
    vector<bool> b_temp;
    vector<int> bc_tag_temp;
    for (iv1 = 0; iv1 < nv1; iv1++)
    {
        if (delete_vertices_temp[iv1] == false)
        { //iv is a required vertex: thus copy in temporary vectors
            b_temp.push_back(boundary_flag[iv1]);
            bc_tag_temp.push_back(bc_tag[iv1]);
            for (int i = 0; i < dim; i++)
            {
                xyz_temp.push_back(xyz[dim * iv1 + i]);
                normal_temp.push_back(normal[dim * iv1 + i]);
            }
        }
    }
    bc_tag.clear();
    boundary_flag.clear();
    xyz.clear();
    normal.clear();
    boundary_flag.insert(boundary_flag.end(), b_temp.begin(), b_temp.end());
    bc_tag.insert(bc_tag.end(), bc_tag_temp.begin(), bc_tag_temp.end());
    xyz.insert(xyz.end(), xyz_temp.begin(), xyz_temp.end());
    normal.insert(normal.end(), normal_temp.begin(), normal_temp.end());
    b_temp.clear();
    xyz_temp.clear();
    normal_temp.clear();
    delete_vertices_temp.clear();
    bc_tag_temp.clear();
    nv = boundary_flag.size();
    cout << "POINTS::delete_periodic_bc_vertices after deleting periodic_bc vertices nv: " << nv << endl;
}

void POINTS::calc_elem_bc_tag(PARAMETERS &parameters)
{
    int iv_orig, iv_new, dim = parameters.dimension;
    double dist = 0.0;
    vector<int> iv_tag_temp;
    vector<int>::iterator it;
    for (int icv = 0; icv < nelem_original; icv++)
    {
        if (elem_boundary_flag_original[icv])
        {
            for (int i1 = 0; i1 < elem_vert_original[icv].size(); i1++)
            {
                iv_orig = elem_vert_original[icv][i1];
                iv_new = iv_original_nearest_vert[iv_orig];
                dist = 0.0;
                for (int i2 = 0; i2 < dim; i2++)
                    dist = dist + ((xyz[dim * iv_new + i2] - xyz_original[dim * iv_orig + i2]) * (xyz[dim * iv_new + i2] - xyz_original[dim * iv_orig + i2]));
                if (sqrt(dist) < 1E-5) //keep vertices which are not removed (thus, dist should be zero)
                    iv_tag_temp.push_back(bc_tag[iv_new]);
            }
            if (iv_tag_temp.size() == 0)
            {
                printf("\n\nERROR from POINTS::calc_elem_bc_tag icv: %i does not have any vertices which are not deleted; elem_vert_original[icv].size(): %lu\n\n", icv, elem_vert_original[icv].size());
                throw bad_exception();
            }
            sort(iv_tag_temp.begin(), iv_tag_temp.end());          //sort vector
            it = unique(iv_tag_temp.begin(), iv_tag_temp.end());   //get indices of duplicate entries
            iv_tag_temp.erase(it, iv_tag_temp.end());              //delete duplicate entries
            iv_tag_temp.resize(distance(iv_tag_temp.begin(), it)); //resize vector to remove empty entry locations
            if (iv_tag_temp.size() > 1)
            {
                printf("\n\nERROR from POINTS::calc_elem_bc_tag icv: %i has vertices belonging to %lu bc_tag; elem_vert_original[icv].size(): %lu\n\n", icv, iv_tag_temp.size(), elem_vert_original[icv].size());
                throw bad_exception();
            }
            elem_bc_tag_original[icv] = iv_tag_temp[0];

            iv_tag_temp.clear();
        }
    }
}

void POINTS::calc_boundary_face_area(PARAMETERS &parameters)
{ //area or length of boundary elements for 3D or 2D (used to compute fluxes at boundaries)
    //algorithm references: http://geomalgorithms.com/a01-_area.html#3D%20Polygons, https://stackoverflow.com/questions/12642256/python-find-area-of-polygon-from-xyz-coordinates
    double crossprod_sum[3], crossprod[3], v1[3], v2[3], normal[3];
    int iv1, iv2, dim = parameters.dimension;
    vector<vector<int>> vert_nb_cv;
    vector<double> elem_normal;
    if (dim == 3)
    {
        calc_vert_nb_cv(parameters, vert_nb_cv, elem_vert_original);
        calc_elem_normal_3D(elem_normal, vert_nb_cv, elem_vert_original, elem_boundary_flag_original);
    }
    for (int icv = 0; icv < elem_vert_original.size(); icv++)
    {
        boundary_face_area_original.push_back(0.0);
        if (elem_boundary_flag_original[icv])
        {
            if (dim == 3)
            { //3D: area of face
                crossprod_sum[0] = 0.0, crossprod_sum[1] = 0.0, crossprod_sum[2] = 0.0;
                for (int i = 0; i < elem_vert_original[icv].size(); i++)
                {
                    iv1 = elem_vert_original[icv][i];
                    if (i < elem_vert_original[icv].size() - 1)
                        iv2 = elem_vert_original[icv][i + 1];
                    else
                        iv2 = elem_vert_original[icv][0]; //last vertex gets connected to vertex no "0"
                    for (int j = 0; j < 3; j++)
                    {
                        v1[j] = xyz_original[3 * iv1 + j];
                        v2[j] = xyz_original[3 * iv2 + j];
                    }
                    cross_product(crossprod, v1, v2);
                    for (int j = 0; j < 3; j++)
                        crossprod_sum[j] += crossprod[j];
                }
                for (int i = 0; i < 3; i++)
                    normal[i] = elem_normal[3 * icv + i];
                boundary_face_area_original[icv] = 0.5 * fabs(crossprod_sum[0] * normal[0] + crossprod_sum[1] * normal[1] + crossprod_sum[2] * normal[2]);
            }
            else
            { //2D: length of edge
                iv1 = elem_vert_original[icv][0];
                iv2 = elem_vert_original[icv][1];
                for (int j = 0; j < dim; j++)
                {
                    v1[j] = xyz_original[dim * iv1 + j];
                    v2[j] = xyz_original[dim * iv2 + j];
                    boundary_face_area_original[icv] = boundary_face_area_original[icv] + (v1[j] - v2[j]) * (v1[j] - v2[j]);
                }
                boundary_face_area_original[icv] = sqrt(boundary_face_area_original[icv]);
            }
        }
    }
    if (dim == 3)
    {
        for (int iv = 0; iv < vert_nb_cv.size(); iv++)
            vert_nb_cv[iv].clear();
        vert_nb_cv.clear();
        elem_normal.clear();
    }
}

void POINTS::delete_corner_edge_vertices(PARAMETERS &parameters)
{ //delete corner vertices for 2D problems and both corner and edge vertices for 3D problems
    clock_t start = clock();
    vector<double> xyz_temp, normal_temp;
    vector<bool> b_temp;
    int dim = parameters.dimension, nv1 = boundary_flag.size();
    for (int iv = 0; iv < nv1; iv++)
    {
        if (corner_edge_vertices[iv] == false)
        { //iv is a required vertex: thus copy in temporary vectors
            b_temp.push_back(boundary_flag[iv]);
            for (int i = 0; i < dim; i++)
            {
                xyz_temp.push_back(xyz[dim * iv + i]);
                normal_temp.push_back(normal[dim * iv + i]);
            }
        }
    }
    boundary_flag.clear();
    xyz.clear();
    normal.clear();
    boundary_flag.insert(boundary_flag.end(), b_temp.begin(), b_temp.end());
    xyz.insert(xyz.end(), xyz_temp.begin(), xyz_temp.end());
    normal.insert(normal.end(), normal_temp.begin(), normal_temp.end());
    b_temp.clear();
    xyz_temp.clear();
    normal_temp.clear();
    // corner_edge_vertices.clear();
    nv = boundary_flag.size();
}

void POINTS::calc_vert_normal(PARAMETERS &parameters)
{
    vector<vector<int>> elem_vert, vert_nb_cv;
    vector<bool> elem_boundary_flag; //, vert_boundary_flag;
    vector<double> elem_normal;
    read_elem_vert_complete_msh(parameters, elem_vert, elem_boundary_flag);
    calc_vert_nb_cv(parameters, vert_nb_cv, elem_vert);
    if (parameters.dimension == 2)
        calc_elem_normal_2D(elem_normal, vert_nb_cv, elem_vert, elem_boundary_flag);
    else
        calc_elem_normal_3D(elem_normal, vert_nb_cv, elem_vert, elem_boundary_flag);
    elem_boundary_flag_original = elem_boundary_flag;

    for (int iv = 0; iv < parameters.dimension * nv; iv++)
        normal.push_back(0.0);
    int nbcv, count;
    double magnitude;
    for (int iv = 0; iv < nv; iv++)
    {
        if (boundary_flag[iv])
        {
            count = 0;
            for (int i1 = 0; i1 < vert_nb_cv[iv].size(); i1++)
            {
                nbcv = vert_nb_cv[iv][i1];
                if (elem_boundary_flag[nbcv])
                {
                    for (int i2 = 0; i2 < parameters.dimension; i2++)
                        normal[parameters.dimension * iv + i2] += elem_normal[parameters.dimension * nbcv + i2];
                    count++;
                }
            }
            for (int i2 = 0; i2 < parameters.dimension; i2++)
                normal[parameters.dimension * iv + i2] /= ((double)count);
            magnitude = 0.0;
            for (int i2 = 0; i2 < parameters.dimension; i2++)
                magnitude += normal[parameters.dimension * iv + i2] * normal[parameters.dimension * iv + i2];
            magnitude = sqrt(magnitude);
            for (int i2 = 0; i2 < parameters.dimension; i2++)
                normal[parameters.dimension * iv + i2] /= magnitude;
        }
    }

    elem_boundary_flag.clear();
    elem_normal.clear();
    int ncv = elem_vert.size(), nv1 = vert_nb_cv.size();
    for (int icv = 0; icv < ncv; icv++)
        elem_vert[icv].clear();
    elem_vert.clear();
    for (int iv = 0; iv < nv1; iv++)
        vert_nb_cv[iv].clear();
    vert_nb_cv.clear();
}

void POINTS::calc_elem_normal_3D(vector<double> &elem_normal, vector<vector<int>> &vert_nb_cv, vector<vector<int>> &elem_vert, vector<bool> &elem_boundary_flag)
{
    int dim = 3;
    int nv = (int)(xyz.size() / dim), ncv = elem_vert.size();
    int iv_nb_cv_internal, iv0, iv1, iv2;
    double normal[3], direction[3], magnitude, d_temp, vec0[3], vec1[3], centroid[3];
    for (int icv = 0; icv < dim * ncv; icv++)
    {
        elem_normal.push_back(0.0);
    }
    for (int icv = 0; icv < ncv; icv++)
    {
        if (elem_boundary_flag[icv])
        {
            if (elem_vert[icv].size() < 3)
            {
                cout << "\n\nERROR from calc_elem_normal_3D boundary element icv: " << icv << " has only " << elem_vert[icv].size() << " vertices\n\n";
                throw bad_exception();
            }
            iv0 = elem_vert[icv][0];
            iv1 = elem_vert[icv][1];
            iv2 = elem_vert[icv][2];
            for (int i1 = 0; i1 < dim; i1++)
            {
                vec0[i1] = xyz[dim * iv1 + i1] - xyz[dim * iv0 + i1];
                vec1[i1] = xyz[dim * iv2 + i1] - xyz[dim * iv1 + i1];
            }
            cross_product(normal, vec0, vec1);
            iv_nb_cv_internal = -1;
            for (int i1 = 0; i1 < vert_nb_cv[iv0].size(); i1++)
            {
                if (!elem_boundary_flag[vert_nb_cv[iv0][i1]])
                { //is internal CV
                    iv_nb_cv_internal = vert_nb_cv[iv0][i1];
                    break;
                }
            }
            if (iv_nb_cv_internal == -1)
            {
                cout << "\n\nERROR from calc_elem_normal_3D boundary vertex iv0: " << iv0 << " does not have a neighboring internal element\n\n";
                cout << "\n\nERROR from calc_elem_normal_3D boundary element icv: " << icv << " does not have a neighboring internal element\n\n";
                throw bad_exception();
            }
            // iv_nb_vert_internal = -1;
            centroid[0] = 0.0;
            centroid[1] = 0.0;
            centroid[2] = 0.0;
            for (int i1 = 0; i1 < elem_vert[iv_nb_cv_internal].size(); i1++)
            {
                // if (elem_vert[iv_nb_cv_internal][i1] != iv0 && elem_vert[iv_nb_cv_internal][i1] != iv1)
                // if (!vert_boundary_flag[elem_vert[iv_nb_cv_internal][i1]])
                // { //internal vertex found
                //     iv_nb_vert_internal = elem_vert[iv_nb_cv_internal][i1];
                //     break;
                // }
                centroid[0] += xyz[dim * elem_vert[iv_nb_cv_internal][i1]];
                centroid[1] += xyz[dim * elem_vert[iv_nb_cv_internal][i1] + 1];
                centroid[2] += xyz[dim * elem_vert[iv_nb_cv_internal][i1] + 2];
            }
            // if (iv_nb_vert_internal == -1)
            // {
            //     cout << "\ncalc_elem_normal_3D internal element iv_nb_cv_internal: " << iv_nb_cv_internal << " does not have a internal vertex\n\n";
            //     throw bad_exception();
            // }
            centroid[0] = centroid[0] / ((double)(elem_vert[iv_nb_cv_internal].size()));
            centroid[1] = centroid[1] / ((double)(elem_vert[iv_nb_cv_internal].size()));
            centroid[2] = centroid[2] / ((double)(elem_vert[iv_nb_cv_internal].size()));
            // direction[0] = xyz[dim * iv1] - xyz[dim * iv_nb_vert_internal];
            // direction[1] = xyz[dim * iv1 + 1] - xyz[dim * iv_nb_vert_internal + 1];
            // direction[2] = xyz[dim * iv1 + 2] - xyz[dim * iv_nb_vert_internal + 2];
            direction[0] = xyz[dim * iv1] - centroid[0];
            direction[1] = xyz[dim * iv1 + 1] - centroid[1];
            direction[2] = xyz[dim * iv1 + 2] - centroid[2];
            if ((direction[0] * normal[0] + direction[1] * normal[1] + direction[2] * normal[2]) < 0)
            { //make outward facing
                normal[0] = -normal[0];
                normal[1] = -normal[1];
                normal[2] = -normal[2];
            }
            magnitude = sqrt(normal[0] * normal[0] + normal[1] * normal[1] + normal[2] * normal[2]);
            normal[0] = normal[0] / magnitude;
            normal[1] = normal[1] / magnitude;
            normal[2] = normal[2] / magnitude;
            elem_normal[dim * icv] = normal[0];
            elem_normal[dim * icv + 1] = normal[1];
            elem_normal[dim * icv + 2] = normal[2];
        }
    }
}

void POINTS::calc_elem_normal_2D(vector<double> &elem_normal, vector<vector<int>> &vert_nb_cv, vector<vector<int>> &elem_vert, vector<bool> &elem_boundary_flag)
{
    int dim = 2;
    int nv = (int)(xyz.size() / dim), ncv = elem_vert.size();
    int iv_nb_cv_internal, iv0, iv1;
    double normal[2], direction[2], magnitude, d_temp, centroid[2];
    for (int icv = 0; icv < dim * ncv; icv++)
    {
        elem_normal.push_back(0.0);
    }
    for (int icv = 0; icv < ncv; icv++)
    {
        if (elem_boundary_flag[icv])
        {
            iv0 = elem_vert[icv][0];
            iv1 = elem_vert[icv][1];
            normal[0] = xyz[dim * iv1 + 1] - xyz[dim * iv0 + 1];
            normal[1] = -(xyz[dim * iv1] - xyz[dim * iv0]);
            iv_nb_cv_internal = -1;
            for (int i1 = 0; i1 < vert_nb_cv[iv0].size(); i1++)
            {
                if (!elem_boundary_flag[vert_nb_cv[iv0][i1]])
                { //is internal CV
                    iv_nb_cv_internal = vert_nb_cv[iv0][i1];
                    break;
                }
            }
            if (iv_nb_cv_internal == -1)
            {
                cout << "\n\nERROR from calc_elem_normal_2D boundary vertex iv0: " << iv0 << " does not have a neighboring internal element\n\n";
                cout << "\n\nERROR from calc_elem_normal_2D boundary element icv: " << icv << " does not have a neighboring internal element\n\n";
                throw bad_exception();
            }
            // iv_nb_vert_internal = -1;
            centroid[0] = 0.0;
            centroid[1] = 0.0;
            for (int i1 = 0; i1 < elem_vert[iv_nb_cv_internal].size(); i1++)
            {
                // if (!vert_boundary_flag[elem_vert[iv_nb_cv_internal][i1]])
                // { //internal vertex found
                //     iv_nb_vert_internal = elem_vert[iv_nb_cv_internal][i1];
                //     break;
                // }
                centroid[0] += xyz[dim * elem_vert[iv_nb_cv_internal][i1]];
                centroid[1] += xyz[dim * elem_vert[iv_nb_cv_internal][i1] + 1];
            }
            centroid[0] = centroid[0] / ((double)(elem_vert[iv_nb_cv_internal].size()));
            centroid[1] = centroid[1] / ((double)(elem_vert[iv_nb_cv_internal].size()));
            // if (iv_nb_vert_internal == -1)
            // {
            //     cout << "\ncalc_elem_normal_2D internal element iv_nb_cv_internal: " << iv_nb_cv_internal << " does not have a internal vertex\n\n";
            //     throw bad_exception();
            // }
            // direction[0] = xyz[dim * iv1] - xyz[dim * iv_nb_vert_internal];
            // direction[1] = xyz[dim * iv1 + 1] - xyz[dim * iv_nb_vert_internal + 1];
            direction[0] = xyz[dim * iv1] - centroid[0];
            direction[1] = xyz[dim * iv1 + 1] - centroid[1];
            if ((direction[0] * normal[0] + direction[1] * normal[1]) < 0)
            { //make outward facing
                normal[0] = -normal[0];
                normal[1] = -normal[1];
            }
            magnitude = sqrt(normal[0] * normal[0] + normal[1] * normal[1]);
            normal[0] = normal[0] / magnitude;
            normal[1] = normal[1] / magnitude;
            elem_normal[dim * icv] = normal[0];
            elem_normal[dim * icv + 1] = normal[1];
        }
    }
}

void POINTS::read_elem_vert_complete_msh(PARAMETERS &parameters, vector<vector<int>> &elem_vert, vector<bool> &elem_boundary_flag)
{                                   //read all non-trivial elements (do not read point, do not read line for 3D)
    int dim = parameters.dimension; //problem dimension (2 or 3)
    int itemp, cv_type, ncv_full, tag, count = 0;
    double dtemp;
    char temp[50];
    vector<int> vec_temp;
    FILE *file;
    file = fopen(parameters.meshfile.c_str(), "r");
    while (true)
    {
        fscanf(file, "%s ", temp);
        if (strcmp(temp, "$Elements") == 0)
            break;
    }
    fscanf(file, "%i ", &ncv_full);
    for (int icv = 0; icv < ncv_full; icv++)
    {
        fscanf(file, "%i ", &itemp);   //cv number
        fscanf(file, "%i ", &cv_type); //cv type
        fscanf(file, "%i ", &itemp);
        fscanf(file, "%i ", &itemp);
        fscanf(file, "%i ", &tag);     //element tag
        if (cv_type == 15)             //corner vertex CV: do not read
            fscanf(file, "%*[^\n]\n"); //skip reading remaining row
        else if (cv_type == 1)
        { //2 node line CV
            if (dim == 2)
            {
                elem_boundary_flag.push_back(true);
                elem_vert.push_back(vec_temp);
                for (int i1 = 0; i1 < 2; i1++)
                {
                    fscanf(file, "%i ", &itemp); //vertex number
                    elem_vert[count].push_back(itemp - 1);
                }
                count++;
            }
            else
            {
                fscanf(file, "%*[^\n]\n"); //skip reading remaining row
            }
        }
        else if (cv_type == 2)
        { //3 node triangular CV
            elem_vert.push_back(vec_temp);
            for (int i1 = 0; i1 < 3; i1++)
            {
                fscanf(file, "%i ", &itemp); //vertex number
                elem_vert[count].push_back(itemp - 1);
            }
            if (dim == 2)
                elem_boundary_flag.push_back(false);
            else
                elem_boundary_flag.push_back(true);
            count++;
        }
        else if (cv_type == 4)
        { //4 node tetrahedral CV
            elem_vert.push_back(vec_temp);
            for (int i1 = 0; i1 < 4; i1++)
            {
                fscanf(file, "%i ", &itemp); //vertex number
                elem_vert[count].push_back(itemp - 1);
            }
            elem_boundary_flag.push_back(false);
            count++;
        }
        else
        {
            cout << "\n\nERROR from read_elem_vert_complete_msh: Unable to identify CV type: " << cv_type << "\n\n";
            throw bad_exception();
        }
    }
    fclose(file);
    elem_vert_original = elem_vert;
    nelem_original = elem_vert_original.size();
}

void POINTS::calc_vert_nb_cv(PARAMETERS &parameters, vector<vector<int>> &vert_nb_cv, vector<vector<int>> &elem_vert)
{
    clock_t start = clock();
    vector<int> vec_temp;
    int iv1, ncv = elem_vert.size();
    for (int iv = 0; iv < nv; iv++)
        vert_nb_cv.push_back(vec_temp);
    for (int icv = 0; icv < ncv; icv++)
    {
        for (int i1 = 0; i1 < elem_vert[icv].size(); i1++)
        {
            iv1 = elem_vert[icv][i1];
            vert_nb_cv[iv1].push_back(icv);
        }
    }
    for (int iv = 0; iv < nv; iv++)
        if (vert_nb_cv[iv].size() == 0)      //no nbcv
            corner_edge_vertices[iv] = true; //hanging vertex: remove it later
}

void POINTS::read_points_xyz_msh(PARAMETERS &parameters)
{
    int dim = parameters.dimension; //problem dimension (2 or 3)
    FILE *file;
    int itemp;
    double dtemp;
    char temp[50];
    file = fopen(parameters.meshfile.c_str(), "r");
    while (true)
    {
        fscanf(file, "%s ", temp);
        if (strcmp(temp, "$Nodes") == 0)
            break;
    }
    fscanf(file, "%i ", &nv);
    for (int iv = 0; iv < nv; iv++)
    {
        fscanf(file, "%i ", &itemp);  //vertex number
        fscanf(file, "%lf ", &dtemp); //x co-ordinate
        xyz.push_back(dtemp);
        fscanf(file, "%lf ", &dtemp); //y co-ordinate
        xyz.push_back(dtemp);
        fscanf(file, "%lf ", &dtemp); //z co-ordinate
        if (dim == 3)
            xyz.push_back(dtemp); //for 2D problems, Z co-ordinate should be zero and should not be stored
    }
    xyz_original = xyz;
    nv_original = nv;
}

void POINTS::read_points_flag_msh(PARAMETERS &parameters)
{
    vector<int>::iterator it;
    for (int iv = 0; iv < nv; iv++)
    { //initialize
        boundary_flag.push_back(false);
        corner_edge_vertices.push_back(false);
    }
    FILE *file;
    int itemp, ncv, cv_type, tag_int, dim = parameters.dimension;
    double dtemp;
    char temp[50];
    file = fopen(parameters.meshfile.c_str(), "r");
    while (true)
    {
        fscanf(file, "%s ", temp);
        if (strcmp(temp, "$Elements") == 0)
            break;
    }
    fscanf(file, "%i ", &ncv);
    if (dim == 2)
    { //2D problem
        for (int icv = 0; icv < ncv; icv++)
        {
            fscanf(file, "%i ", &itemp);   //cv number
            fscanf(file, "%i ", &cv_type); //cv type
            if (cv_type == 1)
            { //important 2 node line CV on the boundary
                fscanf(file, "%i ", &itemp);
                fscanf(file, "%i ", &itemp);
                fscanf(file, "%i ", &tag_int); //element tag
                for (int i = 0; i < 2; i++)
                {                                    //has 2 vertices
                    fscanf(file, "%i ", &itemp);     //vertex number
                    boundary_flag[itemp - 1] = true; //boundary vertex identified
                }
            }
            else if (cv_type == 15)
            { //corner vertex CV: log this vertex number to delete it later
                fscanf(file, "%i ", &itemp);
                fscanf(file, "%i ", &itemp);
                fscanf(file, "%i ", &tag_int); //element tag
                fscanf(file, "%i ", &itemp);   //vertex number
                corner_edge_vertices[itemp - 1] = true;
            }
            else
            {                              //not reading any other kind of CV
                fscanf(file, "%*[^\n]\n"); //skip reading remaining row
            }
        }
    }
    else
    { //3D problem
        int cv_vert_num;
        for (int icv = 0; icv < ncv; icv++)
        {
            fscanf(file, "%i ", &itemp);   //cv number
            fscanf(file, "%i ", &cv_type); //cv type
            if (cv_type == 2)
            { //important 3 node triangle CV on the boundary
                fscanf(file, "%i ", &itemp);
                fscanf(file, "%i ", &itemp);
                fscanf(file, "%i ", &tag_int); //element tag
                for (int i = 0; i < 3; i++)
                {                                    //has 3 vertices
                    fscanf(file, "%i ", &itemp);     //vertex number
                    boundary_flag[itemp - 1] = true; //boundary vertex identified
                }
            }
            else if (cv_type == 3)
            { //important 4 node quad CV on the boundary
                fscanf(file, "%i ", &itemp);
                fscanf(file, "%i ", &itemp);
                fscanf(file, "%i ", &tag_int); //element tag
                for (int i = 0; i < 4; i++)
                {                                    //has 4 vertices
                    fscanf(file, "%i ", &itemp);     //vertex number
                    boundary_flag[itemp - 1] = true; //boundary vertex identified
                }
            }
            else if (cv_type == 1)
            { //2 node line CV: log these 2 vertex numbers to delete them later
                fscanf(file, "%i ", &itemp);
                fscanf(file, "%i ", &itemp);
                fscanf(file, "%i ", &tag_int); //element tag
                fscanf(file, "%i ", &itemp);   //first vertex number
                corner_edge_vertices[itemp - 1] = true;
                fscanf(file, "%i ", &itemp); //second vertex number
                corner_edge_vertices[itemp - 1] = true;
            }
            else
            {                              //not reading any other kind of CV
                fscanf(file, "%*[^\n]\n"); //skip reading remaining row
            }
        }
    }
}


/*!
 * \file fitness_metric.cc
 *
 * \author Ethan Adams
 * \date
 *
 * This file contains the cpp version of FitnessMetric.py
 */

#include "BingoCpp/fitness_metric.h"
#include <iostream>
#include <stdlib.h>
#include <Eigen/Dense>
#include <Eigen/Core>
#include <unsupported/Eigen/NonLinearOptimization>

namespace bingo {

int LMFunctor::operator()(const Eigen::VectorXd &x, Eigen::VectorXd &fvec) {
  agraphIndv.set_constants(x);
  fvec = fit->evaluate_fitness_vector(agraphIndv, *train);
  return 0;
}

int LMFunctor::df(const Eigen::VectorXd &x, Eigen::MatrixXd &fjac) {
  double epsilon;
  epsilon = 1e-5f;

  for (int i = 0; i < x.size(); i++) {
    Eigen::VectorXd xPlus(x);
    xPlus(i) += epsilon;
    Eigen::VectorXd xMinus(x);
    xMinus(i) -= epsilon;
    Eigen::VectorXd fvecPlus(values());
    operator()(xPlus, fvecPlus);
    Eigen::VectorXd fvecMinus(values());
    operator()(xMinus, fvecMinus);
    Eigen::VectorXd fvecDiff(values());
    fvecDiff = (fvecPlus - fvecMinus) / (2.0 * epsilon);
    fjac.block(0, i, values(), 1) = fvecDiff;
  }

  return 0;
}

double FitnessMetric::evaluate_fitness(AcyclicGraph &indv,
                                       TrainingData &train) {
  if (indv.needs_optimization()) {
    optimize_constants(indv, train);
  }

  return ((evaluate_fitness_vector(indv, train)).abs()).mean();
}

void FitnessMetric::optimize_constants(AcyclicGraph &indv,
                                       TrainingData &train) {
  LMFunctor functor;
  functor.train = &train;
  functor.fit = this;
  functor.m = functor.train->size();
  // indv.input_constants();
  functor.n = indv.count_constants();
  functor.agraphIndv = indv;
  Eigen::VectorXd vec = Eigen::VectorXd::Random(functor.n);
  Eigen::LevenbergMarquardt<LMFunctor, double> lm(functor);
  lm.minimize(vec);
  indv.set_constants(vec);
  indv.needs_opt = false;
}

Eigen::ArrayXXd StandardRegression::evaluate_fitness_vector(AcyclicGraph &indv,
    TrainingData &train) {
  ExplicitTrainingData* temp = dynamic_cast<ExplicitTrainingData*>(&train);
  return (indv.evaluate(temp->x)) - temp->y;
}

ImplicitRegression::ImplicitRegression(int required_params, bool normalize_dot,
                                       double acceptable_nans) {
  this->required_params = required_params;
  this->normalize_dot = normalize_dot;
  acceptable_finite_fracion = 1.0 - acceptable_nans;
}

Eigen::ArrayXXd ImplicitRegression::evaluate_fitness_vector(AcyclicGraph &indv,
    TrainingData &train) {
  ImplicitTrainingData* temp = dynamic_cast<ImplicitTrainingData*>(&train);
  std::pair<Eigen::ArrayXXd, Eigen::ArrayXXd> deriv = indv.evaluate_deriv(
        temp->x);
  Eigen::ArrayXXd dot(deriv.second.rows(), deriv.second.cols());
  double infinity = std::numeric_limits<double>::infinity();

  if (normalize_dot) {
    dot = (deriv.second / (deriv.second.square().rowwise().sum().sqrt())) *
          (temp->dx_dt / (temp->dx_dt.square().rowwise().sum().sqrt()));

  } else {
    dot = deriv.second * temp->dx_dt;
  }

  if (required_params != 0) {
    int n_params_used;

    for (int i = 0; i < dot.rows(); ++i) {
      n_params_used = 0;

      for (int j = 0; j < dot.cols(); ++j) {
        if (dot(i, j) > 0) {
          ++n_params_used;
        }
      }

      if (n_params_used >= required_params) {
        return Eigen::ArrayXXd::Constant(deriv.second.rows(), 1, infinity);
      }
    }
  }

  Eigen::ArrayXXd fit(deriv.second.rows(), 1);
  fit = dot.rowwise().sum() / dot.abs().rowwise().sum();
  return fit;
}
} // namespace bingo


/* +------------------------------------------------------------------------+
   |                     Mobile Robot Programming Toolkit (MRPT)            |
   |                          http://www.mrpt.org/                          |
   |                                                                        |
   | Copyright (c) 2005-2018, Individual contributors, see AUTHORS file     |
   | See: http://www.mrpt.org/Authors - All rights reserved.                |
   | Released under BSD License. See details in http://www.mrpt.org/License |
   +------------------------------------------------------------------------+ */

#include "vision-precomp.h"  // Precompiled headers

#include <mrpt/config.h>
#include <mrpt/vision/utils.h>
#include <mrpt/vision/pnp_algos.h>

#include <iostream>

#include <mrpt/math/types_math.h>  // Eigen must be included first via MRPT to enable the plugin system
#include <Eigen/Core>
#include <Eigen/Dense>

#include <mrpt/otherlibs/do_opencv_includes.h>
#if MRPT_HAS_OPENCV
#include <opencv2/core/eigen.hpp>
#endif

#include "dls.h"
#include "epnp.h"
#include "upnp.h"
#include "p3p.h"
#include "ppnp.h"
#include "posit.h"
#include "lhm.h"
#include "rpnp.h"

bool mrpt::vision::pnp::CPnP::dls(
	const Eigen::Ref<Eigen::MatrixXd> obj_pts,
	const Eigen::Ref<Eigen::MatrixXd> img_pts, int n,
	const Eigen::Ref<Eigen::MatrixXd> cam_intrinsic,
	Eigen::Ref<Eigen::MatrixXd> pose_mat)
{
	try
	{
#if MRPT_HAS_OPENCV == 1

		// Input 2d/3d correspondences and camera intrinsic matrix
		Eigen::MatrixXd cam_in_eig, img_pts_eig, obj_pts_eig;

		// Check for consistency of input matrix dimensions
		if (img_pts.rows() != obj_pts.rows() ||
			img_pts.cols() != obj_pts.cols())
			throw(2);
		else if (cam_intrinsic.rows() != 3 || cam_intrinsic.cols() != 3)
			throw(3);

		if (obj_pts.rows() < obj_pts.cols())
		{
			cam_in_eig = cam_intrinsic.transpose();
			img_pts_eig = img_pts.transpose().block(0, 0, n, 2);
			obj_pts_eig = obj_pts.transpose();
		}
		else
		{
			cam_in_eig = cam_intrinsic;
			img_pts_eig = img_pts.block(0, 0, n, 2);
			obj_pts_eig = obj_pts;
		}

		// Output pose
		Eigen::Matrix3d R_eig;
		Eigen::MatrixXd t_eig;

		// Compute pose
		cv::Mat cam_in_cv(3, 3, CV_32F), img_pts_cv(2, n, CV_32F),
			obj_pts_cv(3, n, CV_32F), R_cv(3, 3, CV_32F), t_cv(3, 1, CV_32F);

		cv::eigen2cv(cam_in_eig, cam_in_cv);
		cv::eigen2cv(img_pts_eig, img_pts_cv);
		cv::eigen2cv(obj_pts_eig, obj_pts_cv);

		mrpt::vision::pnp::dls d(obj_pts_cv, img_pts_cv);
		bool ret = d.compute_pose(R_cv, t_cv);

		cv::cv2eigen(R_cv, R_eig);
		cv::cv2eigen(t_cv, t_eig);

		Eigen::Quaterniond q(R_eig);

		pose_mat << t_eig, q.vec();

		return ret;

#else
		throw(-1);
#endif
	}
	catch (int e)
	{
		switch (e)
		{
			case -1:
				std::cout << "Please install OpenCV for DLS-PnP" << std::endl;
				break;
			case 2:
				std::cout << "2d/3d correspondences mismatch\n Check dimension "
							 "of obj_pts and img_pts"
						  << std::endl;
				break;
			case 3:
				std::cout
					<< "Camera intrinsic matrix does not have 3x3 dimensions "
					<< std::endl;
				break;
		}
		return false;
	}
}

bool mrpt::vision::pnp::CPnP::epnp(
	const Eigen::Ref<Eigen::MatrixXd> obj_pts,
	const Eigen::Ref<Eigen::MatrixXd> img_pts, int n,
	const Eigen::Ref<Eigen::MatrixXd> cam_intrinsic,
	Eigen::Ref<Eigen::MatrixXd> pose_mat)
{
	try
	{
#if MRPT_HAS_OPENCV == 1

		// Input 2d/3d correspondences and camera intrinsic matrix
		Eigen::MatrixXd cam_in_eig, img_pts_eig, obj_pts_eig;

		// Check for consistency of input matrix dimensions
		if (img_pts.rows() != obj_pts.rows() ||
			img_pts.cols() != obj_pts.cols())
			throw(2);
		else if (cam_intrinsic.rows() != 3 || cam_intrinsic.cols() != 3)
			throw(3);

		if (obj_pts.rows() < obj_pts.cols())
		{
			cam_in_eig = cam_intrinsic.transpose();
			img_pts_eig = img_pts.transpose().block(0, 0, n, 2);
			obj_pts_eig = obj_pts.transpose();
		}
		else
		{
			cam_in_eig = cam_intrinsic;
			img_pts_eig = img_pts.block(0, 0, n, 2);
			obj_pts_eig = obj_pts;
		}

		// Output pose
		Eigen::Matrix3d R_eig;
		Eigen::MatrixXd t_eig;

		// Compute pose
		cv::Mat cam_in_cv(3, 3, CV_32F), img_pts_cv(2, n, CV_32F),
			obj_pts_cv(3, n, CV_32F), R_cv, t_cv;

		cv::eigen2cv(cam_in_eig, cam_in_cv);
		cv::eigen2cv(img_pts_eig, img_pts_cv);
		cv::eigen2cv(obj_pts_eig, obj_pts_cv);

		mrpt::vision::pnp::epnp e(cam_in_cv, obj_pts_cv, img_pts_cv);
		e.compute_pose(R_cv, t_cv);

		cv::cv2eigen(R_cv, R_eig);
		cv::cv2eigen(t_cv, t_eig);

		Eigen::Quaterniond q(R_eig);

		pose_mat << t_eig, q.vec();

		return true;

#else
		throw(-1);
#endif
	}
	catch (int e)
	{
		switch (e)
		{
			case -1:
				std::cout << "Please install OpenCV for DLS-PnP" << std::endl;
				break;
			case 2:
				std::cout << "2d/3d correspondences mismatch\n Check dimension "
							 "of obj_pts and img_pts"
						  << std::endl;
				break;
			case 3:
				std::cout
					<< "Camera intrinsic matrix does not have 3x3 dimensions "
					<< std::endl;
				break;
		}
		return false;
	}
}

bool mrpt::vision::pnp::CPnP::upnp(
	const Eigen::Ref<Eigen::MatrixXd> obj_pts,
	const Eigen::Ref<Eigen::MatrixXd> img_pts, int n,
	const Eigen::Ref<Eigen::MatrixXd> cam_intrinsic,
	Eigen::Ref<Eigen::MatrixXd> pose_mat)
{
	try
	{
#if MRPT_HAS_OPENCV == 1

		// Input 2d/3d correspondences and camera intrinsic matrix
		Eigen::MatrixXd cam_in_eig, img_pts_eig, obj_pts_eig;

		// Check for consistency of input matrix dimensions
		if (img_pts.rows() != obj_pts.rows() ||
			img_pts.cols() != obj_pts.cols())
			throw(2);
		else if (cam_intrinsic.rows() != 3 || cam_intrinsic.cols() != 3)
			throw(3);

		if (obj_pts.rows() < obj_pts.cols())
		{
			cam_in_eig = cam_intrinsic.transpose();
			img_pts_eig = img_pts.transpose().block(0, 0, n, 2);
			obj_pts_eig = obj_pts.transpose();
		}
		else
		{
			cam_in_eig = cam_intrinsic;
			img_pts_eig = img_pts.block(0, 0, n, 2);
			obj_pts_eig = obj_pts;
		}

		// Output pose
		Eigen::Matrix3d R_eig;
		Eigen::MatrixXd t_eig;

		// Compute pose
		cv::Mat cam_in_cv(3, 3, CV_32F), img_pts_cv(2, n, CV_32F),
			obj_pts_cv(3, n, CV_32F), R_cv, t_cv;

		cv::eigen2cv(cam_in_eig, cam_in_cv);
		cv::eigen2cv(img_pts_eig, img_pts_cv);
		cv::eigen2cv(obj_pts_eig, obj_pts_cv);

		mrpt::vision::pnp::upnp u(cam_in_cv, obj_pts_cv, img_pts_cv);
		u.compute_pose(R_cv, t_cv);

		cv::cv2eigen(R_cv, R_eig);
		cv::cv2eigen(t_cv, t_eig);

		Eigen::Quaterniond q(R_eig);

		pose_mat << t_eig, q.vec();

		return true;
#else
		throw(-1);
#endif
	}
	catch (int e)
	{
		switch (e)
		{
			case -1:
				std::cout << "Please install OpenCV for DLS-PnP" << std::endl;
				break;
			case 2:
				std::cout << "2d/3d correspondences mismatch\n Check dimension "
							 "of obj_pts and img_pts"
						  << std::endl;
				break;
			case 3:
				std::cout
					<< "Camera intrinsic matrix does not have 3x3 dimensions "
					<< std::endl;
				break;
		}
		return false;
	}
}

bool mrpt::vision::pnp::CPnP::p3p(
	const Eigen::Ref<Eigen::MatrixXd> obj_pts,
	const Eigen::Ref<Eigen::MatrixXd> img_pts, int n,
	const Eigen::Ref<Eigen::MatrixXd> cam_intrinsic,
	Eigen::Ref<Eigen::MatrixXd> pose_mat)
{
	try
	{
		// Input 2d/3d correspondences and camera intrinsic matrix
		Eigen::MatrixXd cam_in_eig, img_pts_eig, obj_pts_eig;

		// Check for consistency of input matrix dimensions
		if (img_pts.rows() != obj_pts.rows() ||
			img_pts.cols() != obj_pts.cols())
			throw(2);
		else if (cam_intrinsic.rows() != 3 || cam_intrinsic.cols() != 3)
			throw(3);

		if (obj_pts.rows() < obj_pts.cols())
		{
			cam_in_eig = cam_intrinsic.transpose();
			img_pts_eig = img_pts.transpose().block(0, 0, n, 2);
			obj_pts_eig = obj_pts.transpose();
		}
		else
		{
			cam_in_eig = cam_intrinsic;
			img_pts_eig = img_pts.block(0, 0, n, 2);
			obj_pts_eig = obj_pts;
		}

		// Output pose
		Eigen::Matrix3d R;
		Eigen::Vector3d t;

		// Compute pose
		mrpt::vision::pnp::p3p p(cam_in_eig);
		bool ret = p.solve(R, t, obj_pts_eig, img_pts_eig);

		Eigen::Quaterniond q(R);

		pose_mat << t, q.vec();

		return ret;
	}
	catch (int e)
	{
		switch (e)
		{
			case 2:
				std::cout << "2d/3d correspondences mismatch\n Check dimension "
							 "of obj_pts and img_pts"
						  << std::endl;
				break;
			case 3:
				std::cout
					<< "Camera intrinsic matrix does not have 3x3 dimensions "
					<< std::endl;
				break;
		}
		return false;
	}
}

bool mrpt::vision::pnp::CPnP::rpnp(
	const Eigen::Ref<Eigen::MatrixXd> obj_pts,
	const Eigen::Ref<Eigen::MatrixXd> img_pts, int n,
	const Eigen::Ref<Eigen::MatrixXd> cam_intrinsic,
	Eigen::Ref<Eigen::MatrixXd> pose_mat)
{
	try
	{
		// Input 2d/3d correspondences and camera intrinsic matrix
		Eigen::MatrixXd cam_in_eig, img_pts_eig, obj_pts_eig;

		// Check for consistency of input matrix dimensions
		if (img_pts.rows() != obj_pts.rows() ||
			img_pts.cols() != obj_pts.cols())
			throw(2);
		else if (cam_intrinsic.rows() != 3 || cam_intrinsic.cols() != 3)
			throw(3);

		if (obj_pts.rows() < obj_pts.cols())
		{
			cam_in_eig = cam_intrinsic.transpose();
			img_pts_eig = img_pts.transpose();
			obj_pts_eig = obj_pts.transpose();
		}
		else
		{
			cam_in_eig = cam_intrinsic;
			img_pts_eig = img_pts;
			obj_pts_eig = obj_pts;
		}

		// Output pose
		Eigen::Matrix3d R;
		Eigen::Vector3d t;

		// Compute pose
		mrpt::vision::pnp::rpnp r(obj_pts_eig, img_pts_eig, cam_in_eig, n);
		bool ret = r.compute_pose(R, t);

		Eigen::Quaterniond q(R);

		pose_mat << t, q.vec();

		return ret;
	}
	catch (int e)
	{
		switch (e)
		{
			case 2:
				std::cout << "2d/3d correspondences mismatch\n Check dimension "
							 "of obj_pts and img_pts"
						  << std::endl;
				break;
			case 3:
				std::cout
					<< "Camera intrinsic matrix does not have 3x3 dimensions "
					<< std::endl;
				break;
		}
		return false;
	}
}

bool mrpt::vision::pnp::CPnP::ppnp(
	const Eigen::Ref<Eigen::MatrixXd> obj_pts,
	const Eigen::Ref<Eigen::MatrixXd> img_pts, int n,
	const Eigen::Ref<Eigen::MatrixXd> cam_intrinsic,
	Eigen::Ref<Eigen::MatrixXd> pose_mat)
{
	try
	{
		// Input 2d/3d correspondences and camera intrinsic matrix
		Eigen::MatrixXd cam_in_eig, img_pts_eig, obj_pts_eig;

		// Check for consistency of input matrix dimensions
		if (img_pts.rows() != obj_pts.rows() ||
			img_pts.cols() != obj_pts.cols())
			throw(2);
		else if (cam_intrinsic.rows() != 3 || cam_intrinsic.cols() != 3)
			throw(3);

		if (obj_pts.rows() < obj_pts.cols())
		{
			cam_in_eig = cam_intrinsic.transpose();
			img_pts_eig = img_pts.transpose();
			obj_pts_eig = obj_pts.transpose();
		}
		else
		{
			cam_in_eig = cam_intrinsic;
			img_pts_eig = img_pts;
			obj_pts_eig = obj_pts;
		}

		// Output pose
		Eigen::Matrix3d R;
		Eigen::Vector3d t;

		// Compute pose
		mrpt::vision::pnp::ppnp p(obj_pts_eig, img_pts_eig, cam_in_eig);

		bool ret = p.compute_pose(R, t, n);

		Eigen::Quaterniond q(R);

		pose_mat << t, q.vec();

		return ret;
	}
	catch (int e)
	{
		switch (e)
		{
			case 2:
				std::cout << "2d/3d correspondences mismatch\n Check dimension "
							 "of obj_pts and img_pts"
						  << std::endl;
				break;
			case 3:
				std::cout
					<< "Camera intrinsic matrix does not have 3x3 dimensions "
					<< std::endl;
				break;
		}
		return false;
	}
}

bool mrpt::vision::pnp::CPnP::posit(
	const Eigen::Ref<Eigen::MatrixXd> obj_pts,
	const Eigen::Ref<Eigen::MatrixXd> img_pts, int n,
	const Eigen::Ref<Eigen::MatrixXd> cam_intrinsic,
	Eigen::Ref<Eigen::MatrixXd> pose_mat)
{
	try
	{
		// Input 2d/3d correspondences and camera intrinsic matrix
		Eigen::MatrixXd cam_in_eig, img_pts_eig, obj_pts_eig;

		// Check for consistency of input matrix dimensions
		if (img_pts.rows() != obj_pts.rows() ||
			img_pts.cols() != obj_pts.cols())
			throw(2);
		else if (cam_intrinsic.rows() != 3 || cam_intrinsic.cols() != 3)
			throw(3);

		if (obj_pts.rows() < obj_pts.cols())
		{
			cam_in_eig = cam_intrinsic.transpose();
			img_pts_eig = img_pts.transpose().block(0, 0, n, 2);
			obj_pts_eig = obj_pts.transpose();
		}
		else
		{
			cam_in_eig = cam_intrinsic;
			img_pts_eig = img_pts.block(0, 0, n, 2);
			obj_pts_eig = obj_pts;
		}

		// Output pose
		Eigen::Matrix3d R;
		Eigen::Vector3d t;

		// Compute pose
		mrpt::vision::pnp::posit p(obj_pts_eig, img_pts_eig, cam_in_eig, n);

		bool ret = p.compute_pose(R, t);

		Eigen::Quaterniond q(R);

		pose_mat << t, q.vec();

		return ret;
	}
	catch (int e)
	{
		switch (e)
		{
			case 2:
				std::cout << "2d/3d correspondences mismatch\n Check dimension "
							 "of obj_pts and img_pts"
						  << std::endl;
				break;
			case 3:
				std::cout
					<< "Camera intrinsic matrix does not have 3x3 dimensions "
					<< std::endl;
				break;
		}
		return false;
	}
}

bool mrpt::vision::pnp::CPnP::lhm(
	const Eigen::Ref<Eigen::MatrixXd> obj_pts,
	const Eigen::Ref<Eigen::MatrixXd> img_pts, int n,
	const Eigen::Ref<Eigen::MatrixXd> cam_intrinsic,
	Eigen::Ref<Eigen::MatrixXd> pose_mat)
{
	try
	{
		// Input 2d/3d correspondences and camera intrinsic matrix
		Eigen::MatrixXd cam_in_eig, img_pts_eig, obj_pts_eig;

		// Check for consistency of input matrix dimensions
		if (img_pts.rows() != obj_pts.rows() ||
			img_pts.cols() != obj_pts.cols())
			throw(2);
		else if (cam_intrinsic.rows() != 3 || cam_intrinsic.cols() != 3)
			throw(3);

		if (obj_pts.rows() < obj_pts.cols())
		{
			cam_in_eig = cam_intrinsic.transpose();
			img_pts_eig = img_pts.transpose();
			obj_pts_eig = obj_pts.transpose();
		}
		else
		{
			cam_in_eig = cam_intrinsic;
			img_pts_eig = img_pts;
			obj_pts_eig = obj_pts;
		}

		// Output pose
		Eigen::Matrix3d R;
		Eigen::Vector3d t;

		// Compute pose
		mrpt::vision::pnp::lhm l(obj_pts_eig, img_pts_eig, cam_intrinsic, n);

		bool ret = l.compute_pose(R, t);

		Eigen::Quaterniond q(R);

		pose_mat << t, q.vec();

		return ret;
	}
	catch (int e)
	{
		switch (e)
		{
			case 2:
				std::cout << "2d/3d correspondences mismatch\n Check dimension "
							 "of obj_pts and img_pts"
						  << std::endl;
				break;
			case 3:
				std::cout
					<< "Camera intrinsic matrix does not have 3x3 dimensions "
					<< std::endl;
				break;
		}
		return false;
	}
}


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

#ifndef BOOST_MATH_TOOLS_MINIMA_HPP
#define BOOST_MATH_TOOLS_MINIMA_HPP

#include <utility>
#include <cmath>
#include <boost/math/tools/precision.hpp>
#include <boost/math/policies/policy.hpp>
#include <boost/cstdint.hpp>

namespace boost{ namespace math{ namespace tools{

template <class F, class T>
std::pair<T, T> brent_find_minima(F f, T min, T max, int bits, boost::uintmax_t& max_iter)
{
   BOOST_MATH_STD_USING
   bits = (std::min)(policies::digits<T, policies::policy<> >() / 2, bits);
   T tolerance = static_cast<T>(ldexp(1.0, 1-bits));
   T x;  // minima so far
   T w;  // second best point
   T v;  // previous value of w
   T u;  // most recent evaluation point
   T delta;  // The distance moved in the last step
   T delta2; // The distance moved in the step before last
   T fu, fv, fw, fx;  // function evaluations at u, v, w, x
   T mid; // midpoint of min and max
   T fract1, fract2;  // minimal relative movement in x

   static const T golden = 0.3819660f;  // golden ratio, don't need too much precision here!

   x = w = v = max;
   fw = fv = fx = f(x);
   delta2 = delta = 0;

   uintmax_t count = max_iter;

   do{
      // get midpoint
      mid = (min + max) / 2;
      // work out if we're done already:
      fract1 = tolerance * fabs(x) + tolerance / 4;
      fract2 = 2 * fract1;
      if(fabs(x - mid) <= (fract2 - (max - min) / 2))
         break;

      if(fabs(delta2) > fract1)
      {
         // try and construct a parabolic fit:
         T r = (x - w) * (fx - fv);
         T q = (x - v) * (fx - fw);
         T p = (x - v) * q - (x - w) * r;
         q = 2 * (q - r);
         if(q > 0)
            p = -p;
         q = fabs(q);
         T td = delta2;
         delta2 = delta;
         // determine whether a parabolic step is acceptible or not:
         if((fabs(p) >= fabs(q * td / 2)) || (p <= q * (min - x)) || (p >= q * (max - x)))
         {
            // nope, try golden section instead
            delta2 = (x >= mid) ? min - x : max - x;
            delta = golden * delta2;
         }
         else
         {
            // whew, parabolic fit:
            delta = p / q;
            u = x + delta;
            if(((u - min) < fract2) || ((max- u) < fract2))
               delta = (mid - x) < 0 ? -fabs(fract1) : fabs(fract1);
         }
      }
      else
      {
         // golden section:
         delta2 = (x >= mid) ? min - x : max - x;
         delta = golden * delta2;
      }
      // update current position:
      u = (fabs(delta) >= fract1) ? x + delta : (delta > 0 ? x + fabs(fract1) : x - fabs(fract1));
      fu = f(u);
      if(fu <= fx)
      {
         // good new point is an improvement!
         // update brackets:
         if(u >= x)
            min = x;
         else
            max = x;
         // update control points:
         v = w;
         w = x;
         x = u;
         fv = fw;
         fw = fx;
         fx = fu;
      }
      else
      {
         // Oh dear, point u is worse than what we have already,
         // even so it *must* be better than one of our endpoints:
         if(u < x)
            min = u;
         else
            max = u;
         if((fu <= fw) || (w == x))
         {
            // however it is at least second best:
            v = w;
            w = u;
            fv = fw;
            fw = fu;
         }
         else if((fu <= fv) || (v == x) || (v == w))
         {
            // third best:
            v = u;
            fv = fu;
         }
      }

   }while(--count);

   max_iter -= count;

   return std::make_pair(x, fx);
}

template <class F, class T>
inline std::pair<T, T> brent_find_minima(F f, T min, T max, int digits)
{
   boost::uintmax_t m = (std::numeric_limits<boost::uintmax_t>::max)();
   return brent_find_minima(f, min, max, digits, m);
}

}}} // namespaces

#endif


#include <cmath>
#include <complex>
#include <cstring>
#include <algorithm>
#include <gsl/gsl_errno.h>
#include <gsl/gsl_math.h>
#include <gsl/gsl_odeiv.h>
#include <Eigen/Core>

#include "Const.h"
#include "SpeechSynthesizer.h"

using namespace std;
using namespace Eigen;

static const double NasalCavityArea[] = { 1, 2, 3, 4, 6, 8, 8, 7, 4, 2, 2 };

SpeechSynthesizer::SpeechSynthesizer(int sampling, int freqResolution, vector<double> &tractArea, double dl) :
	m1	(0.125),
	m2	(0.025),
	d1	(0.25),
	d2	(0.05),
	lg	(1.4),
	es	(100),
	eh	(500),
	ks1	(8.0E+4),
	ks2	(8.0E+3),
	kh1	(3 * ks1),
	kh2	(3 * ks2),
	kc	(2.5E+4),
	om1	(sqrt(ks1 / m1)),
	Ag0	(0.05),
	p_s	(7.85E+3),
	x1	(0),
	x2	(0),
	v1	(0),
	v2	(0),
	ug_buf	(freqResolution),
	toVelum		(8.0),
    toSinus		(7.0),
    toConstr	(3.0),
    R_sin		(1.0),
    L_sin		(5.94E-3),
    C_sin		(15.8E-6),
	oral		(4.0, dl, tractArea),
	nasal		(72.0, 1.0, NasalCavityArea, 11),
	r_N(0.2),
	r_vib(3.0),
	dt	(1.0 / sampling),
	df	(1.0 / (freqResolution * dt)),
	N	(freqResolution),
	calcSpan(0.020),
	t(0)
{
	z_in_curr = new double[N];
	h_out_curr = new double[N];
	z_in_prev = new double[N];
	h_out_prev = new double[N];

	for (int i = 0; i < N; i++)
		ug_buf.push_front(0);

	lastCalcTime = -2 * calcSpan;
	CalcTract();
	memcpy(z_in_prev, z_in_curr, N);
	memcpy(h_out_prev, h_out_curr, N);
	lastCalcSpan = calcSpan;

	system.function = Differentiate;
	system.jacobian = nullptr;
	system.dimension = 5;
	system.params = this;
	step = gsl_odeiv_step_alloc(gsl_odeiv_step_rk4, system.dimension);
}

SpeechSynthesizer::~SpeechSynthesizer(void)
{
	delete [] z_in_curr;
	delete [] h_out_curr;
	delete [] z_in_prev;
	delete [] h_out_prev;
}

int SpeechSynthesizer::Differentiate(double t, const double y[], double dydt[], void *params)
{
	SpeechSynthesizer *s = (SpeechSynthesizer*) params;
	const double h_min = s->h_min;
	const double x1 = y[0];
	const double v1 = y[1];
	const double x2 = y[2];
	const double v2 = y[3];
	const double ug = y[4];
	double *dx1 = &dydt[0];
	double *dv1 = &dydt[1];
	double *dx2 = &dydt[2];
	double *dv2 = &dydt[3];
	double *dug = &dydt[4];

	const int n = s->N;
	double *z_in = new double[n];
	double V = 0;

	s->GetTractImpedance(z_in, t);
	for (int i = 0; i < n; i++)
	{
		V += z_in[i] * s->ug_buf[i];
	}

	double xc = -s->Ag0 / (2 * s->lg);
	double h1 = x1 - xc;
	double h2 = x2 - xc;
	if (h1 > 0 && h1 < h_min) h1 = h_min;
	if (h2 > 0 && h2 < h_min) h2 = h_min;
	double Rv = 1.5 * Const::mu / s->lg;
	double Rv1 = Rv * s->d1 / gsl_pow_3(h1);
	double Rv2 = Rv * s->d2 / gsl_pow_3(h2);
	double Rbb = Const::rho / (8 * gsl_pow_2(s->lg));
	double Rc = Rbb * 1.37 / gsl_pow_2(h1) * (ug);
	double R12 = Rbb * (1 / gsl_pow_2(h2) - 1 / gsl_pow_2(h1)) * (ug);
	double Re = -Rbb * 0.5 / gsl_pow_2(h2) * (ug);
	double Lg = Const::rho / (2 * s->lg);
	double Lg1 = Lg * s->d1 / h1;
	double Lg2 = Lg * s->d2 / h2;

	*dug = (h1 <= 0 || h2 <= 0
			? -ug / (t - s->t)
			: (s->p_s - V - (Rc + Rv1 + R12 + Rv2 + Re) * ug) / (Lg1 + Lg2));

	double P1, P2;
	if (h1 <= 0)
	{
		P1 = s->p_s;
		P2 = 0;
	}
	else if (h2 <= 0)
	{
		P1 = s->p_s;
		P2 = P1;
	}
	else
	{
		P1 = s->p_s - (Rc + Rv1 / 2) * ug - Lg1 / 2 * *dug;
		P2 = P1 - (R12 + (Rv1 + Rv2) / 2) * ug - (Lg1 + Lg2) / 2 * *dug;
	}

	double s1 = s->ks1 * x1 * (1 + s->es * gsl_pow_2(x1));
	if (h1 <= 0) s1 += s->kh1 * h1 * (1 + s->eh * gsl_pow_2(h1));
	double s2 = s->ks2 * x2 * (1 + s->es * gsl_pow_2(x2));
	if (h2 <= 0) s2 += s->kh2 * h2 * (1 + s->eh * gsl_pow_2(h2));
	double r1 = 2 * (h1 > 0 ? 0.2 : 1.1) * sqrt(s->ks1 * s->m1);
	double r2 = 2 * (h2 > 0 ? 0.6 : 1.9) * sqrt(s->ks2 * s->m2);

	*dv1 = (-r1 * v1 - s1 - s->kc * (x1 - x2) + P1 * s->lg * s->d1) / s->m1;
	*dv2 = (-r2 * v2 - s2 - s->kc * (x2 - x1) + P2 * s->lg * s->d2) / s->m2;

	*dx1 = v1;
	*dx2 = v2;

	delete [] z_in;

	return GSL_SUCCESS;
}

double SpeechSynthesizer::Step(void)
{
	CalcTract();

	double y[5] = { x1, v1, x2, v2, ug_buf.front() };
	double yerr[5];

	gsl_odeiv_step_apply(step, t, dt, y, yerr, t == 0 ? nullptr : dydt, dydt, &system);
	t += dt;

	if (yerr[0] > 1E-3)
		std::cerr << "large error" << std::endl;

	x1 = y[0];
	v1 = y[1];
	x2 = y[2];
	v2 = y[3];
	ug_buf.pop_back();
	ug_buf.push_front(y[4]);

	double *h_out = new double[N];
	double p = 0;

	GetOutputImpedance(h_out, t);
	for (int i = 0; i < N; i++)
		p += h_out[i] * ug_buf[i];

	delete [] h_out;

	std::printf("%f %.3e %.3e %.3e %.3e %.3e %.3e\n", t, x1 * 2 * lg + Ag0, v1, x2 * 2 * lg + Ag0, v2, ug_buf.front(), p);

	return p;
}

double SpeechSynthesizer::GetTime(void)
{
	return t;
}

void SpeechSynthesizer::CalcTract(void)
{
	double dur = t - lastCalcTime;

	if (dur < calcSpan)
		return;

	lastCalcSpan = dur;
	lastCalcTime = t;

#ifdef _MSC_VER
	complex<double> *Z_in = new complex<double>[N - 1];
	complex<double> *H_out = new complex<double>[N - 1];
#else
	complex<double> Z_in[N - 1];
	complex<double> H_out[N - 1];
#endif

	int i_v = (int) (toVelum / oral.GetElemLength());
	int i_c = i_v + (int) (toConstr / oral.GetElemLength());
	int i_s = (int) (toSinus / nasal.GetElemLength());

    for (int i = 1; i < N; i++)
    {
		double f = i * df;
        double omg = 2 * Const::pi * f;
        complex<double> s(0, omg);

		Matrix2cd K_G = oral.ChainMatrix(f, 0, i_v);
        Matrix2cd K_C = oral.ChainMatrix(f, i_v, i_c);
        Matrix2cd K_L = oral.ChainMatrix(f, i_c);
		double r_L = sqrt(oral.GetEndArea() / Const::pi);
		double ka = omg / Const::c * r_L;
		complex<double> Z_L = Const::rho * Const::c / oral.GetEndArea() * complex<double>(ka * ka / 2, 8.0 * ka / (3.0 * Const::pi));

        complex<double> Z_sin = R_sin + L_sin * s + 1.0 / (C_sin * s);
        Matrix2cd K_N1 = nasal.ChainMatrix(f, 0, i_s);
        Matrix2cd K_N2 = nasal.ChainMatrix(f, i_s);
        Matrix2cd K_sin;
		K_sin << 1, 0, -1.0 / Z_sin, 1;
        Matrix2cd K_N = K_N2 * K_sin * K_N1;
        complex<double> Z_N = 4 * Const::pi * gsl_pow_2(r_N) * Const::rho * Const::c / (gsl_pow_2(Const::c) + gsl_pow_2(r_N * omg)) * (gsl_pow_2(r_N * omg) + Const::c * r_N * s);

        Matrix2cd K_oral = K_L * K_C;
		complex<double> Z_VN = (K_N(1, 1) * Z_N - K_N(0, 1)) / (K_N(0, 0) - K_N(1, 0) * Z_N);

		complex<double> Z_VT = (K_oral(1, 1) * Z_L - K_oral(0, 1)) / (K_oral(0, 0) - K_oral(1, 0) * Z_L);
        Matrix2cd K_cN, K_cT;
		K_cN << 1, 0, -1.0 / Z_VN, 1;
        K_cT << 1, 0, -1.0 / Z_VT, 1;

        Matrix2cd K_fric = K_C * K_cN * K_G;
        Matrix2cd K_tract = K_L * K_fric;
        Matrix2cd K_nasal = K_N * K_cT * K_G;

		//double Z_0 = Const::rho * Const::c / oral.GetEndArea();
        Z_in[i - 1] = (K_tract(1, 1) * Z_L - K_tract(0, 1)) / (K_tract(0, 0) - K_tract(1, 0) * Z_L);

		complex<double> H_vib = oral.GetStartArea() / Const::c * s * r_vib / (Const::c + s * r_vib) * Z_in[i - 1] * oral.GetBeta(f);

        H_out[i - 1] = Z_L / (K_tract(0, 0) - K_tract(1, 0) * Z_L)
            + Z_N / (K_nasal(0, 0) - K_nasal(1, 0) * Z_N)
            + H_vib;
    }

	swap(z_in_curr, z_in_prev);
	swap(h_out_curr, h_out_prev);
	CalcResponse(z_in_curr, Z_in);
	CalcResponse(h_out_curr, H_out);

#ifdef _MSC_VER
	delete [] Z_in;
	delete [] H_out;
#endif
}

inline static complex<double> Exp(const complex<double> z)
{
    double re = real(z);
    double im = imag(z);
    return exp(re) * complex<double>(cos(im), sin(im));
}

void SpeechSynthesizer::CalcResponse(double res[], const complex<double> spec[])
{
    for (int i = 0; i < N; i++)
    {
		double t = i * dt;
		double c_t = 0.54 - 0.46 * cos(Const::pi * (1 + (double)i / N));
		complex<double> r(0, 0);

		for (int j = 1; j < N; j++)
		{
	        double omg = 2 * Const::pi * j * df;
			complex<double> s(0, omg);
			double c_f = 1 / (1 + exp(4 * (8.0 * j / N - 5)));	// designed to decrease around 5/8

			r += c_t * c_f / (2 * Const::pi * N) * spec[j - 1] * Exp(s * t);
		}

		res[i] = real(r);
	}
}

void SpeechSynthesizer::GetTractImpedance(double res[], double t)
{
	double dur = t - lastCalcTime;

	for (int i = 0; i < N; i++)
	{
		res[i] = z_in_curr[i] + (z_in_curr[i] - z_in_prev[i]) / calcSpan * dur;
	}
}
void SpeechSynthesizer::GetOutputImpedance(double res[], double t)
{
	double dur = t - lastCalcTime;

	for (int i = 0; i < N; i++)
	{
		res[i] = h_out_curr[i] + (h_out_curr[i] - h_out_prev[i]) / calcSpan * dur;
	}
}


#include "MyModel_Pixels.h"
#include "DNest4/code/DNest4.h"
#include <stdexcept>
#include <armadillo>

namespace Obscurity
{

// CONSTRUCTOR AND MEMBER FUNCTIONS
MyModel_Pixels::MyModel_Pixels()
:n(ni, nj)
,obscurer_map(ni, nj)
,convolved(ni, nj)
{

}

void MyModel_Pixels::from_prior(DNest4::RNG& rng)
{
    for(size_t j=0; j<nj; ++j)
        for(size_t i=0; i<ni; ++i)
            n(i, j) = rng.randn();

    sd = rng.rand();

    DNest4::Cauchy c1(0.0, 1.0);
    x0 = -std::abs(c1.generate(rng));

    DNest4::Cauchy c2(0.0, 0.1*data.get_t_range());
    timescale = std::abs(c2.generate(rng));

    calculate_obscurer_map();
}

double MyModel_Pixels::calculate_total_flux(double time) const
{
    double speed = 1.0/timescale;
    double offset = x0 + speed*time;

    int j = (int)floor((offset - (x_min + 0.5*dx))/dx);
    return convolved(ni/2, j%nj);
}

double MyModel_Pixels::perturb(DNest4::RNG& rng)
{
    double logH = 0.0;

    if(rng.rand() <= 0.7)
    {
        // Perturb the ns
        int which_i, which_j;
        int reps = 1;
        if(rng.rand() <= 0.5)
            reps = (int)pow(10.0, 3*rng.rand());

        if(reps == 1)
        {
            which_i = rng.rand_int(ni);
            which_j = rng.rand_int(nj);

            logH -= -0.5*pow(n(which_i, which_j), 2);
            n(which_i, which_j) += rng.randh();
            logH += -0.5*pow(n(which_i, which_j), 2);
        }
        else
        {
            for(int i=0; i<reps; ++i)
            {
                which_i = rng.rand_int(ni);
                which_j = rng.rand_int(nj);
                n(which_i, which_j) = rng.randn();
            }
        }
    }
    else
    {
        int which = rng.rand_int(3);

        if(which == 0)
        {
            sd += rng.randh();
            DNest4::wrap(sd, 0.0, 1.0);
            calculate_obscurer_map();
        }
        else if(which == 1)
        {
            DNest4::Cauchy c(0.0, 1.0);
            logH += c.perturb(x0, rng);
            x0 = -std::abs(x0);
        }
        else
        {
            DNest4::Cauchy c(0.0, 0.1*data.get_t_range());
            logH += c.perturb(timescale, rng);
            timescale = std::abs(timescale);
        }
    }

    return logH;
}

double MyModel_Pixels::log_likelihood() const
{
    double logL = 0.0;

    const auto& t = data.get_t();
    const auto& y = data.get_y();
    const auto& sig = data.get_sig();

    double model_prediction;
    for(size_t i=0; i<t.size(); ++i)
    {
        model_prediction = calculate_total_flux(t[i]);
        logL += -0.5*log(2*M_PI) - log(sig[i])
                    - 0.5*pow((y[i] - model_prediction)/sig[i], 2);
    }

    return logL;
}

void MyModel_Pixels::calculate_obscurer_map()
{
    // Obscurer image
    for(size_t j=0; j<nj; ++j)
        for(size_t i=0; i<ni; ++i)
            obscurer_map(i, j) = exp(-sd*n(i, j));

    // FFT of obscurer_map
    arma::cx_mat A = arma::fft2(obscurer_map);

    // (FFT of obscurer map)*(FFT of star map)
    for(size_t j=0; j<nj; ++j)
        for(size_t i=0; i<ni; ++i)
            A(i, j) *= fft_of_star(i, j);

    // obscurer map convolved with star map
    A = arma::ifft2(A);
    for(size_t i=0; i<ni; ++i)
        for(size_t j=0; j<nj; ++j)
            convolved(i, j) = A(i, j).real();
}

void MyModel_Pixels::print(std::ostream& out) const
{
    const auto& t = data.get_t();
    for(size_t i=0; i<t.size(); ++i)
        out<<calculate_total_flux(t[i])<<' ';

    for(size_t i=0; i<ni; ++i)
        for(size_t j=0; j<nj; ++j)
            out<<star(i, j)<<' ';

    for(size_t i=0; i<ni; ++i)
        for(size_t j=0; j<nj; ++j)
            out<<obscurer_map(i, j)<<' ';
}

std::string MyModel_Pixels::description() const
{
    return std::string("");
}

/* STATIC STUFF */

Data                  MyModel_Pixels::data;
arma::mat             MyModel_Pixels::star(MyModel_Pixels::ni, MyModel_Pixels::nj);
arma::cx_mat          MyModel_Pixels::fft_of_star(MyModel_Pixels::ni, MyModel_Pixels::nj);
std::vector<double>   MyModel_Pixels::x(MyModel_Pixels::nj);
std::vector<double>   MyModel_Pixels::y(MyModel_Pixels::ni);

void MyModel_Pixels::initialise()
{
    for(size_t i=0; i<ni; ++i)
        y[i] = y_max - (i + 0.5)*dy;

    for(size_t j=0; j<nj; ++j)
        x[j] = x_min + (j + 0.5)*dx;

    double rsq;

    double limb_darkening_coefficient = 1.0;

    // Argh column-major order
    // Star image
    double tot = 0.0;
    for(size_t j=0; j<nj; ++j)
    {
        for(size_t i=0; i<ni; ++i)
        {
            rsq = x[j]*x[j] + y[i]*y[i];
            if(rsq < 1.0)
            {
                star(i, j) = 1.0 -
                    limb_darkening_coefficient*(1.0 - sqrt(1.0 - rsq));
            }
            else
                star(i, j) = 0.0;
            tot += star(i, j);
        }
    }
    for(size_t j=0; j<nj; ++j)
        for(size_t i=0; i<ni; ++i)
            star(i, j) /= tot;

    arma::mat star2 = star;
    int m, n;
    for(int i=0; i<(int)ni; i++)
    {
        m = DNest4::mod(i - (int)ni/2, (int)ni);
        for(int j=0; j<(int)nj; j++)
        {
            n = DNest4::mod(j - (int)nj/2, (int)nj);
            star2(m, n) = star(i, j);
        }
    }

    fft_of_star = arma::fft2(star2);
}

void MyModel_Pixels::load_data(const char* filename)
{
    data.load(filename);
}

} // namespace Obscurity