#ifndef TRACE_RAYTRACE_
#define TRACE_RAYTRACE_
#include <headers/setup.hpp>
#include <headers/geometry.hpp>
#include <boost/iterator/indirect_iterator.hpp>

namespace raytrace {

    using geometry::Point;
    using geometry::Vector;
    using geometry::Sphere;
    using geometry::Line;
    using geometry::SphPoint;

    struct ExpGeometry
    {
        static constexpr float rho = 1000.0;                           //mm
        static constexpr float d = 60.0;                               //mm
        static constexpr float l1 = 195.0;                             //mm
        static constexpr float l2 = 185.0;                             //mm
        static constexpr float L = 60.0;                               //mm
    };

    class Surface
    {
        private:
            setup::Spline delta_;
            setup::Spline gamma_;                                       //Permettivity = 1 - delta + i * gamma
            double RMSHeight_;                                          //RMS Height [nm]
            double CorrLength_;                                         //Correlation length [um]
            double alpha_;                                              //alpha parameter in PSD ABC model

            double k(double wl) const {return 2 * Constants::pi / wl; }
            double mu0(double th0, double corl, double wl) const {return corl * 1e3 * pow(sin(th0), 2) / (2 * wl); }
            std::complex<double> muc(double corl, double wl) const {return corl * 1e3 * (1.0 - permettivity(wl)) / (2 * wl); }
            double F(double tau, double alpha) const {return 2.0 / sqrt(Constants::pi) * tgamma(alpha + 0.5) / tgamma(alpha) / pow(1 + tau * tau, alpha + 0.5); }
        public:
            Surface(double RMSHeight = Constants::RMSHeight,
            double CorrLength = Constants::CorrLength, double alpha = Constants::alpha, const std::string & str = "Al2O3.txt");
            std::complex<double> permettivity (double wl = Constants::WL) const;
            double Rf (double th0, double wl = Constants::WL) const;
            double TIS (double th0, double rmsh, double corl, double alpha, double wl = Constants::WL) const;
            double Indicatrix1D (double th, double th0, double wl = Constants::WL) const;
            double Indicatrix2D (double th, double phi, double th0, double wl = Constants::WL) const;
            double PSD1D (double p) const noexcept;                              //PSD 1D ([micrometer ^ -1]) [micrometer ^ 3]
            double PSD2D (double p1, double p2) const noexcept;                  //PSD 2D ([micrometer ^ -1]) [micrometer ^ 4]
            double CritAng(double wl = Constants::WL) const {return sqrt(1.0 - real(permettivity(wl))); }
            double RMSHeight() const noexcept {return RMSHeight_; }
            double & RMSHeight() noexcept {return RMSHeight_; }
            double CorrLength() const noexcept {return CorrLength_; }
            double & CorrLength() noexcept {return CorrLength_; }
            double alpha() const noexcept {return alpha_; }
    };

    class ExpSetup
    {
        private:
            double wl_, inc_ang_, src_dist_, det_dist_, src_l_;
            Surface surf_;
            std::vector<Sphere> sphs_;
        public:
            ExpSetup (double inc_ang = 0.0, double RMSHeight = Constants::RMSHeight, double CorrLength = Constants::CorrLength) :
            surf_(Surface(RMSHeight, CorrLength)), sphs_(std::vector<Sphere>{Sphere(0.0, 0.0, ExpGeometry::rho, ExpGeometry::rho, ExpGeometry::d, Sphere::LOWER)}), inc_ang_(inc_ang),
            wl_(Constants::WL), src_dist_(ExpGeometry::l1), det_dist_(ExpGeometry::l2), src_l_(ExpGeometry::L) {}

            const Sphere & substrate() const noexcept {return sphs_.front(); }
            Sphere & substrate() noexcept {return sphs_.front(); }
            const std::vector<Sphere> & spheres() const noexcept {return sphs_; }
            const Surface & surface() const noexcept {return surf_; }
            Surface & surface() noexcept {return surf_; }
            const double & IncAngle() const noexcept {return inc_ang_; }
            double & IncAngle() noexcept {return inc_ang_; }
            const double wavelength() const noexcept {return wl_; }
            const double source_distance() const noexcept {return src_dist_; }
            const double source_length() const noexcept {return src_l_; }
            const double detector_distance() const noexcept {return det_dist_; }
            void add_sphere(const Sphere & sph) noexcept {sphs_.emplace_back(sph); }
            void add_sphere(double x, double y, double height, double radius) noexcept;
            void add_sphere(const std::vector<Sphere> & sphs) noexcept {sphs_.insert(sphs_.end(), sphs.cbegin(), sphs.cend()); }
            void add_sphere(const std::vector<double> & x, const std::vector<double> & y, const std::vector<double> & height, const std::vector<double> & radius);
            void reset_sphere() noexcept {if(sphs_.size() > 1) sphs_.erase(sphs_.cbegin() + 1,sphs_.cend()); }
    };

    class Beam : public Line
    {
        private:
            static constexpr double intersect_limit = 1e-8;
        public:
            using Line::Line;
            double Delta (const Sphere & sph) const noexcept;
            std::vector<SphPoint> Intersect(const Sphere & sph) const;
            std::vector<SphPoint> Intersect(const std::vector<Sphere> & sphs) const;
            bool is_intersect (const Sphere & sph) const noexcept;
            bool is_intersect (const std::vector<Sphere> & sphs) const noexcept;
            double IncAng (const SphPoint & spt) const;
            Vector SpecVec (const SphPoint & spt) const;
            virtual Vector ScatVec(const SphPoint & pt, const Surface & surf, double wl) const = 0;
    };

    class Beam2D : public Beam
    {
        public:
            using Beam::Beam;
            virtual Vector ScatVec(const SphPoint & spt, const Surface & surf, double wl) const override;
    };

    class SphBeam2D : public Beam2D
    {
        public:
            using Beam2D::Beam2D;
            SphBeam2D (const Sphere & sph, double inc_ang, double src_dist);
            SphBeam2D (const ExpSetup & setup) : SphBeam2D(setup.substrate(), setup.IncAngle(), setup.source_distance()) {}
    };

    class PlaneBeam2D : public Beam2D
    {
        public:
            using Beam2D::Beam2D;
            PlaneBeam2D (const Sphere & sph, double inc_ang, double src_dist);
            PlaneBeam2D (const ExpSetup & setup) : PlaneBeam2D(setup.substrate(), setup.IncAngle(), setup.source_distance()) {}
    };

    class Beam3D : public Beam
    {
        public:
            using Beam::Beam;
            virtual Vector ScatVec(const SphPoint & spt, const Surface & surf, double wl) const override;
    };

    class SphBeam3D : public Beam3D
    {
        public:
            using Beam3D::Beam3D;
            SphBeam3D(const Sphere & sph, double inc_ang, double src_dist, double src_len);
            SphBeam3D(const ExpSetup & setup) : SphBeam3D(setup.substrate(), setup.IncAngle(), setup.source_distance(), setup.source_length()) {}
    };

    class PlaneBeam3D : public Beam3D
    {
        public:
            using Beam3D::Beam3D;
            PlaneBeam3D(const Sphere & sph, double inc_ang, double src_dist);
            PlaneBeam3D(const ExpSetup & setup) : PlaneBeam3D(setup.substrate(), setup.IncAngle(), setup.source_distance()) {}
    };

    class TestBeam3D : public Beam3D
    {
        public:
            using Beam3D::Beam3D;
            TestBeam3D(const Sphere & sph, double h, double src_dist);
            TestBeam3D(const ExpSetup & setup, double h) : TestBeam3D(setup.substrate(), h, setup.source_distance()) {}
    };

    struct TracePoint
    {
        Point det_point;
        bool is_scattered;
        bool is_transmitted;
        TracePoint (const Point & pt, bool is_tr, bool is_sc) : det_point(pt), is_transmitted(is_tr), is_scattered(is_sc) {}
    };

    class Trace
    {
        private:
            std::vector<std::unique_ptr<Beam>> trace_;
            bool is_transmitted_ {};
            bool is_scattered_ {};
        public:
            using iterator = boost::indirect_iterator<std::vector<std::unique_ptr<Beam>>::iterator>;
            using const_iterator = boost::indirect_iterator<std::vector<std::unique_ptr<Beam>>::const_iterator>;
            iterator begin() {return iterator(trace_.begin()); }
            iterator end() {return iterator(trace_.end()); }
            const_iterator begin() const {return const_iterator(trace_.begin()); }
            const_iterator end() const {return const_iterator(trace_.end()); }

            template <typename T, typename Object>
            Trace (T && line, const Surface & surf, const Object & sphs, double wl)
            {
                trace_.emplace_back(std::make_unique<std::decay_t<T>>(std::forward<T>(line)));
                is_transmitted_ = trace_.back()->is_intersect(sphs);
                if (is_transmitted_)
                {
                    double sw = setup::dist(setup::gen);
                    auto spts = trace_.back()->Intersect(sphs);
                    assert(spts.size() != 0);
                    auto next_spt_ = *std::min_element(spts.cbegin(), spts.cend(), [](const SphPoint & a, const SphPoint & b) {return a.point().y() < b.point().y(); });
                    is_transmitted_ = sw < surf.Rf(trace_.back()->IncAng(next_spt_), wl);
                    while (is_transmitted_)
                    {
                        if (sw < surf.TIS(trace_.back()->IncAng(next_spt_), surf.RMSHeight(), surf.CorrLength(), surf.alpha(), wl))
                        {
                            trace_.emplace_back(std::make_unique<std::decay_t<T>>(next_spt_.point(), trace_.back()->ScatVec(next_spt_, surf, wl)));
                            is_scattered_ = true;
                        }
                        else
                            trace_.emplace_back(std::make_unique<std::decay_t<T>>(next_spt_.point(), trace_.back()->SpecVec(next_spt_)));
                        spts = trace_.back()->Intersect(sphs);
                        if (spts.size() == 0)
                            break;
                        next_spt_ = *std::min_element(spts.cbegin(), spts.cend(), [](const SphPoint & a, const SphPoint & b) {return a.point().y() < b.point().y(); });
                        sw = setup::dist(setup::gen);
                        is_transmitted_ = sw < surf.Rf(trace_.back()->IncAng(next_spt_), wl);
                    }
                }
            }

            template <typename T>
            Trace (T && line, const ExpSetup & setup) : Trace(std::forward<T>(line), setup.surface(), setup.spheres(), setup.wavelength()) {}

            TracePoint det_res(double det_dist) const noexcept;
            bool is_transmitted() const noexcept {return is_transmitted_; }
            bool is_scattered() const noexcept {return is_scattered_; }
            int size() const noexcept {return trace_.size(); }
    };

    template <typename T>
    Trace make_trace(const ExpSetup & setup)
    {
        return Trace (T(setup), setup);
    }

}

#endif


#include <Eigen/Geometry>
#include <Eigen/LU>
#include <iostream>
#include <numeric>
#include <sps/plane.hpp>
#include <vector>

// This implementation is general but slow because it needs to solve linear systems. It is possible to implement this
// function much faster for specific plane types.
Eigen::VectorXd sps::AbstractPlane::CalcParameters(const Eigen::Vector2d& coords) const
{
    const Eigen::Vector2d              g_c(+0.0, +0.0);
    const std::vector<Eigen::Vector2d> g = {
        {-1.0, -1.0},
        {+1.0, -1.0},
        {+1.0, +1.0},
        {-1.0, +1.0},
    };

    const auto vertex_indices = [&]() -> std::pair<unsigned, unsigned> {
        if (coords(0) >= coords(1) && coords(0) <= -coords(1))
        {
            return {0, 1};
        }
        if (coords(0) >= coords(1) && coords(0) >= -coords(1))
        {
            return {1, 2};
        }
        if (coords(0) <= coords(1) && coords(0) >= -coords(1))
        {
            return {2, 3};
        }
        if (coords(0) <= coords(1) && coords(0) <= -coords(1))
        {
            return {3, 0};
        }

        assert(false);
        return {0, 0};
    }();

    // Prepare to calculate the barycentric coordinates
    const Eigen::Matrix3d R = [&]() {
        Eigen::Matrix3d R = Eigen::Matrix3d::Ones();

        R.block<2, 1>(0, 0) = g_c;
        R.block<2, 1>(0, 1) = g[std::get<0>(vertex_indices)];
        R.block<2, 1>(0, 2) = g[std::get<1>(vertex_indices)];

        return R;
    }();
    const Eigen::Vector3d r(coords(0), coords(1), 1.0);

    // Find the barycentric coordinates
    const Eigen::Vector3d barycentric_coords = Eigen::FullPivLU<Eigen::Matrix3d>(R).solve(r);

    return barycentric_coords(0) * m_center + barycentric_coords(1) * m_vertices[std::get<0>(vertex_indices)] +
           barycentric_coords(2) * m_vertices[std::get<1>(vertex_indices)];
}

Eigen::VectorXd sps::AbstractPlane::CalcGridParameters(const GridCellIndex&              grid_cell,
                                                       const unsigned int                num_candidates,
                                                       const double                      inter_level_scale,
                                                       const std::vector<GridCellIndex>& prev_grid_cells) const
{
    assert(num_candidates % 2 == 1);

    const int    grid_radius         = (num_candidates - 1) / 2;
    const double grid_to_coord_scale = (1.0 / static_cast<double>(grid_radius));

    // Find the coordinates in the top-level grid coordinate system
    auto zoom_transform = Eigen::Affine2d::Identity();
    for (auto& prev_grid_cell : prev_grid_cells)
    {
        const auto prev_unscaled_grid_coords = Eigen::Vector2d(static_cast<double>(std::get<0>(prev_grid_cell)),
                                                               static_cast<double>(std::get<1>(prev_grid_cell)));
        const auto center                    = grid_to_coord_scale * prev_unscaled_grid_coords;
        zoom_transform =
            zoom_transform * Eigen::Translation2d(center) * Eigen::UniformScaling<double>(inter_level_scale);
    }
    const auto unscaled_grid_coords =
        Eigen::Vector2d(static_cast<double>(std::get<0>(grid_cell)), static_cast<double>(std::get<1>(grid_cell)));
    const Eigen::Vector2d finest_level_coords   = grid_to_coord_scale * unscaled_grid_coords;
    const Eigen::Vector2d coarsest_level_coords = zoom_transform * finest_level_coords;

    return CalcParameters(coarsest_level_coords);
}

double sps::AbstractPlane::CalcArea() const
{
    const auto triangle_area_evaluator =
        [](const Eigen::VectorXd& v_0, const Eigen::VectorXd& v_1, const Eigen::VectorXd& v_2) {
            const Eigen::VectorXd r_0      = v_1 - v_0;
            const Eigen::VectorXd r_1      = v_2 - v_0;
            const double          dot_prod = r_0.dot(r_1);

            return 0.5 * std::sqrt(r_0.squaredNorm() * r_1.squaredNorm() - dot_prod * dot_prod);
        };

    std::array<double, 4> areas;

    areas[0] = triangle_area_evaluator(m_center, m_vertices[0], m_vertices[1]);
    areas[1] = triangle_area_evaluator(m_center, m_vertices[1], m_vertices[2]);
    areas[2] = triangle_area_evaluator(m_center, m_vertices[2], m_vertices[3]);
    areas[3] = triangle_area_evaluator(m_center, m_vertices[3], m_vertices[0]);

    return std::accumulate(areas.begin(), areas.end(), 0.0);
}


#pragma once

#include <Eigen/Core>
#include <Eigen/LU>
#include <cppad/cppad.hpp>

// copy of eigen_mat_inv from CppAD to make the reverse pass codegen compatible (not assume double scalar)

namespace tds {
template <class Base>
class atomic_eigen_mat_inv : public CppAD::atomic_base<Base> {
 public:
  // -----------------------------------------------------------
  // type of elements during calculation of derivatives
  typedef Base scalar;
  // type of elements during taping
  typedef CppAD::AD<scalar> ad_scalar;
  // type of matrix during calculation of derivatives
  typedef Eigen::Matrix<scalar, Eigen::Dynamic, Eigen::Dynamic, Eigen::RowMajor>
      matrix;
  // type of matrix during taping
  typedef Eigen::Matrix<ad_scalar, Eigen::Dynamic, Eigen::Dynamic,
                        Eigen::RowMajor>
      ad_matrix;
  /* %$$
  $subhead Constructor$$
  $srccode%cpp% */
  // constructor
  atomic_eigen_mat_inv(void)
      : CppAD::atomic_base<Base>("atom_eigen_mat_inv",
                                 CppAD::atomic_base<Base>::set_sparsity_enum) {}
  /* %$$
  $subhead op$$
  $srccode%cpp% */
  // use atomic operation to invert an AD matrix
  ad_matrix op(const ad_matrix& arg) {
    size_t nr = size_t(arg.rows());
    size_t ny = nr * nr;
    size_t nx = 1 + ny;
    assert(nr == size_t(arg.cols()));
    // -------------------------------------------------------------------
    // packed version of arg
    CPPAD_TESTVECTOR(ad_scalar) packed_arg(nx);
    packed_arg[0] = ad_scalar(nr);
    for (size_t i = 0; i < ny; i++) packed_arg[1 + i] = arg.data()[i];
    // -------------------------------------------------------------------
    // packed version of result = arg^{-1}.
    // This is an atomic_base function call that CppAD uses to
    // store the atomic operation on the tape.
    CPPAD_TESTVECTOR(ad_scalar) packed_result(ny);
    (*this)(packed_arg, packed_result);
    // -------------------------------------------------------------------
    // unpack result matrix
    ad_matrix result(nr, nr);
    for (size_t i = 0; i < ny; i++) result.data()[i] = packed_result[i];
    return result;
  }
  /* %$$
$head Private$$

$subhead Variables$$
$srccode%cpp% */
 private:
  // -------------------------------------------------------------
  // one forward mode vector of matrices for argument and result
  CppAD::vector<matrix> f_arg_, f_result_;
  // one reverse mode vector of matrices for argument and result
  CppAD::vector<matrix> r_arg_, r_result_;
  // -------------------------------------------------------------
  /* %$$
  $subhead forward$$
  $srccode%cpp% */
  // forward mode routine called by CppAD
  virtual bool forward(
      // lowest order Taylor coefficient we are evaluating
      size_t p,
      // highest order Taylor coefficient we are evaluating
      size_t q,
      // which components of x are variables
      const CppAD::vector<bool>& vx,
      // which components of y are variables
      CppAD::vector<bool>& vy,
      // tx [ j * (q+1) + k ] is x_j^k
      const CppAD::vector<scalar>& tx,
      // ty [ i * (q+1) + k ] is y_i^k
      CppAD::vector<scalar>& ty) {
    size_t n_order = q + 1;
    size_t nr = size_t(CppAD::Integer(tx[0 * n_order + 0]));
    size_t ny = nr * nr;
#ifndef NDEBUG
    size_t nx = 1 + ny;
#endif
    assert(vx.size() == 0 || nx == vx.size());
    assert(vx.size() == 0 || ny == vy.size());
    assert(nx * n_order == tx.size());
    assert(ny * n_order == ty.size());
    //
    // -------------------------------------------------------------------
    // make sure f_arg_ and f_result_ are large enough
    assert(f_arg_.size() == f_result_.size());
    if (f_arg_.size() < n_order) {
      f_arg_.resize(n_order);
      f_result_.resize(n_order);
      //
      for (size_t k = 0; k < n_order; k++) {
        f_arg_[k].resize(long(nr), long(nr));
        f_result_[k].resize(long(nr), long(nr));
      }
    }
    // -------------------------------------------------------------------
    // unpack tx into f_arg_
    for (size_t k = 0; k < n_order; k++) {  // unpack arg values for this order
      for (size_t i = 0; i < ny; i++)
        f_arg_[k].data()[i] = tx[(1 + i) * n_order + k];
    }
    // -------------------------------------------------------------------
    // result for each order
    // (we could avoid recalculting f_result_[k] for k=0,...,p-1)
    //
    f_result_[0] = f_arg_[0].inverse();
    for (size_t k = 1; k < n_order; k++) {  // initialize sum
      matrix f_sum = matrix::Zero(long(nr), long(nr));
      // compute sum
      for (size_t ell = 1; ell <= k; ell++)
        f_sum -= f_arg_[ell] * f_result_[k - ell];
      // result_[k] = arg_[0]^{-1} * sum_
      f_result_[k] = f_result_[0] * f_sum;
    }
    // -------------------------------------------------------------------
    // pack result_ into ty
    for (size_t k = 0; k < n_order; k++) {
      for (size_t i = 0; i < ny; i++)
        ty[i * n_order + k] = f_result_[k].data()[i];
    }
    // -------------------------------------------------------------------
    // check if we are computing vy
    if (vx.size() == 0) return true;
    // ------------------------------------------------------------------
    // This is a very dumb algorithm that over estimates which
    // elements of the inverse are variables (which is not efficient).
    bool var = false;
    for (size_t i = 0; i < ny; i++) var |= vx[1 + i];
    for (size_t i = 0; i < ny; i++) vy[i] = var;
    return true;
  }
  /* %$$
  $subhead reverse$$
  $srccode%cpp% */
  // reverse mode routine called by CppAD
  virtual bool reverse(
      // highest order Taylor coefficient that we are computing derivative of
      size_t q,
      // forward mode Taylor coefficients for x variables
      const CppAD::vector<Base>& tx,
      // forward mode Taylor coefficients for y variables
      const CppAD::vector<Base>& ty,
      // upon return, derivative of G[ F[ {x_j^k} ] ] w.r.t {x_j^k}
      CppAD::vector<Base>& px,
      // derivative of G[ {y_i^k} ] w.r.t. {y_i^k}
      const CppAD::vector<Base>& py) {
    size_t n_order = q + 1;
    size_t nr = size_t(CppAD::Integer(tx[0 * n_order + 0]));
    size_t ny = nr * nr;
#ifndef NDEBUG
    size_t nx = 1 + ny;
#endif
    //
    assert(nx * n_order == tx.size());
    assert(ny * n_order == ty.size());
    assert(px.size() == tx.size());
    assert(py.size() == ty.size());
    // -------------------------------------------------------------------
    // make sure f_arg_ is large enough
    assert(f_arg_.size() == f_result_.size());
    // must have previous run forward with order >= n_order
    assert(f_arg_.size() >= n_order);
    // -------------------------------------------------------------------
    // make sure r_arg_, r_result_ are large enough
    assert(r_arg_.size() == r_result_.size());
    if (r_arg_.size() < n_order) {
      r_arg_.resize(n_order);
      r_result_.resize(n_order);
      //
      for (size_t k = 0; k < n_order; k++) {
        r_arg_[k].resize(long(nr), long(nr));
        r_result_[k].resize(long(nr), long(nr));
      }
    }
    // -------------------------------------------------------------------
    // unpack tx into f_arg_
    for (size_t k = 0; k < n_order; k++) {  // unpack arg values for this order
      for (size_t i = 0; i < ny; i++)
        f_arg_[k].data()[i] = tx[(1 + i) * n_order + k];
    }
    // -------------------------------------------------------------------
    // unpack py into r_result_
    for (size_t k = 0; k < n_order; k++) {
      for (size_t i = 0; i < ny; i++)
        r_result_[k].data()[i] = py[i * n_order + k];
    }
    // -------------------------------------------------------------------
    // initialize r_arg_ as zero
    for (size_t k = 0; k < n_order; k++)
      r_arg_[k] = matrix::Zero(long(nr), long(nr));
    // -------------------------------------------------------------------
    // matrix reverse mode calculation
    //
    for (size_t k1 = n_order; k1 > 1; k1--) {
      size_t k = k1 - 1;
      // bar{R}_0 = bar{R}_0 + bar{R}_k (A_0 R_k)^T
      r_result_[0] +=
          r_result_[k] * f_result_[k].transpose() * f_arg_[0].transpose();
      //
      for (size_t ell = 1; ell <= k;
           ell++) {  // bar{A}_l = bar{A}_l - R_0^T bar{R}_k R_{k-l}^T
        r_arg_[ell] -= f_result_[0].transpose() * r_result_[k] *
                       f_result_[k - ell].transpose();
        // bar{R}_{k-l} = bar{R}_{k-1} - (R_0 A_l)^T bar{R}_k
        r_result_[k - ell] -=
            f_arg_[ell].transpose() * f_result_[0].transpose() * r_result_[k];
      }
    }
    r_arg_[0] -=
        f_result_[0].transpose() * r_result_[0] * f_result_[0].transpose();
    // -------------------------------------------------------------------
    // pack r_arg into px
    for (size_t k = 0; k < n_order; k++) {
      for (size_t i = 0; i < ny; i++)
        px[(1 + i) * n_order + k] = r_arg_[k].data()[i];
    }
    //
    return true;
  }
};
}  // namespace tds


#include <scitbx/array_family/boost_python/flex_fwd.h>

#include <boost/python/module.hpp>
#include <boost/python/def.hpp>
#include <boost/python/args.hpp>
#include <scitbx/array_family/boost_python/utils.h>
#include <scitbx/array_family/versa.h>
#include <scitbx/array_family/accessors/c_grid.h>
#include <string>

#include <fftw3.h>

namespace {

  namespace af = scitbx::af;

  void
  complex_to_complex_in_place(
    af::ref<std::complex<double> > const& data,
    int exp_sign)
  {
    SCITBX_ASSERT(exp_sign == FFTW_FORWARD || exp_sign == FFTW_BACKWARD);
    int n = static_cast<int>(data.size());
    fftw_complex *in = reinterpret_cast<fftw_complex*>(data.begin());
    fftw_plan p = fftw_plan_dft_1d(
      n, in, in, exp_sign, FFTW_ESTIMATE);
    fftw_execute(p);
    fftw_destroy_plan(p);
  }

  void
  complex_to_complex_3d_in_place(
    af::ref<std::complex<double>, af::c_grid<3> > const& data,
    int exp_sign)
  {
    SCITBX_ASSERT(exp_sign == FFTW_FORWARD || exp_sign == FFTW_BACKWARD);
    int nx = static_cast<int>(data.accessor()[0]);
    int ny = static_cast<int>(data.accessor()[1]);
    int nz = static_cast<int>(data.accessor()[2]);
    fftw_complex *in = reinterpret_cast<fftw_complex*>(data.begin());
    fftw_plan p = fftw_plan_dft_3d(
      nx, ny, nz, in, in, exp_sign, FFTW_ESTIMATE);
    fftw_execute(p);
    fftw_destroy_plan(p);
  }

  af::versa<std::complex<double>, af::flex_grid<> >
  real_to_complex_in_place(
    af::versa<double, af::flex_grid<> >& data)
  {
    af::boost_python::assert_0_based_1d(data.accessor());
    int m = static_cast<int>(data.accessor().all()[0]);
    int n = static_cast<int>(data.accessor().focus()[0]);
    SCITBX_ASSERT(m == 2*(n/2+1));
    double* in = data.begin();
    fftw_complex *out = reinterpret_cast<fftw_complex*>(in);
    fftw_plan p = fftw_plan_dft_r2c_1d(
      n, in, out, FFTW_ESTIMATE);
    fftw_execute(p);
    fftw_destroy_plan(p);
    return af::versa<std::complex<double>, af::flex_grid<> >(
      data.handle(), af::flex_grid<>(m/2));
  }

  af::versa<double, af::flex_grid<> >
  complex_to_real_in_place(
    af::versa<std::complex<double>, af::flex_grid<> >& data,
    int n)
  {
    af::boost_python::assert_0_based_1d(data.accessor());
    SCITBX_ASSERT(!data.accessor().is_padded());
    int m = static_cast<int>(data.accessor().all()[0]) * 2;
    SCITBX_ASSERT(m == 2*(n/2+1));
    fftw_complex *in = reinterpret_cast<fftw_complex*>(data.begin());
    double* out = reinterpret_cast<double*>(in);
    fftw_plan p = fftw_plan_dft_c2r_1d(
      n, in, out, FFTW_ESTIMATE);
    fftw_execute(p);
    fftw_destroy_plan(p);
    return af::versa<double, af::flex_grid<> >(
      data.handle(), af::flex_grid<>(m).set_focus(n));
  }

  af::versa<std::complex<double>, af::flex_grid<> >
  real_to_complex_3d_in_place(
    af::versa<double, af::flex_grid<> >& data)
  {
    af::boost_python::assert_0_based_3d(data.accessor());
    int mx = static_cast<int>(data.accessor().all()[0]);
    int my = static_cast<int>(data.accessor().all()[1]);
    int mz = static_cast<int>(data.accessor().all()[2]);
    int nx = static_cast<int>(data.accessor().focus()[0]);
    int ny = static_cast<int>(data.accessor().focus()[1]);
    int nz = static_cast<int>(data.accessor().focus()[2]);
    SCITBX_ASSERT(mx == nx);
    SCITBX_ASSERT(my == ny);
    SCITBX_ASSERT(mz == 2*(nz/2+1));
    double* in = data.begin();
    fftw_complex *out = reinterpret_cast<fftw_complex*>(in);
    fftw_plan p = fftw_plan_dft_r2c_3d(
      nx, ny, nz, in, out, FFTW_ESTIMATE);
    fftw_execute(p);
    fftw_destroy_plan(p);
    return af::versa<std::complex<double>, af::flex_grid<> >(
      data.handle(),
      af::flex_grid<>((af::adapt(af::tiny<int, 3>(mx,my,mz/2)))));
  }

  af::versa<double, af::flex_grid<> >
  complex_to_real_3d_in_place(
    af::versa<std::complex<double>, af::flex_grid<> >& data,
    af::tiny<int, 3> const& n)
  {
    af::boost_python::assert_0_based_3d(data.accessor());
    SCITBX_ASSERT(!data.accessor().is_padded());
    int mx = static_cast<int>(data.accessor().all()[0]);
    int my = static_cast<int>(data.accessor().all()[1]);
    int mz = static_cast<int>(data.accessor().all()[2]) * 2;
    int nx = n[0];
    int ny = n[1];
    int nz = n[2];
    SCITBX_ASSERT(mx == nx);
    SCITBX_ASSERT(my == ny);
    SCITBX_ASSERT(mz == 2*(nz/2+1));
    fftw_complex *in = reinterpret_cast<fftw_complex*>(data.begin());
    double* out = reinterpret_cast<double*>(in);
    fftw_plan p = fftw_plan_dft_c2r_3d(
      nx, ny, nz, in, out, FFTW_ESTIMATE);
    fftw_execute(p);
    fftw_destroy_plan(p);
    return af::versa<double, af::flex_grid<> >(
      data.handle(),
      af::flex_grid<>((af::adapt(af::tiny<int, 3>(nx,ny,mz))))
        .set_focus(af::adapt(n)));
  }

  void
  wrap_fftw3()
  {
    using namespace boost::python;
    scope().attr("fftw_version") = std::string(fftw_version);
    def("complex_to_complex_in_place", complex_to_complex_in_place, (
      arg("data"), arg("exp_sign")));
    def("complex_to_complex_3d_in_place", complex_to_complex_3d_in_place, (
      arg("data"), arg("exp_sign")));
    def("real_to_complex_in_place", real_to_complex_in_place, (
      arg("data")));
    def("complex_to_real_in_place", complex_to_real_in_place, (
      arg("data"), arg("n")));
    def("real_to_complex_3d_in_place", real_to_complex_3d_in_place, (
      arg("data")));
    def("complex_to_real_3d_in_place", complex_to_real_3d_in_place, (
      arg("data"), arg("n")));
  }

} // namespace <anonymous>

BOOST_PYTHON_MODULE(fftw3tbx_ext)
{
  wrap_fftw3();
}


#define ARMA_DONT_PRINT_ERRORS

#include <pybind11/pybind11.h>
#include <pybind11/numpy.h>
#include <armadillo>

namespace py = pybind11;

using namespace arma;

typedef py::array_t<double, py::array::f_style | py::array::forcecast> array_tf;
typedef py::array_t<double, py::array::c_style | py::array::forcecast> array_tc;

cube array_to_cube(array_tf m) {

    py::buffer_info _m_buff = m.request();
    int n_rows = _m_buff.shape[0];
    int n_cols = _m_buff.shape[1];
    int n_slices = _m_buff.shape[2];

    cube _m_arma((double *)_m_buff.ptr, n_rows, n_cols, n_slices);

    return _m_arma;
}

mat array_to_mat(array_tf m) {

    py::buffer_info _m_buff = m.request();
    int n_rows = _m_buff.shape[0];
    int n_cols = _m_buff.shape[1];

    mat _m_arma((double *)_m_buff.ptr, n_rows, n_cols);

    return _m_arma;
}

vec array_to_vec(array_tf m) {

    py::buffer_info _m_buff = m.request();
    int n_rows = _m_buff.shape[0];

    vec _m_vec((double *)_m_buff.ptr, n_rows);

    return _m_vec;
}

array_tf cube_to_array(cube m) {

    auto _m_array = array_tf({m.n_rows, m.n_cols, m.n_slices});

    py::buffer_info _m_buff = _m_array.request();
    std::memcpy(_m_buff.ptr, m.memptr(), sizeof(double) * m.n_rows * m.n_cols * m.n_slices);

    return _m_array;
}

array_tf mat_to_array(mat m) {

    auto _m_array = array_tf({m.n_rows, m.n_cols});

    py::buffer_info _m_buff = _m_array.request();
    std::memcpy(_m_buff.ptr, m.memptr(), sizeof(double) * m.n_rows * m.n_cols);

    return _m_array;
}

array_tf vec_to_array(vec m) {

    auto _m_array = array_tf({m.n_rows});

    py::buffer_info _m_buff = _m_array.request();
    std::memcpy(_m_buff.ptr, m.memptr(), sizeof(double) * m.n_rows);

    return _m_array;
}

py::tuple policy_divergence(array_tf _p_K, array_tf _p_kff, array_tf _p_sigma_ctl,
                            array_tf _q_K, array_tf _q_kff, array_tf _q_sigma_ctl,
                            array_tf _mu_x, array_tf _sigma_x,
                            int dm_state, int dm_act, int nb_steps) {

    cube p_K = array_to_cube(_p_K);
    mat p_kff = array_to_mat(_p_kff);
    cube p_sigma_ctl = array_to_cube(_p_sigma_ctl);

    cube q_K = array_to_cube(_q_K);
    mat q_kff = array_to_mat(_q_kff);
    cube q_sigma_ctl = array_to_cube(_q_sigma_ctl);

    mat mu_x  = array_to_mat(_mu_x);
    cube sigma_x = array_to_cube(_sigma_x);

    vec kl(nb_steps);

    for(int i = 0; i < nb_steps; i++) {
        mat q_lambda_ctl = inv_sympd(q_sigma_ctl.slice(i));

        mat diff_K = (q_K.slice(i) - p_K.slice(i)).t() * q_lambda_ctl * (q_K.slice(i) - p_K.slice(i));
        mat diff_crs = (q_K.slice(i) - p_K.slice(i)).t() * q_lambda_ctl * (- q_kff.col(i) + p_kff.col(i));
        mat diff_kff = (- q_kff.col(i) + p_kff.col(i)).t() * q_lambda_ctl * (- q_kff.col(i) + p_kff.col(i));

        kl(i) = as_scalar(0.5 * log( det(q_sigma_ctl.slice(i)) / det(p_sigma_ctl.slice(i)) )
		                  + 0.5 * trace(q_lambda_ctl * p_sigma_ctl.slice(i))
		                  - 0.5 * dm_act
		                  + 0.5 * trace(diff_K * sigma_x.slice(i))
		                  + 0.5 * mu_x.col(i).t() * diff_K * mu_x.col(i)
		                  - mu_x.col(i).t() * diff_crs
		                  + 0.5 * diff_kff);
    }

    array_tf _kl = vec_to_array(kl);

    py::tuple output = py::make_tuple(_kl);
    return output;
}

py::tuple gaussian_divergence(array_tf _mu_p, array_tf _sigma_p,
                              array_tf _mu_q, array_tf _sigma_q,
                              int dm_state, int nb_steps) {

    mat mu_p = array_to_mat(_mu_p);
    cube sigma_p = array_to_cube(_sigma_p);

    mat mu_q = array_to_mat(_mu_q);
    cube sigma_q = array_to_cube(_sigma_q);

    vec kl(nb_steps);

    for(int i = 0; i < nb_steps; i++) {
        mat lambda_q = inv_sympd(sigma_q.slice(i));

        vec diff = mu_q.col(i) - mu_p.col(i);
        double quad = as_scalar(diff.t() * lambda_q * diff);
        double _trace = as_scalar(trace(lambda_q * sigma_p.slice(i)));
        double log_det = as_scalar(log( det(sigma_q.slice(i)) / det(sigma_p.slice(i)) ));

        kl(i) = 0.5 * (_trace + quad + log_det - dm_state);
    }

    array_tf _kl = vec_to_array(kl);

    py::tuple output = py::make_tuple(_kl);
    return output;
}

py::tuple gaussian_interp_w2(array_tf _mu_q, array_tf _sigma_q,
                             array_tf _mu_p, array_tf _sigma_p,
                             double alpha, int dim, int nb_steps) {

    mat mu_q = array_to_mat(_mu_q);
    cube sigma_q = array_to_cube(_sigma_q);

    mat mu_p = array_to_mat(_mu_p);
    cube sigma_p = array_to_cube(_sigma_p);

    mat mu(dim, nb_steps);
    cube sigma(dim, dim, nb_steps);

    for(int i = 0; i < nb_steps; i++) {
        mu.col(i) = (1. - alpha) * mu_q.col(i) + alpha * mu_p.col(i);

        mat chol_sigma_q = real(sqrtmat(sigma_q.slice(i)));
        mat inv_chol_sigma_q = inv_sympd(chol_sigma_q);
        mat _sigma = (1. - alpha) * sigma_q.slice(i)
                     + alpha * real(sqrtmat(chol_sigma_q * sigma_p.slice(i) * chol_sigma_q));

        sigma.slice(i) = inv_chol_sigma_q * _sigma * _sigma * inv_chol_sigma_q;
    }

    // transform outputs to numpy
    array_tf _mu = mat_to_array(mu);
    array_tf _sigma = cube_to_array(sigma);

    py::tuple output = py::make_tuple(_mu, _sigma);
    return output;
}

py::tuple gaussian_interp_kl(array_tf _mu_q, array_tf _sigma_q,
                             array_tf _mu_p, array_tf _sigma_p,
                             double alpha, int dim, int nb_steps) {

    mat mu_q = array_to_mat(_mu_q);
    cube sigma_q = array_to_cube(_sigma_q);

    mat mu_p = array_to_mat(_mu_p);
    cube sigma_p = array_to_cube(_sigma_p);

    mat mu(dim, nb_steps);
    cube sigma(dim, dim, nb_steps);

    for(int i = 0; i < nb_steps; i++) {
        mat lambda_q = inv_sympd(sigma_q.slice(i));
        mat lambda_p = inv_sympd(sigma_p.slice(i));

        sigma.slice(i) = inv_sympd(alpha * lambda_p + (1. - alpha) * lambda_q);
        mu.col(i) = sigma.slice(i) * (alpha * lambda_p * mu_p.col(i) + (1. - alpha) * lambda_q * mu_q.col(i));
    }

    // transform outputs to numpy
    array_tf _mu = mat_to_array(mu);
    array_tf _sigma = cube_to_array(sigma);

    py::tuple output = py::make_tuple(_mu, _sigma);
    return output;
}

double quad_expectation(array_tf _mu, array_tf _sigma_s,
                        array_tf _Q, array_tf _q, double _q0) {

    vec mu  = array_to_vec(_mu);
    mat sigma_s = array_to_mat(_sigma_s);

    mat Q = array_to_mat(_Q);
    vec q = array_to_vec(_q);

	double result = as_scalar(mu.t() * Q * mu) + as_scalar(mu.t() * q) + _q0 + trace(Q * sigma_s);
	return result;
}

py::tuple cubature_forward_pass(array_tf _mu_x0, array_tf _sigma_x0,
                                array_tf _mu_param, array_tf _sigma_param, array_tf _sigma_dyn,
                                array_tf _K, array_tf _kff, array_tf _sigma_ctl,
                                int dm_state, int dm_act, int nb_steps) {

    // inputs
    vec mu_x0 = array_to_vec(_mu_x0);
    mat sigma_x0 = array_to_mat(_sigma_x0);

    mat mu_param = array_to_mat(_mu_param);
    cube sigma_param = array_to_cube(_sigma_param);
    cube sigma_dyn = array_to_cube(_sigma_dyn);

    mat A(dm_state, dm_state);
    mat B(dm_state, dm_act);
    vec c(dm_state);

    // cubature

    // particles for state, action, constant, ...
    // ... sigma_param, sigma_dyn
    int dm_augmented = dm_state + dm_act + 1 + dm_state;
    vec mu_augmented(dm_augmented);
    mat chol_sigma_augmented = 1e-8 * eye(dm_augmented, dm_augmented);

    mat input_cubature_points(dm_augmented, 2 * dm_augmented);
    mat output_cubature_points(dm_state, 2 * dm_augmented);

    cube K = array_to_cube(_K);
    mat kff = array_to_mat(_kff);
    cube sigma_ctl = array_to_cube(_sigma_ctl);

    // outputs
    mat mu_x(dm_state, nb_steps + 1);
    cube sigma_x(dm_state, dm_state, nb_steps + 1);

    mat mu_u(dm_act, nb_steps);
    cube sigma_u(dm_act, dm_act, nb_steps);

    mat mu_xu(dm_state + dm_act, nb_steps + 1);
    cube sigma_xu(dm_state + dm_act, dm_state + dm_act, nb_steps + 1);

    mu_x.col(0) = mu_x0;
    sigma_x.slice(0) = sigma_x0;

    for (int i = 0; i < nb_steps; i++) {
        // get matrices from parameter distribution mean vector
        // reshape here is appled column-wise
        mat At = mu_param(span(0, dm_state * dm_state - 1), i);
        mat Bt = mu_param(span(dm_state * dm_state,
                               dm_state * dm_state + dm_state * dm_act - 1), i);
        vec ct = mu_param(span(dm_state * dm_state + dm_state * dm_act,
                               dm_state * dm_state + dm_state * dm_act + dm_state - 1), i);

        A = reshape(At, size(A));
        B = reshape(Bt, size(B));
        c = reshape(ct, size(c));

        // mu_u = K * mu_x + k
        mu_u.col(i) = K.slice(i) * mu_x.col(i) + kff.col(i);

        // sigma_u = sigma_ctl + K * sigma_x * K_T
        sigma_u.slice(i) = sigma_ctl.slice(i) + K.slice(i) * sigma_x.slice(i) * K.slice(i).t();
        sigma_u.slice(i) = 0.5 * (sigma_u.slice(i) + sigma_u.slice(i).t());
        sigma_u.slice(i) += 1e-8 * eye(dm_act, dm_act);

        // sigma_xu =   [[sigma_x,      sigma_x * K_T],
        //               [K * sigma_x,    sigma_u    ]]
        sigma_xu.slice(i) = join_vert(join_horiz(sigma_x.slice(i), sigma_x.slice(i) * K.slice(i).t()),
                                      join_horiz(K.slice(i) * sigma_x.slice(i), sigma_u.slice(i)));
        sigma_xu.slice(i) = 0.5 * (sigma_xu.slice(i) + sigma_xu.slice(i).t());
        sigma_xu.slice(i) += 1e-8 * eye(dm_state + dm_act, dm_state + dm_act);

        // mu_xu =  [[mu_x],
        //           [mu_u]],
        mu_xu.col(i) = join_vert(mu_x.col(i), mu_u.col(i));

        // form augmented state mean and covariance
        mu_augmented = join_vert(mu_xu.col(i), ones(1), zeros<vec>(dm_state));

        chol_sigma_augmented.fill(0.0);
        chol_sigma_augmented.submat(0, 0, dm_state + dm_act - 1, dm_state + dm_act - 1) = chol(symmatu(sigma_xu.slice(i)), "lower");
        chol_sigma_augmented.submat(dm_state + dm_act + 1, dm_state + dm_act + 1, dm_augmented - 1, dm_augmented - 1) = eye(dm_state, dm_state);

        // calculate cubature points
        input_cubature_points = join_horiz(chol_sigma_augmented, - chol_sigma_augmented);
        input_cubature_points *= sqrt(dm_augmented);
        input_cubature_points.each_col() += mu_augmented;

        // propagate cubature points
         for (int j = 0; j < 2 * dm_augmented; j++){
            vec vec_xu = input_cubature_points.col(j);

            mat mat_xu = kron(vec_xu(span(0, dm_state + dm_act)).t(), eye(dm_state, dm_state));
            mat total_covar = sigma_dyn.slice(i) + mat_xu * sigma_param.slice(i) * mat_xu.t();
            total_covar = 0.5 * (total_covar + total_covar.t());

            mat chol_covar = chol(symmatu(total_covar), "lower");
            output_cubature_points.col(j) = join_horiz(A, B, c, chol_covar) * vec_xu;
        }

        // estimate new mean and covariance
        mu_x.col(i+1) = mean(output_cubature_points, 1);

        output_cubature_points.each_col() -= mu_x.col(i+1);

        sigma_x.slice(i+1).fill(0.0);
        output_cubature_points.each_col([&sigma_x, i] (vec& col){
            sigma_x.slice(i+1) += col * col.t();
        });
        sigma_x.slice(i+1) /= 2 * dm_augmented;
        sigma_x.slice(i+1) = 0.5 * (sigma_x.slice(i+1) + sigma_x.slice(i+1).t());

        if(i == nb_steps - 1) {
            mu_xu.col(i+1) = join_vert(mu_x.col(i+1), zeros<vec>(dm_act));
            sigma_xu.slice(i+1).submat(0, 0, dm_state - 1, dm_state - 1) = sigma_x.slice(i+1);
        }
    }

    // transform outputs to numpy
    array_tf _mu_x = mat_to_array(mu_x);
    array_tf _sigma_x = cube_to_array(sigma_x);
    array_tf _mu_u =  mat_to_array(mu_u);
    array_tf _sigma_u = cube_to_array(sigma_u);
    array_tf _mu_xu =  mat_to_array(mu_xu);
    array_tf _sigma_xu = cube_to_array(sigma_xu);

    py::tuple output = py::make_tuple(_mu_x, _sigma_x, _mu_u, _sigma_u, _mu_xu, _sigma_xu);
    return output;
}

py::tuple policy_augment_cost(array_tf _Cxx, array_tf _cx, array_tf _Cuu,
                              array_tf _cu, array_tf _Cxu, array_tf _c0,
                              array_tf _K, array_tf _kff, array_tf _sigma_ctl,
                              array_tf _alpha, int dm_state, int dm_act, int nb_steps) {

    // inputs
    cube Cxx = array_to_cube(_Cxx);
    mat cx = array_to_mat(_cx);
    cube Cuu = array_to_cube(_Cuu);
    mat cu = array_to_mat(_cu);
    cube Cxu = array_to_cube(_Cxu);
    vec c0 = array_to_vec(_c0);

    cube K = array_to_cube(_K);
    mat kff = array_to_mat(_kff);
    cube sigma_ctl = array_to_cube(_sigma_ctl);

    vec alpha = array_to_vec(_alpha);

    // outputs
    cube agCxx(dm_state, dm_state, nb_steps + 1);
    mat agcx(dm_state, nb_steps + 1);
    cube agCuu(dm_act, dm_act, nb_steps + 1);
    mat agcu(dm_act, nb_steps + 1);
    cube agCxu(dm_state, dm_act, nb_steps + 1);
    vec agc0(nb_steps + 1);

    for (int i = 0; i < nb_steps; i++) {
        mat lambda_ctl = inv_sympd(sigma_ctl.slice(i));

        agCxx.slice(i) = Cxx.slice(i) + 0.5 * alpha(i) * K.slice(i).t() * lambda_ctl * K.slice(i);
        agCuu.slice(i) = Cuu.slice(i) + 0.5 * alpha(i) * lambda_ctl;
        agCxu.slice(i) = Cxu.slice(i) - 0.5 * alpha(i) * K.slice(i).t() * lambda_ctl;
        agcx.col(i) = cx.col(i) + alpha(i) * K.slice(i).t() * lambda_ctl * kff.col(i);
        agcu.col(i) = cu.col(i) - alpha(i) * lambda_ctl * kff.col(i);
        agc0(i) = as_scalar(c0(i) + 0.5 * alpha(i) * log( det(2. * datum::pi * sigma_ctl.slice(i)) )
                            + 0.5 * alpha(i) * kff.col(i).t() * lambda_ctl * kff.col(i));
    }

    // last time step
    agCxx.slice(nb_steps) = Cxx.slice(nb_steps);
    agcx.col(nb_steps) = cx.col(nb_steps);
    agCuu.slice(nb_steps) = Cuu.slice(nb_steps);
    agcu.col(nb_steps) = cu.col(nb_steps);
    agCxu.slice(nb_steps) = Cxu.slice(nb_steps);
    agc0(nb_steps) = c0(nb_steps);

    // transform outputs to numpy
    array_tf _agCxx = cube_to_array(agCxx);
    array_tf _agcx = mat_to_array(agcx);
    array_tf _agCuu =  cube_to_array(agCuu);
    array_tf _agcu = mat_to_array(agcu);
    array_tf _agCxu =  cube_to_array(agCxu);
    array_tf _agc0 = vec_to_array(agc0);

    py::tuple output = py::make_tuple(_agCxx, _agcx, _agCuu, _agcu, _agCxu, _agc0);
    return output;
}

py::tuple policy_backward_pass(array_tf _Cxx, array_tf _cx, array_tf _Cuu,
                               array_tf _cu, array_tf _Cxu, array_tf _c0,
                               array_tf _mu_param, array_tf _sigma_param, array_tf _sigma_dyn,
                               array_tf _alpha, int dm_state, int dm_act, int nb_steps) {

    // inputs
    cube Cxx = array_to_cube(_Cxx);
    mat cx = array_to_mat(_cx);
    cube Cuu = array_to_cube(_Cuu);
    mat cu = array_to_mat(_cu);
    cube Cxu = array_to_cube(_Cxu);
    vec c0 = array_to_vec(_c0);

    mat mu_param = array_to_mat(_mu_param);
    cube sigma_param = array_to_cube(_sigma_param);
    cube sigma_dyn = array_to_cube(_sigma_dyn);

    vec alpha = array_to_vec(_alpha);

    mat A(dm_state, dm_state);
    mat B(dm_state, dm_act);
    vec c(dm_state);

    mat P(dm_state + dm_act + 1, dm_state + dm_act + 1);
    mat Pxx(dm_state, dm_state);
    mat Pxu(dm_state, dm_act);
    mat Puu(dm_act, dm_act);
    vec px(dm_state);
    vec pu(dm_act);
    double p0;

    // outputs
    cube Q(dm_state + dm_act, dm_state + dm_act, nb_steps);
    cube Qxx(dm_state, dm_state, nb_steps);
    cube Qux(dm_act, dm_state, nb_steps);
    cube Quu(dm_act, dm_act, nb_steps);
    cube Quu_inv(dm_act, dm_act, nb_steps);
    mat qx(dm_state, nb_steps);
    mat qu(dm_act, nb_steps);
    vec q0(nb_steps);

    cube V(dm_state, dm_state, nb_steps + 1);
    mat v(dm_state, nb_steps + 1);
    vec v0(nb_steps + 1);

    cube K(dm_act, dm_state, nb_steps);
    mat kff(dm_act, nb_steps);
    cube sigma_ctl(dm_act, dm_act, nb_steps);
    cube lambda_ctl(dm_act, dm_act, nb_steps);

    int _diverge = -1;

    // last time step
    V.slice(nb_steps) = Cxx.slice(nb_steps);
    v.col(nb_steps) = cx.col(nb_steps);
    v0(nb_steps) = c0(nb_steps);

	for(int i = nb_steps - 1; i>= 0; --i)
	{
        // get matrices from parameter distribution mean vector
        // reshape here is appled column-wise
        mat At = mu_param(span(0, dm_state * dm_state - 1), i);
        mat Bt = mu_param(span(dm_state * dm_state, dm_state * dm_state + dm_state * dm_act - 1), i);
        vec ct = mu_param(span(dm_state * dm_state + dm_state * dm_act, dm_state * dm_state + dm_state * dm_act + dm_state - 1), i);

        A = reshape(At, size(A));
        B = reshape(Bt, size(B));
        c = reshape(ct, size(c));

        // extra terms due to parameter distribution
         for (int j = 0; j < dm_state + dm_act + 1; j++){
            for (int k = 0; k < dm_state + dm_act + 1; k++){
                P(j, k) = trace(sigma_param.slice(i).submat(j * dm_state, k * dm_state,
                                                            (j + 1) * dm_state - 1, (k + 1) * dm_state - 1) * V.slice(i+1));
            }
         }

        Pxx = P.submat(0, 0, dm_state - 1, dm_state - 1);
        Puu = P.submat(dm_state, dm_state, dm_state + dm_act - 1, dm_state + dm_act - 1);
        Pxu = P.submat(0, dm_state, dm_state - 1, dm_state + dm_act - 1);

        px = P.submat(0, dm_state + dm_act, dm_state - 1, dm_state + dm_act);
        pu = P.submat(dm_state, dm_state + dm_act, dm_state + dm_act - 1, dm_state + dm_act);
        p0 = P(dm_state + dm_act, dm_state + dm_act);

        Qxx.slice(i) = - (Cxx.slice(i) + A.t() * V.slice(i+1) * A + Pxx) / alpha(i);
        Quu.slice(i) = - (Cuu.slice(i) + B.t() * V.slice(i+1) * B + Puu) / alpha(i);
        Qux.slice(i) = - (Cxu.slice(i) + A.t() * V.slice(i+1) * B + Pxu).t() / alpha(i);

        qu.col(i) = - (cu.col(i) + 2.0 * B.t() * V.slice(i+1) * c + B.t() * v.col(i+1) + 2. * pu) / alpha(i);
        qx.col(i) = - (cx.col(i) + 2.0 * A.t() * V.slice(i+1) * c + A.t() * v.col(i+1) + 2. * px) / alpha(i);
        q0(i) = - as_scalar(c0(i) + v0(i+1) + c.t() * V.slice(i+1) * c +
                            + trace(V.slice(i+1) * sigma_dyn.slice(i)) + v.col(i+1).t() * c + p0) / alpha(i);

        if ((Quu.slice(i)).is_sympd()) {
            _diverge = i;
            break;
        }

        Quu_inv.slice(i) = inv(Quu.slice(i));
        K.slice(i) = - Quu_inv.slice(i) * Qux.slice(i);
        kff.col(i) = - 0.5 * Quu_inv.slice(i) * qu.col(i);

        sigma_ctl.slice(i) = - 0.5 * Quu_inv.slice(i);
        sigma_ctl.slice(i) = 0.5 * (sigma_ctl.slice(i).t() + sigma_ctl.slice(i));

        lambda_ctl.slice(i) = - (Quu.slice(i).t() + Quu.slice(i));
        lambda_ctl.slice(i) = 0.5 * (lambda_ctl.slice(i).t() + lambda_ctl.slice(i));

        V.slice(i) = - alpha(i) * (Qxx.slice(i) + Qux.slice(i).t() * K.slice(i));
        V.slice(i) = 0.5 * (V.slice(i) + V.slice(i).t());

        v.col(i) = - alpha(i) * (qx.col(i) + 2. * Qux.slice(i).t() * kff.col(i));
        v0(i) = - alpha(i) * (as_scalar(0.5 * qu.col(i).t() * kff.col(i)) + q0(i)
                              + 0.5 * (dm_act * log (2. * datum::pi) - log(det(- 2. * Quu.slice(i)))));
	}

    // transform outputs to numpy
    array_tf _Qxx = cube_to_array(Qxx);
    array_tf _Qux = cube_to_array(Qux);
    array_tf _Quu = cube_to_array(Quu);

    array_tf _qx = mat_to_array(qx);
    array_tf _qu = mat_to_array(qu);
    array_tf _q0 = mat_to_array(q0);

    array_tf _V = cube_to_array(V);
    array_tf _v = mat_to_array(v);
    array_tf _v0 = vec_to_array(v0);

    array_tf _K = cube_to_array(K);
    array_tf _kff = mat_to_array(kff);
    array_tf _sigma_ctl = cube_to_array(sigma_ctl);

    py::tuple output =  py::make_tuple(_Qxx, _Qux, _Quu, _qx, _qu, _q0,
                                        _V, _v, _v0,
                                        _K, _kff, _sigma_ctl, _diverge);

    return output;
}

py::tuple parameter_augment_cost(array_tf _mu_nominal, array_tf _sigma_nominal,
                                 double beta, int dm_param, int nb_steps) {

    // inputs
    mat mu_nominal = array_to_mat(_mu_nominal);
    cube sigma_nominal = array_to_cube(_sigma_nominal);

    // outputs
    cube agCxx(dm_param, dm_param, nb_steps);
    mat agcx(dm_param, nb_steps);
    vec agc0(nb_steps);

    for (int i = 0; i < nb_steps; i++) {
        mat lambda_nominal = inv_sympd(sigma_nominal.slice(i));

        agCxx.slice(i) = 0.5 * beta * lambda_nominal;
        agcx.col(i) = - beta * lambda_nominal * mu_nominal.col(i);
        agc0(i) = as_scalar(0.5 * beta * log( det(2. * datum::pi * sigma_nominal.slice(i)) )
                            + 0.5 * beta * mu_nominal.col(i).t() * lambda_nominal * mu_nominal.col(i));
    }

    // transform outputs to numpy
    array_tf _agCxx = cube_to_array(agCxx);
    array_tf _agcx = mat_to_array(agcx);
    array_tf _agc0 = vec_to_array(agc0);

    py::tuple output = py::make_tuple(_agCxx, _agcx, _agc0);
    return output;
}

py::tuple regularized_parameter_augment_cost(array_tf _mu_last, array_tf _sigma_last,
                                             double eta, int dm_param, int nb_steps) {

    // inputs
    mat mu_last = array_to_mat(_mu_last);
    cube sigma_last = array_to_cube(_sigma_last);

    // outputs
    cube agCxx(dm_param, dm_param, nb_steps);
    mat agcx(dm_param, nb_steps);
    vec agc0(nb_steps);

    for (int i = 0; i < nb_steps; i++) {
        mat lambda_last = inv_sympd(sigma_last.slice(i));

        agCxx.slice(i) = 0.5 * eta * lambda_last;
        agcx.col(i) = - eta * lambda_last * mu_last.col(i);
        agc0(i) = as_scalar(0.5 * eta * log( det(2. * datum::pi * sigma_last.slice(i)) )
                            + 0.5 * eta * mu_last.col(i).t() * lambda_last * mu_last.col(i));
    }

    // transform outputs to numpy
    array_tf _agCxx = cube_to_array(agCxx);
    array_tf _agcx = mat_to_array(agcx);
    array_tf _agc0 = vec_to_array(agc0);

    py::tuple output = py::make_tuple(_agCxx, _agcx, _agc0);
    return output;
}

py::tuple parameter_backward_pass(array_tf _mu_x, array_tf _sigma_x,
                                  array_tf _K, array_tf _kff, array_tf _sigma_ctl, array_tf _sigma_dyn,
                                  array_tf _cx, array_tf _Cxx, array_tf _Cuu,
                                  array_tf _cu, array_tf _Cxu, array_tf _c0,
                                  array_tf _agCpp, array_tf _agcp, array_tf _agc0,
                                  double beta, double eta, int dm_state, int dm_act, int dm_param, int nb_steps) {

    // inputs
    mat mu_x = array_to_mat(_mu_x);
    cube sigma_x = array_to_cube(_sigma_x);

    cube K = array_to_cube(_K);
    mat kff = array_to_mat(_kff);
    cube sigma_ctl = array_to_cube(_sigma_ctl);

    cube sigma_dyn = array_to_cube(_sigma_dyn);

    cube Cxx = array_to_cube(_Cxx);
    mat cx = array_to_mat(_cx);
    cube Cuu = array_to_cube(_Cuu);
    mat cu = array_to_mat(_cu);
    cube Cxu = array_to_cube(_Cxu);
    vec c0 = array_to_vec(_c0);

    cube agCpp = array_to_cube(_agCpp);
    mat agcp = array_to_mat(_agcp);
    vec agc0 = array_to_vec(_agc0);

    // recreate state-action-offset dist.
    mat mu_u(dm_act, nb_steps);
    cube sigma_u(dm_act, dm_act, nb_steps);

    mat mu_xu(dm_state + dm_act + 1, nb_steps + 1);
    cube sigma_xu(dm_state + dm_act + 1, dm_state + dm_act + 1, nb_steps + 1);

    for (int i = 0; i < nb_steps; i++) {
        // mu_u = K * mu_x + k
        mu_u.col(i) = K.slice(i) * mu_x.col(i) + kff.col(i);

        // sigma_u = sigma_ctl + K * sigma_x * K_T
        sigma_u.slice(i) = sigma_ctl.slice(i) + K.slice(i) * sigma_x.slice(i) * K.slice(i).t();
        sigma_u.slice(i) = 0.5 * (sigma_u.slice(i) + sigma_u.slice(i).t());
        sigma_u.slice(i) += 1e-8 * eye(dm_act, dm_act);

        // sigma_xu =   [[sigma_x,        sigma_x * K_T,   0.],
        //               [K * sigma_x,    sigma_u,         0.],
        //               [0.              0.               0.]]
        sigma_xu.slice(i) = join_vert( join_horiz(sigma_x.slice(i), sigma_x.slice(i) * K.slice(i).t(), zeros(dm_state, 1)),
                                       join_horiz(K.slice(i) * sigma_x.slice(i), sigma_u.slice(i), zeros(dm_act, 1)),
                                       join_horiz(zeros(1, dm_state), zeros(1, dm_act), zeros(1, 1)) );
        sigma_xu.slice(i) = 0.5 * (sigma_xu.slice(i) + sigma_xu.slice(i).t());
        sigma_xu.slice(i) += 1e-8 * eye(dm_state + dm_act + 1, dm_state + dm_act + 1);

        // mu_xu =  [[mu_x],
        //           [mu_u],
        //           [1],
        mu_xu.col(i) = join_vert(mu_x.col(i), mu_u.col(i), ones(1));
    }

    mu_xu.col(nb_steps) = join_vert(mu_x.col(nb_steps), zeros(dm_act), ones(1));
    sigma_xu.slice(nb_steps).submat(0, 0, dm_state - 1, dm_state - 1) = sigma_x.slice(nb_steps);

    // temp
    mat W(dm_param, dm_param);
    vec w(dm_param);

    mat A(dm_state, dm_state);
    mat B(dm_state, dm_act);
    vec c(dm_state);

    mat A_cl(dm_state, dm_state);
    vec c_cl(dm_state);
    mat sigma_block(dm_state + dm_act + 1, dm_state + dm_act + 1);

    mat P(dm_state + dm_act + 1, dm_state + dm_act + 1);
    mat Pxx(dm_state, dm_state);
    mat Pxu(dm_state, dm_act);
    mat Puu(dm_act, dm_act);
    vec px(dm_state);
    vec pu(dm_act);
    double p0;

    // outputs
    mat mu_optimal(dm_param, nb_steps);
    cube sigma_optimal(dm_param, dm_param, nb_steps);

    cube V(dm_state, dm_state, nb_steps + 1);
    mat v(dm_state, nb_steps + 1);
    vec v0(nb_steps + 1);

    int _diverge = -1;

    // last time step
    V.slice(nb_steps) = - Cxx.slice(nb_steps);
    v.col(nb_steps) = - cx.col(nb_steps);
    v0(nb_steps) = - c0(nb_steps);

	for(int i = nb_steps - 1; i >= 0; --i)
	{
	    mat mat_mu_xu = kron(mu_xu.col(i).t(), eye(dm_state, dm_state));

	    mat Vpp = mat_mu_xu.t() * V.slice(i + 1) * mat_mu_xu + kron(sigma_xu.slice(i), V.slice(i + 1));
	    vec vp =  mat_mu_xu.t() * v.col(i + 1);

        W = 2.0 * (agCpp.slice(i) + Vpp) / (beta + eta);
        W = 0.5 * (W.t() + W);

        w = - (agcp.col(i) + vp) / (beta + eta);

        try {
            sigma_optimal.slice(i) = inv_sympd(W);
        } catch ( const std::runtime_error ) {
            _diverge = i;
            break;
        }
        sigma_optimal.slice(i) = 0.5 * (sigma_optimal.slice(i).t() + sigma_optimal.slice(i));

        mu_optimal.col(i) = sigma_optimal.slice(i) * w;

        mat At = mu_optimal(span(0, dm_state * dm_state - 1), i);
        mat Bt = mu_optimal(span(dm_state * dm_state, dm_state * dm_state + dm_state * dm_act - 1), i);
        vec ct = mu_optimal(span(dm_state * dm_state + dm_state * dm_act, dm_state * dm_state + dm_state * dm_act + dm_state - 1), i);

        A = reshape(At, size(A));
        B = reshape(Bt, size(B));
        c = reshape(ct, size(c));

        // extra terms due to parameter distribution
         for (int j = 0; j < dm_state + dm_act + 1; j++){
            for (int k = 0; k < dm_state + dm_act + 1; k++){
                P(j, k) = trace(sigma_optimal.slice(i).submat(j * dm_state, k * dm_state,
                                                              (j + 1) * dm_state - 1, (k + 1) * dm_state - 1) * V.slice(i+1));
            }
         }

        Pxx = P.submat(0, 0, dm_state - 1, dm_state - 1);
        Puu = P.submat(dm_state, dm_state, dm_state + dm_act - 1, dm_state + dm_act - 1);
        Pxu = P.submat(0, dm_state, dm_state - 1, dm_state + dm_act - 1);

        px = P.submat(0, dm_state + dm_act, dm_state - 1, dm_state + dm_act);
        pu = P.submat(dm_state, dm_state + dm_act, dm_state + dm_act - 1, dm_state + dm_act);
        p0 = P(dm_state + dm_act, dm_state + dm_act);

        A_cl = A + B * K.slice(i);
        c_cl = c + B * kff.col(i);
        sigma_block.submat(dm_state, dm_state, dm_state + dm_act - 1, dm_state + dm_act - 1) = sigma_ctl.slice(i);

        V.slice(i) = (- Cxx.slice(i) + Pxx) + K.slice(i).t() * (- Cuu.slice(i) + Puu) * K.slice(i)
                      + A_cl.t() * V.slice(i + 1) * A_cl + 2. * (- Cxu.slice(i) + Pxu) * K.slice(i);
        V.slice(i) = 0.5 * (V.slice(i) + V.slice(i).t());

        v.col(i) = (- cx.col(i) + 2. * px) + 2. * K.slice(i).t() * (- Cuu.slice(i) + Puu) * kff.col(i)
                    + 2. * (- Cxu.slice(i) + Pxu) * kff.col(i) + K.slice(i).t() * (- cu.col(i) + 2. * pu)
                    + 2. * A_cl.t() * V.slice(i + 1) * c_cl + A_cl.t() * v.col(i + 1);

        v0(i) = as_scalar( (- c0(i) + p0) + kff.col(i).t() * (- Cuu.slice(i) + Puu) * kff.col(i) + kff.col(i).t() * (- cu.col(i) + 2. * pu)
                            - trace(Cuu.slice(i + 1) * sigma_ctl.slice(i)) + v0(i + 1) + trace(V.slice(i + 1) * sigma_dyn.slice(i))
                            + mu_optimal.col(i).t() * kron(sigma_block, V.slice(i + 1)) * mu_optimal.col(i) + trace(kron(sigma_block, V.slice(i + 1)) * sigma_optimal.slice(i))
                            + c_cl.t() * V.slice(i + 1) * c_cl + c_cl.t() * v.col(i + 1) );
	}

    // transform outputs to numpy
    array_tf _V = cube_to_array(V);
    array_tf _v = mat_to_array(v);
    array_tf _v0 = vec_to_array(v0);

    array_tf _mu_optimal = mat_to_array(mu_optimal);
    array_tf _sigma_optimal = cube_to_array(sigma_optimal);

    py::tuple output = py::make_tuple(_V, _v, _v0, _mu_optimal, _sigma_optimal, _diverge);

    return output;
}

PYBIND11_MODULE(core, m)
{
    m.def("policy_divergence", &policy_divergence);
    m.def("gaussian_divergence", &gaussian_divergence);
    m.def("gaussian_interp_w2", &gaussian_interp_w2);
    m.def("gaussian_interp_kl", &gaussian_interp_kl);
    m.def("quad_expectation", &quad_expectation);
    m.def("cubature_forward_pass", &cubature_forward_pass);
    m.def("policy_augment_cost", &policy_augment_cost);
    m.def("policy_backward_pass", &policy_backward_pass);
    m.def("parameter_augment_cost", &parameter_augment_cost);
    m.def("parameter_backward_pass", &parameter_backward_pass);
    m.def("regularized_parameter_augment_cost", &regularized_parameter_augment_cost);
}


/* * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * */
/*                                                                           */
/*  This file is part of the library KASKADE 7                               */
/*    see http://www.zib.de/projects/kaskade7-finite-element-toolbox         */
/*                                                                           */
/*  Copyright (C) 2002-2015 Zuse Institute Berlin                            */
/*                                                                           */
/*  KASKADE 7 is distributed under the terms of the ZIB Academic License.    */
/*    see $KASKADE/academic.txt                                              */
/*                                                                           */
/* * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * */

#include <iostream>

//#define BOOST_DISABLE_ASSERTS

#include <boost/timer/timer.hpp>

#include "dune/grid/config.h"
#include "dune/grid/uggrid.hh"
#include "dune/istl/solvers.hh"

#include "fem/assemble.hh"
#include "fem/lagrangespace.hh"
#include "fem/embedded_errorest.hh"
#include "linalg/trivialpreconditioner.hh"
#include "linalg/jacobiPreconditioner.hh"
#include "linalg/uzawa.hh"
#include "linalg/direct.hh"
#include "io/vtk.hh"
#include "utilities/kaskopt.hh" // property_tree
#include "utilities/gridGeneration.hh"

using namespace Kaskade;
#include "stokes-adaptive.hh"

int main(int argc, char *argv[])
{
  using namespace boost::fusion;

  int verbosity = 1;  // print to console if arguments are changed
  bool dump = false; // do not write properties into file
  std::unique_ptr<boost::property_tree::ptree> pt = getKaskadeOptions(argc, argv, verbosity, dump);

  std::cout << "Start Stokes tutorial programm" << std::endl;
  std::cout << "using addaptive refinement based on embedded error estimation." << std::endl;
  std::cout << "Note: use FixedFractionCriterion as refinement strategy!" << std::endl << std::endl;

  constexpr int dim = 2;
  constexpr int uIdx = 0;
  constexpr int pIdx = 1;

  // command line parameters
  int refinements = getParameter(pt, "refinements", 2);   // initial mesh is refined "refinements" times
  int order = getParameter(pt, "order", 3);               // order for velocity space, order for pressure space is (order-1)
  int steps = getParameter(pt, "steps", 5);               // adaptive refinement steps

  bool direct = getParameter(pt, "direct", true);
  DirectType directType = static_cast<DirectType>(getParameter(pt, "directSolver", 4)); // 4: DirectType::UMFPACK3264

  // grid generation
  using Grid = Dune::UGGrid<dim>;
  using H1Space = FEFunctionSpace<ContinuousLagrangeMapper<double,Grid::LeafGridView> >;
  using Spaces = vector<H1Space const*,H1Space const*>;
  using VariableDescriptions = vector<Variable<SpaceIndex<1>,Components<2>,VariableId<uIdx> >,
                                      Variable<SpaceIndex<0>,Components<1>,VariableId<pIdx> > >;
  using VariableSet = VariableSetDescription<Spaces,VariableDescriptions>;
  using CoefficientVectors = VariableSet::CoefficientVectorRepresentation<>::type;
  using Functional = StokesFunctional<double,VariableSet>;
  using Assembler = VariationalFunctionalAssembler<LinearizationAt<Functional> >;

  Dune::FieldVector<double,dim> x0(0.0), length(1.0);
  GridManager<Grid> gridManager( createRectangle<Grid>(x0,length,1.0));
  gridManager.globalRefine(refinements);

  // construct involved spaces.
  H1Space pressureSpace(gridManager,gridManager.grid().leafGridView(),order-1);
  H1Space velocitySpace(gridManager,gridManager.grid().leafGridView(),order);

  Spaces spaces(&pressureSpace,&velocitySpace);

  // construct variable list.
  std::string varNames[2] = { "u", "p" };

  VariableSet variableSet(spaces,varNames);

  // construct variational functional.
  Functional F;

  // construct Galerkin representation
  Assembler assembler(spaces);
  VariableSet::VariableSet x(variableSet);
  VariableSet::VariableSet dx(variableSet);

  // accuracy in space
  std::vector<std::pair<double,double> > tol(variableSet.noOfVariables);
  tol[0] = std::make_pair(1e0,1e-2);   // u
  tol[1] = std::make_pair(1e0,1e-2);   // p
  bool accurate = 0;

  for (int i=0; i<steps; ++i)    // adaptive refinement loop
  {
    std::cout << "\nstep = " << i+1 << std::endl;
    size_t nnz = assembler.nnz(0,2,0,2,false);
    size_t dof = variableSet.degreesOfFreedom(0,2);
    std::cout << "overall degrees of freedom: " << dof << std::endl;
    //std::cout << "(structurally) nonzero elements: " << nnz << std::endl;

    x = 0;
    assembler.assemble(linearization(F,x));

    if(direct)
    {
      CoefficientVectors solution(VariableSet::CoefficientVectorRepresentation<>::init(spaces));
      CoefficientVectors rhs(assembler.rhs());

      // solve performing one Newton step
      directInverseOperator(AssembledGalerkinOperator<Assembler>(assembler),directType).applyscaleadd(-1.0,rhs,solution);
      x.data = solution.data;
    }
    else // Uzawa Solver
    {
      using VectorOfU = VariableSet::CoefficientVectorRepresentation<uIdx,uIdx+1>::type;
      using VectorOfP = VariableSet::CoefficientVectorRepresentation<pIdx,pIdx+1>::type;
      using Assembler_UU = AssembledGalerkinOperator<Assembler,uIdx,uIdx+1,uIdx,uIdx+1>;
      using Assembler_PU = AssembledGalerkinOperator<Assembler,pIdx,pIdx+1,uIdx,uIdx+1>;
      using Assembler_UP = AssembledGalerkinOperator<Assembler,uIdx,uIdx+1,pIdx,pIdx+1>;
      using PreconAdapt = MatrixRepresentedOperator<MatrixAsTriplet<double>, VectorOfP, VectorOfP>;
      using UzSo = UzawaSolver<VectorOfU,VectorOfP>;

      Assembler_UU A(assembler);
      Assembler_PU B(assembler);
      Assembler_UP Bt(assembler);

      JacobiPreconditioner<Assembler_UU> jacobiPreconditioner(A);
      // inexact inner solver for upper left block
      Dune::CGSolver<VectorOfU> cg(A,jacobiPreconditioner,1e-14,300,0);

      TrivialPreconditioner<PreconAdapt> trivialPreconditioner;

      VectorOfU f(assembler.rhs<uIdx,uIdx+1>());
      VectorOfP g(assembler.rhs<pIdx,pIdx+1>());
      VectorOfU u(VariableSet::CoefficientVectorRepresentation<uIdx,uIdx+1>::init(spaces));
      VectorOfP p(VariableSet::CoefficientVectorRepresentation<pIdx,pIdx+1>::init(spaces));

      UzSo uzawa(A,cg,B,Bt,trivialPreconditioner,1e-4,100,2);

      Dune::InverseOperatorResult res;
      UzSo::Domain solution(vector<VectorOfU,VectorOfP>(u,p));
      UzSo::Range rhs(vector<VectorOfU,VectorOfP>(f,g));
      rhs *= -1.0; // change sign as rhs is -'F while the assembler returns F'
      uzawa.apply(solution,rhs,res);

      at_c<uIdx>(x.data).coefficients() = at_c<0>(at_c<uIdx>(solution.data).data);
      at_c<pIdx>(x.data).coefficients() = at_c<0>(at_c<pIdx>(solution.data).data);
    }

    writeVTKFile(gridManager.grid().leafGridView(),x,"stokes-adaptive", IoOptions(), 1);

    VariableSet::VariableSet e = x;
    projectHierarchically(variableSet,e);
    e -= x;

    accurate = embeddedErrorEstimator(variableSet,e,x,IdentityScaling(),tol,gridManager,1);
    dof = variableSet.degreesOfFreedom(0,2);
    std::cout << "degrees of freedom after refinement = " << dof << std::endl;

  }

  return 0;
}


/*~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~\
|  Phycas: Python software for phylogenetic analysis                          |
|  Copyright (C) 2006 Mark T. Holder, Paul O. Lewis and David L. Swofford     |
|                                                                             |
|  This program is free software; you can redistribute it and/or modify       |
|  it under the terms of the GNU General Public License as published by       |
|  the Free Software Foundation; either version 2 of the License, or          |
|  (at your option) any later version.                                        |
|                                                                             |
|  This program 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 General Public License along    |
|  with this program; if not, write to the Free Software Foundation, Inc.,    |
|  51 Franklin Street, Fifth Floor, Boston, MA 02110-1301 USA.                |
\~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~*/

#if !defined(DIRICHLET_DISTRIBUTION_HPP)
#define DIRICHLET_DISTRIBUTION_HPP

#if defined(_MSC_VER)
#	pragma warning(disable: 4267)	// warning about loss of data when converting size_t to int
#endif

//#include <cmath>
//#include "ncl/nxsdefs.h"

//#include "phycas/src/states_patterns.hpp"

//#include <boost/shared_ptr.hpp>
//#include <boost/format.hpp>
#include "multivariate_probability_distribution.hpp"
//#include "phycas/src/basic_cdf.hpp"
//#include "phycas/src/basic_lot.hpp"
//#include "phycas/src/phycas_string.hpp"
//#if defined(PYTHON_ONLY) && defined(USING_NUMARRAY)
//#	include <boost/python/tuple.hpp>
//#	include <boost/python/numeric.hpp>
//#	include "thirdparty/num_util/num_util.h"
//#endif
//#include "phycas/src/xprobdist.hpp"

namespace phycas
{

/*------------------------------------------------------------------------------------------------------------------------------------------------------------------
|	The Dirichlet distribution, with the number of parameters determined by the number of params supplied to the constructor.
*/
class DirichletDistribution : public MultivariateProbabilityDistribution
	{
	public:
                                                    DirichletDistribution();
                                                    DirichletDistribution(const std::vector<double> & params);
                                                    DirichletDistribution(const DirichletDistribution & other);
		virtual                                     ~DirichletDistribution();

        DirichletDistribution *                     cloneAndSetLot(Lot * other) const;
        DirichletDistribution *                     Clone() const;
		virtual void								SetLot(Lot * other);
		virtual void								ResetLot();
		virtual void								SetSeed(unsigned rnseed);

		virtual bool								IsDiscrete() const;
		virtual std::string 						GetDistributionName() const;
		virtual std::string 						GetDistributionDescription() const;
		std::string                                 GetDescriptionForPython() const;
		virtual std::vector<double>					GetMean() const;
		virtual std::vector<double> 				GetVar() const;
		virtual std::vector<double> 				GetStdDev() const;
		virtual std::vector<double>					Sample() const;
		virtual double								ApproxCDF(const std::vector<double> &x, unsigned nsamples = 10000) const;
		virtual double								GetLnPDF(const std::vector<double> &x) const;
		virtual double								GetRelativeLnPDF(const std::vector<double> &x) const;
		virtual void 								SetMeanAndVariance(const std::vector<double> &m, const std::vector<double> &v);
#		if defined(PYTHON_ONLY)
#			if defined(USING_NUMARRAY)
				void                                AltSetMeanAndVariance(boost::python::numeric::array m, boost::python::numeric::array v);
				boost::python::numeric::array       GetVarCovarMatrix();
#			else
				void                                AltSetMeanAndVariance(std::vector<double> m, std::vector<double> v);
				std::vector<double>					GetVarCovarMatrix();
#			endif
#		endif

		virtual unsigned							GetNParams() const;
		const GammaDistribution &                   GetDistributionOnParameter(unsigned i) const;

    protected:

		void                                        initialize(const double_vect_t & params);

    protected:

		std::vector<double>                         dirParams;
		std::vector<GammaDistribution>              paramDistributions;
		mutable std::vector<double>                 scratchSpace;
	};

} // namespace phycas

#endif


///1
#include <iostream>
#include <vector>
#include <cmath>
#include <tuple>
#include <boost/graph/adjacency_list.hpp>
#include <boost/graph/connected_components.hpp>
#include <boost/pending/disjoint_sets.hpp>
#include <CGAL/Exact_predicates_inexact_constructions_kernel.h>
#include <CGAL/Delaunay_triangulation_2.h>
#include <CGAL/Triangulation_vertex_base_with_info_2.h>
#include <CGAL/Triangulation_face_base_2.h>

typedef CGAL::Exact_predicates_inexact_constructions_kernel K;
typedef std::size_t                                            Index;
typedef CGAL::Triangulation_vertex_base_with_info_2<Index,K>   Vb;
typedef CGAL::Triangulation_face_base_2<K>                     Fb;
typedef CGAL::Triangulation_data_structure_2<Vb,Fb>            Tds;
typedef CGAL::Delaunay_triangulation_2<K,Tds>                  Delaunay;
typedef std::pair<K::Point_2,Index> IPoint;
typedef std::tuple<int,int> Edge;
typedef std::vector<Edge> EdgeV;
typedef Tds::Vertex_handle Vh;

typedef boost::adjacency_list<boost::vecS, boost::vecS, boost::undirectedS> graph;
typedef boost::graph_traits<graph>::vertex_descriptor vertex_desc;
typedef boost::graph_traits<graph>::edge_iterator edge_it;

using namespace std;

// Strategy:
// - Binary search for the best k
//   - For a given k, construct a Graph that contains only points p_k...p_{n-1}
//   - Find the largest connected component
void solve() {
  int n;
  long r;
  cin >> n >> r;
  long r2 = r * r;

  vector<IPoint> points;
  points.reserve(n);
  int x, y;
  for (int i = 0; i < n; ++i) {
    cin >> x >> y;
    points.emplace_back(K::Point_2(x, y), i);
  }

  // Start binary search
  int a = 1;
  int b = n;
  while (a < b - 1) {

    int empireSize = (a + b) / 2;

    // Use delaunay to quickly construct graph of points outside the empire
    Delaunay dt;
    dt.insert(points.begin() + empireSize, points.end());

    graph G(n - empireSize);
    for (auto e = dt.finite_edges_begin(); e != dt.finite_edges_end(); ++e) {
      int u = e->first->vertex((e->second+1)%3)->info();
      int v = e->first->vertex((e->second+2)%3)->info();
      if (dt.segment(e).squared_length() <= r2) {
        boost::add_edge(u - empireSize, v - empireSize, G);
      }
    }

    // Compute largest connected component
    std::vector<int> component_map(n);
    int ncc = boost::connected_components(G, boost::make_iterator_property_map(component_map.begin(), boost::get(boost::vertex_index, G)));

    int largestComponent = 0;
    vector<int> componentSize(ncc, 0);
    for (int i = 0; i < n - empireSize; ++i) {
      int c = component_map[i];
      componentSize[c]++;
      largestComponent = max(largestComponent, int(componentSize[c]));
    }

    // Continue binary search
    if (empireSize <= largestComponent) {
      // Invalid solution, the empire must be at least as large as the rebel alliance (both expand equally fast)
      // => Expand empire
      a = empireSize;
    }
    else {
      b = empireSize;
    }
  }
  cout << a << endl;
}

int main() {
    ios_base::sync_with_stdio(false);
    int t; cin >> t;
    while (t--) {
        solve();
    }
    return 0;
}


//==============================================================================
//         Copyright 2003 - 2011 LASMEA UMR 6602 CNRS/Univ. Clermont II
//         Copyright 2009 - 2011 LRI    UMR 8623 CNRS/Univ Paris Sud XI
//         Copyright 2012 - 2014 MetaScale SAS
//
//          Distributed under the Boost Software License, Version 1.0.
//                 See accompanying file LICENSE.txt or copy at
//                     http://www.boost.org/LICENSE_1_0.txt
//==============================================================================
#ifndef BOOST_SIMD_ARITHMETIC_FUNCTIONS_SIMD_COMMON_HYPOT_HPP_INCLUDED
#define BOOST_SIMD_ARITHMETIC_FUNCTIONS_SIMD_COMMON_HYPOT_HPP_INCLUDED

#include <boost/simd/arithmetic/functions/hypot.hpp>
#include <boost/simd/include/functions/simd/abs.hpp>
#include <boost/simd/include/functions/simd/min.hpp>
#include <boost/simd/include/functions/simd/max.hpp>
#include <boost/simd/include/functions/simd/plus.hpp>
#include <boost/simd/include/functions/simd/unary_minus.hpp>
#include <boost/simd/include/functions/simd/exponent.hpp>
#include <boost/simd/include/functions/simd/ldexp.hpp>
#include <boost/simd/include/functions/simd/sqr.hpp>
#include <boost/simd/include/functions/simd/sqrt.hpp>
#include <boost/simd/include/constants/maxexponentm1.hpp>
#include <boost/simd/include/constants/minexponent.hpp>
#include <boost/dispatch/meta/as_integer.hpp>

#ifndef BOOST_SIMD_NO_INVALIDS
#include <boost/simd/sdk/meta/as_logical.hpp>
#include <boost/simd/include/functions/simd/if_else.hpp>
#include <boost/simd/include/functions/simd/logical_and.hpp>
#include <boost/simd/include/functions/simd/logical_or.hpp>
#include <boost/simd/include/functions/simd/is_inf.hpp>
#include <boost/simd/include/functions/simd/is_nan.hpp>
#include <boost/simd/include/constants/inf.hpp>
#endif

namespace boost { namespace simd { namespace ext
{
  BOOST_DISPATCH_IMPLEMENT          ( hypot_, tag::cpu_
                                    , (A0)(X)
                                    , ((simd_<floating_<A0>,X>))
                                      ((simd_<floating_<A0>,X>))
                                    )
  {
    typedef A0 result_type;

    BOOST_SIMD_FUNCTOR_CALL_REPEAT(2)
    {
      typedef typename dispatch::meta::as_integer<result_type>::type   iA0;
      result_type r =  boost::simd::abs(a0);
      result_type i =  boost::simd::abs(a1);
      iA0 e =  exponent(boost::simd::max(i, r));
      e = boost::simd::min(boost::simd::max(e,Minexponent<A0>()),Maxexponentm1<A0>());
      result_type res =  ldexp(sqrt(sqr(ldexp(r, -e))+sqr(ldexp(i, -e))), e);
      #ifndef BOOST_SIMD_NO_INVALIDS
      typedef typename meta::as_logical<result_type>::type             lA0;
      lA0 test = logical_or(logical_and(is_nan(a0), is_inf(a1)),
                            logical_and(is_nan(a1), is_inf(a0)));
      return if_else(test, Inf<result_type>(), res);
      #else
      return res;
      #endif
    }
  };
} } }

#endif


#include <math.h>
#include <iostream>
#include <Eigen/Cholesky>
#include <Eigen/LU>
#include <Eigen/SVD>

#include "fastQP.h"
#include "gurobiQP.h"

#define _USE_MATH_DEFINES

#define MAX_CONSTRS 1000
#define MAX_STATE   1000
#define MAX_ITER    10

using namespace Eigen;
using namespace std;

//template <typename tA, typename tB, typename tC, typename tD, typename tE, typename tF, typename tG>
//int fastQPThatTakesQinv(vector< MatrixBase<tA>* > QinvblkDiag, const MatrixBase<tB>& f, const MatrixBase<tC>& Aeq, const MatrixBase<tD>& beq, const MatrixBase<tE>& Ain, const MatrixBase<tF>& bin, set<int>& active, MatrixBase<tG>& x)
int fastQPThatTakesQinv(vector< MatrixXd* > QinvblkDiag, const VectorXd& f, const MatrixXd& Aeq, const VectorXd& beq, const MatrixXd& Ain, const VectorXd& bin, set<int>& active, VectorXd& x)
{
  int i,d;
  int iterCnt = 0;

  int M_in = bin.size();
  int M = Aeq.rows();
  int N = Aeq.cols();

  if (f.rows() != N) { cerr << "size of f (" << f.rows() << " by " << f.cols() << ") doesn't match cols of Aeq (" << Aeq.rows() << " by " << Aeq.cols() << ")" << endl; return 2; }
  if (beq.rows() !=M) { cerr << "size of beq doesn't match rows of Aeq" << endl; return 2; }
  if (Ain.cols() !=N) { cerr << "cols of Ain doesn't match cols of Aeq" << endl; return 2; };
  if (bin.rows() != Ain.rows()) { cerr << "bin rows doesn't match Ain rows" << endl; return 2; };
  if (x.rows() != N) { cerr << "x doesn't match Aeq" << endl; return 2; }
  int n_active = active.size();

  MatrixXd Aact = MatrixXd(n_active, N);
  VectorXd bact = VectorXd(n_active);

  MatrixXd QinvAteq(N,M);
  VectorXd minusQinvf(N);

  // calculate a bunch of stuff that is constant during each iteration
  int startrow=0;
//  for (typename vector< MatrixBase<tA>* >::iterator iterQinv=QinvblkDiag.begin(); iterQinv!=QinvblkDiag.end(); iterQinv++) {
//  	MatrixBase<tA> *thisQinv = *iterQinv;
  for (vector< MatrixXd* >::iterator iterQinv=QinvblkDiag.begin(); iterQinv!=QinvblkDiag.end(); iterQinv++) {
  	MatrixXd *thisQinv = *iterQinv;
  	int numRow = thisQinv->rows();
  	int numCol = thisQinv->cols();

  	if (numRow == 1 || numCol == 1) {  // it's a vector
  		d = numRow*numCol;
  		if (M>0) QinvAteq.block(startrow,0,d,M)= thisQinv->asDiagonal()*Aeq.block(0,startrow,M,d).transpose();  // Aeq.transpoODse().block(startrow,0,d,N)
			minusQinvf.segment(startrow,d) = -thisQinv->cwiseProduct(f.segment(startrow,d));
			startrow=startrow+d;
  	} else { // potentially dense matrix
			d = numRow;
			if (numRow!=numCol) {
					cerr << "Q is not square! " << numRow << "x" << numCol << "\n";
					return -2;
			}
			if (M>0) QinvAteq.block(startrow,0,d,M) = thisQinv->operator*(Aeq.block(0,startrow,M,d).transpose());  // Aeq.transpose().block(startrow,0,d,N)
			minusQinvf.segment(startrow,d) = -thisQinv->operator*(f.segment(startrow,d));
			startrow=startrow+d;
		}
  	if (startrow>N) {
			cerr << "Q is too big!" << endl;
			return -2;
		}
  }
  if (startrow!=N) { cerr << "Q is the wrong size.  Got " << startrow << "by" << startrow << " but needed " << N << "by" << N << endl; return -2; }

  MatrixXd A;
  VectorXd b;
  MatrixXd QinvAt;
  VectorXd lam, lamIneq;
  VectorXd violated(M_in);
  VectorXd violation;

  while(1) {
    iterCnt++;

    n_active = active.size();
    Aact.resize(n_active,N);
    bact.resize(n_active);

    i=0;
    for (set<int>::iterator iter=active.begin(); iter!=active.end(); iter++) {
    	if (*iter<0 || *iter>=Ain.rows()) {
    		return -3;  // active set is invalid.  exit quietly, because this is expected behavior in normal operation (e.g. it means I should immediately kick out to gurobi)
    	}
      Aact.row(i) = Ain.row(*iter);
      bact(i++) = bin(*iter);
    }

    A.resize(Aeq.rows() + Aact.rows(),N);
    b.resize(beq.size() + bact.size());
    A << Aeq,Aact;
    b << beq,bact;

    if (A.rows() > 0) {
      //Solve H * [x;lam] = [-f;b] using Schur complements, H = [Q,At';A,0];
      QinvAt.resize(QinvAteq.rows(), QinvAteq.cols() + Aact.rows());

      if (n_active>0) {
				int startrow=0;
				for (vector< MatrixXd* >::iterator iterQinv=QinvblkDiag.begin(); iterQinv!=QinvblkDiag.end(); iterQinv++) {
					MatrixXd* thisQinv = (*iterQinv);
					d = thisQinv->rows();
					int numCol = thisQinv->cols();

					if (numCol == 1) {  // it's a vector
						QinvAt.block(startrow,0,d,M+n_active) << QinvAteq.block(startrow,0,d,M), thisQinv->asDiagonal()*Aact.block(0,startrow,n_active,d).transpose();
					} else { // it's a matrix
						QinvAt.block(startrow,0,d,M+n_active) << QinvAteq.block(startrow,0,d,M), thisQinv->operator*(Aact.block(0,startrow,n_active,d).transpose());
					}

					startrow=startrow+d;
				}
      } else {
      	QinvAt = QinvAteq;
      }

      lam.resize(QinvAt.cols());
      lam = -(A*QinvAt).ldlt().solve(b + (f.transpose()*QinvAt).transpose());
      x = minusQinvf - QinvAt*lam;
      lamIneq = lam.tail(lam.size() - M);
    } else {
      x = minusQinvf;
      lamIneq.resize(0);
    }

    if(Ain.rows() == 0) {
      active.clear();
      break;
    }

    set<int> new_active;

    violation = Ain*x - bin;
    for (i=0; i<M_in; i++)
      if (violation(i) >= 1e-6)
      	new_active.insert(i);

    bool all_pos_mults = true;
    for (i=0; i<n_active; i++) {
    	if (lamIneq(i)<0) {
    		all_pos_mults = false;
    		break;
    	}
    }
    if (new_active.empty() && all_pos_mults) {
    	// existing active was AOK
    	break;
    }

    i=0;
    set<int>::iterator iter=active.begin(), tmp;
    while (iter!=active.end()) { // to accomodating inloop erase
  		tmp = iter++;
    	if (lamIneq(i++)<0) {
    		active.erase(tmp);
    	}
    }
    active.insert(new_active.begin(),new_active.end());

    if (iterCnt > MAX_ITER) {
      //Default to calling this method
//      cout << "FastQP max iter reached." << endl;
//       mexErrMsgIdAndTxt("Drake:approximateIKmex:Error", "Max iter reached. Problem is likely infeasible");
      return -1;
    }
  }
  return iterCnt;
}

//template <typename tA, typename tB, typename tC, typename tD, typename tE, typename tF, typename tG>
//int fastQP(vector< MatrixBase<tA>* > QblkDiag, const MatrixBase<tB>& f, const MatrixBase<tC>& Aeq, const MatrixBase<tD>& beq, const MatrixBase<tE>& Ain, const MatrixBase<tF>& bin, set<int>& active, MatrixBase<tG>& x)
int fastQP(vector< MatrixXd* > QblkDiag, const VectorXd& f, const MatrixXd& Aeq, const VectorXd& beq, const MatrixXd& Ain, const VectorXd& bin, set<int>& active, VectorXd& x)
{
  /* min 1/2 * x'QblkDiag'x + f'x s.t A x = b, Ain x <= bin
   * using active set method.  Iterative solve a linearly constrained
   * quadratic minimization problem where linear constraints include
   * Ain(active,:)x == bin(active).  Quit if all dual variables associated
   * with these equations are positive (i.e. they satisfy KKT conditions).
   *
   * Note:
   * fails if QP is infeasible.
   * active == initial rows of Ain to treat as equations.
   * Frank Permenter - June 6th 2013
   *
   * @retval  if feasible then iterCnt, else -1 for infeasible, -2 for input error
   */

	int N = f.rows();

  MatrixXd* Qinv = new MatrixXd[QblkDiag.size()];
  vector< MatrixXd* > Qinvmap;

	#define REG 1e-13
  // calculate a bunch of stuff that is constant during each iteration
  int startrow=0;
  //typedef typename vector< MatrixBase<tA> >::iterator Qiterator;

  int i=0;
  for (vector< MatrixXd* >::iterator iterQ=QblkDiag.begin(); iterQ!=QblkDiag.end(); iterQ++) {
  	MatrixXd* thisQ = *iterQ;
  	int numRow = thisQ->rows();
  	int numCol = thisQ->cols();

  	if (numCol == 1) {  // it's a vector
  		VectorXd Qdiag_mod = thisQ->operator+(VectorXd::Constant(numRow,REG)); // regularize
  		Qinv[i] = Qdiag_mod.cwiseInverse();
  		Qinvmap.push_back( &Qinv[i] );
  		startrow=startrow+numRow;
		} else { // potentially dense matrix
			if (numRow!=numCol) {
				if (numRow==1)
					cerr << "diagonal Q's must be set as column vectors" << endl;
				else
					cerr << "Q is not square! " << numRow << "x" << numCol << endl;
				return -2;
			}

			MatrixXd Q_mod = thisQ->operator+(REG*MatrixXd::Identity(numRow,numRow));
			Qinv[i] = Q_mod.inverse();
  		Qinvmap.push_back( &Qinv[i] );
  		startrow=startrow+numRow;
		}
//  	cout << "Qinv{" << i << "} = " << Qinv[i] << endl;
		if (startrow>N) {
			cerr << "Q is too big!" << endl;
			return -2;
		}
		i++;
  }
  if (startrow!=N) { cerr << "Q is the wrong size.  Got " << startrow << "by" << startrow << " but needed " << N << "by" << N << endl; return -2; }

  int info = fastQPThatTakesQinv(Qinvmap,f,Aeq,beq,Ain,bin,active,x);

  delete[] Qinv;
  return info;
}

/* Example call (allocate inequality matrix, call function, resize inequalites:
  VectorXd binBnd = VectorXd(2*N);
  AinBnd.setZero();
  int numIneq = boundToIneq(ub,lb,AinBnd,binBnd);
  AinBnd.resize(numIneq,N);
  binBnd.resize(numIneq);
*/
/*
int boundToIneq(const VectorXd& uB,const VectorXd& lB, MatrixXd& Ain, VectorXd& bin)
{
    int rCnt = 0;
    int cCnt = 0;

    if (uB.rows()+lB.rows() > A.rows() ) {
        cerr << "not enough memory allocated";
    }

    if (uB.rows()+lB.rows() > b.rows() ) {
        cerr << "not enough memory allocated";
    }

    for (int i = 0; i < lB.rows(); i++ ) {
        if (!isinf(lB(i))) {
            cout << lB(i);
            cout << i;
            Ain(rCnt,cCnt++) = -1;//lB(i);
            bin(rCnt++) = -lB(i);
        }
    }
    cCnt = 0;
    for (int i = 0; i < uB.rows(); i++ ) {
        if (!isinf(uB(i))) {
            Ain(rCnt,cCnt++) = 1;//uB(i);
            bin(rCnt++) = uB(i);
        }
    }

    //resizing inside function all causes exception (why??)
    //A.resize(rCnt,uB.rows());
    return rCnt;
}
*/

template <typename DerivedA,typename DerivedB>
int myGRBaddconstrs(GRBmodel *model, MatrixBase<DerivedA> const & A, MatrixBase<DerivedB> const & b, char sense, double sparseness_threshold = 1e-14)
{
  int i,j,nnz,error=0;
/*
  // todo: it seems like I should just be able to do something like this:
  SparseMatrix<double,RowMajor> sparseAeq(Aeq.sparseView());
  sparseAeq.makeCompressed();
  error = GRBaddconstrs(model,nq_con,sparseAeq.nonZeros(),sparseAeq.InnerIndices(),sparseAeq.OuterStarts(),sparseAeq.Values(),beq.data(),NULL);
*/

  int *cind = new int[A.cols()];
  double* cval = new double[A.cols()];
  for (i=0; i<A.rows(); i++) {
    nnz=0;
    for (j=0; j<A.cols(); j++) {
      if (abs(A(i,j))>sparseness_threshold) {
        cval[nnz] = A(i,j);
        cind[nnz++] = j;
      }
    }
    error = GRBaddconstr(model,nnz,cind,cval,sense,b(i),NULL);
    if (error) break;
  }

  delete[] cind;
  delete[] cval;
  return error;
}

//template <typename tA, typename tB, typename tC, typename tD, typename tE>
//GRBmodel* gurobiQP(GRBenv *env, vector< MatrixBase<tA>* > QblkDiag, VectorXd& f, const MatrixBase<tB>& Aeq, const MatrixBase<tC>& beq, const MatrixBase<tD>& Ain, const MatrixBase<tE>& bin, VectorXd& lb, VectorXd& ub, set<int>& active, VectorXd& x)
GRBmodel* gurobiQP(GRBenv *env, vector< MatrixXd* > QblkDiag, VectorXd& f, const MatrixXd& Aeq, const VectorXd& beq,
  const MatrixXd& Ain, const VectorXd& bin, VectorXd& lb, VectorXd& ub, set<int>& active, VectorXd& x, double active_set_slack_tolerance)
{
	// Note: f,lb, and ub are VectorXd instead of const MatrixBase templates because i want to be able to call f.data() on them

	// NOTE:  this allocates memory for a new GRBmodel and returns it. (you should delete this object when you're done with it)
	// NOTE:  by convention here, the active set indices correspond to Ain,bin first, then lb, then ub.

  GRBmodel *model = NULL;

  int method;  GRBgetintparam(env,"method",&method);

  int i,j,nparams = f.rows(),Qi,Qj;
  double *lbdata = NULL, *ubdata=NULL;
  if (lb.rows()==nparams) lbdata = lb.data();
  if (ub.rows()==nparams) ubdata = ub.data();
  CGE (GRBnewmodel(env,&model,"QP",nparams,NULL,lbdata,ubdata,NULL,NULL), env);

  int startrow=0,d;
  for (vector< MatrixXd* >::iterator iterQ=QblkDiag.begin(); iterQ!=QblkDiag.end(); iterQ++) {
  	MatrixXd* Q=*iterQ;

    // WARNING:  If there are no constraints, then gurobi clearly solves a different problem: min 1/2 x'Qx + f'x
    //  				 This is very strange; see the solveWGUROBI method in QuadraticProgram
    if (method==2) //&& (Aeq.rows()+Ain.rows()>0))
      *Q = .5* (*Q);

  	if (Q->rows() == 1 || Q->cols() == 1) {  // it's a vector
  		d = Q->rows()*Q->cols();
  	  for (i=0; i<d; i++) {
  	  	Qi=i+startrow;
  	  	CGE (GRBaddqpterms(model,1,&Qi,&Qi,&(Q->operator()(i))), env);
  	  }
  	  startrow=startrow+d;
  	} else { // potentially dense matrix
  		d = Q->rows();
  		if (d!=Q->cols()) {
  			cerr << "Q is not square! " << Q->rows() << "x" << Q->cols() << "\n";
  			return NULL;
  		}

  	  for (i=0; i<d; i++)
    	  for (j=0; j<d; j++) {
    	  	Qi=i+startrow; Qj = j+startrow;
    	  	CGE (GRBaddqpterms(model,1,&Qi,&Qj,&(Q->operator()(i,j))), env);
    	  }
  	  startrow=startrow+d;
  	}
  	if (startrow>nparams) {
  		cerr << "Q is too big!" << endl;
  		return NULL;
  	}
  }

  CGE (GRBsetdblattrarray(model,"Obj",0,nparams,f.data()), env);

  if (Aeq.rows()>0) CGE (myGRBaddconstrs(model,Aeq,beq,GRB_EQUAL, 1e-18), env);
  if (Ain.rows()>0) CGE (myGRBaddconstrs(model,Ain,bin,GRB_LESS_EQUAL, 1e-18), env);

  CGE (GRBupdatemodel(model), env);
  CGE (GRBoptimize(model), env);

  CGE (GRBgetdblattrarray(model, GRB_DBL_ATTR_X, 0, nparams, x.data()), env);

  VectorXd slack(Ain.rows());
  CGE (GRBgetdblattrarray(model, "Slack", Aeq.rows(), Ain.rows(), slack.data()), env);

  int offset=0;
  active.clear();
  for (int i=0; i<Ain.rows(); i++) {
  	if (slack(i)<active_set_slack_tolerance)
  		active.insert(i);
  }
  offset = Ain.rows();
  if (lb.rows()==nparams)
  	for (int i=0; i<nparams; i++)
  		if (x(i)-lb(i) < active_set_slack_tolerance)
  			active.insert(offset+i);
  if (ub.rows()==nparams)
  	for (int i=0; i<nparams; i++)
  		if (ub(i)-x(i) < active_set_slack_tolerance)
  			active.insert(offset+i+nparams);

  return model;
}

GRBmodel* gurobiActiveSetQP(GRBenv *env, vector< MatrixXd* > QblkDiag, VectorXd& f, const MatrixXd& Aeq, const VectorXd& beq,
  const MatrixXd& Ain, const VectorXd& bin, VectorXd& lb, VectorXd& ub, int* &vbasis, int vbasis_len, int* &cbasis, int cbasis_len, VectorXd& x)
{
  // NOTE:  this allocates memory for a new GRBmodel and returns it. (you should delete this object when you're done with it)
  // NOTE:  by convention here, the active set indices correspond to Ain,bin first, then lb, then ub.
  GRBmodel *model = NULL;

  int method;  GRBgetintparam(env,"method",&method);
  if (!(method==0 || method==1)) {
    cerr<< "gurobiActiveSetQP: method should be 0 or 1" << endl;
    return NULL;
  }

  int i,j,nparams = f.rows(),Qi,Qj;
  double *lbdata = NULL, *ubdata=NULL;
  if (lb.rows()==nparams) lbdata = lb.data();
  if (ub.rows()==nparams) ubdata = ub.data();
  CGE (GRBnewmodel(env,&model,"QP",nparams,NULL,lbdata,ubdata,NULL,NULL), env);

  int startrow=0,d;
  for (vector< MatrixXd* >::iterator iterQ=QblkDiag.begin(); iterQ!=QblkDiag.end(); iterQ++) {
    MatrixXd* Q=*iterQ;

	*Q = .5* (*Q);

    if (Q->rows() == 1 || Q->cols() == 1) {  // it's a vector
      d = Q->rows()*Q->cols();
      for (i=0; i<d; i++) {
        Qi=i+startrow;
        CGE (GRBaddqpterms(model,1,&Qi,&Qi,&(Q->operator()(i))), env);
      }
      startrow=startrow+d;
    } else { // potentially dense matrix
      d = Q->rows();
      if (d!=Q->cols()) {
        cerr << "Q is not square! " << Q->rows() << "x" << Q->cols() << "\n";
        return NULL;
      }

      for (i=0; i<d; i++)
        for (j=0; j<d; j++) {
          Qi=i+startrow; Qj = j+startrow;
          CGE (GRBaddqpterms(model,1,&Qi,&Qj,&(Q->operator()(i,j))), env);
        }
      startrow=startrow+d;
    }
    if (startrow>nparams) {
      cerr << "Q is too big!" << endl;
      return NULL;
    }
  }

  CGE (GRBsetdblattrarray(model,"Obj",0,nparams,f.data()), env);

  if (Aeq.rows()>0) CGE (myGRBaddconstrs(model,Aeq,beq,GRB_EQUAL, 1e-18), env);
  if (Ain.rows()>0) CGE (myGRBaddconstrs(model,Ain,bin,GRB_LESS_EQUAL, 1e-18), env);

  CGE (GRBupdatemodel(model), env);

  int numvars; CGE(GRBgetintattr(model,"NumVars",&numvars), env);
  if (numvars == vbasis_len) {
    CGE (GRBsetintattrarray(model, "VBasis", 0, numvars, vbasis), env);
  }
  else {
    delete[] vbasis;
    vbasis = new int[numvars];
  }

  int numconstr; CGE(GRBgetintattr(model,"NumConstrs",&numconstr),env);
  if (numconstr == cbasis_len) {
    CGE (GRBsetintattrarray(model, "CBasis", 0, numconstr, cbasis), env);
  }
  else {
    delete[] cbasis;
    cbasis = new int[numconstr];
  }
  CGE (GRBoptimize(model), env);

  CGE (GRBgetdblattrarray(model, GRB_DBL_ATTR_X, 0, nparams, x.data()), env);

  CGE (GRBgetintattrarray(model, "VBasis", 0, numvars, vbasis), env);
  CGE (GRBgetintattrarray(model, "CBasis", 0, numconstr, cbasis), env);

  return model;
}

/*
template int fastQP(vector< MatrixBase<MatrixXd>* > QblkDiag, const MatrixBase< Map<VectorXd> >&, const MatrixBase< Map<MatrixXd> >&, const MatrixBase< Map<VectorXd> >&, const MatrixBase< Map<MatrixXd> >&, const MatrixBase< Map<VectorXd> >&, set<int>&, MatrixBase< Map<VectorXd> >&);
template GRBmodel* gurobiQP(GRBenv *env, vector< MatrixBase<MatrixXd>* > QblkDiag, VectorXd& f, const MatrixBase< Map<MatrixXd> >& Aeq, const MatrixBase< Map<VectorXd> >& beq, const MatrixBase< Map<MatrixXd> >& Ain, const MatrixBase< Map<VectorXd> >&bin, VectorXd& lb, VectorXd& ub, set<int>&, VectorXd&);
template GRBmodel* gurobiQP(GRBenv *env, vector< MatrixBase<MatrixXd>* > QblkDiag, VectorXd& f, const MatrixBase< MatrixXd >& Aeq, const MatrixBase< VectorXd >& beq, const MatrixBase< MatrixXd >& Ain, const MatrixBase< VectorXd >&bin, VectorXd&lb, VectorXd&ub, set<int>&, VectorXd&);
*/

/*
template int fastQP(vector< MatrixBase< VectorXd > >, const MatrixBase< VectorXd >&, const MatrixBase< Matrix<double,-1,-1,RowMajor,1000,-1> >&, const MatrixBase< Matrix<double,-1,1,0,1000,1> >&, const MatrixBase< Matrix<double,-1,-1,RowMajor,1000,-1> >&, const MatrixBase< Matrix<double,-1,1,0,1000,1> >&, set<int>&, MatrixBase< VectorXd >&);
*/