/* ----------------------------------------------------------------------- *//**
 *
 * @file multilogistic.cpp
 *
 * @brief Multinomial Logistic-Regression functions
 *
 * We implement the iteratively-reweighted-least-squares method.
 *
 *//* ----------------------------------------------------------------------- */

#include <dbconnector/dbconnector.hpp>

#include <modules/shared/HandleTraits.hpp>
#include <modules/prob/boost.hpp>
#include <boost/math/distributions.hpp>
#include <modules/prob/student.hpp>
#include "multilogistic.hpp"

#include <vector>
#include <fstream>

namespace madlib {

// Use Eigen
using namespace dbal::eigen_integration;

namespace modules {

// Import names from other MADlib modules
using dbal::NoSolutionFoundException;

namespace regress {

// Internal functions
AnyType mLogstateToResult(
    const Allocator &inAllocator,
    int ref_category,
    const HandleMap<const ColumnVector, TransparentHandle<double> >& inCoef,
    const ColumnVector &diagonal_of_heissian,
    double logLikelihood,
    double conditionNo);

/**
 * @brief Logistic function
 */
inline double sigma(double x) {
    return 1. / (1. + std::exp(-x));
}

/**
 * @brief Inter- and intra-iteration state for iteratively-reweighted-least-
 *        squares method for logistic regression
 *
 * TransitionState encapsualtes the transition state during the
 * logistic-regression aggregate function. To the database, the state is
 * exposed as a single DOUBLE PRECISION array, to the C++ code it is a proper
 * object containing scalars, a vector, and a matrix.
 *
 * Note: We assume that the DOUBLE PRECISION array is initialized by the
 * database with length at least 4, and all elemenets are 0.
 */
template <class Handle>
class MLogRegrIRLSTransitionState {

    // By §14.5.3/9: "Friend declarations shall not declare partial
    // specializations." We do access protected members in operator+=().
    template <class OtherHandle>
    friend class MLogRegrIRLSTransitionState;

public:

    MLogRegrIRLSTransitionState(const AnyType &inArray)
        : mStorage(inArray.getAs<Handle>()) {
        rebind(static_cast<uint16_t>(mStorage[0]),static_cast<uint16_t>(mStorage[1]));
    }

    /**
     * @brief Convert to backend representation
     *
     * We define this function so that we can use State in the
     * argument list and as a return type.
     */
    inline operator AnyType() const {
        return mStorage;
    }

    /**
     * @brief Initialize the iteratively-reweighted-least-squares state.
     *
     * This function is only called for the first iteration, for the first row.
     */
    inline void initialize(
        const Allocator &inAllocator,
        uint16_t inWidthOfX,
        uint16_t inNumCategories, uint16_t inRefCategory) {

        mStorage = inAllocator.allocateArray<double, dbal::AggregateContext,
            dbal::DoZero, dbal::ThrowBadAlloc>(arraySize(inWidthOfX, inNumCategories));
        rebind(inWidthOfX, inNumCategories);
        widthOfX = inWidthOfX;
        numCategories = inNumCategories;
        ref_category = inRefCategory;
    }

    /**
     * @brief We need to support assigning the previous state
     */
    template <class OtherHandle> MLogRegrIRLSTransitionState &operator=(
        const MLogRegrIRLSTransitionState<OtherHandle> &inOtherState) {

        for (size_t i = 0; i < mStorage.size(); i++)
            mStorage[i] = inOtherState.mStorage[i];
        return *this;
    }

    /**
     * @brief Merge with another State object by copying the intra-iteration
     *     fields
     */
    template <class OtherHandle> MLogRegrIRLSTransitionState &operator+=(
        const MLogRegrIRLSTransitionState<OtherHandle> &inOtherState) {
        if (mStorage.size() != inOtherState.mStorage.size() ||
            widthOfX != inOtherState.widthOfX)
            throw std::logic_error("Internal error: Incompatible transition "
                "states");

        numRows += inOtherState.numRows;
        gradient += inOtherState.gradient;
        X_transp_AX += inOtherState.X_transp_AX;
        logLikelihood += inOtherState.logLikelihood;
        return *this;
    }

    /**
     * @brief Reset the inter-iteration fields.
     */
    inline void reset() {
        numRows = 0;
        gradient.fill(0);
        X_transp_AX.fill(0);
        logLikelihood = 0;
    }

private:
    static inline uint32_t arraySize(const uint16_t inWidthOfX,
        const uint16_t inNumCategories) {
        return 5 + inWidthOfX * inWidthOfX * inNumCategories * inNumCategories
                                 + 2 * inWidthOfX * inNumCategories;
    }

    /**
     * @brief Rebind to a new storage array
     *
     * @param inWidthOfX             The number of independent variables.
     * @param inNumCategories The number of categories of the dependant var

     * Array layout (iteration refers to one aggregate-function call):
     * Inter-iteration components (updated in final function):
     * - 0: widthOfX (number of independant variables)
     * - 1: numCategories (number of categories)
     * - 2: coef (vector of coefficients)
     *
     * Intra-iteration components (updated in transition step):
     * - 2 + widthOfX*numCategories: numRows (number of rows already processed in this iteration)
     * - 3 + widthOfX*numCategories: gradient (X^T A z)
     * - 3 + 2 * widthOfX * inNumCategories: X_transp_AX (X^T A X)
     * - 3 + widthOfX^2*numCategories^2
                         + 2 * widthOfX*numCategories: logLikelihood ( ln(l(c)) )
     * - 4 + widthOfX^2*numCategories^2
                         + 2 * widthOfX*numCategories: ref_category
     */
    void rebind(uint16_t inWidthOfX = 0, uint16_t inNumCategories = 0) {
        widthOfX.rebind(&mStorage[0]);
        numCategories.rebind(&mStorage[1]);
        coef.rebind(&mStorage[2], inWidthOfX*inNumCategories);

        numRows.rebind(&mStorage[2 + inWidthOfX*inNumCategories]);

        gradient.rebind(&mStorage[3 + inWidthOfX*inNumCategories],inWidthOfX*inNumCategories);
        X_transp_AX.rebind(&mStorage[3 + 2 * inWidthOfX*inNumCategories],
            inNumCategories*inWidthOfX, inWidthOfX*inNumCategories);
        logLikelihood.rebind(&mStorage[3 +
             inNumCategories*inNumCategories*inWidthOfX*inWidthOfX
             + 2 * inWidthOfX*inNumCategories]);
        ref_category.rebind(&mStorage[4 +
             inNumCategories*inNumCategories*inWidthOfX*inWidthOfX
             + 2 * inWidthOfX*inNumCategories]);
    }

    Handle mStorage;

public:
    typename HandleTraits<Handle>::ReferenceToUInt16 widthOfX;
    typename HandleTraits<Handle>::ReferenceToUInt16 numCategories;

    typename HandleTraits<Handle>::ColumnVectorTransparentHandleMap coef;
    typename HandleTraits<Handle>::ReferenceToUInt64 numRows;
    typename HandleTraits<Handle>::ColumnVectorTransparentHandleMap gradient;
    typename HandleTraits<Handle>::MatrixTransparentHandleMap X_transp_AX;
    typename HandleTraits<Handle>::ReferenceToDouble logLikelihood;
    typename HandleTraits<Handle>::ReferenceToUInt16 ref_category;
};

/**
 * @brief Inter- and intra-iteration state for robust variance calculations
 *
 * TransitionState encapsualtes the transition state during the
 * logistic-regression robust variance calculation. To the database, the state is
 * exposed as a single DOUBLE PRECISION array, to the C++ code it is a proper
 * object containing scalars, a vector, and a matrix.
 *
 * Note: We assume that the DOUBLE PRECISION array is initialized by the
 * database with length at least 4, and all elemenets are 0.
 */
template <class Handle>
class MLogRegrRobustTransitionState {

    // By §14.5.3/9: "Friend declarations shall not declare partial
    // specializations." We do access protected members in operator+=().
    template <class OtherHandle>
    friend class MLogRegrRobustTransitionState;

public:
    MLogRegrRobustTransitionState(const AnyType &inArray)
        : mStorage(inArray.getAs<Handle>()) {
        rebind(static_cast<uint16_t>(mStorage[0]),static_cast<uint16_t>(mStorage[1]));
    }

    /**
     * @brief Convert to backend representation
     *
     * We define this function so that we can use State in the
     * argument list and as a return type.
     */
    inline operator AnyType() const {
        return mStorage;
    }

    /**
     * @brief Initialize the iteratively-reweighted-least-squares state.
     *
     * This function is only called for the first iteration, for the first row.
     */
    inline void initialize(const Allocator &inAllocator,
        uint16_t inWidthOfX, uint16_t inNumCategories, uint16_t inRefCategory) {
        mStorage = inAllocator.allocateArray<double, dbal::AggregateContext,
            dbal::DoZero, dbal::ThrowBadAlloc>(arraySize(inWidthOfX, inNumCategories));
        rebind(inWidthOfX, inNumCategories);
        widthOfX = inWidthOfX;
        numCategories = inNumCategories;
        ref_category = inRefCategory;
    }

    /**
     * @brief We need to support assigning the previous state
     */
    template <class OtherHandle> MLogRegrRobustTransitionState &operator=(
        const MLogRegrRobustTransitionState<OtherHandle> &inOtherState) {
        for (size_t i = 0; i < mStorage.size(); i++)
            mStorage[i] = inOtherState.mStorage[i];
        return *this;
    }

    /**
     * @brief Merge with another State object by copying the intra-iteration
     *     fields
     */
    template <class OtherHandle> MLogRegrRobustTransitionState &operator+=(
        const MLogRegrRobustTransitionState<OtherHandle> &inOtherState) {
        if (mStorage.size() != inOtherState.mStorage.size() ||
            widthOfX != inOtherState.widthOfX)
            throw std::logic_error("Internal error: Incompatible transition "
                "states");

        numRows += inOtherState.numRows;
        X_transp_AX += inOtherState.X_transp_AX;
        meat += inOtherState.meat;
        return *this;
    }

    /**
     * @brief Reset the inter-iteration fields.
     */
    inline void reset() {
        numRows = 0;
        meat.fill(0);
        X_transp_AX.fill(0);
    }

private:
    static inline uint32_t arraySize(const uint16_t inWidthOfX,
        const uint16_t inNumCategories) {
        return 4 + 2*inWidthOfX * inWidthOfX * inNumCategories * inNumCategories
                                 + inWidthOfX * inNumCategories;
    }

    /**
     * @brief Rebind to a new storage array
     *
     * @param inWidthOfX             The number of independent variables.
     * @param inNumCategories The number of categories of the dependant var

     * Array layout (iteration refers to one aggregate-function call):
     * Inter-iteration components (updated in final function):
     * - 0: widthOfX (number of independant variables)
     * - 1: numCategories (number of categories)
     * - 2: ref_category
     * - 3: coef (vector of coefficients)
     *
     * Intra-iteration components (updated in transition step):
     * - 3 + widthOfX*numCategories: numRows (number of rows already processed in this iteration)
     * - 4 + widthOfX * inNumCategories: X_transp_AX (X^T A X).
     * - 4 + widthOfX^2*numCategories^2: meat  (The meat matrix)
     */
    void rebind(uint16_t inWidthOfX = 0, uint16_t inNumCategories = 0) {
        widthOfX.rebind(&mStorage[0]);
        numCategories.rebind(&mStorage[1]);
        ref_category.rebind(&mStorage[2]);
        coef.rebind(&mStorage[3], inWidthOfX * inNumCategories);
        numRows.rebind(&mStorage[3 + inWidthOfX * inNumCategories]);
        X_transp_AX.rebind(&mStorage[4 + inWidthOfX * inNumCategories],
            inNumCategories * inWidthOfX, inWidthOfX * inNumCategories);
        meat.rebind(&mStorage[4 +
             inNumCategories * inNumCategories * inWidthOfX * inWidthOfX
             + inWidthOfX * inNumCategories], inWidthOfX * inNumCategories, inWidthOfX * inNumCategories);
    }

    Handle mStorage;

public:
    typename HandleTraits<Handle>::ReferenceToUInt16 widthOfX;
    typename HandleTraits<Handle>::ReferenceToUInt16 numCategories;

    typename HandleTraits<Handle>::ColumnVectorTransparentHandleMap coef;
    typename HandleTraits<Handle>::ReferenceToUInt64 numRows;
    typename HandleTraits<Handle>::MatrixTransparentHandleMap X_transp_AX;
    typename HandleTraits<Handle>::MatrixTransparentHandleMap meat;
    typename HandleTraits<Handle>::ReferenceToUInt16 ref_category;
};

/**
 * @brief IRLS Transition
 * @param args
 *
 * Arguments (Matched with PSQL wrapped)
 * - 0: Current State
 * - 1: y value (Integer)
 * - 2: numCategories (Integer)
 * - 3: ref_category (Integer)
 * - 4: X value (Column Vector)
 * - 5: Previous State

 */
AnyType
__mlogregr_irls_step_transition::run(AnyType &args) {
    MLogRegrIRLSTransitionState<MutableArrayHandle<double> > state = args[0];

    // Get x as a vector of double
    MappedColumnVector x = args[4].getAs<MappedColumnVector>();

    // Get the category & numCategories as integer
    int32_t category = args[1].getAs<int32_t>();
    // Number of categories after pivoting (We pivot around the first category)
    int32_t numCategories = (args[2].getAs<int32_t>() - 1);
    int32_t ref_category = args[3].getAs<int32_t>();

    // The following check was added with MADLIB-138.
    if (!x.is_finite())
            throw std::domain_error("Design matrix is not finite.");

    if (state.numRows == 0) {
        if (x.size() > std::numeric_limits<uint16_t>::max())
                throw std::domain_error("Number of independent variables cannot be "
                        "larger than 65535.");

        if (numCategories < 1)
                throw std::domain_error("Number of cateogires must be at least 2");

        // Init the state (requires x.size() and category.size())
        state.initialize(*this,
            static_cast<uint16_t>(x.size()) ,
            static_cast<uint16_t>(numCategories),
            static_cast<uint16_t>(ref_category));

        if (!args[5].isNull()) {
                MLogRegrIRLSTransitionState<ArrayHandle<double> >
                    previousState = args[5];
                state = previousState;
                state.reset();
        }
    }

    /*
     * This check should be done for each iteration. Only checking the first
     * run is not enough.
     */
    if (category > numCategories || category < 0)
        throw std::domain_error("Invalid category. Categories must be integer "
            "values between 0 and (number of categories - 1).");

    if (ref_category > numCategories || ref_category < 0)
        throw std::domain_error("Invalid reference category. Reference category must be integer "
            "value between 0 and (number of categories - 1).");

    // Now do the transition step
    state.numRows++;
    /*     Get y: Convert to 1/0 boolean vector
            Example: Category 4 : 0 0 0 1 0 0
            Storing it in this forms helps us get a nice closed form expression
    */

    //To pivot around the specified reference category
    ColumnVector y(numCategories);
    y.fill(0);
    if (category > ref_category) {
        y(category - 1) = 1;
    } else if (category < ref_category) {
        y(category) = 1;
    }

    /*
    Compute the parameter vector (the 'pi' vector in the documentation)
    for the data point being processed.
    Casting the coefficients into a matrix makes the calculation simple.
    */
    Matrix coef = state.coef;
    coef.resize(numCategories, state.widthOfX);

    //Store the intermediate calculations because we'll reuse them in the LLH
    ColumnVector t1 = x; //t1 is vector of size state.widthOfX
    t1 = coef*x;
    /* Note: The above 2 lines could have been written as:
        ColumnVector t1 = -coef*x;

        but this creates warnings. These warnings are somehow related to the factor
        that x is of an immutable type. The following alternative could resolve
        the warnings (although not necessary the best alternative):
    */

    ColumnVector t2 = t1.array().exp();
    double t3 = 1 + t2.sum();
    ColumnVector pi = t2/t3;

    //The gradient matrix has numCategories rows and widthOfX columns
    Matrix grad = -y*x.transpose() + pi*x.transpose();
    //We cast the gradient into a vector to make the Newton step calculations much easier.
    grad.resize(numCategories*state.widthOfX,1);

    /*
         a is a matrix of size JxJ where J is the number of categories
         a_j1j2 = -pi(j1)*(1-pi(j2))if j1 == j2
         a_j1j2 =  pi(j1)*pi(j2) if j1 != j2
    */
    Matrix a(numCategories,numCategories);
    // Compute the 'a' matrix.
    Matrix piDiag = pi.asDiagonal();
    a = pi * pi.transpose() - piDiag;
    state.gradient.noalias() += grad;

    //Start the Hessian calculations
    Matrix X_transp_AX(numCategories * state.widthOfX, numCategories * state.widthOfX);

    /*
        Again: The following 3 lines could have been written as
        Matrix XXTrans = x * x.transpose();
        but it creates warnings related to the type of x. Here is an easy fix
    */
    Matrix cv_x = x;
    Matrix XXTrans = trans(cv_x);
    XXTrans = cv_x * XXTrans;

    //Eigen doesn't supported outer-products for matrices, so we have to do our own.
    //This operation is also known as a tensor-product.
    for (int i1 = 0; i1 < state.widthOfX; i1++){
         for (int i2 = 0; i2 <state.widthOfX; i2++){
            int rowOffset = numCategories * i1;
            int colOffset = numCategories * i2;

            X_transp_AX.block(rowOffset, colOffset, numCategories,  numCategories) = XXTrans(i1,i2)*a;
        }
    }

    triangularView<Lower>(state.X_transp_AX) += X_transp_AX;

    state.logLikelihood += y.transpose()*t1 - log(t3);

    return state;

}

/**
 * @brief Perform the perliminary aggregation function: Merge transition states
     This merge step is the same as that of logistic regression.
 */
AnyType
__mlogregr_irls_step_merge_states::run(AnyType &args) {
    MLogRegrIRLSTransitionState<MutableArrayHandle<double> > stateLeft = args[0];
    MLogRegrIRLSTransitionState<ArrayHandle<double> > stateRight = args[1];

    // We first handle the trivial case where this function is called with one
    // of the states being the initial state
    if (stateLeft.numRows == 0)
        return stateRight;
    else if (stateRight.numRows == 0)
        return stateLeft;

    // Merge states together and return
    stateLeft += stateRight;
    return stateLeft;
}

/**
 * @brief Perform the logistic-regression final step
 */
AnyType
__mlogregr_irls_step_final::run(AnyType &args) {
    // We request a mutable object.
    // Depending on the backend, this might perform a deep copy.
    MLogRegrIRLSTransitionState<MutableArrayHandle<double> > state = args[0];

    // Aggregates that haven't seen any data just return Null.
    if (state.numRows == 0)
        return Null();

    // See MADLIB-138. At least on certain platforms and with certain versions,
    // LAPACK will run into an infinite loop if pinv() is called for non-finite
    // matrices. We extend the check also to the dependent variables.
    if (!state.X_transp_AX.is_finite() || !state.gradient.is_finite())
        throw NoSolutionFoundException("Over- or underflow in intermediate "
            "calulation. Input data is likely of poor numerical condition.");

    SymmetricPositiveDefiniteEigenDecomposition<Matrix> decomposition(
        -1 * state.X_transp_AX, EigenvaluesOnly, ComputePseudoInverse);

    // Precompute (X^T * A * X)^-1
    Matrix heissianInver = -1 * decomposition.pseudoInverse();

    state.coef.noalias() += heissianInver * state.gradient;

    if(!state.coef.is_finite())
        throw NoSolutionFoundException("Over- or underflow in Newton step, "
            "while updating coefficients. Input data is likely of poor "
            "numerical condition.");

    // We use the intra-iteration field gradient for storing the diagonal
    // of X^T A X, so that we don't have to recompute it in the result function.
    // Likewise, we store the condition number.
    // FIXME: This feels a bit like a hack.
    state.gradient = -1 * heissianInver.diagonal();
    state.X_transp_AX(0,0) = decomposition.conditionNo();

    return state;
}

AnyType
mlogregr_robust_step_transition::run(AnyType &args) {
    using std::endl;

	MLogRegrRobustTransitionState<MutableArrayHandle<double> > state = args[0];
    // Get x as a vector of double
    MappedColumnVector x = args[4].getAs<MappedColumnVector>();
    // Get the category & numCategories as integer
    int16_t category = args[1].getAs<int>();
    // Number of categories after pivoting (We pivot around the first category)
    int16_t numCategories = (args[2].getAs<int>() - 1);
    int32_t ref_category = args[3].getAs<int32_t>();
	MappedColumnVector coefVec = args[5].getAs<MappedColumnVector>();

    // The following check was added with MADLIB-138.
    if (!x.is_finite())
            throw std::domain_error("Design matrix is not finite.");

    if (state.numRows == 0) {
        if (x.size() > std::numeric_limits<uint16_t>::max())
                throw std::domain_error("Number of independent variables cannot be "
                        "larger than 65535.");
        if (numCategories < 1)
                throw std::domain_error("Number of cateogires must be at least 2");
        if (category > numCategories)
                throw std::domain_error("You have entered a category > numCategories"
                    "Categories must be of values {0,1... numCategories-1}");

        // Init the state (requires x.size() and category.size())
        state.initialize(*this,
            static_cast<uint16_t>(x.size()) ,
            static_cast<uint16_t>(numCategories),
            static_cast<uint16_t>(ref_category));

        state.coef = coefVec;
    }

    // Now do the transition step
    state.numRows++;
    /*
        Get y: Convert to 1/0 boolean vector
        Example: Category 4 : 0 0 0 1 0 0
        Storing it in this forms helps us get a nice closed form expression
    */

    //To pivot around the specified reference category
    ColumnVector y(numCategories);
    y.fill(0);
    if (category > ref_category) {
        y(category - 1) = 1;
    } else if (category < ref_category) {
        y(category) = 1;
    }

    /*
    Compute the parameter vector (the 'pi' vector in the documentation)
    for the data point being processed.
    Casting the coefficients into a matrix makes the calculation simple.
    */

    Matrix coef = state.coef;
    coef.resize(numCategories, state.widthOfX);

    //Store the intermediate calculations because we'll reuse them in the LLH
    ColumnVector t1 = x; //t1 is vector of size state.widthOfX
    t1 = coef*x;
    /*
        Note: The above 2 lines could have been written as:
        ColumnVector t1 = -coef*x;

        but this creates warnings. These warnings are somehow related to the factor
        that x is an immutable type.
    */

    ColumnVector t2 = t1.array().exp();
    double t3 = 1 + t2.sum();
    ColumnVector pi = t2/t3;

	//The gradient matrix has numCategories rows and widthOfX columns
    Matrix grad = -y * x.transpose() + pi * x.transpose();
    //We cast the gradient into a vector to make the math easier.
    grad.resize(numCategories * state.widthOfX, 1);

    Matrix GradGradTranspose;
    GradGradTranspose = grad * grad.transpose();
	state.meat += GradGradTranspose;

    /*
         a is a matrix of size JxJ where J is the number of categories
         a_j1j2 = -pi(j1)*(1-pi(j2))if j1 == j2
         a_j1j2 =  pi(j1)*pi(j2) if j1 != j2
    */
    // Compute the 'a' matrix.
    Matrix a(numCategories,numCategories);
    Matrix piDiag = pi.asDiagonal();
    a = pi * pi.transpose() - piDiag;

    //Start the Hessian calculations
    Matrix X_transp_AX(numCategories * state.widthOfX, numCategories * state.widthOfX);

    /*
        Again: The following 3 lines could have been written as
        Matrix XXTrans = x * x.transpose();
        but it creates warnings related to the type of x. Here is an easy fix
    */
    Matrix cv_x = x;
    Matrix XXTrans = trans(cv_x);
    XXTrans = cv_x * XXTrans;

    //Eigen doesn't supported outer-products for matrices, so we have to do our own.
    //This operation is also known as a tensor-product.
    for (int i1 = 0; i1 < state.widthOfX; i1++){
         for (int i2 = 0; i2 <state.widthOfX; i2++){
            int rowOffset = numCategories * i1;
            int colOffset = numCategories * i2;

            X_transp_AX.block(rowOffset, colOffset, numCategories,  numCategories) = XXTrans(i1, i2) * a;
        }
    }

    triangularView<Lower>(state.X_transp_AX) += X_transp_AX;

    return state;

}

/**
 * @brief Perform the perliminary aggregation function: Merge transition states
     This merge step is the same as that of logistic regression.
 */
AnyType
mlogregr_robust_step_merge_states::run(AnyType &args) {
    MLogRegrRobustTransitionState<MutableArrayHandle<double> > stateLeft = args[0];
    MLogRegrRobustTransitionState<ArrayHandle<double> > stateRight = args[1];

    // We first handle the trivial case where this function is called with one
    // of the states being the initial state
    if (stateLeft.numRows == 0)
        return stateRight;
    else if (stateRight.numRows == 0)
        return stateLeft;

    // Merge states together and return
    stateLeft += stateRight;
    return stateLeft;
}

AnyType MLrobuststateToResult(
    const Allocator &inAllocator,
    int ref_category,
    const HandleMap<const ColumnVector, TransparentHandle<double> >& inCoef,
    const ColumnVector &diagonal_of_varianceMat) {

	MutableNativeColumnVector variance(
        inAllocator.allocateArray<double>(inCoef.size()));
    MutableNativeColumnVector stdErr(
        inAllocator.allocateArray<double>(inCoef.size()));
    MutableNativeColumnVector waldZStats(
        inAllocator.allocateArray<double>(inCoef.size()));
    MutableNativeColumnVector waldPValues(
        inAllocator.allocateArray<double>(inCoef.size()));

    for (Index i = 0; i < inCoef.size(); ++i) {
        variance(i) = diagonal_of_varianceMat(i);
        stdErr(i) = std::sqrt(diagonal_of_varianceMat(i));
        waldZStats(i) = inCoef(i) / stdErr(i);
        waldPValues(i) = 2. * prob::cdf(
            prob::normal(), -std::abs(waldZStats(i)));
    }

    // Return all coefficients, standard errors, etc. in a tuple
    AnyType tuple;
    tuple <<  ref_category << inCoef << stdErr << waldZStats << waldPValues;
    return tuple;
}

/**
 * @brief Perform the logistic-regression final step
 */
AnyType
mlogregr_robust_step_final::run(AnyType &args) {
    // We request a mutable object. Depending on the backend, this might perform
    // a deep copy.
    MLogRegrRobustTransitionState<ArrayHandle<double> > state = args[0];
    // Aggregates that haven't seen any data just return Null.
    if (state.numRows == 0)
        return Null();

    // See MADLIB-138. At least on certain platforms and with certain versions,
    // LAPACK will run into an infinite loop if pinv() is called for non-finite
    // matrices. We extend the check also to the dependent variables.
    if (!state.X_transp_AX.is_finite())
        throw NoSolutionFoundException("Over- or underflow in intermediate "
            "calulation. Input data is likely of poor numerical condition.");

    SymmetricPositiveDefiniteEigenDecomposition<Matrix> decomposition(
        -1 * state.X_transp_AX, EigenvaluesOnly, ComputePseudoInverse);

	for(int i = 0; i < state.X_transp_AX.rows(); i++)
	{
		for(int j = 0; j < state.X_transp_AX.cols(); j++)
		{
			elog(INFO, "Bread %i, %i, %f",  i,j, static_cast<float>(state.X_transp_AX(i,j)));
		}
	}

    // Precompute (X^T * A * X)^-1
    Matrix bread = decomposition.pseudoInverse();
	Matrix varianceMat;
    varianceMat = bread * state.meat * bread;

    if(!state.coef.is_finite())
        throw NoSolutionFoundException("Over- or underflow in Newton step, "
            "while updating coefficients. Input data is likely of poor "
            "numerical condition.");

    return MLrobuststateToResult(*this, state.ref_category, state.coef, varianceMat.diagonal());
}

/**
 * @brief Return the difference in log-likelihood between two states
 */
AnyType
__internal_mlogregr_irls_step_distance::run(AnyType &args) {
    MLogRegrIRLSTransitionState<ArrayHandle<double> > stateLeft = args[0];
    MLogRegrIRLSTransitionState<ArrayHandle<double> > stateRight = args[1];

    return std::abs(stateLeft.logLikelihood - stateRight.logLikelihood);
}

/**
 * @brief Return the coefficients and diagnostic statistics of the state
 */
AnyType
__internal_mlogregr_irls_result::run(AnyType &args) {
    MLogRegrIRLSTransitionState<ArrayHandle<double> > state = args[0];

    return mLogstateToResult(*this, state.ref_category, state.coef,
        state.gradient, state.logLikelihood, state.X_transp_AX(0,0));
}

/**
 * @brief Compute the diagnostic statistics
 *
 * This function wraps the common parts of computing the results for IRLS.
 */
AnyType mLogstateToResult(
    const Allocator &inAllocator,
    int ref_category,
    const HandleMap<const ColumnVector, TransparentHandle<double> > &inCoef,
    const ColumnVector &diagonal_of_heissian,
    double logLikelihood,
    double conditionNo) {

    MutableNativeColumnVector stdErr(
        inAllocator.allocateArray<double>(inCoef.size()));
    MutableNativeColumnVector waldZStats(
        inAllocator.allocateArray<double>(inCoef.size()));
    MutableNativeColumnVector waldPValues(
        inAllocator.allocateArray<double>(inCoef.size()));
    MutableNativeColumnVector oddsRatios(
        inAllocator.allocateArray<double>(inCoef.size()));

    for (Index i = 0; i < inCoef.size(); ++i) {
        stdErr(i) = std::sqrt(diagonal_of_heissian(i));
        waldZStats(i) = inCoef(i) / stdErr(i);
        waldPValues(i) = 2. * prob::cdf( prob::normal(),
            -std::abs(waldZStats(i)));
        oddsRatios(i) = std::exp( inCoef(i) );
    }

    // Return all coefficients, standard errors, etc. in a tuple
    AnyType tuple;
    tuple << ref_category << inCoef << logLikelihood << stdErr << waldZStats << waldPValues
        << oddsRatios << conditionNo;
    return tuple;
}

// ---------------------------------------------------------------------------
//             Marginal Effects Multi-Logistic Regression States
// ---------------------------------------------------------------------------
/**
 * @brief State for marginal effects calculation for logistic regression
 *
 * TransitionState encapsualtes the transition state during the
 * marginal effects calculation for the logistic-regression aggregate function.
 * To the database, the state is exposed as a single DOUBLE PRECISION array,
 * to the C++ code it is a proper object containing scalars and vectors.
 *
 * Note: We assume that the DOUBLE PRECISION array is initialized by the
 * database with length at least 5, and all elemenets are 0.
 *
 */

template <class Handle>
class mlogregrMarginalTransitionState {

    // By §14.5.3/9: "Friend declarations shall not declare partial
    // specializations." We do access protected members in operator+=().
    template <class OtherHandle>
    friend class mlogregrMarginalTransitionState;

public:
    mlogregrMarginalTransitionState(const AnyType &inArray)
        : mStorage(inArray.getAs<Handle>()) {
        rebind(static_cast<uint16_t>(mStorage[0]),static_cast<uint16_t>(mStorage[1]));
    }

    /**
     * @brief Convert to backend representation
     *
     * We define this function so that we can use State in the
     * argument list and as a return type.
     */
    inline operator AnyType() const {
        return mStorage;
    }

    /**
     * @brief Initialize the iteratively-reweighted-least-squares state.
     *
     * This function is only called for the first iteration, for the first row.
     */
    inline void initialize(const Allocator &inAllocator,
        uint16_t inWidthOfX, uint16_t inNumCategories, uint16_t inRefCategory) {
        mStorage = inAllocator.allocateArray<double, dbal::AggregateContext,
            dbal::DoZero, dbal::ThrowBadAlloc>(arraySize(inWidthOfX, inNumCategories));
        rebind(inWidthOfX, inNumCategories);
        widthOfX = inWidthOfX;
        numCategories = inNumCategories;
        ref_category = inRefCategory;
    }

    /**
     * @brief We need to support assigning the previous state
     */
    template <class OtherHandle> mlogregrMarginalTransitionState &operator=(
        const mlogregrMarginalTransitionState<OtherHandle> &inOtherState) {
        for (size_t i = 0; i < mStorage.size(); i++)
            mStorage[i] = inOtherState.mStorage[i];
        return *this;
    }

    /**
     * @brief Merge with another State object by copying the intra-iteration
     *     fields
     */
    template <class OtherHandle> mlogregrMarginalTransitionState &operator+=(
        const mlogregrMarginalTransitionState<OtherHandle> &inOtherState) {
        if (mStorage.size() != inOtherState.mStorage.size() ||
            widthOfX != inOtherState.widthOfX)
            throw std::logic_error("Internal error: Incompatible transition "
                "states");

        numRows += inOtherState.numRows;
        margins_matrix += inOtherState.margins_matrix;
        X_bar += inOtherState.X_bar;
        X_transp_AX += inOtherState.X_transp_AX;
        reference_margins+= inOtherState.reference_margins;
        return *this;
    }

    /**
     * @brief Reset the inter-iteration fields.
     */
    inline void reset() {
        numRows = 0;
        margins_matrix.fill(0);
        X_bar.fill(0);
        X_transp_AX.fill(0);
        reference_margins.fill(0);
    }

private:
    static inline uint32_t arraySize(const uint16_t inWidthOfX,
        const uint16_t inNumCategories) {
        return 4 + 3*inWidthOfX * inNumCategories + 1*inWidthOfX +
           inNumCategories * inWidthOfX * inWidthOfX * inNumCategories;
    }

    /**
     * @brief Rebind to a new storage array
     *
     * @param inWidthOfX             The number of independent variables.
     * @param inNumCategories The number of categories of the dependant var

     * Array layout (iteration refers to one aggregate-function call):
     * Inter-iteration components (updated in final function):
     * - 0: widthOfX (number of independant variables)
     * - 1: numCategories (number of categories)
     * - 2: ref_category
     * - 3: coef (vector of coefficients)
     *
     * Intra-iteration components (updated in transition step):
     * - 3 + widthOfX*numCategories: numRows (number of rows already processed in this iteration)
     * - 4 + widthOfX * inNumCategories: margins_matrix
     */
    void rebind(uint16_t inWidthOfX = 0, uint16_t inNumCategories = 0) {
        widthOfX.rebind(&mStorage[0]);
        numCategories.rebind(&mStorage[1]);
        ref_category.rebind(&mStorage[2]);
        coef.rebind(&mStorage[3], inWidthOfX * inNumCategories);
        numRows.rebind(&mStorage[3 + inWidthOfX * inNumCategories]);
        margins_matrix.rebind(&mStorage[4 + inWidthOfX * inNumCategories],
            inNumCategories , inWidthOfX );
        X_bar.rebind(&mStorage[4 + 2*inWidthOfX * inNumCategories], inWidthOfX);
        reference_margins.rebind(&mStorage[4 + 2*inWidthOfX + 2*inWidthOfX*inNumCategories],
            inWidthOfX);
        X_transp_AX.rebind(&mStorage[4 + 3*inWidthOfX * inNumCategories + 1*inWidthOfX],
            inNumCategories*inWidthOfX, inWidthOfX*inNumCategories);
    }

    Handle mStorage;

public:
    typename HandleTraits<Handle>::ReferenceToUInt16 widthOfX;
    typename HandleTraits<Handle>::ReferenceToUInt16 numCategories;

    typename HandleTraits<Handle>::ColumnVectorTransparentHandleMap coef;
    typename HandleTraits<Handle>::ReferenceToUInt64 numRows;
    typename HandleTraits<Handle>::MatrixTransparentHandleMap margins_matrix;
    typename HandleTraits<Handle>::MatrixTransparentHandleMap X_transp_AX;
    typename HandleTraits<Handle>::ColumnVectorTransparentHandleMap reference_margins;
    typename HandleTraits<Handle>::ReferenceToUInt16 ref_category;
    typename HandleTraits<Handle>::ColumnVectorTransparentHandleMap X_bar;
};

AnyType
mlogregr_marginal_step_transition::run(AnyType &args) {
    using std::endl;

	  mlogregrMarginalTransitionState<MutableArrayHandle<double> > state = args[0];
    // Get x as a vector of double
    MappedColumnVector x = args[4].getAs<MappedColumnVector>();
    // Get the category & numCategories as integer
    int16_t category = args[1].getAs<int>();
    // Number of categories after pivoting (We pivot around the first category)
    int16_t numCategories = (args[2].getAs<int>() - 1);
    int32_t ref_category = args[3].getAs<int32_t>();
	  MappedColumnVector coefVec = args[5].getAs<MappedColumnVector>();

    // The following check was added with MADLIB-138.
    if (!x.is_finite())
            throw std::domain_error("Design matrix is not finite.");

    if (state.numRows == 0) {
        if (x.size() > std::numeric_limits<uint16_t>::max())
                throw std::domain_error("Number of independent variables cannot be "
                        "larger than 65535.");
        if (numCategories < 1)
                throw std::domain_error("Number of cateogires must be at least 2");
        if (category > numCategories)
                throw std::domain_error("You have entered a category > numCategories"
                    "Categories must be of values {0,1... numCategories-1}");

        // Init the state (requires x.size() and category.size())
        state.initialize(*this,
            static_cast<uint16_t>(x.size()) ,
            static_cast<uint16_t>(numCategories),
            static_cast<uint16_t>(ref_category));

        state.coef = coefVec;
        state.numCategories = numCategories;
        state.ref_category = ref_category;
    }

    // Now do the transition step
    state.numRows++;
    /*
        Get y: Convert to 1/0 boolean vector
        Example: Category 4 : 0 0 0 1 0 0
        Storing it in this forms helps us get a nice closed form expression
    */

    //To pivot around the specified reference category
    ColumnVector y(numCategories);
    y.fill(0);
    if (category > ref_category) {
        y(category - 1) = 1;
    } else if (category < ref_category) {
        y(category) = 1;
    }

    //    Marginal Effect calculations
    // ----------------------------------------------------------------------
    Matrix coef = state.coef;
    coef.resize(numCategories, state.widthOfX);

    // prob is vector of size # categories
    /*
        Note: The above 2 lines could have been written as:
        ColumnVector prob = -coef*x; but this creates warnings.
        See multilog for details.
    */
    ColumnVector prob(numCategories);
    prob = coef*x;

    // Calculate the odds ratio
    prob = prob.array().exp();
    double prob_sum = prob.sum();

    prob = prob / (1 + prob_sum);

    // Reference category computations. They have been taken out of
    // the output but left in the infrastructure
    double ref_prob = 1 / (1 + prob_sum);

    Matrix probDiag = prob.asDiagonal();

    // Marginal effects (reference calculated separately)
    ColumnVector coef_trans_prob;
    coef_trans_prob = coef.transpose() * prob;
    Matrix margins_matrix = coef;
    margins_matrix.rowwise() -= coef_trans_prob.transpose();
    margins_matrix = probDiag * margins_matrix;

    //    Variance Calculations
    // ----------------------------------------------------------------------
    /*
         a is a matrix of size JxJ where J is the number of categories
         a_j1j2 = -pi(j1)*(1-pi(j2))if j1 == j2
         a_j1j2 =  pi(j1)*pi(j2) if j1 != j2
    */
    Matrix a(numCategories,numCategories);
    a = prob * prob.transpose() - probDiag;

    // Start the Hessian calculations
    Matrix X_transp_AX(numCategories * state.widthOfX, numCategories * state.widthOfX);

    /*
        Again: The following 3 lines could have been written as
        Matrix XXTrans = x * x.transpose();
        but it creates warnings related to the type of x. Here is an easy fix
    */
    Matrix cv_x = x;
    Matrix XXTrans = trans(cv_x);
    XXTrans = cv_x * XXTrans;

    //Eigen doesn't supported outer-products for matrices, so we have to do our own.
    //This operation is also known as a tensor-product.
    for (int i1 = 0; i1 < state.widthOfX; i1++){
         for (int i2 = 0; i2 <state.widthOfX; i2++){
            int rowOffset = numCategories * i1;
            int colOffset = numCategories * i2;

            X_transp_AX.block(rowOffset, colOffset, numCategories,  numCategories) = XXTrans(i1, i2) * a;
        }
    }

    triangularView<Lower>(state.X_transp_AX) += X_transp_AX;
    state.margins_matrix += margins_matrix;
    state.reference_margins += -coef_trans_prob * ref_prob;
    state.X_bar += x; // It is called X_bar but it really is the sum

    return state;

}

/**
 * @brief Perform the perliminary aggregation function: Merge transition states
     This merge step is the same as that of logistic regression.
 */
AnyType
mlogregr_marginal_step_merge_states::run(AnyType &args) {
    mlogregrMarginalTransitionState<MutableArrayHandle<double> > stateLeft = args[0];
    mlogregrMarginalTransitionState<ArrayHandle<double> > stateRight = args[1];

    // We first handle the trivial case where this function is called with one
    // of the states being the initial state
    if (stateLeft.numRows == 0)
        return stateRight;
    else if (stateRight.numRows == 0)
        return stateLeft;

    // Merge states together and return
    stateLeft += stateRight;
    return stateLeft;
}

AnyType mlogregr_marginalstateToResult(
    const Allocator &inAllocator,
    const int numRows,
    const ColumnVector &inCoef,
    const ColumnVector &inMargins,
    const ColumnVector &inVariance
    ) {

    MutableNativeColumnVector margins(
        inAllocator.allocateArray<double>(inMargins.size()));
    MutableNativeColumnVector coef(
        inAllocator.allocateArray<double>(inMargins.size()));
    MutableNativeColumnVector stdErr(
        inAllocator.allocateArray<double>(inMargins.size()));
    MutableNativeColumnVector tStats(
        inAllocator.allocateArray<double>(inMargins.size()));
    MutableNativeColumnVector pValues(
        inAllocator.allocateArray<double>(inMargins.size()));

    for (Index i = 0; i < inMargins.size(); ++i) {
        margins(i) = inMargins(i);
        coef(i) = inCoef(i);
        stdErr(i) = std::sqrt(inVariance(i));
        tStats(i) = margins(i) / stdErr(i);

        // P-values only make sense if numRows > coef.size()
        if (numRows > inCoef.size())
          pValues(i) = 2. * prob::cdf(
              boost::math::complement(
                  prob::students_t(
                      static_cast<double>(numRows - inCoef.size())
                  ),
                  std::fabs(tStats(i))
              ));
    }

    // Return all coefficients, standard errors, etc. in a tuple
    // Note: PValues will return NULL if numRows <= coef.size
    AnyType tuple;
    tuple << margins
          << coef
          << stdErr
          << tStats
        	<< (numRows > inCoef.size()? pValues: Null());
    return tuple;
}

/**
 * @brief Perform the logistic-regression final step
 */
AnyType
mlogregr_marginal_step_final::run(AnyType &args) {
    // We request a mutable object. Depending on the backend, this might perform
    // a deep copy.
    mlogregrMarginalTransitionState<ArrayHandle<double> > state = args[0];
    // Aggregates that haven't seen any data just return Null.
    if (state.numRows == 0)
        return Null();

    if(!state.coef.is_finite())
        throw NoSolutionFoundException("Over- or underflow in Newton step, "
            "while updating coefficients. Input data is likely of poor "
            "numerical condition.");

    // See MADLIB-138. At least on certain platforms and with certain versions,
    // LAPACK will run into an infinite loop if pinv() is called for non-finite
    // matrices. We extend the check also to the dependent variables.
    if (!state.X_transp_AX.is_finite())
        throw NoSolutionFoundException("Over- or underflow in intermediate "
            "calulation. Input data is likely of poor numerical condition.");

    // Include marginal effects of reference variable:
    // FIXME: They have been taken out of the output for now
    //const int size = state.coef.size() + numIndepVars;
    const int size = state.coef.size();

    // Variance-covariance calculation
    // ----------------------------------------------------------
    SymmetricPositiveDefiniteEigenDecomposition<Matrix> decomposition(
        -1 * state.X_transp_AX, EigenvaluesOnly, ComputePseudoInverse);

    // Precompute -(X^T * A * X)^-1
    Matrix V = decomposition.pseudoInverse();

    // Marginal Effect calculation
    // ---------------------------------------------------------
    Matrix margins_matrix;
    margins_matrix = state.margins_matrix / state.numRows;

    // Standard error calculation
    // ----------------------------------------------------------
    ColumnVector x_bar = state.X_bar / state.numRows;
    int numIndepVars = state.coef.size() / state.numCategories;
    int numCategories = state.numCategories;

    Matrix coef = state.coef;
    coef.resize(state.numCategories, state.widthOfX);

    // Variance & Marginal gradient
    ColumnVector marginal_gradient(size);
    ColumnVector variance(size);
    variance.setOnes();

    // Probibility vector at the mean
    ColumnVector p_bar(state.numCategories);
    ColumnVector coef_x_bar(state.numCategories);
    ColumnVector coef_trans_p_bar(state.widthOfX);

    coef_x_bar = coef*x_bar;
    coef_trans_p_bar = coef.transpose() * p_bar;

    p_bar = coef_x_bar;
    p_bar = p_bar.array().exp();
    p_bar = p_bar / (1 + p_bar.sum());

    // Marginal effects at the mean
    Matrix margins_mean_matrix(numCategories, numIndepVars);
    margins_mean_matrix.rowwise() -= coef_trans_p_bar.transpose();
    margins_mean_matrix = p_bar.asDiagonal() * margins_matrix;

    // Compute the variance for each marginal effect
    int index, e_j_J, e_k_K, e_k_K_j_J;
    for (int K=0; K < numIndepVars; K++){
      for (int J=0; J < numCategories; J++){

        for (int k=0; k < numIndepVars; k++){
          for (int j=0; j < numCategories; j++){

            e_j_J = (j==J) ? 1: 0;
            e_k_K = (k==K) ? 1: 0;
            e_k_K_j_J = (j==J && k==K) ? 1: 0;

            index = k*numCategories + j;

            marginal_gradient(index) =
                x_bar(k) * (e_j_J  - p_bar(j)) * margins_mean_matrix(J,K);

            marginal_gradient(index) += p_bar(J) *
                ( e_k_K_j_J - p_bar(j) * e_k_K - x_bar(k) * margins_mean_matrix(j,K));
          }
        }

        // NOTE: Since the earlier Variance calculations are being done by
        // stacking up the indepdent variables for each category separtely
        variance(K + numIndepVars * J) =
                          marginal_gradient.transpose() * V * marginal_gradient;

      }
    }

    // Add in reference variables to all the calculations
    // ----------------------------------------------------------
    ColumnVector coef_with_ref(size);
    ColumnVector margins_with_ref(size);

    // Vectorize the margins_matrix and add the reference variable
    for (int j=0; j < numCategories; j++){
      for (int k=0; k < numIndepVars; k++){

        index = k + numIndepVars *j;
        coef_with_ref(index) = coef(j,k);

        index = k + numIndepVars*j;
        margins_with_ref(index) = margins_matrix(j,k);
      }
    }

    return mlogregr_marginalstateToResult(*this,
                                          state.numRows,
                                          coef_with_ref,
                                          margins_with_ref,
                                          variance);
}

// ------------------------ End of Marginal ------------------------------------

} // namespace regress

} // namespace modules

} // namespace madlib


/**
 * @author     : Zhao Chonyyao (cyzhao@zju.edu.cn)
 * @date       : 2021-04-30
 * @description: elasticity finite element method problem
 * @version    : 1.0
 */
#include <memory>
#include <iomanip>
#include <boost/property_tree/ptree.hpp>

#include "Common/DEFINE_TYPE.h"
#include "Common/error.h"

// TODO: possible bad idea of having dependence to model in problem module
#include "Model/fem/elas_energy.h"
#include "Model/fem/mass_matrix.h"

#include "Problem/energy/basic_energy.h"
#include "Io/io.h"
#include "Geometry/extract_surface.imp"

#include "elas_fem_problem.h"

namespace PhysIKA {
using namespace std;
using namespace Eigen;

template <typename T>
int read_elas_mtr(const char* file, VEC<T>& Young, VEC<T>& Poi, const size_t num_cells)
{
    cout << "read_elas_mtr" << endl;
    Young.resize(num_cells);
    Poi.resize(num_cells);
    ifstream ifs(file);
    if (ifs.fail())
    {
        std::cerr << "[info] "
                  << "can not open file" << file << std::endl;
        return __LINE__;
    }
    size_t cell_id = 0;
    while (!ifs.eof())
    {
        if (cell_id == num_cells)
            break;
        ifs >> Young(cell_id) >> Poi(cell_id);
        ++cell_id;
    }
    ifs.close();
    return 0;
}

template <typename T>
elas_problem_builder<T>::elas_problem_builder(const T* x, const boost::property_tree::ptree& pt)
    : pt_(pt)
{
    if (pt.get<string>("solver_type") == "explicit")
    {
#define SEMI_IMPLICIT
    }

    //TODO: need to check exception
    const string filename = pt.get<string>("filename");
    MAT<T>       nods(1, 1);
    MatrixXi     cells(1, 1);

    const string type = pt.get<string>("type", "tet");

    if (type == "tet")
    {
        IF_ERR(exit, mesh_read_from_vtk<T, 4>(filename.c_str(), nods, cells));
    }
    else if (type == "hex")
        exit_if(mesh_read_from_vtk<T, 8>(filename.c_str(), nods, cells));
    else
    {
        // error_msg("type:<%s> is not supported.", type.c_str());
    }

    const size_t num_nods = nods.cols(), num_cells = cells.cols();
    cout << "V" << nods.rows() << " " << nods.cols() << endl
         << "T " << cells.rows() << " " << cells.cols() << endl;
    if (x != nullptr)
        nods = Map<const MAT<T>>(x, nods.rows(), nods.cols());
    cout << "Boundary Box :\n"
         << nods.rowwise().minCoeff() << endl
         << nods.rowwise().maxCoeff() << endl;

    REST_  = nods;
    cells_ = cells;

    Matrix<T, 3, 3> rot;
    rot << 0, 0, 1, 0, 1, 0, -1, 0, 0;
    cout << "rotatoin matrix is " << rot << endl;
    MAT<T> rotated_nods = rot * nods;

    // const string outdir = argv[3];
    auto phy_paras = pt.get_child("physics");
    //set mtr
    const T      rho       = phy_paras.get<T>("rho", 20);
    const T      Young     = phy_paras.get<T>("Young", 2000.0);
    const T      poi       = phy_paras.get<T>("poi", 0.3);
    const T      gravity   = phy_paras.get<T>("gravity", 9.8);
    const T      dt        = phy_paras.get<T>("dt", 0.01);
    const T      w_pos     = phy_paras.get<T>("w_pos", 1e6);
    const size_t num_frame = phy_paras.get<size_t>("num_frames", 100);

    //set mtr
    VEC<T>       Young_vec(num_cells), Poi_vec(num_cells);
    const string mtr_file = pt.get<string>("mtr_file", "");
    if (mtr_file != "")
        read_elas_mtr(mtr_file.c_str(), Young_vec, Poi_vec, num_cells);

    //read fixed points
    vector<size_t> cons(0);
    const string   cons_file_path = pt.get<string>("cons", "");
    if (cons_file_path != "")
        IF_ERR(exit, read_fixed_verts_from_csv(cons_file_path.c_str(), cons));
    cout << "constrint " << cons.size() << " points" << endl;

    //calc mass vector
    Matrix<T, -1, 1> mass_vec(num_nods);
    // calc_mass_vector<T>(nods, cells, rho, mass_vec);
    if (type == "tet")
        mass_calculator<T, 3, 4, 1, 1, basis_func, quadrature>(nods, cells, rho, mass_vec);
    else if (type == "hex")
        mass_calculator<T, 3, 8, 1, 2, basis_func, quadrature>(nods, cells, rho, mass_vec);

    cout << "build energy" << endl;
    int ELAS = 0;
    int GRAV = 1;
    int KIN  = 2;
    int POS  = 3;
    if (pt_.get<string>("solver_type") == "explicit")
        POS = 2;

    ebf_.resize(POS + 1);
    {
        const string csttt_type = phy_paras.get<string>("csttt", "linear");
        if (pt.get<bool>("rotate", false))
            nods = rotated_nods;
        if (mtr_file != "")
            gen_elas_energy_intf<T>(type, csttt_type, nods, cells, Young_vec, Poi_vec, ebf_[ELAS], &elas_intf_);
        else
            gen_elas_energy_intf<T>(type, csttt_type, nods, cells, Young, poi, ebf_[ELAS], &elas_intf_);
        nods = REST_;
        // to lowercase.
        char axis = pt.get<char>("grav_axis", 'y') | 0x20;
        if (axis > 'z' || axis < 'x')
        {
            // error_msg("grav_axis should be one of x(X), y(Y) or z(Z).");
        }
        ebf_[GRAV] = make_shared<gravity_energy<T, 3>>(num_nods, 1, gravity, mass_vec, axis);

        kinetic_ = pt.get<bool>("dynamics", true) ? make_shared<momentum<T, 3>>(nods.data(), num_nods, mass_vec, dt)
                                                  : nullptr;

        if (pt_.get<string>("solver_type") == "implicit")
            ebf_[KIN] = kinetic_;

        ebf_[POS] = make_shared<position_constraint<T, 3>>(nods.data(), num_nods, w_pos, cons);
    }

    //set constraint

    enum constraint_type
    {
        COLL
    };
    cbf_.resize(COLL + 1);
    collider_  = nullptr;
    cbf_[COLL] = collider_;

    if (pt_.get<string>("solver_type") == "explicit")
    {
        Map<Matrix<T, -1, 1>> position(REST_.data(), REST_.size());
        semi_implicit_ = make_shared<semi_implicit<T>>(dt, mass_vec, position);
    }
}

template <typename T>
std::shared_ptr<Problem<T, 3>> elas_problem_builder<T>::build_problem() const
{
    cout << "assemble energy" << endl;
    shared_ptr<Functional<T, 3>> energy;
    try
    {
        energy = build_energy_t<T, 3>(ebf_);
    }
    catch (std::exception& e)
    {
        cerr << e.what() << endl;
        exit(EXIT_FAILURE);
    }

    shared_ptr<Constraint<T>> constraint;
    cout << "assemble constraint" << endl;
    bool all_null = true;
    for (auto& c : cbf_)
        if (c != nullptr)
            all_null = false;
    if (all_null)
    {
        constraint = nullptr;
        cout << "WARNGING: No hard constraints." << endl;
    }
    else
    {
        try
        {
            constraint = build_constraint_t<T>(cbf_);
        }
        catch (std::exception& e)
        {
            cerr << e.what() << endl;
            exit(EXIT_FAILURE);
        }
    }
    exit_if(constraint != nullptr && energy->Nx() != constraint->Nx(), "energy and constraint has different dimension.");
    return make_shared<Problem<T, 3>>(energy, constraint);
}

template <typename T>
int elas_problem_builder<T>::update_problem(const T* x, const T* v)
{
    if (kinetic_ != nullptr)
        IF_ERR(return, kinetic_->update_location_and_velocity(x, v));
    if (collider_ != nullptr)
        IF_ERR(return, collider_->update(x));
    return 0;
}

template class elas_problem_builder<double>;

template class elas_problem_builder<float>;

}  // namespace PhysIKA


/* ----------------------------------------------------------------------
 *
 *                    *** Smooth Mach Dynamics ***
 *
 * This file is part of the USER-SMD package for LAMMPS.
 * Copyright (2014) Georg C. Ganzenmueller, georg.ganzenmueller@emi.fhg.de
 * Fraunhofer Ernst-Mach Institute for High-Speed Dynamics, EMI,
 * Eckerstrasse 4, D-79104 Freiburg i.Br, Germany.
 *
 * ----------------------------------------------------------------------- */

/* ----------------------------------------------------------------------
 LAMMPS - Large-scale Atomic/Molecular Massively Parallel Simulator
 http://lammps.sandia.gov, Sandia National Laboratories
 Steve Plimpton, sjplimp@sandia.gov

 Copyright (2003) Sandia Corporation.  Under the terms of Contract
 DE-AC04-94AL85000 with Sandia Corporation, the U.S. Government retains
 certain rights in this software.  This software is distributed under
 the GNU General Public License.

 See the README file in the top-level LAMMPS directory.
 ------------------------------------------------------------------------- */

#include <cstdio>
#include <cstdlib>
#include <cstring>
#include <iostream>
#include <Eigen/Eigen>
#include <Eigen/Geometry>
#include "compute_smd_tlsph_shape.h"
#include "atom.h"
#include "update.h"
#include "modify.h"
#include "comm.h"
#include "force.h"
#include "memory.h"
#include "error.h"
#include "pair.h"

using namespace Eigen;
using namespace std;
using namespace LAMMPS_NS;

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

ComputeSmdTlsphShape::ComputeSmdTlsphShape(LAMMPS *lmp, int narg, char **arg) :
                Compute(lmp, narg, arg) {
        if (narg != 3)
                error->all(FLERR, "Illegal compute smd/tlsph_strain command");

        peratom_flag = 1;
        size_peratom_cols = 7;

        nmax = 0;
        strainVector = NULL;
}

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

ComputeSmdTlsphShape::~ComputeSmdTlsphShape() {
        memory->sfree(strainVector);
}

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

void ComputeSmdTlsphShape::init() {

        int count = 0;
        for (int i = 0; i < modify->ncompute; i++)
                if (strcmp(modify->compute[i]->style, "smd/tlsph_strain") == 0)
                        count++;
        if (count > 1 && comm->me == 0)
                error->warning(FLERR, "More than one compute smd/tlsph_strain");
}

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

void ComputeSmdTlsphShape::compute_peratom() {
        double *contact_radius = atom->contact_radius;
        invoked_peratom = update->ntimestep;

        // grow vector array if necessary

        if (atom->nmax > nmax) {
                memory->destroy(strainVector);
                nmax = atom->nmax;
                memory->create(strainVector, nmax, size_peratom_cols, "strainVector");
                array_atom = strainVector;
        }

        int itmp = 0;
        Matrix3d *R = (Matrix3d *) force->pair->extract("smd/tlsph/rotation_ptr", itmp);
        if (R == NULL) {
                error->all(FLERR, "compute smd/tlsph_shape failed to access rotation array");
        }

        Matrix3d *F = (Matrix3d *) force->pair->extract("smd/tlsph/Fincr_ptr", itmp);
        if (F == NULL) {
                error->all(FLERR, "compute smd/tlsph_shape failed to access deformation gradient array");
        }

        int *mask = atom->mask;
        int nlocal = atom->nlocal;
        Matrix3d E, eye;
        eye.setIdentity();
        Quaterniond q;

        for (int i = 0; i < nlocal; i++) {
                if (mask[i] & groupbit) {

                        E = 0.5 * (F[i].transpose() * F[i] - eye); // Green-Lagrange strain
                        strainVector[i][0] = contact_radius[i] * (1.0 + E(0, 0));
                        strainVector[i][1] = contact_radius[i] * (1.0 + E(1, 1));
                        strainVector[i][2] = contact_radius[i] * (1.0 + E(2, 2));

                        q = R[i]; // convert pure rotation matrix to quaternion
                        strainVector[i][3] = q.w();
                        strainVector[i][4] = q.x();
                        strainVector[i][5] = q.y();
                        strainVector[i][6] = q.z();
                } else {
                        for (int j = 0; j < size_peratom_cols; j++) {
                                strainVector[i][j] = 0.0;
                        }
                }
        }
}

/* ----------------------------------------------------------------------
 memory usage of local atom-based array
 ------------------------------------------------------------------------- */

double ComputeSmdTlsphShape::memory_usage() {
        double bytes = size_peratom_cols * nmax * sizeof(double);
        return bytes;
}


#include <iostream>
#include <cmath>
#include <NTL/ZZ.h>

using namespace std;
using namespace NTL;

#include "iterative.h"
#include "recursive.h"

using namespace std;
using namespace NTL;
using namespace chrono;

//#define ITERATIVE

void usage_example_zp(){
    long degree = pow(2,20)-1;

    ZZ prime;
    GenPrime(prime, 128);
    ZZ_p::init(prime);

//  interpolation points:
    ZZ_p* X = new ZZ_p[degree+1];
    ZZ_p* Y = new ZZ_p[degree+1];
    for(unsigned int i=0;i<=degree; i++) {
        random(X[i]);
        random(Y[i]);
    }

    ZZ_pX P;
#ifdef ITERATIVE
    poly_interpolate_zp_iterative(degree, X, Y, P);
#else
    poly_interpolate_zp_recursive(degree, X, Y, P);
#endif

    // EVALUATE
    ZZ_p* X2 = new ZZ_p[degree+1];
    ZZ_p* Y2 = new ZZ_p[degree+1];
    for(unsigned int i=0;i<=degree; i++) {
        random(X[i]);
    }
#ifdef ITERATIVE
    poly_evaluate_zp_iterative(degree, P, X2, Y2);
#else
    poly_evaluate_zp_recursive(degree, P, X2, Y2);
#endif

    delete[] X;
    delete[] Y;
    delete[] X2;
    delete[] Y2;
}

int main() {

    usage_example_zp();

}


#define BOOST_TEST_MODULE "test_dihedral_angle_interaction"

#ifdef BOOST_TEST_DYN_LINK
#include <boost/test/unit_test.hpp>
#else
#include <boost/test/included/unit_test.hpp>
#endif

#include <mjolnir/core/BoundaryCondition.hpp>
#include <mjolnir/core/SimulatorTraits.hpp>
#include <mjolnir/forcefield/local/DihedralAngleInteraction.hpp>
#include <mjolnir/math/constants.hpp>
#include <mjolnir/forcefield/local/ClementiDihedralPotential.hpp>
#include <mjolnir/util/make_unique.hpp>

#include <random>

BOOST_AUTO_TEST_CASE(DihedralAngle_force)
{
    using traits_type         = mjolnir::SimulatorTraits<double, mjolnir::UnlimitedBoundary>;
    using real_type           = traits_type::real_type;
    using coord_type          = traits_type::coordinate_type;
    using boundary_type       = traits_type::boundary_type;
    using system_type         = mjolnir::System<traits_type>;
    using potential_type      = mjolnir::ClementiDihedralPotential<real_type>;
    using dihedral_angle_type = mjolnir::DihedralAngleInteraction<traits_type, potential_type>;

    constexpr real_type tol = 1e-7;

    const real_type k1(1e0);
    const real_type k3(1e0);
    const real_type native(mjolnir::math::constants<real_type>::pi() * 2.0 / 3.0);

    potential_type potential{k1, k3, native};
    dihedral_angle_type interaction("none", {{ {{0,1,2,3}}, potential}});

    const coord_type pos1(1e0, 0e0, 1e0);
    const coord_type pos2(0e0, 0e0, 1e0);
    const coord_type pos3(0e0, 0e0, 0e0);

    system_type sys(4, boundary_type{});

    sys.at(0).mass = 1.0;
    sys.at(1).mass = 1.0;
    sys.at(2).mass = 1.0;
    sys.at(3).mass = 1.0;
    sys.at(0).rmass = 1.0;
    sys.at(1).rmass = 1.0;
    sys.at(2).rmass = 1.0;
    sys.at(3).rmass = 1.0;

    sys.at(0).position = pos1;
    sys.at(1).position = pos2;
    sys.at(2).position = pos3;
    sys.at(3).position = coord_type(0,0,0);
    sys.at(0).velocity = coord_type(0,0,0);
    sys.at(1).velocity = coord_type(0,0,0);
    sys.at(2).velocity = coord_type(0,0,0);
    sys.at(3).velocity = coord_type(0,0,0);
    sys.at(0).force    = coord_type(0,0,0);
    sys.at(1).force    = coord_type(0,0,0);
    sys.at(2).force    = coord_type(0,0,0);
    sys.at(3).force    = coord_type(0,0,0);

    sys.at(0).name  = "X";
    sys.at(1).name  = "X";
    sys.at(2).name  = "X";
    sys.at(3).name  = "X";
    sys.at(0).group = "NONE";
    sys.at(1).group = "NONE";
    sys.at(2).group = "NONE";
    sys.at(3).group = "NONE";

    const real_type dtheta = mjolnir::math::constants<real_type>::pi() / 1800.0;
    for(int i = -1800; i < 1800; ++i)
    {
        BOOST_TEST(mjolnir::math::length(sys[0].position - pos1) == 0.0, boost::test_tools::tolerance(tol));
        BOOST_TEST(mjolnir::math::length(sys[1].position - pos2) == 0.0, boost::test_tools::tolerance(tol));
        BOOST_TEST(mjolnir::math::length(sys[2].position - pos3) == 0.0, boost::test_tools::tolerance(tol));

        BOOST_TEST(mjolnir::math::length(sys[0].velocity) == 0.0, boost::test_tools::tolerance(tol));
        BOOST_TEST(mjolnir::math::length(sys[1].velocity) == 0.0, boost::test_tools::tolerance(tol));
        BOOST_TEST(mjolnir::math::length(sys[2].velocity) == 0.0, boost::test_tools::tolerance(tol));

        sys[0].force = coord_type(0,0,0);
        sys[1].force = coord_type(0,0,0);
        sys[2].force = coord_type(0,0,0);
        sys[3].force = coord_type(0,0,0);

        const real_type theta = i * dtheta;
        const coord_type pos4(std::cos(theta), -std::sin(theta), 0e0);
        sys[3].position = pos4;

        const real_type deriv = potential.derivative(theta);
        const real_type coef = std::abs(deriv);

        interaction.calc_force(sys);

        // magnitude
        // if radius == 1e0, then force strength is equal to dV.
        BOOST_TEST(mjolnir::math::length(sys[1].position - sys[0].position) == 1e0, boost::test_tools::tolerance(tol));
        BOOST_TEST(mjolnir::math::length(sys[2].position - sys[1].position) == 1e0, boost::test_tools::tolerance(tol));
        BOOST_TEST(mjolnir::math::length(sys[3].position - sys[2].position) == 1e0, boost::test_tools::tolerance(tol));

        const real_type force_strength1 = mjolnir::math::length(sys[0].force);
        const real_type force_strength3 = mjolnir::math::length(sys[3].force);
        if(i == 1200)
        {
            BOOST_TEST(coef            == 0.0, boost::test_tools::tolerance(tol));
            BOOST_TEST(force_strength1 == 0.0, boost::test_tools::tolerance(tol));
            BOOST_TEST(force_strength3 == 0.0, boost::test_tools::tolerance(tol));
        }
        else
        {
            BOOST_TEST(coef == force_strength1, boost::test_tools::tolerance(tol));
            BOOST_TEST(coef == force_strength3, boost::test_tools::tolerance(tol));
        }

        // force applied to center particle is equal to sum of others
        const coord_type sum = sys[0].force + sys[1].force + sys[2].force + sys[3].force;
        BOOST_TEST(mjolnir::math::length(sum) == 0.0, boost::test_tools::tolerance(tol));

        // direction
        if(i == 1200) // most stable point
        {
            BOOST_TEST(mjolnir::math::length(sys[0].force) == 0.0, boost::test_tools::tolerance(tol));
            BOOST_TEST(mjolnir::math::length(sys[1].force) == 0.0, boost::test_tools::tolerance(tol));
            BOOST_TEST(mjolnir::math::length(sys[2].force) == 0.0, boost::test_tools::tolerance(tol));
            BOOST_TEST(mjolnir::math::length(sys[3].force) == 0.0, boost::test_tools::tolerance(tol));
        }
        else
        {
            // perpendicular to radius vector
            const real_type normal1 = mjolnir::math::dot_product(sys[0].force, sys[0].position - sys[1].position);
            const real_type normal4 = mjolnir::math::dot_product(sys[3].force, sys[2].position - sys[3].position);
            BOOST_TEST(normal1 == 0.0, boost::test_tools::tolerance(tol));
            BOOST_TEST(normal4 == 0.0, boost::test_tools::tolerance(tol));
        }

        // perpendicular to z axis
        BOOST_TEST(sys[0].force[2] == 0.0, boost::test_tools::tolerance(tol));
        BOOST_TEST(sys[1].force[2] == 0.0, boost::test_tools::tolerance(tol));
        BOOST_TEST(sys[2].force[2] == 0.0, boost::test_tools::tolerance(tol));
        BOOST_TEST(sys[3].force[2] == 0.0, boost::test_tools::tolerance(tol));
    }
}

BOOST_AUTO_TEST_CASE(DihedralAngleInteraction_numerical_diff)
{
    using traits_type         = mjolnir::SimulatorTraits<double, mjolnir::UnlimitedBoundary>;
    using real_type           = traits_type::real_type;
    using coord_type          = traits_type::coordinate_type;
    using boundary_type       = traits_type::boundary_type;
    using system_type         = mjolnir::System<traits_type>;
    using potential_type       = mjolnir::ClementiDihedralPotential<real_type>;
    using dihedral_angle_type = mjolnir::DihedralAngleInteraction<traits_type, potential_type>;

    const real_type k1(1e0);
    const real_type k3(1e0);
    const real_type native(mjolnir::math::constants<real_type>::pi() / 2.0);

    std::mt19937 mt(123456789);
    std::uniform_real_distribution<real_type> uni(-1.0, 1.0);

    potential_type potential{k1, k3, native};
    dihedral_angle_type interaction("none", {{ {{0,1,2,3}}, potential}});

    for(std::size_t i=0; i<100; ++i)
    {
        system_type sys(4, boundary_type{});

        sys.mass (0) = 1.0;
        sys.mass (1) = 1.0;
        sys.mass (2) = 1.0;
        sys.mass (3) = 1.0;
        sys.rmass(0) = 1.0;
        sys.rmass(1) = 1.0;
        sys.rmass(2) = 1.0;
        sys.rmass(3) = 1.0;

        sys.position(0) = coord_type(2.0 + 1e-2 * uni(mt), 0.0 + 1e-2 * uni(mt),  1.0 + 1e-2 * uni(mt));
        sys.position(1) = coord_type(1.0 + 1e-2 * uni(mt), 1.0 + 1e-2 * uni(mt),  0.0 + 1e-2 * uni(mt));
        sys.position(2) = coord_type(0.0 + 1e-2 * uni(mt), 0.0 + 1e-2 * uni(mt),  0.0 + 1e-2 * uni(mt));
        sys.position(3) = coord_type(1.0 + 1e-2 * uni(mt), 1.0 + 1e-2 * uni(mt), -1.0 + 1e-2 * uni(mt));
        sys.velocity(0) = coord_type(0.0, 0.0,  0.0);
        sys.velocity(1) = coord_type(0.0, 0.0,  0.0);
        sys.velocity(2) = coord_type(0.0, 0.0,  0.0);
        sys.velocity(3) = coord_type(0.0, 0.0,  0.0);
        sys.force   (0) = coord_type(0.0, 0.0,  0.0);
        sys.force   (1) = coord_type(0.0, 0.0,  0.0);
        sys.force   (2) = coord_type(0.0, 0.0,  0.0);
        sys.force   (3) = coord_type(0.0, 0.0,  0.0);

        const auto init = sys;

        sys.at(0).name  = "X";
        sys.at(1).name  = "X";
        sys.at(2).name  = "X";
        sys.at(3).name  = "X";
        sys.at(0).group = "NONE";
        sys.at(1).group = "NONE";
        sys.at(2).group = "NONE";
        sys.at(3).group = "NONE";

        constexpr real_type tol = 1e-4;
        constexpr real_type dr  = 1e-4;
        for(std::size_t idx=0; idx<4; ++idx)
        {
            {
                // ----------------------------------------------------------------
                // reset positions
                sys = init;

                // calc U(x-dx)
                const auto E0 = interaction.calc_energy(sys);

                mjolnir::math::X(sys.position(idx)) += dr;

                // calc F(x)
                interaction.calc_force(sys);

                mjolnir::math::X(sys.position(idx)) += dr;

                // calc U(x+dx)
                const auto E1 = interaction.calc_energy(sys);

                // central difference
                const auto dE = (E1 - E0) * 0.5;

                BOOST_TEST(-dE == dr * mjolnir::math::X(sys.force(idx)),
                           boost::test_tools::tolerance(tol));
            }
            {
                // ----------------------------------------------------------------
                // reset positions
                sys = init;

                // calc U(x-dx)
                const auto E0 = interaction.calc_energy(sys);

                mjolnir::math::Y(sys.position(idx)) += dr;

                // calc F(x)
                interaction.calc_force(sys);

                mjolnir::math::Y(sys.position(idx)) += dr;

                // calc U(x+dx)
                const auto E1 = interaction.calc_energy(sys);

                // central difference
                const auto dE = (E1 - E0) * 0.5;

                BOOST_TEST(-dE == dr * mjolnir::math::Y(sys.force(idx)),
                           boost::test_tools::tolerance(tol));
            }
            {
                // ----------------------------------------------------------------
                // reset positions
                sys = init;

                // calc U(x-dx)
                const auto E0 = interaction.calc_energy(sys);

                mjolnir::math::Z(sys.position(idx)) += dr;

                // calc F(x)
                interaction.calc_force(sys);

                mjolnir::math::Z(sys.position(idx)) += dr;

                // calc U(x+dx)
                const auto E1 = interaction.calc_energy(sys);

                // central difference
                const auto dE = (E1 - E0) * 0.5;

                BOOST_TEST(-dE == dr * mjolnir::math::Z(sys.force(idx)),
                           boost::test_tools::tolerance(tol));
            }
        }
    }
}

BOOST_AUTO_TEST_CASE(DihedralAngleInteraction_calc_force_and_energy)
{
    using traits_type         = mjolnir::SimulatorTraits<double, mjolnir::UnlimitedBoundary>;
    using real_type           = traits_type::real_type;
    using coord_type          = traits_type::coordinate_type;
    using boundary_type       = traits_type::boundary_type;
    using system_type         = mjolnir::System<traits_type>;
    using potential_type       = mjolnir::ClementiDihedralPotential<real_type>;
    using dihedral_angle_type = mjolnir::DihedralAngleInteraction<traits_type, potential_type>;

    const real_type k1(1e0);
    const real_type k3(1e0);
    const real_type native(mjolnir::math::constants<real_type>::pi() / 2.0);

    std::mt19937 mt(123456789);
    std::uniform_real_distribution<real_type> uni(-1.0, 1.0);

    potential_type potential{k1, k3, native};
    dihedral_angle_type interaction("none", {{ {{0,1,2,3}}, potential}});

    for(std::size_t i=0; i<100; ++i)
    {
        system_type sys(4, boundary_type{});

        sys.mass (0) = 1.0;
        sys.mass (1) = 1.0;
        sys.mass (2) = 1.0;
        sys.mass (3) = 1.0;
        sys.rmass(0) = 1.0;
        sys.rmass(1) = 1.0;
        sys.rmass(2) = 1.0;
        sys.rmass(3) = 1.0;

        sys.position(0) = coord_type(2.0 + 1e-2 * uni(mt), 0.0 + 1e-2 * uni(mt),  1.0 + 1e-2 * uni(mt));
        sys.position(1) = coord_type(1.0 + 1e-2 * uni(mt), 1.0 + 1e-2 * uni(mt),  0.0 + 1e-2 * uni(mt));
        sys.position(2) = coord_type(0.0 + 1e-2 * uni(mt), 0.0 + 1e-2 * uni(mt),  0.0 + 1e-2 * uni(mt));
        sys.position(3) = coord_type(1.0 + 1e-2 * uni(mt), 1.0 + 1e-2 * uni(mt), -1.0 + 1e-2 * uni(mt));
        sys.velocity(0) = coord_type(0.0, 0.0,  0.0);
        sys.velocity(1) = coord_type(0.0, 0.0,  0.0);
        sys.velocity(2) = coord_type(0.0, 0.0,  0.0);
        sys.velocity(3) = coord_type(0.0, 0.0,  0.0);
        sys.force   (0) = coord_type(0.0, 0.0,  0.0);
        sys.force   (1) = coord_type(0.0, 0.0,  0.0);
        sys.force   (2) = coord_type(0.0, 0.0,  0.0);
        sys.force   (3) = coord_type(0.0, 0.0,  0.0);

        sys.at(0).name  = "X";
        sys.at(1).name  = "X";
        sys.at(2).name  = "X";
        sys.at(3).name  = "X";
        sys.at(0).group = "NONE";
        sys.at(1).group = "NONE";
        sys.at(2).group = "NONE";
        sys.at(3).group = "NONE";

        constexpr real_type tol = 1e-4;
        auto ref_sys = sys;

        const auto energy = interaction.calc_force_and_energy(sys);
        const auto ref_energy = interaction.calc_energy(ref_sys);
        interaction.calc_force(ref_sys);
        BOOST_TEST(ref_energy == energy, boost::test_tools::tolerance(tol));

        for(std::size_t idx=0; idx<sys.size(); ++idx)
        {
            BOOST_TEST(mjolnir::math::X(sys.force(idx)) == mjolnir::math::X(ref_sys.force(idx)), boost::test_tools::tolerance(tol));
            BOOST_TEST(mjolnir::math::Y(sys.force(idx)) == mjolnir::math::Y(ref_sys.force(idx)), boost::test_tools::tolerance(tol));
            BOOST_TEST(mjolnir::math::Z(sys.force(idx)) == mjolnir::math::Z(ref_sys.force(idx)), boost::test_tools::tolerance(tol));
        }
    }
}


//==============================================================================
//         Copyright 2003 - 2011 LASMEA UMR 6602 CNRS/Univ. Clermont II
//         Copyright 2009 - 2011 LRI    UMR 8623 CNRS/Univ Paris Sud XI
//
//          Distributed under the Boost Software License, Version 1.0.
//                 See accompanying file LICENSE.txt or copy at
//                     http://www.boost.org/LICENSE_1_0.txt
//==============================================================================
#ifndef BOOST_SIMD_TOOLBOX_ARITHMETIC_FUNCTIONS_SCALAR_MOD_HPP_INCLUDED
#define BOOST_SIMD_TOOLBOX_ARITHMETIC_FUNCTIONS_SCALAR_MOD_HPP_INCLUDED
#include <boost/simd/toolbox/arithmetic/functions/mod.hpp>
#include <boost/simd/include/functions/scalar/idivfloor.hpp>
/////////////////////////////////////////////////////////////////////////////
// The mod function computes the remainder of dividing x by y.  The
// return value is x-n*y, where n is the value x / y, rounded to -inf.
/////////////////////////////////////////////////////////////////////////////

namespace boost { namespace simd { namespace ext
{
  BOOST_SIMD_FUNCTOR_IMPLEMENTATION( boost::simd::tag::mod_, tag::cpu_
                            , (A0)
                            , (scalar_< arithmetic_<A0> >)
                              (scalar_< arithmetic_<A0> >)
                            )
  {
    typedef A0 result_type;
    BOOST_SIMD_FUNCTOR_CALL_REPEAT(2)
    {
      return a1 ? a0-a1*boost::simd::idivfloor(a0,a1) : a0;
    }
  };
} } }

#endif


#include <Eigen/Dense>
#include <Eigen/Sparse>
#include <boost/bind.hpp>
#include "../library/MPC.h"
#include "../library/Consensus.h"
#include "../library/IOEigen.h"
/**************************** ROS libraries****************************/
#include <ros/ros.h>
#include <trajectory_msgs/MultiDOFJointTrajectory.h>
#include <mav_msgs/conversions.h>
#include <geometry_msgs/PointStamped.h>
#include <geometry_msgs/Twist.h>
#include <gazebo_msgs/GetModelState.h>
#include <gazebo_msgs/GetWorldProperties.h>

/************************** C++ libraries *****************************/
#include <iostream>
#include <algorithm>
#include <vector>
#include <string>
#include <fstream>

//using namespace Foap;
//using namespace IOeigen;

/**************************Functions declarations *********************/
void positionCallback(const geometry_msgs::PointStamped::ConstPtr& msg, int index);
void getCurrentPose(Eigen::MatrixXd &q);
void controlsCallback(const geometry_msgs::Twist::ConstPtr& msg, int index);
void getComputedControls(Eigen::MatrixXd &q);
bool readMatrix(std::string filename, Eigen::MatrixXd &M);

/*************************** Global variables *************************/
#define DEBUG 0

int agent = 3;

bool leader = false;

Eigen::IOFormat OctaveFmt(Eigen::StreamPrecision, 0, ", ", ";\n", "", "", "[", "]");
std::string sep = "\n----------------------------------------\n";

// path and gazebo envoronment input names
std::string path = "resource/";
std::string A_filename = path+"A.dat";
std::string d_filename = path+"d.dat";
std::string simulation_filename = path+"sim.dat";
std::string MPC_filename = path+"MPC.dat";
std::string consensus_filename = path+"consensus.dat";
bool status_file;
std::string robot_name = "hummingbird_"; // same as mav_name in launchfile
std::string obstacle_name = "obstacle";
std::string relativeEntityName = "world" ;

// output filename
std::string output_path = "output/data/";
std::string q_test_filename = output_path+"q_test.dat";//agents position
std::string q_odom_filename = output_path+"q_odom.dat";
std::string adjacency_mat_filename = output_path+"adjacency_mat.dat";
std::string laplacian_mat_filename = output_path+"laplacian_mat.dat";
std::string q0_mat_filename = output_path+"q0_mat.dat";
std::string z_mat_filename = output_path+"z_mat.dat";
std::string q_obst_filename = output_path+"q_obst.dat";// obstacles' position
std::string qp_agent_filename = output_path+"qp_agent_"+std::to_string(agent)+".dat";//agents' velocities
std::string e_consensus_filename = output_path+"e_consensus.dat";//agents' velocities
std::string e_consensus_agents_filename = output_path+"e_consensus_agents.dat";//agents' velocities
std::string qz_agents_filename = output_path+"qz_agents.dat";
std::string qp_computed_controls_filename = output_path+"qp_computed_controls.dat";

// pub & suv variables name
std::string slash("/");
std::string publisher_name = "/command/trajectory";
std::string subscriber_name = "/ground_truth/position/";
std::string topic_controls = "/controls/computed_u";

// ros msg's
std::vector<geometry_msgs::PointStamped> pos_msg;
std::vector<geometry_msgs::Twist> controls_msg;

// system dimentions
int n_robots = 0, dim = 3, n_models = 0, n_obstacles = 0;

// Gains of predefined consensus
double kf=20;

int main(int argc, char **argv){
  ros::init(argc,argv,"uav_"+std::to_string(agent)+"_control");
  ros::NodeHandle nh("~");
  std::string multi_uav_control_path;
  nh.getParam("resources_path", multi_uav_control_path);
  std::cout << multi_uav_control_path << std::endl;
    /************************ Read input paramaters *****************/
    Eigen::MatrixXd A, d_matrix;

    status_file = readMatrix(multi_uav_control_path+A_filename, A);
    std::cout<<A<<std::endl;
    if(!status_file){
        std::cout << "[ERR] Could not read adjacency matrix file " << std::endl;
        std::cout << multi_uav_control_path+A_filename << std::endl;
        return 1;
    }

    status_file =  readMatrix(multi_uav_control_path+d_filename, d_matrix);
    if(!status_file){
        std::cout << "[ERR] Could not read displacements file " << std::endl;
        return 1;
    }

    if(A.cols() != d_matrix.cols()){
      std::cout << "[ERR] Sizes do not match " << std::endl;
        return 1;
    }

    /************************** Get World Properties ******************/
    /*** Get initial position of robots and obstacles from world properties ***/
    // resize q0 & qobs
    Eigen::MatrixXd q0, qobs, qz;
    n_robots = A.cols();
    q0.resize(dim, n_robots);
    qz.resize(dim,n_robots);

    gazebo_msgs::GetModelState getModelState;
    geometry_msgs::Point pp;
    ros::ServiceClient gms_c = nh.serviceClient<gazebo_msgs::GetModelState>("/gazebo/get_model_state") ;
    ros::ServiceClient gwp_c = nh.serviceClient<gazebo_msgs::GetWorldProperties>("/gazebo/get_world_properties");
    gazebo_msgs::GetWorldProperties gwp_s;
    if (gwp_c.call(gwp_s)){
        std::vector<std::string> models = gwp_s.response.model_names;
        n_models = (int)models.size();
        ROS_INFO("Number of models: %i", n_models);
        if(n_models==1){ // only and empty world has been loaded
            ROS_ERROR("An empty world as been loaded");
            return 1;
        }
        // compare if the Adjancency matrix size is equal to the number of robots
        if(n_robots > (n_models-1)){
            ROS_ERROR("Number of robots don't match with the Adjancency matrix");
            return 1;
        }

        // resize qobs
        n_obstacles = n_models-n_robots-1;
        if(n_obstacles>0){
            qobs.resize(dim, n_obstacles);
        }

        //sort the name models to find the robots & obstacles initial positions:
        // models[0] = ground_plane, models[1, n_robots+1] = robots, models[n_robots+2, end] = obstacles
        std::sort(models.begin(), models.end());
        Eigen::VectorXd current_pose(dim); // assume a 3D pose, change dim to add more dimension: e.x: [x, y, z, yaw]
        // get robots position
        for(int i=0; i<n_robots; i++){
            getModelState.request.model_name = models[i+1];
            getModelState.request.relative_entity_name = relativeEntityName ;
            gms_c.call(getModelState);
            pp = getModelState.response.pose.position;
            current_pose << pp.x, pp.y, pp.z;
            q0.col(i) = current_pose;

        }

        // get obstacles positions
        if(n_obstacles>0){
            for(int i=0; i<n_obstacles; i++){
                getModelState.request.model_name = models[n_robots+i+1];
                getModelState.request.relative_entity_name = relativeEntityName ;
                gms_c.call(getModelState);
                pp = getModelState.response.pose.position;
                current_pose << pp.x, pp.y, 1.0; // the obstacles are fixed in z = 1.0
                qobs.col(i) = current_pose;
            }
         // save data to files
    	//IOEigen::writeMatrix(q_obst_filename, qobs);
	}

        ROS_INFO("Number of robots: %i", n_robots);
        ROS_INFO("Number of obstacles: %i", n_obstacles);

        ROS_INFO("Number of current agent: %i", agent);

        std::cout << "Init robot positions: \n" << q0.format(OctaveFmt) << sep;

        if(n_obstacles > 0)
            std::cout << "Init obstacles positions: \n" << qobs.format(OctaveFmt) << sep;
    }
    else{
        ROS_ERROR("Failed to call service get_world_properties");
        return 1;
    }

    /***************************** Pub's & Sub's *********************/
    pos_msg.resize(n_robots);
    std::string topic_pub;
    ros::Publisher pose_pub;
    std::vector<std::string> topic_sub(n_robots);
    std::vector<ros::Subscriber> pose_sub(n_robots);
    trajectory_msgs::MultiDOFJointTrajectory MultiDOF_msg;
    ros::Rate rate(10);
    for(int i=0; i<n_robots; i++){
        topic_sub[i] = slash + robot_name + std::to_string(i) + subscriber_name;
        pose_sub[i] = nh.subscribe<geometry_msgs::PointStamped>(topic_sub[i],1,boost::bind(positionCallback, _1, i));
    }

    topic_pub = slash + robot_name + std::to_string(agent) + publisher_name;
    pose_pub = nh.advertise<trajectory_msgs::MultiDOFJointTrajectory>(topic_pub,1);

    //Distributed topics
    controls_msg.resize(n_robots);
    std::string topic_control_pub;
    ros::Publisher controls_pub;
    std::vector<std::string> controls_topic_sub(n_robots);
    std::vector<ros::Subscriber> controls_sub(n_robots);

    for(int i=0; i<n_robots; i++){
        controls_topic_sub[i] = slash + robot_name + std::to_string(i) + topic_controls;
        controls_sub[i] = nh.subscribe<geometry_msgs::Twist>(controls_topic_sub[i],1,boost::bind(controlsCallback, _1, i));
    }

    topic_control_pub = slash + robot_name + std::to_string(agent) + topic_controls;
    controls_pub = nh.advertise<geometry_msgs::Twist>(topic_control_pub, 1);

    /********************** Compute L, D, qz ******************************/
    Eigen::MatrixXd L, eta, etaA;
    L = Consensus::laplacian(A);
    /******************** System inicialization ***************************/
    Eigen::VectorXd qpz_e = Eigen::VectorXd::Zero(n_robots*dim);
    Eigen::VectorXd e_qz = Eigen::VectorXd::Zero(n_robots*dim);
    Eigen::MatrixXd q_odom(dim, n_robots);
    Eigen::MatrixXd q_controls(dim, n_robots);

    Eigen::MatrixXd qz_centroid;
    Eigen::VectorXd robot_new_pose;
    Eigen::VectorXd q_test, e_c, e_c_L2, e_cons, p_ia, q_c, e_0;

    qz = q0 - d_matrix;

    /************************ Read data parameters *************************/
    Eigen::MatrixXd tmp; // tmp matrix to read paraneters

    int updated = 0, steps = 0, counter = 0, Hp = 0, Hu = 0, Hw = 0;
    //Simulation
    double epsilon, t_old, t_new, dt, t = 0, tf, d, k1, p_gain, lambda;
    //MPC
    double cost_x, cost_u, D, k_d, E, k_e, min_u, max_u, minx, maxx, miny, maxy, minz, maxz;

    double tf2=4;
    Eigen::VectorXd ratios(2);
    Eigen::VectorXd system_params;
    readMatrix(simulation_filename, tmp);
    if(tmp.rows() < 3){
        std::cout << "Default system simulations params inicialization" << std::endl;
        epsilon = 0.05; dt = d = 0.01; tf = 10.0;
    }
    else{
        epsilon = tmp(0,0); dt = d = tmp(1,0); tf = tmp(2,0);
    }

    readMatrix(MPC_filename, tmp);
    if(tmp.rows() < 17){
        std::cout << "Default system MPC params inicialization" << std::endl;
        Hp = 15; Hu = 10; Hw = 1; cost_x = 1000.0; cost_u = 1.0;
        D = 0.5; k_d = 10.0; E = 2.25; k_e = 1.0;
        min_u = -10.0; max_u = 10.0;
        minx = -3.0; maxx = 4.0; miny = -3.0; maxy = 4.0; minz = 0.3; maxz = 2.5;
    }
    else{
        Hp = int(tmp(0,0)); Hu = int(tmp(1,0)); Hw = int(tmp(2,0)); cost_x = tmp(3,0); cost_u = tmp(4,0);
        D = tmp(5,0); k_d = tmp(6,0); E = tmp(7,0); k_e = tmp(8,0);
        min_u = tmp(9,0); max_u = tmp(10,0);
        minx = tmp(11,0); maxx = tmp(12,0); miny = tmp(13,0); maxy = tmp(14,0); minz = tmp(15,0); maxz = tmp(16,0);
    }

    std::cout << "Simulation inicialization" << sep;

    steps = tf/dt;

    /************************ Data storage variables ***********************/
    Eigen::MatrixXd q_odom_data(dim*n_robots, steps+1);
    Eigen::MatrixXd qp_agent_data(dim, steps+1); //Velocities of agent
    Eigen::MatrixXd qp_computed_controls_data(dim*n_robots, steps+1);
    Eigen::MatrixXd e_consensus_data(dim*n_robots, steps+1);// Consensus of virtual system
 	Eigen::MatrixXd qz_data_agents(dim*n_robots, steps+1);
    /*************************** Start simulation ***********************/
    bool trajectory = false;

    /************************************* START SIMULATION ************************************/
    while(ros::ok()){
        if(counter > steps)
            break;

        // hold for updates
        if(updated < 2){
            rate.sleep();
            updated+=1;
            continue;
        }

        ros::spinOnce();

        getCurrentPose(q_odom);

        qz = q_odom - d_matrix;
        e_qz = Consensus::consensus_error(L, qz);
        e_c_L2 = Consensus::norm_consensus_error(e_qz, n_robots);

        q_odom_data.col(counter) = Consensus::matrix2vector(q_odom);
		e_consensus_data.col(counter)=e_qz;
		qz_data_agents.col(counter)=Consensus::matrix2vector(qz);

        if(e_c_L2.maxCoeff()<epsilon)
        {
           ROS_INFO("Consensus reached");
            break;
        }

        MPC mpc_control(dim, Hp, Hu, Hw, dt, cost_x, cost_u, min_u, max_u);

        mpc_control.set_reference(qz, A.col(agent), n_robots, agent);

        for (int i = 0; i < n_robots; i++)
        {
            if (A(i,agent) == 1)
            {
                mpc_control.add_obstacle_constraint_column(q_odom.col(i), q_odom.col(agent), D);
                mpc_control.add_obstacle_penalty_column(q_odom.col(i), q_odom.col(agent), D, k_d);
                mpc_control.add_connectivity_constraint(q_odom.col(i), q_odom.col(agent), E);
                mpc_control.add_connectivity_penalty(q_odom.col(i), q_odom.col(agent), E, k_e);
            }
        }

        for (int i = 0; i < n_obstacles; i++)
        {
            mpc_control.add_obstacle_constraint_column(qobs.col(i), q_odom.col(agent), D);
            mpc_control.add_obstacle_penalty_column(qobs.col(i), q_odom.col(agent), D, k_d);
        }

        //mpc_control.set_lower_constraint(q_odom.col(agent), minz);
        mpc_control.set_box_constraints(q_odom.col(agent), minx, maxx, miny, maxy, minz, maxz);

        Eigen::VectorXd solMPC(dim*Hu);

        int success_mpc;

        success_mpc = mpc_control.compute_control(solMPC);

        Eigen::VectorXd vel(dim);
        for (int i=0; i < dim; i++)
            vel(i) = solMPC[i];

        //Publish computed controls
        geometry_msgs::Twist c_msg;
        c_msg.linear.x = vel(0); c_msg.linear.y = vel(1); c_msg.linear.z = vel(2);
        controls_pub.publish(c_msg);

        //Read published controls
        getComputedControls(q_controls);

        qp_computed_controls_data.col(counter) = Consensus::matrix2vector(q_controls);

        //Distributed control
        MPC mpc_control2(dim, Hp, Hu, Hw, dt, cost_x, cost_u, min_u, max_u);

        mpc_control2.set_reference(qz, A.col(agent), n_robots, agent, q_controls);

        for (int i = 0; i < n_robots; i++)
        {
            if (A(i,agent) == 1)
            {
                mpc_control2.add_obstacle_constraint_column(q_odom.col(i), q_odom.col(agent), q_controls.col(i), D);
                mpc_control2.add_obstacle_penalty_column(q_odom.col(i), q_odom.col(agent), q_controls.col(i), D, k_d);
                mpc_control2.add_connectivity_constraint(q_odom.col(i), q_odom.col(agent), q_controls.col(i), E);
                mpc_control2.add_connectivity_penalty(q_odom.col(i), q_odom.col(agent), q_controls.col(i), E, k_e);
            }
        }

        for (int i = 0; i < n_obstacles; i++)
        {
            mpc_control2.add_obstacle_constraint_column(qobs.col(i), q_odom.col(agent), D);
            mpc_control2.add_obstacle_penalty_column(qobs.col(i), q_odom.col(agent), D, k_d);
        }

        //mpc_control2.set_lower_constraint(q_odom.col(agent), minz);
        mpc_control2.set_box_constraints(q_odom.col(agent), minx, maxx, miny, maxy, minz, maxz);

        success_mpc = mpc_control2.compute_control(solMPC);

        for (int i=0; i < dim; i++)
            vel(i) = solMPC[i];

        qp_agent_data.col(counter) = vel;

        robot_new_pose = q_odom.col(agent) + dt*vel;

        //std::cout << robot_new_pose(0) << " " << robot_new_pose(1) << " " << robot_new_pose(2) << sep;

        // set MultiDOF msg
        trajectory_msgs::MultiDOFJointTrajectory msg;
        msg.header.stamp=ros::Time::now();
        mav_msgs::msgMultiDofJointTrajectoryFromPositionYaw(robot_new_pose, 1 , &msg);
        pose_pub.publish(msg);

        counter++;
        t+=dt;
        rate.sleep();
    }

    std::cout << "Consensus reached in " << counter << " steps " << sep;

    std::cout << "Final error = " << e_c_L2.maxCoeff() << std::endl;
    std::cout << "Error vector" << std::endl << e_c_L2 << sep;

    // save data to files
    if (leader)
    {
        IOEigen::writeMatrix(q_odom_filename, q_odom_data.block(0, 0, dim*n_robots, counter), true);
        IOEigen::writeMatrix(qp_agent_filename, qp_agent_data.block(0, 0, dim, counter), true);
        IOEigen::writeMatrix(e_consensus_filename, e_consensus_data.block(0, 0, dim*n_robots, counter), true);
        IOEigen::writeMatrix(qp_computed_controls_filename, qp_computed_controls_data.block(0, 0, dim*n_robots, counter), true);
    	IOEigen::writeMatrix(qz_agents_filename, qz_data_agents.block(0, 0, dim*n_robots, counter), true);

        system("python /home/cimat/bebop_ws/src/multi_uav_control/src/library/plots.py");
    }
    else
    {
        IOEigen::writeMatrix(qp_agent_filename, qp_agent_data.block(0, 0, dim, counter), true);
    }

    return 0;
}

/*
Multiple pose callback.
Return the pose of the i-robot
*/
void positionCallback(const geometry_msgs::PointStamped::ConstPtr& msg, int index){
        pos_msg[index].point.x = msg->point.x;
        pos_msg[index].point.y = msg->point.y;
        pos_msg[index].point.z = msg->point.z;
}

/*
Get current pose from odometry for each robot
*/
void getCurrentPose(Eigen::MatrixXd &q){
    Eigen::VectorXd current_pose(3);
    for(int i=0; i<n_robots; i++){
        current_pose(0) = pos_msg[i].point.x; current_pose(1) = pos_msg[i].point.y; current_pose(2) = pos_msg[i].point.z;
        q.col(i) = current_pose;
    }
}

/*
Multiple controls callback.
Return the controls of the i-robot
*/
void controlsCallback(const geometry_msgs::Twist::ConstPtr& msg, int index){
        controls_msg[index].linear.x = msg->linear.x;
        controls_msg[index].linear.y = msg->linear.y;
        controls_msg[index].linear.z = msg->linear.z;
}

/*
Get computed controls for each robot
*/
void getComputedControls(Eigen::MatrixXd &q){
    Eigen::VectorXd computed_controls(3);
    for(int i=0; i<n_robots; i++){
        computed_controls(0) = controls_msg[i].linear.x; computed_controls(1) = controls_msg[i].linear.y; computed_controls(2) = controls_msg[i].linear.z;
        q.col(i) = computed_controls;
    }
}

bool readMatrix(std::string filename, Eigen::MatrixXd &M){
    std::ifstream fin(filename.c_str());
    if(!fin.good()){
            return false;
    }
    int rows, cols;
    //read header matrix Name rows cols
    fin >> rows; fin >> cols;
    M.resize(rows,cols);
    for(int i=0; i<rows; i++){
        for(int j=0; j<cols; j++){
            fin >> M(i,j);
        }
    }

    return true;
}