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

///
/// \file prony_like_method_common.hpp
///
#ifndef MXPFIT_PRONY_LIKE_METHOD_COMMON_HPP
#define MXPFIT_PRONY_LIKE_METHOD_COMMON_HPP

#include <cassert>
#include <type_traits>

#include <Eigen/Core>
#include <Eigen/SVD>

#include <mxpfit/exponential_sum.hpp>
#include <mxpfit/matrix_free_gemv.hpp>
#include <mxpfit/vandermonde_least_squares.hpp>

namespace mxpfit
{

namespace detail
{

/// \internal
/// Select roots of Prony polynomials on the unit disk
template <typename T>
struct prony_roots_on_unit_disk
{
    using Scalar        = T;
    using RealScalar    = typename Eigen::NumTraits<Scalar>::Real;
    using ComplexScalar = std::complex<RealScalar>;
    using Index         = Eigen::Index;

    using Vector        = Eigen::Matrix<Scalar, Eigen::Dynamic, 1>;
    using RealVector    = Eigen::Matrix<RealScalar, Eigen::Dynamic, 1>;
    using ComplexVector = Eigen::Matrix<ComplexScalar, Eigen::Dynamic, 1>;

    enum
    {
        IsComplex = Eigen::NumTraits<Scalar>::IsComplex,
    };

    /// \param z  array of roots to be filtered
    /// \param tolerance  small real number to acceptable tolerance for the
    //                    mangnitude of roots.
    static ComplexVector compute(const Eigen::Ref<const ComplexVector>& z,
                                 RealScalar tolerance)
    {
        using Eigen::numext::abs;
        using Eigen::numext::imag;
        static const auto eps     = Eigen::NumTraits<RealScalar>::epsilon();
        constexpr const auto zero = RealScalar();
        constexpr const auto one  = RealScalar(1);

        // Count the number of roots insize unit disk
        Index count = 0;
        for (Index i = 0; i < z.size(); ++i)
        {
            const auto abs_zi = abs(z(i));
            if (abs_zi <= one + tolerance)
            {
                if (IsComplex)
                {
                    ++count;
                }
                else
                {
                    // Discard negative real z(i).
                    const auto xi = real(z(i));
                    const auto yi = imag(z(i));
                    if (!(abs(yi) < eps && xi < zero))
                    {
                        ++count;
                    }
                }
            }
        }

        ComplexVector ret(count);
        Index n = 0;
        for (Index i = 0; i < z.size(); ++i)
        {
            const auto abs_zi = abs(z(i));
            if (abs_zi <= one + tolerance)
            {
                if (IsComplex)
                {
                    ret(n) = z(i);
                    ++n;
                }
                else
                {
                    // Discard negative real z(i).
                    const auto xi = real(z(i));
                    const auto yi = imag(z(i));
                    if (!(abs(yi) < eps && xi < zero))
                    {
                        ret(n) = z(i);
                        ++n;
                    }
                }
            }
        }

        return ret;
    }
};

/// \internal
///
/// Solve overdetermined linear system
///
/// \f[
///    V\boldsymbol{x} = \boldsymbol{b}
/// \f]
///
/// to compute weights of the exponential sum approximation, where \f$V\f$ is a
/// column Vandermonde matrix constructed from roots of a Prony polynomial.
///
/// \param[in] prony_roots  roots of a Prony polynomial
/// \param[in] rhs   a sequence of function/data values on uniform grid
/// \param[out] dst  the solution of overdetermined Vandermonde system
/// \param[in] eps      a prescribed accuracy required for the least-squares
/// \param[in] max_iter maximum number of iterations for CGLS linear solver
///
template <typename VecNodes, typename VecRHS, typename VecDest, typename RealT>
void solve_overdetermined_vandermonde(
    const Eigen::DenseBase<VecNodes>& prony_roots,
    const Eigen::MatrixBase<VecRHS>& rhs, Eigen::DenseBase<VecDest>& dst,
    RealT eps, Eigen::Index max_iter)
{
    EIGEN_STATIC_ASSERT_VECTOR_ONLY(VecNodes);
    EIGEN_STATIC_ASSERT_VECTOR_ONLY(VecRHS);
    EIGEN_STATIC_ASSERT_VECTOR_ONLY(VecDest);

    assert(prony_roots.size() == dst.size());

    using Scalar          = typename VecNodes::Scalar; // usually complex type
    using VandermondeGEMV = MatrixFreeGEMV<VandermondeMatrix<Scalar>>;

    // --- Solve with CGLS method.
    VandermondeMatrix<Scalar> matV(rhs.size(), prony_roots);
    VandermondeGEMV opV(matV);
    VandermondeLeastSquaresSolver<Scalar> solver(opV);
    solver.setTolerance(eps);
    solver.setMaxIterations(max_iter);

    //
    // We need cast the scalar type of RHS vector as the VecRHS::Scalar might be
    // different from VecNodes::Scalar. We should be consider the following
    // case:
    //
    // VecNodes::Scalar -> complex type
    // VecRHS::Scalar   -> real type
    //
    dst = solver.solve(rhs.template cast<Scalar>());

    if (solver.info() == Eigen::NoConvergence)
    {
        //
        // CGLS did not converge.
        // Fall back to least-squares with dense QR factorization.
        //
        auto denseV = matV.toDenseMatrix();
        dst = denseV.jacobiSvd(Eigen::ComputeThinU | Eigen::ComputeThinV)
                  .solve(rhs.template cast<Scalar>());
    }

    return;
}

/// \internal
///
/// ### gen_prony_like_method_result
///
/// Generate and manipulate a result (an exponential sum) of Prony-like method.
///

template <typename T>
struct gen_prony_like_method_result;

template <typename T>
struct gen_prony_like_method_result
{
    using Real       = T;
    using Complex    = std::complex<Real>;
    using ResultType = ExponentialSum<Complex, Complex>;

    template <typename ArrayZ, typename ArrayW>
    static ResultType
    create(const Eigen::ArrayBase<ArrayZ>& z, // exp(-exponents)
           const Eigen::ArrayBase<ArrayW>& w, // weights
           Real x0, Real delta)
    {
        using Index = Eigen::Index;

        using Eigen::numext::abs;
        using Eigen::numext::conj;
        using Eigen::numext::exp;
        using Eigen::numext::imag;
        using Eigen::numext::log;
        using Eigen::numext::real;

        static const auto eps     = Eigen::NumTraits<Real>::epsilon();
        static const auto pi      = 4 * Eigen::numext::atan(Real(1));
        constexpr const auto zero = Real();
        constexpr const auto half = Real(0.5);

        //-------------------------------------------------------------------------
        // The exponents are obtained as a_i = -log(z_i), where {z_i} are the
        // roots of the Prony polynomial.
        //
        // Some z_i might be real and negative: in this case, the corresponding
        // parameter a_i becomes a complex, i.e, a_i = -ln|z_i|+i \pi.
        // However, its complex conjugate a_i^* is not included in the final
        // exponential sum approximation which makes the approximated function
        // non-real. Thus, we disregards those terms.
        //-------------------------------------------------------------------------

        // Count negative, real-valued Prony roots
        Index count = 0;
        for (Index i = 0; i < z.size(); ++i)
        {
            const auto xi = real(z(i));
            const auto yi = imag(z(i));
            if (xi < zero && abs(yi) < eps)
            {
                ++count;
            }
        }

        if (count)
        {
            // Found negative real roots.
            ResultType ret(z.size() + count);

            Index n = 0;
            for (Index i = 0; i < z.size(); ++i)
            {
                const auto xi = real(z(i));
                const auto yi = imag(z(i));
                if (xi < zero && abs(yi) < eps)
                {
                    // Log(z) = -log(xi) + i pi
                    const auto an       = Complex(-log(-xi), -pi) / delta;
                    ret.exponent(n)     = an;
                    ret.exponent(n + 1) = conj(an);

                    const auto wn     = half * w(i) * exp(-x0 * an);
                    ret.weight(n)     = wn;
                    ret.weight(n + 1) = conj(wn);

                    n += 2;
                }
                else
                {
                    const auto an   = -log(z(i)) / delta;
                    ret.exponent(n) = an;
                    ret.weight(n)   = w(i) * exp(-x0 * an);

                    ++n;
                }
            }

            return ret;
        }
        else
        {
            ResultType ret(z.size());
            ret.exponents() = -z.log() / delta;
            if (x0 == Real())
            {
                ret.weights() = w; // Don't forget to copy weights
            }
            else
            {
                ret.weights() = w * (-x0 * ret.exponents()).exp();
            }
            return ret;
        }
    }
};

template <typename T>
struct gen_prony_like_method_result<std::complex<T>>
{
    using Real       = T;
    using Complex    = std::complex<Real>;
    using ResultType = ExponentialSum<Complex, Complex>;

    template <typename ArrayZ, typename ArrayW>
    static ResultType
    create(const Eigen::ArrayBase<ArrayZ>& z, // exp(-exponents)
           const Eigen::ArrayBase<ArrayW>& w, // weights
           Real x0, Real delta)
    {
        ResultType ret(z.size());
        ret.exponents() = -z.log() / delta;
        if (x0 == Real())
        {
            ret.weights() = w; // Don't forget to copy weights
        }
        else
        {
            ret.weights() = w * (-x0 * ret.exponents()).exp();
        }

        return ret;
    }
};

} // namespace detail
} // namespace mxpfit

#endif /* MXPFIT_PRONY_LIKE_METHOD_COMMON_HPP */


/*
//@HEADER
// ************************************************************************
//
// qr_eigen_dense_out_of_place_impl.hpp
//                     		  Pressio
//                             Copyright 2019
//    National Technology & Engineering Solutions of Sandia, LLC (NTESS)
//
// Under the terms of Contract DE-NA0003525 with NTESS, the
// U.S. Government retains certain rights in this software.
//
// Pressio is licensed under BSD-3-Clause terms of use:
//
// Redistribution and use in source and binary forms, with or without
// modification, are permitted provided that the following conditions
// are met:
//
// 1. Redistributions of source code must retain the above copyright
// notice, this list of conditions and the following disclaimer.
//
// 2. Redistributions in binary form must reproduce the above copyright
// notice, this list of conditions and the following disclaimer in the
// documentation and/or other materials provided with the distribution.
//
// 3. Neither the name of the copyright holder nor the names of its
// contributors may be used to endorse or promote products derived
// from this software without specific prior written permission.
//
// THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS
// "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
// LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS
// FOR A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE
// COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT,
// INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES
// (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
// SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION)
// HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT,
// STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING
// IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
// POSSIBILITY OF SUCH DAMAGE.
//
// Questions? Contact Francesco Rizzi (fnrizzi@sandia.gov)
//
// ************************************************************************
//@HEADER
*/

#ifndef QR_IMPL_EIGEN_QR_EIGEN_DENSE_OUT_OF_PLACE_IMPL_HPP_
#define QR_IMPL_EIGEN_QR_EIGEN_DENSE_OUT_OF_PLACE_IMPL_HPP_

#include <Eigen/QR>

namespace pressio{ namespace qr{ namespace impl{

template< typename MatrixType, typename R_t>
class QRHouseholderDenseEigenMatrix<
  MatrixType, R_t,
  ::pressio::mpl::enable_if_t<
    ::pressio::is_dense_matrix_eigen<MatrixType>::value
    >
  >
{
public:
  using sc_t	       = typename ::pressio::Traits<MatrixType>::scalar_type;
  using Q_type = Eigen::Matrix<sc_t, Eigen::Dynamic, Eigen::Dynamic>;
  using factorizer_t = Eigen::HouseholderQR<MatrixType>;

private:
  mutable std::shared_ptr<Q_type> Qmat_	     = {};
  mutable std::shared_ptr<factorizer_t> fct_ = {};

public:
  QRHouseholderDenseEigenMatrix() = default;
  ~QRHouseholderDenseEigenMatrix() = default;

  void computeThinOutOfPlace(const MatrixType & A)
  {
    const auto rows = A.rows();
    const auto cols = A.cols();
    fct_ = std::make_shared<factorizer_t>(A);

    if (!Qmat_ or (Qmat_->rows()!=rows and Qmat_->cols()!=cols)){
      Qmat_ = std::make_shared<Q_type>(rows,cols);
      Qmat_->setZero();
    }

    *Qmat_ = fct_->householderQ() * Q_type::Identity(rows,cols);
  }

  template < typename VectorInType, typename VectorOutType>
  void applyQTranspose(const VectorInType & vecIn, VectorOutType & vecOut) const
  {
    constexpr auto beta  = ::pressio::utils::Constants<sc_t>::zero();
    constexpr auto alpha = ::pressio::utils::Constants<sc_t>::one();
    ::pressio::ops::product(::pressio::transpose(), alpha, *this->Qmat_, vecIn, beta, vecOut);
  }

  template < typename VectorInType, typename VectorOutType>
  void applyRTranspose(const VectorInType & vecIn, VectorOutType & y) const
  {
    // y = R^T vecIn
    auto vecSize = ::pressio::ops::extent(y, 0);
    auto & Rm = fct_->matrixQR().block(0,0,vecSize,vecSize).template triangularView<Eigen::Upper>();
    y = Rm.transpose() * vecIn;
  }

  template <typename VectorType>
  void doLinSolve(const VectorType & rhs, VectorType & y)const
  {
    auto vecSize = ::pressio::ops::extent(y, 0);
    auto & Rm = fct_->matrixQR().block(0,0,vecSize,vecSize).
      template triangularView<Eigen::Upper>();
    y = Rm.solve(rhs);
  }

  const Q_type & QFactor() const {
    return *this->Qmat_;
  }
};

}}} // end namespace pressio::qr::impl
#endif  // QR_IMPL_EIGEN_QR_EIGEN_DENSE_OUT_OF_PLACE_IMPL_HPP_


#include <boost/hana.hpp>
#include <boost/math/constants/constants.hpp>
#include <cmath>
#include <complex>
#include "../include/AlgorithmFactory.h"
#include <iostream>
#include <memory>

#include "../include/basic/SineCosine.h"
#include "../include/basic/Ceil.h"
#include "../include/basic/Floor.h"
#include "../include/basic/Round.h"

namespace constants = boost::math::constants;
namespace hana = boost::hana;
namespace jb = jeanbaptiste;
namespace jbo = jeanbaptiste::options;
namespace jw = jeanbaptiste::windowing;

int main()
{
    std::array<std::complex<double>, 16> complexData =
    {
        std::complex<double>(-1.0f, 0.0f),
		std::complex<double>(-1.0f, 0.0f),
		std::complex<double>(-1.0f, 0.0f),
		std::complex<double>(-1.0f, 0.0f),
		std::complex<double>( 1.0f, 0.0f),
		std::complex<double>( 1.0f, 0.0f),
		std::complex<double>( 1.0f, 0.0f),
		std::complex<double>( 1.0f, 0.0f),
		std::complex<double>(-1.0f, 0.0f),
		std::complex<double>(-1.0f, 0.0f),
		std::complex<double>(-1.0f, 0.0f),
		std::complex<double>(-1.0f, 0.0f),
		std::complex<double>( 1.0f, 0.0f),
		std::complex<double>( 1.0f, 0.0f),
		std::complex<double>( 1.0f, 0.0f),
		std::complex<double>( 1.0f, 0.0f)
    };

    jb::AlgorithmFactory<1, 6, jbo::Radix_2, jbo::Decimation_In_Frequency, jbo::Direction_Forward,
       jbo::Window_None, jbo::Normalization_Square_Root, std::complex<double>> algorithmFactory;

	constexpr auto stage = 4;
    auto algorithm = algorithmFactory.getAlgorithm(stage);
    algorithm->operator()(&complexData[0]);

	std::cout << "Selected algorithm - stage: " << stage << ", samples: " << algorithm->numberOfSamples() << ", frequencies: " << algorithm->numberOfFrequencies() << "\n\n";

	std::cout << "Applying FFT:\n";
    for (auto value : complexData)
       std::cout << std::setw(10) << std::setprecision(5) << value.real() << "\t" << value.imag() << "I\n";

	std::cout << "\nApplying iFFT:\n";
	jb::AlgorithmFactory<1, 6, jbo::Radix_2, jbo::Decimation_In_Frequency, jbo::Direction_Backward,
       jbo::Window_None, jbo::Normalization_Square_Root, std::complex<double>> iAlgorithmFactory;

	auto iAlgorithm = iAlgorithmFactory.getAlgorithm(stage);
    iAlgorithm->operator()(&complexData[0]);

    for (auto value : complexData)
       std::cout << std::setw(10) << std::setprecision(5) << value.real() << "\t" << value.imag() << "I\n";

    std::cin.get();

     return 0;
}


// g++-4.4 bench_gemm.cpp -I .. -O2 -DNDEBUG -lrt -fopenmp && OMP_NUM_THREADS=2  ./a.out
// icpc bench_gemm.cpp -I .. -O3 -DNDEBUG -lrt -openmp  && OMP_NUM_THREADS=2  ./a.out

// Compilation options:
//
// -DSCALAR=std::complex<double>
// -DSCALARA=double or -DSCALARB=double
// -DHAVE_BLAS
// -DDECOUPLED
//

#include <iostream>
#include <Eigen/Core>
#include <bench/BenchTimer.h>

using namespace std;
using namespace Eigen;

#ifndef SCALAR
// #define SCALAR std::complex<float>
#define SCALAR float
#endif

#ifndef SCALARA
#define SCALARA SCALAR
#endif

#ifndef SCALARB
#define SCALARB SCALAR
#endif

typedef SCALAR Scalar;
typedef NumTraits<Scalar>::Real RealScalar;
typedef Matrix<SCALARA,Dynamic,Dynamic> A;
typedef Matrix<SCALARB,Dynamic,Dynamic> B;
typedef Matrix<Scalar,Dynamic,Dynamic> C;
typedef Matrix<RealScalar,Dynamic,Dynamic> M;

#ifdef HAVE_BLAS

extern "C" {
  #include <Eigen/src/misc/blas.h>
}

static float fone = 1;
static float fzero = 0;
static double done = 1;
static double szero = 0;
static std::complex<float> cfone = 1;
static std::complex<float> cfzero = 0;
static std::complex<double> cdone = 1;
static std::complex<double> cdzero = 0;
static char notrans = 'N';
static char trans = 'T';
static char nonunit = 'N';
static char lower = 'L';
static char right = 'R';
static int intone = 1;

void blas_gemm(const MatrixXf& a, const MatrixXf& b, MatrixXf& c)
{
  int M = c.rows(); int N = c.cols(); int K = a.cols();
  int lda = a.rows(); int ldb = b.rows(); int ldc = c.rows();

  sgemm_(&notrans,&notrans,&M,&N,&K,&fone,
         const_cast<float*>(a.data()),&lda,
         const_cast<float*>(b.data()),&ldb,&fone,
         c.data(),&ldc);
}

EIGEN_DONT_INLINE void blas_gemm(const MatrixXd& a, const MatrixXd& b, MatrixXd& c)
{
  int M = c.rows(); int N = c.cols(); int K = a.cols();
  int lda = a.rows(); int ldb = b.rows(); int ldc = c.rows();

  dgemm_(&notrans,&notrans,&M,&N,&K,&done,
         const_cast<double*>(a.data()),&lda,
         const_cast<double*>(b.data()),&ldb,&done,
         c.data(),&ldc);
}

void blas_gemm(const MatrixXcf& a, const MatrixXcf& b, MatrixXcf& c)
{
  int M = c.rows(); int N = c.cols(); int K = a.cols();
  int lda = a.rows(); int ldb = b.rows(); int ldc = c.rows();

  cgemm_(&notrans,&notrans,&M,&N,&K,(float*)&cfone,
         const_cast<float*>((const float*)a.data()),&lda,
         const_cast<float*>((const float*)b.data()),&ldb,(float*)&cfone,
         (float*)c.data(),&ldc);
}

void blas_gemm(const MatrixXcd& a, const MatrixXcd& b, MatrixXcd& c)
{
  int M = c.rows(); int N = c.cols(); int K = a.cols();
  int lda = a.rows(); int ldb = b.rows(); int ldc = c.rows();

  zgemm_(&notrans,&notrans,&M,&N,&K,(double*)&cdone,
         const_cast<double*>((const double*)a.data()),&lda,
         const_cast<double*>((const double*)b.data()),&ldb,(double*)&cdone,
         (double*)c.data(),&ldc);
}

#endif

void matlab_cplx_cplx(const M& ar, const M& ai, const M& br, const M& bi, M& cr, M& ci)
{
  cr.noalias() += ar * br;
  cr.noalias() -= ai * bi;
  ci.noalias() += ar * bi;
  ci.noalias() += ai * br;
  // [cr ci] += [ar ai] * br + [-ai ar] * bi
}

void matlab_real_cplx(const M& a, const M& br, const M& bi, M& cr, M& ci)
{
  cr.noalias() += a * br;
  ci.noalias() += a * bi;
}

void matlab_cplx_real(const M& ar, const M& ai, const M& b, M& cr, M& ci)
{
  cr.noalias() += ar * b;
  ci.noalias() += ai * b;
}

template<typename A, typename B, typename C>
EIGEN_DONT_INLINE void gemm(const A& a, const B& b, C& c)
{
 c.noalias() += a * b;
}

int main(int argc, char ** argv)
{
  std::ptrdiff_t l1 = internal::queryL1CacheSize();
  std::ptrdiff_t l2 = internal::queryTopLevelCacheSize();
  std::cout << "L1 cache size     = " << (l1>0 ? l1/1024 : -1) << " KB\n";
  std::cout << "L2/L3 cache size  = " << (l2>0 ? l2/1024 : -1) << " KB\n";
  typedef internal::gebp_traits<Scalar,Scalar> Traits;
  std::cout << "Register blocking = " << Traits::mr << " x " << Traits::nr << "\n";

  int rep = 1;    // number of repetitions per try
  int tries = 2;  // number of tries, we keep the best

  int s = 2048;
  int m = s;
  int n = s;
  int p = s;
  int cache_size1=-1, cache_size2=l2, cache_size3 = 0;

  bool need_help = false;
  for (int i=1; i<argc;)
  {
    if(argv[i][0]=='-')
    {
      if(argv[i][1]=='s')
      {
        ++i;
        s = atoi(argv[i++]);
        m = n = p = s;
        if(argv[i][0]!='-')
        {
          n = atoi(argv[i++]);
          p = atoi(argv[i++]);
        }
      }
      else if(argv[i][1]=='c')
      {
        ++i;
        cache_size1 = atoi(argv[i++]);
        if(argv[i][0]!='-')
        {
          cache_size2 = atoi(argv[i++]);
          if(argv[i][0]!='-')
            cache_size3 = atoi(argv[i++]);
        }
      }
      else if(argv[i][1]=='t')
      {
        ++i;
        tries = atoi(argv[i++]);
      }
      else if(argv[i][1]=='p')
      {
        ++i;
        rep = atoi(argv[i++]);
      }
    }
    else
    {
      need_help = true;
      break;
    }
  }

  if(need_help)
  {
    std::cout << argv[0] << " -s <matrix sizes> -c <cache sizes> -t <nb tries> -p <nb repeats>\n";
    std::cout << "   <matrix sizes> : size\n";
    std::cout << "   <matrix sizes> : rows columns depth\n";
    return 1;
  }

#if EIGEN_VERSION_AT_LEAST(3,2,90)
  if(cache_size1>0)
    setCpuCacheSizes(cache_size1,cache_size2,cache_size3);
#endif

  A a(m,p); a.setRandom();
  B b(p,n); b.setRandom();
  C c(m,n); c.setOnes();
  C rc = c;

  std::cout << "Matrix sizes = " << m << "x" << p << " * " << p << "x" << n << "\n";
  std::ptrdiff_t mc(m), nc(n), kc(p);
  internal::computeProductBlockingSizes<Scalar,Scalar>(kc, mc, nc);
  std::cout << "blocking size (mc x kc) = " << mc << " x " << kc << "\n";

  C r = c;

  // check the parallel product is correct
  #if defined EIGEN_HAS_OPENMP
  Eigen::initParallel();
  int procs = omp_get_max_threads();
  if(procs>1)
  {
    #ifdef HAVE_BLAS
    blas_gemm(a,b,r);
    #else
    omp_set_num_threads(1);
    r.noalias() += a * b;
    omp_set_num_threads(procs);
    #endif
    c.noalias() += a * b;
    if(!r.isApprox(c)) std::cerr << "Warning, your parallel product is crap!\n\n";
  }
  #elif defined HAVE_BLAS
    blas_gemm(a,b,r);
    c.noalias() += a * b;
    if(!r.isApprox(c)) {
      std::cout << (r  - c).norm() << "\n";
      std::cerr << "Warning, your product is crap!\n\n";
    }
  #else
    if(1.*m*n*p<2000.*2000*2000)
    {
      gemm(a,b,c);
      r.noalias() += a.cast<Scalar>() .lazyProduct( b.cast<Scalar>() );
      if(!r.isApprox(c)) {
        std::cout << (r  - c).norm() << "\n";
        std::cerr << "Warning, your product is crap!\n\n";
      }
    }
  #endif

  #ifdef HAVE_BLAS
  BenchTimer tblas;
  c = rc;
  BENCH(tblas, tries, rep, blas_gemm(a,b,c));
  std::cout << "blas  cpu         " << tblas.best(CPU_TIMER)/rep  << "s  \t" << (double(m)*n*p*rep*2/tblas.best(CPU_TIMER))*1e-9  <<  " GFLOPS \t(" << tblas.total(CPU_TIMER)  << "s)\n";
  std::cout << "blas  real        " << tblas.best(REAL_TIMER)/rep << "s  \t" << (double(m)*n*p*rep*2/tblas.best(REAL_TIMER))*1e-9 <<  " GFLOPS \t(" << tblas.total(REAL_TIMER) << "s)\n";
  #endif

  BenchTimer tmt;
  c = rc;
  BENCH(tmt, tries, rep, gemm(a,b,c));
  std::cout << "eigen cpu         " << tmt.best(CPU_TIMER)/rep  << "s  \t" << (double(m)*n*p*rep*2/tmt.best(CPU_TIMER))*1e-9  <<  " GFLOPS \t(" << tmt.total(CPU_TIMER)  << "s)\n";
  std::cout << "eigen real        " << tmt.best(REAL_TIMER)/rep << "s  \t" << (double(m)*n*p*rep*2/tmt.best(REAL_TIMER))*1e-9 <<  " GFLOPS \t(" << tmt.total(REAL_TIMER) << "s)\n";

  #ifdef EIGEN_HAS_OPENMP
  if(procs>1)
  {
    BenchTimer tmono;
    omp_set_num_threads(1);
    Eigen::setNbThreads(1);
    c = rc;
    BENCH(tmono, tries, rep, gemm(a,b,c));
    std::cout << "eigen mono cpu    " << tmono.best(CPU_TIMER)/rep  << "s  \t" << (double(m)*n*p*rep*2/tmono.best(CPU_TIMER))*1e-9  <<  " GFLOPS \t(" << tmono.total(CPU_TIMER)  << "s)\n";
    std::cout << "eigen mono real   " << tmono.best(REAL_TIMER)/rep << "s  \t" << (double(m)*n*p*rep*2/tmono.best(REAL_TIMER))*1e-9 <<  " GFLOPS \t(" << tmono.total(REAL_TIMER) << "s)\n";
    std::cout << "mt speed up x" << tmono.best(CPU_TIMER) / tmt.best(REAL_TIMER)  << " => " << (100.0*tmono.best(CPU_TIMER) / tmt.best(REAL_TIMER))/procs << "%\n";
  }
  #endif

  if(1.*m*n*p<30*30*30)
  {
      BenchTimer tmt;
      c = rc;
      BENCH(tmt, tries, rep, c.noalias()+=a.lazyProduct(b));
      std::cout << "lazy cpu         " << tmt.best(CPU_TIMER)/rep  << "s  \t" << (double(m)*n*p*rep*2/tmt.best(CPU_TIMER))*1e-9  <<  " GFLOPS \t(" << tmt.total(CPU_TIMER)  << "s)\n";
      std::cout << "lazy real        " << tmt.best(REAL_TIMER)/rep << "s  \t" << (double(m)*n*p*rep*2/tmt.best(REAL_TIMER))*1e-9 <<  " GFLOPS \t(" << tmt.total(REAL_TIMER) << "s)\n";
  }

  #ifdef DECOUPLED
  if((NumTraits<A::Scalar>::IsComplex) && (NumTraits<B::Scalar>::IsComplex))
  {
    M ar(m,p); ar.setRandom();
    M ai(m,p); ai.setRandom();
    M br(p,n); br.setRandom();
    M bi(p,n); bi.setRandom();
    M cr(m,n); cr.setRandom();
    M ci(m,n); ci.setRandom();

    BenchTimer t;
    BENCH(t, tries, rep, matlab_cplx_cplx(ar,ai,br,bi,cr,ci));
    std::cout << "\"matlab\" cpu    " << t.best(CPU_TIMER)/rep  << "s  \t" << (double(m)*n*p*rep*2/t.best(CPU_TIMER))*1e-9  <<  " GFLOPS \t(" << t.total(CPU_TIMER)  << "s)\n";
    std::cout << "\"matlab\" real   " << t.best(REAL_TIMER)/rep << "s  \t" << (double(m)*n*p*rep*2/t.best(REAL_TIMER))*1e-9 <<  " GFLOPS \t(" << t.total(REAL_TIMER) << "s)\n";
  }
  if((!NumTraits<A::Scalar>::IsComplex) && (NumTraits<B::Scalar>::IsComplex))
  {
    M a(m,p);  a.setRandom();
    M br(p,n); br.setRandom();
    M bi(p,n); bi.setRandom();
    M cr(m,n); cr.setRandom();
    M ci(m,n); ci.setRandom();

    BenchTimer t;
    BENCH(t, tries, rep, matlab_real_cplx(a,br,bi,cr,ci));
    std::cout << "\"matlab\" cpu    " << t.best(CPU_TIMER)/rep  << "s  \t" << (double(m)*n*p*rep*2/t.best(CPU_TIMER))*1e-9  <<  " GFLOPS \t(" << t.total(CPU_TIMER)  << "s)\n";
    std::cout << "\"matlab\" real   " << t.best(REAL_TIMER)/rep << "s  \t" << (double(m)*n*p*rep*2/t.best(REAL_TIMER))*1e-9 <<  " GFLOPS \t(" << t.total(REAL_TIMER) << "s)\n";
  }
  if((NumTraits<A::Scalar>::IsComplex) && (!NumTraits<B::Scalar>::IsComplex))
  {
    M ar(m,p); ar.setRandom();
    M ai(m,p); ai.setRandom();
    M b(p,n);  b.setRandom();
    M cr(m,n); cr.setRandom();
    M ci(m,n); ci.setRandom();

    BenchTimer t;
    BENCH(t, tries, rep, matlab_cplx_real(ar,ai,b,cr,ci));
    std::cout << "\"matlab\" cpu    " << t.best(CPU_TIMER)/rep  << "s  \t" << (double(m)*n*p*rep*2/t.best(CPU_TIMER))*1e-9  <<  " GFLOPS \t(" << t.total(CPU_TIMER)  << "s)\n";
    std::cout << "\"matlab\" real   " << t.best(REAL_TIMER)/rep << "s  \t" << (double(m)*n*p*rep*2/t.best(REAL_TIMER))*1e-9 <<  " GFLOPS \t(" << t.total(REAL_TIMER) << "s)\n";
  }
  #endif

  return 0;
}


#include "../../includes/SpecialFunc.h"
#include "../../includes/utils.h"
#include "../../includes/BaseClass.h"
#include <iostream>

#include <stdio.h>
#include <tgmath.h>
#include <math.h>
#include <complex.h>
#include <boost/math/quadrature/trapezoidal.hpp>

#include <pybind11/pybind11.h>
#include <pybind11/complex.h>
#include <pybind11/numpy.h>
#include <pybind11/stl.h>

namespace py = pybind11;

typedef std::vector<double> Vec;
typedef std::complex<double> complex128;
typedef std::vector<complex128> iVec;
typedef py::array_t<double> ndarray;
typedef py::array_t<complex128> Cndarray;
#define J complex128(0.0,1.0)

#define EPS 1e-10
#define Correction 1.//2.*exp(J*1.0472351410849077)

using boost::math::quadrature::trapezoidal;
using boost::math::constants::two_pi;
using boost::math::constants::pi;

struct argument
  {
  int n, m;
  double r, rhon, angle, s, R0, w0, xi, k;
  complex128 Q, beta;
  Vec Offset;
} args, args0 ;

template <typename func_type>
complex128 simpson_rule(func_type f,
                        double a,
                        double b,
                        int n // Number of intervals
                        )
{
    double h = (b - a) / n;

    // Internal sample points, there should be n - 1 of them
    complex128 sum_odds = 0.0;
    for (int i = 1; i < n; i += 2)
    {
        sum_odds += f(a + i * h);
    }
    complex128 sum_evens = 0.0;
    for (int i = 2; i < n; i += 2)
    {
        sum_evens += f(a + i * h);
    }

    return (f(a) + f(b) + 2. * sum_evens + 4. * sum_odds) * h / 3.;
}

iVec
Q_(double r,
  Vec    theta,
  double w0,
  Vec    Offset,
  double k)
{
   iVec result;
   for (auto const& angle: theta)
   {
     result.push_back(1./( 2. * ( r * cos(angle) - Offset[2]/k )/(k * w0*w0 ) - J ) ) ;
   }
   return result;
}

complex128
Q(double r,
  double theta,
  double w0,
  Vec    Offset,
  double k)
{
  return 1./( 2. * ( r * cos(theta) - Offset[2]/k)/(k * w0*w0 ) - J )  ;
}

complex128
ImSimpson(int m, complex128 beta)
{
 auto func = [=](double angle){return exp( beta * cos(angle)  - J *  (double)m * angle) ;};

 complex128 integral = simpson_rule(func,  0.0, two_pi<double>(), 1000);

 return 1./( two_pi<double>()) * integral;
}

complex128
ImTrapz(int m, complex128 beta)
{
 auto func = [=](double angle)->complex128{return exp( beta * cos(angle)  - J *  (double)m * angle) ;};

 complex128 integral = trapezoidal(func, 0.0, two_pi<double>(), EPS);

 return 1./( two_pi<double>()) * integral;
}

complex128
ImHat(double     m,
      complex128 beta,
      double     xi)
{
  return ImTrapz(m, beta) * exp(-beta - J * m * xi ) ;
}

complex128
I_0(argument arg)
{

  complex128 term0 = pow(-J, arg.n) * arg.rhon*arg.rhon,
             term1 = (2. * (double)arg.n + 1.) * _Psi(0,arg.n, arg.rhon),
             term2 = exp(-J * arg.Offset[2]);

  return term0 / term1 * term2;
}

complex128 I_2(argument arg){ return (2. * arg.Q * arg.s*arg.s * arg.rhon * cos(arg.angle) - 1. ) * sin(arg.angle); }

complex128 I_3(argument arg){ return 4. * arg.Q * arg.s*arg.s * arg.Offset[0] * cos(arg.angle); }

complex128 I_4(argument arg){ return 4. * arg.Q * arg.s*arg.s * arg.Offset[1] * cos(arg.angle); }

complex128
I_1(argument arg)
{
  complex128 term0 = J * arg.Q * arg.s*arg.s * pow(arg.R0 - arg.rhon * sin(arg.angle), 2. ),
             term2 = J * arg.rhon * cos(arg.angle),
             term3 = NPnm(arg.n, abs(arg.m), cos(arg.angle)) * sin(arg.angle) ;

  return arg.Q * exp(term0 + term2 ) * term3;
}

complex128
Bnm_integrand(double angle, argument args0)
{
  if (abs(args0.m) >args0.n){return 0;}

  args0.Q = Q(args0.r, angle, args0.w0, args0.Offset, args0.k);

  args0.angle = angle;

  complex128 beta = -2. * J * args0.Q * args0.s*args0.s * args0.R0 * args0.rhon * sin(angle);

  complex128 term0 =  ImHat(args0.m+1, beta, args0.xi) - ImHat(args0.m-1, beta, args0.xi);

  term0 *= I_2(args0);

  term0 -= (I_4(args0) * ImHat(args0.m, beta, args0.xi));

  term0 *= I_1(args0);

  return term0;

}

complex128
Anm_integrand(double angle, argument args0)
{
  if (abs(args0.m) >args0.n){return 0;}

  args0.Q = Q(args0.r, angle, args0.w0, args0.Offset, args0.k);

  args0.angle = angle;

  complex128 beta = -2. * J * args0.Q * args0.s*args0.s * args0.R0 * args0.rhon * sin(angle);

  complex128 term0 =  ImHat(args0.m+1, beta, args0.xi) + ImHat(args0.m-1, beta, args0.xi);

  term0 *= I_2(args0);

  term0 -= (I_3(args0) * ImHat(args0.m, beta, args0.xi));

  term0 *= I_1(args0);

  return term0;

}

complex128
Anm(int    n,
    int    m,
    double k,
    double w0,
    Vec    Offset,
    double Tolerance)
{
  argument args0;
           args0.n      = n;
           args0.m      = m;
           args0.rhon   = (n + 0.5);
           args0.Offset = Offset;
           args0.k      = k;
           args0.R0     = sqrt(Offset[0]*Offset[0] + Offset[1]*Offset[1]);
           args0.xi     = acos(Offset[0]/args0.R0);
           args0.s      = 1./(k*w0);
           args0.w0     = w0;
           args0.r      = args.rhon/k;

  auto func_ = [=](double angle)->complex128 {return Anm_integrand(angle, args0);};

  return -trapezoidal(func_, 0.0, pi<double>(),Tolerance) * I_0(args0) * Correction;

}

complex128
Bnm(int n,
    int m,
    double k,
    double w0,
    Vec Offset,
    double Tolerance)
{
  argument args0;
           args0.n      = n;
           args0.m      = m;
           args0.rhon   = (n + 0.5);
           args0.Offset = Offset;
           args0.k      = k;
           args0.R0     = sqrt(Offset[0]*Offset[0] + Offset[1]*Offset[1]);
           args0.xi     = acos(Offset[0]/args0.R0);
           args0.s      = 1./(k*w0);
           args0.w0     = w0;
           args0.r      = args.rhon/k;

  auto func_ = [=](double angle)->complex128 {return Bnm_integrand(angle, args0);};

  return -trapezoidal(func_, 0.0, pi<double>(),Tolerance) * I_0(args0) * Correction;

}

std::tuple<ndarray,Cndarray>
PyAnm_integrand(int    n,
                int    m,
                int    sampling,
                double k,
                double w0,
                Vec    Offset,
                double Tolerance)
{

 argument args0;
          args0.n      = n;
          args0.m      = m;
          args0.rhon   = (n + 0.5);
          args0.Offset = Offset;
          args0.k      = k;
          args0.R0     = sqrt(Offset[0]*Offset[0] + Offset[1]*Offset[1]);
          args0.xi     = acos(Offset[0]/args0.R0);
          args0.s      = 1./(k*w0);
          args0.w0     = w0;
          args0.r      = args.rhon/k;

  Cndarray _Anm = Cndarray(sampling);
  auto _Anm_data = _Anm.mutable_data();

  Cndarray Angle = Cndarray(sampling);
  auto Angle_data = Angle.mutable_data();

  Vec X = Linspace( 0.0, pi<double>(),sampling);

  for (auto i=0; i<sampling;i++)
  {
    _Anm_data[i] = Anm_integrand(X[i], args);
    Angle_data[i] = X[i];
  }
  return std::make_tuple(Angle, _Anm);
}

std::tuple<ndarray,Cndarray>
PyBnm_integrand(int    n,
                int    m,
                int    sampling,
                double k,
                double w0,
                Vec    Offset,
                double Tolerance)
{

  argument args0;
           args0.n      = n;
           args0.m      = m;
           args0.rhon   = (n + 0.5);
           args0.Offset = Offset;
           args0.k      = k;
           args0.R0     = sqrt(Offset[0]*Offset[0] + Offset[1]*Offset[1]);
           args0.xi     = acos(Offset[0]/args0.R0);
           args0.s      = 1./(k*w0);
           args0.w0     = w0;
           args0.r      = args.rhon/k;

   Cndarray _Bnm = Cndarray(sampling);
   auto _Bnm_data = _Bnm.mutable_data();

   Cndarray Angle = Cndarray(sampling);
   auto Angle_data = Angle.mutable_data();

   Vec X = Linspace( 0.0, pi<double>(),sampling);

   for (auto i=0; i<sampling;i++)
   {
     _Bnm_data[i] = Bnm_integrand(X[i], args);
     Angle_data[i] = X[i];
   }
   return std::make_tuple(Angle, _Bnm);
 }

PYBIND11_MODULE(GaussianBeam, module) {
    module.doc() = "Generalized Lorenz-Mie Theory (GLMT) c++ binding module for light scattering from a spherical scatterer";

    module.def("Anm",
               &Anm,
               py::arg("n"),
               py::arg("m"),
               py::arg("k"),
               py::arg("w0"),
               py::arg("Offset"),
               py::arg("Tolerance") = 1e-8,
               "Compute Anm");

     module.def("Anm_integrand",
                &PyAnm_integrand,
                py::arg("n"),
                py::arg("m"),
                py::arg("sampling"),
                py::arg("k"),
                py::arg("w0"),
                py::arg("Offset"),
                py::arg("Tolerance") = 1e-8,
                "Compute Anm integrand as a function of theta");

    module.def("Bnm",
               &Bnm,
               py::arg("n"),
               py::arg("m"),
               py::arg("k"),
               py::arg("w0"),
               py::arg("Offset"),
               py::arg("Tolerance") = 1e-8,
               "Compute Bnm");

     module.def("Bnm_integrand",
                &PyBnm_integrand,
                py::arg("n"),
                py::arg("m"),
                py::arg("sampling"),
                py::arg("k"),
                py::arg("w0"),
                py::arg("Offset"),
                py::arg("Tolerance") = 1e-8,
                "Compute Bnm integrand as a function of theta");
}

//-


#include <iostream>
#include <optional>

#include <boost/asio.hpp>

#include "vdf.h"

using namespace std;

int gcd_base_bits = 50;
int gcd_128_max_iter = 3;

struct vdf_result {
    // TODO?: add D, T
    form y;
    form p;
};

struct bluebox_request {
    integer discriminant;
    uint64_t iterations;
};

void repeated_square_1weso(const integer& D, uint64_t T, WesolowskiCallback* weso) {
    vdf_original vdfo;
    integer L = root(-D, 4);
    form x = form::generator(D);

    uint64_t num_iterations = 0;
    while (num_iterations <= T) {
        uint64 batch_size = checkpoint_interval;

        square_state_type square_state;
        square_state.pairindex = 0;

        // Overall throughput for blueboxing significantly increases if we run
        // one single-threaded worker per CPU (hyper-)thread, instead of one
        // dual-threaded worker per every two CPU (hyper-)threads.
        uint64 actual_iterations = repeated_square_fast_single_thread(square_state, x, D, L, num_iterations, batch_size, weso);
        if (actual_iterations == ~uint64(0)) {
            // Corruption; f is unchanged. Do the entire batch with the slow
            // algorithm.
            repeated_square_original(vdfo, x, D, L, num_iterations, batch_size, weso);
            actual_iterations = batch_size;
        } else if (actual_iterations < batch_size) {
            // The fast algorithm terminated prematurely for whatever reason. f
            // is still valid. It might terminate prematurely again (e.g. gcd
            // quotient too large), so do one iteration of the slow algorithm.
            // This will also reduce f if the fast algorithm terminated because
            // it was too big.
            repeated_square_original(vdfo, x, D, L, num_iterations + actual_iterations, 1, weso);
            actual_iterations += 1;
        }
        num_iterations += actual_iterations;
    }
}

form prove_1weso(integer D, uint64_t T, form y, form* intermediates) {
    form x = form::generator(D);

    // OneWesolowskiProver(Segment(0, T, x, y), D, intermediates, stopped)
    uint32_t k;
    uint32_t l;
    if (T >= (1 << 16)) {
        ApproximateParameters(T, k, l);
    } else {
        k = 10;
        l = 1;
    }

    // GenerateProof()
    uint64_t k1 = k / 2;
    uint64_t k0 = k - k1;

    PulmarkReducer reducer;

    integer B = GetB(D, x, y);
    integer L = root(-D, 4);
    form id = form::identity(D);
    form p = id;

    for (int64_t j = l - 1; j >= 0; j--) {
        p = FastPowFormNucomp(p, D, integer(1 << k), L, reducer);

        vector<form> ys((1 << k));
        for (uint64_t i = 0; i < (1 << k); i++) {
            ys[i] = id;
        }

        uint64_t limit = T / (k * l);
        if (T % (k * l)) {
            limit++;
        }

        for (uint64_t i = 0; i < limit; i++) {
            if (T >= k * (i * l + j + 1)) {
                // uint64_t b = GetBlock(i*l + j, k, T, B);
                integer res = FastPow(2, T - k * (i*l + j + 1), B);
                mpz_mul_2exp(res.impl, res.impl, k);
                res = res / B;
                auto res_vector = res.to_vector();  // @@ smells
                uint64_t b = res_vector.empty() ? 0 : res_vector[0];

                nucomp_form(ys[b], ys[b], intermediates[i], D, L);
            }
        }

        for (uint64_t b1 = 0; b1 < (1 << k1); b1++) {
            form z = id;
            for (uint64_t b0 = 0; b0 < (1 << k0); b0++) {
                nucomp_form(z, z, ys[b1 * (1 << k0) + b0], D, L);
            }
            z = FastPowFormNucomp(z, D, integer(b1 * (1 << k0)), L, reducer);
            nucomp_form(p, p, z, D, L);
        }

        for (uint64_t b0 = 0; b0 < (1 << k0); b0++) {
            form z = id;
            for (uint64_t b1 = 0; b1 < (1 << k1); b1++) {
                nucomp_form(z, z, ys[b1 * (1 << k0) + b0], D, L);
            }
            z = FastPowFormNucomp(z, D, integer(b0), L, reducer);
            nucomp_form(p, p, z, D, L);
        }
    }
    reducer.reduce(p);

    return p;
}

vdf_result vdf_1weso(integer& D, uint64_t T) {
    cout << "[vdf] Computing repeated square x^2^" << T << "." << endl << flush;

    form x = form::generator(D);
    cout << "[vdf] Base: x.c=" << x.c.impl << endl;
    OneWesolowskiCallback weso_cb(D, x, T);
    auto t0 = chrono::high_resolution_clock::now();
    repeated_square_1weso(D, T, &weso_cb);
    auto t1 = chrono::high_resolution_clock::now();
    auto td = chrono::duration<double>(t1 - t0).count();
    form y = weso_cb.result;
    form *intermediates = weso_cb.forms.get();  // @@ copy/ref, lifetime = weso_cb
    cout << "[vdf] Result: y.a=" << y.a.impl << " y.b=" << y.b.impl << endl;
    if (td > 0) {
        uint32_t kips = T / td / 1000;
        cout << "[vdf] Computed repeated square. kips=" << kips << endl << flush;
    } else {
        cout << "[vdf] Computed repeated square. Too fast for kIPS calculation." << endl << flush;
    }

    cout << "[vdf] Starting to compute Wesolowski proof." << endl << flush;
    form p = prove_1weso(D, T, y, intermediates);
    cout << "[vdf] Proof: p.a=" << p.a.impl << " p.b=" << p.b.impl << endl;
    cout << "[vdf] Finished computing Wesolowski proof." << endl << flush;

    return vdf_result{y, p};
}

string create_response(integer discriminant, uint64_t iterations, form y, form p) {
    int witness_type = 0;
    int discriminant_bits = discriminant.num_bits();
    // = res[BQFC_FORM_SIZE]; reduce + bqfc_serialize(res, a, b, discriminant_bits);
    // ~ core operator: bqfc_compress(a, b) -> a, t, g, b0, b_sign
    std::vector<uint8_t> y_serialized = SerializeForm(y, discriminant_bits);;
    std::vector<uint8_t> p_serialized = SerializeForm(p, discriminant_bits);;

    std::vector<uint8_t> bytes;
    uint8_t int64_bytes[8];

    // Writes the number of iterations
    Int64ToBytes(int64_bytes, iterations);
    VectorAppendArray(bytes, int64_bytes, sizeof(int64_bytes));

    // Writes the y, with prepended size
    Int64ToBytes(int64_bytes, y_serialized.size());
    VectorAppendArray(bytes, int64_bytes, sizeof(int64_bytes));
    VectorAppend(bytes, y_serialized);

    // Writes the witness type
    bytes.push_back(witness_type);

    // Writes the proof
    VectorAppend(bytes, p_serialized);

    return BytesToStr(bytes);
}

tcp::socket connect_client(boost::asio::io_context& io_context, string host, string port) {
    tcp::socket socket(io_context);
    tcp::resolver resolver(io_context);
    boost::asio::connect(socket, resolver.resolve(host, port, boost::asio::ip::resolver_base::address_configured));
    return socket;
}

bluebox_request read_request(tcp::socket &socket, const optional<string> client_id) {
    char buffer[1024];
    uint64_t size;

    cout << "[net] Reading request." << endl;

    memset(buffer, 0, sizeof(buffer));
    boost::asio::read(socket, boost::asio::buffer(buffer, 1));
    // cout << "[net] Prover type: " << buffer[0] << endl;
    assert (buffer[0] == 'S');

    memset(buffer, 0, sizeof(buffer));
    boost::asio::read(socket, boost::asio::buffer(buffer, 3));
    size = stoull(buffer);
    // cout << "[net] Discriminant size: " << size << endl;

    memset(buffer, 0, sizeof(buffer));
    boost::asio::read(socket, boost::asio::buffer(buffer, size));
    integer discriminant(buffer);
    cout << "[net] Discriminant: D=" << discriminant.impl << endl;

    memset(buffer, 0, sizeof(buffer));
    boost::asio::read(socket, boost::asio::buffer(buffer, 1));
    size = buffer[0];

    memset(buffer, 0, sizeof(buffer));
    boost::asio::read(socket, boost::asio::buffer(buffer, size));
    // initial_form is not used for blueboxing. TODO: debug print nevertheless

    if (client_id) {
        string id = client_id.value();
        assert (id.size() < 256);
        uint8_t len = id.size();
        boost::asio::write(socket, boost::asio::buffer("ID", 2));
        boost::asio::write(socket, boost::asio::buffer(&len, 1));
        boost::asio::write(socket, boost::asio::buffer(id, len));
    } else {
        boost::asio::write(socket, boost::asio::buffer("OK", 2));
    }

    memset(buffer, 0, sizeof(buffer));
    boost::asio::read(socket, boost::asio::buffer(buffer, 2));
    size = stoull(buffer);
    // cout << "[net] Iterations size: " << size << endl;

    memset(buffer, 0, sizeof(buffer));
    boost::asio::read(socket, boost::asio::buffer(buffer, size));
    uint64_t iterations = stoull(buffer);
    cout << "[net] Iterations: T=" << iterations << endl;

    return bluebox_request{discriminant, iterations};
}

void interact_ping(tcp::socket& socket) {
    cout << "[net] Sending PING." << endl;
    boost::asio::write(socket, boost::asio::buffer("PING", 4));

    char buffer[1024];
    boost::asio::read(socket, boost::asio::buffer(buffer, 4));
    assert (strncmp(buffer, "PONG", 4) == 0);
    cout << "[net] Received PONG." << endl;
}

void write_response(tcp::socket& socket, string& res) {
    uint8_t res_size[4];
    Int32ToBytes(res_size, res.size());
    boost::asio::write(socket, boost::asio::buffer(res_size, 4));
    boost::asio::write(socket, boost::asio::buffer(res.c_str(), res.size()));

    cout << "[net] Sending STOP." << endl;
    boost::asio::write(socket, boost::asio::buffer("STOP", 4));

    cout << "[net] Waiting for ACK." << endl;
    char buffer[1024];
    boost::asio::read(socket, boost::asio::buffer(buffer, 3));
    assert (strncmp(buffer, "ACK", 3) == 0);

    cout << "[net] Sent response." << endl;
}

void dump_response(string &res) {
    uint8_t res_size[4];
    Int32ToBytes(res_size, res.size());

    cout << "[net] Response: size=0x";
    {
        auto flags = cout.flags();
        cout << hex << setfill('0')
            << setw(2) << +res_size[0]
            << setw(2) << +res_size[1]
            << setw(2) << +res_size[2]
            << setw(2) << +res_size[3];
        cout.flags(flags);
    }
    cout << " data=" << res << endl;
}

int run_vdf(int argc, char *argv[]) {
    integer D(argv[2]);
    uint64_t T = atoi(argv[3]);

    cout << "Discriminant: D=" << D.impl << endl;
    cout << "Iterations: T=" << T << endl;

    vdf_result r = vdf_1weso(D, T);

    string res = create_response(D, T, r.y, r.p);
    dump_response(res);

    return 0;
}

int run_client(int argc, char *argv[]) {
    auto host = argv[2];
    auto port = argv[3];
    auto client_id = argc > 3 ? optional<string>{argv[4]} : nullopt;

    boost::asio::io_context io_context;
    tcp::socket socket = connect_client(io_context, host, port);

    bluebox_request req = read_request(socket, client_id);

    future<vdf_result> vdf_future = async(vdf_1weso, ref(req.discriminant), req.iterations);
    future_status vdf_status;
    do {
        switch (vdf_status = vdf_future.wait_for(chrono::minutes(3)); vdf_status) {
            case future_status::timeout: interact_ping(socket); break;
        }
    } while (vdf_status != future_status::ready);
    vdf_result r = vdf_future.get();

    string res = create_response(req.discriminant, req.iterations, r.y, r.p);
    dump_response(res);

    write_response(socket, res);

    return 0;
}

int main(int argc, char *argv[]) try {
    if (argc < 3 || (argv[1][0] != 'p' && argv[1][0] != 'c')) {
        cerr << "Usage: " << argv[0] << " <mode> <args>" << endl;
        cerr << "  " << argv[0] << " p <discriminant> <iterations>" << endl;
        cerr << "  " << argv[0] << " c <host> <port> [<client_id>]" << endl;
        return 64;
    }

    init_gmp();
    if (hasAVX2()) {
        gcd_base_bits = 63;
        gcd_128_max_iter = 2;
    }
    set_rounding_mode();

    switch (argv[1][0]) {
    case 'p':
        return run_vdf(argc, argv);
    case 'c':
        return run_client(argc, argv);
    }
    cerr << "nyi" << endl;
    return 64;
} catch (exception &e) {
    cerr << "[cli] Exception: " << e.what() << endl;
    return 1;
}

// TODO: better cli parsing
// TODO: progress! (bar?), new proto
// TODO: proper logging?
// TODO: looped client? (would make timelord-launcher obsolete)


#include "solvers/SolverBoxBPP.h"

#include "rigidbody/RigidBodySystem.h"
#include "rigidbody/RigidBody.h"
#include "contact/Contact.h"

#include <Eigen/Dense>

namespace
{
    enum eIndexSet { kFree = 0, kLower, kUpper, kIgnore };

    static inline void multAndSub(const JBlock& G, const Eigen::Vector3f& x, const Eigen::Vector3f& y, const float a, Eigen::Ref<Eigen::VectorXf> b)
    {
            b -= a * G.col(0) * x(0);
            b -= a * G.col(1) * x(1);
            b -= a * G.col(2) * x(2);
            b -= a * G.col(3) * y(0);
            b -= a * G.col(4) * y(1);
            b -= a * G.col(5) * y(2);
    }

    static inline unsigned int initContacts(std::vector<Contact*>& contacts)
    {
        // Count the number of constraint rows
        // and assign each contact an index.
        //
        unsigned int rows = 0;
        for(auto c : contacts)
        {
            c->index = rows;
            rows += c->J0.rows();
        }

        return rows;
    }

    // Update the box bounds, lower and upper, of the constraint impulses.
    // The value in Contact::lambda is used for updating the bounds.
    //
    static inline void updateBounds(std::vector<Contact*>& contacts, Eigen::VectorXf& lower, Eigen::VectorXf& upper)
    {
        for(auto c : contacts)
        {
            const unsigned int dim = c->J0.rows();
            // Non-interpenetration row.
            //
            lower(c->index) = 0.0f;
            upper(c->index) = std::numeric_limits<float>::max();

            // Friction rows.
            // Compute the box bounds as [-mu*lambda_n, mu*lambda_n]
            // which is an approximation of the isotropic Coulomb friction cone.
            //
            for(unsigned int i = 1; i < dim; ++i)
            {
                lower(c->index + i) = -c->mu * c->lambda(0);
                upper(c->index + i) = c->mu * c->lambda(0);
            }
        }
    }

    // Build the rhs vector of the Schur complement linear system.
    //
    static inline void buildRHS(const std::vector<Contact*>& contacts, float h, Eigen::VectorXf& b)
    {
        for(auto c : contacts)
        {
            const float gamma = h * c->k / (h * c->k + c->b); // error reduction parameter
            const unsigned int dim = c->J0.rows();
            b.segment(c->index, dim) = -gamma * c->phi / h;

            multAndSub(c->J0, c->body0->xdot, c->body0->omega, 1.0f, b.segment(c->index, dim));
            multAndSub(c->J1, c->body1->xdot, c->body1->omega, 1.0f, b.segment(c->index, dim));

            if( !c->body0->fixed )
            {
                multAndSub(c->J0Minv, c->body0->f, c->body0->tau, h, b.segment(c->index, dim));
            }
            if( !c->body1->fixed )
            {
                multAndSub(c->J1Minv, c->body1->f, c->body1->tau, h, b.segment(c->index, dim));
            }
        }
    }

    // Build the Schur complement system using the contact constraints.
    //
    static inline void buildMatrix(const std::vector<Contact*>& contacts, float h, Eigen::MatrixXf& A)
    {
        for(auto c : contacts)
        {
            const unsigned int dim = c->J0.rows();
            const float eps = 1.0f / (h * h * c->k + h * c->b);    // constraint force mixing

            A.block(c->index, c->index, dim, dim) = 1e-10f * Eigen::MatrixXf::Identity(dim, dim);
            A(c->index, c->index) += eps;

            if( !c->body0->fixed )
            {
                A.block(c->index, c->index, dim, dim) += c->J0Minv * c->J0.transpose();
                for(auto cc : c->body0->contacts)
                {
                    if( cc != c )
                    {
                        const int ddim = c->J0.rows();
                        if( cc->body0 == c->body0 )
                        {
                            A.block(c->index, cc->index, dim, ddim) += c->J0Minv * cc->J0.transpose();
                        }
                        else
                        {
                            A.block(c->index, cc->index, dim, ddim) += c->J0Minv * cc->J1.transpose();
                        }
                    }
                }
            }

            if( !c->body1->fixed )
            {
                A.block(c->index, c->index, dim, dim) += c->J1Minv * c->J1.transpose();

                for(auto cc : c->body1->contacts)
                {
                    if( cc != c )
                    {
                        const int ddim = c->J1.rows();
                        if( cc->body0 == c->body1 )
                        {
                            A.block(c->index, cc->index, dim, ddim) += c->J1Minv * cc->J0.transpose();
                        }
                        else
                        {
                            A.block(c->index, cc->index, dim, ddim) += c->J1Minv * cc->J1.transpose();
                        }
                    }
                }
            }
        }

    }

    // Pivoting rules.
    // The index set is updates based on the LCP variables x and v.
    // Specifically, 'free' variables are pivoted to the 'tight' set if the lower or upper bounds are violated.
    // Similarly, 'tight' variables are pivoted to the 'free' set if the residual velocity v has the wrong sign.
    //
    static inline unsigned int pivot(Eigen::VectorXi& idx, Eigen::VectorXf& x, const Eigen::VectorXf& l, const Eigen::VectorXf& u, const Eigen::VectorXf& v)
    {
        static const float tol = 1e-5f;
        unsigned int numPivots = 0;
        const unsigned int n = idx.rows();
        for (unsigned int j = 0; j < n; ++j)
        {
            // By default, assume all variables belong to the 'free' set.
            //
            int new_idx = kFree;

            if (idx[j] == kIgnore) continue;

            if(idx[j] == kFree && x[j] <= l[j])     // case: free variable, lower bound
            {
                new_idx = kLower;
            }
            else if(idx[j] == kFree && x[j] >= u[j])  // case: free variable, upper bound
            {
                new_idx = kUpper;
            }
            else if(idx[j] == kLower && v[j] > -tol)    // case: lower tight variable, +ve velocity
            {
                new_idx = kLower;
            }
            else if(idx[j] == kUpper && v[j] < tol)     // case: upper tight variable, -ve velocity
            {
                new_idx = kUpper;
            }

            if(new_idx != idx[j])
            {
                ++numPivots;
                idx[j] = new_idx;
            }
        }
        return numPivots;
    }

    // Returns the indices of tight variables in tightIdx.
    //
    static inline unsigned int tightIndices(const Eigen::VectorXi& idx, std::vector<int>& tightIdx)
    {
        const unsigned int n = idx.rows();
        unsigned int numTight = 0;
        tightIdx.clear();
        for(unsigned int i = 0; i < n; ++i)
        {
            if( idx[i] == kLower || idx[i] == kUpper )
            {
                ++numTight;
                tightIdx.push_back(i);
            }
        }
        return numTight;
    }

    // Returns the indices of free variables in freeIdx.
    //
    static inline unsigned int freeIndices(const Eigen::VectorXi& idx, std::vector<int>& freeIdx)
    {
        const unsigned int n = idx.rows();
        unsigned int numFree = 0;
        freeIdx.clear();
        for(unsigned int i = 0; i < n; ++i)
        {
            if( idx[i] == kFree )
            {
                ++numFree;
                freeIdx.push_back(i);
            }
        }
        return numFree;
    }

    // Solve the principal sub-problem comprising the free variables:
    //      Aff * xf = bf - Aft * xt
    //
    //  where Aff is the sub-matrix of free variables, bf are the corresponding entries in the rhs vector,
    //  Aft is the sub-matrix that couples the tight variables and the free variables, and
    //  xt are the tight variables whose values is determined by the lower and upper bounds (l and u).
    //
    // Inputs:
    //    A - the lead matrix
    //    b - the rhs vector
    //    idx - the index set of all variables
    //    l - the lower bounds
    //    u - the upper bounds
    //
    // Outputs:
    //    x - the solution of constraint impulses (free and tight variables)
    //
    static inline void solvePrincipalSubproblem(const Eigen::MatrixXf& A,
                                                const Eigen::VectorXf& b,
                                                const Eigen::VectorXi& idx,
                                                const Eigen::VectorXf& l,
                                                const Eigen::VectorXf& u,
                                                Eigen::VectorXf& x)
    {
        std::vector<int> freeIdx, tightIdx;
        const unsigned int numTight = tightIndices(idx, tightIdx);
        const unsigned int numFree = freeIndices(idx, freeIdx);
        if( numFree > 0 )
        {
            Eigen::MatrixXf Aff(numFree, numFree);
            Eigen::VectorXf bf(numFree);

            // Build sub-matrix using free indices
            //
            for(unsigned int j = 0; j < numFree; ++j)
            {
                for(unsigned int i = 0; i < numFree; ++i)
                {
                    Aff(i,j) = A(freeIdx[i], freeIdx[j]);
                }
            }

            // Build rhs vector using free indices
            //
            for (unsigned int i = 0; i < numFree; ++i)
            {
                bf(i) = b(freeIdx[i]);
            }

            // Update rhs vector with impulses from tight indices
            //  e.g.     bf -= A_ft * x_t
            //
            for(unsigned int j = 0; j < numTight; ++j)
            {
                for(unsigned int i = 0; i < numFree; ++i)
                {
                    if( idx[tightIdx[j]] == kLower )
                    {
                        bf(i) -= A(freeIdx[i], tightIdx[j]) * l(tightIdx[j]);
                    }
                    else if( idx[tightIdx[j]] == kUpper )
                    {
                        bf(i) -= A(freeIdx[i], tightIdx[j]) * u(tightIdx[j]);
                    }
                }
            }

            // Cholesky solve for the principal sub-problem.
            //
            Eigen::LDLT<Eigen::MatrixXf> ldlt(Aff);
            const Eigen::VectorXf xf = ldlt.solve(bf);

            // Update free variables with the solution
            for(unsigned int i = 0; i < numFree; ++i)
            {
                x(freeIdx[i]) = xf(i);
            }

            // Update tight variables with values from lower/upper bounds.
            //
            for(unsigned int i = 0; i < numTight; ++i)
            {
                if( idx[tightIdx[i]] == kLower )
                {
                    x(tightIdx[i]) = l(tightIdx[i]);
                }
                else if( idx[tightIdx[i]] == kUpper )
                {
                    x(tightIdx[i]) = u(tightIdx[i]);
                }
            }
        }

    }

    static inline void updateContacts(const Eigen::VectorXf& x, std::vector<Contact*>& contacts)
    {
        // Distribute impulses to the contacts
        //
        for(auto c : contacts)
        {
            const unsigned int dim = c->J0.rows();
            c->lambda = x.segment(c->index, dim);
        }
    }

}

SolverBoxBPP::SolverBoxBPP(RigidBodySystem* _rigidBodySystem) : Solver(_rigidBodySystem)
{

}

void SolverBoxBPP::solve(float h)
{
    auto contacts = m_rigidBodySystem->getContacts();

    const unsigned int numContacts = contacts.size();
    if( numContacts > 0 )
    {
        const unsigned int dim = initContacts(contacts);

        Eigen::MatrixXf A = Eigen::MatrixXf::Zero(dim,dim);
        Eigen::VectorXf b = Eigen::VectorXf::Zero(dim);
        Eigen::VectorXf x = Eigen::VectorXf::Zero(dim);

        // Update box bounds.
        //
        Eigen::VectorXf lower = Eigen::VectorXf::Zero(dim);
        Eigen::VectorXf upper = Eigen::VectorXf::Zero(dim);

		updateBounds(contacts, lower, upper);

		// Initialize the index set.
        // All variables are initially set to 'free'.
        //
		Eigen::VectorXi idx = Eigen::VectorXi::Constant(dim, kFree);

		// Ignore variables where lower and upper are zero.
		//
		for (int i = 0; i < dim; ++i)
		{
			static const float tol = 1e-5f;
			if (std::abs(upper(i)) < tol && std::abs(lower(i) < tol))
				idx(i) = kIgnore;
		}

        // Construct the lead matrix and rhs vector.
        //
        buildMatrix(contacts, h, A);
        buildRHS(contacts, h, b);

        // Perform an initial solve to update friction box bounds.
        //
        solvePrincipalSubproblem(A, b, idx, lower, upper, x);
        updateContacts(x, contacts);
        updateBounds(contacts, lower, upper);
        idx.setConstant(kFree);

        // Block pivoting iterations.
        //
        for(unsigned int iter = 0; iter < m_maxIter; ++iter)
        {
            // Solve the principal sub-problem:
            //    Aff * xf = bf - Aft * xt
            //
            solvePrincipalSubproblem(A, b, idx, lower, upper, x);

            // Compute residual velocity v.
            //
            const Eigen::VectorXf v = A*x - b;

            // Pivot.
            //
            const int numPivots = pivot(idx, x, lower, upper, v);
            if( numPivots == 0 )
                break;      // Done
        }

        updateContacts(x, contacts);
        updateBounds(contacts, lower, upper);
    }
}


    // KALMANFILTER3D_UPDATE - updates a EKF with 9 state variables
    // Called each time a new mobile-to-anchor measurement is available

    // INPUTS:
    // X - the 9x1 a priori state vector - [x,xdot,ax_bias,y,ydot,ay_bias,z,zdot,az_bias]'
    // P - the 9x9 a priori covariance matrix
    // T - the delta time between updates
    // tao_acc - variance of acceleratmetor's gaussian noise
    // tao_bias - variance of acceleratmetor bias's noise
    // sigma_r - the estimated standard deviation of the new range measurement
    // r_meas - the actual range measurement
    // OUTPUTS:
    // X - the new state vector estimate
    // P - the new covariance matrix estimate
    // error - the difference between the range measurement and the estimated
    // range measurement (perhaps useful for outlier filtering.)

    // determine whether to perform filter using SNR
    // According to experiments on 2017-03-24,
    // Time Domain UWB range info have incorrect vPeak when satuatured,
    // probably because of Int16 overflow
    // float SNR = 20 * std::log( range_info.vPeak / ( range_info.noise + 0.1) );
    //   if(SNR < m_snr_threshold){
    //       //ROS_WARN("Anchor %d SNR %f too small, discard.", anchor_id, SNR);
    //       return false;
    //   }
    // X(0) = m_position.point.x;
    // X(1) = m_velocity.point.x;
    // X(2) = m_acc_bias.point.x;
    // X(3) = m_position.point.y;
    // X(4) = m_velocity.point.y;
    // X(5) = m_acc_bias.point.y;
    // X(6) = m_position.point.z;
    // X(7) = m_velocity.point.z;
    // X(8) = m_acc_bias.point.z;

    // ------------------------------------------------------------------
#include <Eigen/Dense>
#include <Eigen/Core>
#include <Eigen/Geometry>
#include <bits/stdc++.h>
#include <ros/ros.h>
#include <math.h>
#include <stdio.h>
#include <dwm1001/anchor.h>
#include <sensor_msgs/Imu.h>
#include "std_msgs/MultiArrayLayout.h"
#include "std_msgs/MultiArrayDimension.h"
#include "std_msgs/Float32MultiArray.h"
#include "geometry_msgs/Vector3Stamped.h"
#include "geometry_msgs/PoseStamped.h"
#include "mavros_msgs/Altitude.h"
#include "sensor_msgs/Range.h"
#include "geometry_msgs/TwistStamped.h"

// Declarations
double r_pred;
//int count =0 ;
Eigen::MatrixXd H(1,9);
double R;
Eigen::VectorXd K(9);
double error,error_threshold;
double precisionRangeErrEst,precisionRangeMm;
//
//geometry_msgs::Vector3Stamped pose;
geometry_msgs::PoseStamped pose;
geometry_msgs::PoseStamped imu_msg;
float ax,ay,az;
Eigen::Quaternionf q;
Eigen::Quaternionf q_;
Eigen::MatrixXf R_mat(3,3);
Eigen::MatrixXf imuacc(3,1);
Eigen::MatrixXf acc(3,1);
//
Eigen::VectorXd X(9);
Eigen::VectorXd X_p(9);
Eigen::VectorXd X_e(9);
Eigen::MatrixXd F(9,9);
Eigen::MatrixXd block_F(3,3);
Eigen::MatrixXd Q(9,9);
Eigen::VectorXd u(9);
Eigen::MatrixXd B(9,9);
Eigen::MatrixXd P(9,9);
Eigen::MatrixXd P_p(9,9);
Eigen::MatrixXd M(9,9);
Eigen::MatrixXd block_B(3,3);
Eigen::MatrixXd block_Q(3,3);

double r_vel;
Eigen::MatrixXd r_meas_vel(3,1);
Eigen::MatrixXd r_pred_vel(3,1);
Eigen::MatrixXd H_vel(3,9);
Eigen::MatrixXd R_vel(3,3);
Eigen::MatrixXd K_vel(9,3);
Eigen::MatrixXd P_p_vel(9,9);
Eigen::VectorXd X_p_vel(9);

double tao_acc;
double tao_bias,error_1;
double sigma_r,sigma_a,r_meas;
double T,yaw_off,yaw,yaw_rad;
double m_R_scale;
double m_last_range_time;
std_msgs::Float32MultiArray output;

double m_kalman_sigma_a,T_sq,m_tao_acc_sqrt,m_tao_bias_sqrt,T_cub,m_z_damping_factor,m_Q_scale;
double m_snr_threshold;
double x,y,z;
double ax_,ay_;

double r_h,h_offset,h_meas;
Eigen::VectorXd X_h(9);
Eigen::MatrixXd P_h(9,9);
Eigen::VectorXd K_h(9);
Eigen::MatrixXd H_h(1,9);
Eigen::MatrixXd nine_cov = Eigen::MatrixXd::Identity(9,9);
double trans_x,trans_y;
Eigen::MatrixXf rpy(3,1);
int set_=0,set_p=0;

void setState(Eigen::VectorXd Y)
{
    X = Y;
}
void setCovariance(Eigen::MatrixXd S)
{
    P = S;
}

void param(ros::NodeHandle& nh)
{
    nh.getParam("KalmanFilter/start_x", x);
    nh.getParam("KalmanFilter/start_y", y);
    nh.getParam("KalmanFilter/start_z", z);
    nh.getParam("KalmanFilter/error_threshold", error_threshold);
    nh.getParam("KalmanFilter/m_R_scale", m_R_scale);
    nh.getParam("KalmanFilter/precisionRangeErrEst", precisionRangeErrEst);
    nh.getParam("KalmanFilter/m_kalman_sigma_a", m_kalman_sigma_a);
    nh.getParam("KalmanFilter/precisionRangeMm", precisionRangeMm);
    nh.getParam("KalmanFilter/m_tao_acc_sqrt", m_tao_acc_sqrt);
    nh.getParam("KalmanFilter/m_tao_bias_sqrt", m_tao_bias_sqrt );
    nh.getParam("KalmanFilter/m_z_damping_factor", m_z_damping_factor);
    nh.getParam("KalmanFilter/m_Q_scale", m_Q_scale);
    nh.getParam("KalmanFilter/r_h", r_h);
    nh.getParam("KalmanFilter/r_vel", r_vel);
    nh.getParam("KalmanFilter/trans_x", trans_x);
    nh.getParam("KalmanFilter/trans_y", trans_y);
    nh.getParam("KalmanFilter/yaw_offset", yaw_off);
}

void Initialize(ros::NodeHandle& nh)
{

    int static count_ =0 ;
    if(count_ == 0)
    {   param(nh);
        count_++; }
    else{
    std::string pose_topic("initialize");
    // nav_msgs::Path edge;
    // sharedEdge = ros::topic::waitForMessage<nav_msgs::Path>("/path_planned/edge",create_path);
    // if(sharedEdge != NULL){
    // edge = *sharedEdge;
    // }
    boost::shared_ptr<dwm1001::anchor const> final_msg;
    ros::Duration one_second(0.5);
    final_msg = ros::topic::waitForMessage<dwm1001::anchor>(pose_topic,nh,one_second);
    if(final_msg != NULL){
    //p_msg  = *final_msg ;
    x = final_msg->x;
    y = final_msg->y;
    z = final_msg->z;
    }

    }
    //ROS_WARN("m_Q_scale %f", m_Q_scale);
    // ROS_WARN("x %f", x);

    nine_cov(0,0) = 0.001;
    nine_cov(3,3) = 0.001;
    nine_cov(6,6) = 0.001;
    // set cov of vel
    nine_cov(1,1) = 0.01;
    nine_cov(4,4) = 0.01;
    nine_cov(7,7) = 0.01;
    // set cov of acc_bias
    nine_cov(2,2) = 0.1;
    nine_cov(5,5) = 0.1;
    nine_cov(8,8) = 0.1;
    setCovariance(nine_cov);
    X<<x,0,0,y,0,0,z,0,0;
    m_last_range_time = ros::Time::now().toSec();
    H_h << 0, 0, 0,
           0, 0, 0,
           1, 0, 0;
    H_vel << 0, 1, 0, 0, 0, 0, 0, 0, 0,
             0, 0, 0, 0, 1, 0, 0, 0, 0,
             0, 0, 0, 0, 0, 0, 0, 1, 0;

}

void Initialize_2()
{

    std::string pose_topic("initialize");
    boost::shared_ptr<dwm1001::anchor const> final_msg;
    ros::Duration one_second(0.5);
    final_msg = ros::topic::waitForMessage<dwm1001::anchor>(pose_topic,one_second);
    if(final_msg != NULL){
    //p_msg  = *final_msg ;
    x = final_msg->x;
    y = final_msg->y;
    z = final_msg->z;
    }
    ROS_WARN("initialized");
    //ROS_WARN("m_Q_scale %f", m_Q_scale);
    // ROS_WARN("x %f", x);
/*
    Eigen::MatrixXd nine_cov = Eigen::MatrixXd::Identity(9,9);
    nine_cov(0,0) = 0.001;
    nine_cov(3,3) = 0.001;
    nine_cov(6,6) = 0.001;
    // set cov of vel
    nine_cov(1,1) = 0.01;
    nine_cov(4,4) = 0.01;
    nine_cov(7,7) = 0.01;
    // set cov of acc_bias
    nine_cov(2,2) = 0.1;
    nine_cov(5,5) = 0.1;
    nine_cov(8,8) = 0.1;*/
    setCovariance(nine_cov);
    X<<x,0,0,y,0,0,z,0,0;
    m_last_range_time = ros::Time::now().toSec();

}

void correction_step(const dwm1001::anchor::ConstPtr& msg)
{
    r_pred = std::sqrt( std::pow(X(0) - msg->x, 2) +
                        std::pow(X(3) - msg->y, 2) +
                        std::pow(X(6) - msg->z, 2) ) + 1e-5;

    r_meas = msg->range;
    // H is the linearized measurement matrix
    H << (X(0) - msg->x)/r_pred, 0, 0,
         (X(3) - msg->y)/r_pred, 0, 0,
         (X(6) - msg->z)/r_pred, 0, 0;

    // K is the Kalman Gain
    sigma_r = double(precisionRangeErrEst) / 1000.0;
    R = std::pow(sigma_r,2) * m_R_scale;
    //ROS_WARN("runing");
    K= P*H.transpose() / ( (H*P*H.transpose())(0,0) + R );
    // Update P for the a posteriori covariance matrix
    //std::cout<<K(6)<<'\n';

    P_p = ( Eigen::MatrixXd::Identity(9,9) - K*H ) * P;
    // Return the measurement innovation
    error_1 = std::fabs(r_meas - r_pred);
    // Update the state
    X_p = X + K * (r_meas - r_pred);
    // decide to take the range info or not.
    if(error_1 < error_threshold){
        //ROS_WARN("\n sucess too large: %f", error);
        setState(X_p);
        setCovariance(P_p);
        //set_= 0;
        return ;
    } else {

        ROS_WARN("Anchor id , Update too large: predicted %f measured-- %f", r_pred, r_meas);
        std::cout << msg->device_id << "\n";
        set_=set_+1;
        if(set_>5)
            {Initialize_2();
              set_=0;
            }
        return ;
    }

}

void convert_NED(sensor_msgs::Imu imu)
{
            int static count_ =0 ;
    if(count_ == 0)
       {ROS_WARN("<--------imu_step---------->");
        count_++; }
    //conerts the acceleration from body frame to earth frame(NED)
    q = Eigen::Quaternionf(imu.orientation.w, imu.orientation.x, imu.orientation.y, imu.orientation.z);
    R_mat= q.toRotationMatrix();
    //rpy = q.toRotationMatrix().eulerAngles(0, 1, 2);
    //ROS_WARN("%f",rpy(2,0)*180/3.14);
    //yaw = rpy(2,0) - (yaw_off*0.017454);

    ax=imu.linear_acceleration.x;
    ay=imu.linear_acceleration.y;
    az=imu.linear_acceleration.z;
    imuacc << ax,ay,az;
    acc= R_mat*imuacc ;
    imu_msg.pose.orientation.w = imu.orientation.w;
    imu_msg.pose.orientation.x = imu.orientation.x;
    imu_msg.pose.orientation.y = imu.orientation.y;
    imu_msg.pose.orientation.z = imu.orientation.z;

}

void prediction_step(const sensor_msgs::Imu::ConstPtr& msg)
{
    convert_NED(*msg);
    ax=acc(0,0);
    ay=acc(1,0);
    az=acc(2,0)-9.8;

    // yaw_rad = 0.0174533*yaw_off;
    // ax_= ax*std::cos(yaw_rad) + ay*std::sin(yaw_rad);
    // ay_=-ax*std::sin(yaw_rad) + ay*std::cos(yaw_rad);
    //ROS_WARN("Acceleration:## %f, ## %f, ## %f", ax,ay,az);
    T = msg->header.stamp.toSec() - m_last_range_time;
    if(T>1){
        T = 1;
    } else if(T<0){
        T = 0.02;
    }

    sigma_a = m_kalman_sigma_a;
   // r_meas = double(precisionRangeMm) / 1000.0;

    T_sq = std::pow(T,2);
    T_cub = std::pow(T,3);

    // F is a 9x9 State Transition Matrix
    F = Eigen::MatrixXd::Zero(9,9);

    block_F << 1, T, -T_sq/2.0,
               0, 1, -T,
               0, 0, 1;
    F.block<3,3>(0,0) = block_F;
    F.block<3,3>(3,3) = block_F;
    F.block<3,3>(6,6) = block_F;

    // Q is the acceleration model
    tao_acc = m_tao_acc_sqrt * m_tao_acc_sqrt;
    tao_bias = m_tao_bias_sqrt * m_tao_bias_sqrt;
    Q = Eigen::MatrixXd::Zero(9,9);

    block_Q << (T_cub*tao_acc/3.0)+(T_cub*T_sq)*tao_bias/20.0, (T_sq*tao_acc/2)+(T_sq*T_sq)*tao_bias/8.0  ,-T_cub*tao_bias/6,
               (T_sq*tao_acc/2.0)+(T_sq*T_sq)*tao_bias/8.0 ,   T*tao_acc+(T_cub*tao_bias/3)               ,-T_sq*tao_bias/2,
               -T_cub*tao_bias/6.0,                         -T_sq*tao_bias/2                          ,T*tao_bias         ;
    Q.block<3,3>(0,0) = block_Q;
    Q.block<3,3>(3,3) = block_Q;
    Q.block<3,3>(6,6) = block_Q * m_z_damping_factor;
    Q *= m_Q_scale;
    //std::cout << block_Q*m_Q_scale << '\n'<<'\n';
    //ROS_WARN("covariance:## %f, ## %f, ## %f",(T_cub*tao_acc/3.0)+(T_cub*T_sq)*tao_bias/20.0 , T*tao_acc+(T_cub*tao_bias/3), T*tao_bias);

    u << ax, 0, 0,
         ay, 0, 0,
         az, 0, 0;
    B = Eigen::MatrixXd::Zero(9,9);

    block_B << T_sq/2.0,  0,  0,
               T      ,  0,  0,
               0      ,  0,  0;
    B.block<3,3>(0,0) = block_B;
    B.block<3,3>(3,3) = block_B;
    B.block<3,3>(6,6) = block_B;

    // X is the predicted state vector and the predicted covariance matrix
    X_e = F*X + B * u;
    //if(count%250 == 0)
    //std::cout << P << '\n'<<'\n';
    // M is the predicted covariance matrix
    M = F*P*F.transpose() + Q;

    // time update
    m_last_range_time = msg->header.stamp.toSec();
    //ROS_WARN("runing---12");
    error = X.squaredNorm()-X_e.squaredNorm();
    //ROS_WARN("predicted %f#%f$%f",X_e(0),X_e(3),X_e(6));
    if(isnan(X(0))==1)
        {Initialize_2();
            return;}
    if(error < error_threshold){
        //ROS_WARN("\n sucess too large: %f", error);
        set_p=0;
        setState(X_e);
        setCovariance(M);
        return ;
    } else {

        ROS_WARN("\n Estimate too large: %f", error);
       set_p=set_p+1;
        if(set_p>5)
            {Initialize_2();
             set_p=0;
            }
        return ;
    }

}

void anchor1_cb(const dwm1001::anchor::ConstPtr& msg)
{
    correction_step(msg);
    //ROS_WARN("runing");
}
void anchor2_cb(const dwm1001::anchor::ConstPtr& msg)
{
    correction_step(msg);
}
void anchor3_cb(const dwm1001::anchor::ConstPtr& msg)
{
    correction_step(msg);
}
void anchor4_cb(const dwm1001::anchor::ConstPtr& msg)
{
    correction_step(msg);
}
// void anchor5_cb(const dwm1001::anchor::ConstPtr& msg)
// {
//  correction_step(msg);
// }
// void anchor6_cb(const dwm1001::anchor::ConstPtr& msg)
// {
//  correction_step(msg);
// }
//x,y,z,error_threshold,sigma_r,m_R_scale,precisionRangeErrEst,m_kalman_sigma_a,precisionRangeMm,
//m_tao_acc_sqrt,m_tao_bias_sqrt,m_z_damping_factor,m_Q_scale
void correction_step_vel(const geometry_msgs::TwistStamped::ConstPtr& msg)
{
        int static count_ =0 ;
    if(count_ == 0)
       {ROS_WARN("<--------velocity_step---------->");
        count_++; }

    r_pred_vel << X(1),X(4),X(7);

    r_meas_vel << msg->twist.linear.x,msg->twist.linear.y,msg->twist.linear.z;
    // K is the Kalman Gain
    R_vel = Eigen::MatrixXd::Identity(3,3) * (r_vel*r_vel);
    //ROS_WARN("R_matrix:## %f",R);
    R_vel = (H_vel * P * H_vel.transpose()) + R_vel ;
    K_vel = P * H_vel.transpose() * R_vel.inverse();
    // Update P for the a posteriori covariance matrix
    P_p_vel = ( Eigen::MatrixXd::Identity(9,9) - K_vel * H_vel ) * P;
    // Update the state
    X_p_vel = X + K_vel * (r_meas_vel - r_pred_vel);
    // decide to take the range info or not.
        //ROS_WARN("\n sucess velocity: %f", );

        setState(X_p_vel);
        setCovariance(P_p_vel);
        return ;

}

void height_cb(const sensor_msgs::Range::ConstPtr& msg)
{
    int static count_ =0 ;
    if(count_ == 0)
       {ROS_WARN("<--------height_step---------->");
        count_++; }

    h_meas = msg->range;
    // K is the Kalman Gain
    r_h = r_h*r_h;
    //ROS_WARN("R_matrix:## %f",R);

   K_h = P*H_h.transpose() / ( (H_h*P*H_h.transpose())(0,0) + r_h );

    // K is the Kalman Gain
    // Update P for the a posteriori covariance matrix

    P_h = ( Eigen::MatrixXd::Identity(9,9) - K_h*H_h ) * P;
    // Return the measurement innovation
    // Update the state
    X_h = X + K_h * (h_meas - X(6));
    //std::cout<<K_h<<"--"<<h_meas<<"---"<<X_h(6)<<'\n';
        //ROS_WARN("\n height: %f", msg->local);
        setState(X_h);
        setCovariance(P_h);

}

int main(int argc, char** argv){

    ros::init(argc,argv,"KalmanFilter");
    ros::NodeHandle nh;
    //Initializinging the parameters
    Initialize(nh);

    //Subscriber and Publisher for the data. remapped in the launch file to the topic required
    ros::Subscriber Imu=nh.subscribe("imu",100,prediction_step);
    ros::Subscriber anchor_1 = nh.subscribe("anchor_1",10,anchor1_cb);
    ros::Subscriber anchor_2 = nh.subscribe("anchor_2",10,anchor2_cb);
    ros::Subscriber anchor_3 = nh.subscribe("anchor_3",10,anchor3_cb);
    ros::Subscriber anchor_4 = nh.subscribe("anchor_4",10,anchor4_cb);
    // ros::Subscriber anchor_5 = nh.subscribe("anchor_5",10,anchor5_cb);
    // ros::Subscriber anchor_6 = nh.subscribe("anchor_6",10,anchor6_cb);
    ros::Subscriber height_sb = nh.subscribe("height",10,height_cb);
    ros::Subscriber velocity_sb = nh.subscribe("velocity",10,correction_step_vel);

    ros::Publisher fused = nh.advertise<std_msgs::Float32MultiArray>("Filtered_data", 100);
    ros::Publisher fused_pose = nh.advertise<geometry_msgs::PoseStamped>("Filtered_pose", 10);

    ros::Rate loop_rate(50);
    while(ros::ok()){
        //param(nh);
  //    output.data.clear();
        // for (int i = 0; i < 9; i++)
        //  output.data.push_back(X(i));

        pose.header.stamp = ros::Time::now();
        pose.header.frame_id = "base";

        pose.pose.position.x=X(0)+trans_x;
        pose.pose.position.y=X(3)+trans_y;
        pose.pose.position.z=X(6);
        pose.pose.orientation = imu_msg.pose.orientation;

        //q_ = Eigen::AngleAxisf(rpy(0,0), Eigen::Vector3f::UnitX())*Eigen::AngleAxisf(rpy(1,0), Eigen::Vector3f::UnitY())*Eigen::AngleAxisf(yaw, Eigen::Vector3f::UnitZ());

        // pose.pose.orientation.w = q_.w();
        // pose.pose.orientation.x = q_.x();
        // pose.pose.orientation.y = q_.y();
        // pose.pose.orientation.z = q_.z();
        //fused.publish(output);
        fused_pose.publish(pose);
        ros::spinOnce();
        loop_rate.sleep();
    }
    return 0;
}


/*  _______________________________________________________________________

    PECOS: Parallel Environment for Creation Of Stochastics
    Copyright (c) 2011, Sandia National Laboratories.
    This software is distributed under the GNU Lesser General Public License.
    For more information, see the README file in the top Pecos directory.
    _______________________________________________________________________ */

//- Class:	 RandomVariable
//- Description:
//- Owner:       Mike Eldred
//- Revised by:
//- Version:

#ifndef RANDOM_VARIABLE_HPP
#define RANDOM_VARIABLE_HPP

#include "pecos_data_types.hpp"
#include <boost/math/distributions.hpp>
#include <boost/math/special_functions/sqrt1pm1.hpp> // includes expm1,log1p

namespace bmth = boost::math;
namespace bmp  = bmth::policies;

namespace Pecos {

// -----------------------------------
// Non-default boost math/policy types
// -----------------------------------

// continuous random variable types:
typedef bmth::
  normal_distribution< Real,
                       bmp::policy< bmp::overflow_error<bmp::ignore_error> > >
  normal_dist;
typedef bmth::
  lognormal_distribution< Real,
                       bmp::policy< bmp::overflow_error<bmp::ignore_error> > >
  lognormal_dist;
typedef bmth::
  triangular_distribution< Real,
                       bmp::policy< bmp::overflow_error<bmp::ignore_error> > >
  triangular_dist;
typedef bmth::
  exponential_distribution< Real,
                       bmp::policy< bmp::overflow_error<bmp::ignore_error> > >
  exponential_dist;
typedef bmth::
  beta_distribution< Real,
                       bmp::policy< bmp::overflow_error<bmp::ignore_error> > >
  beta_dist;
typedef bmth::
  gamma_distribution< Real,
                       bmp::policy< bmp::overflow_error<bmp::ignore_error> > >
  gamma_dist;
typedef bmth::
  inverse_gamma_distribution< Real,
                       bmp::policy< bmp::overflow_error<bmp::ignore_error> > >
  inv_gamma_dist;
typedef bmth::
  extreme_value_distribution< Real,
                       bmp::policy< bmp::overflow_error<bmp::ignore_error> > >
  extreme_value_dist;
typedef bmth::
  weibull_distribution< Real,
                       bmp::policy< bmp::overflow_error<bmp::ignore_error> > >
  weibull_dist;
// discrete random variable types:
typedef bmth::
  poisson_distribution< Real,
                       bmp::policy< bmp::overflow_error<bmp::ignore_error> > >
  poisson_dist;
typedef bmth::
  binomial_distribution< Real,
                       bmp::policy< bmp::overflow_error<bmp::ignore_error> > >
  binomial_dist;
typedef bmth::
  negative_binomial_distribution< Real,
                       bmp::policy< bmp::overflow_error<bmp::ignore_error> > >
  negative_binomial_dist;
typedef bmth::
  geometric_distribution< Real,
                       bmp::policy< bmp::overflow_error<bmp::ignore_error> > >
  geometric_dist;
typedef bmth::
  hypergeometric_distribution< Real,
                       bmp::policy< bmp::overflow_error<bmp::ignore_error> > >
  hypergeometric_dist;
// distributions used in statistical utilities (e.g., confidence intervals):
typedef bmth::
  chi_squared_distribution< Real,
                       bmp::policy< bmp::overflow_error<bmp::ignore_error> > >
  chi_squared_dist;
typedef bmth::
  students_t_distribution< Real,
                       bmp::policy< bmp::overflow_error<bmp::ignore_error> > >
  students_t_dist;
typedef bmth::
  fisher_f_distribution< Real,
                       bmp::policy< bmp::overflow_error<bmp::ignore_error> > >
  fisher_f_dist;

/// base class for random variable hierarchy

/** This class enables cdf(), ccdf(), inverse_cdf(), inverse_ccdf(),
    pdf(), pdf_gradient(), pdf_hessian(), and related random variable
    utilities from contained distribution parameters. */

// ExtremeValueBase: Gumbel, Frechet, Weibull -> alpha, beta, no bounds
// EmpiricalBase: histogram bin, KDE, ...

class RandomVariable
{
public:

  //
  //- Heading: Constructors, destructor, and operator=
  //

  /// default constructor
  RandomVariable();
  /// standard constructor for envelope
  RandomVariable(short ran_var_type);
  /// copy constructor
  RandomVariable(const RandomVariable& ran_var);

  /// destructor
  virtual ~RandomVariable();

  /// assignment operator
  RandomVariable operator=(const RandomVariable& ran_var);

  //
  //- Heading: Virtual functions
  //

  /// return the cumulative distribution function value of the random
  /// variable at x
  virtual Real cdf(Real x) const;
  /// return the x value corresponding to a cumulative probability
  virtual Real inverse_cdf(Real p_cdf) const;

  /// return the complementary cumulative distribution function value
  /// of the random variable at x
  virtual Real ccdf(Real x) const;
  /// return the x value corresponding to a complementary cumulative probability
  virtual Real inverse_ccdf(Real p_ccdf) const;

  /// return the value of the random variable's probability density
  /// function at x
  virtual Real pdf(Real x) const;
  /// return the gradient of the random variable's probability density
  /// function at x
  virtual Real pdf_gradient(Real x) const;
  /// return the hessian of the random variable's probability density
  /// function at x
  virtual Real pdf_hessian(Real x) const;
  /// return the value of the natural log of the random variable's probability
  /// density function at x (useful for calculations of log density in Bayesian
  /// methods)
  virtual Real log_pdf(Real x) const;
  /// return the gradient of the natural log of the random variable's
  /// probability density function at x (useful for defining MCMC proposal
  /// distributions in Bayesian methods)
  virtual Real log_pdf_gradient(Real x) const;
  /// return the Hessian of the natural log of the random variable's probability
  /// density function at x (useful for defining MCMC proposal distributions in
  /// Bayesian methods)
  virtual Real log_pdf_hessian(Real x) const;

  /// return the value of a standardized random variable's probability density
  /// function at x
  virtual Real standard_pdf(Real z) const;
  /// return the natural log of a standardized random variable's probability
  /// density function at x (useful for calculations of log density in
  /// Bayesian methods)
  virtual Real log_standard_pdf(Real z) const;
  /// return the gradient of the natural log of a standardized random
  /// variable's probability density function at x (useful for
  /// calculations of log density in Bayesian methods)
  virtual Real log_standard_pdf_gradient(Real z) const;
  /// return the Hessian of the natural log of a standardized random
  /// variable's probability density function at x (useful for
  /// calculations of log density in Bayesian methods)
  virtual Real log_standard_pdf_hessian(Real z) const;

  /// scale variable value x from current to standardized distribution
  virtual Real to_standard(Real x) const;
  /// scale variable value z from standardized to current distribution
  virtual Real from_standard(Real z) const;

  /// return the value of the named distribution parameter
  virtual Real parameter(short dist_param) const;
  /// update the value of the named distribution parameter
  virtual void parameter(short dist_param, Real val);

  /// return the distribution mean
  virtual Real mean() const;
  /// return the distribution mode
  virtual Real median() const;
  /// return the distribution mode
  virtual Real mode() const;
  /// return the distribution variance
  virtual Real standard_deviation() const;
  /// return the distribution variance
  virtual Real variance() const;

  /// return the distribution mean and standard deviation as a pair
  /** default is only overridden when more efficient to compute together */
  virtual RealRealPair moments() const;
  /// return the distribution lower and upper bounds as a pair
  virtual RealRealPair bounds() const;

  /// compute the coefficient of variation (used to compute selected
  /// correlation warping factors); defined for semi-infinite distributions
  /// with nonzero mean (lognormal, exponential, gamma, frechet, weibull)
  /** default is only overridden when more efficient to compute together */
  virtual Real coefficient_of_variation() const;
  /// compute the warping factor for correlation between the current
  /// variable and the one passed in (used in NatafTransformation)
  virtual Real correlation_warping_factor(const RandomVariable& rv,
					  Real corr) const;

  /// compute the design Jacobian from differentiating the X->Z mapping with
  /// respect to the distibution parameter s
  virtual Real dx_ds(short dist_param, short u_type, Real x, Real z) const;
  /// compute the mapping-specific factor that is multiplied by dz/ds for
  /// contributions to the dx/ds design Jacobian in the case of correlated
  /// random variables (dz/ds is evaluated numerically and multiplied by
  /// this analytic factor)
  virtual Real dz_ds_factor(short u_type, Real x, Real z) const;

  //
  //- Heading: Member functions
  //

  /// set ranVarType
  void type(short ran_var_type);
  /// get ranVarType
  short type() const;

  /// returns ranVarRep for access to derived class member functions
  /// that are not mapped to the base level
  RandomVariable* random_variable_rep() const;

protected:

  //
  //- Heading: Constructors
  //

  /// constructor initializes the base class part of letter classes
  /// (BaseConstructor overloading avoids infinite recursion in the
  /// derived class constructors - Coplien, p. 139)
  RandomVariable(BaseConstructor);

  //
  //- Heading: Member functions
  //

  //
  //- Heading: Data
  //

  /// enumeration value indicating type of random variable
  short ranVarType;

private:

  //
  //- Heading: Member functions
  //

  /// Used only by the standard envelope constructor to initialize
  /// ranVarRep to the appropriate derived type.
  RandomVariable* get_random_variable(short ran_var_type);

  //
  //- Heading: Data members
  //

  /// pointer to the letter (initialized only for the envelope)
  RandomVariable* ranVarRep;
  /// number of objects sharing ranVarRep
  int referenceCount;
};

inline void RandomVariable::type(short ran_var_type)
{
  if (ranVarRep) ranVarRep->ranVarType = ran_var_type;
  else           ranVarType = ran_var_type;
}

inline short RandomVariable::type() const
{ return (ranVarRep) ? ranVarRep->ranVarType : ranVarType; }

inline RandomVariable* RandomVariable::random_variable_rep() const
{ return ranVarRep; }

} // namespace Pecos

#endif