/**
* This file is part of Fast-Planner.
*
* Copyright 2019 Boyu Zhou, Aerial Robotics Group, Hong Kong University of Science and Technology, <uav.ust.hk>
* Developed by Boyu Zhou <bzhouai at connect dot ust dot hk>, <uv dot boyuzhou at gmail dot com>
* for more information see <https://github.com/HKUST-Aerial-Robotics/Fast-Planner>.
* If you use this code, please cite the respective publications as
* listed on the above website.
*
* Fast-Planner is free software: you can redistribute it and/or modify
* it under the terms of the GNU Lesser General Public License as published by
* the Free Software Foundation, either version 3 of the License, or
* (at your option) any later version.
*
* Fast-Planner is distributed in the hope that it will be useful,
* but WITHOUT ANY WARRANTY; without even the implied warranty of
* MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
* GNU General Public License for more details.
*
* You should have received a copy of the GNU Lesser General Public License
* along with Fast-Planner. If not, see <http://www.gnu.org/licenses/>.
*/

#include <Eigen/Eigen>
#include <cmath>
#include <iostream>
#include <plan_env/raycast.h>

int signum(int x) {
  return x == 0 ? 0 : x < 0 ? -1 : 1;
}

double mod(double value, double modulus) {
  return fmod(fmod(value, modulus) + modulus, modulus);
}

double intbound(double s, double ds) {
  // Find the smallest positive t such that s+t*ds is an integer.
  if (ds < 0) {
    return intbound(-s, -ds);
  } else {
    s = mod(s, 1);
    // problem is now s+t*ds = 1
    return (1 - s) / ds;
  }
}

void Raycast(const Eigen::Vector3d& start, const Eigen::Vector3d& end, const Eigen::Vector3d& min,
             const Eigen::Vector3d& max, int& output_points_cnt, Eigen::Vector3d* output) {
  //    std::cout << start << ' ' << end << std::endl;
  // From "A Fast Voxel Traversal Algorithm for Ray Tracing"
  // by John Amanatides and Andrew Woo, 1987
  // <http://www.cse.yorku.ca/~amana/research/grid.pdf>
  // <http://citeseer.ist.psu.edu/viewdoc/summary?doi=10.1.1.42.3443>
  // Extensions to the described algorithm:
  //   • Imposed a distance limit.
  //   • The face passed through to reach the current cube is provided to
  //     the callback.

  // The foundation of this algorithm is a parameterized representation of
  // the provided ray,
  //                    origin + t * direction,
  // except that t is not actually stored; rather, at any given point in the
  // traversal, we keep track of the *greater* t values which we would have
  // if we took a step sufficient to cross a cube boundary along that axis
  // (i.e. change the integer part of the coordinate) in the variables
  // tMaxX, tMaxY, and tMaxZ.

  // Cube containing origin point.
  int x = (int)std::floor(start.x());
  int y = (int)std::floor(start.y());
  int z = (int)std::floor(start.z());
  int endX = (int)std::floor(end.x());
  int endY = (int)std::floor(end.y());
  int endZ = (int)std::floor(end.z());
  Eigen::Vector3d direction = (end - start);
  double maxDist = direction.squaredNorm();

  // Break out direction vector.
  double dx = endX - x;
  double dy = endY - y;
  double dz = endZ - z;

  // Direction to increment x,y,z when stepping.
  int stepX = (int)signum((int)dx);
  int stepY = (int)signum((int)dy);
  int stepZ = (int)signum((int)dz);

  // See description above. The initial values depend on the fractional
  // part of the origin.
  double tMaxX = intbound(start.x(), dx);
  double tMaxY = intbound(start.y(), dy);
  double tMaxZ = intbound(start.z(), dz);

  // The change in t when taking a step (always positive).
  double tDeltaX = ((double)stepX) / dx;
  double tDeltaY = ((double)stepY) / dy;
  double tDeltaZ = ((double)stepZ) / dz;

  // Avoids an infinite loop.
  if (stepX == 0 && stepY == 0 && stepZ == 0) return;

  double dist = 0;
  while (true) {
    if (x >= min.x() && x < max.x() && y >= min.y() && y < max.y() && z >= min.z() && z < max.z()) {
      output[output_points_cnt](0) = x;
      output[output_points_cnt](1) = y;
      output[output_points_cnt](2) = z;

      output_points_cnt++;
      dist = sqrt((x - start(0)) * (x - start(0)) + (y - start(1)) * (y - start(1)) +
                  (z - start(2)) * (z - start(2)));

      if (dist > maxDist) return;

      /*            if (output_points_cnt > 1500) {
                      std::cerr << "Error, too many racyast voxels." <<
         std::endl;
                      throw std::out_of_range("Too many raycast voxels");
                  }*/
    }

    if (x == endX && y == endY && z == endZ) break;

    // tMaxX stores the t-value at which we cross a cube boundary along the
    // X axis, and similarly for Y and Z. Therefore, choosing the least tMax
    // chooses the closest cube boundary. Only the first case of the four
    // has been commented in detail.
    if (tMaxX < tMaxY) {
      if (tMaxX < tMaxZ) {
        // Update which cube we are now in.
        x += stepX;
        // Adjust tMaxX to the next X-oriented boundary crossing.
        tMaxX += tDeltaX;
      } else {
        z += stepZ;
        tMaxZ += tDeltaZ;
      }
    } else {
      if (tMaxY < tMaxZ) {
        y += stepY;
        tMaxY += tDeltaY;
      } else {
        z += stepZ;
        tMaxZ += tDeltaZ;
      }
    }
  }
}

void Raycast(const Eigen::Vector3d& start, const Eigen::Vector3d& end, const Eigen::Vector3d& min,
             const Eigen::Vector3d& max, std::vector<Eigen::Vector3d>* output) {
  //    std::cout << start << ' ' << end << std::endl;
  // From "A Fast Voxel Traversal Algorithm for Ray Tracing"
  // by John Amanatides and Andrew Woo, 1987
  // <http://www.cse.yorku.ca/~amana/research/grid.pdf>
  // <http://citeseer.ist.psu.edu/viewdoc/summary?doi=10.1.1.42.3443>
  // Extensions to the described algorithm:
  //   • Imposed a distance limit.
  //   • The face passed through to reach the current cube is provided to
  //     the callback.

  // The foundation of this algorithm is a parameterized representation of
  // the provided ray,
  //                    origin + t * direction,
  // except that t is not actually stored; rather, at any given point in the
  // traversal, we keep track of the *greater* t values which we would have
  // if we took a step sufficient to cross a cube boundary along that axis
  // (i.e. change the integer part of the coordinate) in the variables
  // tMaxX, tMaxY, and tMaxZ.

  // Cube containing origin point.
  int x = (int)std::floor(start.x());
  int y = (int)std::floor(start.y());
  int z = (int)std::floor(start.z());
  int endX = (int)std::floor(end.x());
  int endY = (int)std::floor(end.y());
  int endZ = (int)std::floor(end.z());
  Eigen::Vector3d direction = (end - start);
  double maxDist = direction.squaredNorm();

  // Break out direction vector.
  double dx = endX - x;
  double dy = endY - y;
  double dz = endZ - z;

  // Direction to increment x,y,z when stepping.
  int stepX = (int)signum((int)dx);
  int stepY = (int)signum((int)dy);
  int stepZ = (int)signum((int)dz);

  // See description above. The initial values depend on the fractional
  // part of the origin.
  double tMaxX = intbound(start.x(), dx);
  double tMaxY = intbound(start.y(), dy);
  double tMaxZ = intbound(start.z(), dz);

  // The change in t when taking a step (always positive).
  double tDeltaX = ((double)stepX) / dx;
  double tDeltaY = ((double)stepY) / dy;
  double tDeltaZ = ((double)stepZ) / dz;

  output->clear();

  // Avoids an infinite loop.
  if (stepX == 0 && stepY == 0 && stepZ == 0) return;

  double dist = 0;
  while (true) {
    if (x >= min.x() && x < max.x() && y >= min.y() && y < max.y() && z >= min.z() && z < max.z()) {
      output->push_back(Eigen::Vector3d(x, y, z));

      dist = (Eigen::Vector3d(x, y, z) - start).squaredNorm();

      if (dist > maxDist) return;

      if (output->size() > 1500) {
        std::cerr << "Error, too many racyast voxels." << std::endl;
        throw std::out_of_range("Too many raycast voxels");
      }
    }

    if (x == endX && y == endY && z == endZ) break;

    // tMaxX stores the t-value at which we cross a cube boundary along the
    // X axis, and similarly for Y and Z. Therefore, choosing the least tMax
    // chooses the closest cube boundary. Only the first case of the four
    // has been commented in detail.
    if (tMaxX < tMaxY) {
      if (tMaxX < tMaxZ) {
        // Update which cube we are now in.
        x += stepX;
        // Adjust tMaxX to the next X-oriented boundary crossing.
        tMaxX += tDeltaX;
      } else {
        z += stepZ;
        tMaxZ += tDeltaZ;
      }
    } else {
      if (tMaxY < tMaxZ) {
        y += stepY;
        tMaxY += tDeltaY;
      } else {
        z += stepZ;
        tMaxZ += tDeltaZ;
      }
    }
  }
}

bool RayCaster::setInput(const Eigen::Vector3d& start,
                         const Eigen::Vector3d& end /* , const Eigen::Vector3d& min,
                         const Eigen::Vector3d& max */) {
  start_ = start;
  end_ = end;
  // max_ = max;
  // min_ = min;

  x_ = (int)std::floor(start_.x());
  y_ = (int)std::floor(start_.y());
  z_ = (int)std::floor(start_.z());
  endX_ = (int)std::floor(end_.x());
  endY_ = (int)std::floor(end_.y());
  endZ_ = (int)std::floor(end_.z());
  direction_ = (end_ - start_);
  maxDist_ = direction_.squaredNorm();

  // Break out direction vector.
  dx_ = endX_ - x_;
  dy_ = endY_ - y_;
  dz_ = endZ_ - z_;

  // Direction to increment x,y,z when stepping.
  stepX_ = (int)signum((int)dx_);
  stepY_ = (int)signum((int)dy_);
  stepZ_ = (int)signum((int)dz_);

  // See description above. The initial values depend on the fractional
  // part of the origin.
  tMaxX_ = intbound(start_.x(), dx_);
  tMaxY_ = intbound(start_.y(), dy_);
  tMaxZ_ = intbound(start_.z(), dz_);

  // The change in t when taking a step (always positive).
  tDeltaX_ = ((double)stepX_) / dx_;
  tDeltaY_ = ((double)stepY_) / dy_;
  tDeltaZ_ = ((double)stepZ_) / dz_;

  dist_ = 0;

  step_num_ = 0;

  // Avoids an infinite loop.
  if (stepX_ == 0 && stepY_ == 0 && stepZ_ == 0)
    return false;
  else
    return true;
}

bool RayCaster::step(Eigen::Vector3d& ray_pt) {
  // if (x_ >= min_.x() && x_ < max_.x() && y_ >= min_.y() && y_ < max_.y() &&
  // z_ >= min_.z() && z_ <
  // max_.z())
  ray_pt = Eigen::Vector3d(x_, y_, z_);

  // step_num_++;

  // dist_ = (Eigen::Vector3d(x_, y_, z_) - start_).squaredNorm();

  if (x_ == endX_ && y_ == endY_ && z_ == endZ_) {
    return false;
  }

  // if (dist_ > maxDist_)
  // {
  //   return false;
  // }

  // tMaxX stores the t-value at which we cross a cube boundary along the
  // X axis, and similarly for Y and Z. Therefore, choosing the least tMax
  // chooses the closest cube boundary. Only the first case of the four
  // has been commented in detail.
  if (tMaxX_ < tMaxY_) {
    if (tMaxX_ < tMaxZ_) {
      // Update which cube we are now in.
      x_ += stepX_;
      // Adjust tMaxX to the next X-oriented boundary crossing.
      tMaxX_ += tDeltaX_;
    } else {
      z_ += stepZ_;
      tMaxZ_ += tDeltaZ_;
    }
  } else {
    if (tMaxY_ < tMaxZ_) {
      y_ += stepY_;
      tMaxY_ += tDeltaY_;
    } else {
      z_ += stepZ_;
      tMaxZ_ += tDeltaZ_;
    }
  }

  return true;
}


/*
 * TreeIkSolverVel_wdls.hpp
 *
 *  Created on: Nov 28, 2008
 *      Author: rubensmits
 */

#ifndef TREEIKSOLVERVEL_WDLS_HPP_
#define TREEIKSOLVERVEL_WDLS_HPP_

#include "treeiksolver.hpp"
#include "treejnttojacsolver.hpp"
#include <Eigen/Core>

namespace KDL {

    class TreeIkSolverVel_wdls: public TreeIkSolverVel {
    public:
        static const int E_SVD_FAILED = -100; //! Child SVD failed

        TreeIkSolverVel_wdls(const Tree& tree, const std::vector<std::string>& endpoints);
        virtual ~TreeIkSolverVel_wdls();

        virtual double CartToJnt(const JntArray& q_in, const Twists& v_in, JntArray& qdot_out);

        /*
         * Set the joint space weighting matrix
         *
         * @param weight_js joint space weighting symmetric matrix,
         * default : identity.  M_q : This matrix being used as a
         * weight for the norm of the joint space speed it HAS TO BE
         * symmetric and positive definite. We can actually deal with
         * matrices containing a symmetric and positive definite block
         * and 0s otherwise. Taking a diagonal matrix as an example, a
         * 0 on the diagonal means that the corresponding joints will
         * not contribute to the motion of the system. On the other
         * hand, the bigger the value, the most the corresponding
         * joint will contribute to the overall motion. The obtained
         * solution q_dot will actually minimize the weighted norm
         * sqrt(q_dot'*(M_q^-2)*q_dot). In the special case we deal
         * with, it does not make sense to invert M_q but what is
         * important is the physical meaning of all this : a joint
         * that has a zero weight in M_q will not contribute to the
         * motion of the system and this is equivalent to saying that
         * it gets an infinite weight in the norm computation.  For
         * more detailed explanation : vincent.padois@upmc.fr
         */
        void setWeightJS(const Eigen::MatrixXd& Mq);
        const Eigen::MatrixXd& getWeightJS() const {return Wq;}

        /*
         * Set the task space weighting matrix
         *
         * @param weight_ts task space weighting symmetric matrix,
         * default: identity M_x : This matrix being used as a weight
         * for the norm of the error (in terms of task space speed) it
         * HAS TO BE symmetric and positive definite. We can actually
         * deal with matrices containing a symmetric and positive
         * definite block and 0s otherwise. Taking a diagonal matrix
         * as an example, a 0 on the diagonal means that the
         * corresponding task coordinate will not be taken into
         * account (ie the corresponding error can be really big). If
         * the rank of the jacobian is equal to the number of task
         * space coordinates which do not have a 0 weight in M_x, the
         * weighting will actually not impact the results (ie there is
         * an exact solution to the velocity inverse kinematics
         * problem). In cases without an exact solution, the bigger
         * the value, the most the corresponding task coordinate will
         * be taken into account (ie the more the corresponding error
         * will be reduced). The obtained solution will minimize the
         * weighted norm sqrt(|x_dot-Jq_dot|'*(M_x^2)*|x_dot-Jq_dot|).
         * For more detailed explanation : vincent.padois@upmc.fr
         */
        void setWeightTS(const Eigen::MatrixXd& Mx);
        const Eigen::MatrixXd& getWeightTS() const {return Wy;}

        void setLambda(const double& lambda);
        double getLambda () const {return lambda;}

    private:
        Tree tree;
        TreeJntToJacSolver jnttojacsolver;
        Jacobians jacobians;

        Eigen::MatrixXd J, Wy, Wq, J_Wq, Wy_J_Wq, U, V, Wy_U, Wq_V;
        Eigen::VectorXd t, Wy_t, qdot, tmp, S;
        double lambda;
    };

}

#endif /* TREEIKSOLVERVEL_WDLS_HPP_ */


// written by bastiaan konings schuiling 2008 - 2014
// this work is public domain. the code is undocumented, scruffy, untested, and should generally not be used for anything important.
// i do not offer support, so don't ask. to be used for inspiration :)

#include "bluntmath.hpp"

#include <cmath>

#include <boost/random.hpp>

#include <boost/thread/mutex.hpp>

namespace blunted {

  unsigned int fastrandseed;
  unsigned int max_uint;

  boost::mutex randMutex;

  typedef boost::mt19937 BaseGenerator;
  typedef boost::uniform_real<float> Distribution;
  typedef boost::variate_generator<BaseGenerator, Distribution> Generator;
  BaseGenerator base;
  Distribution dist;
  Generator rng(base, dist);

  real clamp(const real value, const real min, const real max) {
    assert(max >= min);
    if (min > value) return min;
    if (max < value) return max;
    return value;
  }

  real NormalizedClamp(const real value, const real min, const real max) {
    assert(max > min);
    real banana = clamp(value, min, max);
    banana = (banana - min) / (max - min);
    return banana;
  }

  real invsqrt(real fvalue) {
    return 1. / sqrt(fvalue);
  }

  float dot_product(real v1[3], real v2[3]) {
    return (v1[0] * v2[0] + v1[1] * v2[1] + v1[2] * v2[2]);
  }

  void normalize(real v[3]) {
    real f = 1.0f / sqrt(dot_product(v, v));

    v[0] *= f;
    v[1] *= f;
    v[2] *= f;
  }

  bool sign(real n) {
    return n >= 0;
  }

  signed int signSide(real n) {
    return n >= 0 ? 1 : -1;
  }

  bool is_odd(int n) {
    return n & 1;
  }

  void randomseed() {
    rng.engine().seed(static_cast<unsigned int>(std::time(0)));
  }

  inline real boostrandom() {
    return rng();
  }

  real random(real min, real max) {
    float stretch = max - min;

    randMutex.lock();
    real value = min + (boostrandom() * stretch);
    randMutex.unlock();
    return value;
  }

  int pot(int x) {
    int val = 1;
    while (val < x) {
      val *= 2;
    }
    return val;
  }

  real ModulateIntoRange(real min, real max, real value) {
    real step = max - min;
    real newValue = value;
    while (newValue < min) newValue += step;
    while (newValue > max) newValue -= step;
    return newValue;
  }

}


/* author: jakob fischer (mail@jakobfischer.eu)
 * description:
 * Tool for solving a ODE for a reaction system given as an jrnf-file. The concentration and time
 * are read and written from / to a comma seperated file. Every row / line represents one time step.
 * the first column contains the time of this step, the second the mean of the quadratic change
 * (<f(x)> ;dx/dt = f(x)). Then follow the concentrations of all species.
 *
 * One version of the program (option "simulate" calculates effective reaction rates directly
 * from formation enthalpies and activation energies given in the network description.
 * (\beta = 1/(k_b T) = 1). This allows higher precision and usage of an explicit solver which
 * is over all faster for small networks (< 50 species).
 *
 * TODO There seems to be a problem with exceptionally short steps (distances in the
 *      <times> vector given to integrate_times) as well as with very wide steps. This
 *      occurs when the "write_log" option lead to output step sizes below 1e-7 or
 *      above 1e7. The solution for now is to first use linear stepping, then
 *      "write_log_abs" and at larger timescale linear stepping again.
 *      Giving a minimal and maximal step size as parameter might be a good
 *      improvement for future versions!
 */

#include <iostream>
#include <fstream>
#include <ctime>
#include <utility>

#include <boost/array.hpp>
#include <boost/numeric/odeint.hpp>

#include <algorithm>
#include "tools/cl_para.h"
#include "net_tools/reaction_network.h"
#include "net_tools/reaction_network_fileop.h"
#include "net_tools/network_tools.h"

using namespace std;
using namespace boost::numeric::odeint;

#ifndef GIT_VERSION
#define GIT_VERSION "no version"
#endif

// version string increment if mayor changes happen
const string odeint_rnet_version_string="1x00x00";

// various ublas types needed for implicit solver
typedef boost::numeric::ublas::vector< double > vector_type;
typedef boost::numeric::ublas::matrix< double > matrix_type;
typedef boost::numeric::ublas::matrix< int > matrix_type_int;

/*
 * Main class that bundles the functionality to load reaction networks and
 * files containing species' concentration. Every object is associated with
 * an network and concentration file at creation. The class offers methods
 * for solving the ODE defined by the network and write the results to the
 * concentration file.
 */

class reaction_network_system {
    double beta;                // beta = 1/(k_b T)
    std::vector<species> sp;    // The reaction
    std::vector<reaction> re;   // network
    // initial time, concentration loaded with concentration file
    double initial_t;
    vector_type initial_con;
    std::string fn_concentration;      // filename

    matrix_type_int N, N_in, N_out;    // stoichiometric matrices
    // List that contains for each column in N_in / N_out (= for each reaction)
    // a vector of those rows (species) that are nonzero.
    std::vector< std::vector<size_t> > N_in_list, N_out_list;
    // Intermediate values / thermodynamic values
    vector_type Ea, mu0, e_bs_in, e_bs_out, e_m_bEa;

public:

    /*
     * Simple calculation of concentration change (rhs / f(x)) from species
     * concentration and reaction constants (<k>, <k_b>).
     */

    struct fast_system {
        reaction_network_system& rns;

        fast_system(reaction_network_system& rns_) : rns(rns_)  {}

        // <x> - state / concentration;  dxdt - change / rhs (output)
        void operator()( const vector_type &x , vector_type &dxdt , double /* t */ ) {

            for(size_t i=0; i<rns.sp.size(); ++i)
                dxdt[i]=0;

            for(size_t i=0; i<rns.re.size(); ++i) {
                double rate_f=rns.re[i].get_k();
	        double rate_b=rns.re[i].get_k_b();

                for(size_t j=0; j<rns.re[i].get_no_educt_s(); ++j)
                    rate_f *= x[rns.re[i].get_educt_id(j)];

                for(size_t j=0; j<rns.re[i].get_no_product_s(); ++j)
                    rate_b *= x[rns.re[i].get_product_id(j)];

                for(size_t j=0; j<rns.re[i].get_no_educt_s(); ++j)
                    dxdt[rns.re[i].get_educt_id(j)] -= (rate_f - rate_b);

                for(size_t j=0; j<rns.re[i].get_no_product_s(); ++j)
                    dxdt[rns.re[i].get_product_id(j)] += (rate_f - rate_b);
             }

            for(size_t i=0; i<rns.sp.size(); ++i)
                if(rns.sp[i].is_constant())
	            dxdt[i]=0;
         }
    };

    /*
     * Calculation of concentration change (rhs) from species concentration
     * and thermodynamic values (energies).
     */

    struct stiff_system {
        reaction_network_system& rns;

        stiff_system(reaction_network_system& rns_) : rns(rns_)  {}

        // Function to calculate reaction rates of state <x> into <rates>.
        void calculate_rates(const vector_type &x, vector_type &rates) {
            vector_type a2(rns.re.size()), a3(rns.re.size());
            for(size_t j=0; j<rns.N.size2(); ++j) {
                a2(j) = 1;
                a3(j) = 1;

                for(size_t k=0; k<x.size(); ++k) {
                    if(rns.N_in(k,j) != 0)
                        a2(j) *= pow(x(k), rns.N_in(k,j));

                    if(rns.N_out(k,j) != 0)
                        a3(j) *= pow(x(k), rns.N_out(k,j));
                }
            }

            rates = element_prod(rns.e_m_bEa, element_prod(a2, rns.e_bs_in) - element_prod(a3, rns.e_bs_out));
        }

        // <x> - state / concentration;  <J> - Jacobi matrix (output);  dfdt - time partial differential (output)
        void operator()( const vector_type &x , vector_type &dxdt , double /* t */ ) {
            vector_type a2(rns.re.size()), a3(rns.re.size());
            for(size_t j=0; j<rns.N.size2(); ++j) {
                a2(j) = 1;
                a3(j) = 1;

                for(size_t l=0; l<rns.N_in_list[j].size(); ++l)
                        a2(j) *= pow(x(rns.N_in_list[j][l]), rns.N_in(rns.N_in_list[j][l],j));

                for(size_t l=0; l<rns.N_out_list[j].size(); ++l)
                        a3(j) *= pow(x(rns.N_out_list[j][l]), rns.N_out(rns.N_out_list[j][l],j));

            }

            dxdt = prod(rns.N, element_prod(rns.e_m_bEa, element_prod(a2, rns.e_bs_in) - element_prod(a3, rns.e_bs_out)));

            for(size_t k=0; k<x.size(); ++k)
                if(rns.sp[k].is_constant())
                    dxdt(k) = 0;
        }
    };

    /*
     * Calculation of Jacobi matrix from species concentration and thermodynamic
     * values (energies).
     */

    struct stiff_system_jacobi {
        const reaction_network_system& rns;

        stiff_system_jacobi(const reaction_network_system& rns_) : rns(rns_)  {}

        // <x> - state / concentration;  <J> - Jacobi matrix (output);  dfdt - time partial differential (output)
        void operator()( const vector_type & x  , matrix_type &J , const double & /* t */ , vector_type &dfdt ) {
            for(size_t i=0; i<x.size(); ++i) {
                for(size_t l=0; l<x.size(); ++l) {
                    double s(0);

                    for(size_t k=0; k<rns.re.size(); ++k) {
                        double p1(1), p2(1);    // p1 = \prod_{j \neq k} {x_j}^N_{jk}^\mathrm{in}
                                                // p2 = \prod_{j \neq k} {x_j}^N_{jk}^\mathrm{out}

                        for(size_t m=0; m<rns.N_in_list[k].size(); ++m) {
                            size_t j=rns.N_in_list[k][m];

                            if(j != l)
                                p1 *= x(j);
                        }

                        for(size_t m=0; m<rns.N_out_list[k].size(); ++m) {
                            size_t j=rns.N_out_list[k][m];

                            if(j != l && rns.N_out(j,k) != 0)
                                p2 *= x(j);
                        }

                        s += rns.N(i,k)*rns.e_m_bEa(k)*(rns.e_bs_in(k)*rns.N_in(l,k)*pow(x(l), rns.N_in(l,k)-1)*p1 -
                                                        rns.e_bs_out(k)*rns.N_out(l,k)*pow(x(l),rns.N_out(l,k)-1)*p2);
                    }

                    J(i,l) = s;
                    if(rns.sp[i].is_constant())
                        J(i,l) = 0;
                }
            }

            dfdt = boost::numeric::ublas::zero_vector<double>(x.size());
        }
    };

    /*
     * Constructor that loads the network file as well as the concentration file
     * (for initial concentration). Also all temporary values that are necessary to
     * solve (calculate rhs + Jacobi matrix) the ode faster are calculated.
     */

    reaction_network_system(const std::string& fn_network, const std::string& fn_concentration_)
        : fn_concentration(fn_concentration_), beta(1) {
        // load network into <re> and <sp>
        read_jrnf_reaction_n(fn_network, sp, re);

	// Remove 1-0 and 0-1 reactions. Such reactions are there to ballance flow through the boundary conditions
        std::remove_if (re.begin(), re.end(),
                        [] (reaction& re) -> bool {
                            return re.get_no_educt() == 1 && re.get_no_product() == 0 ||
                                   re.get_no_educt() == 0 && re.get_no_product() == 1; });

        // Calculate stoichiometric matrices, activation energies and their logarithms...
        N_in = N_out = boost::numeric::ublas::zero_matrix<double>(sp.size(), re.size());
        Ea =  e_bs_in = e_bs_out = e_m_bEa = boost::numeric::ublas::zero_vector<double>(re.size());
        mu0 = boost::numeric::ublas::zero_vector<double>(sp.size());

        // Load species standard gibbs energy / chemical energy into mu0
        for(size_t i=0; i<sp.size(); ++i)
            mu0(i) = sp[i].get_energy();

        // Load / calculate reaction specific quantities like activation Energy
        // <Ea> and joint energy of educts <m0_educts> or products <mu0_products>.
        for(size_t i=0; i<re.size(); ++i) {
            double mu0_educts(0), mu0_products(0);

            for(size_t j=0; j<re[i].get_no_educt_s(); ++j) {
                size_t id(re[i].get_educt_id(j)),  mul(re[i].get_educt_mul(j));
                N_in(id,i) += mul;
                mu0_educts += sp[id].get_energy()*mul;
            }

            for(size_t j=0; j<re[i].get_no_product_s(); ++j) {
                size_t id(re[i].get_product_id(j)),  mul(re[i].get_product_mul(j));
                N_out(id,i) += mul;
                mu0_products += sp[id].get_energy()*mul;
            }

            Ea(i) = re[i].get_activation()+max(mu0_educts,mu0_products);
            e_bs_in(i) = exp(beta*mu0_educts);
            e_bs_out(i) = exp(beta*mu0_products);
            e_m_bEa(i) = exp(-beta*Ea(i));
        }

        N = N_out-N_in;

        // Build lists (for speedup)
        for(size_t i=0; i<re.size(); ++i) {
            N_in_list.push_back(std::vector<size_t>());
            N_out_list.push_back(std::vector<size_t>());

            for(size_t j=0; j<sp.size(); ++j) {
                if(N_in(j,i) != 0)
                    N_in_list[i].push_back(j);

                if(N_out(j,i) != 0)
                    N_out_list[i].push_back(j);
            }
        }

        // Read last line of concentration file and write initial concentration to initial_con
        // and initial time to initial_t. When done, open the same file for appending...
        initial_con = boost::numeric::ublas::zero_vector<double>(sp.size());

        std::ifstream  data(fn_concentration.c_str());

        if(!data.good()) {
           std::cout << "Could not open concentration file: " << fn_concentration << std::endl;
           return;
        }

        std::string line;
        std::getline(data,line);       // dont want the header
        double last_msd=0;

        // Get line of actual data
        while(!std::getline(data,line).eof()) {
            std::stringstream ls(line);
            std::string cell;

            // read in value by value and count its position to place it in the
            // right variables
            size_t cnt=0;
            while(std::getline(ls,cell,',')) {
                std::stringstream in(cell);

                if(cnt == 0)
                    in >> initial_t ;
                else if(cnt == 1)
                    in >> last_msd;
                else if(cnt <= sp.size()+1)
                    in >> initial_con(cnt-2);
                else
                    std::cout << "Error at reading csv / concentration file!" << std::endl;

                ++cnt;
            }
        }

        std::cout << "Simulating file: " << fn_concentration << std::endl;
        std::cout << "Loaded concentration file with starting time " << initial_t << std::endl;
        data.close();
    }

    /*
     * To check some differential equations or the algorithm it might be usefull
     * to calculate the right hand side of the differential equation for a given
     * concentration vector.
     */

    void print_rhs() {
        vector_type x(initial_con);
        vector_type dxdt=boost::numeric::ublas::zero_vector<double>(sp.size());

        // calculate rhs - operator saves it do dxdt
        stiff_system(*this)(x, dxdt, 0);

        cout << "righthandside:" << endl;
        for(size_t i=0; i<dxdt.size(); ++i)
            cout << "/  " << dxdt(i) << "  /";
        cout << endl;
    }

    /*
     * Method for calling the system that uses energetics for calculating effective
     * rates and allows usage of a stiff solver - implicit stepper. System (f(x))
     * is defined in "stiff_system" (rhs) and "stiff_system_jacobi (Jacobi matrix).
     *
     * Tmax           - time up to that ode is solved
     * deltaT         - initial step size for addaptive stepper (does not relate to time output is written)
     * solve_implicit - do use implicit solver?
     * wint           - number of time output is written to output file between time of last
     *                  entry in concentration file and Tmax
     * write_log = 0  - linearly spaced times for output
     *           = 1  - logarithmically spaced (relative -> log(t-t0) is equidist.)
     *           = 2  - logarithmically spaced (absolute -> log(t) is equidist.)
     */

    void run(double Tmax=25000, double deltaT=0.1, bool solve_implicit=true,
             size_t wint=500, size_t write_log=0) {
        fstream out(fn_concentration.c_str(), std::ios_base::out | std::ios_base::app);
        out.precision(25);

        size_t t0 = time(NULL);

        // Init solver and start
        vector_type x(initial_con), last_con(initial_con);
        double last_write(initial_t);

        // Function that is called by stepper and writes concentration (<vec>) to file.
        auto write_state = [this, &out, t0, wint, deltaT, Tmax, &last_con, &last_write]( const vector_type &vec , const double t ) {
            // calculate mean square distance from last write
            double last_msd=0;
            for(size_t i=0; i<vec.size(); ++i)
                 last_msd += (vec[i]-last_con[i])*(vec[i]-last_con[i]);

            if(t != last_write)
                last_msd /= (vec.size()*(t-last_write));
            last_write=t;
            last_con=vec;

            // write time and msd ...
            out << t << "," << last_msd;
            // and concentrations
            for(size_t l=0; l<vec.size(); ++l)
                out << "," << vec(l);

            out << std::endl;
        };

        size_t step_no = 0;  // Number of steps done by integrator (for diagnostics)
        // Calculate vector of time points at which "write_state" will be called.
        std::vector<double> times( wint );
        for( size_t i=0 ; i<wint ; ++i )
            if(write_log == 0)
                // linearly spaced output times
                times[i] = initial_t + double(i+1)/(wint)*(Tmax-initial_t);
            else if(write_log == 1 || initial_t == 0)
                // logarithmically spaced output times (relative to t0)
                times[i] = initial_t + (exp(double(i+1))-1)/(exp(double(wint))-1)*(Tmax-initial_t);
            else
                // if t0!=0 and one want's logarithmically spaced times (relative to t=0)!
                times[i] = initial_t*pow(exp(1/double(wint)*log(Tmax/initial_t)),double(i+1));

        // Call integrator - returns number of steps. Either implicit or explicit stepper is used.
        step_no = solve_implicit ?
                      integrate_times( make_controlled< rosenbrock4< double > >( 1.0e-6 , 1.0e-6 ) ,
                                       make_pair( stiff_system(*this) , stiff_system_jacobi(*this) ) ,
                                       x , times, deltaT, write_state) :
                      integrate_times( make_controlled< runge_kutta_dopri5< vector_type > >( 1.0e-6 , 1.0e-6 ) ,
                                       stiff_system(*this) , x, times, deltaT, write_state);

        // print time + steps needed and close file
        size_t t1 = time(NULL);
        std::cout << "Run took " << t1-t0 << " seconds and " << step_no << " steps!" << std::endl;
        out.close();
   }

   /*
    * The fastint integration (calculating forward and backward rates directly
    * through reaction constants) needs the networks to be exanded what this
    * method ensures (multiple occurance of same educt product in same reaction
    * has to be explicit - "A + A" instead of "2 A").
    */

    void initialize_fastint() {
        for(size_t i=0; i<re.size(); ++i) {
            for(size_t j=0; j<re[i].get_no_educt_s(); ++j) {
                if(re[i].get_educt_mul(j) > 1.1) {
                    re[i].add_educt_s(re[i].get_educt_id(j), re[i].get_educt_mul(j)-1);
                    re[i].set_educt_mul(j,1.0);
                }
            }

            for(size_t j=0; j<re[i].get_no_product_s(); ++j) {
                if(re[i].get_product_mul(j) > 1.1) {
                    re[i].add_product_s(re[i].get_product_id(j), re[i].get_product_mul(j)-1);
                    re[i].set_product_mul(j,1.0);
                }
            }
        }
    }

    /*
     * Method for calling the simple / fast integrator that uses explicit stepping and
     * the reaction constants defined in the network object. System (f(x)) is defined
     * in "fast_system".
     * For description of method's parameters see method "run" above.
     */

    void run_fastint(double Tmax=25000, double deltaT=0.1, size_t wint=500, size_t write_log=0) {
        fstream out(fn_concentration.c_str(), std::ios_base::out | std::ios_base::app);
        out.precision(25);

        size_t t0 = time(NULL);

        vector_type x(initial_con), last_con(initial_con);
        double last_write(initial_t);

        // Function that is called by stepper and writes concentration (<vec>) to file.
        auto write_state = [this, &out, t0, wint, deltaT, Tmax, &last_con, &last_write]( const vector_type &vec , const double t ) {
            // calculate mean square distance from last write
            double last_msd=0;
            for(size_t i=0; i<vec.size(); ++i)
                 last_msd += (vec[i]-last_con[i])*(vec[i]-last_con[i]);
            if(t != last_write)
                last_msd /= (vec.size()*(t-last_write));
            last_write=t;
            last_con=vec;

            // write time + msd...
            out << t << "," << last_msd;
            // ...and all concentrations
            for(size_t l=0; l<vec.size(); ++l)
                out << "," << vec(l);

            // If msd (change of concentration) per time and per species is less
            // than 1e-20 but not zero this integrator stops early because the
            // assumption is that the network is converged. (Might not be true in
            // the strict sense for all imaginable networks that can be simulated!)
            if(last_msd > 1e-44 && last_msd < 1e-20) {
                std::cout << "Reached msd < 1e-20 condition - exiting early!" << std::endl;
                size_t t1 = time(NULL);
                std::cout << "Run took " << t1-t0 << " seconds!" << std::endl;
                exit(0);
            }

            out << std::endl;
        };

        // Calculate vector of time points at which "write_state" will be called.
        // For details see same code in method "run" above.
        std::vector<double> times( wint );
        for( size_t i=0 ; i<wint ; ++i )
            if(write_log == 0)
                times[i] = initial_t + double(i+1)/(wint)*(Tmax-initial_t);
            else if(write_log == 1 || initial_t == 0)
                times[i] = initial_t + (exp(double(i+1))-1)/(exp(double(wint))-1)*(Tmax-initial_t);
            else
                times[i] = initial_t*pow(exp(1/double(wint)*log(Tmax/initial_t)),double(i+1));

        // Call integrator - returns number of steps
        size_t step_no = integrate_times( make_dense_output< runge_kutta_dopri5< vector_type > >( 1.0e-6 , 1.0e-6 ) ,
                         fast_system(*this) , x, times, deltaT, write_state);

        // print time + steps needed and close file
        size_t t1 = time(NULL);
        std::cout << "Run took " << t1-t0 << " seconds and " << step_no << " steps!" << std::endl;
        out.close();
   }
};

/*
 * Main-Function. Manages commandline parameter and prints help screen.
 * All the real work is done through using the reaction_network_system
 * class (above).
 */

int main(int argc, const char *argv []){
    srand(time(0));
    std::cout << "odeint_rnet version " << odeint_rnet_version_string << " (commit:"
              << GIT_VERSION << ")" << std::endl;

    cl_para cl(argc, argv);  // create class to query command line parameters

    // Prints the right hand side of the ODE (f(x)) for diagnostic purposes.

    if(cl.have_param("print_rhs")) {
        std::string fn_network, fn_concentration;

        if(cl.have_param("net"))
            fn_network=cl.get_param("net");
        else {
            cout << "You have to give the name of reaction network by 'net'!" << endl;
            return 1;
        }

        if(cl.have_param("con"))
            fn_concentration=cl.get_param("con");
        else {
            cout << "You have to give the name of the concentration file by 'con'!" << endl;
            return 1;
        }

        // load network and concentration file and call function to print right hand side
        reaction_network_system rns = reaction_network_system(fn_network, fn_concentration);
        rns.print_rhs();
    }

    // Simulates / integrates a reaction network while holding the boundary point
    // species (constant=true) constant. Effective reaction rates are calculated
    // directly from chemical energies. Stiff solver is available with option
    // "solve_implicit". Very slow for larger networks (>50 species).

    if(cl.have_param("simulate")) {
        std::string fn_network, fn_concentration;

        if(cl.have_param("net"))
            fn_network=cl.get_param("net");
        else {
            cout << "You have to give the name of reaction network by 'net'!" << endl;
            return 1;
        }

        if(cl.have_param("con"))
            fn_concentration=cl.get_param("con");
        else {
            cout << "You have to give the name of the concentration file by 'con'!" << endl;
            return 1;
        }

        // Initial time-step for integrator
        double deltaT = cl.have_param("deltaT") ? cl.get_param_d("deltaT") : 0.1;
        // ODE is integrated up to <Tmax>
        double Tmax = cl.have_param("Tmax") ? cl.get_param_d("Tmax") : 25000;
        // Number of times the output is written between initial time and <Tmax>
        double wint = cl.have_param("wint") ? cl.get_param_d("wint") : 500;
        // Use implicit solver?
        bool solve_implicit=cl.have_param("solve_implicit");
        // Is output given logarithmically spaced? (Or linearly?)
        size_t write_log=0;
        if(cl.have_param("write_log"))
            write_log=1;
        if(cl.have_param("write_log_abs"))
            write_log=2;

        std::cout << "Parameters are deltaT=" << deltaT << "  and Tmax=" << Tmax
                  << "   wint=" << wint << std::endl;

        if(write_log == 1)
            cout << "Period of output will be equidistant on logscale (from t=t_0)." << endl;

        if(write_log == 2)
            cout << "Period of output will be equidistant on logscale (from t=0)." << endl;

        if(solve_implicit)
            cout << "Stiff solver is used!" << endl;

        // Load reaction network and concentration file
        reaction_network_system rns = reaction_network_system(fn_network, fn_concentration);

        // Simulate ODE (and write results to file)
        // The method gives diagnostic feedback to the user through console output
        rns.run(Tmax, deltaT, solve_implicit, wint, write_log);
    }

    // Simulates / integrates a reaction network while holding the boundary point
    // species (constant=true) constant. Official parameter to call this is
    // "fastint", the option to call using "simsim" is maintained for backward
    // compatibility. This integrator is faster as the one above, especially for
    // larger (>50 species) networks. It uses the reaction constants present in
    // the network and no stiff solver.

    if(cl.have_param("fastint") || cl.have_param("simsim")) {
        std::string fn_network, fn_concentration;

        if(cl.have_param("net"))
            fn_network=cl.get_param("net");
        else {
            cout << "You have to give the name of reaction network by 'net'!" << endl;
            return 1;
        }

        if(cl.have_param("con"))
            fn_concentration=cl.get_param("con");
        else {
            cout << "You have to give the name of the concentration file by 'con'!" << endl;
            return 1;
        }

        // Initial time-step for integrator
        double deltaT = cl.have_param("deltaT") ? cl.get_param_d("deltaT") : 0.1;
        // ODE is integrated up to <Tmax>
        double Tmax = cl.have_param("Tmax") ? cl.get_param_d("Tmax") : 25000;
        // Number of times the output is written between initial time and <Tmax>
        double wint = cl.have_param("wint") ? cl.get_param_d("wint") : 500;
        // Is output given logarithmically spaced? (Or linearly?)
        size_t write_log=0;
        if(cl.have_param("write_log"))
            write_log=1;
        if(cl.have_param("write_log_abs"))
            write_log=2;

        std::cout << "Parameters are deltaT=" << deltaT << "  and Tmax=" << Tmax
                  << "   wint=" << wint << std::endl;

        if(write_log == 1)
            cout << "Period of output will be equidistant on logscale (from t=t_0)." << endl;

        if(write_log == 2)
            cout << "Period of output will be equidistant on logscale (from t=0)." << endl;

        // Load reaction network and concentration file
        reaction_network_system rns = reaction_network_system(fn_network, fn_concentration);

        // Simulate ODE (and write results to file)
        // The method gives diagnostic feedback to the user through console output
        rns.initialize_fastint();
        rns.run_fastint(Tmax, deltaT, wint, write_log);
    }

    // User interface / help dialogue
    if(cl.have_param("help") || cl.have_param("info")) {
        cout << "          odeint_rnet" << endl;
        cout << "          ===========" << endl;
        cout << " call with parameter 'info' or 'help' for showing this screen" << endl;
        cout << "current version is " << odeint_rnet_version_string << " (commit:" << GIT_VERSION << ")" << endl;
        cout << endl;
        cout << "-'print_rhs': load reaction network 'net', (last) concentration in 'con'! and" << endl;
        cout << "prints  right hand side of ode for diagnostic purposes." << endl;
	cout << endl;
        cout << "-'simulate': load reaction network 'net' and simulate file 'con'!. Parameters are:" << endl;
        cout << "   x'deltaT': Integration interval (relevant for integrator, not for output!)" << endl;
        cout << "   x'Tmax': Time up to which the system is simulated" << endl;
        cout << "   x'wint': number of output times between Tstart and Tmax" << endl;
        cout << "   x'write_log': write output logarithmically spaced (from t=t_0)" << endl;
        cout << "   x'write_log_abs': write output logarithmically spaced from (t=0)" << endl;
        cout << "   x'solve_implicit': use stiff solver (might be slower for big nets!)" << endl;
        cout << endl;
        cout << "-'fastint': load reaction network 'net' and simulate file 'con'!. Other than above" << endl;
        cout << "   this integrator does not offer a stiff solver. It also does not use energies to" << endl;
        cout << "   calculate effective rates directly but uses reaction constants. Parameters are:" << endl;
        cout << "   x'deltaT': Integration interval (relevant for integrator, not for output!)" << endl;
        cout << "   x'Tmax': Time up to which the system is simulated" << endl;
        cout << "   x'wint': number of output times between Tstart and Tmax" << endl;
        cout << "   x'write_log': write output logarithmically spaced (from t=t_0)" << endl;
        cout << "   x'write_log_abs': write output logarithmically spaced from (t=0)" << endl;
	cout << endl;
    }
}


#include <armadillo>

// pybind11
#include <pybind11/numpy.h>
#include <pybind11/stl.h>
#include <pybind11/pybind11.h>
namespace py = pybind11;

// Carma
#include <carma/carma.h>

struct _neigh {
    arma::mat data;
    arma::colvec weights;
    arma::uword n_rows;
    _neigh(const arma::mat& m, const arma::mat& w):
        data(w.n_elem, m.n_cols, arma::fill::zeros),
        weights(w.n_elem, arma::fill::zeros),
        n_rows(0) {}
};

typedef _neigh neigh_t;

void neigh_vec(neigh_t& n,
               const arma::mat& m,
               const arma::uword m_nrow,
               const arma::uword m_ncol,
               const arma::mat& w,
               const arma::uword m_b,
               const arma::uword m_i,
               const arma::uword m_j) {

    arma::uword w_leg_i = w.n_rows / 2, w_leg_j = w.n_cols / 2;

    // copy values
    arma::uword k = 0;
    for (arma::uword i = 0; i < w.n_rows; ++i)
        for (arma::uword j = 0; j < w.n_cols; ++j)
            if (m_i + i >= w_leg_i && m_j + j >= w_leg_j &&
                m_i + i < w_leg_i + m_nrow &&
                m_j + j < w_leg_j + m_ncol &&
                arma::is_finite(m(m_j + m_i * m_ncol, 0))) {

                n.data(k, m_b) = m((m_j + j - w_leg_j) +
                    (m_i + i - w_leg_i) * m_ncol, m_b);
                n.weights(k++) = w(i, j);
            }
    n.n_rows = k;
}

arma::colvec nm_post_mean_x(const arma::colvec& x,
                            const arma::mat& sigma,
                            const arma::colvec& mu0,
                            const arma::mat& sigma0) {

    // inverse sigma0
    arma::mat inv_sum_weights(arma::size(sigma0));
    inv_sum_weights = arma::inv(sigma + sigma0);

    return sigma * inv_sum_weights * mu0 + sigma0 * inv_sum_weights * x;
}

arma::mat bayes_smoother(const arma::mat& m,
                         const arma::uword m_nrow,
                         const arma::uword m_ncol,
                         const arma::mat& w,
                         const arma::mat& sigma,
                         bool covar_sigma0) {

    // initialize result matrix
    arma::mat res(arma::size(m), arma::fill::none);
    res.fill(arma::datum::nan);

    // prior mean vector (neighbourhood)
    arma::colvec mu0(m.n_cols, arma::fill::zeros);

    // prior co-variance matrix (neighbourhood)
    arma::mat sigma0(arma::size(sigma), arma::fill::zeros);

    // neighbourhood
    neigh_t neigh(m, w);

    // compute values for each pixel
    for (arma::uword i = 0; i < m_nrow; ++i)
        for (arma::uword j = 0; j < m_ncol; ++j) {

            // fill neighbours values
            for (arma::uword b = 0; b < m.n_cols; ++b)
                neigh_vec(neigh, m, m_nrow, m_ncol, w, b, i, j);

            if (neigh.n_rows == 0) continue;

            // compute prior mean
            mu0 = arma::mean(neigh.data.rows(0, neigh.n_rows - 1), 0).as_col();

            // compute prior sigma
            sigma0 = arma::cov(neigh.data.rows(0, neigh.n_rows - 1), 0);

            // prior sigma covariance
            if (!covar_sigma0) {

                // clear non main diagonal cells
                sigma0.elem(arma::trimatu_ind(
                        arma::size(sigma0), 1)).fill(0.0);
                sigma0.elem(arma::trimatl_ind(
                        arma::size(sigma0), -1)).fill(0.0);
            }

            // evaluate multivariate bayesian
            res.row(j + i * m_ncol) =
                nm_post_mean_x(m.row(j + i * m_ncol).as_col(),
                               sigma, mu0, sigma0).as_row();
        }
    return res;
}

arma::mat kernel_smoother(const arma::mat& m,
                          const arma::uword m_nrow,
                          const arma::uword m_ncol,
                          const arma::mat& w,
                          const bool normalised) {

    // initialize result matrix
    arma::mat res(arma::size(m), arma::fill::none);
    res.fill(arma::datum::nan);

    // neighbourhood
    neigh_t neigh(m, w);

    // compute values for each pixel
    for (arma::uword b = 0; b < m.n_cols; ++b)
        for (arma::uword i = 0; i < m_nrow; ++i)
            for (arma::uword j = 0; j < m_ncol; ++j) {

                // fill neighbours values
                neigh_vec(neigh, m, m_nrow, m_ncol, w, b, i, j);

                if (neigh.n_rows == 0) continue;

                // normalise weight values
                if (normalised)
                    neigh.weights = neigh.weights /
                        arma::sum(neigh.weights.subvec(0, neigh.n_rows - 1));

                // compute kernel neighbourhood weighted mean
                res(j + i * m_ncol, b) = arma::as_scalar(
                    neigh.weights.subvec(0, neigh.n_rows - 1).as_row() *
                        neigh.data.col(b).subvec(0, neigh.n_rows - 1));
            }
    return res;
}

arma::mat bilinear_smoother(const arma::mat& m,
                            const arma::uword m_nrow,
                            const arma::uword m_ncol,
                            const arma::mat& w,
                            double tau) {

    // initialize result matrix
    arma::mat res(arma::size(m), arma::fill::none);
    res.fill(arma::datum::nan);

    // neighbourhood
    neigh_t neigh(m, w);

    // compute values for each pixel
    for (arma::uword b = 0; b < m.n_cols; ++b)
        for (arma::uword i = 0; i < m_nrow; ++i)
            for (arma::uword j = 0; j < m_ncol; ++j) {

                // fill neighbours values
                neigh_vec(neigh, m, m_nrow, m_ncol, w, b, i, j);

                if (neigh.n_rows == 0) continue;

                // compute bilinear weight
                arma::colvec bln_weight = neigh.weights % arma::normpdf(
                    neigh.data.col(b) - m(j + i * m_ncol, b), 0, tau);

                // normalise weight values
                bln_weight = bln_weight /
                    arma::sum(bln_weight.subvec(0, neigh.n_rows - 1));

                // compute kernel neighbourhood weighted mean
                res(j + i * m_ncol, b) = arma::as_scalar(
                    bln_weight.subvec(0, neigh.n_rows - 1).as_row() *
                    neigh.data.col(b).subvec(0, neigh.n_rows - 1));
            }
    return res;
}

void PyInit_smoothing(py::module &m) {
    m.def("bayes_smoother", &bayes_smoother, "Bayes Smoother");

    m.def("kernel_smoother", &kernel_smoother, "Kernel Smoother");

    m.def("bilinear_smoother", &bilinear_smoother, "Bilinear Smoother");
}


// Method for converting transfer function to system of first order differential equations for evaluation numerically is in Ogata, 4th ed, page 78.
// Use Mathematica StateSpaceModel instead of Ogata as basis for conversion of TF to State Space.

// Step response achieved by setting u(t) = 1, in derivative term, see MMA State Space model in Plasticity.nb, subsection on Utilities for preparing C++ code.

// Chop p0 in calling routine
// Bound p1 >= 1e-7 in calling routine

#include <iostream>
#include <vector>
#include <cassert>
#include <string>
#include <map>

#include <gsl/gsl_poly.h>
#include <gsl/gsl_complex_math.h>
#include <gsl/gsl_integration.h>
#include <gsl/gsl_errno.h>
#include <gsl/gsl_matrix.h>
#include <gsl/gsl_odeiv2.h>
#include <gsl/gsl_spline.h>

#include <boost/math/tools/polynomial.hpp>

#include "fmt/format.h"
#include "Performance.h"

// steps for interpolation, 5000 comes very close to Mathematica numerical results, check timing
const int steps = 5000;

unsigned debugPerformance = 0;
void debugPerformanceOn(unsigned d) {debugPerformance=d;}

double 	MaxRootRealPart(const std::vector<double>& coeff);
double 	H2sq(const std::vector<double>& num, const std::vector<double>& den);
double 	integrandH2(double w, void *p);
int 	deriv (double t, const double x[], double f[], void *p);
double 	stepPerformance(const std::vector<double>& num, const std::vector<double>& den, double tmax, signalType s);
double 	integrandStep(double y, void *p);

struct my_params {const std::vector<double>& num; const std::vector<double>& den;};
struct stepParam {gsl_spline *spline; gsl_interp_accel *acc;};

// Do chopping to set param to zero if close and check bounds before call

double performance(const std::vector<double>& num, const std::vector<double>& den,
					double gamma, double tmax, signalType s)
{
	if (MaxRootRealPart(den) > -1e-6) return 1e20;
	if (debugPerformance){
		double sf = stepPerformance(num, den, tmax, s);
		double hf = H2sq(num,den);
		std::map<signalType,std::string> signal =
			{{signalType::output, "output"}, {signalType::controlOpen,"cntrlO"}, {signalType::controlClosed,"cntrlC"}};
		std::cout << fmt::format("step = {}, h2 = {} for {}\n", sf, hf, signal[s]);
		return sf + gamma*hf;
	}
	else
		return stepPerformance(num, den, tmax, s) + gamma*H2sq(num,den);
}

// coeff of polynomial from low order to high order terms
double MaxRootRealPart(const std::vector<double>& coeff)
{
	auto n = coeff.size();
	// std::cout << "n = " << n << std::endl;
    std::vector<double> z(2*(n-1));
	gsl_poly_complex_workspace * w
	  = gsl_poly_complex_workspace_alloc (n);
	if (GSL_SUCCESS != gsl_poly_complex_solve(coeff.data(), n, w, z.data())){
		gsl_poly_complex_workspace_free (w);
		return 1.0;		// positive value is marked as unstable
	}
	gsl_poly_complex_workspace_free (w);

	double max = -1e20;
	for (unsigned i = 0; i < n-1; i++) {
	  // std::cout << fmt::format("z{} = {:+.18f} {:+.18f}\n", i, z[2*i], z[2*i+1]);
	  if (z[2*i] > max) max = z[2*i];
	}
	return max;
}

// In the following, check that integrand gets small at large frequency, required for H2 integral to converge. If not then if order of numerator = order of denominator, then integration will be infinite, because at high frequency, high order terms in numerator and denominator dominate, and so given same order of those terms, the system does not go to zero at high frequency. Correct for this by redefining the num/den transfer function as (num/den - high order num coeff/high order den coeff). This subtraction removes the Dirac delta impulse component of the dynamics, which corresponds in frequency space to a uniform addition to the absolute value of the tf at all frequencies equal to the amount subtracted off. Does not change dynamics, except to remove impulse at time zero acting over infinitesimal duration.

// returns square of H2 value
double H2sq(const std::vector<double>& num, const std::vector<double>& den)
{
	bool sizeFix = false;
	boost::math::tools::polynomial<double> poly_newnum;
	boost::math::tools::polynomial<double> poly_newden;
	if (num.size() == den.size()){			// deleting dirac impulse for equal-sized num & den
		double numHiOrder = num.back();
		double denHiOrder = den.back();
		boost::math::tools::polynomial<double> polyn(num.begin(), num.end());
		boost::math::tools::polynomial<double> polyd(den.begin(), den.end());
		poly_newnum = denHiOrder*polyn - numHiOrder*polyd;
		poly_newden = denHiOrder*polyd;
		// high order term in new numerator should be zero and automatically dropped, reducing size
		if (poly_newnum.data().size() != polyn.data().size() - 1){
			if (poly_newnum.data().back() < 1e-5)
				poly_newnum.data().pop_back();
			else
				assert(poly_newnum.data().size() == polyn.data().size() - 1);
		}
		sizeFix = true;
// 		std::cout << "Num = " << polyn << std::endl;
// 		std::cout << "Den = " << polyd << std::endl;
// 		std::cout << "NN  = " << poly_newnum << std::endl;
// 		std::cout << "ND  = " << poly_newden << std::endl;
	}

	gsl_function F;
	F.function = &integrandH2;
    // if size fixed above, use newnum and newden
    const std::vector<double>& n = (sizeFix) ? poly_newnum.data() : num;
    const std::vector<double>& d = (sizeFix) ? poly_newden.data() : den;
    my_params params {n, d};
    F.params = &params;
	if (integrandH2(1e10, F.params) > 1e-3) return 1e20; 	// should not happen, because den.size > num.size
	double result, error;
	gsl_integration_workspace *w = gsl_integration_workspace_alloc (1000);
	// std::cout << "h2 int start" << std::endl;
	if (GSL_SUCCESS != gsl_integration_qagi(&F, 0, 1e-7, 1000, w, &result, &error))
		result = 1e30;
	// std::cout << "h2 int end" << std::endl;
	gsl_integration_workspace_free(w);
	return result / (2.0*M_PI);
}

double integrandH2(double w, void *p)
{
	gsl_complex s;
	my_params *params = static_cast<my_params *>(p);
	const std::vector<double>& num = params->num;
	const std::vector<double>& den = params->den;
	GSL_SET_COMPLEX(&s, 0, w);
    int numSize = static_cast<int>(num.size());
    int denSize = static_cast<int>(den.size());
	return gsl_complex_abs2(gsl_complex_div(gsl_poly_complex_eval(num.data(), numSize, s),
								gsl_poly_complex_eval(den.data(), denSize, s)));
}

// deriv works for both output signal and control signal, see MMA file
// Possible speed up by optimizing deriv, see callgrind source code line by line weighting of time cost.

int deriv(double t, const double x[], double f[], void *p)
{
	(void)(t); /* avoid unused parameter warning */
	my_params *params = static_cast<my_params *>(p);
	const std::vector<double>& den = params->den;
	auto lastrow = den.size()-2;
	f[lastrow] = 0;
	for (unsigned i = 0; i < lastrow; ++i){
		f[i] = x[i+1];
		f[lastrow] -= den[i]*x[i];
	}
	f[lastrow] -= den[lastrow]*x[lastrow];
	f[lastrow] /= den.back();
	// next line for input, u(t) = 1 for step input
	f[lastrow] += 1;
	return GSL_SUCCESS;
}

// must make different ycoeff for output and control signals
// for output, same coeff for open and closed loops
// for control signals, diff coeff for open and closed loops
// see MMA file

double stepPerformance(const std::vector<double>& num, const std::vector<double>& den, double tmax, signalType s)
{
	double time[steps+1];
	double y[steps+1];
	my_params params = {num, den};
	auto dim = den.size()-1;	// dimensions of state space model for dynamics
	gsl_odeiv2_system sys = {deriv, NULL, dim, &params};
    // see GSL docs for alternative algorithms
	gsl_odeiv2_driver *d =
		gsl_odeiv2_driver_alloc_y_new (&sys, gsl_odeiv2_step_rkf45, 1e-6, 1e-6, 0.0);
	double t = 0.0;
	double denBack = den.back();

	// coefficients to get output, initialize with values for each case, max dim is 4, so use that
	double ycoeff[4];
	unsigned long ydim;		// number of output coefficients to get output y, varies by problem, set explicitly
	double yinputCoeff = 0;	// must add yinputCoeff * input to output; input=1 for the step response

	if (s == signalType::output){ // case of output signal, same for open and closed loops
		ydim = 3;
		ycoeff[0] = num[0]/denBack; ycoeff[1] = num[1]/denBack; ycoeff[2] = num[2]/denBack;
	}
	else {
		// cases of control signal, # coeff is always dimension of problem, dim
		ydim = dim;
		double numBack = num.back();
		yinputCoeff = numBack / denBack;
		if (s == signalType::controlOpen) { // control signal, open loop
			for (unsigned i = 0; i < ydim; ++i)
				ycoeff[i] = (num[i] - den[i]*yinputCoeff) / denBack;
		}
		else if (s == signalType::controlClosed) { // control signal, closed loop
			for (unsigned i = 0; i < ydim; ++i){
				ycoeff[i] = (num[i] - den[i]*yinputCoeff) / denBack;
				if (debugPerformance >= 2) std::cout << ycoeff[i] << " ";
			}
			if (debugPerformance >= 2) std::cout << yinputCoeff << std::endl;
			if (debugPerformance >= 2){
				for (auto& n : den) std::cout << -n/denBack << " ";
				std::cout << std::endl;
				for (auto& n : num) std::cout << n << " ";
				std::cout << std::endl;
				for (auto& n : den) std::cout << n << " ";
				std::cout << std::endl;
			}
		}
		else {
			assert(false);	// should not be here, signal must be one of above types
		}
	}

	time[0] = 0.0;
	// at time zero with step, dominated by infinite freq, so if order of den > num, then at time zero,
	// initial value is zero, if order den=num, then ratio of highest order terms,
	// order num>den should not happen
	y[0] = (den.size() > num.size()) ? 0.0 : num.back()/den.back();
    std::vector<double> x(den.size()-1,0); // 0 at end causes zero initialization
	for (int i = 1; i <= steps; i++){
		// add one to tmax, so that interpolation extends past boundary for integration
		// otherwise, can get an error when interpolating close to the upper boundary
	  	double ti = i * (tmax+1.0) / static_cast<double>(steps);
	  	int status = gsl_odeiv2_driver_apply(d, &t, ti, x.data());
		assert(status == GSL_SUCCESS);
		time[i] = t;
		y[i] = 0;
		for (unsigned j = 0; j < ydim; ++j)
			y[i] += ycoeff[j] * x[j];
		y[i] += yinputCoeff;
		// if (i % 100 == 0) std::cout << fmt::format("{:7.3f} {:8.6f}\n", time[i], y[i]);
		if (debugPerformance >= 3 && i % 100 == 0 && s == signalType::controlClosed)
			std::cout << fmt::format("{:7.3f} {:8.6f}\n", time[i], y[i]);
	}
	gsl_odeiv2_driver_free (d);

	gsl_interp_accel *acc = gsl_interp_accel_alloc();
    gsl_spline *spline = gsl_spline_alloc(gsl_interp_cspline, steps);
    if (GSL_SUCCESS != gsl_spline_init (spline, time, y, steps)){
    	gsl_spline_free (spline);
    	gsl_interp_accel_free (acc);
    	return 1e20;
    }

    // std::cout << gsl_spline_eval(spline, 19.9689, acc) << std::endl;

	stepParam integrParam = {spline, acc};
    gsl_function F;
	F.function = &integrandStep;
	F.params = &integrParam;
	double result, error;
	gsl_integration_workspace *w = gsl_integration_workspace_alloc (1000);
	// std::cout << "step int start" << std::endl;
	// using 1e-6 for abs and rel error
	double errtol = 1e-6;
	if (GSL_SUCCESS != gsl_integration_qag(&F, 0.0, tmax, errtol, errtol, 1000, 6, w, &result, &error)){
		gsl_integration_cquad_workspace *ctable = gsl_integration_cquad_workspace_alloc(200);
		if (GSL_SUCCESS != gsl_integration_cquad(&F, 0, tmax, errtol, errtol, ctable, &result, &error, NULL)){
     		boost::math::tools::polynomial<double> poly_newnum(num.begin(), num.end());
			boost::math::tools::polynomial<double> poly_newden(den.begin(), den.end());
			if (true){
				std::cout << "Step integration error, with num = " << poly_newnum << std::endl;
				std::cout << "                         and den = " << poly_newden << std::endl;
				std::cout << "                      and result = " << result << std::endl;
			}
    		result = 1e20;
    	}
    	gsl_integration_cquad_workspace_free(ctable);
	}
	// std::cout << "step int end" << std::endl;
	gsl_integration_workspace_free(w);

    gsl_spline_free (spline);
    gsl_interp_accel_free (acc);

	return result;
}

double integrandStep(double t, void *p)
{
	stepParam *params = static_cast<stepParam *>(p);
	double z = 1.0 - gsl_spline_eval(params->spline, t, params->acc);
	return z*z;
}


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

namespace ppl {
namespace mcmc {

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

    tree_cache += 4 * n_params;

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

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

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

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

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

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

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

    return output;
}

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

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

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

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

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

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

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

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

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

    while (1) {

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

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

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

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

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

    }

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

    return eps;
}

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

    // start timing warmup
    stopwatch_warmup.start();

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

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

        logger.printProgress(i);

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

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

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

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

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

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

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

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

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

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

            double log_sum_weight_subtree = math::neg_inf<double>;

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

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

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

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

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

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

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

            rho = rho_b + rho_f;

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

            if (!valid) break;

        } // end tree doubling for-loop

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

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

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

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

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

    } // end for-loop to sample 1 point

    // stop timing sampling
    stopwatch_sampling.stop();

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

} // namespace mcmc

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

} // namespace ppl


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

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

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

namespace biggles
{

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

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

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

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

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

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

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

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

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

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

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

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

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

}