#include <gtest/gtest.h>

#include "scheme/numeric/util.hh"

#include "scheme/nest/pmap/Rotation1DMap.hh"
#include "scheme/nest/NEST.hh"

#include <Eigen/Dense>
#include <random>

#include <boost/lexical_cast.hpp>

#include <fstream>

namespace scheme { namespace nest { namespace pmap { namespace test {

using namespace Eigen;
using std::cout;
using std::endl;

Vector3d get_axis( Matrix3d const & m ){
	AngleAxisd aa( m );
	if( aa.axis().dot( Vector3d(1,1,1) ) < 0 ) return -aa.axis();
	else                                       return  aa.axis();
}
double get_angle( Matrix3d const & m ){
	AngleAxisd aa( m );
	if( aa.axis().dot( Vector3d(1,1,1) ) < 0 ) return -aa.angle();
	else                                       return  aa.angle();
}

TEST( Rotation1DMap , flip_test ){

	{
		NEST<1,Matrix3d,Rotation1DMap> nest;
		Vector3d axis(1,0,0), flip(0,1,0);
		nest.set_axis( axis );
		nest.set_flip_axis( flip );
		{ Matrix3d m = nest.set_and_get(0,0); ASSERT_DOUBLE_EQ(  0.0     , get_angle(m) );                                               }
		{ Matrix3d m = nest.set_and_get(1,0); ASSERT_DOUBLE_EQ(  M_PI    , get_angle(m) ); ASSERT_GE( get_axis(m).dot( flip ), 0.9999 ); }
		{ Matrix3d m = nest.set_and_get(0,1); ASSERT_DOUBLE_EQ( -M_PI/2.0, get_angle(m) ); ASSERT_GE( get_axis(m).dot( axis ), 0.9999 ); }
		{ Matrix3d m = nest.set_and_get(1,1); ASSERT_DOUBLE_EQ(  M_PI/2.0, get_angle(m) ); ASSERT_GE( get_axis(m).dot( axis ), 0.9999 ); }
		{
			Matrix3d m = nest.set_and_get(2,1);
			ASSERT_DOUBLE_EQ(  M_PI    , get_angle(m) );
			ASSERT_GE( get_axis(m).dot( Vector3d(0,sqrt(2)/2,-sqrt(2)/2) ), 0.9999 );
		}
		{
			Matrix3d m = nest.set_and_get(3,1);
			ASSERT_DOUBLE_EQ(  M_PI, get_angle(m) );
			ASSERT_GE( get_axis(m).dot( Vector3d(0,sqrt(2)/2,sqrt(2)/2) ), 0.9999 );
		}
		#ifndef NDEBUG
		#ifndef CXX14
		ASSERT_DEATH( nest.set_and_get(4,1), ".*" );
		#endif
		#endif
	}

}

TEST( Rotation1DMap , basic_test ){

	{
		NEST<1,Matrix3d,Rotation1DMap> nest;
		Vector3d axis(1,0,0); nest.set_axis( axis );
		{ Matrix3d m = nest.set_and_get(0,0); ASSERT_DOUBLE_EQ(  0.0     , get_angle(m) );                                               }
		{ Matrix3d m = nest.set_and_get(0,1); ASSERT_DOUBLE_EQ( -M_PI/2.0, get_angle(m) ); ASSERT_GE( get_axis(m).dot( axis ), 0.9999 ); }
		{ Matrix3d m = nest.set_and_get(1,1); ASSERT_DOUBLE_EQ(  M_PI/2.0, get_angle(m) ); ASSERT_GE( get_axis(m).dot( axis ), 0.9999 ); }
	}
	{
		NEST<1,Matrix3d,Rotation1DMap> nest;
		Vector3d axis(1,1,1); nest.set_axis( axis );
		{ Matrix3d m = nest.set_and_get(0,0); ASSERT_DOUBLE_EQ(  0.0     , get_angle(m) );                                               }
		{ Matrix3d m = nest.set_and_get(0,1); ASSERT_DOUBLE_EQ( -M_PI/2.0, get_angle(m) ); ASSERT_GE( get_axis(m).dot( axis ), 0.9999 ); }
		{ Matrix3d m = nest.set_and_get(1,1); ASSERT_DOUBLE_EQ(  M_PI/2.0, get_angle(m) ); ASSERT_GE( get_axis(m).dot( axis ), 0.9999 ); }
	}
	{
		NEST<1,Matrix3d,Rotation1DMap> nest(0,M_PI,1);
		Vector3d axis(1,1,1); nest.set_axis( axis );
		{ Matrix3d m = nest.set_and_get(0,0); ASSERT_DOUBLE_EQ(    M_PI/2.0, get_angle(m) ); ASSERT_GE( get_axis(m).dot( axis ), 0.9999 ); }
		{ Matrix3d m = nest.set_and_get(0,1); ASSERT_DOUBLE_EQ(    M_PI/4.0, get_angle(m) ); ASSERT_GE( get_axis(m).dot( axis ), 0.9999 ); }
		{ Matrix3d m = nest.set_and_get(1,1); ASSERT_DOUBLE_EQ(  3*M_PI/4.0, get_angle(m) ); ASSERT_GE( get_axis(m).dot( axis ), 0.9999 ); }
	}

}

TEST( Rotation1DMap , lookup ){
	int NITER = 1000;
	#ifdef NDEBUG
	NITER *= 30;
	#endif

	NEST<1,Matrix3d,Rotation1DMap> nest;
	// cout << nest.get_index( Matrix3d::Identity(), 0 ) << endl;

	NestBase<> *nestp = new NEST<1,Matrix3d,Rotation1DMap>();
	Matrix3d m = Matrix3d::Identity();
	ASSERT_EQ( 0, nestp->virtual_get_index( &m , 0 ) );

	std::mt19937 rng(0);
	std::normal_distribution<> gauss;
	std::uniform_real_distribution<> uniform;
	for(int i = 0; i < NITER; ++i){
		Vector3d axis( gauss(rng), gauss(rng), gauss(rng) );
		axis.normalize();
		nest.set_axis(axis);
		nest.lb = uniform(rng) * 2*M_PI - M_PI;
		nest.ub = nest.lb + uniform(rng) * (2*M_PI-nest.lb-M_PI);
		BOOST_VERIFY( -M_PI <= nest.lb && nest.lb <= M_PI );
		BOOST_VERIFY( -M_PI <= nest.ub && nest.ub <= M_PI );
		BOOST_VERIFY( nest.lb < nest.ub );
		nest.nside = (uint64_t)(3*uniform(rng)+1.0);

		for(int resl = 0; resl < 5; ++resl){
			for(uint64_t index = 0; index < nest.size(resl); ++index){
				Matrix3d m = nest.set_and_get(index,resl);
				ASSERT_EQ( nest.get_index(m,resl), index );
			}
		}

	}
}

TEST( Rotation1DMap , lookup_flip ){
	int NITER = 1000;
	#ifdef NDEBUG
	NITER *= 30;
	#endif

	NEST<1,Matrix3d,Rotation1DMap> nest;
	// cout << nest.get_index( Matrix3d::Identity(), 0 ) << endl;

	NestBase<> *nestp = new NEST<1,Matrix3d,Rotation1DMap>();
	Matrix3d m = Matrix3d::Identity();
	ASSERT_EQ( 0, nestp->virtual_get_index( &m , 0 ) );

	std::mt19937 rng(0);
	std::normal_distribution<> gauss;
	std::uniform_real_distribution<> uniform;
	for(int i = 0; i < NITER; ++i){
		Vector3d axis( gauss(rng), gauss(rng), gauss(rng) );
		axis.normalize();
		nest.set_axis(axis);
		nest.set_flip_axis( axis.cross( Vector3d( gauss(rng), gauss(rng), gauss(rng) ) ) );
		nest.lb = uniform(rng) * 2*M_PI - M_PI;
		nest.ub = nest.lb + uniform(rng) * (2*M_PI-nest.lb-M_PI);
		BOOST_VERIFY( -M_PI <= nest.lb && nest.lb <= M_PI );
		BOOST_VERIFY( -M_PI <= nest.ub && nest.ub <= M_PI );
		BOOST_VERIFY( nest.lb < nest.ub );
		nest.nside = (uint64_t)(3*uniform(rng)+1.0);

		for(int resl = 0; resl < 5; ++resl){
			for(uint64_t index = 0; index < nest.size(resl); ++index){
				Matrix3d m = nest.set_and_get(index,resl);
				ASSERT_EQ( nest.get_index(m,resl), index );
			}
		}

	}
}

}}}}


#include <iostream>
#include <optional>

#include <boost/asio.hpp>

#include "vdf.h"

using namespace std;

int gcd_base_bits = 50;
int gcd_128_max_iter = 3;

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

struct bluebox_request {
    integer discriminant;
    uint64_t iterations;
};

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

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

        square_state_type square_state;
        square_state.pairindex = 0;

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

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

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

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

    PulmarkReducer reducer;

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

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

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

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

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

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

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

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

    return p;
}

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

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

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

    return vdf_result{y, p};
}

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

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

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

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

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

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

    return BytesToStr(bytes);
}

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

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

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

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

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

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

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

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

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

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

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

    return bluebox_request{discriminant, iterations};
}

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

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

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

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

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

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

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

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

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

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

    vdf_result r = vdf_1weso(D, T);

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

    return 0;
}

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

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

    bluebox_request req = read_request(socket, client_id);

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

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

    write_response(socket, res);

    return 0;
}

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

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

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

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


#include "solvers/SolverBoxBPP.h"

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

#include <Eigen/Dense>

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

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

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

        return rows;
    }

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

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

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

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

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

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

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

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

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

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

    }

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

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

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

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

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

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

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

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

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

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

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

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

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

    }

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

}

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

}

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

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

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

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

		updateBounds(contacts, lower, upper);

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

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

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

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

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

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

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

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


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

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

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

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

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

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

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

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

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

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

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

void Initialize(ros::NodeHandle& nh)
{

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

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

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

}

void Initialize_2()
{

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

}

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

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

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

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

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

}

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

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

}

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

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

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

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

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

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

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

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

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

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

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

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

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

}

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

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

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

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

}

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

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

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

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

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

}

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

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

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

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

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

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

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

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

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


// Copyright (c) 2020 fortiss GmbH
//
// Authors: Julian Bernhard, Klemens Esterle, Patrick Hart and
// Tobias Kessler
//
// This work is licensed under the terms of the MIT license.
// For a copy, see <https://opensource.org/licenses/MIT>.

#ifndef BARK_MODELS_EXECUTION_MPC_COST_FUNCTOR_HPP_
#define BARK_MODELS_EXECUTION_MPC_COST_FUNCTOR_HPP_

#include <Eigen/Dense>
#include <iostream>
#include "bark/models/execution/mpc/common.hpp"

namespace bark {
namespace models {
namespace execution {
using namespace bark::models::dynamic;

template <typename T>
inline Eigen::Matrix<T, Eigen::Dynamic, Eigen::Dynamic> KinematicModel(
    T const* const* parameters,
    bark::models::dynamic::Trajectory desired_trajectory,
    const OptimizationSettings& optimization_settings) {
  Eigen::Matrix<T, Eigen::Dynamic, Eigen::Dynamic> traj =
      desired_trajectory.cast<T>();

  // TODO: Lf als Parameter
  T Lf = T(2.6);

  for (int i = 1; i < optimization_settings.num_optimization_steps; i++) {
    Eigen::Matrix<T, Eigen::Dynamic, Eigen::Dynamic> f_dot(1, 5);

    // fill matrix
    f_dot << T(1),  // (0) t' = 1
        traj(i - 1, StateDefinition::VEL_POSITION) *
            cos(traj(
                i - 1,
                StateDefinition::THETA_POSITION)),  // (1) x_dot = v*cos(psi)
        traj(i - 1, StateDefinition::VEL_POSITION) *
            sin(traj(
                i - 1,
                StateDefinition::THETA_POSITION)),  // (2) y_dot = v*sin(psi)
        tan(parameters[1][i - 1]) / Lf,  // (3) theta_dot = tan(delta)/Lf
        parameters[0][i - 1];            // (4) v_dot = a

    traj.block(i, 0, 1, 5) =
        traj.block(i - 1, 0, 1, 5) + T(optimization_settings.dt) * f_dot;
  }
  return traj;
}

struct CostFunctor {
  template <typename T>
  bool operator()(T const* const* parameters, T* residual) {
    // initialize
    T cost = T(0.0);
    Eigen::Matrix<T, Eigen::Dynamic, Eigen::Dynamic> traj = KinematicModel(
        parameters, desired_discrete_states_, optimization_settings_);
    Eigen::Matrix<T, Eigen::Dynamic, Eigen::Dynamic> desired_discrete_states =
        desired_discrete_states_.cast<T>();
    Eigen::Matrix<T, Eigen::Dynamic, Eigen::Dynamic> weights =
        weights_.cast<T>();

    // offsets
    for (int i = 0; i < optimization_settings_.num_optimization_steps - 1;
         i++) {
      T xDiff, yDiff, thetaDiff, velDiff;
      xDiff = traj(i, StateDefinition::X_POSITION) -
              desired_discrete_states(i, StateDefinition::X_POSITION);
      yDiff = traj(i, StateDefinition::Y_POSITION) -
              desired_discrete_states(i, StateDefinition::Y_POSITION);
      thetaDiff = traj(i, StateDefinition::THETA_POSITION) -
                  desired_discrete_states(i, StateDefinition::THETA_POSITION);
      velDiff = traj(i, StateDefinition::VEL_POSITION) -
                desired_discrete_states(i, StateDefinition::VEL_POSITION);

      cost += weights(i, StateDefinition::X_POSITION) * xDiff * xDiff;
      cost += weights(i, StateDefinition::Y_POSITION) * yDiff * yDiff;
      cost +=
          weights(i, StateDefinition::THETA_POSITION) * thetaDiff * thetaDiff;
      cost += weights(i, StateDefinition::VEL_POSITION) * velDiff * velDiff;
    }

    // differential values
    // TODO(@all): seperate weights
    T weights_acc = T(1e1);
    T weights_delta = T(1);
    for (int i = 0; i < optimization_settings_.num_optimization_steps - 1;
         i++) {
      cost +=
          weights_acc * parameters[0][i] * parameters[0][i];  // acceleration
      cost += weights_delta * parameters[1][i] * parameters[1][i];  // steering
    }

    residual[0] = cost / (weights.sum() +
                          (weights_acc + weights_delta) *
                              T(optimization_settings_.num_optimization_steps));
    return true;
  }

  bool set_desired_states(
      const bark::models::dynamic::Trajectory& desired_states) {
    desired_discrete_states_ = desired_states;
    return true;
  }

  bool set_optimization_settings(
      const OptimizationSettings& optimization_settings) {
    optimization_settings_ = optimization_settings;
    return true;
  }

  bool set_weights(
      const Eigen::Matrix<double, Eigen::Dynamic, Eigen::Dynamic> weights) {
    weights_ = weights;
    return true;
  }

  bark::models::dynamic::Trajectory desired_discrete_states_;
  Eigen::Matrix<double, Eigen::Dynamic, Eigen::Dynamic> weights_;
  OptimizationSettings optimization_settings_;
};

}  // namespace execution
}  // namespace models
}  // namespace bark

#endif  // BARK_MODELS_EXECUTION_MPC_COST_FUNCTOR_HPP_


/*  _______________________________________________________________________

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

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

#ifndef RANDOM_VARIABLE_HPP
#define RANDOM_VARIABLE_HPP

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

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

namespace Pecos {

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

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

/// base class for random variable hierarchy

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

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

class RandomVariable
{
public:

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

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

  /// destructor
  virtual ~RandomVariable();

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

  //
  //- Heading: Virtual functions
  //

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

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

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

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

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

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

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

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

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

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

  //
  //- Heading: Member functions
  //

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

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

protected:

  //
  //- Heading: Constructors
  //

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

  //
  //- Heading: Member functions
  //

  //
  //- Heading: Data
  //

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

private:

  //
  //- Heading: Member functions
  //

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

  //
  //- Heading: Data members
  //

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

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

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

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

} // namespace Pecos

#endif


/**
* @Author: Marcel Ruhf <marcelruhf>
* @Email:  m.ruhf@protonmail.ch
* @Project: SoccerRL
* @Filename: Functions.cpp
* @License: Licensed under the Apache 2.0 license (see LICENSE.md)
* @Copyright: Copyright (c) 2017 Marcel Ruhf
*/

#include <cmath>
#include <string>
#include <boost/property_tree/ptree.hpp>
#include <boost/property_tree/json_parser.hpp>
#include <boost/optional.hpp>
#include <opencv2/opencv.hpp>
#include <MarkerFinder.hpp>
#include <RobotTracker.hpp>
#include <BallTracker.hpp>
#include <Constants.hpp>
#include <Functions.hpp>

namespace mr
{
    template <typename T>
    int signum(T val) {
        return (T(0) < val) - (val < T(0));
    }

    double distEuclidPixels(cv::Point2f p1, cv::Point2f p2)
    {
        return std::sqrt(std::pow(p2.x - p1.x, 2) + std::pow(p2.y - p1.y, 2));
    }

    double distEuclidMM(cv::Point2f p1, cv::Point2f p2)
    {
        double distPx = distEuclidPixels(p1, p2);
        return (distPx / PIXELS_PER_MM);
    }

    boost::optional<cv::Point2f> initialize_goal_centre(const cv::Mat& frame)
    {
        MarkerFinder marker;
        marker.preprocess(frame);
        boost::optional<std::vector<cv::Point2f>> goalCentreCorners = marker.getCorners(GOAL_CENTRE_MARKER_ID);
        if (goalCentreCorners)
        {
            cv::Moments mu = cv::moments(*goalCentreCorners, true);
            return cv::Point2f(mu.m10/mu.m00, mu.m01/mu.m00);
        }
        return boost::optional<cv::Point2f>{};
    }

    void get_vars(int vars_array[4], const cv::Mat& src, RobotTracker& robot, BallTracker& ball)
    {
        //double len1 = std::hypot(robot_centroid.x, robot_centroid.y);
        //double len2 = std::hypot(ball_centroid.x, ball_centroid.y);
        //double dot = robot_centroid.ddot(ball_centroid);
        //double a = dot/(len1*len2);
        //double arb = acos(a) / (180 * 3.141592653589793);
        //root.add("arb", arb);

        // Determine distance between robot and ball in terms of the X axis
        robot.preprocess(src);
        ball.setImage(src);
        boost::optional<cv::Point2f> robot_centroid = robot.getCentrePoint();
        boost::optional<cv::Point2f> ball_centroid = ball.getCentrePoint();
        int ball_velocity = ball.getVelocity();
        if (!ball_centroid || !robot_centroid)
        {
            vars_array = {0};
        }
        else
        {
            double drbx     = (ball_centroid->x - robot_centroid->x) / PIXELS_PER_MM;
            vars_array[0]   = static_cast<int>(std::round(drbx));
            vars_array[1]   = static_cast<int>(std::round(ball_velocity));
            vars_array[2]   = static_cast<int>(std::round(ball_centroid->x));
            vars_array[3]   = static_cast<int>(std::round(ball_centroid->y));
        }
    }
}


#include <boost/math/special_functions/legendre.hpp>


#include <iostream>
#include <complex>
#include <cmath>
#include <Eigen/Dense>
#include "mex.h"
#include <unsupported/Eigen/MatrixFunctions>
#include <vector>

typedef Eigen::Matrix<double, 2, 5> Matrix2x5d;
typedef Eigen::Matrix<double, 5, 5> Matrix5d;

Matrix5d F;
Matrix5d V;
Matrix2x5d J;

std::complex<double> dmdT1;
std::complex<double> dmdT2;
std::complex<double> dmdM0;
std::complex<double> dmdB0;
std::complex<double> dmdP0;
const std::complex<double> i1(0.0, 1.0);
const double PI = 3.14159265358979323846;
Eigen::Matrix4d Rz;
Eigen::Matrix4d dRz;
Eigen::Matrix4d Rx;
Eigen::Matrix4d E;
Eigen::Matrix4d dEdT1;
Eigen::Matrix4d dEdT2;
Eigen::Matrix4d dEdM0;

void update_Rz(double p);

void update_dRz(double p, double dp);

void update_Rx(double a);

void update_E(double R1, double R2, double TR, double M0);

void update_dEdT1(double R1, double TR, double M0);

void update_dEdT2(double R2, double TR);

void update_dEdM0(double R1, double TR);

void add2FIM();

//Main Function that generates the EPG states given the acquisition parameters
void mex_isochromat_CRLB_theta(double *signal_real, double *signal_imag, double *dmdT1_real, double *dmdT1_imag, double *dmdT2_real, double *dmdT2_imag, double *Mz, double *rCRLB, double *N, double *theta, double *RF, double *TRssfp, double *T1, double *T2, double *M0, double *B0, double *phi0)
{

    double R1 = 1 / T1[0];
	double R2 = 1 / T2[0];

	int length = (int) N[0];

    ///// MEMORY ALLOCATION AND INITIALIZATION
    // signal vectors
    Eigen::Vector4d Mplus(0, 0, 0, 1);  //Magnetization vector after the RF excitation
    Eigen::Vector4d Mminus(0, 0, 0, 1); //Magnetization vector before the RF excitation
    Mminus(2) = M0[0];

    Eigen::VectorXd aux_FA(length);
    aux_FA.setZero();
    Eigen::VectorXd TR(length);
    TR.setOnes();
    TR *= TRssfp[0];
    Eigen::VectorXd TE(length);
    TE.setOnes();
    TE *= TRssfp[0]/2.0;

    // rotation and relaxation matrices
    Rz.setZero();
    Rz(2,2) = 1;
    Rz(3,3) = 1;

    Rx.setZero();
    Rx(0,0) = 1;
    Rx(3,3) = 1;

    E.setZero();
    E(3,3) = 1;
    dEdT1.setZero();
    dEdT2.setZero();
    dEdM0.setZero();

    // for CRLB calculation
    F.setZero();

    Eigen::Vector4d Mplus_dOdT1(0, 0, 0, 0);
    Eigen::Vector4d Mminus_dOdT1(0, 0, 0, 0);
    Eigen::Vector4d Mplus_dOdT2(0, 0, 0, 0);
    Eigen::Vector4d Mminus_dOdT2(0, 0, 0, 0);
    Eigen::Vector4d Mplus_dOdM0(0, 0, 0, 0);
    Eigen::Vector4d Mminus_dOdM0(0, 0, 0, 0);
    Mminus_dOdM0(2) = M0[0];
    Eigen::Vector4d Mplus_dOdB0(0, 0, 0, 0);
    Eigen::Vector4d Mminus_dOdB0(0, 0, 0, 0);
    Eigen::Vector4d Mplus_dOdP0(0, 0, 0, 0);
    Eigen::Vector4d Mminus_dOdP0(0, 0, 0, 0);

    double FPangle;

    /////ISOCHROMAT DYNAMICS SIMULATION
    for (int j=0; j<length; j++){

        //transformation from theta to alpha
        if (j==0) {aux_FA(j) = theta[j];}
        else {aux_FA(j) = theta[j] + theta[j-1];}

        //apply rotations
        update_Rz(-RF[j]); update_Rx(aux_FA(j));
        Mplus = Rx * Rz * Mminus;
        Mplus_dOdT1 = Rx * Rz * Mminus_dOdT1;
        Mplus_dOdT2 = Rx * Rz * Mminus_dOdT2;
        Mplus_dOdM0 = Rx * Rz * Mminus_dOdM0;
        Mplus_dOdB0 = Rx * Rz * Mminus_dOdB0;
        Mplus_dOdP0 = Rx * Rz * Mminus_dOdP0;
        update_Rz(RF[j]);
        Mplus = Rz * Mplus;
        Mplus_dOdT1 = Rz * Mplus_dOdT1;
        Mplus_dOdT2 = Rz * Mplus_dOdT2;
        Mplus_dOdM0 = Rz * Mplus_dOdM0;
        Mplus_dOdB0 = Rz * Mplus_dOdB0;
        Mplus_dOdP0 = Rz * Mplus_dOdP0;

        //CRLB calculation
        FPangle = 2*PI*B0[0] * TE(j) * pow(10.0,-3.0);
        update_Rz(FPangle + phi0[0]);
        Mminus       = Rz * Mplus;
        Mminus_dOdT1 = Rz * Mplus_dOdT1;
        Mminus_dOdT2 = Rz * Mplus_dOdT2;
        Mminus_dOdM0 = Rz * Mplus_dOdM0;

        update_dRz(FPangle + phi0[0], 2*PI*TE(j)*pow(10.0,-3.0) );
        Mminus_dOdB0 = dRz * Mplus + Rz * Mplus_dOdB0;

        update_dRz(FPangle + phi0[0], 1.0);
        Mminus_dOdP0 = dRz * Mplus + Rz * Mplus_dOdP0;

        dmdT1 = exp(-R2 * TE(j)) * (Mminus_dOdT1(0) + i1*Mminus_dOdT1(1));
        dmdT2 = TE(j) * exp(-R2 * TE(j)) * pow(R2,2) * (Mminus(0) + i1*Mminus(1)) + exp(-R2 * TE(j)) * (Mminus_dOdT2(0) + i1*Mminus_dOdT2(1));
        dmdM0 = exp(-R2 * TE(j)) * (Mminus_dOdM0(0) + i1*Mminus_dOdM0(1));
        dmdB0 = exp(-R2 * TE(j)) * (Mminus_dOdB0(0) + i1*Mminus_dOdB0(1));
        dmdP0 = exp(-R2 * TE(j)) * (Mminus_dOdP0(0) + i1*Mminus_dOdP0(1));

        add2FIM();

        signal_real[j] = exp(-R2 * TE(j)) * Mminus(0);
        signal_imag[j] = exp(-R2 * TE(j)) * Mminus(1);
        dmdT1_real[j]  = dmdT1.real();
        dmdT1_imag[j]  = dmdT1.imag();
        dmdT2_real[j]  = dmdT2.real();
        dmdT2_imag[j]  = dmdT2.imag();
        Mz[j] = Mminus(2) * exp(-R1 * TR(j)) + (1 - exp(-R1 * TR(j))) * std::abs(M0[0]);

        //apply relaxation
        update_E(R1, R2, TR(j), std::abs(M0[0]));
    	update_dEdT1(R1, TR(j), std::abs(M0[0]));
        update_dEdT2(R2, TR(j));
        update_dEdM0(R1, TR(j));
        Mminus = E * Mplus;
        Mminus_dOdT1 = E * Mplus_dOdT1 + dEdT1 * Mplus;
        Mminus_dOdT2 = E * Mplus_dOdT2 + dEdT2 * Mplus;
        Mminus_dOdM0 = E * Mplus_dOdM0 + dEdM0 * Mplus;
        Mminus_dOdB0 = E * Mplus_dOdB0;
        Mminus_dOdP0 = E * Mplus_dOdP0;

        //apply rotation due to free precession
        FPangle = 2*PI*B0[0] * TR(j) * pow(10.0,-3.0);
        update_Rz(FPangle);
        update_dRz(FPangle, 2*PI*TR(j)*pow(10.0,-3.0) );
        Mminus = Rz * Mminus;
        Mminus_dOdT1 = Rz * Mminus_dOdT1;
        Mminus_dOdT2 = Rz * Mminus_dOdT2;
        Mminus_dOdM0 = Rz * Mminus_dOdM0;
        Mminus_dOdB0 = dRz * E * Mplus + Rz * Mminus_dOdB0;
        Mminus_dOdP0 = Rz * Mminus_dOdP0;

}

    V = F.inverse();
    rCRLB[0] = V(0,0);
    rCRLB[1] = V(1,1);
    rCRLB[2] = V(2,2);
    rCRLB[3] = V(3,3);
    rCRLB[4] = V(4,4);
}

void mexFunction(int nlhs, mxArray *plhs[], int nrhs, const mxArray *prhs[])
{
    /* Declare arrays necessary to compute signals */
    double *N;         //units []
    double *theta;     //units [rad]
    double *RF;        //units [rad]
    double *TRssfp;    //units [ms]
    double *T1;        //units [ms]
    double *T2;        //units [ms]
    double *M0;        //a.u.
    double *B0;        //units [Hz]
    double *phi0;      //units [rad]

    double *signal_real;
    double *signal_imag;
    double *dmdT1_real;
    double *dmdT1_imag;
    double *dmdT2_real;
    double *dmdT2_imag;
    double *Mz;
    double *rCRLB;

    /* Check for proper number of arguments */

    /* Create pointers to the inputs */
    N      = mxGetPr(prhs[0]);
    theta  = mxGetPr(prhs[1]);
    RF     = mxGetPr(prhs[2]);
    TRssfp = mxGetPr(prhs[3]);
    T1     = mxGetPr(prhs[4]);
    T2     = mxGetPr(prhs[5]);
    M0     = mxGetPr(prhs[6]);
    B0     = mxGetPr(prhs[7]);
    phi0   = mxGetPr(prhs[8]);

    /* Create pointers to the outputs */
    plhs[0] = mxCreateDoubleMatrix((mwSize)*N, 1, mxCOMPLEX);
    plhs[1] = mxCreateDoubleMatrix((mwSize)*N, 1, mxCOMPLEX);
    plhs[2] = mxCreateDoubleMatrix((mwSize)*N, 1, mxCOMPLEX);
    plhs[3] = mxCreateDoubleMatrix((mwSize)*N, 1, mxREAL);
    plhs[4] = mxCreateDoubleMatrix(1, 5, mxREAL);

    // Check this link to see how to return complex arrays:
    // http://matlab.izmiran.ru/help/techdoc/matlab_external/ch04cre9.html
    signal_real = mxGetPr(plhs[0]);
    signal_imag = mxGetPi(plhs[0]);
    dmdT1_real  = mxGetPr(plhs[1]);
    dmdT1_imag  = mxGetPi(plhs[1]);
    dmdT2_real  = mxGetPr(plhs[2]);
    dmdT2_imag  = mxGetPi(plhs[2]);
    Mz          = mxGetPr(plhs[3]);
    rCRLB       = mxGetPr(plhs[4]);

    /* Call the computational routine */
    mex_isochromat_CRLB_theta(signal_real, signal_imag, dmdT1_real, dmdT1_imag, dmdT2_real, dmdT2_imag, Mz, rCRLB, N, theta, RF, TRssfp, T1, T2, M0, B0, phi0);

}

/////AUXILIAR FUNCTIONS
//Update augmented rotation matrix around z-axis
void update_Rz(double p)
{
    Rz(0,0) = cos(p);
    Rz(0,1) = -sin(p);
    Rz(1,0) = sin(p);
    Rz(1,1) = cos(p);
}

//Update derivative of the augmented rotation matrix around z-axis
void update_dRz(double p, double dp)
{
    dRz(0,0) = -dp * sin(p);
    dRz(0,1) = -dp * cos(p);
    dRz(1,0) =  dp * cos(p);
    dRz(1,1) = -dp * sin(p);
}

//Update augmented rotation matrix around x-axis
void update_Rx(double a)
{
    Rx(1,1) = cos(a);
    Rx(1,2) = sin(a);
    Rx(2,1) = -sin(a);
    Rx(2,2) = cos(a);
}

//Update augmented relaxation matrix and its derivatives
void update_E(double R1, double R2, double TR, double M0)
{
	E(0,0) = exp(-R2 * TR);
	E(1,1) = exp(-R2 * TR);
	E(2,2) = exp(-R1 * TR);
	E(2,3) = (1 - exp(-R1 * TR)) * M0;
}

void update_dEdT1(double R1, double TR, double M0)
{
    dEdT1(2,2) = TR * exp(-R1 * TR) * pow(R1,2.0);
    dEdT1(2,3) = -dEdT1(2,2) * M0;
}

void update_dEdT2(double R2, double TR)
{
    dEdT2(0,0) = TR * exp(-R2 * TR) * pow(R2,2.0);
    dEdT2(1,1) = dEdT2(0,0);
}

void update_dEdM0(double R1, double TR)
{
    dEdM0(2,3) = 1 - exp(-R1 * TR);
}

//Update Fisher Information Matrix
void add2FIM()
{
    J(0,0) = dmdT1.real();
    J(1,0) = dmdT1.imag();
    J(0,1) = dmdT2.real();
    J(1,1) = dmdT2.imag();
    J(0,2) = dmdM0.real();
    J(1,2) = dmdM0.imag();
    J(0,3) = dmdB0.real();
    J(1,3) = dmdB0.imag();
    J(0,4) = dmdP0.real();
    J(1,4) = dmdP0.imag();
    F += J.transpose() * J;
}


/**
* This file is part of CubeSLAM
*
* Copyright (C) 2018  Shichao Yang (Carnegie Mellon Univ)
*/

#include "Thirdparty/g2o/g2o/types/types_six_dof_expmap.h"
#include "detect_3d_cuboid/matrix_utils.h"

#include "g2o_Object.h"

#include <Eigen/Core>
#include <Eigen/Geometry>
#include <Eigen/Dense>
#include <math.h>
#include <algorithm> // std::swap

namespace g2o
{

using namespace Eigen;
using namespace std;

SE3Quat exptwist_norollpitch(const Vector6d &update)
{
    Vector3d omega;
    for (int i = 0; i < 3; i++)
        omega[i] = update[i];
    Vector3d upsilon;
    for (int i = 0; i < 3; i++)
        upsilon[i] = update[i + 3];

    double theta = omega.norm();
    Matrix3d Omega = skew(omega);

    Matrix3d R;
    R << cos(omega(2)), -sin(omega(2)), 0,
        sin(omega(2)), cos(omega(2)), 0,
        0, 0, 1;

    Matrix3d V;
    if (theta < 0.00001)
    {
        V = R;
    }
    else
    {
        Matrix3d Omega2 = Omega * Omega;

        V = (Matrix3d::Identity() + (1 - cos(theta)) / (theta * theta) * Omega + (theta - sin(theta)) / (pow(theta, 3)) * Omega2);
    }

    return SE3Quat(Quaterniond(R), V * upsilon);
}

void VertexCuboid::oplusImpl(const double *update_)
{
    Eigen::Map<const Vector9d> update(update_);

    g2o::cuboid newcube;
    if (whether_fixrotation)
    {
        newcube.pose.setTranslation(_estimate.pose.translation() + update.segment<3>(3));
    }
    else if (whether_fixrollpitch) //NOTE this only works for cuboid already has parallel to ground. otherwise update_z will also change final RPY
    {
        Vector9d update2 = update;
        update2(0) = 0;
        update2(1) = 0;
        newcube.pose = _estimate.pose * exptwist_norollpitch(update2.head<6>()); //NOTE object pose is from object to world!!!!
    }
    else
        newcube.pose = _estimate.pose * SE3Quat::exp(update.head<6>());

    if (whether_fixheight) // use previous height
        newcube.setTranslation(Vector3d(newcube.translation()(0), _estimate.translation()(1), newcube.translation()(2)));

    if (fixedscale(0) > 0) // if fixed scale is set, use it.
        newcube.scale = fixedscale;
    else
        newcube.scale = _estimate.scale + update.tail<3>();

    setEstimate(newcube);
}

// similar as above
void VertexCuboidFixScale::oplusImpl(const double *update_)
{
    Eigen::Map<const Vector6d> update(update_);

    g2o::cuboid newcube;
    if (whether_fixrotation)
    {
        newcube.pose.setRotation(_estimate.pose.rotation());
        newcube.pose.setTranslation(_estimate.pose.translation() + update.tail<3>());
    }
    else if (whether_fixrollpitch)
    {
        Vector6d update2 = update;
        update2(0) = 0;
        update2(1) = 0;
        newcube.pose = _estimate.pose * exptwist_norollpitch(update2);
    }
    else
        newcube.pose = _estimate.pose * SE3Quat::exp(update);

    if (whether_fixheight)
        newcube.setTranslation(Vector3d(newcube.translation()(0), _estimate.translation()(1), newcube.translation()(2)));

    if (fixedscale(0) > 0)
        newcube.scale = fixedscale;
    else
        newcube.scale = _estimate.scale;

    setEstimate(newcube);
}

void EdgeSE3CuboidFixScaleProj::computeError()
{
    const VertexSE3Expmap *SE3Vertex = dynamic_cast<const VertexSE3Expmap *>(_vertices[0]);              //  world to camera pose
    const VertexCuboidFixScale *cuboidVertex = dynamic_cast<const VertexCuboidFixScale *>(_vertices[1]); //  object pose to world

    SE3Quat cam_pose_Tcw = SE3Vertex->estimate();
    cuboid global_cube = cuboidVertex->estimate();

    Vector4d rect_project = global_cube.projectOntoImageBbox(cam_pose_Tcw, Kalib); // center, width, height
    _error = rect_project - _measurement;
}

double EdgeSE3CuboidFixScaleProj::get_error_norm()
{
    computeError();
    return _error.norm();
}

void EdgeSE3CuboidProj::computeError()
{
    const VertexSE3Expmap *SE3Vertex = dynamic_cast<const VertexSE3Expmap *>(_vertices[0]); //  world to camera pose
    const VertexCuboid *cuboidVertex = dynamic_cast<const VertexCuboid *>(_vertices[1]);    //  object pose to world

    SE3Quat cam_pose_Tcw = SE3Vertex->estimate();
    cuboid global_cube = cuboidVertex->estimate();

    Vector4d rect_project = global_cube.projectOntoImageBbox(cam_pose_Tcw, Kalib); // center, width, height

    _error = rect_project - _measurement;
}

double EdgeSE3CuboidProj::get_error_norm()
{
    computeError();
    return _error.norm();
}

void EdgeDynamicPointCuboidCamera::computeError()
{
    const VertexSE3Expmap *SE3Vertex = dynamic_cast<const VertexSE3Expmap *>(_vertices[0]);              // world to camera pose
    const VertexCuboidFixScale *cuboidVertex = dynamic_cast<const VertexCuboidFixScale *>(_vertices[1]); // object to world pose
    const VertexSBAPointXYZ *pointVertex = dynamic_cast<const VertexSBAPointXYZ *>(_vertices[2]);        // point to object pose

    Vector3d localpt = SE3Vertex->estimate() * (cuboidVertex->estimate().pose * pointVertex->estimate());

    Vector2d projected = Vector2d(Kalib(0, 2) + Kalib(0, 0) * localpt(0) / localpt(2), Kalib(1, 2) + Kalib(1, 1) * localpt(1) / localpt(2));
    _error = _measurement - projected;
}

void EdgeDynamicPointCuboidCamera::linearizeOplus()
{
    const VertexSE3Expmap *SE3Vertex = dynamic_cast<const VertexSE3Expmap *>(_vertices[0]);              // world to camera pose
    const VertexCuboidFixScale *cuboidVertex = dynamic_cast<const VertexCuboidFixScale *>(_vertices[1]); // object to world pose
    const VertexSBAPointXYZ *pointVertex = dynamic_cast<const VertexSBAPointXYZ *>(_vertices[2]);        // point to object pose

    Vector3d objectpt = pointVertex->estimate();
    SE3Quat combinedT = SE3Vertex->estimate() * cuboidVertex->estimate().pose;
    Vector3d camerapt = combinedT * objectpt;

    double fx = Kalib(0, 0);
    double fy = Kalib(1, 1);

    double x = camerapt[0];
    double y = camerapt[1];
    double z = camerapt[2];
    double z_2 = z * z;

    Matrix<double, 2, 3> projptVscamerapt; //2d projected pixel / 3D local camera pt
    projptVscamerapt(0, 0) = fx / z;
    projptVscamerapt(0, 1) = 0;
    projptVscamerapt(0, 2) = -x * fx / z_2;

    projptVscamerapt(1, 0) = 0;
    projptVscamerapt(1, 1) = fy / z;
    projptVscamerapt(1, 2) = -y * fy / z_2;

    // jacobian of point
    _jacobianOplus[2] = -projptVscamerapt * combinedT.rotation().toRotationMatrix();

    // jacobian of camera
    _jacobianOplus[0](0, 0) = x * y / z_2 * fx;
    _jacobianOplus[0](0, 1) = -(1 + (x * x / z_2)) * fx;
    _jacobianOplus[0](0, 2) = y / z * fx;
    _jacobianOplus[0](0, 3) = -1. / z * fx;
    _jacobianOplus[0](0, 4) = 0;
    _jacobianOplus[0](0, 5) = x / z_2 * fx;

    _jacobianOplus[0](1, 0) = (1 + y * y / z_2) * fy;
    _jacobianOplus[0](1, 1) = -x * y / z_2 * fy;
    _jacobianOplus[0](1, 2) = -x / z * fy;
    _jacobianOplus[0](1, 3) = 0;
    _jacobianOplus[0](1, 4) = -1. / z * fy;
    _jacobianOplus[0](1, 5) = y / z_2 * fy;

    // jacobian of object pose.   obj twist  [angle position]
    Matrix<double, 3, 6> skewjaco;
    skewjaco.leftCols<3>() = -skew(objectpt);
    skewjaco.rightCols<3>() = Matrix3d::Identity();
    _jacobianOplus[1] = _jacobianOplus[2] * skewjaco; //2*6
    if (cuboidVertex->whether_fixrollpitch)           //zero gradient for roll/pitch
    {
        _jacobianOplus[1](0, 0) = 0;
        _jacobianOplus[1](0, 1) = 0;
        _jacobianOplus[1](1, 0) = 0;
        _jacobianOplus[1](1, 1) = 0;
    }
    if (cuboidVertex->whether_fixrotation)
    {
        _jacobianOplus[1](0, 0) = 0;
        _jacobianOplus[1](0, 1) = 0;
        _jacobianOplus[1](1, 0) = 0;
        _jacobianOplus[1](1, 1) = 0;
        _jacobianOplus[1](0, 2) = 0;
        _jacobianOplus[1](1, 2) = 0;
    }
}

Vector2d EdgeDynamicPointCuboidCamera::computeError_debug()
{
    computeError();
    return _error;
}

void EdgeObjectMotion::computeError()
{
    const VertexCuboidFixScale *cuboidVertexfrom = dynamic_cast<const VertexCuboidFixScale *>(_vertices[0]);   // object to world pose
    const VertexCuboidFixScale *cuboidVertexto = dynamic_cast<const VertexCuboidFixScale *>(_vertices[1]);     // object to world pose
    const VelocityPlanarVelocity *velocityVertex = dynamic_cast<const VelocityPlanarVelocity *>(_vertices[2]); // object to world pose

    if (cuboidVertexfrom == nullptr || cuboidVertexto == nullptr || velocityVertex == nullptr)
        cout << "Bad casting when compute Edge motion error!!!!!!!!!!!!!" << endl;

    // predict motion x y yaw and compute measurement.
    SE3Quat posefrom = cuboidVertexfrom->estimate().pose;
    double yaw_from = posefrom.toXYZPRYVector()(5);

    SE3Quat poseto = cuboidVertexto->estimate().pose;
    double yaw_to = poseto.toXYZPRYVector()(5);

    Vector2d velocity = velocityVertex->estimate(); //v w   linear velocity and steer angle

    const double vehicle_length = 2.71; // front and back wheels distance
    // vehicle motion model is applied to back wheel center
    Vector3d trans_back_pred = posefrom.translation() + (velocity(0) * delta_t - vehicle_length * 0.5) * Vector3d(cos(yaw_from), sin(yaw_from), 0);
    double yaw_pred = yaw_from + tan(velocity(1)) * delta_t / vehicle_length * velocity(0);

    // as mentioned in paper: my object frame is at the center. the motion model applies to back wheen center. have offset.
    Vector3d trans_pred = trans_back_pred + vehicle_length * 0.5 * Vector3d(cos(yaw_pred), sin(yaw_pred), 0);

    _error = Vector3d(poseto.translation()[0], poseto.translation()[1], yaw_to) - Vector3d(trans_pred(0), trans_pred(1), yaw_pred);
    if (_error[2] > 2.0 * M_PI)
        _error[2] -= 2.0 * M_PI;
    if (_error[2] < -2.0 * M_PI)
        _error[2] += 2.0 * M_PI;
}

Vector3d EdgeObjectMotion::computeError_debug()
{
    computeError();
    return _error;
}

Vector3d cuboid::point_boundary_error(const Vector3d &world_point, const double max_outside_margin_ratio, double point_scale) const
{
    // transform the point to local object frame  TODO actually can compute gradient analytically...
    Vector3d local_pt = point_scale * (this->pose.inverse() * world_point).cwiseAbs(); // change global point to local cuboid body frame.  make it positive.
    Vector3d error;

    // if point is within the cube, error=0, otherwise penalty how far it is outside cube
    for (int i = 0; i < 3; i++)
    {
        if (local_pt(i) < this->scale(i))
            error(i) = 0;
        else if (local_pt(i) < (max_outside_margin_ratio + 1) * this->scale(i))
            error(i) = local_pt(i) - this->scale(i);
        else
            error(i) = max_outside_margin_ratio * this->scale(i); // if points two far, give a constant error, don't optimize.
    }

    return error;
}

void EdgePointCuboidOnlyObject::computeError()
{
    const VertexCuboid *cuboidVertex = dynamic_cast<const VertexCuboid *>(_vertices[0]); // world to camera pose

    _error.setZero();

    const g2o::cuboid &estimate_cube = cuboidVertex->estimate();

    Vector3d point_edge_error;
    point_edge_error.setZero();
    for (size_t i = 0; i < object_points.size(); i++) // use abs  otherwise   pos neg will counteract by different pts.     maybe each edge one pt?
        point_edge_error += estimate_cube.point_boundary_error(object_points[i], max_outside_margin_ratio).cwiseAbs();
    if (object_points.size() > 0)
        point_edge_error = point_edge_error / object_points.size();

    point_edge_error = point_edge_error.array() / estimate_cube.scale.array(); //scale it

    // add prior shape dimension error?
    double prior_weight = 0.2;
    Vector3d prior_shape_error = estimate_cube.scale; // setZero?  or penalize large box! or set a range?
    if (prior_object_half_size(0) > 0)                // if prior shape is being set, such as KITTI, then give large weight for shape error
    {
        prior_weight = 50.0;
        prior_shape_error = ((estimate_cube.scale - prior_object_half_size).array() / prior_object_half_size.array()).cwiseAbs();
    }

    _error = 1.0 * point_edge_error + prior_weight * prior_shape_error;
}

Vector3d EdgePointCuboidOnlyObject::computeError_debug()
{
    computeError();
    return _error;
}

// similar as above
void EdgePointCuboidOnlyObjectFixScale::computeError()
{
    const VertexCuboidFixScale *cuboidVertex = dynamic_cast<const VertexCuboidFixScale *>(_vertices[0]); // world to camera pose

    _error.setZero();

    const g2o::cuboid &estimate_cube = cuboidVertex->estimate();

    Vector3d point_edge_error;
    point_edge_error.setZero();
    for (size_t i = 0; i < object_points.size(); i++)
        point_edge_error += estimate_cube.point_boundary_error(object_points[i], max_outside_margin_ratio).cwiseAbs();
    if (object_points.size() > 0)
        point_edge_error = point_edge_error / object_points.size();

    point_edge_error = point_edge_error.array() / estimate_cube.scale.array();

    _error = 1.0 * point_edge_error;
}

void EdgePointCuboid::computeError()
{
    const VertexSBAPointXYZ *pointVertex = dynamic_cast<const VertexSBAPointXYZ *>(_vertices[0]); // point position
    const VertexCuboid *cuboidVertex = dynamic_cast<const VertexCuboid *>(_vertices[1]);          //  object pose to world
    const g2o::cuboid estimate_cube = cuboidVertex->estimate();

    Vector3d point_edge_error = estimate_cube.point_boundary_error(pointVertex->estimate(), max_outside_margin_ratio).cwiseAbs(); // abs to add shape error
    point_edge_error = point_edge_error.array() / estimate_cube.scale.array();

    // add prior shape dimension error?
    double prior_weight = 0.2;
    Vector3d prior_shape_error = estimate_cube.scale; // setZero?  or penalize large box! or set a range?
    _error = 1.0 * point_edge_error + prior_weight * prior_shape_error;
}

void EdgePointCuboidFixScale::computeError()
{
    const VertexSBAPointXYZ *pointVertex = dynamic_cast<const VertexSBAPointXYZ *>(_vertices[0]);        // point position
    const VertexCuboidFixScale *cuboidVertex = dynamic_cast<const VertexCuboidFixScale *>(_vertices[1]); //  object pose to world
    const g2o::cuboid estimate_cube = cuboidVertex->estimate();

    Vector3d point_edge_error = estimate_cube.point_boundary_error(pointVertex->estimate(), max_outside_margin_ratio).cwiseAbs(); // abs to add shape error
    point_edge_error = point_edge_error.array() / estimate_cube.scale.array();

    _error = 1.0 * point_edge_error;
}

void UnaryLocalPoint::computeError()
{
    // transform the point to local object frame
    const VertexSBAPointXYZ *pointVertex = dynamic_cast<const VertexSBAPointXYZ *>(_vertices[0]); // point position
    Vector3d local_pt = pointVertex->estimate().cwiseAbs();                                       // make it positive.
    Vector3d point_edge_error;

    // if point is within the cube, point_edge_error=0, otherwise penalty how far it is outside cube
    for (int i = 0; i < 3; i++)
    {
        if (local_pt(i) < objectscale(i))
            point_edge_error(i) = 0;
        else if (local_pt(i) < (max_outside_margin_ratio + 1) * objectscale(i))
            point_edge_error(i) = local_pt(i) - objectscale(i);
        else
            point_edge_error(i) = max_outside_margin_ratio * objectscale(i); // if points two far, give a constant error, don't optimize.
    }

    _error = point_edge_error.array() / objectscale.array();
}

} // namespace g2o


#include <iostream>
#include <vector>
#include <Eigen/Dense>

#include "roadnet.hpp"

using namespace dubins_traffic;

int main()
{
    Eigen::Vector3d transform( 1.0, 0.0, 1.6 );
    dubins_traffic::RoadNetwork rd( 2.0, transform, 2, 3 );
    std::cout << rd << std::endl;

    double x, y;
    rd.map_point( 0, 0, x, y );
    std::cerr << "(0, 0) -> (" << x << ", " << y << ")" << std::endl;

    rd.map_point( 1, 2, x, y );
    std::cerr << "(1, 2) -> (" << x << ", " << y << ")" << std::endl;

    for (size_t idx = 0; idx < rd.number_of_segments(); idx++)
        std::cerr << rd.mapped_segment( idx ).transpose() << std::endl;

    return 0;
}


//==================================================================================================
/*!
  @file

  @copyright 2016 NumScale SAS

  Distributed under the Boost Software License, Version 1.0.
  (See accompanying file LICENSE.md or copy at http://boost.org/LICENSE_1_0.txt)
*/
//==================================================================================================
#ifndef BOOST_SIMD_ARCH_COMMON_SIMD_FUNCTION_SINC_HPP_INCLUDED
#define BOOST_SIMD_ARCH_COMMON_SIMD_FUNCTION_SINC_HPP_INCLUDED

#include <boost/simd/detail/overload.hpp>
#include <boost/simd/constant/one.hpp>
#include <boost/simd/function/divides.hpp>
#include <boost/simd/function/if_else.hpp>
#include <boost/simd/function/is_eqz.hpp>
#include <boost/simd/function/sin.hpp>

#if !defined(BOOST_SIMD_NO_DENORMALS)
#include <boost/simd/constant/eps.hpp>
#include <boost/simd/function/abs.hpp>
#include <boost/simd/function/is_less.hpp>
#endif

#if !defined(BOOST_SIMD_NO_INFINITIES)
#include <boost/simd/function/if_zero_else.hpp>
#include <boost/simd/function/is_inf.hpp>
#endif

namespace boost { namespace simd { namespace ext
{
   namespace bd = boost::dispatch;
   namespace bs = boost::simd;
   BOOST_DISPATCH_OVERLOAD_IF( sinc_
                          , (typename A0, typename X)
                          , (detail::is_native<X>)
                          , bd::cpu_
                          , bs::pack_<bd::floating_<A0>, X>
                          )
   {
      BOOST_FORCEINLINE A0 operator()( const A0& a0) const BOOST_NOEXCEPT
      {
        A0 r1 =  bs::sin(a0)/a0;
        #if !defined(BOOST_SIMD_NO_DENORMALS)
        r1 = bs::if_else ( bs::is_less(bs::abs(a0), bs::Eps<A0>())
                          , bs::One<A0>()
                          , r1
                          );
        #else
        r1 = bs::if_else(bs::is_eqz(a0), bs::One<A0>(), r1);
        #endif
        #if !defined(BOOST_SIMD_NO_INFINITIES)
        r1 = bs::if_zero_else(bs::is_inf(a0), r1);
        #endif
        return r1;
      }
   };

} } }

#endif


#include <iostream>

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

#include <Eigen/Core>
#include <opencv2/opencv.hpp>
#include <opencv2/core/eigen.hpp>
// #include <cuda_runtime.h>

namespace py = pybind11;
using namespace std;

using Eigen::Dynamic;
using Eigen::RowMajor;
using Eigen::Unaligned;

const double DEGREE2RADIAN = M_PI / 180.0;

void process_image(
    const cv::Mat &pano,
    cv::Mat &pers,
    const cv::Mat_<double> &rot,
    const cv::Mat_<double> &K);

/*
 Create instrinsic parameter
 */
cv::Mat_<double> get_intrinsic_param(cv::Size &size_per)
{
    double fov = 90.0 * DEGREE2RADIAN;
    double focal_length = (double)size_per.width/(2. * tan(fov / 2.));
    cv::Mat_<double> K = (
        cv::Mat_<double>(3, 3) <<
        focal_length, 0., (double)size_per.width/2.,
        0., focal_length, (double)size_per.height/2.,
        0., 0., 1.);
    return K;
}

/*
 Get rotation matrix from the direction given
 - Input: an array of double (3 rotation angles)
 - Output: rotation matrix
 */
cv::Mat_<double> angle2RotMat(array<double, 3> &rot_angle)
{
	cv::Mat_<double> Th2c_key = (cv::Mat_<double>(3,3) <<
		1., 0., 0.,
		0., cos(rot_angle[0]), sin(rot_angle[0]),
		0., -sin(rot_angle[0]), cos(rot_angle[0]))
		*(cv::Mat_<double>(3,3) <<
		cos(rot_angle[1]), 0., -sin(rot_angle[1]),
		0., 1., 0.,
		sin(rot_angle[1]), 0., cos(rot_angle[1]))
		*(cv::Mat_<double>(3,3) <<
        cos(rot_angle[2]), sin(rot_angle[2]), 	0.,
		-sin(rot_angle[2]), cos(rot_angle[2]), 	0.,
		0., 0., 1.);
    return Th2c_key;
}

/*
 Get x, y pixel location of panorama image
 from theta (rotation around y-axis) and
 phi (rotation around x-axis)
 */
cv::Mat get_image(
    array<Eigen::Matrix<unsigned char, Dynamic, Dynamic, RowMajor>, 3> &src,
    array<double, 3> &angles)
{
    cv::Mat rgb[3];
    cv::Mat im_pano(src[0].rows(), src[0].cols(), CV_8UC3);
    for (int i=0; i<src.size(); i++)
    {
        cv::Mat channel(src[i].rows(), src[i].cols(), CV_8UC1, src[i].data());
        rgb[i] = channel;
    }
    cv::merge(rgb, 3, im_pano);

    // Get intrinsic parameter:
    const int height = 360; // (int)im_pano.rows/4;
    const int width = 640; // (int)im_pano.cols/4;
    cv::Size size_per = cv::Size(width, height);
    cv::Mat_<double> K = get_intrinsic_param(size_per);
    //cout << "Intrinsic Param: " << K << endl;

    cv::Mat_<double> rotation = angle2RotMat(angles); // make it front for now
    //cout << "Rotation Matrix: " << rotation << endl;

    // Create perspective image
    cv::Mat im_perspective = cv::Mat::zeros(cv::Size(width, height), CV_8UC3);
    process_image(im_pano, im_perspective, K, rotation);

    return im_perspective;
}

PYBIND11_MODULE(extension, m)
{
    m.def("get_image", &get_image);

    // Cuffer protocol for return value
    py::class_<cv::Mat>(m, "Image", py::buffer_protocol())
        .def_buffer([](cv::Mat& im) -> py::buffer_info{
            return py::buffer_info(
                // pointer to buffer
                im.data,
                //size of one scalar
                sizeof(unsigned char),
                // Python struct-style format descriptor
                py::format_descriptor<unsigned char>::format(),
                // Number of dimensions
                3,
                // Buffer dimensions
                { im.rows, im.cols, im.channels() },
                // Strides (in bytes) for each index
                {
                    sizeof(unsigned char) * im.channels() * im.cols,
                    sizeof(unsigned char) * im.channels(),
                    sizeof(unsigned char)
                }
            );
        });
}