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

#include <cmath>
#include <cstdio>

#include <Eigen/Core>

#include <drake/common/autodiff.h>  // IWYU pragma: keep
#include <drake/common/default_scalars.h>
#include <drake/common/drake_assert.h>
#include <drake/common/unused.h>
#include <drake/solvers/snopt_solver.h>
#include <drake/systems/analysis/simulator.h>
#include <drake/systems/framework/context.h>
#include <drake/systems/framework/continuous_state.h>
#include <drake/systems/framework/vector_system.h>
#include <drake/systems/trajectory_optimization/direct_collocation.h>

namespace shambhala {
namespace systems {

template <typename T>
class SimpleContinuousTimeSystem final : public drake::systems::VectorSystem<T> {
 public:
  DRAKE_NO_COPY_NO_MOVE_NO_ASSIGN(SimpleContinuousTimeSystem);
  SimpleContinuousTimeSystem() : drake::systems::VectorSystem<T>(::drake::systems::SystemTypeTag<SimpleContinuousTimeSystem>{}, 2, 1) { // n_in, n_out
    this->DeclareContinuousState(6); // n_state
  }

  // Scalar-converting copy constructor.  See @ref system_scalar_conversion.
  template <typename U>
  explicit SimpleContinuousTimeSystem(const SimpleContinuousTimeSystem<U>&) : SimpleContinuousTimeSystem<T>() {}

 private:
  // xdot = -x + x^3
  void DoCalcVectorTimeDerivatives(
      const drake::systems::Context<T>& context,
      const Eigen::VectorBlock<const drake::VectorX<T>>& input,
      const Eigen::VectorBlock<const drake::VectorX<T>>& state,
      Eigen::VectorBlock<drake::VectorX<T>>* derivatives) const override {

    // DS Kinetic 60
    const double g = 9.8;
    const double rho = 1.225 * .6; // ISA at 5km

    const double m = 2; // http://www.dskinetic.com/k60.aspx
    const double S = .23;

    // const double Vc2 = m*g/(.5*rho*S);
    // const double lambda = Vc2/g;

    const double cD0 = .005; // https://www.tngtech.com/fileadmin/Public/Images/BigTechday/BTD10/Folien/Folien_SpencerLisenby.pdf
    const double k = .006;

    const auto& speed = state(0);
    const auto& pitch = state(1);
    const auto& yaw   = state(2);
    const auto& z     = state(3);

    const auto& cL    = input(0);
    const auto& roll  = input(1);

    const double windspeed_gradient = .027; // jet stream at 5km: .02

    const auto zd = speed*sin(pitch);
    const auto W = windspeed_gradient*z;
    const auto Wd = windspeed_gradient*zd;

    const auto cD = cD0 + k*cL*cL;
    const auto D = .5*cD*rho*S*speed*speed;
    const auto L = .5*cL*rho*S*speed*speed;

    (*derivatives)(0) = 1/(m                 )*(     -D      - m*g*sin(pitch) + m*Wd*cos(pitch)*sin(yaw));
    (*derivatives)(1) = 1/(m*speed           )*( L*cos(roll) - m*g*cos(pitch) - m*Wd*sin(pitch)*sin(yaw));
    (*derivatives)(2) = 1/(m*speed*cos(pitch))*( L*sin(roll)                  + m*Wd           *cos(yaw));
    (*derivatives)(3) = zd;
    (*derivatives)(4) = speed*cos(pitch)*cos(yaw);
    (*derivatives)(5) = speed*cos(pitch)*sin(yaw) - W;
  }

  // y = x
  void DoCalcVectorOutput(
      const drake::systems::Context<T>& context,
      const Eigen::VectorBlock<const drake::VectorX<T>>& input,
      const Eigen::VectorBlock<const drake::VectorX<T>>& state,
      Eigen::VectorBlock<drake::VectorX<T>>* output) const /*override*/ {
    drake::unused(context, input);
    *output = state;
  }
};

}  // namespace systems
}  // namespace shambhala

DRAKE_DEFINE_CLASS_TEMPLATE_INSTANTIATIONS_ON_DEFAULT_SCALARS(
    class ::shambhala::systems::SimpleContinuousTimeSystem)

int main() {
  // Create the simple system.
  shambhala::systems::SimpleContinuousTimeSystem<double> system;

  auto context = system.CreateDefaultContext();
  const int N = 201;
  const double dt_min = 9./N;
  const double dt_max = 13./N;
  drake::systems::trajectory_optimization::DirectCollocation dircol(
      &system, *context, N, dt_min, dt_max);
  dircol.AddEqualTimeIntervalsConstraints();

  // design limits
  auto tau = 4*acos(0);

  dircol.AddConstraintToAllKnotPoints(dircol.input()(0) >= 0); // 0 <= cL <= 1.2
  dircol.AddConstraintToAllKnotPoints(dircol.input()(0) <= 1.2);

  dircol.AddConstraintToAllKnotPoints(dircol.state()(0) >= 0);
  dircol.AddConstraintToAllKnotPoints(dircol.state()(0) <= 140); // DS Kinetic 60 speed record 140 m/s

  dircol.AddConstraintToAllKnotPoints(dircol.state()(1) >= -tau/2); // pitch looks forward
  dircol.AddConstraintToAllKnotPoints(dircol.state()(1) <=  tau/2);

  dircol.AddConstraintToAllKnotPoints(dircol.state()(2) >= -1.5*tau); // we want one turn in either direction
  dircol.AddConstraintToAllKnotPoints(dircol.state()(2) <=  1.5*tau);  // and allow an extra half turn

  dircol.AddConstraintToAllKnotPoints(dircol.state()(3) >= -500); // max 1km altitude differential (arbitrary)
  dircol.AddConstraintToAllKnotPoints(dircol.state()(3) <=  500);

  // planning target

  // dircol.AddConstraint(dircol.initial_state()(4) == 0); // x_N = x_0
  dircol.AddConstraint(dircol.initial_state()(5) == 0);

  dircol.AddConstraint(dircol.final_state()(0) >= dircol.initial_state()(0));
  dircol.AddConstraint(dircol.final_state()(1) == dircol.initial_state()(1));
  dircol.AddConstraint(dircol.final_state()(2) == dircol.initial_state()(2) + 1*tau);
  dircol.AddConstraint(dircol.final_state()(3) >= dircol.initial_state()(3));

  auto zmin = dircol.NewContinuousVariables(1, "zmin")(0);
  auto zmax = dircol.NewContinuousVariables(1, "zmax")(0);
  dircol.AddConstraintToAllKnotPoints(dircol.state()(3) >= zmin);
  dircol.AddConstraintToAllKnotPoints(dircol.state()(3) <= zmax);
  dircol.AddLinearCost(zmax-zmin);
  dircol.AddLinearConstraint(zmax == -zmin);

  // solver spoonfeeding

  dircol.AddConstraintToAllKnotPoints(dircol.state()(0) >= 1); // nonzero speed!

  dircol.AddConstraintToAllKnotPoints(dircol.input()(1) >= -tau/4); // do not turn the plane upside down
  dircol.AddConstraintToAllKnotPoints(dircol.input()(1) <=  tau/4);

  dircol.SetSolverOption(drake::solvers::SnoptSolver::id(), "Iterations limit", 10000000);
  dircol.SetSolverOption(drake::solvers::SnoptSolver::id(), "Major iterations limit", 10000);

  // solving

  auto result = dircol.Solve();
  if (result != drake::solvers::SolutionResult::kSolutionFound) {
    auto s = drake::solvers::to_string(result);
    printf("solving failed: %s\n", s.c_str());
    return 1;
  }

  // printing

  {
    auto inputs = dircol.ReconstructInputTrajectory();
    auto traj = dircol.ReconstructStateTrajectory();
    auto timestamps = traj.get_segment_times();
    printf("[\n");
    for (size_t i = 0; i < timestamps.size(); i++) {
      auto t = timestamps[i];
      auto V = traj.value(timestamps[i]).coeff(0);
      auto pitch = traj.value(timestamps[i]).coeff(1);
      auto yaw = traj.value(timestamps[i]).coeff(2);
      auto altitude = traj.value(timestamps[i]).coeff(3);
      auto x = traj.value(timestamps[i]).coeff(4);
      auto y = traj.value(timestamps[i]).coeff(5);
      auto cL = inputs.value(timestamps[i]).coeff(0);
      auto roll = inputs.value(timestamps[i]).coeff(1);
      printf("(%03.3f,  %03.3f, %03.3f, %03.3f, %03.3f, %03.3f, %03.3f,  %03.3f, %03.3f), # %3.3f\n", t, V, pitch, yaw, altitude, x, y, cL, roll, 9.8*altitude + .5*V*V);
    }
    printf("]\n");
  }
}


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

/// \file
/// Defines the class Wedge2D.

#pragma once

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

#include "DataStructures/Tensor/Tensor.hpp"
#include "Domain/OrientationMap.hpp"
#include "Utilities/TypeTraits.hpp"

namespace PUP {
class er;
}  // namespace PUP

namespace CoordinateMaps {

/*!
 * \ingroup CoordinateMapsGroup
 *
 * \brief Two dimensional map from the logical square to a wedge, which is used
 * by domains which use disks or annuli.
 *
 * \details The coordinate map is constructed by
 * linearly interpolating between a bulged arc which is circumscribed by a
 * circular arc of radius `radius_outer` and a bulged arc which is
 * circumscribed by a circular arc of radius `radius_inner`. These arcs can be
 * made to be straight or circular, based on the value of `circularity_inner`
 * or `circularity_outer`, which can take on values between 0 (for straight)
 * and 1 (for circular). These arcs extend \f$\pi/2\f$ in angle, and can be
 * oriented along the +/- x or y axis. The choice of using either equiangular
 * or equidistant coordinates along the arcs is specifiable with
 * `with_equiangular_map`. The default logical coordinate that points in the
 * angular direction is \f$\eta\f$. We introduce the auxiliary variable
 * \f$\mathrm{H}\f$ which is a function of \f$\eta\f$. If we are using
 * equiangular coordinates, we have:
 *
 * \f[\mathrm{H}(\eta) = \textrm{tan}(\eta\pi/4)\f]
 *
 * With derivative:
 *
 * \f[\mathrm{H}'(\eta) = \frac{\pi}{4}(1+\mathrm{H}^2)\f]
 *
 * If we are using equidistant coordinates, we have:
 *
 * \f[\mathrm{H}(\eta) = \eta\f]
 *
 * with derivative:
 *
 * <center>\f$\mathrm{H}'(\eta) = 1\f$</center>
 *
 * We also define the variable \f$\rho\f$, given by:
 *
 * \f[\rho = \sqrt{1+\mathrm{H}^2}\f]
 *
 * In terms of the the circularity \f$c\f$ and the radius \f$R\f$,
 * the mapping is:
 *
 * \f[\vec{x}(\xi,\eta) =
 * \frac{1}{2}\left\{(1-\xi)\Big[(1-c_{inner})\frac{R_{inner}}{\sqrt 2}
 * + c_{inner}\frac{R_{inner}}{\rho}\Big] +
 * (1+\xi)\Big[(1-c_{outer})\frac{R_{outer}}{\sqrt 2} +c_{outer}
 * \frac{R_{outer}}{\rho}\Big] \right\}\begin{bmatrix}
 * 1\\
 * \mathrm{H}\\
 * \end{bmatrix}\f]
 *
 * We will define the variables \f$T(\xi)\f$ and \f$A(\xi)\f$, the trapezoid
 * and annulus factors: \f[T(\xi) = T_0 + T_1\xi\f] \f[A(\xi) = A_0 +
 * A_1\xi\f]
 * Where \f{align*}T_0 &= \frac{1}{2} \big\{ (1-c_{outer})R_{outer} +
 * (1-c_{inner})R_{inner}\big\}\\
 * T_1 &= \partial_{\xi} T = \frac{1}{2} \big\{ (1-c_{outer})R_{outer} -
 * (1-c_{inner})R_{inner}\big\}\\
 * A_0 &= \frac{1}{2} \big\{ c_{outer}R_{outer} + c_{inner}R_{inner}\big\}\\
 * A_1 &= \partial_{\xi} A = \frac{1}{2} \big\{ c_{outer}R_{outer} -
 * c_{inner}R_{inner}\big\}\f}
 *
 * The map can then be rewritten as:
 * \f[\vec{x}(\xi,\eta) = \left\{\frac{T(\xi)}{\sqrt 2} +
 * \frac{A(\xi)}{\rho}\right\}\begin{bmatrix}
 * 1\\
 * \mathrm{H}\\
 * \end{bmatrix}\f]
 *
 *
 * The Jacobian is: \f[J =
 * \begin{bmatrix}
 * \frac{T_1}{\sqrt 2} + \frac{A_1}{\rho} &
 * \mathrm{H}\mathrm{H}'\frac{A(\xi)}{\rho^3} \\
 * \mathrm{H}\partial_{\xi}x &
 * \mathrm{H}\partial_{\eta}x + \mathrm{H}'x\\
 * \end{bmatrix}
 * \f]
 *
 * The inverse Jacobian is: \f[J^{-1} =
 * \frac{1}{x}\begin{bmatrix}
 * \frac{1}{\partial_{\xi}x}\Big\{
 * \frac{T(\xi)}{\sqrt 2}+\frac{A(\xi)}{\rho^3}
 * \Big\} & \mathrm{H}\frac{1}{\partial_{\xi}x}\frac{A(\xi)}{\rho^3}\\
 * -\mathrm{H}\mathrm{H}'^{-1} & \mathrm{H}'^{-1}\\
 * \end{bmatrix}
 * \f]
 *
 * For a more detailed discussion, see the
 * documentation for Wedge3D, where the default logical coordinate that points
 * in the radial direction is \f$\zeta\f$ (In Wedge2D the logical coordinate
 * that points in the radial direction is \f$\xi\f$), and one may set either
 * of the two other logical coordinates to zero to obtain an equivalent
 * Wedge2D map.
 */

class Wedge2D {
 public:
  static constexpr size_t dim = 2;

  Wedge2D(double radius_inner, double radius_outer, double circularity_inner,
          double circularity_outer, OrientationMap<2> orientation_of_wedge,
          bool with_equiangular_map) noexcept;

  Wedge2D() = default;
  ~Wedge2D() = default;
  Wedge2D(Wedge2D&&) = default;
  Wedge2D& operator=(Wedge2D&&) = default;
  Wedge2D(const Wedge2D&) = default;
  Wedge2D& operator=(const Wedge2D&) = default;

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

  /// Returns invalid if \f$x<=0\f$ (for a \f$+x\f$-oriented `Wedge2D`).
  boost::optional<std::array<double, 2>> inverse(
      const std::array<double, 2>& target_coords) const noexcept;

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

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

  // clang-tidy: google runtime references
  void pup(PUP::er& p);  // NOLINT
 private:
  friend bool operator==(const Wedge2D& lhs, const Wedge2D& rhs) noexcept;

  double radius_inner_{};
  double radius_outer_{};
  double circularity_inner_{};
  double circularity_outer_{};
  OrientationMap<2> orientation_of_wedge_{};
  bool with_equiangular_map_ = false;
  double scaled_trapezoid_zero_{std::numeric_limits<double>::signaling_NaN()};
  double annulus_zero_{std::numeric_limits<double>::signaling_NaN()};
  double scaled_trapezoid_rate_{std::numeric_limits<double>::signaling_NaN()};
  double annulus_rate_{std::numeric_limits<double>::signaling_NaN()};
};
bool operator!=(const Wedge2D& lhs, const Wedge2D& rhs) noexcept;
}  // namespace CoordinateMaps


#ifndef STAN_MATH_PRIM_SCAL_PROB_WEIBULL_CDF_HPP
#define STAN_MATH_PRIM_SCAL_PROB_WEIBULL_CDF_HPP

#include <stan/math/prim/scal/meta/is_constant_struct.hpp>
#include <stan/math/prim/scal/meta/partials_return_type.hpp>
#include <stan/math/prim/scal/meta/operands_and_partials.hpp>
#include <stan/math/prim/scal/err/check_consistent_sizes.hpp>
#include <stan/math/prim/scal/err/check_finite.hpp>
#include <stan/math/prim/scal/err/check_nonnegative.hpp>
#include <stan/math/prim/scal/err/check_not_nan.hpp>
#include <stan/math/prim/scal/err/check_positive_finite.hpp>
#include <stan/math/prim/scal/fun/multiply_log.hpp>
#include <stan/math/prim/scal/fun/value_of.hpp>
#include <stan/math/prim/scal/meta/length.hpp>
#include <stan/math/prim/scal/fun/constants.hpp>
#include <stan/math/prim/scal/meta/include_summand.hpp>
#include <stan/math/prim/scal/meta/VectorBuilder.hpp>
#include <stan/math/prim/scal/meta/scalar_seq_view.hpp>
#include <boost/random/weibull_distribution.hpp>
#include <boost/random/variate_generator.hpp>
#include <cmath>

namespace stan {
  namespace math {

    /**
     * Returns the Weibull cumulative distribution function for the given
     * location and scale. Given containers of matching sizes, returns the
     * product of probabilities.
     *
     * @tparam T_y type of real parameter
     * @tparam T_shape type of shape parameter
     * @tparam T_scale type of scale paramater
     * @param y real parameter
     * @param alpha shape parameter
     * @param sigma scale parameter
     * @return probability or product of probabilities
     * @throw std::domain_error if y is negative, alpha sigma is nonpositive
     */
    template <typename T_y, typename T_shape, typename T_scale>
    typename return_type<T_y, T_shape, T_scale>::type
    weibull_cdf(const T_y& y, const T_shape& alpha, const T_scale& sigma) {
      typedef typename stan::partials_return_type<T_y, T_shape, T_scale>::type
        T_partials_return;

      static const char* function("weibull_cdf");

      using boost::math::tools::promote_args;
      using std::log;
      using std::exp;

      if (!(stan::length(y) && stan::length(alpha) && stan::length(sigma)))
        return 1.0;

      T_partials_return cdf(1.0);
      check_nonnegative(function, "Random variable", y);
      check_positive_finite(function, "Shape parameter", alpha);
      check_positive_finite(function, "Scale parameter", sigma);

      operands_and_partials<T_y, T_shape, T_scale>
        ops_partials(y, alpha, sigma);

      scalar_seq_view<T_y> y_vec(y);
      scalar_seq_view<T_scale> sigma_vec(sigma);
      scalar_seq_view<T_shape> alpha_vec(alpha);
      size_t N = max_size(y, sigma, alpha);
      for (size_t n = 0; n < N; n++) {
        const T_partials_return y_dbl = value_of(y_vec[n]);
        const T_partials_return sigma_dbl = value_of(sigma_vec[n]);
        const T_partials_return alpha_dbl = value_of(alpha_vec[n]);
        const T_partials_return pow_ = pow(y_dbl / sigma_dbl, alpha_dbl);
        const T_partials_return exp_ = exp(-pow_);
        const T_partials_return cdf_ = 1.0 - exp_;

        cdf *= cdf_;

        const T_partials_return rep_deriv = exp_ * pow_ / cdf_;
        if (!is_constant_struct<T_y>::value)
          ops_partials.edge1_.partials_[n] += rep_deriv * alpha_dbl / y_dbl;
        if (!is_constant_struct<T_shape>::value)
          ops_partials.edge2_.partials_[n]
            += rep_deriv * log(y_dbl / sigma_dbl);
        if (!is_constant_struct<T_scale>::value)
          ops_partials.edge3_.partials_[n] -= rep_deriv * alpha_dbl / sigma_dbl;
      }

      if (!is_constant_struct<T_y>::value) {
        for (size_t n = 0; n < stan::length(y); ++n)
          ops_partials.edge1_.partials_[n] *= cdf;
      }
      if (!is_constant_struct<T_shape>::value) {
        for (size_t n = 0; n < stan::length(alpha); ++n)
          ops_partials.edge2_.partials_[n] *= cdf;
      }
      if (!is_constant_struct<T_scale>::value) {
        for (size_t n = 0; n < stan::length(sigma); ++n)
          ops_partials.edge3_.partials_[n] *= cdf;
      }
      return ops_partials.build(cdf);
    }

  }
}
#endif


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

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

 * See LICENSE for the license information

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

/**
 *  @file   TangentPreintegration.cpp
 *  @author Frank Dellaert
 *  @author Adam Bry
 **/

#include "TangentPreintegration.h"
#include <gtsam/base/numericalDerivative.h>
#include <boost/make_shared.hpp>

using namespace std;

namespace gtsam {

//------------------------------------------------------------------------------
TangentPreintegration::TangentPreintegration(const boost::shared_ptr<Params>& p,
    const Bias& biasHat) :
    PreintegrationBase(p, biasHat) {
  resetIntegration();
}

//------------------------------------------------------------------------------
void TangentPreintegration::resetIntegration() {
  deltaTij_ = 0.0;
  preintegrated_.setZero();
  preintegrated_H_biasAcc_.setZero();
  preintegrated_H_biasOmega_.setZero();
}

//------------------------------------------------------------------------------
bool TangentPreintegration::equals(const TangentPreintegration& other,
    double tol) const {
  return p_->equals(*other.p_, tol) && std::abs(deltaTij_ - other.deltaTij_) < tol
      && biasHat_.equals(other.biasHat_, tol)
      && equal_with_abs_tol(preintegrated_, other.preintegrated_, tol)
      && equal_with_abs_tol(preintegrated_H_biasAcc_,
          other.preintegrated_H_biasAcc_, tol)
      && equal_with_abs_tol(preintegrated_H_biasOmega_,
          other.preintegrated_H_biasOmega_, tol);
}

//------------------------------------------------------------------------------
// See extensive discussion in ImuFactor.lyx
Vector9 TangentPreintegration::UpdatePreintegrated(const Vector3& a_body,
    const Vector3& w_body, double dt, const Vector9& preintegrated,
    OptionalJacobian<9, 9> A, OptionalJacobian<9, 3> B,
    OptionalJacobian<9, 3> C) {
  const auto theta = preintegrated.segment<3>(0);
  const auto position = preintegrated.segment<3>(3);
  const auto velocity = preintegrated.segment<3>(6);

  // This functor allows for saving computation when exponential map and its
  // derivatives are needed at the same location in so<3>
  so3::DexpFunctor local(theta);

  // Calculate exact mean propagation
  Matrix3 w_tangent_H_theta, invH;
  const Vector3 w_tangent = // angular velocity mapped back to tangent space
      local.applyInvDexp(w_body, A ? &w_tangent_H_theta : 0, C ? &invH : 0);
  const Rot3 R(local.expmap());  // nRb: rotation of body in nav frame
  const Vector3 a_nav = R * a_body;
  const double dt22 = 0.5 * dt * dt;

  Vector9 preintegratedPlus;
  preintegratedPlus <<                          // new preintegrated vector:
      theta + w_tangent * dt,                   // theta
      position + velocity * dt + a_nav * dt22,  // position
      velocity + a_nav * dt;                    // velocity

  if (A) {
    // Exact derivative of R*a with respect to theta:
    const Matrix3 a_nav_H_theta = R.matrix() * skewSymmetric(-a_body) * local.dexp();

    A->setIdentity();
    A->block<3, 3>(0, 0).noalias() += w_tangent_H_theta * dt;  // theta
    A->block<3, 3>(3, 0) = a_nav_H_theta * dt22;  // position wrpt theta...
    A->block<3, 3>(3, 6) = I_3x3 * dt;            // .. and velocity
    A->block<3, 3>(6, 0) = a_nav_H_theta * dt;    // velocity wrpt theta
  }
  if (B) {
    B->block<3, 3>(0, 0) = Z_3x3;
    B->block<3, 3>(3, 0) = R.matrix() * dt22;
    B->block<3, 3>(6, 0) = R.matrix() * dt;
  }
  if (C) {
    C->block<3, 3>(0, 0) = invH * dt;
    C->block<3, 3>(3, 0) = Z_3x3;
    C->block<3, 3>(6, 0) = Z_3x3;
  }

  return preintegratedPlus;
}

//------------------------------------------------------------------------------
void TangentPreintegration::update(const Vector3& measuredAcc,
    const Vector3& measuredOmega, const double dt, Matrix9* A, Matrix93* B,
    Matrix93* C) {

  // std::cout << "TangentPreintegration::update" << std::endl;
  // Correct for bias in the sensor frame
  Vector3 acc = biasHat_.correctAccelerometer(measuredAcc);
  Vector3 omega = biasHat_.correctGyroscope(measuredOmega);

  // Possibly correct for sensor pose by converting to body frame
  Matrix3 D_correctedAcc_acc, D_correctedAcc_omega, D_correctedOmega_omega;
  if (p().body_P_sensor)
    boost::tie(acc, omega) = correctMeasurementsBySensorPose(acc, omega,
        D_correctedAcc_acc, D_correctedAcc_omega, D_correctedOmega_omega);

  // Do update
  deltaTij_ += dt;
  preintegrated_ = UpdatePreintegrated(acc, omega, dt, preintegrated_, A, B, C);

  if (p().body_P_sensor) {
    // More complicated derivatives in case of non-trivial sensor pose
    *C *= D_correctedOmega_omega;
    if (!p().body_P_sensor->translation().isZero())
      *C += *B * D_correctedAcc_omega;
    *B *= D_correctedAcc_acc; // NOTE(frank): needs to be last
  }

  // new_H_biasAcc = new_H_old * old_H_biasAcc + new_H_acc * acc_H_biasAcc
  // where acc_H_biasAcc = -I_3x3, hence
  // new_H_biasAcc = new_H_old * old_H_biasAcc - new_H_acc
  preintegrated_H_biasAcc_ = (*A) * preintegrated_H_biasAcc_ - (*B);

  // new_H_biasOmega = new_H_old * old_H_biasOmega + new_H_omega * omega_H_biasOmega
  // where omega_H_biasOmega = -I_3x3, hence
  // new_H_biasOmega = new_H_old * old_H_biasOmega - new_H_omega
  preintegrated_H_biasOmega_ = (*A) * preintegrated_H_biasOmega_ - (*C);
}

//------------------------------------------------------------------------------
Vector9 TangentPreintegration::biasCorrectedDelta(
    const imuBias::ConstantBias& bias_i, OptionalJacobian<9, 6> H) const {
  // We correct for a change between bias_i and the biasHat_ used to integrate
  // This is a simple linear correction with obvious derivatives
  const imuBias::ConstantBias biasIncr = bias_i - biasHat_;
  const Vector9 biasCorrected = preintegrated()
      + preintegrated_H_biasAcc_ * biasIncr.accelerometer()
      + preintegrated_H_biasOmega_ * biasIncr.gyroscope();

  if (H) {
    (*H) << preintegrated_H_biasAcc_, preintegrated_H_biasOmega_;
  }
  return biasCorrected;
}

//------------------------------------------------------------------------------
// sugar for derivative blocks
#define D_R_R(H) (H)->block<3,3>(0,0)
#define D_R_t(H) (H)->block<3,3>(0,3)
#define D_R_v(H) (H)->block<3,3>(0,6)
#define D_t_R(H) (H)->block<3,3>(3,0)
#define D_t_t(H) (H)->block<3,3>(3,3)
#define D_t_v(H) (H)->block<3,3>(3,6)
#define D_v_R(H) (H)->block<3,3>(6,0)
#define D_v_t(H) (H)->block<3,3>(6,3)
#define D_v_v(H) (H)->block<3,3>(6,6)

//------------------------------------------------------------------------------
Vector9 TangentPreintegration::Compose(const Vector9& zeta01,
    const Vector9& zeta12, double deltaT12, OptionalJacobian<9, 9> H1,
    OptionalJacobian<9, 9> H2) {
  const auto t01 = zeta01.segment<3>(0);
  const auto p01 = zeta01.segment<3>(3);
  const auto v01 = zeta01.segment<3>(6);

  const auto t12 = zeta12.segment<3>(0);
  const auto p12 = zeta12.segment<3>(3);
  const auto v12 = zeta12.segment<3>(6);

  Matrix3 R01_H_t01, R12_H_t12;
  const Rot3 R01 = Rot3::Expmap(t01, R01_H_t01);
  const Rot3 R12 = Rot3::Expmap(t12, R12_H_t12);

  Matrix3 R02_H_R01, R02_H_R12; // NOTE(frank): R02_H_R12 == Identity
  const Rot3 R02 = R01.compose(R12, R02_H_R01, R02_H_R12);

  Matrix3 t02_H_R02;
  Vector9 zeta02;
  const Matrix3 R = R01.matrix();
  zeta02 << Rot3::Logmap(R02, t02_H_R02), // theta
  p01 + v01 * deltaT12 + R * p12, // position
  v01 + R * v12; // velocity

  if (H1) {
    H1->setIdentity();
    D_R_R(H1) = t02_H_R02 * R02_H_R01 * R01_H_t01;
    D_t_R(H1) = R * skewSymmetric(-p12) * R01_H_t01;
    D_t_v(H1) = I_3x3 * deltaT12;
    D_v_R(H1) = R * skewSymmetric(-v12) * R01_H_t01;
  }

  if (H2) {
    H2->setZero();
    D_R_R(H2) = t02_H_R02 * R02_H_R12 * R12_H_t12;
    D_t_t(H2) = R;
    D_v_v(H2) = R;
  }

  return zeta02;
}

//------------------------------------------------------------------------------
void TangentPreintegration::mergeWith(const TangentPreintegration& pim12,
    Matrix9* H1, Matrix9* H2) {
  if (!matchesParamsWith(pim12)) {
    throw std::domain_error(
        "Cannot merge pre-integrated measurements with different params");
  }

  if (params()->body_P_sensor) {
    throw std::domain_error(
        "Cannot merge pre-integrated measurements with sensor pose yet");
  }

  const double t01 = deltaTij();
  const double t12 = pim12.deltaTij();
  deltaTij_ = t01 + t12;

  const Vector9 zeta01 = preintegrated();
  Vector9 zeta12 = pim12.preintegrated(); // will be modified.

  const imuBias::ConstantBias bias_incr_for_12 = biasHat() - pim12.biasHat();
  zeta12 += pim12.preintegrated_H_biasOmega_ * bias_incr_for_12.gyroscope()
      + pim12.preintegrated_H_biasAcc_ * bias_incr_for_12.accelerometer();

  preintegrated_ = TangentPreintegration::Compose(zeta01, zeta12, t12, H1, H2);

  preintegrated_H_biasAcc_ = (*H1) * preintegrated_H_biasAcc_
      + (*H2) * pim12.preintegrated_H_biasAcc_;

  preintegrated_H_biasOmega_ = (*H1) * preintegrated_H_biasOmega_
      + (*H2) * pim12.preintegrated_H_biasOmega_;
}

//------------------------------------------------------------------------------

}// namespace gtsam


#include "viscosity3d.h"
#include "array3.h"
//#include "sparse/sparsematrix.h"
//#include "sparse/cgsolver.h"
#include "pcgsolver/pcg_solver.h"
#include <fstream>
#include <cmath>
//#include "wallclocktime.h"

//#include <Eigen/Sparse>

//include ViennaCL stuff for multithreaded (OpenMP) iterative solvers
/*
#define VIENNACL_WITH_EIGEN 1

#include "viennacl/linalg/ilu.hpp"
#include "viennacl/linalg/cg.hpp"
#include "viennacl/linalg/bicgstab.hpp"
#include "viennacl/linalg/gmres.hpp"
#include "viennacl/linalg/norm_2.hpp"
#include "viennacl/linalg/prod.hpp"
*/

static int u_ind(int i, int j, int k, int nx, int ny) {
   return i + j*(nx + 1) + k*(nx + 1)*ny;
}

static int v_ind(int i, int j, int k, int nx, int ny, int nz){
   return i + j*nx + k*nx*(ny + 1) + (nx + 1)*ny*nz;
}

static int w_ind(int i, int j, int k, int nx, int ny, int nz){
   return i + j*nx + k*nx*ny + (nx + 1)*ny*nz + nx*(ny + 1)*nz;
}

Array3c u_state;//(nx+1,ny,nz,0);
Array3c v_state;//(nx,ny+1,nz,0);
Array3c w_state;//(nx,ny,nz+1,0);

SparseMatrixd matrix;//(dim,dim,15);
SparseMatrixd matrix2;//(dim,dim
std::vector<double> rhs;//(dim);
std::vector<double> soln;//(dim);

void advance_viscosity_implicit_weighted(Array3f& u, Array3f& v, Array3f& w,
   const Array3f& vol_u, const Array3f& vol_v, const Array3f& vol_w,
   const Array3f& vol_c, const Array3f& vol_ex, const Array3f& vol_ey, const Array3f& vol_ez,
   const Array3f& solid_phi,
   const Array3f& viscosity, float dt, float dx) {
   float over_dx = 1.0f / dx;
   int nx = solid_phi.ni;
   int ny = solid_phi.nj;
   int nz = solid_phi.nk;
   printf("Creating state arrays.\n");
   std::cout << "Phi-size:" << nx << " " << ny << " " << nz << std::endl;
   int dim = (nx + 1)*ny*nz + nx*(ny + 1)*nz + nx*ny*(nz + 1);
   if (u_state.ni != u.ni) {
      printf("Creating matrices and vectors.\n");
      u_state.resize(nx + 1, ny, nz);
      v_state.resize(nx, ny + 1, nz);
      w_state.resize(nx, ny, nz + 1);
      matrix.resize(dim);
      matrix2.resize(dim);
      rhs.resize(dim);
      soln.resize(dim);
      printf("Done that for good.");
   }

   u_state.assign(0);
   v_state.assign(0);
   w_state.assign(0);
   matrix.zero();
   //matrix2.zero();
   rhs.assign(dim, 0);
   soln.assign(dim, 0);

   const int SOLID = 3;
   const int FLUID = 2;
   //const int AIR = 1;

   //check if interpolated velocity positions are inside solid
   for (int k = 0; k < nz; ++k) for (int j = 0; j < ny; ++j) for (int i = 0; i < nx + 1; ++i) {
      if (i - 1 < 0 || i >= nx || solid_phi(i - 1, j, k) + solid_phi(i, j, k) <= 0)
         u_state(i, j, k) = SOLID;
      else
         u_state(i, j, k) = FLUID;
   }

   for (int k = 0; k < nz; ++k) for (int j = 0; j < ny + 1; ++j) for (int i = 0; i < nx; ++i) {
      if (j - 1 < 0 || j >= ny || solid_phi(i, j - 1, k) + solid_phi(i, j, k) <= 0)
         v_state(i, j, k) = SOLID;
      else
         v_state(i, j, k) = FLUID;
   }

   for (int k = 0; k < nz + 1; ++k) for (int j = 0; j < ny; ++j) for (int i = 0; i < nx; ++i) {
      if (k - 1 < 0 || k >= nz || solid_phi(i, j, k - 1) + solid_phi(i, j, k) <= 0)
         w_state(i, j, k) = SOLID;
      else
         w_state(i, j, k) = FLUID;
   }

   float factor = dt*sqr(over_dx);
   //u-terms
   //2u_xx+ v_xy +uyy + u_zz + w_xz
   printf("Building u-components.\n");
   for (int k = 1; k < nz; ++k) for (int j = 1; j < ny; ++j) for (int i = 1; i < nx; ++i) {

      if (u_state(i, j, k) == FLUID) {
         int index = u_ind(i, j, k, nx, ny);

         rhs[index] = vol_u(i, j, k)*u(i, j, k);
         matrix.set_element(index, index, vol_u(i, j, k));

         float visc_right = viscosity(i, j, k);
         float visc_left = viscosity(i - 1, j, k);
         float vol_right = vol_c(i, j, k);
         float vol_left = vol_c(i - 1, j, k);

         float visc_top = 0.25f*(viscosity(i - 1, j + 1, k) + viscosity(i - 1, j, k) + viscosity(i, j + 1, k) + viscosity(i, j, k));
         float visc_bottom = 0.25f*(viscosity(i - 1, j, k) + viscosity(i - 1, j - 1, k) + viscosity(i, j, k) + viscosity(i, j - 1, k));
         float vol_top = vol_ez(i, j + 1, k);
         float vol_bottom = vol_ez(i, j, k);

         float visc_front = 0.25f*(viscosity(i - 1, j, k + 1) + viscosity(i - 1, j, k) + viscosity(i, j, k + 1) + viscosity(i, j, k));
         float visc_back = 0.25f*(viscosity(i - 1, j, k) + viscosity(i - 1, j, k - 1) + viscosity(i, j, k) + viscosity(i, j, k - 1));
         float vol_front = vol_ey(i, j, k + 1);
         float vol_back = vol_ey(i, j, k);

         //u_x_right
         matrix.add_to_element(index, index, 2 * factor*visc_right*vol_right);
         if (u_state(i + 1, j, k) == FLUID) {
            matrix.add_to_element(index, u_ind(i + 1, j, k, nx, ny), -2 * factor*visc_right*vol_right);
         }
         else if (u_state(i + 1, j, k) == SOLID)
            rhs[index] -= -2 * factor*visc_right*vol_right*u(i + 1, j, k);

         //u_x_left
         matrix.add_to_element(index, index, 2 * factor*visc_left*vol_left);
         if (u_state(i - 1, j, k) == FLUID)
            matrix.add_to_element(index, u_ind(i - 1, j, k, nx, ny), -2 * factor*visc_left*vol_left);
         else if (u_state(i - 1, j, k) == SOLID)
            rhs[index] -= -2 * factor*visc_left*vol_left*u(i - 1, j, k);

         //u_y_top
         matrix.add_to_element(index, index, +factor*visc_top*vol_top);
         if (u_state(i, j + 1, k) == FLUID)
            matrix.add_to_element(index, u_ind(i, j + 1, k, nx, ny), -factor*visc_top*vol_top);
         else if (u_state(i, j + 1, k) == SOLID)
            rhs[index] -= -u(i, j + 1, k)*factor*visc_top*vol_top;

         //u_y_bottom
         matrix.add_to_element(index, index, +factor*visc_bottom*vol_bottom);
         if (u_state(i, j - 1, k) == FLUID)
            matrix.add_to_element(index, u_ind(i, j - 1, k, nx, ny), -factor*visc_bottom*vol_bottom);
         else if (u_state(i, j - 1, k) == SOLID)
            rhs[index] -= -u(i, j - 1, k)*factor*visc_bottom*vol_bottom;

         //u_z_front
         matrix.add_to_element(index, index, +factor*visc_front*vol_front);
         if (u_state(i, j, k + 1) == FLUID)
            matrix.add_to_element(index, u_ind(i, j, k + 1, nx, ny), -factor*visc_front*vol_front);
         else if (u_state(i, j, k + 1) == SOLID)
            rhs[index] -= -u(i, j, k + 1)*factor*visc_front*vol_front;

         //u_z_back
         matrix.add_to_element(index, index, +factor*visc_back*vol_back);
         if (u_state(i, j, k - 1) == FLUID)
            matrix.add_to_element(index, u_ind(i, j, k - 1, nx, ny), -factor*visc_back*vol_back);
         else if (u_state(i, j, k - 1) == SOLID)
            rhs[index] -= -u(i, j, k - 1)*factor*visc_back*vol_back;

         //v_x_top
         if (v_state(i, j + 1, k) == FLUID)
            matrix.add_to_element(index, v_ind(i, j + 1, k, nx, ny, nz), -factor*visc_top*vol_top);
         else if (v_state(i, j + 1, k) == SOLID)
            rhs[index] -= -v(i, j + 1, k)*factor*visc_top*vol_top;

         if (v_state(i - 1, j + 1, k) == FLUID)
            matrix.add_to_element(index, v_ind(i - 1, j + 1, k, nx, ny, nz), factor*visc_top*vol_top);
         else if (v_state(i - 1, j + 1, k) == SOLID)
            rhs[index] -= v(i - 1, j + 1, k)*factor*visc_top*vol_top;

         //v_x_bottom
         if (v_state(i, j, k) == FLUID)
            matrix.add_to_element(index, v_ind(i, j, k, nx, ny, nz), +factor*visc_bottom*vol_bottom);
         else if (v_state(i, j, k) == SOLID)
            rhs[index] -= v(i, j, k)*factor*visc_bottom*vol_bottom;

         if (v_state(i - 1, j, k) == FLUID)
            matrix.add_to_element(index, v_ind(i - 1, j, k, nx, ny, nz), -factor*visc_bottom*vol_bottom);
         else if (v_state(i - 1, j, k) == SOLID)
            rhs[index] -= -v(i - 1, j, k)*factor*visc_bottom*vol_bottom;

         //w_x_front
         if (w_state(i, j, k + 1) == FLUID)
            matrix.add_to_element(index, w_ind(i, j, k + 1, nx, ny, nz), -factor*visc_front*vol_front);
         else if (w_state(i, j, k + 1) == SOLID)
            rhs[index] -= -w(i, j, k + 1)*factor*visc_front*vol_front;

         if (w_state(i - 1, j, k + 1) == FLUID)
            matrix.add_to_element(index, w_ind(i - 1, j, k + 1, nx, ny, nz), factor*visc_front*vol_front);
         else if (w_state(i - 1, j, k + 1) == SOLID)
            rhs[index] -= w(i - 1, j, k + 1)*factor*visc_front*vol_front;

         //w_x_back
         if (w_state(i, j, k) == FLUID)
            matrix.add_to_element(index, w_ind(i, j, k, nx, ny, nz), +factor*visc_back*vol_back);
         else if (w_state(i, j, k) == SOLID)
            rhs[index] -= w(i, j, k)*factor*visc_back*vol_back;

         if (w_state(i - 1, j, k) == FLUID)
            matrix.add_to_element(index, w_ind(i - 1, j, k, nx, ny, nz), -factor*visc_back*vol_back);
         else if (w_state(i - 1, j, k) == SOLID)
            rhs[index] -= -w(i - 1, j, k)*factor*visc_back*vol_back;
      }
   }

   //v-terms
   //vxx + 2vyy + vzz + u_yx + w_yz
   printf("Building v-components.\n");
   for (int k = 1; k < nz; ++k) for (int j = 1; j < ny; ++j) for (int i = 1; i < nx; ++i) {
      if (v_state(i, j, k) == FLUID) {
         int index = v_ind(i, j, k, nx, ny, nz);

         rhs[index] = vol_v(i, j, k)*v(i, j, k);
         matrix.set_element(index, index, vol_v(i, j, k));

         float visc_right = 0.25f*(viscosity(i, j - 1, k) + viscosity(i + 1, j - 1, k) + viscosity(i, j, k) + viscosity(i + 1, j, k));
         float visc_left = 0.25f*(viscosity(i, j - 1, k) + viscosity(i - 1, j - 1, k) + viscosity(i, j, k) + viscosity(i - 1, j, k));
         float vol_right = vol_ez(i + 1, j, k);
         float vol_left = vol_ez(i, j, k);

         float visc_top = viscosity(i, j, k);
         float visc_bottom = viscosity(i, j - 1, k);
         float vol_top = vol_c(i, j, k);
         float vol_bottom = vol_c(i, j - 1, k);

         float visc_front = 0.25f*(viscosity(i, j - 1, k) + viscosity(i, j - 1, k + 1) + viscosity(i, j, k) + viscosity(i, j, k + 1));
         float visc_back = 0.25f*(viscosity(i, j - 1, k) + viscosity(i, j - 1, k - 1) + viscosity(i, j, k) + viscosity(i, j, k - 1));
         float vol_front = vol_ex(i, j, k + 1);
         float vol_back = vol_ex(i, j, k);

         //v_x_right
         matrix.add_to_element(index, index, +factor*visc_right*vol_right);
         if (v_state(i + 1, j, k) == FLUID)
            matrix.add_to_element(index, v_ind(i + 1, j, k, nx, ny, nz), -factor*visc_right*vol_right);
         else if (v_state(i + 1, j, k) == SOLID)
            rhs[index] -= -v(i + 1, j, k)*factor*visc_right*vol_right;

         //v_x_left
         matrix.add_to_element(index, index, +factor*visc_left*vol_left);
         if (v_state(i - 1, j, k) == FLUID)
            matrix.add_to_element(index, v_ind(i - 1, j, k, nx, ny, nz), -factor*visc_left*vol_left);
         else if (v_state(i - 1, j, k) == SOLID)
            rhs[index] -= -v(i - 1, j, k)*factor*visc_left*vol_left;

         //vy_top
         matrix.add_to_element(index, index, +2 * factor*visc_top*vol_top);
         if (v_state(i, j + 1, k) == FLUID)
            matrix.add_to_element(index, v_ind(i, j + 1, k, nx, ny, nz), -2 * factor*visc_top*vol_top);
         else if (v_state(i, j + 1, k) == SOLID)
            rhs[index] -= -2 * factor*visc_top*vol_top*v(i, j + 1, k);

         //vy_bottom
         matrix.add_to_element(index, index, +2 * factor*visc_bottom*vol_bottom);
         if (v_state(i, j - 1, k) == FLUID)
            matrix.add_to_element(index, v_ind(i, j - 1, k, nx, ny, nz), -2 * factor*visc_bottom*vol_bottom);
         else if (v_state(i, j - 1, k) == SOLID)
            rhs[index] -= -2 * factor*visc_bottom*vol_bottom*v(i, j - 1, k);

         //v_z_front
         matrix.add_to_element(index, index, +factor*visc_front*vol_front);
         if (v_state(i, j, k + 1) == FLUID)
            matrix.add_to_element(index, v_ind(i, j, k + 1, nx, ny, nz), -factor*visc_front*vol_front);
         else if (v_state(i, j, k + 1) == SOLID)
            rhs[index] -= -v(i, j, k + 1)*factor*visc_front*vol_front;

         //v_z_back
         matrix.add_to_element(index, index, +factor*visc_back*vol_back);
         if (v_state(i, j, k - 1) == FLUID)
            matrix.add_to_element(index, v_ind(i, j, k - 1, nx, ny, nz), -factor*visc_back*vol_back);
         else if (v_state(i, j, k - 1) == SOLID)
            rhs[index] -= -v(i, j, k - 1)*factor*visc_back*vol_back;

         //u_y_right
         if (u_state(i + 1, j, k) == FLUID)
            matrix.add_to_element(index, u_ind(i + 1, j, k, nx, ny), -factor*visc_right*vol_right);
         else if (u_state(i + 1, j, k) == SOLID)
            rhs[index] -= -u(i + 1, j, k)*factor*visc_right*vol_right;

         if (u_state(i + 1, j - 1, k) == FLUID)
            matrix.add_to_element(index, u_ind(i + 1, j - 1, k, nx, ny), factor*visc_right*vol_right);
         else if (u_state(i + 1, j - 1, k) == SOLID)
            rhs[index] -= u(i + 1, j - 1, k)*factor*visc_right*vol_right;

         //u_y_left
         if (u_state(i, j, k) == FLUID)
            matrix.add_to_element(index, u_ind(i, j, k, nx, ny), factor*visc_left*vol_left);
         else if (u_state(i, j, k) == SOLID)
            rhs[index] -= u(i, j, k)*factor*visc_left*vol_left;

         if (u_state(i, j - 1, k) == FLUID)
            matrix.add_to_element(index, u_ind(i, j - 1, k, nx, ny), -factor*visc_left*vol_left);
         else if (u_state(i, j - 1, k) == SOLID)
            rhs[index] -= -u(i, j - 1, k)*factor*visc_left*vol_left;

         //w_y_front
         if (w_state(i, j, k + 1) == FLUID)
            matrix.add_to_element(index, w_ind(i, j, k + 1, nx, ny, nz), -factor*visc_front*vol_front);
         else if (w_state(i, j, k + 1) == SOLID)
            rhs[index] -= -w(i, j, k + 1)*factor*visc_front*vol_front;

         if (w_state(i, j - 1, k + 1) == FLUID)
            matrix.add_to_element(index, w_ind(i, j - 1, k + 1, nx, ny, nz), factor*visc_front*vol_front);
         else if (w_state(i, j - 1, k + 1) == SOLID)
            rhs[index] -= w(i, j - 1, k + 1)*factor*visc_front*vol_front;

         //w_y_back
         if (w_state(i, j, k) == FLUID)
            matrix.add_to_element(index, w_ind(i, j, k, nx, ny, nz), factor*visc_back*vol_back);
         else if (w_state(i, j, k) == SOLID)
            rhs[index] -= w(i, j, k)*factor*visc_back*vol_back;

         if (w_state(i, j - 1, k) == FLUID)
            matrix.add_to_element(index, w_ind(i, j - 1, k, nx, ny, nz), -factor*visc_back*vol_back);
         else if (w_state(i, j - 1, k) == SOLID)
            rhs[index] -= -w(i, j - 1, k)*factor*visc_back*vol_back;
      }
   }

   //w-terms
   //wxx+ wyy+ 2wzz + u_zx + v_zy
   printf("Building w-components.\n");
   for (int k = 1; k < nz; ++k) for (int j = 1; j < ny; ++j) for (int i = 1; i < nx; ++i) {
      if (w_state(i, j, k) == FLUID) {
         int index = w_ind(i, j, k, nx, ny, nz);
         rhs[index] = vol_w(i, j, k)*w(i, j, k);
         matrix.set_element(index, index, vol_w(i, j, k));

         float visc_right = 0.25f*(viscosity(i, j, k) + viscosity(i, j, k - 1) + viscosity(i + 1, j, k) + viscosity(i + 1, j, k - 1));
         float visc_left = 0.25f*(viscosity(i, j, k) + viscosity(i, j, k - 1) + viscosity(i - 1, j, k) + viscosity(i - 1, j, k - 1));
         float vol_right = vol_ey(i + 1, j, k);
         float vol_left = vol_ey(i, j, k);

         float visc_top = 0.25f*(viscosity(i, j, k) + viscosity(i, j, k - 1) + viscosity(i, j + 1, k) + viscosity(i, j + 1, k - 1));;
         float visc_bottom = 0.25f*(viscosity(i, j, k) + viscosity(i, j, k - 1) + viscosity(i, j - 1, k) + viscosity(i, j - 1, k - 1));;
         float vol_top = vol_ex(i, j + 1, k);
         float vol_bottom = vol_ex(i, j, k);

         float visc_front = viscosity(i, j, k);
         float visc_back = viscosity(i, j, k - 1);
         float vol_front = vol_c(i, j, k);
         float vol_back = vol_c(i, j, k - 1);

         //w_x_right
         matrix.add_to_element(index, index, +factor*visc_right*vol_right);
         if (w_state(i + 1, j, k) == FLUID)
            matrix.add_to_element(index, w_ind(i + 1, j, k, nx, ny, nz), -factor*visc_right*vol_right);
         else if (w_state(i + 1, j, k) == SOLID)
            rhs[index] -= -factor*visc_right*vol_right*w(i + 1, j, k);

         //w_x_left
         matrix.add_to_element(index, index, factor*visc_left*vol_left);
         if (w_state(i - 1, j, k) == FLUID)
            matrix.add_to_element(index, w_ind(i - 1, j, k, nx, ny, nz), -factor*visc_left*vol_left);
         else if (w_state(i - 1, j, k) == SOLID)
            rhs[index] -= -factor*visc_left*vol_left*w(i - 1, j, k);

         //w_y_top
         matrix.add_to_element(index, index, +factor*visc_top*vol_top);
         if (w_state(i, j + 1, k) == FLUID)
            matrix.add_to_element(index, w_ind(i, j + 1, k, nx, ny, nz), -factor*visc_top*vol_top);
         else if (w_state(i, j + 1, k) == SOLID)
            rhs[index] -= -factor*visc_top*vol_top*w(i, j + 1, k);

         //w_y_bottom
         matrix.add_to_element(index, index, factor*visc_bottom*vol_bottom);
         if (w_state(i, j - 1, k) == FLUID)
            matrix.add_to_element(index, w_ind(i, j - 1, k, nx, ny, nz), -factor*visc_bottom*vol_bottom);
         else if (w_state(i, j - 1, k) == SOLID)
            rhs[index] -= -factor*visc_bottom*vol_bottom*w(i, j - 1, k);

         //w_z_front
         matrix.add_to_element(index, index, +2 * factor*visc_front*vol_front);
         if (w_state(i, j, k + 1) == FLUID)
            matrix.add_to_element(index, w_ind(i, j, k + 1, nx, ny, nz), -2 * factor*visc_front*vol_front);
         else if (w_state(i, j, k + 1) == SOLID)
            rhs[index] -= -2 * factor*visc_front*vol_front*w(i, j, k + 1);

         //w_z_back
         matrix.add_to_element(index, index, +2 * factor*visc_back*vol_back);
         if (w_state(i, j, k - 1) == FLUID)
            matrix.add_to_element(index, w_ind(i, j, k - 1, nx, ny, nz), -2 * factor*visc_back*vol_back);
         else if (w_state(i, j, k - 1) == SOLID)
            rhs[index] -= -2 * factor*visc_back*vol_back*w(i, j, k - 1);

         //u_z_right
         if (u_state(i + 1, j, k) == FLUID)
            matrix.add_to_element(index, u_ind(i + 1, j, k, nx, ny), -factor*visc_right*vol_right);
         else if (u_state(i + 1, j, k) == SOLID)
            rhs[index] -= -u(i + 1, j, k)*factor*visc_right*vol_right;

         if (u_state(i + 1, j, k - 1) == FLUID)
            matrix.add_to_element(index, u_ind(i + 1, j, k - 1, nx, ny), factor*visc_right*vol_right);
         else if (u_state(i + 1, j, k - 1) == SOLID)
            rhs[index] -= u(i + 1, j, k - 1)*factor*visc_right*vol_right;

         //u_z_left
         if (u_state(i, j, k) == FLUID)
            matrix.add_to_element(index, u_ind(i, j, k, nx, ny), factor*visc_left*vol_left);
         else if (u_state(i, j, k) == SOLID)
            rhs[index] -= u(i, j, k)*factor*visc_left*vol_left;

         if (u_state(i, j, k - 1) == FLUID)
            matrix.add_to_element(index, u_ind(i, j, k - 1, nx, ny), -factor*visc_left*vol_left);
         else if (u_state(i, j, k - 1) == SOLID)
            rhs[index] -= -u(i, j, k - 1)*factor*visc_left*vol_left;

         //v_z_top
         if (v_state(i, j + 1, k) == FLUID)
            matrix.add_to_element(index, v_ind(i, j + 1, k, nx, ny, nz), -factor*visc_top*vol_top);
         else if (v_state(i, j + 1, k) == SOLID)
            rhs[index] -= -v(i, j + 1, k)*factor*visc_top*vol_top;

         if (v_state(i, j + 1, k - 1) == FLUID)
            matrix.add_to_element(index, v_ind(i, j + 1, k - 1, nx, ny, nz), factor*visc_top*vol_top);
         else if (v_state(i, j + 1, k - 1) == SOLID)
            rhs[index] -= v(i, j + 1, k - 1)*factor*visc_top*vol_top;

         //v_z_bottom
         if (v_state(i, j, k) == FLUID)
            matrix.add_to_element(index, v_ind(i, j, k, nx, ny, nz), +factor*visc_bottom*vol_bottom);
         else if (v_state(i, j, k) == SOLID)
            rhs[index] -= v(i, j, k)*factor*visc_bottom*vol_bottom;

         if (v_state(i, j, k - 1) == FLUID)
            matrix.add_to_element(index, v_ind(i, j, k - 1, nx, ny, nz), -factor*visc_bottom*vol_bottom);
         else if (v_state(i, j, k - 1) == SOLID)
            rhs[index] -= -v(i, j, k - 1)*factor*visc_bottom*vol_bottom;

      }
   }

   ////strip out near zero entries to speed this thing up!
   //printf("Stripping out near-zeros.\n");
   //SparseMatrixd matrix2(matrix.m,matrix.n,15);

   //std::ofstream outfile("matrix.m");
   //matrix.write_matlab(outfile, "A_mat");
   //outfile.close();

   /*
   std::vector<Eigen::Triplet<double> > entries;
   Eigen::SparseMatrix<double> eigen_system(dim, dim);
   Eigen::VectorXd eigen_rhs(dim);

   //copy to eigen matrix
   for (unsigned int row = 0; row < matrix.n; ++row){
   for (unsigned int col = 0; col < matrix.index[row].size(); ++col) {
   int index = matrix.index[row][col];
   double val = matrix(row, index);
   //if(std::abs(val) > 1e-10)
   entries.push_back(Eigen::Triplet<double>(row, index, val));
   //double val = matrix(row,index);
   //if(std::abs(val) > 1e-10)
   //   matrix2.set_element(row,index, val);
   }
   if (matrix.index[row].size() == 0) {
   entries.push_back(Eigen::Triplet<double>(row, row, 1));
   }
   }
   eigen_system.setFromTriplets(entries.begin(), entries.end());

   double solve_start_time, solve_end_time;
   try {
   //try solving with ViennaCL, for a nice multithreaded change of pace.
   std::cout << "Running ViennaCL BiCGStab Viscosity solve." << std::endl;
   //solve_start_time = get_time_in_seconds();

   //copy the data

   viennacl::vector<double> vcl_rhs(rhs.size()), vcl_soln(rhs.size());
   viennacl::compressed_matrix<double> vcl_sparsemat(eigen_system.rows(), eigen_system.cols());

   viennacl::copy(rhs, vcl_rhs);
   viennacl::copy(eigen_system, vcl_sparsemat);
   std::cout << "..Preconditioner\n";
   viennacl::linalg::ilut_tag ilut_config(20U, 0.01, false);
   viennacl::linalg::ilut_precond< viennacl::compressed_matrix<double> > vcl_ilut(vcl_sparsemat, ilut_config);

   std::cout << "..Running solve\n";
   viennacl::linalg::bicgstab_tag custom_tag(1e-10, 1000);
   vcl_soln = viennacl::linalg::solve(vcl_sparsemat, vcl_rhs, custom_tag, vcl_ilut);

   //solve_end_time = get_time_in_seconds();

   viennacl::vector<double> residual(rhs.size());
   residual = viennacl::linalg::prod(vcl_sparsemat, vcl_soln) - vcl_rhs;
   std::cout << "Residual: " << viennacl::linalg::norm_2(residual) << std::endl;

   viennacl::copy(vcl_soln, soln);
   }
   catch (...) {
   std::cout << "ViennaCL FAILED*****\n";
   std::cout << "ViennaCL failed, trying Eigen BiCGStab.\n";
   //solve_start_time = get_time_in_seconds();
   //iterative solver
   Eigen::IncompleteLUT<double> precon;
   Eigen::BiCGSTAB<Eigen::SparseMatrix<double>, Eigen::IncompleteLUT<double> > solver2;
   solver2.compute(eigen_system);
   solver2.setMaxIterations(1000);
   solver2.setTolerance(1e-9);
   Eigen::VectorXd eigen_rhs(dim);
   for (int i = 0; i < dim; ++i) eigen_rhs[i] = rhs[i];

   Eigen::VectorXd solution = solver2.solve(eigen_rhs);
   //solve_end_time = get_time_in_seconds();
   Eigen::VectorXd residual = eigen_system * solution - eigen_rhs;
   std::cout << "Residual magnitude: " << residual.norm() << std::endl;
   if (solver2.info() != Eigen::Success) {
   std::cout << "Solve failed.\n";
   //assert(false);
   //exit(1);
   }
   else {
   std::cout << "Solve succeeded!\n";
   for (int i = 0; i < soln.size(); ++i) {
   soln[i] = solution(i);
   }
   }
   }
   */

   printf("Solving sparse system.\n");
   PCGSolver<double> solver;
   double res_out;
   int iter_out;
   solver.set_solver_parameters(1e-9, 10000, 0.97, 0.1);

   printf("Launching CG\n");
   solver.solve(matrix, rhs, soln, res_out, iter_out);

   std::cout << "Finished with residual :" << res_out << " after iterations: " << iter_out << std::endl;

   if (iter_out >= 1000) {
      printf("\n\n\n***************FAILED******************\n\n\n");
      std::cout << "Residual :" << res_out << " iterations: " << iter_out << std::endl;
      exit(1);
   }

   printf("Copying back.\n");
   for (int k = 0; k < nz; ++k)
   for (int j = 0; j < ny; ++j)
   for (int i = 0; i < nx + 1; ++i) {
      if (u_state(i, j, k) == FLUID) {
         u(i, j, k) = (float)soln[u_ind(i, j, k, nx, ny)];
      }
      else {
         u(i, j, k) = 0;
      }
   }

   for (int k = 0; k < nz; ++k)
   for (int j = 0; j < ny + 1; ++j)
   for (int i = 0; i < nx; ++i) {
      if (v_state(i, j, k) == FLUID) {
         v(i, j, k) = (float)soln[v_ind(i, j, k, nx, ny, nz)];
      }
      else
         v(i, j, k) = 0;
   }

   for (int k = 0; k < nz + 1; ++k)
   for (int j = 0; j < ny; ++j)
   for (int i = 0; i < nx; ++i) {
      if (w_state(i, j, k) == FLUID) {
         w(i, j, k) = (float)soln[w_ind(i, j, k, nx, ny, nz)];
      }
      else
         w(i, j, k) = 0;
   }
   printf("Done copying back.\n");

}


#include <iostream>
#include <vector>
#include <array>
#include <functional>

#include <vexcl/devlist.hpp>
#include <vexcl/vector.hpp>
#include <vexcl/multivector.hpp>
#include <vexcl/generator.hpp>
#include <vexcl/element_index.hpp>

// http://headmyshoulder.github.com/odeint-v2
#include <boost/numeric/odeint.hpp>

namespace odeint = boost::numeric::odeint;

typedef double value_type;
typedef vex::symbolic< value_type > sym_value;
typedef std::array<sym_value, 3> sym_state;

// System  function for Lorenz attractor ensemble ODE.
// [1] http://headmyshoulder.github.com/odeint-v2/doc/boost_numeric_odeint/tutorial/chaotic_systems_and_lyapunov_exponents.html
// This is only used to record operations chain for autogenerated kernel.
struct sys_func
{
    const value_type sigma;
    const value_type b;
    const sym_value &R;

    sys_func(value_type sigma, value_type b, const sym_value &R)
        : sigma(sigma), b(b), R(R) {}

    template <class Sig>
    struct result {
        typedef void type;
    };

    void operator()( const sym_state &x , sym_state &dxdt , value_type ) const
    {
        dxdt[0] = sigma * (x[1] - x[0]);
        dxdt[1] = R * x[0] - x[1] - x[0] * x[2];
        dxdt[2] = x[0] * x[1] - b * x[2];
    }
};

int main( int argc , char **argv )
{
    size_t n;
    const value_type dt = 0.01;
    const value_type t_max = 100.0;

    using namespace std;

    n = argc > 1 ? atoi( argv[1] ) : 1024;

    vex::Context ctx( vex::Filter::DoublePrecision && vex::Filter::Env );
    cout << ctx << endl;

    // Custom kernel body will be recorded here:
    std::ostringstream body;
    vex::generator::set_recorder(body);

    // State types that would become kernel parameters:
    sym_state sym_S = {
        sym_value(sym_value::VectorParameter),
        sym_value(sym_value::VectorParameter),
        sym_value(sym_value::VectorParameter)
    };

    // Const kernel parameter.
    sym_value sym_R(sym_value::VectorParameter, sym_value::Const);

    /* Odeint is modern C++ library for ODE solution. We can use its collection
     * of ODE steppers to generate effective kernel which would compute one
     * iteration of a 4th order Runge-Kutta method. For that, we instantiate
     * appropriate odeint stepper with vex::symbolic<double> as a value type.
     * One iteration of the stepper would provide us with sequence of
     * expressions suitable for generation of requiered kernel.  This technique
     * may be used with any generic algorithms for generation of customized and
     * effective kernels.
     */
    // Symbolic stepper:
    odeint::runge_kutta4<
            sym_state , value_type , sym_state , value_type ,
            odeint::range_algebra , odeint::default_operations
            > sym_stepper;

    sys_func sys(10.0, 8.0 / 3.0, sym_R);
    sym_stepper.do_step(std::ref(sys), sym_S, 0, dt);

    auto kernel = vex::generator::build_kernel(ctx, "lorenz", body.str(),
            sym_S[0], sym_S[1], sym_S[2], sym_R
            );

    // Real state initialization:
    value_type Rmin = 0.1;
    value_type Rmax = 50.0;
    value_type dR   = (Rmax - Rmin) / (n - 1);

    vex::vector<value_type> X(ctx, n);
    vex::vector<value_type> Y(ctx, n);
    vex::vector<value_type> Z(ctx, n);
    vex::vector<value_type> R(ctx, n);

    X = 10.0;
    Y = 10.0;
    Z = 10.0;
    R = Rmin + dR * vex::element_index();

    // Integration loop:
    for(value_type t = 0; t < t_max; t += dt)
        kernel(X, Y, Z, R);

    std::vector< value_type > result( n );
    vex::copy( X , result );
    cout << result[0] << endl;
}

// vim: et


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

#include "Molassembler/Shapes/InertialMoments.h"

#include "Molassembler/Shapes/CoordinateSystemTransformation.h"
#include "Molassembler/Shapes/ContinuousMeasures.h"
#include <Eigen/Eigenvalues>

#include "Molassembler/Temple/Functional.h"
#include "Molassembler/Temple/Adaptors/Iota.h"

namespace Scine {
namespace Molassembler {
namespace Shapes {
namespace Detail {

//! Determine degeneracy of intertial moments
unsigned degeneracy(const Eigen::Vector3d& inertialMoments) {
  constexpr double degeneracyEpsilon = 0.05;
  unsigned mdeg = 0;
  if(
    std::fabs(
      (inertialMoments(2) - inertialMoments(1)) / inertialMoments(2)
    ) <= degeneracyEpsilon
  ) {
    mdeg += 1;
  }
  if(
    std::fabs(
      (inertialMoments(1) - inertialMoments(0)) / inertialMoments(2)
    ) <= degeneracyEpsilon
  ) {
    mdeg += 2;
  }

  return 1 + (mdeg + 1) / 2;
}

} // namespace Detail

InertialMoments principalInertialMoments(
  const InertialPositionsType& normalizedPositions
) {
  Eigen::Matrix3d inertialMatrix = Eigen::Matrix3d::Zero(3, 3);
  const unsigned N = normalizedPositions.cols();

  for(unsigned i = 0; i < N; ++i) {
    const auto& vec = normalizedPositions.col(i);
    inertialMatrix(0, 0) += vec.y() * vec.y() + vec.z() * vec.z();
    inertialMatrix(1, 1) += vec.x() * vec.x() + vec.z() * vec.z();
    inertialMatrix(2, 2) += vec.x() * vec.x() + vec.y() * vec.y();
    inertialMatrix(1, 0) -= vec.x() * vec.y(); // xy
    inertialMatrix(2, 0) -= vec.x() * vec.z(); // xz
    inertialMatrix(2, 1) -= vec.y() * vec.z(); // yz
  }

  // Decompose the inertial matrix to get principal axes and inertial moments
  Eigen::SelfAdjointEigenSolver<Eigen::Matrix3d> decomposition(inertialMatrix);

  InertialMoments result;
  result.moments = decomposition.eigenvalues();
  result.axes = decomposition.eigenvectors();
  return result;
}

Top standardizeTop(Eigen::Ref<InertialPositionsType> normalizedPositions) {
  const unsigned N = normalizedPositions.cols();
  assert(N > 1);

  InertialMoments moments = principalInertialMoments(normalizedPositions);

  const unsigned degeneracy = Detail::degeneracy(moments.moments);

  auto rotateEverything = [&](const CoordinateSystem& sourceSystem) {
    const CoordinateSystem defaultCoordinateSystem {};
    assert(sourceSystem.isRightHanded());
    const auto R = rotationMatrix(sourceSystem, defaultCoordinateSystem);
    // Rotate coordinates
    normalizedPositions = R * normalizedPositions;
    // Rotate inertial moment axes
    moments.axes = R * moments.axes;
  };

  if(moments.moments(0) < 0.1 && degeneracy == 2) {
    // The top is linear: If IA << IB = IC and IA ~ 0. We rotate IA to z
    const CoordinateSystem inertialMomentSystem {
      moments.axes.col(1),
      moments.axes.col(2)
    };
    rotateEverything(inertialMomentSystem);
    assert(moments.axes.col(0).cwiseAbs().isApprox(Eigen::Vector3d::UnitZ(), 1e-10));
    return Top::Line;
  }

  if(degeneracy == 1) {
    /* The top is asymmetric. Rotate the axis with the
     * highest moment of inertia to coincide with z, and the one with second most
     * to coincide with x.
     *
     * To better define orientation, we could look for Cn axes. This is done in
     * another function. No need to burden this function with that here.
     */
    CoordinateSystem inertialMomentSystem {
      moments.axes.col(1), // second highest becomes x
      moments.axes.col(2).cross(moments.axes.col(1)) // y = z.cross(x)
    };
    assert(inertialMomentSystem.z.isApprox(moments.axes.col(2), 1e-10));
    rotateEverything(inertialMomentSystem);
    // Make sure rotation went as intended
    assert(moments.axes.col(2).cwiseAbs().isApprox(Eigen::Vector3d::UnitZ(), 1e-10));
    assert(moments.axes.col(1).cwiseAbs().isApprox(Eigen::Vector3d::UnitX(), 1e-10));

    return Top::Asymmetric;
  }

  if(degeneracy == 2) {
    /* The top is symmetric. The subsets are:
     * - Oblate (disc): IA = IB < IC
     * - Prolate (rugby football): IA < IB = IC
     *
     * We rotate the unique axis to coincide with z (it's probably the site of
     * the highest-order Cn or Sn, and one of the degenerate axes to coincide
     * with x. There could be a C2 on x.
     *
     * This is most likely rare and should occur only for largely undistorted
     * structures. Perhaps we can flowchart point groups here?
     */
    // Calculate Ray's asymmetry parameter
    const double A = 1 / moments.moments(0);
    const double B = 1 / moments.moments(1);
    const double C = 1 / moments.moments(2);
    const double kappa = (2 * B - A - C) / (A - C);
    assert(-1 <= kappa && kappa <= 1);
    if(kappa < 0) {
      // Prolate top. IA is unique
      CoordinateSystem inertialMomentSystem {
        moments.axes.col(1),
        moments.axes.col(2)
      };
      rotateEverything(inertialMomentSystem);
      assert(moments.axes.col(0).cwiseAbs().isApprox(Eigen::Vector3d::UnitZ(), 1e-10));
      return Top::Prolate;
    }

    // Oblate top. IC is unique
    CoordinateSystem inertialMomentSystem {
      moments.axes.col(0),
      moments.axes.col(1)
    };
    rotateEverything(inertialMomentSystem);
    assert(moments.axes.col(2).cwiseAbs().isApprox(Eigen::Vector3d::UnitZ(), 1e-10));
    return Top::Oblate;
  }

  assert(degeneracy == 3);
  /* The top is spherical (IA = IB = IC).
   *
   * Note that there is no reason to rotate anything on the basis of the axes
   * from the inertial moment analysis since any choice of axes gives the
   * spherical symmetry. We can't use those to rotate the system.
   *
   * Rotate an arbitrary position to +z instead (good for Td
   * and Oh octahedral, less so for Oh cubic and Ih, which should be less
   * common)
   */
  unsigned selectedIndex = 0;
  for(; selectedIndex < N; ++selectedIndex) {
    /* As long as the position isn't close to the centroid and it's not exactly
     * the -z vector, we can rotate it
     */
    if(
      normalizedPositions.col(selectedIndex).norm() > 0.2
      && !normalizedPositions.col(selectedIndex).normalized().isApprox(
        -Eigen::Vector3d::UnitZ(),
        1e-10
      )
    ) {
      break;
    }
  }
  assert(selectedIndex != N);

  // Determine axis of rotation as sum of z and position coordinate
  const Eigen::Vector3d rotationAxis = (
    normalizedPositions.col(selectedIndex).normalized()
    + Eigen::Vector3d::UnitZ()
  ).normalized();
  const Eigen::Matrix3d rotationMatrix = Eigen::AngleAxisd(M_PI, rotationAxis).toRotationMatrix();

  // Rotate all coordinates
  for(unsigned i = 0; i < N; ++i) {
    normalizedPositions.col(i) = rotationMatrix * normalizedPositions.col(i);
  }
  // Check that everything went as planned
  assert(normalizedPositions.col(selectedIndex).normalized().cwiseAbs().isApprox(Eigen::Vector3d::UnitZ(), 1e-10));

  return Top::Spherical;
}

unsigned reorientAsymmetricTop(Eigen::Ref<InertialPositionsType> normalizedPositions) {
  const unsigned P = normalizedPositions.cols();
  const auto& axes = Eigen::Matrix3d::Identity();

  struct AxisBest {
    unsigned order = 1;
    double csm = 1; // This functions much like a detection threshold below
    unsigned axisIndex;

    AxisBest(unsigned index) : axisIndex(index) {}

    bool operator < (const AxisBest& other) const {
      return order > other.order;
    }
  };

  auto orderedAxisBest = Temple::sorted(
    Temple::map(
      Temple::iota<unsigned>(3),
      [&](const unsigned axisIndex) -> AxisBest {
        const Eigen::Vector3d axis = axes.col(axisIndex);
        AxisBest best {axisIndex};
        for(unsigned n = 2; n <= P; ++n) {
          const double axisCSM = Continuous::Fixed::element(normalizedPositions, Elements::Rotation::Cn(axis, n));
          if(axisCSM < best.csm) {
            best.order = n;
            best.csm = axisCSM;
          }
        }
        return best;
      }
    )
  );

  if(orderedAxisBest.front().order > 1) {
    /* Only mess with the coordinate frame if any sort of axis was found.
     * We want the second-highest order axis on x, highest order axis along z,
     * doesn't really matter if +z or -z
     */
    const CoordinateSystem highestOrderSystem {
      axes.col(orderedAxisBest.at(1).axisIndex),
      axes.col(orderedAxisBest.back().axisIndex)
    };

    normalizedPositions = rotationMatrix(highestOrderSystem, {}) * normalizedPositions;
  }

  return orderedAxisBest.front().order;
}

} // namespace Shapes
} // namespace Molassembler
} // namespace Scine


// Demonstration of calculating Elmore delay with a graph
// to accompany "Analyzing On-Chip Interconnect with Modern C++"
// Author: Jeff Trull <edaskel@att.net>

/*
Copyright (c) 2014 Jeffrey E. Trull

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.
*/

#include <iostream>
#include <fstream>

#include <boost/graph/visitors.hpp>
#include <boost/graph/undirected_dfs.hpp>
#include <boost/graph/graphviz.hpp>
#include <boost/units/systems/si/time.hpp>

#include "ckt_graph.h"

// For Elmore we need two visitors:
// 1) To calculate total capacitance at and "below" each node
// 2) To sum resistor delays to each node from a chosen input

template<typename Graph>
struct cap_summing_visitor : boost::default_dfs_visitor {
    typedef typename boost::graph_traits<Graph>::edge_descriptor edge_t;
    typedef typename boost::graph_traits<Graph>::vertex_descriptor vertex_t;

    // This edge variant visitor returns the correct contribution for an edge
    // to a downstream node: the value on the edge, if a capacitor, or
    // the supplied downstream capacitance, if a resistor.
    struct cap_summer : boost::static_visitor<capacitor_value_t> {
        capacitor_value_t operator()(resistor_value_t const&) const {
            return downstream_;   // count downstream, not edge
        }
        capacitor_value_t operator()(capacitor_value_t const& c) const {
            return c;             // count edge but not downstream (i.e., lump to gnd)
        }
        cap_summer(capacitor_value_t downstream) : downstream_(downstream) {}
    private:
        capacitor_value_t downstream_;
    };

    void tree_edge(edge_t e, Graph const& g) {
        // remember the predecessor so later we can avoid summing up capacitance from back edges
        predecessors_[target(e, g)] = source(e, g);
    }

    void finish_vertex(vertex_t u, Graph const& g) {
        // We want to sum only *downstream* capacitance
        // However, circuits are undirected, and the downstreamness of our search
        // is only due to the order in which we encountered vertices
        // a "tree edge" is downstream, and a "back edge" is upstream

        downstream_caps_[u] = capacitor_value_t();
        for (auto e : boost::make_iterator_range(out_edges(u, g))) {
            // detect back edge, i.e., an edge to an already-visited node
            // this won't work correctly for circuits with resistive loops
            if ((predecessors_[target(e, g)] == u) || (target(e, g) == g.gnd())) {
                // our first time at this node, OR the node is gnd, in which case this is normal
                capacitor_value_t contr = boost::apply_visitor(cap_summer(downstream_caps_[target(e, g)]), g[e]);
                downstream_caps_[u] += contr;
            }
        }
    }

    // finish_edge never seems to get called

    cap_summing_visitor(std::vector<capacitor_value_t>& capsvec)
        : downstream_caps_(capsvec), predecessors_(capsvec.size()) {}

private:
    std::vector<capacitor_value_t>& downstream_caps_;
    std::vector<vertex_t> predecessors_;
};

typedef quantity<si::time> delay_t;

// adds up delays.  To be run on filtered (R-only) graph
template<typename Graph>
struct delay_calculating_visitor : boost::default_dfs_visitor {
    typedef typename boost::graph_traits<Graph>::edge_descriptor edge_t;
    typedef typename boost::graph_traits<Graph>::vertex_descriptor vertex_t;

    void start_vertex(vertex_t u, Graph const&) {
        // this hook is called once at the beginning;
        // we can use it to initialize the delay from the tree root
        delays_[u] = delay_t();
    }

    struct res_summer : boost::static_visitor<resistor_value_t> {
        resistor_value_t operator()(resistor_value_t const& r) const {
            return r;
        }
        resistor_value_t operator()(capacitor_value_t const&) const {
            assert(0);   // program bug - we should operate on a filtered graph
            return resistor_value_t();
        }
    };

    void tree_edge(edge_t e, Graph const& g) {
        delays_[target(e, g)] = delays_[source(e, g)] +
            boost::apply_visitor(res_summer(), g[e]) *
            downstream_caps_[target(e, g)];
    }

    void back_edge(edge_t, Graph const&) {
        // Resistive loops will break this algorithm - you could put a check here
    }

    delay_calculating_visitor(std::vector<capacitor_value_t>& capsvec,
                              std::vector<delay_t>& delays)
        : downstream_caps_(capsvec), delays_(delays) {}

private:
    std::vector<capacitor_value_t>& downstream_caps_;   // bottom-up
    std::vector<delay_t>&          delays_;            // top-down
};

int main() {
    using namespace boost;
    using namespace std;

    quantity<si::resistance> kohm(1.0 * units::si::kilo  * units::si::ohms);
    quantity<si::capacitance> ff(1.0 * units::si::femto * units::si::farads);

    // Coupling test case
    ckt_graph_t coupling_test;
    auto gnd = coupling_test.gnd();

    auto vagg = add_vertex("vagg", coupling_test);   // driver voltage source
    auto n1   = add_vertex("n1",   coupling_test);
    add_edge(vagg, n1, 0.1*kohm,   coupling_test);   // driver impedance
    auto n2   = add_vertex("n2",   coupling_test);   // central node - where coupling occurs
    add_edge(n1, n2, 1.0*kohm,     coupling_test);   // first "pi" model, aggressor side
    add_edge(n1, gnd, 50.0*ff,     coupling_test);   // caps for first "pi"
    add_edge(n2, gnd, 50.0*ff,     coupling_test);
    auto n3   = add_vertex("n3",   coupling_test);   // aggressor-side receiver
    add_edge(n2, n3, 1.0*kohm,     coupling_test);   // second "pi" model, aggressor side
    add_edge(n2, gnd, 50.0*ff,     coupling_test);   // caps for first "pi"
    add_edge(n3, gnd, 50.0*ff,     coupling_test);
    add_edge(n3, gnd, 20.0*ff,     coupling_test);   // aggressor side receiver load

    // then the same thing over again for the victim side
    auto vvic = add_vertex("vvic", coupling_test);   // driver voltage source
    auto n5   = add_vertex("n5",   coupling_test);
    add_edge(vvic, n5, 0.1*kohm,   coupling_test);   // driver impedance
    auto n6   = add_vertex("n6",   coupling_test);   // central node
    add_edge(n5, n6, 1.0*kohm,     coupling_test);   // first "pi" model
    add_edge(n5, gnd, 50.0*ff,     coupling_test);
    add_edge(n6, gnd, 50.0*ff,     coupling_test);
    auto n7   = add_vertex("n7",   coupling_test);   // victim-side receiver
    add_edge(n6, n7, 1.0*kohm,     coupling_test);   // second "pi" model
    add_edge(n6, gnd, 50.0*ff,     coupling_test);
    add_edge(n7, gnd, 50.0*ff,     coupling_test);
    add_edge(n7, gnd, 20.0*ff,     coupling_test);   // victim side receiver load

    // coupling capacitor between the two signal traces
    add_edge(n2, n6, 100.0*ff,     coupling_test);

    // debug output via Dot
    /*
    ofstream dbg("/tmp/coupling_test.dot");
    write_graphviz(dbg, coupling_test,
                   make_label_writer(get(vertex_bundle, coupling_test)),
                   make_label_writer(get(edge_bundle, coupling_test)));
    */

    // calculate Elmore delay
    // Create visitors
    vector<capacitor_value_t> downstream_caps(num_vertices(coupling_test));
    cap_summing_visitor<ckt_graph_t> capvis(downstream_caps);

    // undirected DFS requires edge color map
    map<ckt_graph_t::edge_descriptor, default_color_type> edge_colors;
    auto cpmap = make_assoc_property_map(edge_colors);

    // perform first pass and sum capacitances
    undirected_dfs(coupling_test, edge_color_map(cpmap).visitor(capvis).root_vertex(vagg));

    typedef filtered_graph<ckt_graph_t, resistors_only> resonly_graph_t;
    resonly_graph_t res_graph(coupling_test, resistors_only(&coupling_test));

    vector<delay_t>          delays(num_vertices(coupling_test));
    delay_calculating_visitor<resonly_graph_t>
                                     delvis(downstream_caps, delays);

    // create *vertex* color map for depth_first_visit
    auto vindex_map = typed_identity_property_map<size_t>();
    vector<default_color_type> colorvec(num_vertices(coupling_test));   // underlying storage
    auto cvpmap = make_iterator_property_map(colorvec.begin(), vindex_map);

    depth_first_visit(res_graph, vagg, delvis, cvpmap);

    cout << "Elmore delay of aggressor net: " << delays.at(n3) << endl;

    // If you want to calculate the delay from a different node you have to recalculate
    // both capacitance and delays starting with that new node
}


/**
* \file NormalInitializer.hpp
*
* \brief Glorot normal initializer, as described in http://jmlr.org/proceedings/papers/v9/glorot10a/glorot10a.pdf
*
* \date   Jun 20, 2018
* \author Mathias Bøgh Stokholm
*/

#ifndef NEURAL_NORMALINITIALIZER_HPP
#define NEURAL_NORMALINITIALIZER_HPP

#include <Eigen/Core>
#include <random>

namespace neural {
    /**
     * @brief Glorot normal initializer (also called Xavier normal initializer)
     * Draws samples from a truncated normal distribution centered on 0 with the std deviation set according to the
     * number of inputs and outputs.
     * @tparam Dtype: The data type to generate samples as
     * @tparam FanIn: The number of input units of the tensor to generate weights for
     * @tparam FanOut: The number of output units of the tensor to generate weights for
     */
    template <typename Dtype, unsigned int FanIn, unsigned int FanOut>
    class GlorotNormal {
    public:
        GlorotNormal(): m_distribution(0.0, std::sqrt(2.0 / (FanIn + FanOut))) {}
        GlorotNormal(const GlorotNormal&) = default;

        Dtype operator()(Eigen::DenseIndex element_location, Eigen::DenseIndex /*unused*/ = 0) const {
            // FIXME: This is massively hacky and ugly
            auto* nonConstThis = const_cast<GlorotNormal*>(this);
            return Dtype(nonConstThis->m_distribution(nonConstThis->m_generator));
        }

    private:
        std::mt19937 m_generator;
        std::normal_distribution<double> m_distribution;
    };
}

#endif //NEURAL_NORMALINITIALIZER_HPP


//
// Copyright © 2017 Arm Ltd. All rights reserved.
// SPDX-License-Identifier: MIT
//

#include "DetectionPostProcess.hpp"

#include <armnn/utility/Assert.hpp>

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

#include <algorithm>
#include <numeric>

namespace armnn
{

std::vector<unsigned int> GenerateRangeK(unsigned int k)
{
    std::vector<unsigned int> range(k);
    std::iota(range.begin(), range.end(), 0);
    return range;
}

void TopKSort(unsigned int k, unsigned int* indices, const float* values, unsigned int numElement)
{
    std::partial_sort(indices, indices + k, indices + numElement,
                      [&values](unsigned int i, unsigned int j) { return values[i] > values[j]; });
}

float IntersectionOverUnion(const float* boxI, const float* boxJ)
{
    // Box-corner format: ymin, xmin, ymax, xmax.
    const int yMin = 0;
    const int xMin = 1;
    const int yMax = 2;
    const int xMax = 3;
    float areaI = (boxI[yMax] - boxI[yMin]) * (boxI[xMax] - boxI[xMin]);
    float areaJ = (boxJ[yMax] - boxJ[yMin]) * (boxJ[xMax] - boxJ[xMin]);
    float yMinIntersection = std::max(boxI[yMin], boxJ[yMin]);
    float xMinIntersection = std::max(boxI[xMin], boxJ[xMin]);
    float yMaxIntersection = std::min(boxI[yMax], boxJ[yMax]);
    float xMaxIntersection = std::min(boxI[xMax], boxJ[xMax]);
    float areaIntersection = std::max(yMaxIntersection - yMinIntersection, 0.0f) *
                                std::max(xMaxIntersection - xMinIntersection, 0.0f);
    float areaUnion = areaI + areaJ - areaIntersection;
    return areaIntersection / areaUnion;
}

std::vector<unsigned int> NonMaxSuppression(unsigned int numBoxes,
                                            const std::vector<float>& boxCorners,
                                            const std::vector<float>& scores,
                                            float nmsScoreThreshold,
                                            unsigned int maxDetection,
                                            float nmsIouThreshold)
{
    // Select boxes that have scores above a given threshold.
    std::vector<float> scoresAboveThreshold;
    std::vector<unsigned int> indicesAboveThreshold;
    for (unsigned int i = 0; i < numBoxes; ++i)
    {
        if (scores[i] >= nmsScoreThreshold)
        {
            scoresAboveThreshold.push_back(scores[i]);
            indicesAboveThreshold.push_back(i);
        }
    }

    // Sort the indices based on scores.
    unsigned int numAboveThreshold = boost::numeric_cast<unsigned int>(scoresAboveThreshold.size());
    std::vector<unsigned int> sortedIndices = GenerateRangeK(numAboveThreshold);
    TopKSort(numAboveThreshold, sortedIndices.data(), scoresAboveThreshold.data(), numAboveThreshold);

    // Number of output cannot be more than max detections specified in the option.
    unsigned int numOutput = std::min(maxDetection, numAboveThreshold);
    std::vector<unsigned int> outputIndices;
    std::vector<bool> visited(numAboveThreshold, false);

    // Prune out the boxes with high intersection over union by keeping the box with higher score.
    for (unsigned int i = 0; i < numAboveThreshold; ++i)
    {
        if (outputIndices.size() >= numOutput)
        {
            break;
        }
        if (!visited[sortedIndices[i]])
        {
            outputIndices.push_back(indicesAboveThreshold[sortedIndices[i]]);
        }
        for (unsigned int j = i + 1; j < numAboveThreshold; ++j)
        {
            unsigned int iIndex = indicesAboveThreshold[sortedIndices[i]] * 4;
            unsigned int jIndex = indicesAboveThreshold[sortedIndices[j]] * 4;
            if (IntersectionOverUnion(&boxCorners[iIndex], &boxCorners[jIndex]) > nmsIouThreshold)
            {
                visited[sortedIndices[j]] = true;
            }
        }
    }
    return outputIndices;
}

void AllocateOutputData(unsigned int numOutput,
                        unsigned int numSelected,
                        const std::vector<float>& boxCorners,
                        const std::vector<unsigned int>& outputIndices,
                        const std::vector<unsigned int>& selectedBoxes,
                        const std::vector<unsigned int>& selectedClasses,
                        const std::vector<float>& selectedScores,
                        float* detectionBoxes,
                        float* detectionScores,
                        float* detectionClasses,
                        float* numDetections)
{
    for (unsigned int i = 0; i < numOutput; ++i)
        {
            unsigned int boxIndex = i * 4;
            if (i < numSelected)
            {
                unsigned int boxCornorIndex = selectedBoxes[outputIndices[i]] * 4;
                detectionScores[i] = selectedScores[outputIndices[i]];
                detectionClasses[i] = boost::numeric_cast<float>(selectedClasses[outputIndices[i]]);
                detectionBoxes[boxIndex] = boxCorners[boxCornorIndex];
                detectionBoxes[boxIndex + 1] = boxCorners[boxCornorIndex + 1];
                detectionBoxes[boxIndex + 2] = boxCorners[boxCornorIndex + 2];
                detectionBoxes[boxIndex + 3] = boxCorners[boxCornorIndex + 3];
            }
            else
            {
                detectionScores[i] = 0.0f;
                detectionClasses[i] = 0.0f;
                detectionBoxes[boxIndex] = 0.0f;
                detectionBoxes[boxIndex + 1] = 0.0f;
                detectionBoxes[boxIndex + 2] = 0.0f;
                detectionBoxes[boxIndex + 3] = 0.0f;
            }
        }
        numDetections[0] = boost::numeric_cast<float>(numSelected);
}

void DetectionPostProcess(const TensorInfo& boxEncodingsInfo,
                          const TensorInfo& scoresInfo,
                          const TensorInfo& anchorsInfo,
                          const TensorInfo& detectionBoxesInfo,
                          const TensorInfo& detectionClassesInfo,
                          const TensorInfo& detectionScoresInfo,
                          const TensorInfo& numDetectionsInfo,
                          const DetectionPostProcessDescriptor& desc,
                          Decoder<float>& boxEncodings,
                          Decoder<float>& scores,
                          Decoder<float>& anchors,
                          float* detectionBoxes,
                          float* detectionClasses,
                          float* detectionScores,
                          float* numDetections)
{
    IgnoreUnused(anchorsInfo, detectionClassesInfo, detectionScoresInfo, numDetectionsInfo);

    // Transform center-size format which is (ycenter, xcenter, height, width) to box-corner format,
    // which represents the lower left corner and the upper right corner (ymin, xmin, ymax, xmax)
    std::vector<float> boxCorners(boxEncodingsInfo.GetNumElements());

    const unsigned int numBoxes  = boxEncodingsInfo.GetShape()[1];
    const unsigned int numScores = scoresInfo.GetNumElements();

    for (unsigned int i = 0; i < numBoxes; ++i)
    {
        // Y
        float boxEncodingY = boxEncodings.Get();
        float anchorY      = anchors.Get();

        ++boxEncodings;
        ++anchors;

        // X
        float boxEncodingX = boxEncodings.Get();
        float anchorX      = anchors.Get();

        ++boxEncodings;
        ++anchors;

        // H
        float boxEncodingH = boxEncodings.Get();
        float anchorH      = anchors.Get();

        ++boxEncodings;
        ++anchors;

        // W
        float boxEncodingW = boxEncodings.Get();
        float anchorW      = anchors.Get();

        ++boxEncodings;
        ++anchors;

        float yCentre = boxEncodingY / desc.m_ScaleY * anchorH + anchorY;
        float xCentre = boxEncodingX / desc.m_ScaleX * anchorW + anchorX;

        float halfH = 0.5f * expf(boxEncodingH / desc.m_ScaleH) * anchorH;
        float halfW = 0.5f * expf(boxEncodingW / desc.m_ScaleW) * anchorW;

        unsigned int indexY = i * 4;
        unsigned int indexX = indexY + 1;
        unsigned int indexH = indexX + 1;
        unsigned int indexW = indexH + 1;

        // ymin
        boxCorners[indexY] = yCentre - halfH;
        // xmin
        boxCorners[indexX] = xCentre - halfW;
        // ymax
        boxCorners[indexH] = yCentre + halfH;
        // xmax
        boxCorners[indexW] = xCentre + halfW;

        ARMNN_ASSERT(boxCorners[indexY] < boxCorners[indexH]);
        ARMNN_ASSERT(boxCorners[indexX] < boxCorners[indexW]);
    }

    unsigned int numClassesWithBg = desc.m_NumClasses + 1;

    // Decode scores
    std::vector<float> decodedScores;
    decodedScores.reserve(numScores);

    for (unsigned int i = 0u; i < numScores; ++i)
    {
        decodedScores.emplace_back(scores.Get());
        ++scores;
    }

    // Perform Non Max Suppression.
    if (desc.m_UseRegularNms)
    {
        // Perform Regular NMS.
        // For each class, perform NMS and select max detection numbers of the highest score across all classes.
        std::vector<float> classScores(numBoxes);

        std::vector<unsigned int> selectedBoxesAfterNms;
        selectedBoxesAfterNms.reserve(numBoxes);

        std::vector<float> selectedScoresAfterNms;
        selectedBoxesAfterNms.reserve(numScores);

        std::vector<unsigned int> selectedClasses;

        for (unsigned int c = 0; c < desc.m_NumClasses; ++c)
        {
            // For each boxes, get scores of the boxes for the class c.
            for (unsigned int i = 0; i < numBoxes; ++i)
            {
                classScores[i] = decodedScores[i * numClassesWithBg + c + 1];
            }
            std::vector<unsigned int> selectedIndices = NonMaxSuppression(numBoxes,
                                                                          boxCorners,
                                                                          classScores,
                                                                          desc.m_NmsScoreThreshold,
                                                                          desc.m_DetectionsPerClass,
                                                                          desc.m_NmsIouThreshold);

            for (unsigned int i = 0; i < selectedIndices.size(); ++i)
            {
                selectedBoxesAfterNms.push_back(selectedIndices[i]);
                selectedScoresAfterNms.push_back(classScores[selectedIndices[i]]);
                selectedClasses.push_back(c);
            }
        }

        // Select max detection numbers of the highest score across all classes
        unsigned int numSelected = boost::numeric_cast<unsigned int>(selectedBoxesAfterNms.size());
        unsigned int numOutput = std::min(desc.m_MaxDetections,  numSelected);

        // Sort the max scores among the selected indices.
        std::vector<unsigned int> outputIndices = GenerateRangeK(numSelected);
        TopKSort(numOutput, outputIndices.data(), selectedScoresAfterNms.data(), numSelected);

        AllocateOutputData(detectionBoxesInfo.GetShape()[1], numOutput, boxCorners, outputIndices,
                           selectedBoxesAfterNms, selectedClasses, selectedScoresAfterNms,
                           detectionBoxes, detectionScores, detectionClasses, numDetections);
    }
    else
    {
        // Perform Fast NMS.
        // Select max scores of boxes and perform NMS on max scores,
        // select max detection numbers of the highest score
        unsigned int numClassesPerBox = std::min(desc.m_MaxClassesPerDetection, desc.m_NumClasses);
        std::vector<float> maxScores;
        std::vector<unsigned int>boxIndices;
        std::vector<unsigned int>maxScoreClasses;

        for (unsigned int box = 0; box < numBoxes; ++box)
        {
            unsigned int scoreIndex = box * numClassesWithBg + 1;

            // Get the max scores of the box.
            std::vector<unsigned int> maxScoreIndices = GenerateRangeK(desc.m_NumClasses);
            TopKSort(numClassesPerBox, maxScoreIndices.data(),
                decodedScores.data() + scoreIndex, desc.m_NumClasses);

            for (unsigned int i = 0; i < numClassesPerBox; ++i)
            {
                maxScores.push_back(decodedScores[scoreIndex + maxScoreIndices[i]]);
                maxScoreClasses.push_back(maxScoreIndices[i]);
                boxIndices.push_back(box);
            }
        }

        // Perform NMS on max scores
        std::vector<unsigned int> selectedIndices = NonMaxSuppression(numBoxes, boxCorners, maxScores,
                                                                      desc.m_NmsScoreThreshold,
                                                                      desc.m_MaxDetections,
                                                                      desc.m_NmsIouThreshold);

        unsigned int numSelected = boost::numeric_cast<unsigned int>(selectedIndices.size());
        unsigned int numOutput = std::min(desc.m_MaxDetections,  numSelected);

        AllocateOutputData(detectionBoxesInfo.GetShape()[1], numOutput, boxCorners, selectedIndices,
                           boxIndices, maxScoreClasses, maxScores,
                           detectionBoxes, detectionScores, detectionClasses, numDetections);
    }
}

} // namespace armnn