#ifndef INCLUDED_JPEZY_ENCODER_HPP
#define INCLUDED_JPEZY_ENCODER_HPP
#define _USE_MATH_DEFINES
#include "../jpezy.hpp"
#include "huffman_table.hpp"
#include "jpezy_writer.hpp"
#include <boost/range/algorithm/copy.hpp>
#include <boost/range/algorithm/fill.hpp>
#include <chrono>
#include <experimental/iterator>
#include <iostream>
#include <srook/algorithm/for_each.hpp>
#include <srook/array.hpp>
#include <srook/io/bofstream.hpp>
#include <srook/math/constants/algorithm/sqrt.hpp>
#include <srook/mpl/constant_sequence/math/make_costable.hpp>
#include <srook/iterator/range_access.hpp>
#include <vector>

namespace jpezy {

template <class T>
struct encoder {
    SROOK_CONSTEXPR encoder(const property& pr_, const std::vector<T>& r_, const std::vector<T>& g_, const std::vector<T>& b_)
        : pr(pr_)
        , r{ r_ }
        , g{ g_ }
        , b{ b_ }
        , Y_block{}
        , Cb_block{}
        , Cr_block{}
        , DCT_data{}
        , pre_DC{}
    {
        SROOK_ST_ASSERT(rgb_size == static_cast<std::size_t>(RGB::ELEMENT_SIZE));
    }

    template <class MODE_TAG = COLOR_MODE>
    std::size_t encode(const char* output_file)
    {
        int size = pr.get<property::At::HSize>() * pr.get<property::At::VSize>() * 3;
        if (size < 10240)
            size = 10240;

        jpezy_writer jpeg(size, pr, output_file);

        {
            raii_messenger mes("Write JPEG Header ...");
            jpeg.write_header();
        }

        {
            raii_messenger mes("Encoding ...");

            const int VUnits = (pr.get<property::At::VSize>() / 16) + ((pr.get<property::At::VSize>() % 16) ? 1 : 0);
            const int HUnits = (pr.get<property::At::HSize>() / 16) + ((pr.get<property::At::HSize>() % 16) ? 1 : 0);

            for (int y = 0; y < VUnits; ++y) {
                for (int x = 0; x < HUnits; ++x) {
                    make_YCC(x, y);
                    SROOK_IF_CONSTEXPR (std::is_same_v<MODE_TAG, GRAY_MODE>) {
                        boost::fill(Cb_block, 0);
                        boost::fill(Cr_block, 0);
                    }
                    make_MCU(jpeg.get_stream());
                }
            }
        }

        {
            raii_messenger mes("Write EOI ...");
            jpeg.write_eoi();
        }

        jpeg.output_file();
        return jpeg.get_stream().wrote_size();
    }

private:
    using value_type = int;

    static SROOK_CONSTEXPR value_type block = 8;
public:
	static SROOK_CONSTEXPR value_type block_size = block;
private:
    static SROOK_CONSTEXPR value_type blocks_size = block * block;
    static SROOK_CONSTEXPR value_type rgb_size = 3;
    static SROOK_CONSTEXPR value_type mcu_size = 4;

    void make_YCC(int ux, int uy) SROOK_NOEXCEPT_TRUE
    {
        SROOK_DECLTYPE(Y_block) Crblock{}, Cbblock{};

        for (std::size_t i = 0; i < mcu_size; ++i) {
            SROOK_DEDUCED_TYPENAME SROOK_DECLTYPE(Y_block)::value_type::iterator
                yp = srook::begin(Y_block[i]),
                cbp = srook::begin(Cbblock[i]),
                crp = srook::begin(Crblock[i]);

            for (std::size_t sy = uy * (block * 2) + ((i > 1) ? block : 0), sy_t = sy + block; sy < sy_t; ++sy) {
                const std::size_t ii = sy < pr.get<property::At::VSize>() ? sy : pr.get<property::At::VSize>() - 1;

                for (std::size_t sx = ux * (block * 2) + ((i & 1) ? block : 0), sx_t = sx + block; sx < sx_t; ++sx) {
                    const std::size_t jj = sx < pr.get<property::At::HSize>() ? sx : pr.get<property::At::HSize>() - 1;
                    const std::size_t index = ii * pr.get<property::At::HSize>() + jj;

                    const srook::byte &rv = r[index], gv = g[index], bv = b[index];

                    *yp++ = RGB::Y(rv, gv, bv);
                    *cbp++ = RGB::Cb(rv, gv, bv);
                    *crp++ = RGB::Cr(rv, gv, bv);
                }
            }
        }

        int n = 0;
        for (int i = 0; i < mcu_size; ++i) {
            switch (i) {
            case 0:
                n = 0;
                break;
            case 1:
                n = 4;
                break;
            case 2:
                n = 32;
                break;
            case 3:
                n = 36;
                break;
            default:
                break;
            }
            for (int y = 0; y < block; y += 2) {
                for (int x = 0; x < block; x += 2) {
                    const int index = y * block + x;
                    Cr_block[n] = Crblock[i][index];
                    Cb_block[n] = Cbblock[i][index];
                    ++n;
                }
                n += mcu_size;
            }
        }
    }

    template <class U, std::size_t s>
    void DCT(srook::array<U, s>& pic) SROOK_NOEXCEPT_TRUE
    {
        SROOK_CONSTEXPR_OR_CONST double dis_sqrt = 1.0 / srook::sqrt(2.0);

        for (int i = 0; i < block; ++i) {
            const double cv = i ? 1.0 : dis_sqrt;

            for (int j = 0; j < block; ++j) {
                const double cu = j ? 1.0 : dis_sqrt;
                double sum = 0;

                for (int y = 0; y < block; ++y) {
                    for (int x = 0; x < block; ++x) {
                        sum += pic[y * block + x] * cos_table[j * block_size + x] * cos_table[i * block_size + y];
                    }
                }
                DCT_data[i * block + j] = int(sum * cu * cv / 4);
            }
        }
    }

    void quantization(int cs) SROOK_NOEXCEPT_TRUE
    {
        const srook::array<int, 64>& qt = cs ? CQuantumTb : YQuantumTb;
        srook::for_each(srook::algorithm::make_counter(DCT_data), [&qt](auto& v, std::size_t i) { v /= qt[i]; });
    }

    template <class U, std::size_t dcs, std::size_t acs>
    void encode_huffman(int cs, srook::io::jpeg::bofstream& ofs, const huffmanCode_Tb<U, dcs>& dcT, const huffmanCode_Tb<U, acs>& acT, int eob_idx, int zrl_idx)
    {
        const SROOK_DECLTYPE(DCT_data)& dct_data = DCT_data;

        // DC
        const int diff = dct_data[0] - pre_DC[cs];
        pre_DC[cs] = dct_data[0];

        std::size_t di = 0;
        for (int abs_diff = std::abs(diff); abs_diff > 0; abs_diff >>= 1, ++di)
            ;
        if (di > static_cast<signed>(dcT.size()))
            throw std::runtime_error(__func__);

        (ofs | srook::io::jpeg::bofstream::Bits(dcT.size_tb[di])) << dcT.code_tb[di];
        if (di)
            (ofs | srook::io::jpeg::bofstream::Bits(di)) << (diff < 0 ? diff - 1 : diff);

        // AC
        for (std::size_t n = 1, run = 0; n < blocks_size; ++n) {
            int abs_coefficient = std::abs(dct_data[ZZ[n]]);

            if (abs_coefficient) {
                while (run > 15) {
                    (ofs | srook::io::jpeg::bofstream::Bits(acT.size_tb[zrl_idx])) << acT.code_tb[zrl_idx];
                    run -= 16;
                }
                int s = 0;
                for (; abs_coefficient > 0; abs_coefficient >>= 1, ++s)
                    ;

                const int a_di = run * 10 + s + (run == 15);
                if (a_di >= static_cast<signed>(acT.size()))
                    throw std::runtime_error(__func__);

                (ofs | srook::io::jpeg::bofstream::Bits(acT.size_tb[a_di])) << acT.code_tb[a_di];

                int v = dct_data[ZZ[n]];
                if (v < 0)
                    --v;

                (ofs | srook::io::jpeg::bofstream::Bits(s)) << v;
                run = 0;
            } else {
                if (n == 63)
                    (ofs | srook::io::jpeg::bofstream::Bits(acT.size_tb[eob_idx])) << acT.code_tb[eob_idx];
                else
                    ++run;
            }
        }
    }

    void make_MCU(srook::io::jpeg::bofstream& ofps)
    {
        for (auto& v : Y_block) {
            DCT(v);
            quantization(0);
            encode_huffman(0, ofps, YDcHuffmanT, YAcHuffmanT, YEOBidx, YZRLidx);
        }

        DCT(Cb_block);
        quantization(1);
        encode_huffman(1, ofps, CDcHuffmanT, CAcHuffmanT, CEOBidx, CZRLidx);

        DCT(Cr_block);
        quantization(2);
        encode_huffman(2, ofps, CDcHuffmanT, CAcHuffmanT, CEOBidx, CZRLidx);
    }

    struct RGB {
        static SROOK_CONSTEXPR int Y(srook::byte r, srook::byte g, srook::byte b)
        {
            return int((0.2990 * srook::to_integer<int>(r)) + (0.5870 * srook::to_integer<int>(g)) + (0.1140 * srook::to_integer<int>(b)) - 128);
        }
        static SROOK_CONSTEXPR int Cb(srook::byte r, srook::byte g, srook::byte b)
        {
            return int(-(0.1687 * srook::to_integer<int>(r)) - (0.3313 * srook::to_integer<int>(g)) + (0.5000 * srook::to_integer<int>(b)));
        }
        static SROOK_CONSTEXPR int Cr(srook::byte r, srook::byte g, srook::byte b)
        {
            return int((0.5000 * srook::to_integer<int>(r)) - (0.4187 * srook::to_integer<int>(g)) - (0.0813 * srook::to_integer<int>(b)));
        }
        enum {
            R,
            G,
            B,
            ELEMENT_SIZE
        };
    };

    const property& pr;
    const std::vector<T> r, g, b;

    srook::array<srook::array<value_type, blocks_size>, mcu_size> Y_block;
    srook::array<value_type, blocks_size> Cb_block, Cr_block;

    static SROOK_CONSTEXPR srook::array<const double, block * block> cos_table = srook::constant_sequence::math::unwrap_costable::array<srook::constant_sequence::math::make_costable_t<8,8>>::value;

    srook::array<value_type, blocks_size> DCT_data;
    srook::array<value_type, rgb_size> pre_DC;
};

template <class T>
encoder(const property&, const std::vector<T>&, const std::vector<T>&, const std::vector<T>&) -> encoder<T>;

} // namespace jpezy

#endif


#define PY_SSIZE_T_CLEAN
#include <Python.h>
#include <vector>
#include <boost/timer/timer.hpp>
#include <GeomLib.hpp>
#include <MPILib/include/MPINetworkCode.hpp>
#include <MPILib/include/RateAlgorithmCode.hpp>
#include <MPILib/include/SimulationRunParameter.hpp>
#include <MPILib/include/report/handler/InactiveReportHandler.hpp>
#include <MPILib/include/WilsonCowanAlgorithm.hpp>
#include <MPILib/include/RateFunctorCode.hpp>
#include <MPILib/include/utilities/ProgressBar.hpp>
#include <MPILib/include/BasicDefinitions.hpp>
#include <MPILib/include/WilsonCowanParameter.hpp>
#include <MPILib/include/WilsonCowanAlgorithm.hpp>
#include <MPILib/include/MiindTvbModelAbstract.hpp>

MPILib::Rate External_RateFunction(MPILib::Time t){
	return 1.0;
}
/* This class is designed to be used in conjunction with
 * Miind_WilsonCowan.py only. Any updates to this shared library
 * should be communicated to the TVB team and a new .so file should be provided
 * to be included in the tvb-library source tree (tvb-library/tvb/simulator/models/).
 */
class MiindWilsonCowan : public MPILib::MiindTvbModelAbstract<double, MPILib::utilities::CircularDistribution>{
private:
	std::vector<double> E_Initials = std::vector<double> (); // initial excitatory values
	std::vector<double> I_Initials = std::vector<double> (); // initial inhibitory values

public:

	MiindWilsonCowan(int num_nodes, double simulation_length, double dt) :
		MiindTvbModelAbstract(num_nodes, simulation_length) {
			this->_time_step = dt;
		}

	// In general, we won't need to set initial values but we implement it here
	// so that we can match TVB's Wilson Cowan example.
	void setInitialValues(std::vector<double> E_vals, std::vector<double> I_vals) {
		for(int i=0; i<_num_nodes; i++) {
			E_Initials.push_back(E_vals[i]);
			I_Initials.push_back(I_vals[i]);
		}
	}

	void init(std::vector<double> params)
	{
		// (TVB : tau_e) population time constant
		Time   E_tau       = params[4];
		// (TVB : c_e) maximum rate reached by the sigmoid function
		Rate   E_max_rate  = params[6];
		// (TVB : a_e) noise term modulates the input (a term in Wilson Cowan)
		double E_noise     = params[10];
		// (TVB : b_e) bias term shifts the input (theta term in Wilson Cowan)
		double E_bias      = params[12];
		// (TVB : r_e) smoothing term over output which models refractory dynamics
		double E_smoothing = params[8];
		// (TVB : tau_i) population time constant
		Time   I_tau       = params[5];
		// (TVB : c_i) maximum rate reached by the sigmoid function
		Rate   I_max_rate  = params[7];
		// (TVB : a_i) noise term modulates the input
		double I_noise     = params[11];
		// (TVB : b_i) bias term shifts the input
		double I_bias      = params[13];
		// (TVB : r_e) smoothing term over output which models refractory dynamics
		double I_smoothing = params[9];
		// (TVB : c_ee) Weight of E->E connection
		double E_E_Weight  = params[0];
		// (TVB : -c_ie) Weight of E->I connection
		double I_E_Weight  = -params[1];
		 // (TVB : c_ei) Weight of I->E connection
		double E_I_Weight  = params[2];
		// (TVB : -c_ii) Weight of I->I connection
		double I_I_Weight  = -params[3];
		// (TVB : P) Some additional drive to excitatory pop
		double P_E_Weight  = params[14];
		// (TVB : Q) Some additional drive to inhibitory pop
		double Q_I_Weight  = params[15];

		try {
			std::vector<NodeId> E_ids = std::vector<NodeId>();
			std::vector<NodeId> I_ids = std::vector<NodeId>();

			for(int i=0; i<_num_nodes; i++) {
				MPILib::WilsonCowanParameter E_param = MPILib::WilsonCowanParameter(
								E_tau, E_max_rate, E_noise, E_bias, E_Initials[i], E_smoothing);
				MPILib::WilsonCowanAlgorithm E_alg(E_param);

				MPILib::WilsonCowanParameter I_param = MPILib::WilsonCowanParameter(
								I_tau, I_max_rate, I_noise, I_bias, I_Initials[i], I_smoothing);
				MPILib::WilsonCowanAlgorithm I_alg(I_param);

				MPILib::Rate RateFunction_P(MPILib::Time);
				MPILib::RateFunctor<double> rate_functor_p(External_RateFunction);

				MPILib::Rate RateFunction_Q(MPILib::Time);
				MPILib::RateFunctor<double> rate_functor_q(External_RateFunction);

				MPILib::NodeId id_E = network.addNode(E_alg, MPILib::EXCITATORY_DIRECT);
				MPILib::NodeId id_I = network.addNode(I_alg, MPILib::INHIBITORY_DIRECT);
				MPILib::NodeId id_P = network.addNode(rate_functor_p, MPILib::NEUTRAL);
				MPILib::NodeId id_Q = network.addNode(rate_functor_q, MPILib::NEUTRAL);

				E_ids.push_back(id_E);
				I_ids.push_back(id_I);

				network.makeFirstInputOfSecond(id_E,id_E,E_E_Weight);
				network.makeFirstInputOfSecond(id_I,id_E,I_E_Weight);
				network.makeFirstInputOfSecond(id_E,id_I,E_I_Weight);
				network.makeFirstInputOfSecond(id_I,id_I,I_I_Weight);
				network.makeFirstInputOfSecond(id_P,id_E,P_E_Weight);
				network.makeFirstInputOfSecond(id_Q,id_I,Q_I_Weight);

			}

			// Set each node to have an external successor (for coupling input from TVB)
			for(auto& id : E_ids) {
				network.setNodeExternalSuccessor(id);
				network.setNodeExternalPrecursor(id, 1);
			}

			for(auto& id : I_ids) {
				network.setNodeExternalSuccessor(id);
				network.setNodeExternalPrecursor(id, 1);
			}

			std::string sim_name = "miind_wc";
			MPILib::report::handler::InactiveReportHandler handler =
											MPILib::report::handler::InactiveReportHandler();

			SimulationRunParameter par_run( handler,(_simulation_length/_time_step)+1,0,
											_simulation_length,_time_step,_time_step,sim_name,_time_step);

			network.configureSimulation(par_run);

		} catch(std::exception& exc){
			std::cout << exc.what() << std::endl;
		}
	}
};

static MiindWilsonCowan *model;

static PyObject *miind_init(PyObject *self, PyObject *args)
{
    int nodes;
    double sim_time;
    double time_step;

    if (!PyArg_ParseTuple(args, "idd", &nodes, &sim_time, &time_step))
	return NULL;

    model = new MiindWilsonCowan(nodes, sim_time, time_step);

    Py_INCREF(Py_None);
    return Py_None;
}

static PyObject *miind_initParams(PyObject *self, PyObject *args)
{
    PyObject *float_list;
    int pr_length;

    if (!PyArg_ParseTuple(args, "O", &float_list))
        return NULL;
    pr_length = PyObject_Length(float_list);
    if (pr_length < 0)
        return NULL;

    std::vector<double> params(pr_length);

    for (int index = 0; index < pr_length; index++) {
        PyObject *item;
        item = PyList_GetItem(float_list, index);
        if (!PyFloat_Check(item))
            params[index] = 0.0;
        params[index] = PyFloat_AsDouble(item);
    }

    model->init(params);

    Py_INCREF(Py_None);
    return Py_None;
}

static PyObject *miind_setInitialValues(PyObject *self, PyObject *args)
{
    PyObject *float_listE;
    PyObject *float_listI;
    int pr_length;

    if (!PyArg_ParseTuple(args, "OO", &float_listE, &float_listI))
        return NULL;
    pr_length = PyObject_Length(float_listE);
    if (pr_length < 0)
        return NULL;

    std::vector<double> insE(pr_length);
    std::vector<double> insI(pr_length);

    for (int index = 0; index < pr_length; index++) {
        PyObject *item;
        item = PyList_GetItem(float_listE, index);
        if (!PyFloat_Check(item))
            insE[index] = 0.0;
        insE[index] = PyFloat_AsDouble(item);
    }

    for (int index = 0; index < pr_length; index++) {
        PyObject *item;
        item = PyList_GetItem(float_listI, index);
        if (!PyFloat_Check(item))
            insI[index] = 0.0;
        insI[index] = PyFloat_AsDouble(item);
    }

    model->setInitialValues(insE,insI);

    Py_INCREF(Py_None);
    return Py_None;
}

static PyObject *miind_getTimeStep(PyObject *self, PyObject *args)
{
    return Py_BuildValue("d", model->getTimeStep());
}

static PyObject *miind_getSimulationLength(PyObject *self, PyObject *args)
{
    return Py_BuildValue("d", model->getSimulationLength());
}

static PyObject *miind_startSimulation(PyObject *self, PyObject *args)
{
    model->startSimulation();
    Py_INCREF(Py_None);
    return Py_None;
}

static PyObject *miind_evolveSingleStep(PyObject *self, PyObject *args)
{
    PyObject *float_list;
    int pr_length;

    if (!PyArg_ParseTuple(args, "O", &float_list))
        return NULL;
    pr_length = PyObject_Length(float_list);
    if (pr_length < 0)
        return NULL;

    std::vector<double> activities(pr_length);

    for (int index = 0; index < pr_length; index++) {
        PyObject *item;
        item = PyList_GetItem(float_list, index);
        if (!PyFloat_Check(item))
            activities[index] = 0.0;
        activities[index] = PyFloat_AsDouble(item);
    }

    std::vector<double> out_activities = model->evolveSingleStep(activities);

    PyObject* tuple = PyTuple_New(pr_length);

    for (int index = 0; index < pr_length; index++) {
        PyTuple_SetItem(tuple, index, Py_BuildValue("d", out_activities[index]));
    }

    return tuple;
}

static PyObject *miind_endSimulation(PyObject *self, PyObject *args)
{
    model->endSimulation();
    Py_INCREF(Py_None);
    return Py_None;
}

static PyMethodDef MiindModelMethods[] = {
    {"init",  miind_init, METH_VARARGS, "Init Miind model."},
    {"initParams",  miind_initParams, METH_VARARGS, "Init Miind model parameters."},
    {"getTimeStep",  miind_getTimeStep, METH_VARARGS, "Get time step."},
    {"getSimulationLength",  miind_getSimulationLength, METH_VARARGS, "Get sim time."},
    {"startSimulation",  miind_startSimulation, METH_VARARGS, "Start simulation."},
    {"evolveSingleStep",  miind_evolveSingleStep, METH_VARARGS, "Evolve one time step."},
    {"endSimulation",  miind_endSimulation, METH_VARARGS, "Clean up."},
    {"setInitialValues", miind_setInitialValues, METH_VARARGS, "Set initial values."},
    {NULL, NULL, 0, NULL}        /* Sentinel */
};

static struct PyModuleDef miindmodule = {
    PyModuleDef_HEAD_INIT,
    "libmiindwc",   /* name of module */
    NULL, /* module documentation, may be NULL */
    -1,       /* size of per-interpreter state of the module,
                 or -1 if the module keeps state in global variables. */
    MiindModelMethods
};

PyMODINIT_FUNC
PyInit_libmiindwc(void)
{
    return PyModule_Create(&miindmodule);
}


/* * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * *
 * Filename    : sphdec.cpp                                                                            *
 * Project     : Planewalker - Schnorr-Euchner sphere decoder simulation for space-time lattice codes  *
 * Authors     : Pasi Pyrrö, Oliver Gnilke                                                             *
 * Version     : 1.0                                                                                *
 * Copyright   : Aalto University ~ School of Science ~ Department of Mathematics and Systems Analysis *
 * Date        : 9.1.2017                                                                              *
 * Language    : C++ (2011 or newer standard)                                                          *
 * Description : The core algorithm implementations for the simulation + wrapper function for them     *
 * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * */

#define ARMA_NO_DEBUG /* disable Armadillo bound checks for addiotional speed */

#include <iostream>
#include <armadillo> /* linear algebra library */
#include <vector>
#include <complex>
#include <string>

#include "sphdec.hpp"
#include "misc.hpp"
#include "algorithms.hpp"

using namespace std;
using namespace arma;

/* Decision feedback equalization on xt[i]
   i.e. Use the Babai nearest plane algorithm as a starting point in dimension i */
inline void dfe(int i, int q, vec &xt, const vec &y, vec &delta, const vec &ksi, const mat &R, const vector<int> &S) {
    /* calculate the "Babai coeffient" and "round" it to nearest symbol in the symbol set */
    xt[i] = nearest_symbol((y[i]-ksi[i])/R(i,i), S);
    if (xt[i] <= S[0]) { /* only feasible way is positive direction */
        delta[i] = 2;
    } else if (xt[i] >= S[q - 1]) { /* only feasible way is negative direction */
        delta[i] = -2;
    } else { /* there's room to enumerate within signal set */
        delta[i] = 2*sesd_sign(y[i]-ksi[i]-R(i,i)*xt[i]);
    }
}

/* Schnorr-Euchner enumeration step for the sphere decoder i.e. zig-zag around the dfe point.
   Makes sure we loop over the points in dimension i in non descending order.
   This enables early stopping criteria when we find a point in dimension i outside the radius,
   i.e. no points in that sub search tree are inside the sphere because they can only be further away or at the same distance. */
inline void se_enum(int i, vec &xt, vec &delta) {
    /* add delta to xt and pick new delta (difference to next symbol in S we want to try) */
    xt[i] = xt[i] + delta[i];
    delta[i] = -delta[i] - 2*sesd_sign(delta[i]);
}

/* Perform some basic checks to ensure that the sphere decoder works as intended */
inline bool check(double radius, const mat &R, const vec &y) {
    if (radius <= 0){
        log_msg("sphdec: negative squared initial radius given!", "Error");
        log_msg("aborting simulation round...");
        return false;
    }

    if (R.n_rows != R.n_cols) {
        log_msg("sphdec: R is not a square matrix!", "Error");
        log_msg("aborting simulation round...");
        return false;
    }

    if (y.n_elem != R.n_cols) {
        log_msg("sphdec: vector y dimension mismatch!", "Error");
        log_msg("aborting simulation round...");
        return false;
    }
    return true;
}

/* Basic sphere decoder algorithm */
vector<int> sphdec(const vec &y, const mat &R, const vector<int> &S, int &counter, double radius){

    /* Initialize */
    int k = params["no_of_matrices"];
    int q = params["x-PAM"];
    int i = k-1; /* We start from dimension k-1 and iterate all the way down to dimension 0 */

    vector<int> x(k); /* point to decode */
    vec xt(k), ksi(k), delta(k), dist(k);
    xt.zeros(); ksi.zeros(); delta.zeros(); dist.zeros();
    counter = 0; /* counts how many search tree nodes (loop iterations) we went through */

    vec distances(k, fill::zeros);

    double xidist = 0.0;
    bool found = false;

    if (!check(radius, R, y))
        return vector<int>(0);

    /* Initialize xt[k-1] and delta[k-1] */
    dfe(i, q, xt, y, delta, ksi, R, S);

    while (!exit_flag) {

        counter++;
        /* Step 3. */
        xidist = pow(y[i]-ksi[i]-R(i,i)*xt[i], 2);

        distances[i] = dist[i] + xidist;

        /***** Uncomment line below to debug *****/
        // cout << i << ": " << vec2str(xt, k) << ", xidist = " << xidist + dist[i] << ", C = " << radius << endl;

        if (radius < dist[i] + xidist) { // current point xt is outside the sphere
            // Step 4.
            if (i == k-1) {
                break;
            } else {
                // Step 6.
                i++;
                se_enum(i, xt, delta);
            }
        } else { // we are inside the sphere
            if (xt[i] < S[0] || xt[i] > S[q - 1]){ // we are outside the signal set boundaries
                if ((xt[i] < S[0] && (xt[i] + delta[i]) > S[q - 1]) || (xt[i] > S[q - 1] && (xt[i] + delta[i]) < S[0])){
                    // Step 4.
                    if (i == k-1) {
                        break;
                    } else {
                        i++;
                        se_enum(i, xt, delta);
                    }
                } else {
                    se_enum(i, xt, delta);
                }
            } else {
                if (i > 0) {
                    ksi[i-1] = 0;
                    for (int j = i; j < k; j++)
                        ksi[i-1] += R(i-1, j)*xt[j];
                    dist[i-1] = dist[i] + xidist;
                    i--;
                    dfe(i, q, xt, y, delta, ksi, R, S);
                } else { // lattice point is found (Step 5)
                    radius = dist[0] + xidist;
                    found = true;
                    x = conv_to<vector<int>>::from(xt);
                    i++;
                    se_enum(i, xt, delta);
                }
            }
        }
    }
    if (found)
        return x;
    else {
        log_msg("Initial squared radius used: " + to_string(radius), "Alert");
        log_msg("distances considered: " + vec2str(distances, distances.size()), "Alert");
        log_msg("sphdec: point not found!", "Alert");
        return vector<int>(0);
    }
}

/* Sphere decoder algorithm that considers the spherical shaping of the codebook */
vector<int> sphdec_spherical_shaping(const vec &y, const mat &HR, const mat &R, const vector<int> &S,
                                     int &counter, double P, double radius){

    /* Initialize */
    int k = params["no_of_matrices"];
    int q = params["x-PAM"];
    int i = k-1; /* We start from dimension k-1 and iterate all the way down to dimension 0 */

    vector<int> x(k); /* point to decode */
    vec xt(k), ksi(k), delta(k), dist(k), curr(k), ener(k);
    xt.zeros(); ksi.zeros(); delta.zeros(); dist.zeros(); curr.zeros(); ener.zeros();
    counter = 0; /* counts how many search tree nodes (loop iterations) we went through */

    vec distances(k, fill::zeros);

    double xidist = 0.0, xiener = 0.0;
    bool found = false;

    if (!check(radius, HR, y))
        return vector<int>(0);

    /* Initialize xt[k-1] and delta[k-1] */
    dfe(i, q, xt, y, delta, ksi, HR, S);

    while (!exit_flag) {

        counter++;
        // Step 3.
        xidist = pow(y[i]-ksi[i]-HR(i,i)*xt[i], 2);

        distances[i] = dist[i] + xidist;

        /***** Uncomment line below to debug *****/
        // cout << i << ": " << vec2str(xt, k) << ", xidist = " << xidist + dist[i] << ", C = " << radius << endl;

        if (radius < dist[i] + xidist) { /* current point xt is outside the sphere */
            // Step 4.
            if (i == k-1) {
                break;
            } else {
                // Step 6.
                i++;
                se_enum(i, xt, delta);
            }
        } else { /* we are inside the sphere */

            xiener = pow(xt[i]*R(i,i) + curr[i], 2);

            // we are outside the signal set boundaries
            if (xt[i] < S[0] || xt[i] > S[q - 1] || ener[i] + xiener > P + 10e-6) {
                if ((xt[i] < S[0] && (xt[i] + delta[i]) > S[q - 1]) || (xt[i] > S[q - 1] && (xt[i] + delta[i]) < S[0])){
                    // Step 4.
                    if (i == k-1) {
                        break;
                    } else {
                        i++;
                        se_enum(i, xt, delta);
                    }
                } else
                    se_enum(i, xt, delta);
            } else {
                if (i > 0) {
                    ksi[i-1] = 0;
                    curr[i-1] = 0;
                    for (int j = i; j < k; j++){
                        ksi[i-1] += HR(i-1, j)*xt[j];
                        curr[i-1] += xt[j]*R(i-1,j);
                    }
                    dist[i-1] = dist[i] + xidist;
                    ener[i-1] = ener[i] + xiener;
                    i--;
                    dfe(i, q, xt, y, delta, ksi, HR, S);
                } else { // lattice point is found (Step 5)
                    radius = dist[0] + xidist;
                    found = true;
                    x = conv_to<vector<int>>::from(xt);
                    i++;
                    se_enum(i, xt, delta);
                }
            }
        }
    }
    if (found)
        return x;
    else {
        log_msg("Initial squared radius used: " + to_string(radius), "Alert");
        // mat tmp = y-HR*xt;
        // cx_mat distasd = cx_mat(tmp, mat(tmp.n_rows, tmp.n_cols, fill::zeros));
        // log_msg("Current squared distance: " + to_string(frob_norm_squared(distasd)), "Alert");
        // log_msg("xidist: " + to_string(xidist), "Alert");
        log_msg("distances considered: " + vec2str(distances, distances.size()), "Alert");
        log_msg("sphdec: point not found!", "Alert");
        return vector<int>(0);
    }
}

/* Wrapper function for the sphere decoder to handle complex to real matrix conversion,
   QR-decomposition and other mappings */
vector<int> sphdec_wrapper(const vector<cx_mat> &bases, const mat Rorig, const cx_mat &H,
                           const cx_mat &X, const cx_mat &N, const vector<int> &symbset, int &visited_nodes, double radius) {

    /* read simulation parameters */
    int n = params["no_of_receiver_antennas"];
    int t = params["time_slots"];
    int k = params["no_of_matrices"];
    double P = dparams["spherical_shaping_max_power"];

    vector<int> x(k);                          /* decoded point */
    cx_mat Y(n, t);                            /* received code block */
    mat B(2*t*n, k), Q, R;                     /* real matrices for QR-decompostion */
    vec y(2*t*n);                              /* raw input vector for the sphere decoder (unmapped) */
    vec y2;                                    /* input vector for the sphere decoder (mapped to same space a R) */

    Y = H*X + N;                               /* Calculate simulated code block that we would receive */
    y = to_real_vector(Y);                     /* convert Y to real vector */

    for (int i = 0; i < k; i++)
        B.col(i) = to_real_vector(H*bases[i]); /* B = (HX1 HX2 ... HXk) = generator matrix of faded lattice */

    qr_econ(Q, R, B);                          /* QR-decomposition of B (omits zero rows in R) */
    process_qr(Q, R);                          /* Make sure R has positive diagonal elements */

    y2 = Q.st()*y;                             /* Map y to same basis as R */

    /* decide which sphere decoder algorithm to use */
    if (P <= 0)
        x = sphdec(y2, R, symbset, visited_nodes, radius);
    else
        x = sphdec_spherical_shaping(y2, R, Rorig, symbset, visited_nodes, P, radius);

    return x;
}


#if HAVE_CONFIG_H
# include <config.h>
#endif

#include <boost/bind.hpp>
using boost::cref;

#include "Cosmology.h"
#include "PowerSpectrum.h"
#include "Quadrature.h"
#include "CDE.h"
#include "BSPTN.h"
#include "SPT.h"
#include "Spline.h"
#include "SpecialFunctions.h"
#include "LinearPS.h"

#include <cerrno>
#include <cstdio>
#include <cstdlib>
#include <sys/stat.h>
#include <math.h>       /* pow */

/////////////////CLUSTERING DARK ENERGY (AND MG RESUMMED SPECTRUM) LIBRARY ///////////////////
// Load Linear PS
CDE::CDE(const Cosmology& C, const PowerSpectrum& P_l, double epsrel_)
: C(C), P_l(P_l)
{
    epsrel = epsrel;
}

/* Tree level cross bispectrum for  isosceles or equilateral configurations */
// vars specifies omega_m, cs^2, w , p3 and scale factor
double CDE::Bcde(int a, double vars[], double k, double p, double x) const {
IOW iow;
iow.initn_cde_btree(vars[4], k, p, x, vars[0], vars[1], vars[2],vars[3]);
double d =  sqrt(pow2(k) + pow2(p) + 2*p*k*x);
double omega_m = vars[0]/pow3(vars[4]);
double A = -3.*(1.+vars[2]);
double omegaf = pow(vars[4],A);
double omega_q= (1.-vars[0])*omegaf;
double cs2 = vars[1];
double norm = 2./pow4(dnorm_spt);
double terms[3];
if(d < 1e-6)
    return 0;
else
    switch(a) {
      // theta_m density_m density_m
        case 1:
						return    norm*(P_l(k)*P_l(p)*G1_nk*F1p_nk[0]*F2A_nk[0] +  P_l(d)*P_l(k)*G1_nk*F2D_nk[0]*F1kmp_nk[0]
                          + P_l(p)*P_l(d)*G2_nk[0]*F1p_nk[0]*F1kmp_nk[0]);
            break;
      // theta_q density_m density_m
        case 2:
            return    norm*(P_l(k)*P_l(p)*G1q_nk*F1p_nk[0]*F2A_nk[0] +  P_l(d)*P_l(k)*G1q_nk*F2D_nk[0]*F1kmp_nk[0]
                      + P_l(p)*P_l(d)*G2q_nk[0]*F1p_nk[0]*F1kmp_nk[0]);
            break;
      // theta_m \delta_q \delta_q
        case 3:
            return    norm*(P_l(k)*P_l(p)*G1_nk*F1pq_nk[0]*F2Aq_nk[0] +  P_l(d)*P_l(k)*G1_nk*F2Dq_nk[0]*F1kmpq_nk[0]
                      + P_l(p)*P_l(d)*G2_nk[0]*F1pq_nk[0]*F1kmpq_nk[0]);
            break;
       // theta_q \delta_q \delta_q
        case 4:
            return    norm*(P_l(k)*P_l(p)*G1q_nk*F1p_nk[0]*F2A_nk[0] +  P_l(d)*P_l(k)*G1q_nk*F2D_nk[0]*F1kmp_nk[0]
                      + P_l(p)*P_l(d)*G2q_nk[0]*F1p_nk[0]*F1kmp_nk[0]);
           break;
      // theta_m \delta_m \delta_q
        case 5:
            return    norm*(P_l(k)*P_l(p)*G1_nk*F1p_nk[0]*F2Aq_nk[0] +  P_l(d)*P_l(k)*G1_nk*F2D_nk[0]*F1kmpq_nk[0]
                      + P_l(p)*P_l(d)*G2_nk[0]*F1p_nk[0]*F1kmpq_nk[0]);
            break;
     // theta_q \delta_m \delta_q
        case 6:
            return    norm*(P_l(k)*P_l(p)*G1q_nk*F1p_nk[0]*F2Aq_nk[0] +  P_l(d)*P_l(k)*G1q_nk*F2D_nk[0]*F1kmpq_nk[0]
                      + P_l(p)*P_l(d)*G2q_nk[0]*F1p_nk[0]*F1kmpq_nk[0]);
           break;
     // theta_m \delta_t \delta_t [\delta_t = \omega_m*(\delta_m + \omega_q/\omega_m*(1+cs^2)*\delta_q) ]
        case 7:
        // theta_m density_m density_m
  				 terms[0]= norm*(P_l(k)*P_l(p)*G1_nk*F1p_nk[0]*F2A_nk[0] +  P_l(d)*P_l(k)*G1_nk*F2D_nk[0]*F1kmp_nk[0]
                         + P_l(p)*P_l(d)*G2_nk[0]*F1p_nk[0]*F1kmp_nk[0]);

        // theta_m \delta_q \delta_q
           terms[1]= norm*(P_l(k)*P_l(p)*G1_nk*F1pq_nk[0]*F2Aq_nk[0] +  P_l(d)*P_l(k)*G1_nk*F2Dq_nk[0]*F1kmpq_nk[0]
                     + P_l(p)*P_l(d)*G2_nk[0]*F1pq_nk[0]*F1kmpq_nk[0]);

        // theta_m \delta_m \delta_q
           terms[2]= norm*(P_l(k)*P_l(p)*G1_nk*F1p_nk[0]*F2Aq_nk[0] +  P_l(d)*P_l(k)*G1_nk*F2D_nk[0]*F1kmpq_nk[0]
                     + P_l(p)*P_l(d)*G2_nk[0]*F1p_nk[0]*F1kmpq_nk[0]);

          return pow2(omega_m)*(terms[0] + 2.*omega_q/omega_m*(1.+cs2)*terms[2] + pow2(omega_q/omega_m*(1.+cs2))*terms[1]);
           break;

           // theta_q \delta_t \delta_t [\delta_t = \omega_m*(\delta_m + \omega_q/\omega_m*(1+cs^2)*\delta_q) ]
       case 8:
          // theta_q density_m density_m
            terms[0]= norm*(P_l(k)*P_l(p)*G1q_nk*F1p_nk[0]*F2A_nk[0] +  P_l(d)*P_l(k)*G1q_nk*F2D_nk[0]*F1kmp_nk[0]
                      + P_l(p)*P_l(d)*G2q_nk[0]*F1p_nk[0]*F1kmp_nk[0]);

          // theta_q \delta_m \delta_q
            terms[1]=  norm*(P_l(k)*P_l(p)*G1q_nk*F1p_nk[0]*F2Aq_nk[0] +  P_l(d)*P_l(k)*G1q_nk*F2D_nk[0]*F1kmpq_nk[0]
                      + P_l(p)*P_l(d)*G2q_nk[0]*F1p_nk[0]*F1kmpq_nk[0]);

          // theta_q \delta_q \delta_q
            terms[2]=  norm*(P_l(k)*P_l(p)*G1q_nk*F1p_nk[0]*F2A_nk[0] +  P_l(d)*P_l(k)*G1q_nk*F2D_nk[0]*F1kmp_nk[0]
                      + P_l(p)*P_l(d)*G2q_nk[0]*F1p_nk[0]*F1kmp_nk[0]);

            return pow2(omega_m)*(terms[0] + 2.*omega_q/omega_m*(1.+cs2)*terms[2] + pow2(omega_q/omega_m*(1.+cs2))*terms[1]);
            break;
        case 9:
          // d_m d_m d_m
          return    norm*(P_l(k)*P_l(p)*F1_nk*F1p_nk[0]*F2A_nk[0] +  P_l(d)*P_l(k)*F1_nk*F2D_nk[0]*F1kmp_nk[0]
                        + P_l(p)*P_l(d)*F2_nk[0]*F1p_nk[0]*F1kmp_nk[0]);
          break;
      default:
            warning("SPT: invalid indices, a = %d\n", a);
            return 0;
    }
}

// 1-loop resummed spectrum

// Spline linear growth factors and no wiggle spectra
Spline F1_spline, G1_spline, EHUNWcde, RENWcde;
double ANWcde,sigmav;

double KMIN_cde = 1e-4;
double KMAX_cde = 50.;
void CDE::F1_init(int a, double scalef, double omega0, double par1, double par2, double par3) const{
  IOW iow;
  vector<double> kval_table, F1_table, G1_table;
  int n3 = 300;
  for(int i = 0; i<n3; i++){
  double k = KMIN_cde*exp(i*log(KMAX_cde/(KMIN_cde))/(n3-1.));
  iow.initn_lin(a, scalef, k, omega0, par1, par2, par3);
  double ling1=F1_nk;
  double ling2=G1_nk;

  kval_table.push_back(k);
  F1_table.push_back(ling1);
  G1_table.push_back(ling2);

  }
  F1_spline = LinearSpline(kval_table,F1_table);
  G1_spline = LinearSpline(kval_table,G1_table);
}

static double ANWcde_integrand(double q, double k){
return pow2(F1_spline(k))*RENWcde(k)*(1.-j0(k*q))*pow2(q);
}

static double nowig_integrand(const PowerSpectrum& P_l, double k, double q){
double lambda = 0.25*pow(k/0.05,0.04);
return P_l(q)/EHUNWcde(q)*exp(-pow2(log(k/q)/lambda)/2.)/q;
}

////////////////
// FOR RSD ONLY

static double cde_exp(const PowerSpectrum& P_l, double q){
  return  P_l(q)*pow2(F1_spline(q))/(6.*pow2(M_PI));
}

static void sigv_init(const PowerSpectrum& P_l){
    sigmav = Integrate(bind(cde_exp,cref(P_l), _1), KMIN_cde, KMAX_cde, 1e-3);
  }
//////////////

// Initialize the NW power spectrum, A^{nw,0l} and sigma_v

// run this before running other kernel initializations (initn or initn_cde)
// a chooses cde(1) or mg(0)
void CDE::initcde(int a, double scalef, double omega0, double par1, double par2, double par3) const{
NoWigglePS now(C, 0. , EisensteinHu);
IOW iow;

// linear spectrum initialization in range [1e-4,50]
F1_init(a, scalef, omega0, par1, par2, par3);

// sigma_v iinitialization
//sigv_init(cref(P_l));

// EHUNWcde initialization
int n3 = 500;
vector<double> kval_table, ehu_table, now_table;;
for(int i = 0; i<n3; i++ ){
double  k = KMIN_cde*exp(i*log(KMAX_cde/KMIN_cde)/(n3-1.));
double  ehups = now.Evaluate(k);
  kval_table.push_back(k);
  ehu_table.push_back(ehups);
    }

EHUNWcde = LinearSpline(kval_table,ehu_table);

// Renormalized linear spectrum and ANWcde initialization
const double qmin = 10.;
const double qmax = 300.;
double c[2] = {qmin,KMIN_cde};
double d[2] = {qmax,KMAX_cde};

for(int i = 0; i<n3; i++ ){
  double k = kval_table[i];

  double lambda = 0.25*pow(k/0.05,0.04);
  double nowps = EHUNWcde(k)/sqrt(2*M_PI*pow2(lambda))*Integrate<ExpSub>(bind(nowig_integrand,cref(P_l),k,_1),KMIN_cde,KMAX_cde,epsrel);
  now_table.push_back(nowps);
    }

  RENWcde = LinearSpline(kval_table,now_table);
  ANWcde = 1./((pow3(qmax)-pow3(qmin))*pow2(M_PI))*Integrate<2>(bind(ANWcde_integrand,_1,_2),c,d,1e-3,1e-3);

}

// NUMERICAL KERNEL RESUMMED 1-LOOP TERMS
//Integrating over angle - P22 numerical:
static double F2F(int y, int a,  const PowerSpectrum& P_l, double k, double r){
	double temp_ps;
	double myresult=0.;
	switch (a) {
			case 1:
			for( int i = 0; i < n1; i++ )
				{
        double d = 1+ r*r - 2*r*x128[i];
		 		if(d < 1e-5){
			 	temp_ps=0.;
		 		}
		  	else {
		  	temp_ps = RENWcde(k*sqrt(d)) * pow2(F2_nk[i*n2 + y]);
				}
        myresult += w128[i] * temp_ps;
				}
        return 2. * r * r * myresult;
        break;
			case 2:
      myresult = 0.;
			for( int i = 0; i < n1; i++ )
				{
        double d = 1+ r*r - 2*r*x128[i];
				if(d < 1e-5){
				temp_ps=0.;
				}
				else {
				temp_ps = RENWcde(k*sqrt(d)) * G2_nk[i*n2 + y]*F2_nk[i*n2 + y];
				}
        myresult += w128[i] * temp_ps;
				}
        return  2. * r * r * myresult;
        break;
			case 3:
			for( int i = 0; i < n1; i++ )
				{
				double d = 1+ r*r - 2*r*x128[i];
				if(d < 1e-5){
				temp_ps=0.;
				}
				else {
				temp_ps = RENWcde(k*sqrt(d)) * pow2(G2_nk[i*n2 + y]);
				}
        myresult += w128[i] * temp_ps;
				}
        return  2. * r * r * myresult;
        break;
			}
	}

  //Integrating over angle -P13 numerical:
  static double F3F(int y, int a,  const PowerSpectrum& P_l, double k, double r){
    double myresult = 0.;
    // Set the limits of angular integrations
  	switch (a) {
  			case 1:
  			 for( int i = 0; i < n1; i++ )
  				{
  				myresult += w128[i] * F1_nk * F3_nk[i*n2 + y];
  				}
        return   6. * r * r * myresult;
        break;
  			case 2:
  			 for( int i = 0; i < n1; i++ )
  				{
  					myresult += w128[i] * (G1_nk*F3_nk[i*n2 + y] + F1_nk*G3_nk[i*n2 + y]);
  				}
        return   3. * r * r * myresult;
        break;
  			case 3:
  				for( int i = 0; i < n1; i++ )
  				{
  					myresult += w128[i] * G1_nk * G3_nk[i*n2 + y];
  				}
          return    6. * r * r * myresult;
          break;

  		}
    }

/*Selection function for numerical kernel magnitude */

inline int num_selec_mag(double kmin, double kmax, double y){
      double YMIN =QMINp/kmax;
      double YMAX = QMAXp/kmin;
    	int mag_int;
  //  	mag_int = (y-YMIN)*(n2*1.-1)/(YMAX-YMIN); // linear
    //mag_int = sqrt((y-QMINp/kmax)*kmin/QMAXp)*(n2-1); //quadratic
    mag_int = (int)round((n2-1)*log(y/YMIN)/log(YMAX/YMIN)); //exponential
    		return mag_int;
}

/* P^{(22)} */
// a = 1 : delta delta
// a = 2 : delta theta
// a = 3 : theta theta
double CDE::P22nres(double kmin, double kmax,  int a, double k) const {
    	double KMAX = QMAXp/k;
    	double KMIN = QMINp/k;
      int y1 = num_selec_mag(kmin, kmax, KMIN);
      int y2 = num_selec_mag(kmin, kmax, KMAX);
      double y[n2];
      double integrand[n2];
      for (int i = y1; i<=y2; i++){
      y[i] = QMINp/kmax * exp(i*log(QMAXp*kmax/(QMINp*kmin))/(n2*1.-1.)); // exponential sampling
  //    y[i] = i*1./(n2-1.)*(QMAXp/kmin-QMINp/kmax)+QMINp/kmax; // linear sampling
      integrand[i] = RENWcde(k*y[i]) * F2F(i, a, cref(P_l), k, y[i]);
}
double res = 0.;
  for( int i = y1+1; i <= y2; ++ i ){
res += 0.5 * (y[i] - y[i-1])*(integrand[i] + integrand[i-1]);
}
return  k*k*k/(4*M_PI*M_PI)/pow4(dnorm_spt) * res;
}

double CDE::P13nres(double kmin, double kmax,  int a, double k) const {
  double KMAX = QMAXp/k;
  double KMIN = QMINp/k;
  int y1 = num_selec_mag(kmin, kmax, KMIN);
  int y2 = num_selec_mag(kmin, kmax, KMAX);
  double y[n2];
  double integrand[n2];
  for (int i = y1; i<=y2; i++){
  y[i] = QMINp/kmax * exp(i*log(QMAXp*kmax/(QMINp*kmin))/(n2*1.-1.));
//  y[i] = i*1./(n2-1.)*(QMAXp/kmin-QMINp/kmax)+QMINp/kmax; // linear sampling
  integrand[i] = RENWcde(k*y[i]) * F3F(i, a, cref(P_l), k, y[i]);
  }
double res = 0.;
  for( int i = y1+1; i <= y2; ++ i ){
res += 0.5 * (y[i] - y[i-1])*(integrand[i] + integrand[i-1]);
}
  return  k*k*k/(4*M_PI*M_PI)/pow4(dnorm_spt) * RENWcde(k) * res;
}

double CDE::PNWloopn(double kmin, double kmax, double k, int a ) const{
  switch(a) {
      case 1:
          return   pow2(F1_spline(k)/dnorm_spt)*RENWcde(k) + P13nres(kmin, kmax, a, k) + P22nres(kmin, kmax, a, k);
          break;
      case 2:
          return  F1_spline(k)*G1_spline(k)/pow2(dnorm_spt)*RENWcde(k) + P13nres(kmin, kmax, a, k) + P22nres(kmin, kmax, a, k);
          break;
      case 3:
          return  pow2(G1_spline(k)/dnorm_spt)*RENWcde(k) + P13nres(kmin, kmax, a, k) + P22nres(kmin, kmax, a, k);
          break;
      default:
          warning("CDE: invalid indices, a = %d\n", a);
          return 0;
  }
}

/* Resummed 1-loop spectra a la Fonseca */

double CDE::Presumn(double kmin, double kmax, double k, int a) const{
  SPT spt(C,P_l,epsrel);
  switch(a){
  case 1:
   return   PNWloopn(kmin,kmax,k,1) + exp(-0.5*pow2(k)*ANWcde/pow2(dnorm_spt))*(spt.PLOOPn(kmin,kmax,1,k)-PNWloopn(kmin,kmax,k,1) + 0.5*pow2(k)*ANWcde/pow2(dnorm_spt)*pow2(F1_spline(k)/dnorm_spt)*(P_l(k) - RENWcde(k)));
      break;
  case 2:
   return PNWloopn(kmin,kmax,k,2) + exp(-0.5*pow2(k)*ANWcde/pow2(dnorm_spt))*(spt.PLOOPn(kmin,kmax,2,k)-PNWloopn(kmin,kmax,k,2) + 0.5*pow2(k)*ANWcde/pow2(dnorm_spt)*G1_spline(k)*F1_spline(k)/pow2(dnorm_spt)*(P_l(k) - RENWcde(k)));
      break;
  case 3:
      return PNWloopn(kmin,kmax,k,3) + exp(-0.5*pow2(k)*ANWcde/pow2(dnorm_spt))*(spt.PLOOPn(kmin,kmax,3,k)-PNWloopn(kmin,kmax,k,3) + 0.5*pow2(k)*ANWcde/pow2(dnorm_spt)*pow2(G1_spline(k)/dnorm_spt)*(P_l(k) - RENWcde(k)));
      break;
  case 4:
      return pow2(F1_spline(k)/dnorm_spt)*RENWcde(k);
      //IR CHECK
  case 5:
      return P13nres(kmin, kmax, 1, k) + P22nres(kmin, kmax, 1, k) + exp(-0.5*pow2(k)*ANWcde/pow2(dnorm_spt))*(spt.P22n(kmin,kmax,1,k) + spt.P13n(kmin,kmax,1,k) -(P13nres(kmin, kmax, 1, k) + P22nres(kmin, kmax, 1, k)));
  default:
          warning("CDE: invalid indices, a = %d\n", a);
          return 0;
}
}


/* ----------------------------------------------------------------------------
 * Copyright 2020, Massachusetts Institute of Technology, * Cambridge, MA 02139
 * All Rights Reserved
 * Authors: Yulun Tian, et al. (see README for the full author list)
 * See LICENSE for the license information
 * -------------------------------------------------------------------------- */

#include <DPGO/DPGO_utils.h>
#include <DPGO/PGOAgent.h>
#include <DPGO/QuadraticOptimizer.h>

#include <Eigen/CholmodSupport>
#include <algorithm>
#include <cassert>
#include <chrono>
#include <cmath>
#include <iostream>
#include <random>

#include "StieVariable.h"

using std::lock_guard;
using std::unique_lock;
using std::set;
using std::thread;
using std::vector;

namespace DPGO {

PGOAgent::PGOAgent(unsigned ID, const PGOAgentParameters &params)
    : mID(ID), d(params.d), r(params.r), n(1),
      mParams(params), mState(PGOAgentState::WAIT_FOR_DATA),
      mStatus(ID, mState, 0, 0, false, 0),
      mRobustCost(params.robustCostType, params.robustCostParams),
      mProblemPtr(nullptr),
      mInstanceNumber(0), mIterationNumber(0), mNumPosesReceived(0),
      mLogger(params.logDirectory) {
  if (mParams.verbose) {
    std::cout << "Initializing PGO agent..." << std::endl;
    std::cout << params << std::endl;
  }

  // Initialize X
  X = Matrix::Zero(r, d + 1);
  X.block(0, 0, d, d) = Matrix::Identity(d, d);
  if (mID == 0) setLiftingMatrix(fixedStiefelVariable(d, r));
  resetTeamStatus();
}

PGOAgent::~PGOAgent() {
  // Make sure that optimization thread is not running, before exiting
  endOptimizationLoop();
}

void PGOAgent::setX(const Matrix &Xin) {
  lock_guard<mutex> lock(mPosesMutex);
  assert(mState != PGOAgentState::WAIT_FOR_DATA);
  assert(Xin.rows() == relaxation_rank());
  assert(Xin.cols() == (dimension() + 1) * num_poses());
  mState = PGOAgentState::INITIALIZED;
  X = Xin;
  if (mParams.acceleration) {
    initializeAcceleration();
  }
  if (mParams.verbose) {
    printf("Robot %u resets trajectory estimates. New trajectory length = %u\n", getID(), num_poses());
  }
}

bool PGOAgent::getX(Matrix &Mout) {
  lock_guard<mutex> lock(mPosesMutex);
  Mout = X;
  return true;
}

bool PGOAgent::getSharedPose(unsigned int index, Matrix &Mout) {
  if (mState != PGOAgentState::INITIALIZED)
    return false;
  lock_guard<mutex> lock(mPosesMutex);
  if (index >= num_poses()) return false;
  Mout = X.block(0, index * (d + 1), r, d + 1);
  return true;
}

bool PGOAgent::getAuxSharedPose(unsigned int index, Matrix &Mout) {
  assert(mParams.acceleration);
  if (mState != PGOAgentState::INITIALIZED)
    return false;
  lock_guard<mutex> lock(mPosesMutex);
  if (index >= num_poses()) return false;
  Mout = Y.block(0, index * (d + 1), r, d + 1);
  return true;
}

bool PGOAgent::getSharedPoseDict(PoseDict &map) {
  if (mState != PGOAgentState::INITIALIZED)
    return false;
  map.clear();
  lock_guard<mutex> lock(mPosesMutex);
  for (const auto &mSharedPose: localSharedPoseIDs) {
    unsigned idx = std::get<1>(mSharedPose);
    map[mSharedPose] = X.block(0, idx * (d + 1), r, d + 1);
  }
  return true;
}

bool PGOAgent::getAuxSharedPoseDict(PoseDict &map) {
  assert(mParams.acceleration);
  if (mState != PGOAgentState::INITIALIZED)
    return false;
  map.clear();
  lock_guard<mutex> lock(mPosesMutex);
  for (const auto &mSharedPose: localSharedPoseIDs) {
    unsigned idx = std::get<1>(mSharedPose);
    map[mSharedPose] = Y.block(0, idx * (d + 1), r, d + 1);
  }
  return true;
}

void PGOAgent::setLiftingMatrix(const Matrix &M) {
  assert(M.rows() == r);
  assert(M.cols() == d);
  YLift.emplace(M);
}

void PGOAgent::setPoseGraph(
    const std::vector<RelativeSEMeasurement> &inputOdometry,
    const std::vector<RelativeSEMeasurement> &inputPrivateLoopClosures,
    const std::vector<RelativeSEMeasurement> &inputSharedLoopClosures,
    const Matrix &TInit) {
  assert(!isOptimizationRunning());
  assert(mState == PGOAgentState::WAIT_FOR_DATA);
  assert(n == 1);

  if (inputOdometry.empty()) return;

  for (const auto &edge: inputOdometry) {
    addOdometry(edge);
  }
  for (const auto &edge: inputPrivateLoopClosures) {
    addPrivateLoopClosure(edge);
  }
  for (const auto &edge: inputSharedLoopClosures) {
    addSharedLoopClosure(edge);
  }

  // Check validity of initial trajectory estimate, if provided
  bool local_init = true;
  unsigned expected_rows = dimension();
  unsigned expected_cols = (dimension() + 1) * num_poses();
  if (TInit.rows() > 0 && TInit.cols() > 0) {
    if (TInit.rows() == expected_rows && TInit.cols() == expected_cols) {
      local_init = false;
    } else {
      local_init = true;
      printf("Error: provided initial trajectory has wrong dimension! "
             "Expect (%u,%u), received (%ld, %ld). Using local initialization. \n",
             expected_rows, expected_cols, TInit.rows(), TInit.cols());
    }
  }

  // Create new optimization problem
  mProblemPtr = new QuadraticProblem(num_poses(), dimension(), relaxation_rank());

  // Robot can construct the quadratic cost matrix now, as it does not depend on neighbor values
  constructQMatrix();

  // Initialize trajectory estimate in an arbitrary frame
  if (!local_init) {
    if (mParams.verbose) printf("Using provided trajectory initialization.\n");
    TLocalInit.emplace(TInit.block(0, 0, expected_rows, expected_cols));
  } else {
    if (mParams.verbose) printf("Using internal trajectory initialization.\n");
    localInitialization();
  }

  // Waiting for initialization in the GLOBAL frame
  mState = PGOAgentState::WAIT_FOR_INITIALIZATION;

  // If I am the first robot or if cross-robot initialization if off,
  // I will consider myself as initialized in the global frame
  if (mID == 0 || !mParams.multirobot_initialization) {
    X = YLift.value() * TLocalInit.value();  // Lift to correct relaxation rank
    XInit.emplace(X);
    mState = PGOAgentState::INITIALIZED;
    if (mParams.acceleration) {
      initializeAcceleration();
    }

    // Save initial trajectory
    if (mParams.logData) {
      mLogger.logTrajectory(dimension(), num_poses(), TLocalInit.value(), "trajectory_initial.csv");
    }
  }
}

void PGOAgent::addOdometry(const RelativeSEMeasurement &factor) {
  assert(mState != PGOAgentState::INITIALIZED);
  // check that this is a odometry measurement
  assert(factor.r1 == mID);
  assert(factor.r2 == mID);
  assert(factor.p1 + 1 == factor.p2);
  assert(factor.R.rows() == d && factor.R.cols() == d);
  assert(factor.t.rows() == d && factor.t.cols() == 1);

  // update number of poses
  n = std::max(n, (unsigned) factor.p2 + 1);

  lock_guard<mutex> mLock(mMeasurementsMutex);
  odometry.push_back(factor);
}

void PGOAgent::addPrivateLoopClosure(const RelativeSEMeasurement &factor) {
  assert(mState != PGOAgentState::INITIALIZED);
  assert(factor.r1 == mID);
  assert(factor.r2 == mID);
  assert(factor.R.rows() == d && factor.R.cols() == d);
  assert(factor.t.rows() == d && factor.t.cols() == 1);

  // update number of poses
  n = std::max(n, (unsigned) std::max(factor.p1 + 1, factor.p2 + 1));

  lock_guard<mutex> lock(mMeasurementsMutex);
  privateLoopClosures.push_back(factor);
}

void PGOAgent::addSharedLoopClosure(const RelativeSEMeasurement &factor) {
  assert(mState != PGOAgentState::INITIALIZED);
  assert(factor.R.rows() == d && factor.R.cols() == d);
  assert(factor.t.rows() == d && factor.t.cols() == 1);

  if (factor.r1 == mID) {
    assert(factor.r2 != mID);
    n = std::max(n, (unsigned) factor.p1 + 1);
    localSharedPoseIDs.insert(std::make_pair(mID, factor.p1));
    neighborSharedPoseIDs.insert(std::make_pair(factor.r2, factor.p2));
    neighborRobotIDs.insert(factor.r2);
  } else {
    assert(factor.r2 == mID);
    n = std::max(n, (unsigned) factor.p2 + 1);
    localSharedPoseIDs.insert(std::make_pair(mID, factor.p2));
    neighborSharedPoseIDs.insert(std::make_pair(factor.r1, factor.p1));
    neighborRobotIDs.insert(factor.r1);
  }

  lock_guard<mutex> lock(mMeasurementsMutex);
  sharedLoopClosures.push_back(factor);
}

Matrix PGOAgent::computeNeighborTransform(const PoseID &nID, const Matrix &var) {
  assert(YLift);
  assert(var.rows() == r);
  assert(var.cols() == d + 1);

  // Find the corresponding inter-robot loop closure
  RelativeSEMeasurement &m = findSharedLoopClosureWithNeighbor(nID);

  // Notations:
  // world1: world frame before alignment
  // world2: world frame after alignment
  // frame1 : frame associated to my public pose
  // frame2 : frame associated to neighbor's public pose
  Matrix dT = Matrix::Identity(d + 1, d + 1);
  dT.block(0, 0, d, d) = m.R;
  dT.block(0, d, d, 1) = m.t;
  Matrix T_world2_frame2 = Matrix::Identity(d + 1, d + 1);
  T_world2_frame2.block(0, 0, d, d + 1) =
      YLift.value().transpose() *
          var;  // Round the received neighbor pose value back to SE(d)
  Matrix T = TLocalInit.value();
  Matrix T_frame1_frame2 = Matrix::Identity(d + 1, d + 1);
  Matrix T_world1_frame1 = Matrix::Identity(d + 1, d + 1);
  if (m.r1 == nID.first) {
    // Incoming edge
    T_frame1_frame2 = dT.inverse();
    T_world1_frame1.block(0, 0, d, d + 1) =
        T.block(0, m.p2 * (d + 1), d, d + 1);
  } else {
    // Outgoing edge
    T_frame1_frame2 = dT;
    T_world1_frame1.block(0, 0, d, d + 1) =
        T.block(0, m.p1 * (d + 1), d, d + 1);
  }
  Matrix T_world2_frame1 = T_world2_frame2 * T_frame1_frame2.inverse();
  Matrix T_world2_world1 = T_world2_frame1 * T_world1_frame1.inverse();
  checkRotationMatrix(T_world2_world1.block(0, 0, d, d));
  return T_world2_world1;
}

Matrix PGOAgent::computeRobustNeighborTransformTwoStage(unsigned int neighborID, const PoseDict &poseDict) {
  std::vector<Matrix> RVec;
  std::vector<Vector> tVec;
  for (const auto &it: poseDict) {
    const PoseID nID = it.first;
    const auto var = it.second;
    if (neighborSharedPoseIDs.find(nID) != neighborSharedPoseIDs.end()) {
      const auto T = computeNeighborTransform(nID, var);
      RVec.emplace_back(T.block(0, 0, d, d));
      tVec.emplace_back(T.block(0, d, d, 1));
    }
  }
  int m = (int) RVec.size();
  const Vector kappa = Vector::Ones(m);
  const Vector tau = Vector::Ones(m);
  Matrix ROpt;
  Vector tOpt;
  std::vector<size_t> inlierIndices;
  // Perform robust single rotation averaging
  double maxRotationError = angular2ChordalSO3(0.5);  // approximately 30 deg
  robustSingleRotationAveraging(ROpt, inlierIndices, RVec, kappa, maxRotationError);
  int inlierSize = (int) inlierIndices.size();
  printf("[RobustRelativeTransform] This robot %u, neighbor %u: finds %i inliers out of %i measurements.\n",
         getID(),
         neighborID,
         inlierSize,
         m);
  if (inlierSize == 0) {
    throw std::runtime_error("Robust single rotation averaging returns empty inlier set!");
  }
  // Perform single translation averaging on the inlier set
  std::vector<Vector> tVecInliers;
  for (const auto index: inlierIndices) {
    tVecInliers.emplace_back(tVec[index]);
  }
  singleTranslationAveraging(tOpt, tVecInliers);
  // Return transformation as matrix
  Matrix TOpt = Matrix::Identity(dimension() + 1, dimension() + 1);
  TOpt.block(0, 0, d, d) = ROpt;
  TOpt.block(0, d, d, 1) = tOpt;
  return TOpt;
}

Matrix PGOAgent::computeRobustNeighborTransform(unsigned int neighborID, const PoseDict &poseDict) {
  std::vector<Matrix> RVec;
  std::vector<Vector> tVec;
  for (const auto &it: poseDict) {
    const PoseID nID = it.first;
    const auto var = it.second;
    if (neighborSharedPoseIDs.find(nID) != neighborSharedPoseIDs.end()) {
      const auto T = computeNeighborTransform(nID, var);
      RVec.emplace_back(T.block(0, 0, d, d));
      tVec.emplace_back(T.block(0, d, d, 1));
    }
  }
  int m = (int) RVec.size();
  const Vector kappa = 1.82 * Vector::Ones(m);  // rotation stddev approximately 30 degree
  const Vector tau = 0.01 * Vector::Ones(m);  // translation stddev 10 m
  const double cbar = RobustCost::computeErrorThresholdAtQuantile(0.9, 3);
  Matrix ROpt;
  Vector tOpt;
  std::vector<size_t> inlierIndices;
  robustSinglePoseAveraging(ROpt, tOpt, inlierIndices, RVec, tVec, kappa, tau, cbar);
  int inlierSize = (int) inlierIndices.size();
  printf("[RobustRelativeTransform] This robot %u, neighbor %u: finds %i inliers out of %i measurements.\n",
         getID(),
         neighborID,
         inlierSize,
         m);
  if (inlierSize == 0) {
    throw std::runtime_error("Robust single pose averaging returns empty inlier set!");
  }
  // Return transformation as matrix
  Matrix TOpt = Matrix::Identity(dimension() + 1, dimension() + 1);
  TOpt.block(0, 0, d, d) = ROpt;
  TOpt.block(0, d, d, 1) = tOpt;
  return TOpt;
}

void PGOAgent::initializeInGlobalFrame(unsigned neighborID, const PoseDict &poseDict) {
  // Require the lifting matrix to initialize
  assert(YLift);

  // Halt optimization
  bool optimizationHalted = false;
  if (isOptimizationRunning()) {
    if (mParams.verbose)
      printf("Robot %u halting optimization thread...\n", getID());
    optimizationHalted = true;
    endOptimizationLoop();
  }

  // Halt insertion of new poses
  lock_guard<mutex> tLock(mPosesMutex);

  // Halt insertion of new measurements
  lock_guard<mutex> mLock(mMeasurementsMutex);

  // Clear cache
  lock_guard<mutex> nLock(mNeighborPosesMutex);
  neighborPoseDict.clear();
  neighborAuxPoseDict.clear();

  // Compute relative transform to neighbor's frame of reference
  Matrix T_world2_world1;
  try {
    T_world2_world1 = computeRobustNeighborTransformTwoStage(neighborID, poseDict);
  } catch (const std::runtime_error& e) {
    printf("Robust initialization is not successful! Abort and wait to try again...\n");
    return;
  }

  // Apply global transformation to local trajectory estimate
  Matrix T = TLocalInit.value();
  Matrix T_world1_frame = Matrix::Identity(d + 1, d + 1);
  Matrix T_world2_frame = Matrix::Identity(d + 1, d + 1);
  for (size_t i = 0; i < num_poses(); ++i) {
    T_world1_frame.block(0, 0, d, d + 1) =
        T.block(0, i * (d + 1), d, d + 1);
    T_world2_frame = T_world2_world1 * T_world1_frame;
    T.block(0, i * (d + 1), d, d + 1) =
        T_world2_frame.block(0, 0, d, d + 1);
  }

  // Lift back to correct relaxation rank
  X = YLift.value() * T;
  XInit.emplace(X);

  // Mark this agent as initialized
  mState = PGOAgentState::INITIALIZED;

  // Initialize auxiliary variables
  if (mParams.acceleration) {
    initializeAcceleration();
  }

  // Log initial trajectory
  if (mParams.logData) {
    mLogger.logTrajectory(dimension(), num_poses(), T, "trajectory_initial.csv");
  }

  if (optimizationHalted) startOptimizationLoop(mRate);
}

void PGOAgent::updateNeighborPoses(unsigned neighborID, const PoseDict &poseDict) {
  assert(neighborID != mID);
  // Initialize this robot in the global frame, if not initialized
  const auto neighborState = getNeighborStatus(neighborID).state;
  if (mState == PGOAgentState::WAIT_FOR_INITIALIZATION && neighborState == PGOAgentState::INITIALIZED) {
    initializeInGlobalFrame(neighborID, poseDict);
  }
  // Save neighbor public poses in local cache
  for (const auto &it: poseDict) {
    const auto nID = it.first;
    const auto var = it.second;
    assert(nID.first == neighborID);
    assert(var.rows() == r);
    assert(var.cols() == d + 1);
    mNumPosesReceived++;
    if (neighborSharedPoseIDs.find(nID) == neighborSharedPoseIDs.end())
      continue;
    // Only save poses from neighbors if this agent is initialized
    // and if the sending agent is also initialized
    if (mState == PGOAgentState::INITIALIZED && neighborState == PGOAgentState::INITIALIZED) {
      lock_guard<mutex> lock(mNeighborPosesMutex);
      neighborPoseDict[nID] = var;
    }
  }
}

void PGOAgent::updateAuxNeighborPoses(unsigned neighborID, const PoseDict &poseDict) {
  assert(mParams.acceleration);
  assert(neighborID != mID);
  for (const auto &it: poseDict) {
    const auto nID = it.first;
    const auto var = it.second;
    assert(nID.first == neighborID);
    assert(var.rows() == r);
    assert(var.cols() == d + 1);
    mNumPosesReceived++;
    if (neighborSharedPoseIDs.find(nID) == neighborSharedPoseIDs.end())
      continue;
    // Only save poses from neighbors if this agent is initialized
    // and if the sending agent is also initialized
    if (mState == PGOAgentState::INITIALIZED && getNeighborStatus(neighborID).state == PGOAgentState::INITIALIZED) {
      lock_guard<mutex> lock(mNeighborPosesMutex);
      neighborAuxPoseDict[nID] = var;
    }
  }
}

bool PGOAgent::getTrajectoryInLocalFrame(Matrix &Trajectory) {
  if (mState != PGOAgentState::INITIALIZED) {
    return false;
  }
  lock_guard<mutex> lock(mPosesMutex);

  Matrix T = X.block(0, 0, r, d).transpose() * X;
  Matrix t0 = T.block(0, d, d, 1);

  for (unsigned i = 0; i < n; ++i) {
    T.block(0, i * (d + 1), d, d) =
        projectToRotationGroup(T.block(0, i * (d + 1), d, d));
    T.block(0, i * (d + 1) + d, d, 1) = T.block(0, i * (d + 1) + d, d, 1) - t0;
  }

  Trajectory = T;
  return true;
}

bool PGOAgent::getTrajectoryInGlobalFrame(Matrix &Trajectory) {
  if (!globalAnchor) return false;
  assert(globalAnchor.value().rows() == relaxation_rank());
  assert(globalAnchor.value().cols() == dimension() + 1);
  if (mState != PGOAgentState::INITIALIZED) return false;
  lock_guard<mutex> lock(mPosesMutex);

  Matrix T = globalAnchor.value().block(0, 0, r, d).transpose() * X;
  Matrix t0 = globalAnchor.value().block(0, 0, r, d).transpose() *
      globalAnchor.value().block(0, d, r, 1);

  for (unsigned i = 0; i < n; ++i) {
    T.block(0, i * (d + 1), d, d) =
        projectToRotationGroup(T.block(0, i * (d + 1), d, d));
    T.block(0, i * (d + 1) + d, d, 1) = T.block(0, i * (d + 1) + d, d, 1) - t0;
  }

  Trajectory = T;
  return true;
}

bool PGOAgent::getPoseInGlobalFrame(unsigned int poseID, Matrix &T) {
  if (!globalAnchor) return false;
  assert(globalAnchor.value().rows() == relaxation_rank());
  assert(globalAnchor.value().cols() == dimension() + 1);
  if (mState != PGOAgentState::INITIALIZED) return false;
  lock_guard<mutex> lock(mPosesMutex);
  if (poseID < 0 || poseID >= num_poses()) return false;
  Matrix Ya = globalAnchor.value().block(0, 0, r, d);
  Matrix pa = globalAnchor.value().block(0, d, r, 1);
  Matrix t0 = Ya.transpose() * pa;
  Matrix Xi = X.block(0, poseID * (d + 1), r, d + 1);
  Matrix Ti = Ya.transpose() * Xi;
  Ti.block(0, d, d, 1) -= t0;
  assert(Ti.rows() == d);
  assert(Ti.cols() == d + 1);
  T = Ti;
  return true;
}

bool PGOAgent::getNeighborPoseInGlobalFrame(unsigned int neighborID, unsigned int poseID, Matrix &T) {
  if (!globalAnchor) return false;
  assert(globalAnchor.value().rows() == relaxation_rank());
  assert(globalAnchor.value().cols() == dimension() + 1);
  if (mState != PGOAgentState::INITIALIZED) return false;
  lock_guard<mutex> lock(mNeighborPosesMutex);
  PoseID nID = std::make_pair(neighborID, poseID);
  if (neighborPoseDict.find(nID) != neighborPoseDict.end()) {
    Matrix Ya = globalAnchor.value().block(0, 0, r, d);
    Matrix pa = globalAnchor.value().block(0, d, r, 1);
    Matrix t0 = Ya.transpose() * pa;
    Matrix Xi = neighborPoseDict.at(nID);
    assert(Xi.rows() == r);
    assert(Xi.cols() == d + 1);
    Matrix Ti = Ya.transpose() * Xi;
    Ti.block(0, d, d, 1) -= t0;
    assert(Ti.rows() == d);
    assert(Ti.cols() == d + 1);
    T = Ti;
    return true;
  }
  return false;
}

std::vector<unsigned> PGOAgent::getNeighborPublicPoses(
    const unsigned &neighborID) const {
  // Check that neighborID is indeed a neighbor of this agent
  assert(neighborRobotIDs.find(neighborID) != neighborRobotIDs.end());
  std::vector<unsigned> poseIndices;
  for (PoseID pair: neighborSharedPoseIDs) {
    if (pair.first == neighborID) {
      poseIndices.push_back(pair.second);
    }
  }
  return poseIndices;
}

std::vector<unsigned> PGOAgent::getNeighbors() const {
  std::vector<unsigned> v(neighborRobotIDs.size());
  std::copy(neighborRobotIDs.begin(), neighborRobotIDs.end(), v.begin());
  return v;
}

void PGOAgent::reset() {
  // Terminate optimization thread if running
  endOptimizationLoop();

  if (mParams.logData) {
    // Save measurements (including final weights)
    std::vector<RelativeSEMeasurement> measurements = odometry;
    measurements.insert(measurements.end(), privateLoopClosures.begin(), privateLoopClosures.end());
    measurements.insert(measurements.end(), sharedLoopClosures.begin(), sharedLoopClosures.end());
    mLogger.logMeasurements(measurements, "measurements.csv");

    // Save trajectory estimates after rounding
    Matrix T;
    if (getTrajectoryInGlobalFrame(T)) {
      mLogger.logTrajectory(dimension(), num_poses(), T, "trajectory_optimized.csv");
      std::cout << "Saved optimized trajectory to " << mParams.logDirectory << std::endl;
    }

    // Save solution before rounding
    writeMatrixToFile(X, mParams.logDirectory + "X.txt");
  }

  mInstanceNumber++;
  mIterationNumber = 0;
  mNumPosesReceived = 0;

  // Assume that the old lifting matrix can still be used
  mState = PGOAgentState::WAIT_FOR_DATA;
  mStatus = PGOAgentStatus(getID(), mState, mInstanceNumber, mIterationNumber, false, 0);

  odometry.clear();
  privateLoopClosures.clear();
  sharedLoopClosures.clear();

  neighborPoseDict.clear();
  neighborAuxPoseDict.clear();
  localSharedPoseIDs.clear();
  neighborSharedPoseIDs.clear();
  neighborRobotIDs.clear();
  resetTeamStatus();

  if (mProblemPtr) {
    delete mProblemPtr;
    mProblemPtr = nullptr;
  }
  mRobustCost.reset();
  globalAnchor.reset();
  TLocalInit.reset();
  XInit.reset();

  mOptimizationRequested = false;
  mPublishPublicPosesRequested = false;
  mPublishWeightsRequested = false;

  n = 1;
  X = Matrix::Zero(r, d + 1);
  X.block(0, 0, d, d) = Matrix::Identity(d, d);
}

void PGOAgent::iterate(bool doOptimization) {
  mIterationNumber++;

  // Save early stopped solution
  if (mIterationNumber == 50 && mParams.logData) {
    Matrix T;
    if (getTrajectoryInGlobalFrame(T)) {
      mLogger.logTrajectory(dimension(), num_poses(), T, "trajectory_early_stop.csv");
    }
  }

  // Update measurement weights (GNC)
  if (shouldUpdateLoopClosureWeights()) {
    updateLoopClosuresWeights();
    mRobustCost.update();
    // If warm start is disabled, reset trajectory estimate to initial guess
    if (!mParams.robustOptWarmStart) {
      assert(XInit);
      X = XInit.value();
      printf("Warm start is disabled. Robot %u resets trajectory estimates.\n", getID());
    }
    // Reset acceleration
    if (mParams.acceleration) {
      initializeAcceleration();
    }

  }

  // Perform iteration
  if (mState == PGOAgentState::INITIALIZED) {
    // Save current iterate
    XPrev = X;

    // lock pose update
    unique_lock<mutex> tLock(mPosesMutex);

    // lock measurements
    unique_lock<mutex> mLock(mMeasurementsMutex);

    // lock neighbor pose update
    unique_lock<mutex> lock(mNeighborPosesMutex);

    bool success;
    if (mParams.acceleration) {
      updateGamma();
      updateAlpha();
      updateY();
      success = updateX(doOptimization, true);
      updateV();
      // Check restart condition
      if (shouldRestart()) {
        restartNesterovAcceleration(doOptimization);
      }
      mPublishPublicPosesRequested = true;
    } else {
      success = updateX(doOptimization, false);
      if (doOptimization) {
        mPublishPublicPosesRequested = true;
      }
    }

    if (doOptimization) {
      mStatus.agentID = getID();
      mStatus.state = mState;
      mStatus.instanceNumber = instance_number();
      mStatus.iterationNumber = iteration_number();
      mStatus.relativeChange = sqrt((X - XPrev).squaredNorm() / num_poses());
      // Check local termination condition
      bool readyToTerminate = true;
      if (!success) readyToTerminate = false;
      if (mStatus.relativeChange > mParams.relChangeTol) readyToTerminate = false;
      double ratio = computeConvergedLoopClosureRatio();
      if (ratio < mParams.robustOptMinConvergenceRatio) readyToTerminate = false;
      mStatus.readyToTerminate = readyToTerminate;
    }
  }
}

void PGOAgent::constructQMatrix() {
  vector<RelativeSEMeasurement> privateMeasurements = odometry;
  privateMeasurements.insert(privateMeasurements.end(), privateLoopClosures.begin(), privateLoopClosures.end());

  // Initialize Q with private measurements
  SparseMatrix Q = constructConnectionLaplacianSE(privateMeasurements);

  // Initialize relative SE matrix in homogeneous form
  Matrix T = Matrix::Zero(d + 1, d + 1);

  // Initialize aggregate weight matrix
  Matrix Omega = Matrix::Zero(d + 1, d + 1);

  // Go through shared loop closures
  for (const auto &m: sharedLoopClosures) {
    // Set relative SE matrix (homogeneous form)
    T.block(0, 0, d, d) = m.R;
    T.block(0, d, d, 1) = m.t;
    T(d, d) = 1;

    // Set aggregate weight matrix
    for (unsigned row = 0; row < d; ++row) {
      Omega(row, row) = m.weight * m.kappa;
    }
    Omega(d, d) = m.weight * m.tau;

    if (m.r1 == mID) {
      // First pose belongs to this robot
      // Hence, this is an outgoing edge in the pose graph
      assert(m.r2 != mID);

      // Modify quadratic cost
      size_t idx = m.p1;

      Matrix W = T * Omega * T.transpose();

      for (size_t col = 0; col < d + 1; ++col) {
        for (size_t row = 0; row < d + 1; ++row) {
          Q.coeffRef(idx * (d + 1) + row, idx * (d + 1) + col) += W(row, col);
        }
      }

    } else {
      // Second pose belongs to this robot
      // Hence, this is an incoming edge in the pose graph
      assert(m.r2 == mID);

      // Modify quadratic cost
      size_t idx = m.p2;

      for (size_t col = 0; col < d + 1; ++col) {
        for (size_t row = 0; row < d + 1; ++row) {
          Q.coeffRef(idx * (d + 1) + row, idx * (d + 1) + col) +=
              Omega(row, col);
        }
      }
    }
  }

  assert(mProblemPtr);
  mProblemPtr->setQ(Q);
}

bool PGOAgent::constructGMatrix(const PoseDict &poseDict) {
  SparseMatrix G(relaxation_rank(), (dimension() + 1) * num_poses());

  for (const auto &m: sharedLoopClosures) {
    // Construct relative SE matrix in homogeneous form
    Matrix T = Matrix::Zero(d + 1, d + 1);
    T.block(0, 0, d, d) = m.R;
    T.block(0, d, d, 1) = m.t;
    T(d, d) = 1;

    // Construct aggregate weight matrix
    Matrix Omega = Matrix::Zero(d + 1, d + 1);
    for (unsigned row = 0; row < d; ++row) {
      Omega(row, row) = m.weight * m.kappa;
    }
    Omega(d, d) = m.weight * m.tau;

    if (m.r1 == mID) {
      // First pose belongs to this robot
      // Hence, this is an outgoing edge in the pose graph
      assert(m.r2 != mID);

      // Read neighbor's pose
      const PoseID nID = std::make_pair(m.r2, m.p2);
      auto KVpair = poseDict.find(nID);
      if (KVpair == poseDict.end()) {
        if (mParams.verbose) {
          printf("constructGMatrix: robot %u cannot find neighbor pose (%u, %u)\n",
                 getID(), nID.first, nID.second);
        }
        return false;
      }
      Matrix Xj = KVpair->second;

      size_t idx = m.p1;

      // Modify linear cost
      Matrix L = -Xj * Omega * T.transpose();
      for (size_t col = 0; col < d + 1; ++col) {
        for (size_t row = 0; row < r; ++row) {
          G.coeffRef(row, idx * (d + 1) + col) += L(row, col);
        }
      }

    } else {
      // Second pose belongs to this robot
      // Hence, this is an incoming edge in the pose graph
      assert(m.r2 == mID);

      // Read neighbor's pose
      const PoseID nID = std::make_pair(m.r1, m.p1);
      auto KVpair = poseDict.find(nID);
      if (KVpair == poseDict.end()) {
        if (mParams.verbose) {
          printf("constructGMatrix: robot %u cannot find neighbor pose (%u, %u)\n",
                 getID(), nID.first, nID.second);
        }
        return false;
      }
      Matrix Xi = KVpair->second;

      size_t idx = m.p2;

      // Modify linear cost
      Matrix L = -Xi * T * Omega;
      for (size_t col = 0; col < d + 1; ++col) {
        for (size_t row = 0; row < r; ++row) {
          G.coeffRef(row, idx * (d + 1) + col) += L(row, col);
        }
      }
    }
  }

  assert(mProblemPtr);
  mProblemPtr->setG(G);
  return true;
}

void PGOAgent::startOptimizationLoop(double freq) {
  // Asynchronous updates currently restricted to non-accelerated updates
  assert(!mParams.acceleration);

  if (isOptimizationRunning()) {
    if (mParams.verbose)
      printf("startOptimizationLoop: optimization thread already running! \n");
    return;
  }

  mRate = freq;

  mOptimizationThread = new thread(&PGOAgent::runOptimizationLoop, this);
}

void PGOAgent::runOptimizationLoop() {
  if (mParams.verbose)
    printf("Robot %u optimization thread running at %f Hz.\n", getID(), mRate);

  // Create exponential distribution with the desired rate
  std::random_device rd;  // Will be used to obtain a seed for the random number engine
  std::mt19937 rng(rd());  // Standard mersenne_twister_engine seeded with rd()
  std::exponential_distribution<double> ExponentialDistribution(mRate);

  while (true) {
    double sleepUs =
        1e6 * ExponentialDistribution(rng);  // sleeping time in microsecond

    usleep(sleepUs);

    iterate(true);

    // Check if finish requested
    if (mEndLoopRequested) {
      break;
    }
  }
}

void PGOAgent::endOptimizationLoop() {
  if (!isOptimizationRunning()) return;

  mEndLoopRequested = true;

  // wait for thread to finish
  mOptimizationThread->join();

  delete mOptimizationThread;

  mOptimizationThread = nullptr;

  mEndLoopRequested = false;  // reset request flag

  if (mParams.verbose)
    printf("Robot %u optimization thread exited. \n", getID());
}

bool PGOAgent::isOptimizationRunning() {
  return mOptimizationThread != nullptr;
}

RelativeSEMeasurement &PGOAgent::findSharedLoopClosureWithNeighbor(const PoseID &nID) {
  const unsigned neighborID = nID.first;
  const unsigned neighborPose = nID.second;
  for (auto &m: sharedLoopClosures) {
    if ((m.r1 == neighborID && m.p1 == neighborPose) ||
        (m.r2 == neighborID && m.p2 == neighborPose)) {
      return m;
    }
  }

  // The desired measurement is not found. Throw a runtime error.
  throw std::runtime_error("Cannot find shared loop closure with neighbor.");
}

RelativeSEMeasurement &PGOAgent::findSharedLoopClosure(const PoseID &srcID, const PoseID &dstID) {
  for (auto &m: sharedLoopClosures) {
    if (m.r1 == srcID.first && m.p1 == srcID.second && dstID.first == m.r2 && dstID.second == m.p2) {
      return m;
    }
  }

  // The desired measurement is not found. Throw a runtime error.
  throw std::runtime_error("Cannot find specified shared loop closure.");
}

void PGOAgent::localInitialization() {
  std::vector<RelativeSEMeasurement> measurements = odometry;
  measurements.insert(measurements.end(), privateLoopClosures.begin(), privateLoopClosures.end());

  Matrix T0;
  if (mParams.robustCostType == RobustCostType::L2) {
    T0 = chordalInitialization(dimension(), num_poses(), measurements);
  } else {
    // In robust mode, we do not trust the loop closures and hence initialize from odometry
    T0 = odometryInitialization(dimension(), num_poses(), odometry);
  }

  assert(T0.rows() == d);
  assert(T0.cols() == (d + 1) * n);
  TLocalInit.emplace(T0);
}

Matrix PGOAgent::localPoseGraphOptimization() {
  // Compute initialization if necessary
  if (!TLocalInit)
    localInitialization();

  // Compute connection laplacian
  std::vector<RelativeSEMeasurement> measurements = odometry;
  measurements.insert(measurements.end(), privateLoopClosures.begin(), privateLoopClosures.end());
  SparseMatrix Q = constructConnectionLaplacianSE(measurements);

  // Form optimization problem
  QuadraticProblem problem(n, d, d);
  problem.setQ(Q);

  // Initialize optimizer object
  QuadraticOptimizer optimizer(&problem);
  optimizer.setVerbose(mParams.verbose);
  optimizer.setTrustRegionInitialRadius(10);
  optimizer.setTrustRegionIterations(10);
  optimizer.setTrustRegionTolerance(1e-1);
  optimizer.setTrustRegionMaxInnerIterations(50);

  // Optimize
  Matrix Topt = optimizer.optimize(TLocalInit.value());
  if (mParams.verbose) printf("Optimization time: %f sec.\n", optimizer.getOptResult().elapsedMs / 1e3);
  return Topt;
}

bool PGOAgent::getLiftingMatrix(Matrix &M) const {
  assert(mID == 0);
  if (YLift.has_value()) {
    M = YLift.value();
    return true;
  }
  return false;
}

void PGOAgent::setGlobalAnchor(const Matrix &M) {
  assert(M.rows() == relaxation_rank());
  assert(M.cols() == dimension() + 1);
  globalAnchor.emplace(M);
}

bool PGOAgent::shouldTerminate() {
  // terminate if reached maximum iterations
  if (iteration_number() > mParams.maxNumIters) {
    printf("Reached maximum iterations.\n");
    return true;
  }

  for (size_t robot = 0; robot < mParams.numRobots; ++robot) {
    PGOAgentStatus robotStatus = mTeamStatus[robot];
    assert(robotStatus.agentID == robot);
    if (robotStatus.state != PGOAgentState::INITIALIZED) {
      return false;
    }
  }

  // Check if all agents are ready to terminate optimization
  for (size_t robot = 0; robot < mParams.numRobots; ++robot) {
    PGOAgentStatus robotStatus = mTeamStatus[robot];
    if (!robotStatus.readyToTerminate) {
      return false;
    }
  }

  return true;
}

bool PGOAgent::shouldRestart() const {
  if (mParams.acceleration) {
    return ((mIterationNumber + 1) % mParams.restartInterval == 0);
  }
  return false;
}

void PGOAgent::restartNesterovAcceleration(bool doOptimization) {
  if (mParams.acceleration && mState == PGOAgentState::INITIALIZED) {
    if (mParams.verbose) {
      printf("Robot %u restarts Nesteorv acceleration.\n", getID());
    }
    X = XPrev;
    updateX(doOptimization, false);
    V = X;
    Y = X;
    gamma = 0;
    alpha = 0;
  }
}

void PGOAgent::initializeAcceleration() {
  assert(mParams.acceleration);
  if (mState == PGOAgentState::INITIALIZED) {
    XPrev = X;
    gamma = 0;
    alpha = 0;
    V = X;
    Y = X;
  }
}

void PGOAgent::updateGamma() {
  assert(mParams.acceleration);
  assert(mState == PGOAgentState::INITIALIZED);
  gamma = (1 + sqrt(1 + 4 * pow(mParams.numRobots, 2) * pow(gamma, 2))) / (2 * mParams.numRobots);
}

void PGOAgent::updateAlpha() {
  assert(mParams.acceleration);
  assert(mState == PGOAgentState::INITIALIZED);
  alpha = 1 / (gamma * mParams.numRobots);
}

void PGOAgent::updateY() {
  assert(mParams.acceleration);
  assert(mState == PGOAgentState::INITIALIZED);
  LiftedSEManifold manifold(relaxation_rank(), dimension(), num_poses());
  Matrix M = (1 - alpha) * X + alpha * V;
  Y = manifold.project(M);
}

void PGOAgent::updateV() {
  assert(mParams.acceleration);
  assert(mState == PGOAgentState::INITIALIZED);
  LiftedSEManifold manifold(relaxation_rank(), dimension(), num_poses());
  Matrix M = V + gamma * (X - Y);
  V = manifold.project(M);
}

bool PGOAgent::updateX(bool doOptimization, bool acceleration) {
  if (!doOptimization) {
    if (acceleration) {
      X = Y;
    }
    return true;
  }

  if (mParams.verbose)
    printf("Robot %u optimize at iteration %u... \n", getID(), iteration_number());

  if (acceleration)
    assert(mParams.acceleration);

  assert(mState == PGOAgentState::INITIALIZED);

  // Update quadratic cost matrix (unless using L2 cost function since it does not change measurement weights)
  if (mParams.robustCostType != RobustCostType::L2) {
    constructQMatrix();
  }

  // Construct linear cost matrix (depending on neighbor robots' poses)
  bool hasG;
  if (acceleration) {
    hasG = constructGMatrix(neighborAuxPoseDict);
  } else {
    hasG = constructGMatrix(neighborPoseDict);
  }

  // Skip update if G matrix is not constructed successfully
  if (!hasG) {
    if (mParams.verbose) {
      printf("Robot %u could not construct G matrix. Skip update...\n", getID());
    }
    return false;
  }

  // Initialize optimizer
  QuadraticOptimizer optimizer(mProblemPtr);
  optimizer.setVerbose(mParams.verbose);
  optimizer.setAlgorithm(mParams.algorithm);
  optimizer.setTrustRegionTolerance(1e-2); // Force optimizer to make progress
  optimizer.setTrustRegionIterations(1);
  optimizer.setTrustRegionMaxInnerIterations(10);
  optimizer.setTrustRegionInitialRadius(100);

  // Starting solution
  Matrix XInit;
  if (acceleration) {
    XInit = Y;
  } else {
    XInit = X;
  }
  assert(XInit.rows() == relaxation_rank());
  assert(XInit.cols() == (dimension() + 1) * num_poses());

  // Optimize!
  X = optimizer.optimize(XInit);
  assert(X.rows() == relaxation_rank());
  assert(X.cols() == (dimension() + 1) * num_poses());

  // Print optimization statistics
  const auto &result = optimizer.getOptResult();
  if (mParams.verbose) {
    printf("df: %f, gn0: %f, gn1: %f, df/gn0: %f\n",
           result.fInit - result.fOpt,
           result.gradNormInit,
           result.gradNormOpt,
           (result.fInit - result.fOpt) / result.gradNormInit);
  }

  return true;
}

void PGOAgent::resetTeamStatus() {
  mTeamStatus.clear();
  for (unsigned robot = 0; robot < mParams.numRobots; ++robot) {
    mTeamStatus.emplace_back(robot);
  }
}

bool PGOAgent::shouldUpdateLoopClosureWeights() const {
  // No need to update weight if using L2 cost
  if (mParams.robustCostType == RobustCostType::L2) return false;

  return ((mIterationNumber + 1) % mParams.robustOptInnerIters == 0);
}

void PGOAgent::updateLoopClosuresWeights() {
  assert(mState == PGOAgentState::INITIALIZED);

  // Update private loop closures
  for (auto &m: privateLoopClosures) {
    if (m.isKnownInlier) continue;
    Matrix Y1 = X.block(0, m.p1 * (d + 1), r, d);
    Matrix p1 = X.block(0, m.p1 * (d + 1) + d, r, 1);
    Matrix Y2 = X.block(0, m.p2 * (d + 1), r, d);
    Matrix p2 = X.block(0, m.p2 * (d + 1) + d, r, 1);
    double residual = std::sqrt(computeMeasurementError(m, Y1, p1, Y2, p2));
    double weight = mRobustCost.weight(residual);
    m.weight = weight;
    if (mParams.verbose) {
      printf("Agent %u update edge: (%zu, %zu) -> (%zu, %zu), residual = %f, weight = %f \n",
             getID(), m.r1, m.p1, m.r2, m.p2, residual, weight);
    }
  }

  // Update shared loop closures
  // Agent i is only responsible for updating edge weights with agent j, where j > i
  for (auto &m: sharedLoopClosures) {
    if (m.isKnownInlier) continue;
    Matrix Y1, Y2, p1, p2;
    if (m.r1 == getID()) {
      if (m.r2 < getID()) continue;

      Y1 = X.block(0, m.p1 * (d + 1), r, d);
      p1 = X.block(0, m.p1 * (d + 1) + d, r, 1);
      const PoseID nbrPoseID = std::make_pair(m.r2, m.p2);
      auto KVpair = neighborPoseDict.find(nbrPoseID);
      if (KVpair == neighborPoseDict.end()) {
        printf("Agent %u cannot update edge: (%zu, %zu) -> (%zu, %zu). \n",
               getID(), m.r1, m.p1, m.r2, m.p2);
        continue;
      }
      Matrix X2 = KVpair->second;
      Y2 = X2.block(0, 0, r, d);
      p2 = X2.block(0, d, r, 1);
    } else {
      if (m.r1 < getID()) continue;

      Y2 = X.block(0, m.p2 * (d + 1), r, d);
      p2 = X.block(0, m.p2 * (d + 1) + d, r, 1);
      const PoseID nbrPoseID = std::make_pair(m.r1, m.p1);
      auto KVpair = neighborPoseDict.find(nbrPoseID);
      if (KVpair == neighborPoseDict.end()) {
        printf("Agent %u cannot update edge: (%zu, %zu) -> (%zu, %zu). \n",
               getID(), m.r1, m.p1, m.r2, m.p2);
        continue;
      }
      Matrix X1 = KVpair->second;
      Y1 = X1.block(0, 0, r, d);
      p1 = X1.block(0, d, r, 1);
    }
    double residual = std::sqrt(computeMeasurementError(m, Y1, p1, Y2, p2));
    double weight = mRobustCost.weight(residual);
    m.weight = weight;
    if (mParams.verbose) {
      printf("Agent %u update edge: (%zu, %zu) -> (%zu, %zu), residual = %f, weight = %f \n",
             getID(), m.r1, m.p1, m.r2, m.p2, residual, weight);
    }
  }
  mPublishWeightsRequested = true;
}

double PGOAgent::computeConvergedLoopClosureRatio() {
  // Currently, this function is only meaningful for GNC_TLS
  if (mParams.robustCostType != RobustCostType::GNC_TLS) {
    return 1.0;
  }

  double totalCount = 0;
  double acceptCount = 0;
  double rejectCount = 0;
  for (const auto &m: privateLoopClosures) {
    if (m.isKnownInlier) continue;
    if (m.weight == 1) {
      acceptCount += 1;
    } else if (m.weight == 0) {
      rejectCount += 1;
    }
    totalCount += 1;
  }
  for (const auto &m: sharedLoopClosures) {
    if (m.isKnownInlier) continue;
    if (m.weight == 1) {
      acceptCount += 1;
    } else if (m.weight == 0) {
      rejectCount += 1;
    }
    totalCount += 1;
  }
  double convergedCount = acceptCount + rejectCount;

  if (mParams.verbose) {
    printf(
        "Robot %u :\n "
        "accepted loop closures: %i\n "
        "rejected loop closures: %i\n "
        "undecided loop closures: %i\n",
        getID(),
        (int) acceptCount,
        (int) rejectCount,
        (int) (totalCount - convergedCount));
  }

  return convergedCount / totalCount;
}

bool PGOAgent::isDuplicateMeasurement(const RelativeSEMeasurement &m,
                                      const vector<RelativeSEMeasurement> &measurements) {
  for (const RelativeSEMeasurement &m2: measurements) {
    if (m.r1 == m2.r1 && m.r2 == m2.r2 && m.p1 == m2.p1 && m.p2 == m2.p2) {
      return true;
    }
  }
  return false;
}

}  // namespace DPGO


/*
Copyright (c) 2019 Matthew H. Reilly (kb1vc)
All rights reserved.

Redistribution and use in source and binary forms, with or without
modification, are permitted provided that the following conditions are
met:

    Redistributions of source code must retain the above copyright
    notice, this list of conditions and the following disclaimer.

    Redistributions in binary form must reproduce the above copyright
    notice, this list of conditions and the following disclaimer in
    the documentation and/or other materials provided with the
    distribution.

THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS
"AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR
A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT
HOLDER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL,
SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT
LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE,
DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY
THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
(INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
*/
#ifndef POINT_HDR_DEF
#define POINT_HDR_DEF
#include <string>
#include <regex>
#include <cmath>
#include <boost/format.hpp>

/**
 * \class GeoProf::Point
 *
 * \brief Point implements navigation operations on a geographic point
 * on the surface of the earth. The Point object includes methods to
 * calculate the distance and bearing from one point to another, or the
 * location (point) at a given distance and bearing from a starting location.
 *
 * \author $Author: kb1vc $
 *
 * \date $Date: 2005/04/14 14:16:20 $
 *
 * Contact: kb1vc@kb1vc.org
 *
 */
namespace GeoProf {
  class Point {
  public:
    /**
     * @brief Create a point object from latitude, and longitude
     *
     * @param _lat the latitude of the point -- negative values are south of the equator
     * @param _lon the longitude of the point -- negative values are west of Grenwich
     *
     * Note that bearings between the poles are unlikely to make sense.
     * In places where every direction is "south" or "north" the math
     * is often beyond sensible limits.  If your application requires
     * pole-to-pole paths, then look elsewhere, or just point south.
     */
    Point(double _lat = 0.0, double _lon = 0.0) {
      lat = _lat;
      lon = _lon;
    }

    /**
     * @brief Create a point object from another point object
     *
     * @param orig the point that we're cloning.
     */
    Point(const Point & orig) {
      lat = orig.lat;
      lon = orig.lon;
    }

    /**
     * @brief Create a point from a Maidenhead Grid specifier
     *
     * @param grid A Maidenhead Grid specifier of the form XXnnxx (capital
     * letter pair, digit pair, letter pair)
     */
    Point(const std::string & grid) {
      grid2Pt(grid);
    }

    /**
     * @brief Return degrees latitude (negative is south of the equator)
     *
     * @return latitude
     */
    double getLatitude() const { return lat; }

    /**
     * @brief Return degrees longitude (negative is west of the Grenwich)
     *
     * @return longitude
     */
    double getLongitude() const { return lon; }

    /**
     * @brief Calculate the bearing (in degrees -- 0 is north) from this point to anothe point
     * along the shorter great circle route
     *
     * @param other the other point
     * @return bearing in degrees.
     */
    double bearingTo(const Point & other) const;

    /**
     * @brief Calculate the great circle distance (in meters) from this point to another point
     *
     * @param other the other point
     * @return great circle distance in meters
     */
    double distanceTo(const Point & other) const;

    /**
     * @brief Calculate the bearing (in degrees -- 0 is north) and distance (in meters)
     * from this point to another point
     * along the shorter great circle route
     *
     * @param other the other point
     * @param bearing (output) direction to the other point along the shorter great circle path
     * @param reverse_bearing direction of travel from other point to this point (in degrees, 0 is north)
     * @param distance (output) distance in meters to the other point along the shorter great circle path
     */
    void bearingDistanceTo(const Point & other, double & bearing, double & reverse_bearing, double & distance) const;

    /**
     * @brief Return the point at which one would arrive after traveling on the
     * great circle path from this point at the specified bearing and for the specified distance.
     *
     * @param bearing direction of travel (in degrees, 0 is north)
     * @param distance of travel in meters
     * @param next (output) the point at which we'll arrive.
     */
    void stepTo(double bearing, double distance, Point & next) const;

    /**
     * @brief convert this point to its Maidenhead Grid specifier
     *
     * @param grid (output) The Maidenhead Grid string XXnnxx
     */
    void pt2Grid(std::string & grid) const;

    /**
     * @brief Set this point to the location of a Maidenhead Grid specifier
     *
     * @param grid A Maidenhead Grid specifier of the form XXnnxx (capital
     * letter pair, digit pair, letter pair)
     */
    void grid2Pt(const std::string & grid);

    /**
     * @brief Set this point from a string describing the location in
     * degrees-minutes-seconds of latitude and logitude. Specification is
     * in lat degrees/minutes/seconds and lon degrees/minutes/seconds as
     * doubles.  North/South/East/West
     *
     * @param lat_d latitude degrees
     * @param lat_m latitude minutes
     * @param lat_s latitude seconds
     * @param ns true if latitude is north of the equator
     * @param lon_d longitude degrees
     * @param lon_m longitude minutes
     * @param lon_s longitude seconds
     * @param ew true if latitude is east of Grenwich
     */
    void dms2Pt(double lat_d, double lat_m, double lat_s, char ns,
		 double lon_d, double lon_m, double lon_s, char ew);

    /**
     * @brief Correct the calculated bearing and distance by iterating
     * around the calculated values to find the minimum error.
     *
     * @param other the other point
     * @param bearing direction to the other point along the shorter great circle path
     * @param distance distance in meters to the other point along the shorter great circle path
     * @param new_bearing (output) corrected bearing
     * @param new_distance (output) corrected range
     *
     */
    void correctBearingDistanceTo(const Point & other,
				  const double bearing,
				  const double distance,
				  double & new_bearing,
				  double & new_distance) const;

    std::string toString() const {
      char ns = ' ';
      if (lat > 0.) {
	ns = 'N';
      }
      if (lat < 0.0) {
	ns = 'S';
      }

      char ew = ' ';
      if (lon > 0.0) {
	ew = 'E';
      }
      if (lon < 0.0) {
	ew = 'W';
      }

      return (boost::format("%g %c %g %c") % lat % ns % lon % ew).str();
    }
  private:
    /// Latitude South is negative, North is positive.
    double lat;

    /// Longitude West is negative, East is positive.
    double lon;

    static std::regex grid_regexp; //  ("[A-R][A-R][0-9][0-9][A-X][A-X]", std::regex_constants::icase);

    /**
     * @brief Validate a string as a Maidenhead Grid specifier
     *
     * @param grid A Maidenhead Grid specifier of the form XXnnxx (capital
     * letter pair, digit pair, letter pair)
     * @return true if this is a properly formatted grid, false otherwise.
     */
    bool checkGrid(const std::string & grid) const;

    /**
     * @brief Helper for translating char positions in a grid
     * locator into double offsets from 0 degrees.
     */
    double gridDiff(char v, char s, double mul) const;

    /**
     * @brief four quadrant arc tangent returning result in the range 0..2pi
     *
     * @param y rise
     * @param x run
     * @return arctan in range 0..2pi
     */
    double atan2Pt(double y, double x) const;

    /**
     * @brief translate an angle into the range 0..2pi
     *
     * @param ang angle in radians
     * @return angle in range 0..2pi
     */
    double inSpan(double ang) const;

    // recursive helper to correctBearingDistanceTo
    bool recCorrectBearingDistanceTo(const Point & other,
				     const double bearing,
				     const double distance,
				     const double b_span,
				     const double d_span,
				     double & new_bearing,
				     double & new_distance
				     ) const;

    // Useful constants
    static const double clarke_66_al;    /*Clarke 1866 ellipsoid*/
    static const double clarke_66_bl;
    static const double rad_per_deg;
  };

}
#endif


#pragma once

#include <gmock/gmock-matchers.h>

#include <Eigen/Geometry>

#include <type_traits>

template <typename Matrix>
struct SameMatrixMatcher : testing::MatcherInterface<Matrix const &>
{
    using Scalar = typename Matrix::Scalar;
    SameMatrixMatcher(Matrix const & expected, Scalar epsilon) :
        mExpected(expected),
        mEpsilon(epsilon)
    {}

    bool MatchAndExplain(Matrix const & matrix, testing::MatchResultListener* listener) const override
    {
        Matrix difference = mExpected - matrix;
        auto differenceMagnitude = difference.norm();
        *listener << "different by " << differenceMagnitude;
        return differenceMagnitude < mEpsilon;
    }

    void DescribeTo(std::ostream* os) const override
    {
        *os << "same as\n" << mExpected;
    }

    void DescribeNegationTo(std::ostream* os) const override
    {
        *os << "different from\n" << mExpected;
    }
private:
    Matrix const mExpected;
    Scalar const mEpsilon;
};

template <typename Quaternion>
struct SameOrientationMatcher : testing::MatcherInterface<Quaternion const &>
{
    using Scalar = typename Quaternion::Scalar;

    SameOrientationMatcher(Quaternion const & expected, Scalar epsilon) :
        mExpected(expected),
        mEpsilon(epsilon)
    {}

    bool MatchAndExplain(Quaternion const & quaternion, testing::MatchResultListener* listener) const override
    {
        auto differenceMagnitude = 1 - std::abs(mExpected.dot(quaternion));
        *listener << "different by " << differenceMagnitude;
        return differenceMagnitude < mEpsilon;
    }

    void DescribeTo(std::ostream* os) const override
    {
        *os << "same as [ " << mExpected.coeffs().transpose() << " ]";
    }

    void DescribeNegationTo(std::ostream* os) const override
    {
        *os << "different from [ " << mExpected.coeffs().transpose() << " ]";
    }
private:
    Quaternion const mExpected;
    Scalar const mEpsilon;
};

template<typename T>
inline auto IsSameMatrix(T const & expected, typename T::Scalar epsilon) {
    return testing::MakePolymorphicMatcher(SameMatrixMatcher<typename T::PlainMatrix>(expected, epsilon));
}

template<typename T>
inline auto IsSameOrientation(T const & expected, typename T::Scalar epsilon) {
    return testing::MakePolymorphicMatcher(SameOrientationMatcher<Eigen::Quaternion<typename T::Scalar>>(expected, epsilon));
}