/**
 * \file      ekf-flexibility-estimator-base.hpp
 * \author    Mehdi Benallegue
 * \date      2013
 * \brief     Declare the class of the flexibility estimation using the extended
 *            Kalman Filter.
 *
 * \details
 *
 *
 */

#ifndef FLEXIBILITYESTIMATION_EKFFLEXIBILITYESTIMATORBASE_H
#define FLEXIBILITYESTIMATION_EKFFLEXIBILITYESTIMATORBASE_H

#include <boost/utility.hpp>

#include <state-observation/api.h>
#include <state-observation/observer/extended-kalman-filter.hpp>

#include <state-observation/flexibility-estimation/flexibility-estimator-base.hpp>

namespace stateObservation
{
namespace flexibilityEstimation
{
/**
 * \class  EKFFlexibilityEstimatorBase
 * \brief  This class is the base class of the flexibility estimators that
 *         use an extended Kalman Filter. Several methods require to be overloaded
 *         to derive an implementation from this base class.
 *
 */

class STATE_OBSERVATION_DLLAPI EKFFlexibilityEstimatorBase : public FlexibilityEstimatorBase
{
public:
  /// The constructor.
  ///  \li stateSize : size of the state vector
  ///  \li measurementSize : size of the measurements vector
  ///  \li inputSize : size of the input vector
  ///  \li dx gives the derivation step for a finite differences derivation method

  EKFFlexibilityEstimatorBase(Index stateSize,
                              Index measurementSize,
                              Index inputSize,
                              const Vector & dx = Vector::Zero(0));

  /// virtual destructor
  virtual ~EKFFlexibilityEstimatorBase();

  /// Sets a value of the flexibility x_k provided from another source
  /// can be used for initialization of the estimator
  /// This is a pure virtual function that requires to be overloaded in
  /// implementation
  virtual void setFlexibilityGuess(const Matrix &) = 0;

  /// Sets the covariance matrix of the flexibility Guess
  virtual void setFlexibilityCovariance(const Matrix & P);

  /// Gets the covariance matrix of the flexibility
  virtual Matrix getFlexibilityCovariance() const;

  /// Sets the covariance matrices for the process noises
  /// \li Q process noise
  virtual void setProcessNoiseCovariance(const Matrix & Q);

  /// Sets the covariance matrices for the sensor noises
  /// \li R sensor noise
  virtual void setMeasurementNoiseCovariance(const Matrix & R);

  /// gets the covariance matrices for the process noises
  virtual Matrix getProcessNoiseCovariance() const;

  /// gets the covariance matrices for the sensor noises
  virtual Matrix getMeasurementNoiseCovariance() const;

  /// Sets the value of the next sensor measurement y_{k+1}
  virtual void setMeasurement(const Vector & y);

  virtual Vector getMeasurement();

  /// Sets the value of the next input for the state process dynamics
  /// i.e. : gives u_k such that x_{k+1} = f(x_k,u_k)
  virtual void setInput(const Vector & u);

  /// Sets the value of the next  measurement
  /// i.e. : gives u_{k+1} such that y_{k+1}=h(x_{k+1},u_{k+1})
  virtual void setMeasurementInput(const Vector & u);

  virtual Vector getInput();

  virtual Vector getMeasurementInput();

  /// Gets an estimation of the flexibility in the form of a state vector \hat{x_{k+1}}
  virtual const Vector & getFlexibilityVector();

  /// Gets an estimation of the flexibility in the form of a homogeneous matrix
  virtual Matrix4 getFlexibility() = 0;

  /// Gets a const reference on the extended Kalman filter
  virtual const stateObservation::ExtendedKalmanFilter & getEKF() const;

  /// Gets a reference on the extended Kalman filter
  virtual stateObservation::ExtendedKalmanFilter & getEKF();

  /// Gets the state size
  /// this method is pure virtual and reauires to be overloaded in implementation
  virtual Index getStateSize() const = 0;

  /// Gets the measurements size
  /// this method is pure virtual and reauires to be overloaded in implementation
  virtual Index getMeasurementSize() const = 0;

  /// Gets the input size
  /// this method is pure virtual and reauires to be overloaded in implementation
  virtual Index getInputSize() const = 0;

  /// Gets a simulation of the
  virtual Vector getSimulatedMeasurement();

  /// Resets the covariance matrices to their original values
  virtual void resetCovarianceMatrices() = 0;

  /// Get the last vector of inovation of the Kalman filter
  virtual Vector getInnovation();

  /// Get the simulated measurement of the predicted state
  virtual Vector getPredictedMeasurement();

  /// Get the predicted state
  virtual Vector getPrediction();

  /// Get the last simulated measurement
  virtual Vector getLastPredictedMeasurement();

  /// Get the last predicted state
  virtual Vector getLastPrediction();

protected:
  virtual void setJacobians(const Matrix & A, const Matrix & C);

  virtual void useFiniteDifferencesJacobians(Vector dx);

  stateObservation::ExtendedKalmanFilter ekf_;

  bool finiteDifferencesJacobians_;

  Vector dx_;

  Vector lastX_;

  TimeIndex k_;

private:
};
} // namespace flexibilityEstimation
} // namespace stateObservation
#endif // FLEXIBILITYESTIMATION_EKFFLEXIBILITYESTIMATORBASE_H


/**
 * @file   NumpyEigenConverter.hpp
 * @author Paul Furgale <paul.furgale@utoronto.ca>
 * @date   Fri Feb  4 11:17:25 2011
 *
 * @brief  Classes to support conversion from numpy arrays in Python
 *         to Eigen3 matrices in c++
 *
 *
 */

#ifndef NUMPY_EIGEN_CONVERTER_HPP
#define NUMPY_EIGEN_CONVERTER_HPP

#include <numpy_eigen/boost_python_headers.hpp>
//#include <iostream>

#define PY_ARRAY_UNIQUE_SYMBOL NP_Eigen_AS
// #include <numpy/npy_no_deprecated_api.h>
#include <numpy/arrayobject.h>

#include "type_traits.hpp"
#include <boost/lexical_cast.hpp>
#include "copy_routines.hpp"

/**
 * @class NumpyEigenConverter
 * @tparam the Eigen3 matrix type this class is specialized for
 *
 * adapted from http://misspent.wordpress.com/2009/09/27/how-to-write-boost-python-converters/
 * General help available http://docs.scipy.org/doc/numpy/reference/c-api.array.html
 *
 * To use:
 *
 * #include <NumpyEigenConverter.hpp>
 *
 *
 * BOOST_PYTHON_MODULE(libmy_module_python)
 * {
 *   // The converters will cause a segfault unless import_array() is called before the first one
 *   import_array();
 *   NumpyEigenConverter<Eigen::Matrix< double, 1, 1 > >::register_converter();
 *   NumpyEigenConverter<Eigen::Matrix< double, 2, 1 > >::register_converter();
 * }
 *
 */
template<typename EIGEN_MATRIX_T>
struct NumpyEigenConverter
{

  typedef EIGEN_MATRIX_T matrix_t;
  typedef typename matrix_t::Scalar scalar_t;

  enum {
    RowsAtCompileTime = matrix_t::RowsAtCompileTime,
    ColsAtCompileTime = matrix_t::ColsAtCompileTime,
    MaxRowsAtCompileTime = matrix_t::MaxRowsAtCompileTime,
    MaxColsAtCompileTime = matrix_t::MaxColsAtCompileTime,
    NpyType = TypeToNumPy<scalar_t>::NpyType,
    //Flags = ei_compute_matrix_flags<_Scalar, _Rows, _Cols, _Options, _MaxRows, _MaxCols>::ret,
    //CoeffReadCost = NumTraits<Scalar>::ReadCost,
    Options = matrix_t::Options
    //InnerStrideAtCompileTime = 1,
    //OuterStrideAtCompileTime = (Options&RowMajor) ? ColsAtCompileTime : RowsAtCompileTime
  };

  static std::string castSizeOption(int option)
  {
    if(option == Eigen::Dynamic)
      return "Dynamic";
    else
      return boost::lexical_cast<std::string>(option);
  }

  static std::string toString()
  {
    return std::string() + "Eigen::Matrix<" + TypeToNumPy<scalar_t>::typeString() + ", " +
      castSizeOption(RowsAtCompileTime) + ", " +
      castSizeOption(ColsAtCompileTime) + ", " +
      boost::lexical_cast<std::string>((int)Options) + ", " +
      castSizeOption(MaxRowsAtCompileTime) + ", " +
      castSizeOption(MaxColsAtCompileTime) + ">";
  }

  // The "Convert from C to Python" API
  static PyObject * convert(const matrix_t & M)
  {
    PyObject * P = NULL;
    if(RowsAtCompileTime == 1 || ColsAtCompileTime == 1)
      {
	// Create a 1D array
	npy_intp dimensions[1];
	dimensions[0] = M.size();
	P = PyArray_SimpleNew(1, dimensions, TypeToNumPy<scalar_t>::NpyType);
	numpyTypeDemuxer< CopyEigenToNumpyVector<const matrix_t> >(&M,P);
      }
    else
      {
	// create a 2D array.
	npy_intp dimensions[2];
	dimensions[0] = M.rows();
	dimensions[1] = M.cols();
	P = PyArray_SimpleNew(2, dimensions, TypeToNumPy<scalar_t>::NpyType);
	numpyTypeDemuxer< CopyEigenToNumpyMatrix<const matrix_t> >(&M,P);
      }

    // incrementing the reference seems to cause a memory leak.
    // boost::python::incref(P);
    // This agrees with the sample code found here:
    // http://mail.python.org/pipermail/cplusplus-sig/2008-October/013825.html
    return P;
  }

  static bool isDimensionValid(int requestedSize, int sizeAtCompileTime, int maxSizeAtCompileTime)
  {
    bool valid = true;
    if(sizeAtCompileTime == Eigen::Dynamic)
      {
	// Check for dynamic fixed size
	// http://eigen.tuxfamily.org/dox-devel/TutorialMatrixClass.html#TutorialMatrixOptTemplParams
	if(!(maxSizeAtCompileTime == Eigen::Dynamic || requestedSize <= maxSizeAtCompileTime))
	  {
	    valid = false;
	  }
      }
    else if(sizeAtCompileTime != requestedSize)
      {
	valid = false;
      }
    return valid;
  }

  static void checkMatrixSizes(PyObject * obj_ptr)
  {
    int rows = PyArray_DIM(obj_ptr, 0);
    int cols = PyArray_DIM(obj_ptr, 1);

    bool rowsValid = isDimensionValid(rows, RowsAtCompileTime, MaxRowsAtCompileTime);
    bool colsValid = isDimensionValid(cols, ColsAtCompileTime, MaxColsAtCompileTime);
    if(!rowsValid || !colsValid)
    {
	THROW_TYPE_ERROR("Can not convert " << npyArrayTypeString(obj_ptr) << " to " << toString()
			 << ". Mismatched sizes.");
      }
  }

  static void checkRowVectorSizes(PyObject * obj_ptr, int cols)
  {
    if(!isDimensionValid(cols, ColsAtCompileTime, MaxColsAtCompileTime))
      {
	THROW_TYPE_ERROR("Can not convert " << npyArrayTypeString(obj_ptr) << " to " << toString()
			 << ". Mismatched sizes.");
      }
  }

  static void checkColumnVectorSizes(PyObject * obj_ptr, int rows)
  {
    // Check if the type can accomidate one column.
    if(ColsAtCompileTime == Eigen::Dynamic || ColsAtCompileTime == 1)
      {
	if(!isDimensionValid(rows, RowsAtCompileTime, MaxRowsAtCompileTime))
	  {
	    THROW_TYPE_ERROR("Can not convert " << npyArrayTypeString(obj_ptr) << " to " << toString()
			     << ". Mismatched sizes.");
	  }
      }
    else
      {
	THROW_TYPE_ERROR("Can not convert " << npyArrayTypeString(obj_ptr) << " to " << toString()
			 << ". Mismatched sizes.");
      }

  }

  static void checkVectorSizes(PyObject * obj_ptr)
  {
	int size = PyArray_DIM(obj_ptr, 0);

    // If the number of rows is fixed at 1, assume that is the sense of the vector.
    // Otherwise, assume it is a column.
    if(RowsAtCompileTime == 1)
      {
	checkRowVectorSizes(obj_ptr, size);
      }
    else
      {
	checkColumnVectorSizes(obj_ptr, size);
      }
  }

  static void* convertible(PyObject *obj_ptr)
  {
    // Check for a null pointer.
    if(!obj_ptr)
      {
        //THROW_TYPE_ERROR("PyObject pointer was null");
        return 0;
      }

    // Make sure this is a numpy array.
    if (!PyArray_Check(obj_ptr))
      {
        //THROW_TYPE_ERROR("Conversion is only defined for numpy array and matrix types");
        return 0;
      }

    // Check the type of the array.
    int npyType = PyArray_ObjectType(obj_ptr, 0);

    if(!TypeToNumPy<scalar_t>::canConvert(npyType))
      {
        //THROW_TYPE_ERROR("Can not convert " << npyArrayTypeString(obj_ptr) << " to " << toString()
        //                 << ". Mismatched types.");
        return 0;
      }

    // Check the array dimensions.
    int nd = PyArray_NDIM(obj_ptr);

    if(nd != 1 && nd != 2)
      {
	THROW_TYPE_ERROR("Conversion is only valid for arrays with 1 or 2 dimensions. Argument has " << nd << " dimensions");
      }

    if(nd == 1)
      {
	checkVectorSizes(obj_ptr);
      }
    else
      {
	// Two-dimensional matrix type.
	checkMatrixSizes(obj_ptr);
      }

    return obj_ptr;
  }

  static void construct(PyObject *obj_ptr, boost::python::converter::rvalue_from_python_stage1_data *data)
  {
    boost::python::converter::rvalue_from_python_storage<matrix_t> * matData = reinterpret_cast<boost::python::converter::rvalue_from_python_storage<matrix_t> * >(data);
    void* storage = matData->storage.bytes;

    // Make sure storage is 16byte aligned. With help from code from Memory.h
    void * aligned = reinterpret_cast<void*>((reinterpret_cast<size_t>(storage) & ~(size_t(15))) + 16);

    matrix_t * Mp = new (aligned) matrix_t();
    // Stash the memory chunk pointer for later use by boost.python
    // This signals boost::python that the new value must be deleted eventually
    data->convertible = storage;

    // std::cout << "Creating aligned pointer " << aligned << " from storage " << storage << std::endl;
    // std::cout << "matrix size: " << sizeof(matrix_t) << std::endl;
    // std::cout << "referent size: " << boost::python::detail::referent_size< matrix_t & >::value << std::endl;
    // std::cout << "sizeof(storage): " << sizeof(matData->storage) << std::endl;
    // std::cout << "sizeof(bytes): " << sizeof(matData->storage.bytes) << std::endl;

    matrix_t & M = *Mp;

    int nd = PyArray_NDIM(obj_ptr);
    if(nd == 1)
      {
	int size = PyArray_DIM(obj_ptr, 0);
	// This is a vector type
	if(RowsAtCompileTime == 1)
	  {
	    // Row Vector
	    M.resize(1,size);
	  }
	else
	  {
	    // Column Vector
	    M.resize(size,1);
	  }
	numpyTypeDemuxer< CopyNumpyToEigenVector<matrix_t> >(&M,obj_ptr);
      }
    else
      {
	int rows = PyArray_DIM(obj_ptr, 0);
	int cols = PyArray_DIM(obj_ptr, 1);

	M.resize(rows,cols);
	numpyTypeDemuxer< CopyNumpyToEigenMatrix<matrix_t> >(&M,obj_ptr);
      }

  }

  // The registration function.
  static void register_converter()
  {
    boost::python::to_python_converter<matrix_t,NumpyEigenConverter>();
    boost::python::converter::registry::push_back(
						  &NumpyEigenConverter::convertible,
						  &NumpyEigenConverter::construct,
						  boost::python::type_id<matrix_t>());

  }

};

#endif /* NUMPY_EIGEN_CONVERTER_HPP */


#include "linear_elastic_energy.hpp"
#include <Eigen/Dense>

using namespace std;
using namespace Eigen;

void LinearElasticEnergy::precompute(const TriMesh& mesh)
{
    A_.resize(mesh.f.rows());
    DmInverse_.resize(mesh.f.rows());

    for (int idx=0; idx<mesh.f.rows(); idx++) {
        int i0 = mesh.f(idx, 0);
        int i1 = mesh.f(idx, 1);
        int i2 = mesh.f(idx, 2);

        Mat2d Dm;
        Dm.col(0) = mesh.u.segment<2>(2*i1) - mesh.u.segment<2>(2*i0);
        Dm.col(1) = mesh.u.segment<2>(2*i2) - mesh.u.segment<2>(2*i0);

        DmInverse_[idx] = Dm.inverse();

        A_[idx] = 0.5 * (Vec3d() << Dm.col(0), 0.0).finished().cross(
                        (Vec3d() << Dm.col(1), 0.0).finished()).norm();
    }

    lambda_ = VecXd::Zero(3*mesh.f.rows());
}

void LinearElasticEnergy::getForceAndHessian(const TriMesh& mesh,
                                                const VecXd& x,
                                                VecXd& F,
                                                SparseMatrixd& dFdx,
                                                SparseMatrixd& dFdv) const {
    assert(F.size() >= x.size());

    for (int idx=0; idx<mesh.f.rows(); idx++) {
    }
}

void LinearElasticEnergy::getHessianPattern(const TriMesh& mesh, vector<SparseTripletd> &triplets) const {
    for (int idx=0; idx<mesh.f.rows(); idx++) {
        int idxs[3] = { mesh.f(idx,0), mesh.f(idx,1), mesh.f(idx,2) };
        for (size_t j=0; j<3; ++j)
            for (size_t k=0; k<3; ++k)
                for (size_t l=0; l<3; ++l)
                    for (size_t n=0; n<3; ++n)
                        triplets.push_back(SparseTripletd(3*idxs[j]+l,3*idxs[k]+n, 1.0));
    }
}

void LinearElasticEnergy::perVertexCount(const TriMesh& mesh, std::vector<int>& counts) const {
    for (int idx=0; idx<mesh.f.rows(); idx++) {
        for (int j=0; j<3; j++) {
            counts[mesh.f(idx,j)]++;
        }
    }
}

void LinearElasticEnergy::update(const TriMesh& mesh, const VecXd& x, double dt, VecXd& dx) {
    const double a = (1.0 / ksx_) / (dt * dt);

    for (int idx=0; idx<mesh.f.rows(); idx++) {
        const int idxs[3] = { mesh.f(idx,0), mesh.f(idx,1), mesh.f(idx,2) };

        Mat3x2d Dw;
        for (int i=0; i<2; i++) {
            Dw.col(i) = x.segment<3>(3*idxs[i+1]) - x.segment<3>(3*idxs[0]);// +
                            //dx.segment<3>(3*idxs[i+1]) - dx.segment<3>(3*idxs[0]); // Jacobi
        }

        // Deformation gradient
        Mat3x2d F = Dw * DmInverse_[idx];

        const double wu_mag = F.col(0).norm();
        const double wv_mag = F.col(1).norm();

        Vec3d wu = F.col(0) / wu_mag;
        Vec3d wv = F.col(1) / wv_mag;

        // C(x) = A ( || wu || - bu  || wv || - bv    wu . wv )
        Vec3d C = A_[idx] * Vec3d(wu_mag - 1.0, wv_mag - 1.0, F.col(0).dot(F.col(1)));

        // grad C is 3, 3x3 matrices, where each is w.r.t a different vertex
        // Each column is the derivative of the respective C
        Mat3d gradC[3];

        gradC[1].col(0) = A_[idx] * DmInverse_[idx](0,0) * wu;
        gradC[1].col(1) = A_[idx] * DmInverse_[idx](0,1) * wv;
        gradC[1].col(2) = A_[idx] * (DmInverse_[idx](0,0) * F.col(1) + DmInverse_[idx](0,1) * F.col(0));
        gradC[2].col(0) = A_[idx] * DmInverse_[idx](1,0) * wu;
        gradC[2].col(1) = A_[idx] * DmInverse_[idx](1,1) * wv;
        gradC[2].col(2) = A_[idx] * (DmInverse_[idx](1,0) * F.col(1) + DmInverse_[idx](1,1) * F.col(0));
        gradC[0] = -(gradC[1] + gradC[2]);

	double den1 = (gradC[0].col(0).squaredNorm() / mesh.m[3*idxs[0]] +
	               gradC[1].col(0).squaredNorm() / mesh.m[3*idxs[1]] +
	               gradC[2].col(0).squaredNorm() / mesh.m[3*idxs[2]] + a);
	double den2 = (gradC[0].col(1).squaredNorm() / mesh.m[3*idxs[0]] +
	               gradC[1].col(1).squaredNorm() / mesh.m[3*idxs[1]] +
	               gradC[2].col(1).squaredNorm() / mesh.m[3*idxs[2]] + a);
	double den3 = (gradC[0].col(2).squaredNorm() / mesh.m[3*idxs[0]] +
	               gradC[1].col(2).squaredNorm() / mesh.m[3*idxs[1]] +
	               gradC[2].col(2).squaredNorm() / mesh.m[3*idxs[2]] + a);

	Vec3d dl = (-C - lambda_.segment<3>(3*idx) * a).array() / Eigen::Array3d(den1, den2, den3);

	dx.segment<3>(3*idxs[0]) += gradC[0] * dl / mesh.m[3*idxs[0]];
	dx.segment<3>(3*idxs[1]) += gradC[1] * dl / mesh.m[3*idxs[1]];
	dx.segment<3>(3*idxs[2]) += gradC[2] * dl / mesh.m[3*idxs[2]];

	lambda_.segment<3>(3*idx) += dl;
    }
}


#include <UnitTest++/UnitTest++.h>
#include <stdexcept>
#include <iostream>
#include <cmath>
#include <limits>
#include <boost/filesystem.hpp>
#include <fstream>

#include "coela_utility/src/string_utils.h"
#include "coela_utility/src/histogram_container.h"

#include "coela_core/src/ccd_image.h"
#include "../psf_characterisation.h"

using namespace coela;
using namespace std;

SUITE(psf_characterisation)
{
    using namespace psf_characterisation;

    string test_suite_output_dir= string(UnitTestSuite::GetSuiteName()) + "_tests/";
    TEST(Run_Standard_Suite_Setup) {
        cout << "*** \""<< UnitTestSuite::GetSuiteName() <<"\" unit tests running ***" <<endl;
        boost::filesystem::create_directories(test_suite_output_dir);
    }

    struct cone_shaped_PSF_image {
        double cone_peak;
        PixelArray2d<double> cone_img;
        PixelPosition cone_centre;
        PixelIndex cone_centre_pixel;

        cone_shaped_PSF_image() {

            cone_img = PixelArray2d<double> (100, 100, 0);
            cone_peak = cone_img.range().x_dim()/2.0 - 5.0;
            cone_centre = PixelPosition(cone_img.range().x_dim()/2.0 +0.5 ,
                                        cone_img.range().y_dim()/2.0 +0.5);
            cone_centre_pixel = PixelPosition::pixel_containing_point(cone_centre);

            //Cone shape determined by:
            // y = x (x=radius)
            // for x < cone_peak
            //y = 0; thereafter

            for (PixelIterator i(cone_img.range()); i!=i.end; ++i) {
                cone_img(i) =
                    max(0.0,
                        cone_peak - coord_distance(PixelPosition::centre_of_pixel(i), cone_centre)
                       );
            }

        }

    };

    TEST_FIXTURE(cone_shaped_PSF_image, write_image_for_visual_check) {
        cone_img.write_to_file(test_suite_output_dir+"cone_img.fits");
    }

    TEST_FIXTURE(cone_shaped_PSF_image, sanity_check) {
        //Therefore, enclosed flux at x< cone peak
        // EF = 1/3*PI*x^3

        double analytical_sum_flux=M_PI*cone_peak*cone_peak*cone_peak / 3.0;

//        cout<<"Analytic sum: " <<analytical_sum_flux<<"; actual: "<< cone_img.sum()<<endl;

        CHECK_CLOSE(analytical_sum_flux, cone_img.sum(),  analytical_sum_flux/100.0);
    }

    TEST_FIXTURE(cone_shaped_PSF_image, test_get_radial_data_unmasked) {
        using namespace psf_characterisation::detail;
        std::vector<std::pair <double, double>  > radial_data =
            get_radial_data(cone_img, cone_centre, cone_img.range().x_dim());
        //Oversized range should be ok.

        CHECK(radial_data.empty()==false);

        sort(radial_data.begin(), radial_data.end(), double_pair_first_member_predicate);
        ofstream datfile(string(test_suite_output_dir+"raw_radial_data.txt").c_str());
        for (size_t i=0; i!=radial_data.size(); ++i) {
            datfile<< radial_data[i].first <<" "<<radial_data[i].second<<endl;

        }
        datfile.close();
    }

    TEST_FIXTURE(cone_shaped_PSF_image, test_get_radial_data_unmasked_and_bin) {
        using namespace psf_characterisation::detail;
        std::vector<std::pair <double, double>  > radial_data =
            get_radial_data(cone_img, cone_centre, cone_img.range().x_dim());
        //Oversized range should be ok.

        vector<psf_profile_point> proFileInfo =
            bin_average_radial_data(radial_data);

//        cout<<"Get " << proFileInfo.size()<<" bins"<<endl;
        CHECK(proFileInfo.size() > cone_img.range().x_dim()/2.0);
        //expect a few extra along the diagonal
    }

    TEST_FIXTURE(cone_shaped_PSF_image, test_FWHM_estimation) {
        double est_fwhm = estimate_fwhm_in_image_pix(cone_img,
                          cone_centre,
                          cone_peak,
                          cone_img.range().x_dim()/2.0
                                                    );

        //45 degree angle slope ->
        //half-width at base = max height.
        //Therefore; full-width at half-height = max height.

        double calculated_FWHM = cone_peak;
//        cout<<"FWHM est: " << est_fwhm<<endl;
        CHECK_CLOSE(calculated_FWHM, est_fwhm, cone_peak / 20);
    }

    //Doesn't quite work, because we end up trying to interpolate past known points... not worth fixing, throws a sensible exception.
//    TEST_FIXTURE(cone_shaped_PSF_image, test_full_enclosed_flux_radius_estimation){
//        double full_flux_est_radius =
//        estimate_radius_to_enclose_flux(cone_img,
//                cone_centre,
//                cone_img.sum(),
//                cone_img.range().x_dim() );
//
////        cout<<"Est radius 100pc flux: "<<full_flux_est_radius<<"; expected: "<< cone_peak<<endl;
//        CHECK_CLOSE(cone_peak, full_flux_est_radius, cone_peak*0.01);
//    }

    TEST_FIXTURE(cone_shaped_PSF_image, test_10p_enclosed_flux_radius_estimation) {
//        double flux_10p_est_radius =
//        estimate_radius_to_enclose_flux(cone_img,
//                cone_centre,
//                cone_img.sum()*0.1,
//                cone_img.range().x_dim() );

//        cout<<"Est radius 10pc flux: "<<full_flux_est_radius<<endl;
//        cout<<"; expected: "<< cone_peak<<endl;
//        CHECK_CLOSE(cone_peak, full_flux_est_radius, cone_peak*0.01);
    }

    TEST_FIXTURE(cone_shaped_PSF_image, test_FWEF_estimation) {
        double est_fwhef = estimate_FWHEF_in_image_pix(cone_img,
                           cone_centre,
                           cone_img.sum(),
                           cone_img.range().x_dim()
                                                      );

//       cout<<"est FWHEF: "<<est_fwhef<<endl;

        //EF = PI*X^2*A - (2/3 * PI * X^3 / 3)
        double x= est_fwhef /2.0;

        double analytic_ef  = M_PI*x*x*cone_peak - (2.0/3.0) *M_PI*x*x*x ;

//       cout<<"Est HWHEF "<<x<<endl;
//       cout<<"Resulting analytic flux proportion: " <<analytic_ef / cone_img.sum()<<endl;

        CHECK_CLOSE(0.5, analytic_ef/ cone_img.sum(), 0.02);
    }

    TEST_FIXTURE(cone_shaped_PSF_image, test_fully_encircled_flux_estimation) {

//        PixelArray2d<double>& cone_img_ref = cone_img;

        double  fully_encircled =
            estimate_encircled_flux_at_pixel_radius(
                cone_img,
                cone_centre,
                cone_img.range().x_dim()/2.0 - 3.5
            );

//        cout<<"Estimate total flux via encircled_flux_at_pixel_radius: "<< fully_encircled<<endl;
//        cout<<"Actual total flux "<< cone_img.sum()<<endl;
        CHECK_CLOSE(cone_img.sum(), fully_encircled, fully_encircled*0.01);
    }

    TEST_FIXTURE(cone_shaped_PSF_image, test_encircled_flux_estimation) {

        double ap_radius =   cone_img.range().x_dim()/5.0;

        double  EF = estimate_encircled_flux_at_pixel_radius(cone_img,
                     cone_centre,
                     ap_radius
                                                            );

        double x = ap_radius;
        double analytic_ef= M_PI*x*x*cone_peak - (2.0/3.0) *M_PI*x*x*x ;

//        cout<<"Estimate encircled_flux_at_pixel_radius: "<< ap_radius<<": "<< EF<<endl;
//
//        cout<<"Analytic calculation: "<< analytic_ef<<endl;
//
//        cout<<"Percent diff: " << (EF - analytic_ef )/ analytic_ef<<endl;

        CHECK_CLOSE(analytic_ef, EF, analytic_ef*0.035);
    }

}


#include <Eigen/Core>
#include <iostream>
using namespace Eigen;
using namespace std;

template<typename Derived>
Eigen::VectorBlock<Derived, 2>
firstTwo(MatrixBase<Derived>& v)
{
  return Eigen::VectorBlock<Derived, 2>(v.derived(), 0);
}

template<typename Derived>
const Eigen::VectorBlock<const Derived, 2>
firstTwo(const MatrixBase<Derived>& v)
{
  return Eigen::VectorBlock<const Derived, 2>(v.derived(), 0);
}

int main(int, char**)
{
  Matrix<int,1,6> v; v << 1,2,3,4,5,6;
  cout << firstTwo(4*v) << endl; // calls the const version
  firstTwo(v) *= 2;              // calls the non-const version
  cout << "Now the vector v is:" << endl << v << endl;
  return 0;
}


#pragma once

#include <Eigen/Core>
#include <string>

///
/// \brief Contains data from the IMU mesaurements.
///
class ImuMeasurement{
public:
  EIGEN_MAKE_ALIGNED_OPERATOR_NEW

  uint64_t t;       ///< ROS time message received (nanoseconds).

  Eigen::Vector3d I_a_WI;  ///< Raw acceleration from the IMU (m/s/s)
  Eigen::Vector3d I_w_WI;  ///< Raw angular velocity from the IMU (deg/s)

  ~ImuMeasurement() {}
  ImuMeasurement();
  ImuMeasurement(const uint64_t _t,
                 const Eigen::Vector3d& _I_a_WI,
                 const Eigen::Vector3d& _I_w_WI);
  friend std::ostream& operator<< (std::ostream& stream, const ImuMeasurement& meas);
};


//
// Created by Vasiliy Ershov on 08/11/2016.
//

#ifndef PROJECT_GAMMA_POISSON_MODEL_HPP
#define PROJECT_GAMMA_POISSON_MODEL_HPP

#include <common/utils/parallel/openmp_wrapper.h>
#include <boost/math/special_functions/binomial.hpp>
#include <boost/math/special_functions/gamma.hpp>
#include <boost/math/special_functions/trigamma.hpp>
#include <vector>
#include "kmer_data.hpp"
#include "thread_utils.h"
#include "valid_hkmer_generator.hpp"
//

namespace n_gamma_poisson_model {

struct QualFunc {
  double alpha_;
  double beta_;

  double operator()(double x) const { return alpha_ * x + beta_; }

  double GenomicLogLikelihood(double x) const {
    const double val = (*this)(x);
    const double exp_point = exp(val);
    return val - (std::isfinite(exp_point) ? log(1 + exp_point) : val);
  }
};

class GammaDistribution {
 private:
  double shape_;
  double rate_;
  double log_gamma_at_shape_;

 public:
  GammaDistribution(const GammaDistribution&) = default;

  GammaDistribution& operator=(const GammaDistribution&) = default;

  GammaDistribution(const double shape = 1, const double rate = 1)
      : shape_(shape), rate_(rate) {
    log_gamma_at_shape_ = boost::math::lgamma(shape_);
  }

  inline double GetShape() const { return shape_; }

  inline double GetRate() const { return rate_; }

  inline double LogGammaAtShape() const { return log_gamma_at_shape_; }
};

class GammaMixture {
 private:
  GammaDistribution first_;
  GammaDistribution second_;
  double first_weight_;

 public:
  GammaMixture() : first_(1, 1), second_(1, 1), first_weight_(-1) {}

  GammaMixture(const GammaDistribution& first,
                const GammaDistribution& second,
                double firstWeight)
      : first_(first), second_(second), first_weight_(firstWeight) {}

  const GammaDistribution& GetFirst() const { return first_; }

  const GammaDistribution& GetSecond() const { return second_; }

  double GetFirstWeight() const { return first_weight_; }
};

class PoissonGammaDistribution {
 private:
  const GammaDistribution& prior_;
  static std::array<double, 100000> log_gamma_integer_cache_;

 private:
  inline double IntLogGamma(size_t count) const {
    if (count < log_gamma_integer_cache_.size()) {
      return log_gamma_integer_cache_[count];
    } else {
      return boost::math::lgamma(((double)count) + 1);
    }
  }

 public:
  PoissonGammaDistribution(const GammaDistribution& prior) : prior_(prior) {}

  inline double PartialLogLikelihood(size_t count) const {
    const double a = prior_.GetShape();
    const double b = prior_.GetRate();

    double ll = 0.0;
    ll += a * log(b) - (a + (double)count) * log(b + 1);
    ll +=
        boost::math::lgamma(prior_.GetShape() + (double)count) - prior_.LogGammaAtShape();
    return ll;
  }

  inline double LogLikelihood(size_t count) const {
    const double a = prior_.GetShape();
    const double b = prior_.GetRate();

    double ll = 0.0;
    ll += a * log(b) - (a + (double)count) * log(b + 1);
    ll += boost::math::lgamma(prior_.GetShape() + ((double)count)) - IntLogGamma(count) -
          prior_.LogGammaAtShape();

    return ll;
  }

  inline double Quantile(double p) const {
    const double a = prior_.GetShape();
    const double b = prior_.GetRate();
    return boost::math::ibeta_inva(a, 1.0 / (1.0 + b), 1.0 - p);
  }

  inline double Cumulative(size_t count) const {
    const double a = prior_.GetShape();
    const double b = prior_.GetRate();

    return 1.0 - boost::math::ibeta((double)count + 1, a, 1.0 / (1.0 + b));
  }
};

constexpr int RunSizeLimit = 8;

class ParametricClusterModel {
 private:
  GammaMixture prior_;
  QualFunc qual_func_;
  double count_threshold_;
  std::array<double, RunSizeLimit> alphas_;

 public:
 public:
  ParametricClusterModel() : count_threshold_(100000) {}

  double ErrorRate(const int runSize) const {
    auto idx = runSize - 1;
    idx = std::max(idx, 0);
    idx = std::min(idx, (int)(alphas_.size() - 1));
    return alphas_[idx] * (runSize == 0 ? 0.5 : 1);
  }

  double ExpectedErrorRate(const hammer::HKMer& from,
                           const hammer::HKMer& to) const {
    double errRate = 0;
    for (unsigned i = 0; i < hammer::K; ++i) {
      errRate +=
          std::abs(from[i].len - to[i].len) * log(ErrorRate(from[i].len));
      //      errRate += std::abs(from[i].len - to[i].len) *
      //      log(ErrorRate(from[i].len)) - log(1.0 - ErrorRate(from[i].len));
    }
    return exp(errRate);
  }

  ParametricClusterModel(const GammaMixture& prior,
                         const QualFunc& qualFunc,
                         const double countThreshold,
                         const std::array<double, RunSizeLimit>& alphas)
      : prior_(prior), qual_func_(qualFunc), count_threshold_(countThreshold) {
    std::copy(alphas.begin(), alphas.end(), alphas_.begin());
    for (unsigned i = 0; i < RunSizeLimit; ++i) {
      INFO("Run length " << i << " estimated error rate " << alphas_[i]);
    }
  }

  ParametricClusterModel(const ParametricClusterModel& other) = default;

  ParametricClusterModel& operator=(const ParametricClusterModel&) = default;

  double QualityLogPrior(double qual) const {
    return std::max(std::min(qual_func_.GenomicLogLikelihood(qual), -1e-10), -1000.0);
  }

  bool NeedSubcluster(const hammer::KMerStat& stat) const {
    return qual_func_.GenomicLogLikelihood(stat.qual) > -0.1 &&
           stat.count >= count_threshold_;
  }

  const GammaDistribution& GenomicPrior() const { return prior_.GetFirst(); }

  const GammaDistribution& NoisePrior() const { return prior_.GetSecond(); }

  double GenerateLogLikelihood(double expectedNoiseCount,
                               size_t noiseCount) const {
    const auto& prior = NoisePrior();

    GammaDistribution posterior(prior.GetShape() + expectedNoiseCount,
                                 prior.GetRate() + 1);
    return PoissonGammaDistribution(posterior).LogLikelihood(noiseCount);
  }

  double GenomicLogLikelihood(size_t count) const {
    const auto& prior = GenomicPrior();
    const double a = prior.GetShape();
    const double b = prior.GetRate();

    double ll = a * log(b) - (a + (double)count) * log(b + 1);
    ll += boost::math::lgamma(prior.GetShape() + ((double)count)) -
          prior.LogGammaAtShape() - boost::math::lgamma(((double)count) + 1);
    return ll;
  }
};

// this class estimate prior distribution.
class TClusterModelEstimator {
 private:
  const KMerData& data_;
  double threshold_;
  unsigned num_threads_;
  size_t max_terations_;
  bool calc_likelihood_;

 private:

  struct TClusterSufficientStat {
    double count_ = 0;
    double qualtiy_ = 0;
    double genomic_class_prob_ = 0;
  };

  struct TQualityStat {
    double quality_ = 0;
    double class_ = 0;

    TQualityStat(double quality, double cls) : quality_(quality), class_(cls) {}
  };

  struct TRunErrorStats {
    const KMerData* data_;
    std::array<double, RunSizeLimit> error_counts_;
    std::array<double, RunSizeLimit> total_count_;

    TRunErrorStats(const KMerData& data)
        : data_(&data){

          };

    std::array<double, RunSizeLimit> EstimateAlphas(
        size_t priorSize = 100) const {
      const double priors[] = {0.002, 0.004, 0.01, 0.02,
                               0.035, 0.05,  0.09, 0.11};

      std::array<double, RunSizeLimit> alphas;
      for (unsigned i = 0; i < RunSizeLimit; ++i) {
        alphas[i] = (error_counts_[i] + priors[i] * (double)priorSize) /
                    (total_count_[i] + (double)priorSize);
      }
      alphas[0] *= 2;
      return alphas;
    };

    TRunErrorStats& operator+=(const TRunErrorStats& other) {
      if (this != &other) {
        for (unsigned i = 0; i < RunSizeLimit; ++i) {
          error_counts_[i] += other.error_counts_[i];
          total_count_[i] += other.total_count_[i];
        }
      }
      return *this;
    }

    void Add(const std::vector<size_t>& indices, size_t centerIdx) {
      const auto& center = (*data_)[centerIdx].kmer;

      for (auto idx : indices) {
        if (idx == centerIdx) {
          continue;
        }
        double errKmerCount = (double)(*data_)[idx].count;
        const auto& errKmer = (*data_)[idx].kmer;
        for (unsigned i = 0; i < hammer::K; ++i) {
          if (center[i].len > RunSizeLimit) {
            continue;
          }
          const int len = center[i].len - 1;
          total_count_[len] += errKmerCount;
          if (center[i].len != errKmer[i].len) {
            error_counts_[len] += errKmerCount;
          }
        }
      }
      for (unsigned i = 0; i < hammer::K; ++i) {
        if (center[i].len > RunSizeLimit) {
          continue;
        }
        total_count_[center[i].len - 1] += (*data_)[centerIdx].count;
      }
    }
  };

  inline void Expectation(const PoissonGammaDistribution& first,
                          const PoissonGammaDistribution& second,
                          const QualFunc& qualFunc,
                          TClusterSufficientStat& center) const {
    const double logPrior = qualFunc.GenomicLogLikelihood(center.qualtiy_) +
                            log(boost::math::gamma_q(center.count_, threshold_));

    const double firstLL = first.PartialLogLikelihood((size_t)center.count_) + logPrior;
    const double secondLL = second.PartialLogLikelihood((size_t)center.count_) +
                            log(std::max(1.0 - exp(logPrior), 1e-20));

    const double posterior = 1.0 / (1.0 + exp(secondLL - firstLL));
    center.genomic_class_prob_ = posterior;
  }

  inline void QualityExpectation(const QualFunc& qualFunc,
                                 TClusterSufficientStat& center) const {
    center.genomic_class_prob_ =
        exp(qualFunc.GenomicLogLikelihood(center.qualtiy_));
  }

  inline TClusterSufficientStat Create(const size_t centerIdx) const {
    TClusterSufficientStat stat;
    stat.genomic_class_prob_ =
        data_[centerIdx].count > 0
            ? boost::math::gamma_q(data_[centerIdx].count, threshold_)
            : 0;
    stat.count_ = data_[centerIdx].count;
    stat.qualtiy_ = data_[centerIdx].qual;
    return stat;
  }

  std::vector<TClusterSufficientStat> CreateSufficientStats(
      const std::vector<size_t>& clusterCenters) const {
    std::vector<TClusterSufficientStat> clusterSufficientStat;
    clusterSufficientStat.reserve(clusterCenters.size());

    for (size_t i = 0; i < clusterCenters.size(); ++i) {
      const size_t centerIdx = clusterCenters[i];
      auto stat = Create(centerIdx);
      if (stat.count_ > 0) {
        clusterSufficientStat.push_back(stat);
      }
    }
    return clusterSufficientStat;
  }

  std::vector<TQualityStat> CreateQualityStats(
      const std::vector<std::vector<size_t>>& clusters,
      const std::vector<size_t>& clusterCenters) const {
    std::vector<TQualityStat> qualities;
    qualities.reserve(clusterCenters.size());

    for (size_t i = 0; i < clusterCenters.size(); ++i) {
      const size_t centerIdx = clusterCenters[i];
      if (data_[centerIdx].count >= threshold_) {
        for (auto idx : clusters[i]) {
          if (idx != centerIdx) {
            qualities.push_back(TQualityStat(data_[idx].qual, 0));
          }
        }
        qualities.push_back(TQualityStat(data_[centerIdx].qual, 1));
      }
    }
    return qualities;
  }

  template <bool WEIGHTED = true>
  class TCountsStat {
   private:
    double count_ = 0;
    double count2_ = 0;
    double weight_ = 0;

   public:
    void Add(const TClusterSufficientStat& stat) {
      const double w = (WEIGHTED ? stat.genomic_class_prob_ : 1.0);
      count_ += w * stat.count_;
      count2_ += w * stat.count_ * stat.count_;
      weight_ += w;
    }

    TCountsStat& operator+=(const TCountsStat& other) {
      if (this != &other) {
        count_ += other.count_;
        count2_ += other.count2_;
        weight_ += other.weight_;
      }
      return *this;
    }

    double GetWeightedSum() const { return count_; }

    double GetWeightedSum2() const { return count2_; }

    double GetWeight() const { return weight_; }
  };

  class TLogGammaStat {
   private:
    double genomic_shape_;
    double non_genomic_shape_;
    double genomic_log_gamma_sum_ = 0;
    double non_genomic_log_gamma_sum_ = 0;

   public:
    TLogGammaStat(double genomicShape, double nonGenomicShape)
        : genomic_shape_(genomicShape), non_genomic_log_gamma_sum_(nonGenomicShape) {}

    void Add(const TClusterSufficientStat& stat) {
      genomic_log_gamma_sum_ += stat.genomic_class_prob_ *
                            boost::math::lgamma(stat.count_ + genomic_shape_);
      non_genomic_log_gamma_sum_ +=
          (1.0 - stat.genomic_class_prob_) *
          boost::math::lgamma(stat.count_ + non_genomic_shape_);
    }

    TLogGammaStat& operator+=(const TLogGammaStat& other) {
      if (this != &other) {
        genomic_log_gamma_sum_ += other.genomic_log_gamma_sum_;
        non_genomic_log_gamma_sum_ += other.non_genomic_log_gamma_sum_;
      }
      return *this;
    }

    double GetGenomicLogGammaSum() const { return genomic_log_gamma_sum_; }

    double GetNonGenomicLogGammaSum() const { return non_genomic_log_gamma_sum_; }
  };

  class TQualityLogitLinearRegressionPoint {
   private:
    // p(genomic) = exp(Alpha qual + beta) / (1.0 + exp(Alpha qual + beta))
    QualFunc func_;

    double likelihood_ = 0;

    double der_alpha_ = 0;
    double der_beta_ = 0;

    double der2_alpha_ = 0;
    double der2_beta_ = 0;
    double der2_alpha_beta_ = 0;

   public:
    TQualityLogitLinearRegressionPoint(QualFunc func) : func_(func) {}

    void Add(const TClusterSufficientStat& statistic) {
      Add(statistic.genomic_class_prob_, statistic.qualtiy_);
    }

    void Add(const TQualityStat& statistic) {
      Add(statistic.class_, statistic.quality_);
    }

    void Add(const double firstClassProb, double qual) {
      const double val = func_(qual);
      const double expPoint = exp(val);
      const double p =
          std::isfinite(expPoint) ? expPoint / (1.0 + expPoint) : 1.0;

      der_alpha_ += (firstClassProb - p) * qual;
      der_beta_ += firstClassProb - p;

      der2_alpha_ -= sqr(qual) * p * (1 - p);
      der2_beta_ -= p * (1 - p);
      der2_alpha_beta_ -= qual * p * (1 - p);

      likelihood_ += firstClassProb * val -
                    (std::isfinite(expPoint) ? log(1 + expPoint) : val);
    }

    TQualityLogitLinearRegressionPoint& operator+=(
        const TQualityLogitLinearRegressionPoint& other) {
      if (this != &other) {
        likelihood_ += other.likelihood_;

        der_alpha_ += other.der_alpha_;
        der_beta_ += other.der_beta_;

        der2_alpha_ += other.der2_alpha_;
        der2_beta_ += other.der2_beta_;
        der2_alpha_beta_ += other.der2_alpha_beta_;
      }
      return *this;
    }

    double GetLikelihood() const { return likelihood_; }

    double GetDerAlpha() const { return der_alpha_; }

    double GetDerBeta() const { return der_beta_; }

    double GetDer2Alpha() const { return der2_alpha_; }

    double GetDer2Beta() const { return der2_beta_; }

    double GetDer2AlphaBeta() const { return der2_alpha_beta_; }
  };

  QualFunc Update(const QualFunc& current,
                   const TQualityLogitLinearRegressionPoint& pointStats) const {
    const double dera = pointStats.GetDerAlpha();
    const double derb = pointStats.GetDerBeta();

    const double daa = pointStats.GetDer2Alpha() + 1e-3;
    const double dbb = pointStats.GetDer2Beta() + 1e-3;
    const double dab = pointStats.GetDer2AlphaBeta();
    const double det = daa * dbb - sqr(dab);

    double stepAlpha = (dbb * dera - dab * derb) / det;
    double stepBeta = (daa * derb - dab * dera) / det;

    INFO("Quality estimation iteration gradient: " << dera << " " << derb);
    INFO("Quality estimation likelihood: " << pointStats.GetLikelihood());

    return {current.alpha_ - stepAlpha, current.beta_ - stepBeta};
  }

  class TGammaDerivativesStats {
   private:
    double first_class_shift_;
    double second_class_shift_;

    double digamma_sum_first_ = 0;
    double trigamma_sum_first_ = 0;

    double digamma_sum_second_ = 0;
    double trigamma_sum_second_ = 0;

   public:
    TGammaDerivativesStats(double firstShift, double secondShift)
        : first_class_shift_(firstShift), second_class_shift_(secondShift) {}

    void Add(const TClusterSufficientStat& statistic) {
      const double p = statistic.genomic_class_prob_;
      digamma_sum_first_ +=
          p > 1e-3 ? p * boost::math::digamma(statistic.count_ + first_class_shift_)
                   : 0;
      trigamma_sum_first_ +=
          p > 1e-3
              ? p * boost::math::trigamma(statistic.count_ + first_class_shift_)
              : 0;

      digamma_sum_second_ +=
          p < (1.0 - 1e-3) ? (1.0 - p) * boost::math::digamma(statistic.count_ +
                                                              second_class_shift_)
                           : 0;
      trigamma_sum_second_ +=
          p < (1.0 - 1e-3) ? (1.0 - p) * boost::math::trigamma(statistic.count_ +
                                                               second_class_shift_)
                           : 0;
    }

    TGammaDerivativesStats& operator+=(const TGammaDerivativesStats& other) {
      if (this != &other) {
        digamma_sum_first_ += other.digamma_sum_first_;
        trigamma_sum_first_ += other.trigamma_sum_first_;

        digamma_sum_second_ = other.digamma_sum_second_;
        trigamma_sum_second_ += other.trigamma_sum_second_;
      }
      return *this;
    }

    double GetDigammaSumFirst() const { return digamma_sum_first_; }

    double GetTrigammaSumFirst() const { return trigamma_sum_first_; }

    double GetDigammaSumSecond() const { return digamma_sum_second_; }

    double GetTrigammaSumSecond() const { return trigamma_sum_second_; }
  };

  static inline double sqr(double x) { return x * x; }

  struct TDirection {
    double Direction;
    double GradientNorm;
    double Mu;
  };

  static TDirection MoveDirection(double shape, const double weightedSum,
                                  const double weight, const double digammaSum,
                                  const double trigammaSum,
                                  double regularizer = 1e-4) {
    const double mu = weight / weightedSum;
    const double digammaAtShape = boost::math::digamma(shape);
    const double trigammaAtShape = boost::math::trigamma(shape);

    const double b = mu * shape;

    const double der =
        weight * (log(b) - log(b + 1) - digammaAtShape) + digammaSum;
    const double der2 =
        trigammaSum + weight * (1.0 / shape - mu / (b + 1) - trigammaAtShape);

    return {-der / (der2 + regularizer), std::abs(der), mu};
  }

  double Likelihood(GammaDistribution& prior, double weightedSum,
                    double weight, double lgammaSum) {
    const double a = prior.GetShape();
    const double b = prior.GetRate();
    return weight * a * (log(b) - log(b + 1)) + weightedSum * log(b + 1) +
           lgammaSum - weight * prior.LogGammaAtShape();
  }

 public:
  TClusterModelEstimator(const KMerData& data, double threshold,
                         unsigned num_threads = 16,
                         size_t maxIterations = 40,
                         bool calcLikelihood = false)
      : data_(data),
        threshold_(threshold),
        num_threads_(num_threads),
        max_terations_(maxIterations),
        calc_likelihood_(calcLikelihood) {}

  static inline GammaDistribution Update(const GammaDistribution& point,
                                          const TDirection& direction,
                                          double minShape = 0.01) {
    double shape = std::max(point.GetShape() + direction.Direction, minShape);
    double rate = shape * direction.Mu;
    return GammaDistribution(shape, rate);
  }

  static GammaDistribution MomentMethodEstimator(const double sum,
                                                  const double sum2,
                                                  const double weight) {
    const double m = sum / weight;
    const double var = sum2 / weight - m * m;
    const double rate = 1.0 / std::max(var / m - 1, 1e-3);
    const double shape = m * rate;
    return GammaDistribution(shape, rate);
  }

  ParametricClusterModel Estimate(
      const std::vector<std::vector<size_t>>& clusters,
      const std::vector<size_t>& clusterCenter, const bool useEM = false,
      const size_t sample = 0) {
    if (sample && clusters.size() > sample) {
      std::vector<std::vector<size_t>> sampledClusters;
      std::vector<size_t> sampledCenters;
    }

    const auto qualityFunc = [&]() -> QualFunc {
      auto qualStats = CreateQualityStats(clusters, clusterCenter);

      QualFunc cursor = {-1e-5, 0.0};

      for (unsigned i = 0; i < 15; ++i) {
        const auto qualDerStats =
            n_computation_utils::TAdditiveStatisticsCalcer<
                TQualityStat, TQualityLogitLinearRegressionPoint>(qualStats,
                                                                  num_threads_)
                .Calculate([&]() -> TQualityLogitLinearRegressionPoint {
                  return TQualityLogitLinearRegressionPoint(cursor);
                });

        cursor = Update(cursor, qualDerStats);

        if ((std::abs(qualDerStats.GetDerAlpha()) +
             std::abs(qualDerStats.GetDerBeta())) < 1e-2) {
          break;
        }
      }

      INFO("Quality function: " << cursor.alpha_ << "q + " << cursor.beta_);
      return cursor;
    }();

    auto alphas = [&]() -> std::array<double, RunSizeLimit> {
      TRunErrorStats errorStats =
          n_computation_utils::ParallelStatisticsCalcer<TRunErrorStats>(
              num_threads_)
              .Calculate(
                  clusters.size(),
                  [&]() -> TRunErrorStats { return TRunErrorStats(data_); },
                  [&](TRunErrorStats& stat, size_t k) {
                    if (data_[clusterCenter[k]].count >= threshold_) {
                      stat.Add(clusters[k], clusterCenter[k]);
                    }
                  });
      return errorStats.EstimateAlphas();
    }();

    std::vector<TClusterSufficientStat> clusterSufficientStat =
        CreateSufficientStats(clusterCenter);

    const auto totalStats =
        n_computation_utils::TAdditiveStatisticsCalcer<TClusterSufficientStat,
                                                     TCountsStat<false>>(
            clusterSufficientStat, num_threads_)
            .Calculate([]() -> TCountsStat<false> {
              return TCountsStat<false>();
            });

#pragma omp parallel for num_threads(num_threads_)
    for (size_t k = 0; k < clusterSufficientStat.size(); ++k) {
      QualityExpectation(qualityFunc, clusterSufficientStat[k]);
    }

    auto countsStats =
        n_computation_utils::TAdditiveStatisticsCalcer<TClusterSufficientStat,
                                                     TCountsStat<true>>(
            clusterSufficientStat, num_threads_)
            .Calculate([]() -> TCountsStat<true> {
              return TCountsStat<true>();
            });

    GammaDistribution genomicPrior = [&]() -> GammaDistribution {
      const double m = countsStats.GetWeightedSum() / countsStats.GetWeight();
      const double var =
          countsStats.GetWeightedSum2() / countsStats.GetWeight() - m * m;
      const double rate = 1.0 / std::max(var / m - 1, 1e-3);
      const double shape = m * rate;
      return GammaDistribution(shape, rate);
    }();

    GammaDistribution nonGenomicPrior = [&]() -> GammaDistribution {
      const double m =
          (totalStats.GetWeightedSum() - countsStats.GetWeightedSum()) /
          (totalStats.GetWeight() - countsStats.GetWeight());
      const double var =
          (totalStats.GetWeightedSum2() - countsStats.GetWeightedSum2()) /
              (totalStats.GetWeight() - countsStats.GetWeight()) -
          m * m;
      const double rate = 1.0 / std::max(var / m - 1, 1e-3);
      const double shape = m * rate;
      return GammaDistribution(shape, rate);
    }();

    for (unsigned i = 0, steps = 0; i < max_terations_; ++i, ++steps) {
      auto gammaDerStats =
          n_computation_utils::TAdditiveStatisticsCalcer<TClusterSufficientStat,
                                                       TGammaDerivativesStats>(
              clusterSufficientStat, num_threads_)
              .Calculate([&]() -> TGammaDerivativesStats {
                return TGammaDerivativesStats(genomicPrior.GetShape(),
                                              nonGenomicPrior.GetShape());
              });

      auto genomicDirection = MoveDirection(
          genomicPrior.GetShape(), countsStats.GetWeightedSum(),
          countsStats.GetWeight(), gammaDerStats.GetDigammaSumFirst(),
          gammaDerStats.GetTrigammaSumFirst());

      auto nonGenomicDirection = MoveDirection(
          nonGenomicPrior.GetShape(),
          totalStats.GetWeightedSum() - countsStats.GetWeightedSum(),
          totalStats.GetWeight() - countsStats.GetWeight(),
          gammaDerStats.GetDigammaSumSecond(),
          gammaDerStats.GetTrigammaSumSecond());

      auto gradientNorm =
          genomicDirection.GradientNorm + nonGenomicDirection.GradientNorm;

      INFO("Iteration #" << i << " gradient norm " << gradientNorm);

      genomicPrior = Update(genomicPrior, genomicDirection);
      nonGenomicPrior = Update(nonGenomicPrior, nonGenomicDirection);

      if (calc_likelihood_) {
        auto logGammaStats =
            n_computation_utils::TAdditiveStatisticsCalcer<TClusterSufficientStat,
                                                         TLogGammaStat>(
                clusterSufficientStat, num_threads_)
                .Calculate([&]() -> TLogGammaStat {
                  return TLogGammaStat(genomicPrior.GetShape(),
                                       nonGenomicPrior.GetShape());
                });

        INFO("Genomic likelihood: " << Likelihood(
                 genomicPrior, countsStats.GetWeightedSum(),
                 countsStats.GetWeight(),
                 logGammaStats.GetGenomicLogGammaSum()));

        INFO("NonGenomic likelihood: " << Likelihood(
                 nonGenomicPrior,
                 totalStats.GetWeightedSum() - countsStats.GetWeightedSum(),
                 totalStats.GetWeight() - countsStats.GetWeight(),
                 logGammaStats.GetNonGenomicLogGammaSum()));
      }

      {
        INFO("Genomic gamma prior estimation step: shape "
             << genomicPrior.GetShape() << " and rate "
             << genomicPrior.GetRate());
        INFO("Nongenomic gamma prior estimation step: shape "
             << nonGenomicPrior.GetShape() << " and rate "
             << nonGenomicPrior.GetRate());
      }

      double shapeDiff = std::abs(genomicDirection.Direction) +
                         std::abs(nonGenomicDirection.Direction);

      if (useEM) {
        if ((shapeDiff < 1e-2) || gradientNorm < 1e-1 ||
            (steps == 5 && (i < max_terations_ - 10))) {
          PoissonGammaDistribution genomic(genomicPrior);
          PoissonGammaDistribution nonGenomic(nonGenomicPrior);
#pragma omp parallel for num_threads(num_threads_)
          for (size_t k = 0; k < clusterSufficientStat.size(); ++k) {
            Expectation(genomic, nonGenomic, qualityFunc,
                        clusterSufficientStat[k]);
          }

          countsStats = n_computation_utils::TAdditiveStatisticsCalcer<
                            TClusterSufficientStat, TCountsStat<true>>(
                            clusterSufficientStat, num_threads_)
                            .Calculate([]() -> TCountsStat<true> {
                              return TCountsStat<true>();
                            });
          steps = 0;
        }
      } else {
        if ((shapeDiff < 1e-4) || gradientNorm < 1e-2) {
          break;
        }
      }
    }

    INFO("Genomic gamma prior genomic estimated with shape "
         << genomicPrior.GetShape() << " and rate " << genomicPrior.GetRate());
    INFO("Nongenomic Gamma prior estimated with shape "
         << nonGenomicPrior.GetShape() << " and rate "
         << nonGenomicPrior.GetRate());

    return ParametricClusterModel(
        GammaMixture(genomicPrior, nonGenomicPrior,
                      countsStats.GetWeight() / totalStats.GetWeight()),
        qualityFunc, threshold_, alphas);
  }

  static GammaDistribution EstimatePrior(const std::vector<size_t>& counts) {
    const size_t observations = counts.size();
    double sum = 0;
    double sum2 = 0;
    for (auto count : counts) {
      sum += (double)count;
      sum2 += (double)count * (double)count;
    }

    GammaDistribution prior =
        TClusterModelEstimator::MomentMethodEstimator(sum, sum2, (double)observations);

    for (unsigned i = 0, steps = 0; i < 10; ++i, ++steps) {
      double digammaSum = 0;
      double trigammaSum = 0;
      for (auto count : counts) {
        digammaSum += boost::math::digamma((double)count + prior.GetShape());
        trigammaSum += boost::math::trigamma((double)count + prior.GetShape());
      }

      auto direction = MoveDirection(prior.GetShape(), sum, (double)observations,
                                     digammaSum, trigammaSum);

      const double shapeDiff = std::abs(direction.Direction);
      if (shapeDiff < 1e-3 || (direction.GradientNorm < 1e-4)) {
        break;
      }
      prior = Update(prior, direction, 1e-2);
    }
    return prior;
  }
};

}  // namespace NGammaPoissonModel

#endif  // PROJECT_GAMMA_POISSON_MODEL_HPP


#ifndef MLT_MODELS_REGRESSORS_REGRESSOR_HPP
#define MLT_MODELS_REGRESSORS_REGRESSOR_HPP

#include <Eigen/Core>

#include "../base.hpp"

namespace mlt {
namespace models {
namespace regressors {
	template <class ConcreteType>
	class Regressor : public Predictor<ConcreteType, MatrixXd> {
	public:
		inline auto score(Features input, Target target) const {
			return (1 - ((y - _self().predict(input)).array().pow(2).rowwise().sum().array() / (y.colwise() - (y.rowwise().mean())).array().pow(2).rowwise().sum().array())).mean();
		}

	protected:
		Regressor() = default;
		Regressor(const Regressor&) = default;
		Regressor(Regressor&&) = default;
		Regressor& operator=(const Regressor&) = default;
		~Regressor() = default;

		inline auto _to_target_matrix(Target target) {
			return target;
		}
	};
}
}
}
#endif


#define EIGEN_USE_MKL_ALL

#include <iostream> // for standard output
#include <fstream> // for file input
#include <iomanip> // some nice printing functions
#include <random> // for randomness
#include <boost/algorithm/string.hpp> // string manipulation library
#include <boost/filesystem.hpp> // filesystem path manipulation library
#include <boost/program_options.hpp> // options parsing library
#include <eigen3/Eigen/Dense> // linear algebra library

#include "constants.h"
#include "qp-math.h"
#include "gates.h"
#include "nv-math.h"
#include "nv-control.h"

using namespace std;
using namespace Eigen;
namespace fs = boost::filesystem;
namespace po = boost::program_options;

int main(const int arg_num, const char *arg_vec[]) {

  // -------------------------------------------------------------------------------------
  // Set input options
  // -------------------------------------------------------------------------------------

  const uint help_text_length = 90;

  unsigned long long int seed;

  po::options_description general("General options", help_text_length);
  general.add_options()
    ("help,h", "produce help message")
    ("seed", po::value<unsigned long long int>(&seed)->default_value(0),
     "seed for random number generator")
    ;

  bool pair_search;
  bool coherence_scan;
  bool coherence_signal;
  bool angular_coherence_signal;
  bool rotation;
  bool coupling;
  bool iswap;
  bool swap;
  bool identity;
  bool swap_nvst;
  bool larmor_identity;
  bool initialize;
  bool initialize_x;
  bool initialize_larmor;
  bool testing_mode;

  po::options_description simulations("Available simulations",help_text_length);
  simulations.add_options()
    ("pair_search",
     po::value<bool>(&pair_search)->default_value(false)->implicit_value(true),
     "identify and characterize addressable larmor pairs")
    ("scan", po::value<bool>(&coherence_scan)->default_value(false)->implicit_value(true),
     "perform NV coherence scan for effective larmor frequencies")
    ("signal",
     po::value<bool>(&coherence_signal)->default_value(false)->implicit_value(true),
     "measure NV coherence signal as a function of the coupling constant"
     " at the frequency of a target nucleus")
    ("angular_signal", po::value<bool>(&angular_coherence_signal)
     ->default_value(false)->implicit_value(true),
     "measure NV coherence signal as a function of the coupling constant"
     " and angle of a control field at the frequency of a target nucleus")
    ("rotate", po::value<bool>(&rotation)->default_value(false)->implicit_value(true),
     "rotate an individual nucleus")
    ("couple", po::value<bool>(&coupling)->default_value(false)->implicit_value(true),
     "couple an individual nucleus to NV center")
    ("iswap", po::value<bool>(&iswap)->default_value(false)->implicit_value(true),
     "compute iSWAP fidelities")
    ("swap", po::value<bool>(&swap)->default_value(false)->implicit_value(true),
     "compute SWAP fidelities")
    ("identity", po::value<bool>(&identity)->default_value(false)->implicit_value(true),
     "compute fidelity of an identity operation on a single spin")
    ("swap_nvst", po::value<bool>(&swap_nvst)->default_value(false)->implicit_value(true),
     "compute SWAP_NVST fidelity")
    ("larmor_identity",
     po::value<bool>(&larmor_identity)->default_value(false)->implicit_value(true),
     "compute fidelity of an identity operation on a larmor qubit")
    ("initialize",
     po::value<bool>(&initialize)->default_value(false)->implicit_value(true),
     "compute fidelity of a deterministic initialization of a thermalized nucleus"
     " into |u> or |d>")
    ("initialize_x",
     po::value<bool>(&initialize_x)->default_value(false)->implicit_value(true),
     "compute fidelity of a probabalistic initialization of a thermalized nucleus"
     " into |u> +/- |d>")
    ("initialize_larmor",
     po::value<bool>(&initialize_larmor)->default_value(false)->implicit_value(true),
     "compute fidelity of a probabalistic initialization of a larmor pair"
     " from |dd> into |ud> +/- |du>")
    ("test" ,po::value<bool>(&testing_mode)->default_value(false)->implicit_value(true),
     "enable testing mode")
    ;

  double c13_abundance;
  double c13_factor;
  double nuclear_isolation;
  double larmor_isolation;
  double larmor_isolation_in_kHz;
  double min_hyperfine_xy;
  double min_hyperfine_xy_in_kHz;
  uint max_cluster_size;
  double hyperfine_cutoff;
  double hyperfine_cutoff_in_kHz;
  int ms;
  uint k_DD_int;
  axy_harmonic k_DD;
  double static_Bz;
  double static_Bz_in_gauss;
  double identity_time;
  double scale_factor;
  double integration_factor;
  bool no_nn;

  po::options_description simulation_options("Simulation options",help_text_length);
  simulation_options.add_options()
    ("c13_factor", po::value<double>(&c13_factor)->default_value(1),
     "abundance of C-13 relative to its natural abundance")
    ("nuclear_isolation", po::value<double>(&nuclear_isolation)->default_value(100),
     "maximum internuclear coupling factor for an 'isolated' nucleus (Hz)")
    ("larmor_isolation", po::value<double>(&larmor_isolation_in_kHz)->default_value(0),
     "isolation factor of parallel component of hyperfine field"
     " for an 'isolated' nucleus (kHz)")
    ("min_hyperfine_xy", po::value<double>(&min_hyperfine_xy_in_kHz)->default_value(0),
     "minimum magnitude of hyperfine field perpendicular to the NV axis"
     " for a larmor pair to be counted in a pair search (kHz)")
    ("max_cluster_size", po::value<uint>(&max_cluster_size)->default_value(6),
     "maximum allowable size of C-13 clusters")
    ("hyperfine_cutoff", po::value<double>(&hyperfine_cutoff_in_kHz)->default_value(10),
     "set cutoff scale for hyperfine field (kHz)")
    ("ms", po::value<int>(&ms)->default_value(1),
     "NV center spin state used with |0> for an effective two-level system (+/-1)")
    ("k_DD", po::value<uint>(&k_DD_int)->default_value(1),
     "resonance harmonic used in spin addressing (1 or 3)")
    ("static_Bz", po::value<double>(&static_Bz_in_gauss)->default_value(140.1,"140.1"),
     "strength of static magnetic field along the NV axis (gauss)")
    ("identity_time", po::value<double>(&identity_time)->default_value(1),
     "time of identity operation (s)")
    ("scale_factor", po::value<double>(&scale_factor)->default_value(10),
     "factor used to define different scales (i.e. if a << b, then a = b/scale_factor)")
    ("integration_factor", po::value<double>(&integration_factor)->default_value(10),
     "factor used to determine size of integration step size")
    ("no_nn" ,po::value<bool>(&no_nn)->default_value(false)->implicit_value(true),
     "turn off internuclear couplings")
    ;

  vector<uint> target_nuclei;
  bool target_pairs;
  double angle;
  double angle_over_pi;
  double target_polar;
  double target_pitch_over_pi;
  double target_azimuth;
  double target_azimuth_over_pi;
  double nv_polar;
  double nv_pitch_over_pi;
  double nv_azimuth;
  double nv_azimuth_over_pi;

  po::options_description addressing_options("Nuclear spin addressing options"
                                             " (all angles are in units of pi radians)",
                                             help_text_length);
  addressing_options.add_options()
    ("target", po::value<vector<uint>>(&target_nuclei)->multitoken(),
     "indices of nuclei to target")
    ("target_pairs",
     po::value<bool>(&target_pairs)->default_value(false)->implicit_value(true),
     "target only one nucleus in each larmor pair")
    ("angle", po::value<double>(&angle_over_pi)->default_value(1),
     "angle of operation to perform")
    ("target_pitch", po::value<double>(&target_pitch_over_pi)->default_value(0),
     "pitch (angle above x-y plane) of target rotation axis")
    ("target_azimuth", po::value<double>(&target_azimuth_over_pi)->default_value(0),
     "azimuthal angle of target rotation axis")
    ("nv_pitch", po::value<double>(&nv_pitch_over_pi)->default_value(0.5),
     "pitch (angle above x-y plane) of NV rotation axis")
    ("nv_azimuth", po::value<double>(&nv_azimuth_over_pi)->default_value(0),
     "azimuthal angle of NV rotation axis")
    ;

  uint coherence_bins;
  double scan_time;
  double scan_time_in_ms;
  double f_DD;
  uint angular_resolution;
  double signal_gB_factor;
  double max_f_factor;

  po::options_description scan_options("Coherence measurement options",help_text_length);
  scan_options.add_options()
    ("bins", po::value<uint>(&coherence_bins)->default_value(500),
     "number of bins in coherence measurements")
    ("scan_time", po::value<double>(&scan_time_in_ms)->default_value(1),
     "time for each measurement in the coherence scan (microseconds)")
    ("f_DD", po::value<double>(&f_DD)->default_value(0.06,"0.06"),
     "magnitude of fourier component used in coherence scanning")
    ("angular_resolution", po::value<uint>(&angular_resolution)->default_value(20),
     "angular resolution for magnetic field direction during signal measurement"
     " (1/[pi radians])")
    ("signal_gB_factor", po::value<double>(&signal_gB_factor)->default_value(100),
     "sets signal field strength to static_gBz/signal_gB_factor")
    ("max_f_factor", po::value<double>(&max_f_factor)->default_value(0.5),
     "factor scaling maximum value of f_DD in coherence signal measurement")
    ;

  string lattice_file;

  po::options_description file_io("File IO",help_text_length);
  file_io.add_options()
    ("lattice_file", po::value<string>(&lattice_file),
     "input file defining system configuration")
    ;

  bool print_lattice;
  bool target_info;

  po::options_description print_options("Available printing options",help_text_length);
  print_options.add_options()
    ("print_lattice",
     po::value<bool>(&print_lattice)->default_value(false)->implicit_value(true),
     "print C-13 lattice to a file")
    ("target_info",
     po::value<bool>(&target_info)->default_value(false)->implicit_value(true),
     "print information about target nuclei")
    ;

  po::options_description all("Allowed options");
  all.add(general);
  all.add(simulations);
  all.add(simulation_options);
  all.add(addressing_options);
  all.add(scan_options);
  all.add(file_io);
  all.add(print_options);

  // collect inputs
  po::variables_map inputs;
  po::store(parse_command_line(arg_num, arg_vec, all), inputs);
  po::notify(inputs);

  // if requested, print help text
  if (inputs.count("help")) {
    cout << all;
    return 0;
  }

  // -------------------------------------------------------------------------------------
  // Run a sanity check on inputs
  // -------------------------------------------------------------------------------------

  // determine whether certain options were used
  bool using_input_lattice = inputs.count("lattice_file");
  bool set_hyperfine_cutoff = !inputs["hyperfine_cutoff"].defaulted();
  bool set_c13_factor = !inputs["c13_factor"].defaulted();
  bool set_target_nuclei = inputs.count("target");

  // make sure we are either printing something, or performing a simulation
  if (!testing_mode) {
    const bool printing = print_lattice || target_info;
    if (!printing) {
      if (int(pair_search)
          + int(coherence_scan)
          + int(coherence_signal)
          + int(angular_coherence_signal)
          + int(rotation)
          + int(coupling)
          + int(iswap)
          + int(swap)
          + int(identity)
          + int(swap_nvst)
          + int(larmor_identity)
          + int(initialize)
          + int(initialize_x)
          + int(initialize_larmor)
          != 1) {
        cout << "Please choose one simulation to perform\n";
        return -1;
      }
    }
  }

  // check lattice options
  assert(!(print_lattice && using_input_lattice));
  assert(!(using_input_lattice && !fs::exists(lattice_file)));
  assert(!(using_input_lattice && set_hyperfine_cutoff));
  assert(!(using_input_lattice && set_c13_factor));

  // check targeting options
  assert(!(set_target_nuclei && target_pairs));

  // verify validity of other values
  assert(hyperfine_cutoff_in_kHz > 0);
  assert(c13_factor >= 0);
  assert(nuclear_isolation >= 0);
  assert(larmor_isolation_in_kHz >= 0);
  assert(min_hyperfine_xy_in_kHz >= 0);
  assert(max_cluster_size > 0);
  assert(ms == 1 || ms == -1);
  assert((k_DD_int == 1) || (k_DD_int == 3));
  assert(scale_factor > 1);
  assert(integration_factor > 1);

  if (coherence_scan) {
    assert(coherence_bins > 0);
    assert(scan_time_in_ms > 0);
  }
  assert(angular_resolution > 2);
  assert(signal_gB_factor > 1);
  assert(max_f_factor <= 1);

  // set some variables based on iputs
  c13_abundance = c13_factor*c13_natural_abundance;
  larmor_isolation = larmor_isolation_in_kHz * 1e3;
  min_hyperfine_xy = min_hyperfine_xy_in_kHz * 1e3;
  hyperfine_cutoff = hyperfine_cutoff_in_kHz * 1e3;
  k_DD = (k_DD_int == 1 ? first : third);
  static_Bz = static_Bz_in_gauss * gauss;
  scan_time = scan_time_in_ms * 1e-3;

  angle = angle_over_pi * pi;
  target_polar = pi/2 - target_pitch_over_pi * pi;
  target_azimuth = target_azimuth_over_pi * pi;
  nv_polar = pi/2 - nv_pitch_over_pi * pi;
  nv_azimuth = nv_azimuth_over_pi * pi;

  uniform_real_distribution<double> rnd(0.0,1.0); // uniform distribution on [0,1)
  mt19937_64 generator(seed); // use and seed the 64-bit Mersenne Twister 19937 generator

  // -------------------------------------------------------------------------------------
  // Construct lattice of nuclei
  // -------------------------------------------------------------------------------------

  vector<Vector3d> nuclei;

  if (!using_input_lattice) { // place nuclei at lattice sites

    // number of fcc cells to simulate out from the origin
    const int cell_radius
      = round(pow(abs(g_e*g_C13)/(4*pi*a0*a0*a0*hyperfine_cutoff),1.0/3));

    // set positions of nuclei at lattice sites
    for (uint b: {0,1}) {
      for (int l = -2*cell_radius; l <= 2*cell_radius; l++) {
        for (int m = -2*cell_radius; m <= 2*cell_radius; m++) {
          for (int n = -2*cell_radius; n <= 2*cell_radius; n++) {
            if (rnd(generator) < c13_abundance) { // check for C-13 isotopic abundance
              if (l != 0 || m != 0 || n != 0) { // don't place C-13 nucleus on NV sites
                const Vector3d pos = b*ao+l*a1+m*a2+n*a3;
                // only place C-13 nuclei with a hyperfine field strength above the cutoff
                if (hyperfine(pos).norm() > hyperfine_cutoff) {
                  nuclei.push_back(pos);
                }
              }
            }
          }
        }
      }
    }
    cout << "Placed " << nuclei.size() << " C-13 nuclei\n\n";
    if (nuclei.size() == 0) return 0;

    if (print_lattice) {
      for (uint i = 0; i < nuclei.size(); i++) {
        cout << nuclei.at(i)(0) << " "
             << nuclei.at(i)(1) << " "
             << nuclei.at(i)(2) << endl;
      }
    }

  } else { // if using_input_lattice, read in the lattice

    string line;
    ifstream lattice(lattice_file);

    // get C-13 positions
    double x,y,z;
    while (getline(lattice,line,' ')) {
      x = stod(line);
      getline(lattice,line,' ');
      y = stod(line);
      getline(lattice,line);
      z = stod(line);
      nuclei.push_back((Vector3d() << x,y,z).finished());
    }
    lattice.close();

    // assert that no C-13 nuclei lie at the NV lattice sites
    for (uint i = 0; i < nuclei.size(); i++) {
      if ((nuclei.at(i) == n_pos) || (nuclei.at(i) == e_pos)) {
        cout << "The input lattice places a C-13 nucleus at one of the NV lattice sites!"
             << endl;
        return -1;
      }
    }
  }

  // -------------------------------------------------------------------------------------
  // Characterize nuclei and targets
  // -------------------------------------------------------------------------------------

  // identify nuclei which cannot be addressed
  vector<uint> addressable_nuclei;
  vector<uint> unaddressable_nuclei;
  for (uint n = 0; n < nuclei.size(); n++) {
    if (can_address(nuclei,n)) {
      addressable_nuclei.push_back(n);
    } else {
      unaddressable_nuclei.push_back(n);
    }
  }
  if (unaddressable_nuclei.size() > 0) {
    cout << "The following nuclei cannot be addressed:";
    for (uint n: unaddressable_nuclei) {
      cout << " " << n;
    }
    cout << endl << endl;
  }

  // identify larmor pairs
  vector<uint> larmor_nuclei;
  vector<vector<uint>> larmor_pairs;
  for (uint i = 0; i < addressable_nuclei.size(); i++) {
    const uint n_i = addressable_nuclei.at(i);
    for (uint j = i+1; j < addressable_nuclei.size(); j++) {
      const uint n_j = addressable_nuclei.at(j);
      if (is_larmor_pair(nuclei,n_i,n_j)) {
        larmor_nuclei.push_back(n_i);
        larmor_nuclei.push_back(n_j);
        larmor_pairs.push_back({n_i,n_j});
        continue;
      }
    }
  }

  // if we are only doing a search for isolated larmor pairs, do it now
  if (pair_search) {
    cout << "Starting search for isolated larmor pairs..." << endl;
    cout << "idx1 idx2 hyperfine_xy" << endl;
    for (vector<uint> larmor_pair: larmor_pairs) {
      const uint ln_1 = larmor_pair.at(0);
      const uint ln_2 = larmor_pair.at(1);
      const double A_z = hyperfine_z(nuclei.at(ln_1));
      const double A_xy = hyperfine_xy(nuclei.at(ln_1));

      if (A_xy <= min_hyperfine_xy) continue;

      bool isolated_pair = (strong_field_coupling(nuclei, ln_1, ln_2)
                            <= nuclear_isolation);
      if (!isolated_pair) continue;

      for (uint n = 0; n < nuclei.size(); n++) {
        if (n == ln_1 || n == ln_2) continue;
        if (strong_field_coupling(nuclei, n, ln_1) > nuclear_isolation ||
            strong_field_coupling(nuclei, n, ln_2) > nuclear_isolation ||
            (abs(hyperfine_z(nuclei.at(n)) - A_z) < larmor_isolation)) {
          isolated_pair = false;
          break;
        }
      }
      if (!isolated_pair) continue;

      cout << ln_1 << " " << ln_2 << " " << A_xy << endl;
    }
    return 0;
  }

  // determine which nuclei to target
  if (!set_target_nuclei) {
    if (target_pairs) target_nuclei = larmor_nuclei;
    else target_nuclei = addressable_nuclei;

  } else { // if (set_target_nuclei)
    vector<uint> unaddressable_targets;
    for (uint t_i = 0; t_i < target_nuclei.size(); t_i++) {
      if (in_vector(target_nuclei.at(t_i),unaddressable_nuclei)) {
        unaddressable_targets.push_back(target_nuclei.at(t_i));
        target_nuclei.erase(target_nuclei.begin()+t_i);
        t_i--;
      }
    }
    if (unaddressable_targets.size() > 0) {
      cout << "(WARNING) Ignoring following target nuclei:";
      for (uint n: unaddressable_targets) {
        cout << " " << n;
      }
      cout << endl << endl;
    }
  }
  if (target_nuclei.size() == 0) {
    cout << "There are no target nuclei." << endl;
    return 0;
  }

  // identify targeted larmor pairs
  vector<vector<uint>> targeted_larmor_pairs;
  for (uint i = 0; i < target_nuclei.size(); i++) {
    const uint n_i = target_nuclei.at(i);
    for (uint j = i+1; j < target_nuclei.size(); j++) {
      const uint n_j = target_nuclei.at(j);
      if (is_larmor_pair(nuclei,n_i,n_j)) {
        targeted_larmor_pairs.push_back({n_i,n_j});
        continue;
      }
    }
  }

  if (swap_nvst || larmor_identity) {
    if (targeted_larmor_pairs.size() == 0) {
      if (larmor_pairs.size() == 0) {
        cout << "There are no larmor pairs in this system\n";
      } else {
        cout << "No larmor pairs are targeted\n";
      }
      return -1;
    }
  }

  // -------------------------------------------------------------------------------------
  // Cluster C-13 nuclei
  // -------------------------------------------------------------------------------------

  // unless we are performing a coherence scan, we will be grouping together clusters by
  //  the larmor frequencies of the nuclei, so first we check whether doing so is possible
  //  for the given min_cluster_size_cap
  const uint min_cluster_size_cap = smallest_possible_cluster_size(nuclei);
  const bool cluster_larmor_pairs =
    !(coherence_scan || (max_cluster_size < min_cluster_size_cap));
  if (!coherence_scan) {
    cout << "The minimum cluster size cap is " << min_cluster_size_cap << endl;
    if (!testing_mode && (max_cluster_size < min_cluster_size_cap)) return -1;
  }

  const vector<vector<uint>> clusters = cluster_nuclei(nuclei, max_cluster_size,
                                                       cluster_larmor_pairs);
  const double cluster_coupling = get_cluster_coupling(nuclei,clusters);

  cout << "Nuclei grouped into " << clusters.size() << " clusters"
       << " with a coupling factor of " << cluster_coupling << " Hz\n";

  // collect and print histogram of cluster sizes
  vector<uint> size_hist(largest_cluster_size(clusters));
  for (uint i = 0; i < clusters.size(); i++) {
    size_hist.at(clusters.at(i).size()-1) += 1;
  }
  cout << "Cluster size histogram:\n";
  for (uint i = 0; i < size_hist.size(); i++) {
    cout << "  " << i+1 << ": " << size_hist.at(i) << endl;
  }
  cout << endl;

  // initialize nv_system object
  const nv_system nv(nuclei, clusters, ms, g_C13*static_Bz, k_DD,
                     scale_factor, integration_factor, no_nn);

  // -------------------------------------------------------------------------------------
  // Print info about target nuclei
  // -------------------------------------------------------------------------------------

  if (target_info) {
    cout << "Target info:\n\n";
    for (uint n: target_nuclei) {
      cout << "index: " << n << endl
           << "position (nm): "
           << in_crystal_basis(nv.nuclei.at(n)).transpose() * a0/2 / nm << endl
           << "effective_larmor (kHz): " << effective_larmor(nv,n).norm() * 1e-3 << endl
           << "hyperfine (kHz): " << hyperfine(nv,n).norm() * 1e-3 << endl
           << "hyperfine_perp (kHz): " << hyperfine_perp(nv,n).norm() * 1e-3 << endl
           << endl;
    }
    for (uint i = 0; i < target_nuclei.size(); i++) {
      const Vector3d pos_i = nv.nuclei.at(target_nuclei.at(i));
      for (uint j = i + 1; j < target_nuclei.size(); j++) {
        const Vector3d pos_j = nv.nuclei.at(target_nuclei.at(j));
        cout << "indices: " << i << " " << j << endl
             << "displacement (nm): "
             << in_crystal_basis(pos_j-pos_i).transpose() * a0/2 / nm << endl
             << "coupling (Hz): " << coupling_strength(pos_i,pos_j) << endl
             << "strong field coupling (Hz): "
             <<  strong_field_coupling(pos_i,pos_j) << endl
             << endl;
      }
    }
  }

  // -------------------------------------------------------------------------------------
  // Coherence scan
  // -------------------------------------------------------------------------------------

  if (coherence_scan) {
    cout << endl
         << "Larmor and hyperfine frequency data:" << endl
         << "# format: w_larmor A_perp" << endl;

    double w_max = 0, w_min = DBL_MAX;
    for (uint i = 0; i < nv.nuclei.size(); i++) {
      const double A_perp = hyperfine_perp(nv,i).norm();
      const double w_larmor = effective_larmor(nv,i).norm();
      cout << i << " " << w_larmor << " " << A_perp << endl;

      if (w_larmor < w_min) w_min = w_larmor;
      if (w_larmor > w_max) w_max = w_larmor;
    }

    cout << endl
         << "Coherence scan results:" << endl
         << "# format: w_scan coherence" << endl;
    const double w_range = max(w_max - w_min, 1e5);
    const double w_start = max(w_min - w_range/10, 0.);
    const double w_end = w_max + w_range/10;
    for (uint i = 0; i < coherence_bins; i++) {
      const double w_scan = w_start + (i+0.5)*(w_end-w_start)/coherence_bins;
      const double coherence = coherence_measurement(nv, w_scan, f_DD, scan_time);
      cout << w_scan << " " << coherence << endl;
    }

  }

  // -------------------------------------------------------------------------------------
  // Coherence signal
  // -------------------------------------------------------------------------------------

  if (coherence_signal || angular_coherence_signal) {
    cout << "Coherence signal results:" << endl
         << "# k_DD: " << nv.k_DD << endl;
    cout << "# format: f_DD coherence(s)" << endl;

    for (uint target: target_nuclei) {
      // only perform this measurement once for larmor pairs
      vector<uint> targets = {target};
      for (uint nn: target_nuclei) {
        if (nn == target) continue;
        if (is_larmor_pair(nv,nn,target)) targets.push_back(nn);
      }
      if (targets.size() == 2 && targets.at(0) > targets.at(1)) continue;

      cout << "# targets:";
      for (uint nn: targets) cout << " " << nn;
      cout << endl;
      if (angular_coherence_signal) {
        cout << "# azimuths/pi:";
        for (uint nn: targets) {
          const double angular_pos = atan2(dot(nv.nuclei.at(nn),yhat),
                                           dot(nv.nuclei.at(nn),xhat));
          cout << " " << angular_pos/pi;
        }
        cout << endl;
      }

      // determine control field to use during coherence signal measurement
      const double w_signal = effective_larmor(nv,target).norm();
      const double phi_dec = [&]() -> double {
        if (!angular_coherence_signal) return angle;
        else return 0;
      }();
      const control_fields controls(nv.static_gBz/signal_gB_factor*xhat,
                                    w_signal, phi_dec);
      cout << "# phi_dec/pi: " << phi_dec/pi << endl;

      // fix coherence signal measurement time
      const double measurement_time = 4*pi / (axy_f_max(nv.k_DD) * max_f_factor
                                              * hyperfine_perp(nv,target).norm() / 4);
      cout << "# measurement_time (ms): " << measurement_time * 1e3 << endl;

      for (uint ii = 0; ii < coherence_bins; ii++) {
        const double f_DD = (ii+0.5)/coherence_bins * axy_f_max(nv.k_DD) * max_f_factor;

        if (!angular_coherence_signal) {
          const double coherence =
            coherence_measurement(nv, w_signal, f_DD, measurement_time, controls);
          cout << f_DD << " " << coherence << endl;

        } else {
          cout << f_DD;
          for (uint jj = 0; jj < angular_resolution; jj++) {
            const double phi_DD = (jj+0.5)/angular_resolution * pi;
            const double coherence =
              coherence_measurement(nv, w_signal, f_DD, measurement_time,
                                    controls, phi_DD);
            cout << " " << coherence;
          }
          cout << endl;
        }
      }
    }
  }

  // -------------------------------------------------------------------------------------
  // Individual addressing -- control
  // -------------------------------------------------------------------------------------

  if (rotation) {
    cout << "target fidelity time pulses\n";
    for (uint target: target_nuclei) {
      const Vector3d target_axis = axis(target_polar,target_azimuth);
      vector<protocol> P(2);
      for (bool exact : {true,false}) {
        P.at(exact) = rotate_target(nv, target, angle, target_axis, exact);
      }
      const uint subsystem_target = get_index_in_subsystem(nv, target);
      cout << target << " "
           << protocol_fidelity(P, {subsystem_target}) << " "
           << P.at(false).time << " "
           << P.at(false).pulses << endl;
    }
  }

  // -------------------------------------------------------------------------------------
  // Individual addressing -- NV coupling
  // -------------------------------------------------------------------------------------

  if (coupling) {
    cout << "target fidelity time pulses\n";
    const Vector3d nv_axis = axis(nv_polar, nv_azimuth);
    for (uint target: target_nuclei) {
      const Vector3d target_axis = axis(target_polar, target_azimuth);
      vector<protocol> P(2);
      for (bool exact : {true,false}) {
        P.at(exact) = couple_target(nv, target, angle, nv_axis, target_axis, exact);
      }
      const uint subsystem_target = get_index_in_subsystem(nv, target);
      cout << target << " "
           << protocol_fidelity(P, {0, subsystem_target}) << " "
           << P.at(false).time << " "
           << P.at(false).pulses << endl;
    }
  }

  // -------------------------------------------------------------------------------------
  // NV/nucleus iSWAP fidelity
  // -------------------------------------------------------------------------------------

  if (iswap) {
    cout << "target fidelity time pulses\n";
    for (uint target: target_nuclei) {
      vector<protocol> P(2);
      for (bool exact : {true,false}) {
        P.at(exact) = iSWAP(nv, target, exact);
      }
      const uint subsystem_target = get_index_in_subsystem(nv, target);
      cout << target << " "
           << protocol_fidelity(P, {0, subsystem_target}) << " "
           << P.at(false).time << " "
           << P.at(false).pulses << endl;
    }
  }

  // -------------------------------------------------------------------------------------
  // SWAP fidelity: NV electron spin and single nuclear spin
  // -------------------------------------------------------------------------------------

  if (swap) {
    cout << "target fidelity time pulses\n";
    for (uint target: target_nuclei) {
      vector<protocol> P(2);
      for (bool exact : {true,false}) {
        P.at(exact) = SWAP(nv, target, exact);
      }
      const uint subsystem_target = get_index_in_subsystem(nv, target);
      cout << target << " "
           << protocol_fidelity(P, {0, subsystem_target}) << " "
           << P.at(false).time << " "
           << P.at(false).pulses << endl;
    }
  }

  // -------------------------------------------------------------------------------------
  // Fidelity of identity operation on a single spin
  // -------------------------------------------------------------------------------------

  if (identity) {
    cout << "target fidelity\n";
    for (uint target: target_nuclei) {
      vector<protocol> P(2);
      for (bool exact : {true,false}) {
        P.at(exact) = target_identity(nv, target, identity_time, exact);
      }
      const uint cluster_target = get_index_in_cluster(nv, target);
      cout << target << " "
           << protocol_fidelity(P, {cluster_target}) << endl;
    }
  }

  // -------------------------------------------------------------------------------------
  // SWAP fidelity: NV electron spin and singlet-triplet subspace of two nuclear spins
  // -------------------------------------------------------------------------------------

  if (swap_nvst) {
    cout << "idx1 idx2 fidelity time pulses\n";
    for (vector<uint> idxs: targeted_larmor_pairs) {
      const uint idx1 = idxs.at(0);
      const uint idx2 = idxs.at(1);
      vector<protocol> P(2);
      for (bool exact : {true,false}) {
        P.at(exact) = SWAP_NVST(nv, idx1, idx2, exact);
      }
      const uint ss_idx1 = get_index_in_subsystem(nv,idx1);
      const uint ss_idx2 = get_index_in_subsystem(nv,idx2);
      cout << idx1 << " " << idx2 << " "
           << protocol_fidelity(P, {0, ss_idx1, ss_idx2}) << " "
           << P.at(false).time << " "
           << P.at(false).pulses << endl;
    }
  }

  // -------------------------------------------------------------------------------------
  // Fidelity of identity operation on a larmor qubit
  // -------------------------------------------------------------------------------------

  if (larmor_identity) {
    cout << "idx1 idx2 fidelity\n";
    const bool targeting_pair = true;
    for (vector<uint> idxs: targeted_larmor_pairs) {
      const uint idx1 = idxs.at(0);
      const uint idx2 = idxs.at(1);
      vector<protocol> P(2);
      for (bool exact : {true,false}) {
        P.at(exact) = target_identity(nv, idx1, identity_time, exact, targeting_pair);
      }
      const uint cluster_idx1 = get_index_in_cluster(nv,idx1);
      const uint cluster_idx2 = get_index_in_cluster(nv,idx2);
      cout << idx1 << " " << idx2 << " "
           << protocol_fidelity(P, {cluster_idx1, cluster_idx2}) << endl;
    }
  }

  // -------------------------------------------------------------------------------------
  // Fidelity of deterministically initializing a thermalized nucleus into |u> or |d>
  // -------------------------------------------------------------------------------------

  if (initialize) {
    cout << "target fidelity time pulses\n";
    for (uint target: target_nuclei) {
      vector<protocol> P(2);
      for (bool exact : {true,false}) {
        P.at(exact) = initialize_spin(nv, target, exact);
      }
      const uint subsystem_target = get_index_in_subsystem(nv, target);
      cout << target << " "
           << protocol_fidelity(P, {0, subsystem_target}) << " "
           << P.at(false).time << " "
           << P.at(false).pulses << endl;
    }
  }

  // -------------------------------------------------------------------------------------
  // Fidelity of probabalistically initializing a thermalized nucleus into |u> +/- |d>
  // -------------------------------------------------------------------------------------

  if (initialize_x) {
    cout << "target fidelity time pulses\n";
    for (uint target: target_nuclei) {
      vector<protocol> P(2);
      for (bool exact : {true,false}) {
        P.at(exact) = initialize_spin_X(nv, target, exact);
      }
      const uint subsystem_target = get_index_in_subsystem(nv, target);
      cout << target << " "
           << protocol_fidelity(P, {0, subsystem_target}) << " "
           << P.at(false).time << " "
           << P.at(false).pulses << endl;
    }
  }

  // -------------------------------------------------------------------------------------
  // Fidelity of probabalistically initializing a larmor pair from |dd> into |ud> +/- |du>
  // -------------------------------------------------------------------------------------

  if (initialize_larmor) {
    cout << "idx1 idx2 fidelity time pulses\n";
    for (vector<uint> idxs: targeted_larmor_pairs) {
      const uint idx1 = idxs.at(0);
      const uint idx2 = idxs.at(1);
      vector<protocol> P(2);
      for (bool exact : {true,false}) {
        P.at(exact) = initialize_larmor_qubit(nv, idx1, idx2, exact);
      }
      const uint ss_idx1 = get_index_in_subsystem(nv,idx1);
      const uint ss_idx2 = get_index_in_subsystem(nv,idx2);
      cout << idx1 << " " << idx2 << " "
           << protocol_fidelity(P, {0, ss_idx1, ss_idx2}) << " "
           << P.at(false).time << " "
           << P.at(false).pulses << endl;
    }
  }

}


#include <boost/program_options.hpp>
#include <frovedis.hpp>
#include <frovedis/ml/clustering/gmm.hpp>

using namespace boost;
using namespace frovedis;
using namespace std;

template <class T>
void do_gmm(const string& input, const string& output, int k,
            const string& cov_type, const string& init_params,
            int n_init, int num_iteration,
            double eps, long seed, bool binary) {
  rowmajor_matrix<T> mat;
  time_spent t(DEBUG);
  if (binary) {
    mat = make_rowmajor_matrix_loadbinary<T>(input);
    t.show("load matrix: ");
  } else {
    mat = make_rowmajor_matrix_load<T>(input);
    t.show("load matrix: ");
  }
  auto gmm_model = frovedis::gaussian_mixture<T>(k, cov_type,
                   eps, num_iteration, init_params, seed);
  gmm_model.fit(mat);
  t.show("train: ");
  LOG(DEBUG) << "number of loops until convergence: " << gmm_model.n_iter_()
             << std::endl;
  LOG(DEBUG) << "likelihood: " << gmm_model.lower_bound_() << std::endl;
  binary ? gmm_model.savebinary(output) : gmm_model.save(output);
  t.show("model save: ");
}

template <class T>
void do_assign(const string& input, const string& input_cluster,
               const string& output, int k, bool binary) {
  auto gmm_model = frovedis::gaussian_mixture<T>(k);
  if (!directory_exists(output)) make_directory(output);
  if (binary) {
    time_spent t(DEBUG);
    auto mat = make_rowmajor_matrix_loadbinary<T>(input);
    t.show("load matrix: ");
    gmm_model.loadbinary(input_cluster);
    t.show("load model: ");
    auto pred = gmm_model.predict(mat);
    auto prob = gmm_model.predict_proba(mat);
    t.show("prediction time: ");
    auto score = gmm_model.score(mat);
    std::cout << "score: " << score << std::endl;
    make_dvector_scatter(pred).savebinary(output + "/prediction");
    prob.savebinary(output + "/probability");
    t.show("prediction save: ");
  } else {
    time_spent t(DEBUG);
    auto mat = make_rowmajor_matrix_load<T>(input);
    t.show("load matrix: ");
    gmm_model.load(input_cluster);
    t.show("load model: ");
    auto pred = gmm_model.predict(mat);
    auto prob = gmm_model.predict_proba(mat);
    t.show("prediction time: ");
    auto score = gmm_model.score(mat);
    std::cout << "score: " << score << std::endl;
    make_dvector_scatter(pred).saveline(output + "/predict");
    prob.save(output + "/predict_prob");
    t.show("prediction save: ");
  }
}

int main(int argc, char* argv[]) {
  use_frovedis use(argc, argv);

  using namespace boost::program_options;

  options_description opt("option");
  // clang-format off
  opt.add_options()("help,h", "print help")
    ("assign,a", "assign data to cluster mode")
    ("input,i", value<string>(),"input matrix")
    ("cluster,c", value<string>(), "input model,cov,pi for assignment")
    ("output,o", value<string>(), "output centroids or cluster")
    ("k,k", value<int>(), "number of clusters")
    ("cov-type,v", value<string>(), "covariance type (default: full)")
    ("num-init,t", value<int>(),"number of time for running with different centroid seeds (default: 1)")
    ("num-iteration,n", value<int>(),"maximum number of iteration (default: 300)")
    ("eps,e", value<double>(),"epsilon to stop the iteration (default: 0.001)")
    ("init-params,p", value<string>(), "initialization method [kmeans or random] (default: kmeans)")
    ("seed,r", value<long>(), "seed for init randomizer (default: 123)")
    ("float", "for float type input")
    ("double","for double type input (default)")
    ("verbose", "set loglevel to DEBUG")
    ("verbose2", "set loglevel to TRACE")
    ("binary,b", "use binary input/output");
  // clang-format on
  variables_map argmap;
  store(command_line_parser(argc, argv).options(opt).allow_unregistered().run(),
        argmap);
  notify(argmap);

  string input, output, input_cluster;
  string cov_type = "full", init_params = "kmeans";
  int k = 0;
  int num_iteration = 300, n_init = 1;
  double eps = 0.001;
  long seed = 123;
  bool assign = false;
  bool binary = false;

  if (argmap.count("help")) {
    cerr << opt << endl;
    exit(1);
  }

  if (argmap.count("input")) {
    input = argmap["input"].as<string>();
  } else {
    cerr << "input is not specified" << endl;
    cerr << opt << endl;
    exit(1);
  }

  if (argmap.count("output")) {
    output = argmap["output"].as<string>();
  } else {
    cerr << "output is not specified" << endl;
    cerr << opt << endl;
    exit(1);
  }

  if (argmap.count("assign")) {
    assign = true;
  }

  if (argmap.count("cluster")) {
    input_cluster = argmap["cluster"].as<string>();
  } else {
    if (assign == true) {
      cerr << "cluster is not specified" << endl;
      cerr << opt << endl;
      exit(1);
    }
  }

  if (argmap.count("k")) {
    k = argmap["k"].as<int>();
  } else {
    if (assign == false) {
      cerr << "number of cluster is not specified" << endl;
      cerr << opt << endl;
      exit(1);
    }
  }

  if(argmap.count("cov-type")) {
    cov_type = argmap["cov-type"].as<string>();
  }

  if(argmap.count("init-params")){
    init_params = argmap["init-params"].as<string>();
  }

  if (argmap.count("num-iteration")) {
    num_iteration = argmap["num-iteration"].as<int>();
  }

  if (argmap.count("num-init")) {
    n_init = argmap["num-init"].as<int>();
  }

  if (argmap.count("epsilon")) {
    eps = argmap["epsilon"].as<double>();
  }

  if (argmap.count("seed")) {
    seed = argmap["seed"].as<long>();
  }
  if (argmap.count("binary")) {
    binary = true;
  }
  if (argmap.count("verbose")) {
    set_loglevel(DEBUG);
  }
  if (argmap.count("verbose2")) {
    set_loglevel(TRACE);
  }

  if (assign) {
    if(argmap.count("double"))
      do_assign<double>(input, input_cluster, output, k, binary);
    else
      do_assign<float>(input, input_cluster, output, k, binary);
  }
  else {
    if(argmap.count("double"))
      do_gmm<double>(input, output, k, cov_type, init_params,
                     n_init, num_iteration, eps, seed, binary);
    else
      do_gmm<float>(input, output, k, cov_type, init_params,
                    n_init, num_iteration, eps, seed, binary);
  }
}


/**
 * @file systematic.hpp
 * @author Vahid Bastani
 *
 * Systematic resampling method
 */
#ifndef SSMPACK_FILTER_RESAMPLER_SYSTEMATIC
#define SSMPACK_FILTER_RESAMPLER_SYSTEMATIC

#include "ssmkit/filter/resampler/base.hpp"
#include "ssmkit/random/generator.hpp"

#include <armadillo>

#include <random>

namespace ssmkit {
namespace filter {
namespace resampler {

/** Implements systematic resampling method
 */
template <class Criterion>
class Systematic : public BaseResampler<Systematic<Criterion>> {
  friend class BaseResampler<Systematic<Criterion>>;

 private:
  std::uniform_real_distribution<double> uniform_;

 protected:
  arma::vec generateOrderedNumbers(const int &num_par) {
    double u0 = uniform_(random::Generator::get().getGenerator());
    arma::vec u(num_par);
    int k = 0;
    u.imbue([&u0, &num_par, &k]() { return (k++ + u0) / num_par; });
    return u;
  }

 public:
  Systematic(Criterion criterion)
      : BaseResampler<Systematic<Criterion>>(criterion) {}
};

template<class Criterion>
Systematic<Criterion> makeSystematic(Criterion criterion){
  return Systematic<Criterion>(criterion);
}

} // namespace resampler
} // namespace filter
} // namespace ssmkit
#endif // SSMPACK_FILTER_RESAMPLER_SYSTEMATIC