#ifndef OPENGV2_BSPLINEREAL_HPP
#define OPENGV2_BSPLINEREAL_HPP

#include <Eigen/Eigen>
#include <Eigen/StdVector>
#include <Eigen/Sparse>
#include <Eigen/SparseCholesky>
#include <cmath>
#include <algorithm>
#include <numeric>
#include <iostream>

namespace opengv2 {
    template<int dim>
    class BsplineReal {
    public:
        EIGEN_MAKE_ALIGNED_OPERATOR_NEW

        explicit BsplineReal(int p = dim,
                             const std::vector<Eigen::Matrix<double, dim, 1>, Eigen::aligned_allocator<Eigen::Matrix<double, dim, 1>>> &Q = std::vector<Eigen::Matrix<double, dim, 1>, Eigen::aligned_allocator<Eigen::Matrix<double, dim, 1>>>(),
                             int controlPointsNum = -1,
                             const std::vector<double> &u = std::vector<double>(),
                             const std::vector<Eigen::Matrix<double, dim, 1>, Eigen::aligned_allocator<Eigen::Matrix<double, dim, 1>>> &dQ = std::vector<Eigen::Matrix<double, dim, 1>, Eigen::aligned_allocator<Eigen::Matrix<double, dim, 1>>>(),
                         bool derWithMagnitude = false, double lambda = 1) : degree(p), lambda(lambda) {
            // if dataPoints given, then do approximation
            if (!Q.empty()) {
                dataPoints = Q;

                // Data registration
                if (!u.empty()) {
                    correspondingUs = u;
                } else {
                    // Data registration using chord length
                    // Range: [0,1]
                    correspondingUs.resize(dataPoints.size());

                    double d = 0;
                    for (int i = 1; i < dataPoints.size(); ++i) {
                        d += (dataPoints[i] - dataPoints[i - 1]).norm();
                    }

                    correspondingUs.front() = 0;
                    correspondingUs.back() = 1;
                    for (int i = 1; i < dataPoints.size() - 1; ++i) {
                        correspondingUs[i] =
                                correspondingUs[i - 1] + (dataPoints[i] - dataPoints[i - 1]).norm() / d;
                    }
                }

                // derivative
                if (!dQ.empty()) {
                    derivatives = dQ;

                    if (!derWithMagnitude) {
                        // derivatives is unit direction, magnitude estimation by speed estimation
                        for (int i = 1; i <= dataPoints.size() - 2; ++i) {
                            double speed = (dataPoints[i] - dataPoints[i - 1]).norm() /
                                           (correspondingUs[i] - correspondingUs[i - 1]) +
                                           (dataPoints[i + 1] - dataPoints[i]).norm() /
                                           (correspondingUs[i + 1] - correspondingUs[i]);
                            derivatives[i] *= speed / 2;
                        }
                        derivatives.front() *=
                                (dataPoints[1] - dataPoints[0]).norm() / (correspondingUs[1] - correspondingUs[0]);
                        derivatives.back() *=
                                (dataPoints[dataPoints.size() - 1] - dataPoints[dataPoints.size() - 2]).norm() /
                                (correspondingUs[dataPoints.size() - 1] - correspondingUs[dataPoints.size() - 2]);
                    }
                }

                // approximation
                if (controlPointsNum < 0)
                    controlPointsNum = std::max(int(Q.size() / 3), degree + 1);
                if (controlPointsNum <= degree) {
                    std::cerr << "Constructor: control points number should greater than degree!" << std::endl;
                    return;
                }
                approximation(controlPointsNum);
            }
        };

        /*
         * Approximation
         * return: status: failed(-1), success(0)
         */
        int approximation(int controlPointsNum) {
            /*** knot initialization (9.68) in NURBS book ***/
            knotVector.resize(controlPointsNum + degree + 1);
            std::fill_n(knotVector.begin(), degree + 1, correspondingUs.front());
            std::fill_n(knotVector.rbegin(), degree + 1, correspondingUs.back());
            double d = dataPoints.size() / double(controlPointsNum - degree);
            for (int j = 1; j <= controlPointsNum - 1 - degree; j++) {
                int i = floor(j * d);
                double alpha = j * d - i;
                knotVector[degree + j] = (1 - alpha) * correspondingUs[i - 1] + alpha * correspondingUs[i];
            }

            return optimization();
        };

        /*
         * Compute nonzero basis functions and their derivatives up to derivativeLimit.
         * N_{i-p,p}(u),...,N_{i,p}(u)
         * Output: two-dimensional array, ders. ders[k][j] is the kth derivative of the function N_{i-p+j,p}, where 0<=k<=n and 0<=j<=p.
         */
        int dersBasisFuns(double u, size_t spanIdx, int derivativeLimit, std::vector<std::vector<double>> &ders) const {
            // check
            if (knotVector.empty()) {
                std::cerr << "Function dersBasisFuns: knotVector is empty!" << std::endl;
                return -1;
            }

            ders.resize(derivativeLimit + 1);
            for (auto &it:ders) {
                it.assign(degree + 1, 0);
            }

            double ndu[degree + 1][degree + 1]; // store the basis functions and knot differences

            // store (in an alternating fashion) the two most recently computed rows a_{k,j} and a_{k-1,j}
            double a[2][degree + 1];

            std::vector<double> left, right;
            left.resize(degree + 1);
            right.resize(degree + 1);

            ndu[0][0] = 1;
            for (int j = 1; j <= degree; j++) {
                left[j] = u - knotVector[spanIdx + 1 - j];
                right[j] = knotVector[spanIdx + j] - u;
                double saved = 0.0;
                for (int r = 0; r < j; ++r) {
                    ndu[j][r] = right[r + 1] + left[j - r];
                    double temp = ndu[r][j - 1] / ndu[j][r];

                    ndu[r][j] = saved + right[r + 1] * temp;
                    saved = left[j - r] * temp;
                }
                ndu[j][j] = saved;
            }

            /* Load the basis functions */
            for (int j = 0; j <= degree; j++)
                ders[0][j] = ndu[j][degree];

            if (derivativeLimit > 0) {
                /*** This section computes the derivatives (Eq.[2.9]) ***/
                /* Load over function index */
                for (int r = 0; r <= degree; r++) {
                    int s1 = 0, s2 = 1; // Alternate rows in array a
                    a[0][0] = 1;

                    // loop to compute kth derivative
                    for (int k = 1; k <= derivativeLimit; k++) {
                        double d = 0;
                        int rk = r - k, pk = degree - k;
                        if (r >= k) {
                            a[s2][0] = a[s1][0] / ndu[pk + 1][rk];
                            d = a[s2][0] * ndu[rk][pk];
                        }

                        int j1, j2;
                        if (rk >= -1)
                            j1 = 1;
                        else
                            j1 = -rk;

                        if (r - 1 <= pk)
                            j2 = k - 1;
                        else
                            j2 = degree - r;

                        for (int j = j1; j <= j2; j++) {
                            a[s2][j] = (a[s1][j] - a[s1][j - 1]) / ndu[pk + 1][rk + j];
                            d += a[s2][j] * ndu[rk + j][pk];
                        }

                        if (r <= pk) {
                            a[s2][k] = -a[s1][k - 1] / ndu[pk + 1][r];
                            d += a[s2][k] * ndu[r][pk];
                        }

                        ders[k][r] = d;

                        // switch rows
                        std::swap(s1, s2);
                    }
                }

                /* Multiply through by correct factors (Eq. [2.9]) */
                int r = degree;
                for (int k = 1; k <= derivativeLimit; ++k) {
                    for (int j = 0; j <= degree; j++)
                        ders[k][j] *= r;

                    r *= (degree - k);
                }
            }

            return 0;
        }

        /*
         * Determine the knot span index i \in [p,n], whole is [0,n+p+1]
         * u \in [u_{i}, u_{i+1}) , special at u==u_{n+1}
         */
        size_t findSpan(double u) const {
            // assume it's totally p-smooth, which means there is only one knot for each boundary
            // Notes: the size of knot vector is n+p+2
            size_t n = knotVector.size() - 2 - degree;

            // special case
            if (u == knotVector[n + 1])
                return n;

            // binary search
            size_t low = degree;
            size_t high = n + 1;
            size_t mid = (low + high) / 2;
            while (u < knotVector[mid] || u >= knotVector[mid + 1]) {
                if (u < knotVector[mid])
                    high = mid;
                else
                    low = mid;

                mid = (low + high) / 2;
            }

            return mid;
        };

        // shape won't change
        int insertKnot(double u) {
            if (knotVector.empty()) {
                std::cerr << "Function insertKnot: knotVector is empty!" << std::endl;
                return -1;
            }
            if (controlPoints.empty()) {
                std::cerr << "Function insertKnot: controlPoints is empty!" << std::endl;
                return -1;
            }

            size_t spanIdx = findSpan(u);
            std::vector<Eigen::Matrix<double, dim, 1>, Eigen::aligned_allocator<Eigen::Matrix<double, dim,
                    1 >>> Q(controlPoints.size() + 1);

            std::copy(controlPoints.begin(), controlPoints.begin() + spanIdx - degree + 1, Q.begin());
            std::copy(controlPoints.begin() + spanIdx, controlPoints.end(), Q.begin() + spanIdx + 1);
            for (size_t i = spanIdx - degree + 1; i <= spanIdx; i++) {
                double alpha = (u - knotVector[i]) / (knotVector[i + degree] - knotVector[i]);
                Q[i] = alpha * controlPoints[i] + (1 - alpha) * controlPoints[i - 1];
            }

            //update knotVector
            knotVector.insert(knotVector.begin() + spanIdx + 1, u);
            //update controlPoints
            controlPoints = std::move(Q);

            return 0;
        };

        // assume X: ascending sorted
        int refineKnotVect(std::vector<double> &X) {
            if (knotVector.empty()) {
                std::cerr << "Function refineKnotVect: knotVector is empty!" << std::endl;
                return -1;
            }
            if (controlPoints.empty()) {
                std::cerr << "Function refineKnotVect: controlPoints is empty!" << std::endl;
                return -1;
            }

            std::sort(X.begin(), X.end());

            std::vector<Eigen::Matrix<double, dim, 1>, Eigen::aligned_allocator<Eigen::Matrix<double, dim, 1>>> Q(
                    controlPoints.size() + X.size());
            std::vector<double> Ubar(knotVector.size() + X.size());

            size_t n = knotVector.size() - 2 - degree;
            size_t r = X.size() - 1;
            size_t m = knotVector.size() - 1;
            size_t a = findSpan(X.front());
            size_t b = findSpan(X.back());
            b++;

            for (size_t j = 0; j <= a - degree; j++)
                Q[j] = controlPoints[j];
            for (size_t j = b - 1; j <= n; j++)
                Q[j + X.size()] = controlPoints[j];
            for (size_t j = 0; j <= a; j++)
                Ubar[j] = knotVector[j];
            for (size_t j = b + degree; j <= m; j++)
                Ubar[j + X.size()] = knotVector[j];

            int i = b + degree - 1;
            int k = b + degree + r;
            for (int j = r; j >= 0; j--) {
                while (X[j] <= knotVector[i] && i > a) {
                    Q[k - degree - 1] = controlPoints[i - degree - 1];
                    Ubar[k] = knotVector[i];
                    k--;
                    i--;
                }

                Q[k - degree - 1] = Q[k - degree];
                for (int l = 1; l <= degree; l++) {
                    int idx = k - degree + l;
                    double alpha = Ubar[k + l] - X[j];
                    if (std::abs(alpha) == 0.0)
                        Q[idx - 1] = Q[idx];
                    else {
                        alpha /= (Ubar[k + l] - knotVector[i - degree + l]);
                        Q[idx - 1] = alpha * Q[idx - 1] + (1 - alpha) * Q[idx];
                    }
                }
                Ubar[k] = X[j];
                k--;
            }

            //update knotVector
            knotVector = std::move(Ubar);
            //update controlPoints
            controlPoints = std::move(Q);
            return 0;
        };

        // optimize control points w.r.t distance error after knot vector change
        int optimization() {
            /*** check ***/
            if (dataPoints.empty()) {
                std::cerr << "Function optimize: dataPoints is empty!" << std::endl;
                return -1;
            }
            if (dataPoints.size() != derivatives.size() && derivatives.size() != 0) {
                std::cerr
                        << "Function optimization: dataPoints ans derivatives should be same size (NAN for unknow) OR derivatives should be empty!"
                        << std::endl;
                return -1;
            }
            if (knotVector.empty()) {
                std::cerr << "Function optimization: knotVector is empty!" << std::endl;
                return -1;
            }

            int controlPointsNum = knotVector.size() - degree - 1;

            // initialize controlPoints
            controlPoints.resize(controlPointsNum);
            controlPoints.front() = dataPoints.front();
            controlPoints.back() = dataPoints.back();

            // support for derivative lacking
            int derivativesCounter = 0;
            std::vector<int> map(derivatives.size(), -1);
            for (size_t k = 0; k < derivatives.size(); k++) {
                if (!std::isnan(derivatives[k][0])) {
                    map[k] = derivativesCounter++;
                }
            }

            /*** Problem ***/
            Eigen::SparseMatrix<double> N(dataPoints.size(), controlPointsNum);
            Eigen::SparseMatrix<double> M(dataPoints.size(), controlPointsNum);
            std::vector<std::vector<double>> ders;
            for (int k = 0; k < dataPoints.size(); ++k) {
                size_t spanIdx = findSpan(correspondingUs[k]);
                dersBasisFuns(correspondingUs[k], spanIdx, derivativesCounter > 0 ? 1 : 0, ders);

                for (int j = 0; j <= degree; ++j) {
                    if (ders[0][j] != 0)
                        N.insert(k, spanIdx - degree + j) = ders[0][j]; // N_{i-p+j,p}^{0}(\bar{u}_k)

                    if (derivativesCounter > 0) {
                        if (ders[1][j] != 0)
                            M.insert(k, spanIdx - degree + j) = ders[1][j]; // N_{i-p+j,p}^{1}(\bar{u}_k)
                    }
                }
            }

            // formulating B
            std::vector<Eigen::Matrix<double, dim, 1>, Eigen::aligned_allocator<Eigen::Matrix<double, dim, 1>>> B;
            B.assign(controlPointsNum - 2, Eigen::Matrix<double, dim, 1>::Zero());
            for (int l = 1; l <= controlPointsNum - 2; ++l) {
                // for l-th col of N
                for (Eigen::SparseMatrix<double>::InnerIterator it(N, l); it; ++it) {
                    // N_{l,p}(\bar{u}_k), valid for column-major
                    int k = it.row();
                    if (k != 0 && k != dataPoints.size() - 1) {
                        B[l - 1] += it.value() * (dataPoints[k] - N.coeff(k, 0) * dataPoints.front() -
                                                  N.coeff(k, controlPointsNum - 1) * dataPoints.back());
                    }
                }

                if (derivativesCounter > 0) {
                    // for l-th col of M
                    for (Eigen::SparseMatrix<double>::InnerIterator it(M, l); it; ++it) {
                        // N_{l,p}^{1}(\bar{u}_k), valid for column-major
                        int k = it.row();
                        if (map[k] > -1) {
                            B[l - 1] += lambda * it.value() * (derivatives[k] - M.coeff(k, 0) * dataPoints.front() -
                                                               M.coeff(k, controlPointsNum - 1) *
                                                               dataPoints.back());
                        }
                    }
                }
            }

            // formulating M,N
            Eigen::SparseMatrix<double> Nc = N.block(1, 1, dataPoints.size() - 2, controlPointsNum - 2);
            Eigen::SparseMatrix<double> A = Nc.transpose() * Nc;

            if (derivativesCounter > 0) {
                Eigen::SparseMatrix<double> Mc(derivativesCounter, controlPointsNum - 2);
                for (int i = 1; i <= controlPointsNum - 2; ++i) {
                    for (Eigen::SparseMatrix<double>::InnerIterator it(M, i); it; ++it) {// i-th col
                        if (map[it.row()] > -1) {
                            Mc.insert(map[it.row()], it.col() - 1) = it.value();
                        }
                    }
                }
                A += lambda * Mc.transpose() * Mc;
            }

            // solving
            A.makeCompressed();
            Eigen::SimplicialLDLT<Eigen::SparseMatrix<double>> solver;
            solver.compute(A);
            if (solver.info() != Eigen::Success) {
                std::cerr << "Function optimization: decomposition failed!" << std::endl;
                return -1;
            }
            Eigen::VectorXd b(controlPointsNum - 2);
            for (int i = 0; i < dim; i++) {
                for (int j = 0; j < B.size(); ++j) {
                    b[j] = B[j][i];
                }
                Eigen::VectorXd x = solver.solve(b);
                if (solver.info() != Eigen::Success) {
                    std::cerr << "Function optimization: solving failed!" << std::endl;
                    return -1;
                }
                for (int j = 1; j <= controlPointsNum - 2; ++j) {
                    controlPoints[j][i] = x[j - 1];
                }
            }

            return 0;
        };

        /*
         * Evaluate up to k-th derivative at u
         */
        int evaluate(double u, int derivativeLimit,
                     std::vector<Eigen::Matrix<double, dim, 1>, Eigen::aligned_allocator<Eigen::Matrix<double, dim, 1>>> &Ders) const {
            if (controlPoints.empty()) {
                std::cerr << "Function evaluate: controlPoints is empty!" << std::endl;
                return -1;
            }

            Ders.assign(derivativeLimit + 1, Eigen::Matrix<double, dim, 1>::Zero());
            std::vector<std::vector<double>> ders;
            size_t spanIdx = findSpan(u);

            dersBasisFuns(u, spanIdx, derivativeLimit, ders);

            for (size_t i = 0; i < ders.size(); i++) {
                for (size_t j = 0; j < ders[i].size(); ++j) {
                    Ders[i] += ders[i][j] * controlPoints[spanIdx - degree + j];
                }
            }

            return 0;
        }

        inline std::vector<Eigen::Matrix<double, dim, 1>, Eigen::aligned_allocator<Eigen::Matrix<double, dim, 1>>> &
        getCP() {
            return controlPoints;
        };

        inline const std::vector<double> &getCorrespondingUs() {
            return correspondingUs;
        }

        inline const std::vector<double> &getKnotVector() {
            return knotVector;
        }

    protected:
        int degree;

        std::vector<double> knotVector;
        std::vector<Eigen::Matrix<double, dim, 1>, Eigen::aligned_allocator<Eigen::Matrix<double, dim, 1>>> controlPoints;

        // optional
        std::vector<Eigen::Matrix<double, dim, 1>, Eigen::aligned_allocator<Eigen::Matrix<double, dim, 1>>> dataPoints;
        std::vector<double> correspondingUs; //corresponding u for each data point
        std::vector<Eigen::Matrix<double, dim, 1>, Eigen::aligned_allocator<Eigen::Matrix<double, dim, 1>>> derivatives; //corresponding derivatives for each data point
        double lambda;
    };
}

#endif //OPENGV2_BSPLINEREAL_HPP


#include <vector>

#include <boost/shared_ptr.hpp>
#include <gflags/gflags.h>
#include <glog/logging.h>

#include <cmath>

#include "caffe/blob.hpp"
#include "caffe/common.hpp"
#include "caffe/layers/st_layer.hpp"
#include "caffe/util/math_functions.hpp"

namespace caffe {

template <typename Dtype>
void SpatialTransformerLayer<Dtype>::LayerSetUp(const vector<Blob<Dtype>*>& bottom,
      const vector<Blob<Dtype>*>& top) {

	string prefix = "\t\tSpatial Transformer Layer:: LayerSetUp: \t";

	if(this->layer_param_.st_param().transform_type() == "affine") {
		transform_type_ = "affine";
	} else {
		CHECK(false) << prefix << "Transformation type only supports affine now!" << std::endl;
	}

	if(this->layer_param_.st_param().sampler_type() == "bilinear") {
		sampler_type_ = "bilinear";
	} else {
		CHECK(false) << prefix << "Sampler type only supports bilinear now!" << std::endl;
	}

	if(this->layer_param_.st_param().to_compute_du()) {
		to_compute_dU_ = true;
	}

	std::cout<<prefix<<"Getting output_H_ and output_W_"<<std::endl;

	output_H_ = bottom[0]->shape(2);
	if(this->layer_param_.st_param().has_output_h()) {
		output_H_ = this->layer_param_.st_param().output_h();
	}
	output_W_ = bottom[0]->shape(3);
	if(this->layer_param_.st_param().has_output_w()) {
		output_W_ = this->layer_param_.st_param().output_w();
	}

	std::cout<<prefix<<"output_H_ = "<<output_H_<<", output_W_ = "<<output_W_<<std::endl;

	std::cout<<prefix<<"Getting pre-defined parameters"<<std::endl;

	is_pre_defined_theta[0] = false;
	if(this->layer_param_.st_param().has_theta_1_1()) {
		is_pre_defined_theta[0] = true;
		++ pre_defined_count;
		pre_defined_theta[0] = this->layer_param_.st_param().theta_1_1();
		std::cout<<prefix<<"Getting pre-defined theta[1][1] = "<<pre_defined_theta[0]<<std::endl;
	}

	is_pre_defined_theta[1] = false;
	if(this->layer_param_.st_param().has_theta_1_2()) {
		is_pre_defined_theta[1] = true;
		++ pre_defined_count;
		pre_defined_theta[1] = this->layer_param_.st_param().theta_1_2();
		std::cout<<prefix<<"Getting pre-defined theta[1][2] = "<<pre_defined_theta[1]<<std::endl;
	}

	is_pre_defined_theta[2] = false;
	if(this->layer_param_.st_param().has_theta_1_3()) {
		is_pre_defined_theta[2] = true;
		++ pre_defined_count;
		pre_defined_theta[2] = this->layer_param_.st_param().theta_1_3();
		std::cout<<prefix<<"Getting pre-defined theta[1][3] = "<<pre_defined_theta[2]<<std::endl;
	}

	is_pre_defined_theta[3] = false;
	if(this->layer_param_.st_param().has_theta_2_1()) {
		is_pre_defined_theta[3] = true;
		++ pre_defined_count;
		pre_defined_theta[3] = this->layer_param_.st_param().theta_2_1();
		std::cout<<prefix<<"Getting pre-defined theta[2][1] = "<<pre_defined_theta[3]<<std::endl;
	}

	is_pre_defined_theta[4] = false;
	if(this->layer_param_.st_param().has_theta_2_2()) {
		is_pre_defined_theta[4] = true;
		++ pre_defined_count;
		pre_defined_theta[4] = this->layer_param_.st_param().theta_2_2();
		std::cout<<prefix<<"Getting pre-defined theta[2][2] = "<<pre_defined_theta[4]<<std::endl;
	}

	is_pre_defined_theta[5] = false;
	if(this->layer_param_.st_param().has_theta_2_3()) {
		is_pre_defined_theta[5] = true;
		++ pre_defined_count;
		pre_defined_theta[5] = this->layer_param_.st_param().theta_2_3();
		std::cout<<prefix<<"Getting pre-defined theta[2][3] = "<<pre_defined_theta[5]<<std::endl;
	}

	// check the validation for the parameter theta
	CHECK(bottom[1]->count(1) + pre_defined_count == 6) << "The dimension of theta is not six!"
			<< " Only " << bottom[1]->count(1) << " + " << pre_defined_count << std::endl;
	CHECK(bottom[1]->shape(0) == bottom[0]->shape(0)) << "The first dimension of theta and " <<
			"U should be the same" << std::endl;

	// initialize the matrix for output grid
	std::cout<<prefix<<"Initializing the matrix for output grid"<<std::endl;

	vector<int> shape_output(2);
	shape_output[0] = output_H_ * output_W_; shape_output[1] = 3;
	output_grid.Reshape(shape_output);

	Dtype* data = output_grid.mutable_cpu_data();
	for(int i=0; i<output_H_ * output_W_; ++i) {
		data[3 * i] = (i / output_W_) * 1.0 / output_H_ * 2 - 1;
		data[3 * i + 1] = (i % output_W_) * 1.0 / output_W_ * 2 - 1;
		data[3 * i + 2] = 1;
	}

	// initialize the matrix for input grid
	std::cout<<prefix<<"Initializing the matrix for input grid"<<std::endl;

	vector<int> shape_input(3);
	shape_input[0] = bottom[1]->shape(0); shape_input[1] = output_H_ * output_W_; shape_input[2] = 2;
	input_grid.Reshape(shape_input);

	std::cout<<prefix<<"Initialization finished."<<std::endl;
}

template <typename Dtype>
void SpatialTransformerLayer<Dtype>::Reshape(const vector<Blob<Dtype>*>& bottom,
      const vector<Blob<Dtype>*>& top) {

	string prefix = "\t\tSpatial Transformer Layer:: Reshape: \t";

	if(global_debug) std::cout<<prefix<<"Starting!"<<std::endl;

	N = bottom[0]->shape(0);
	C = bottom[0]->shape(1);
	H = bottom[0]->shape(2);
	W = bottom[0]->shape(3);

	// reshape V
	vector<int> shape(4);

	shape[0] = N;
	shape[1] = C;
	shape[2] = output_H_;
	shape[3] = output_W_;

	top[0]->Reshape(shape);

	// reshape dTheta_tmp
	vector<int> dTheta_tmp_shape(4);

	dTheta_tmp_shape[0] = N;
	dTheta_tmp_shape[1] = 2;
	dTheta_tmp_shape[2] = 3;
	dTheta_tmp_shape[3] = output_H_ * output_W_ * C;

	dTheta_tmp.Reshape(dTheta_tmp_shape);

	// init all_ones_2
	vector<int> all_ones_2_shape(1);
	all_ones_2_shape[0] = output_H_ * output_W_ * C;
	all_ones_2.Reshape(all_ones_2_shape);

	// reshape full_theta
	vector<int> full_theta_shape(2);
	full_theta_shape[0] = N;
	full_theta_shape[1] = 6;
	full_theta.Reshape(full_theta_shape);

	if(global_debug) std::cout<<prefix<<"Finished."<<std::endl;
}

template <typename Dtype>
Dtype SpatialTransformerLayer<Dtype>::transform_forward_cpu(const Dtype* pic, Dtype px, Dtype py) {

	bool debug = false;

	string prefix = "\t\tSpatial Transformer Layer:: transform_forward_cpu: \t";

	if(debug) std::cout<<prefix<<"Starting!\t"<<std::endl;
	if(debug) std::cout<<prefix<<"(px, py) = ("<<px<<", "<<py<<")"<<std::endl;

	Dtype res = (Dtype)0.;

	Dtype x = (px + 1) / 2 * H; Dtype y = (py + 1) / 2 * W;

	if(debug) std::cout<<prefix<<"(x, y) = ("<<x<<", "<<y<<")"<<std::endl;

	int m, n; Dtype w;

	m = floor(x); n = floor(y); w = 0;
	if(debug) std::cout<<prefix<<"1: (m, n) = ("<<m<<", "<<n<<")"<<std::endl;

	if(m >= 0 && m < H && n >= 0 && n < W) {
		w = max(0, 1 - abs(x - m)) * max(0, 1 - abs(y - n));
		res += w * pic[m * W + n];
		if(debug) std::cout<<prefix<<"w = "<<w<<", pic[m, n] = "<<pic[m * W + n]<<std::endl;
	}

	m = floor(x) + 1; n = floor(y); w = 0;
	if(debug) std::cout<<prefix<<"2: (m, n) = ("<<m<<", "<<n<<")"<<std::endl;

	if(m >= 0 && m < H && n >= 0 && n < W) {
		w = max(0, 1 - abs(x - m)) * max(0, 1 - abs(y - n));
		res += w * pic[m * W + n];
		if(debug) std::cout<<prefix<<"w = "<<w<<", pic[m, n] = "<<pic[m * W + n]<<std::endl;
	}

	m = floor(x); n = floor(y) + 1; w = 0;
	if(debug) std::cout<<prefix<<"3: (m, n) = ("<<m<<", "<<n<<")"<<std::endl;

	if(m >= 0 && m < H && n >= 0 && n < W) {
		w = max(0, 1 - abs(x - m)) * max(0, 1 - abs(y - n));
		res += w * pic[m * W + n];
		if(debug) std::cout<<prefix<<"w = "<<w<<", pic[m, n] = "<<pic[m * W + n]<<std::endl;
	}

	m = floor(x) + 1; n = floor(y) + 1; w = 0;
	if(debug) std::cout<<prefix<<"4: (m, n) = ("<<m<<", "<<n<<")"<<std::endl;

	if(m >= 0 && m < H && n >= 0 && n < W) {
		w = max(0, 1 - abs(x - m)) * max(0, 1 - abs(y - n));
		res += w * pic[m * W + n];
		if(debug) std::cout<<prefix<<"w = "<<w<<", pic[m, n] = "<<pic[m * W + n]<<std::endl;
	}

	if(debug) std::cout<<prefix<<"Finished. \tres = "<<res<<std::endl;

	return res;
}

template <typename Dtype>
void SpatialTransformerLayer<Dtype>::Forward_cpu(const vector<Blob<Dtype>*>& bottom,
    const vector<Blob<Dtype>*>& top) {

	string prefix = "\t\tSpatial Transformer Layer:: Forward_cpu: \t";

	// CHECK(false) << "Don't use the CPU implementation! If you really want to, delete the" <<
	//		" CHECK in st_layer.cpp file. Line number: 240-241." << std::endl;

	if(global_debug) std::cout<<prefix<<"Starting!"<<std::endl;

	const Dtype* U = bottom[0]->cpu_data();
	const Dtype* theta = bottom[1]->cpu_data();
	const Dtype* output_grid_data = output_grid.cpu_data();

	Dtype* input_grid_data = input_grid.mutable_cpu_data();
	Dtype* V = top[0]->mutable_cpu_data();

	caffe_set(input_grid.count(), (Dtype)0, input_grid_data);
	caffe_set(top[0]->count(), (Dtype)0, V);

	// for each input
	for(int i = 0; i < N; ++i) {

		Dtype* coordinates = input_grid_data + (output_H_ * output_W_ * 2) * i;

		caffe_cpu_gemm<Dtype>(CblasNoTrans, CblasTrans, output_H_ * output_W_, 2, 3, (Dtype)1.,
		      output_grid_data, theta + 6 * i, (Dtype)0., coordinates);

		int row_idx; Dtype px, py;

		for(int j = 0; j < C; ++j)
			for(int s = 0; s < output_H_; ++s)
				for(int t = 0; t < output_W_; ++t) {

					row_idx = output_W_ * s + t;

					px = coordinates[row_idx * 2];
					py = coordinates[row_idx * 2 + 1];

					V[top[0]->offset(i, j, s, t)] = transform_forward_cpu(
							U + bottom[0]->offset(i, j, 0, 0), px, py);
				}
	}

	if(global_debug) std::cout<<prefix<<"Finished."<<std::endl;
}

template <typename Dtype>
void SpatialTransformerLayer<Dtype>::transform_backward_cpu(Dtype dV, const Dtype* U, const Dtype px,
		const Dtype py, Dtype* dU, Dtype& dpx, Dtype& dpy) {

	bool debug = false;

	string prefix = "\t\tSpatial Transformer Layer:: transform_backward_cpu: \t";

	if(debug) std::cout<<prefix<<"Starting!"<<std::endl;

	Dtype x = (px + 1) / 2 * H; Dtype y = (py + 1) / 2 * W;
	if(debug) std::cout<<prefix<<"(x, y) = ("<<x<<", "<<y<<")"<<std::endl;

	int m, n; Dtype w;

	m = floor(x); n = floor(y); w = 0;
	if(debug) std::cout<<prefix<<"(m, n) = ("<<m<<", "<<n<<")"<<std::endl;

	if(m >= 0 && m < H && n >= 0 && n < W) {
		w = max(0, 1 - abs(x - m)) * max(0, 1 - abs(y - n));

		dU[m * W + n] += w * dV;

		if(abs(x - m) < 1) {
			if(m >= x) {
				dpx += max(0, 1 - abs(y - n)) * U[m * W + n] * dV * H / 2;
				if(debug) std::cout<<prefix<<"dpx += "<<max(0, 1 - abs(y - n))<<" * "<<U[m * W + n]<<" * "<<dV<<" * "<<H / 2<<std::endl;
			} else {
				dpx -= max(0, 1 - abs(y - n)) * U[m * W + n] * dV * H / 2;
				if(debug) std::cout<<prefix<<"dpx -= "<<max(0, 1 - abs(y - n))<<" * "<<U[m * W + n]<<" * "<<dV<<" * "<<H / 2<<std::endl;
			}
		}

		if(abs(y - n) < 1) {
			if(n >= y) {
				dpy += max(0, 1 - abs(x - m)) * U[m * W + n] * dV * W / 2;
				if(debug) std::cout<<prefix<<"dpy += "<<max(0, 1 - abs(x - m))<<" * "<<U[m * W + n]<<" * "<<dV<<" * "<<W / 2<<std::endl;
			} else {
				dpy -= max(0, 1 - abs(x - m)) * U[m * W + n] * dV * W / 2;
				if(debug) std::cout<<prefix<<"dpy -= "<<max(0, 1 - abs(x - m))<<" * "<<U[m * W + n]<<" * "<<dV<<" * "<<W / 2<<std::endl;
			}
		}
	}

	m = floor(x) + 1; n = floor(y); w = 0;
	if(debug) std::cout<<prefix<<"(m, n) = ("<<m<<", "<<n<<")"<<std::endl;

	if(m >= 0 && m < H && n >= 0 && n < W) {
		w = max(0, 1 - abs(x - m)) * max(0, 1 - abs(y - n));

		dU[m * W + n] += w * dV;

		if(abs(x - m) < 1) {
			if(m >= x) {
				dpx += max(0, 1 - abs(y - n)) * U[m * W + n] * dV * H / 2;
				if(debug) std::cout<<prefix<<"dpx += "<<max(0, 1 - abs(y - n))<<" * "<<U[m * W + n]<<" * "<<dV<<" * "<<H / 2<<std::endl;
			} else {
				dpx -= max(0, 1 - abs(y - n)) * U[m * W + n] * dV * H / 2;
				if(debug) std::cout<<prefix<<"dpx -= "<<max(0, 1 - abs(y - n))<<" * "<<U[m * W + n]<<" * "<<dV<<" * "<<H / 2<<std::endl;
			}
		}

		if(abs(y - n) < 1) {
			if(n >= y) {
				dpy += max(0, 1 - abs(x - m)) * U[m * W + n] * dV * W / 2;
				if(debug) std::cout<<prefix<<"dpy += "<<max(0, 1 - abs(x - m))<<" * "<<U[m * W + n]<<" * "<<dV<<" * "<<W / 2<<std::endl;
			} else {
				dpy -= max(0, 1 - abs(x - m)) * U[m * W + n] * dV * W / 2;
				if(debug) std::cout<<prefix<<"dpy -= "<<max(0, 1 - abs(x - m))<<" * "<<U[m * W + n]<<" * "<<dV<<" * "<<W / 2<<std::endl;
			}
		}
	}

	m = floor(x); n = floor(y) + 1; w = 0;
	if(debug) std::cout<<prefix<<"(m, n) = ("<<m<<", "<<n<<")"<<std::endl;

	if(m >= 0 && m < H && n >= 0 && n < W) {
		w = max(0, 1 - abs(x - m)) * max(0, 1 - abs(y - n));

		dU[m * W + n] += w * dV;

		if(abs(x - m) < 1) {
			if(m >= x) {
				dpx += max(0, 1 - abs(y - n)) * U[m * W + n] * dV * H / 2;
				if(debug) std::cout<<prefix<<"dpx += "<<max(0, 1 - abs(y - n))<<" * "<<U[m * W + n]<<" * "<<dV<<" * "<<H / 2<<std::endl;
			} else {
				dpx -= max(0, 1 - abs(y - n)) * U[m * W + n] * dV * H / 2;
				if(debug) std::cout<<prefix<<"dpx -= "<<max(0, 1 - abs(y - n))<<" * "<<U[m * W + n]<<" * "<<dV<<" * "<<H / 2<<std::endl;
			}
		}

		if(abs(y - n) < 1) {
			if(n >= y) {
				dpy += max(0, 1 - abs(x - m)) * U[m * W + n] * dV * W / 2;
				if(debug) std::cout<<prefix<<"dpy += "<<max(0, 1 - abs(x - m))<<" * "<<U[m * W + n]<<" * "<<dV<<" * "<<W / 2<<std::endl;
			} else {
				dpy -= max(0, 1 - abs(x - m)) * U[m * W + n] * dV * W / 2;
				if(debug) std::cout<<prefix<<"dpy -= "<<max(0, 1 - abs(x - m))<<" * "<<U[m * W + n]<<" * "<<dV<<" * "<<W / 2<<std::endl;
			}
		}
	}

	m = floor(x) + 1; n = floor(y) + 1; w = 0;
	if(debug) std::cout<<prefix<<"(m, n) = ("<<m<<", "<<n<<")"<<std::endl;

	if(m >= 0 && m < H && n >= 0 && n < W) {
		w = max(0, 1 - abs(x - m)) * max(0, 1 - abs(y - n));

		dU[m * W + n] += w * dV;

		if(abs(x - m) < 1) {
			if(m >= x) {
				dpx += max(0, 1 - abs(y - n)) * U[m * W + n] * dV * H / 2;
				if(debug) std::cout<<prefix<<"dpx += "<<max(0, 1 - abs(y - n))<<" * "<<U[m * W + n]<<" * "<<dV<<" * "<<H / 2<<std::endl;
			} else {
				dpx -= max(0, 1 - abs(y - n)) * U[m * W + n] * dV * H / 2;
				if(debug) std::cout<<prefix<<"dpx -= "<<max(0, 1 - abs(y - n))<<" * "<<U[m * W + n]<<" * "<<dV<<" * "<<H / 2<<std::endl;
			}
		}

		if(abs(y - n) < 1) {
			if(n >= y) {
				dpy += max(0, 1 - abs(x - m)) * U[m * W + n] * dV * W / 2;
				if(debug) std::cout<<prefix<<"dpy += "<<max(0, 1 - abs(x - m))<<" * "<<U[m * W + n]<<" * "<<dV<<" * "<<W / 2<<std::endl;
			} else {
				dpy -= max(0, 1 - abs(x - m)) * U[m * W + n] * dV * W / 2;
				if(debug) std::cout<<prefix<<"dpy -= "<<max(0, 1 - abs(x - m))<<" * "<<U[m * W + n]<<" * "<<dV<<" * "<<W / 2<<std::endl;
			}
		}
	}

	if(debug) std::cout<<prefix<<"Finished."<<std::endl;
}

template <typename Dtype>
void SpatialTransformerLayer<Dtype>::Backward_cpu(const vector<Blob<Dtype>*>& top,
    const vector<bool>& propagate_down,
    const vector<Blob<Dtype>*>& bottom) {

		string prefix = "\t\tSpatial Transformer Layer:: Backward_cpu: \t";

		// CHECK(false) << "Don't use the CPU implementation! If you really want to, delete the" <<
		//		" CHECK in st_layer.cpp file. Line number: 420-421." << std::endl;

		if(global_debug) std::cout<<prefix<<"Starting!"<<std::endl;

		const Dtype* dV = top[0]->cpu_diff();
		const Dtype* input_grid_data = input_grid.cpu_data();
		const Dtype* U = bottom[0]->cpu_data();

		Dtype* dU = bottom[0]->mutable_cpu_diff();
		Dtype* dTheta = bottom[1]->mutable_cpu_diff();
		Dtype* input_grid_diff = input_grid.mutable_cpu_diff();

		caffe_set(bottom[0]->count(), (Dtype)0, dU);
		caffe_set(bottom[1]->count(), (Dtype)0, dTheta);
		caffe_set(input_grid.count(), (Dtype)0, input_grid_diff);

		for(int i = 0; i < N; ++i) {

			const Dtype* coordinates = input_grid_data + (output_H_ * output_W_ * 2) * i;
			Dtype* coordinates_diff = input_grid_diff + (output_H_ * output_W_ * 2) * i;

			int row_idx; Dtype px, py, dpx, dpy, delta_dpx, delta_dpy;

			for(int s = 0; s < output_H_; ++s)
				for(int t = 0; t < output_W_; ++t) {

					row_idx = output_W_ * s + t;

					px = coordinates[row_idx * 2];
					py = coordinates[row_idx * 2 + 1];

					for(int j = 0; j < C; ++j) {

						delta_dpx = delta_dpy = (Dtype)0.;

						transform_backward_cpu(dV[top[0]->offset(i, j, s, t)], U + bottom[0]->offset(i, j, 0, 0),
								px, py, dU + bottom[0]->offset(i, j, 0, 0), delta_dpx, delta_dpy);

						coordinates_diff[row_idx * 2] += delta_dpx;
						coordinates_diff[row_idx * 2 + 1] += delta_dpy;
					}

					dpx = coordinates_diff[row_idx * 2];
					dpy = coordinates_diff[row_idx * 2 + 1];

					dTheta[6 * i] += dpx * (s * 1.0 / output_H_ * 2 - 1);
					dTheta[6 * i + 1] += dpx * (t * 1.0 / output_W_ * 2 - 1);
					dTheta[6 * i + 2] += dpx;
					dTheta[6 * i + 3] += dpy * (s * 1.0 / output_H_ * 2 - 1);
					dTheta[6 * i + 4] += dpy * (t * 1.0 / output_W_ * 2 - 1);
					dTheta[6 * i + 5] += dpy;
				}
		}

		if(global_debug) std::cout<<prefix<<"Finished."<<std::endl;
}

#ifdef CPU_ONLY
STUB_GPU(SpatialTransformerLayer);
#endif

INSTANTIATE_CLASS(SpatialTransformerLayer);
REGISTER_LAYER_CLASS(SpatialTransformer);

}  // namespace caffe


#ifndef PCP_ALGORITHM_BILATERAL_FILTER_HPP
#define PCP_ALGORITHM_BILATERAL_FILTER_HPP

/**
 * @file
 * @ingroup algorithm
 */

#include "pcp/common/norm.hpp"
#include "pcp/common/vector3d.hpp"
#include "pcp/kdtree/linked_kdtree.hpp"
#include "pcp/traits/output_iterator_traits.hpp"
#include "pcp/traits/point_map.hpp"

#include <Eigen/Core>
#include <algorithm>
#include <array>
#include <execution>
#include <vector>

namespace pcp {
namespace algorithm {
namespace bilateral {

/**
 * @brief
 * Parameters for the bilateral filtering algorithm
 */
struct params_t
{
    double sigmaf = 1.;  ///< Standard deviation for the support region of spatial weight function f
    double sigmag = 0.1; ///< Standard deviation for the influence weight function g
    std::size_t K = 1u;  ///< Number of iterations of bilateral filtering
};

namespace detail {

template <
    class KdTree,
    class CoordinateMap,
    class PointMap,
    class NormalMap,
    class SpatialWeightFunction,
    class InfluenceWeightFunction,
    class ProjectionFunction,
    class ScalarType>
std::invoke_result_t<PointMap, std::size_t> compute_pi(
    std::size_t const i,
    ScalarType const sigmaf,
    ScalarType const sigmag,
    KdTree const& kdtree,
    CoordinateMap const& coordinate_map,
    PointMap const& point_map,
    NormalMap const& normal_map,
    SpatialWeightFunction const& f,
    InfluenceWeightFunction const& g,
    ProjectionFunction const& projection)
{
    using scalar_type = ScalarType;
    using point_type  = std::invoke_result_t<PointMap, std::size_t>;
    using normal_type = std::invoke_result_t<NormalMap, std::size_t>;

    scalar_type constexpr two  = scalar_type{2.};
    scalar_type constexpr zero = scalar_type{0.};

    auto const s  = point_map(i);
    auto const ci = coordinate_map(i);

    sphere_a<scalar_type> support_region{};
    support_region.position = ci;
    support_region.radius   = two * sigmaf;

    auto const neighbors = kdtree.range_search(support_region);

    scalar_type k = zero;
    point_type sprime{zero, zero, zero};
    for (auto it = neighbors.begin(); it != neighbors.end(); ++it)
    {
        point_type const p           = point_map(*it);
        normal_type const np         = normal_map(*it);
        point_type const s_projected = projection(p, np, s);

        scalar_type const rf = common::norm(s - p);
        scalar_type const rg = common::norm(s_projected - s);

        scalar_type const wf = f(sigmaf, rf);
        scalar_type const wg = g(sigmag, rg);

        scalar_type const w = wf * wg;
        k += w;
        common::basic_vector3d_t<scalar_type> translation{
            w * s_projected.x(),
            w * s_projected.y(),
            w * s_projected.z()};
        sprime = sprime + translation;
    }
    sprime = sprime / k;

    return sprime;
}

template <
    class KdTree,
    class CoordinateMap,
    class PointMap,
    class NormalMap,
    class SpatialWeightFunction,
    class SpatialWeightFunctionDerivative,
    class InfluenceWeightFunction,
    class InfluenceWeightFunctionDerivative,
    class ProjectionFunction,
    class ScalarType>
std::invoke_result_t<NormalMap, std::size_t> compute_ni(
    std::size_t const i,
    ScalarType const sigmaf,
    ScalarType const sigmag,
    KdTree const& kdtree,
    CoordinateMap const& coordinate_map,
    PointMap const& point_map,
    NormalMap const& normal_map,
    SpatialWeightFunction const& f,
    SpatialWeightFunctionDerivative const& df,
    InfluenceWeightFunction const& g,
    InfluenceWeightFunctionDerivative const& dg,
    ProjectionFunction const& projection)
{
    using scalar_type           = ScalarType;
    using point_type            = std::invoke_result_t<PointMap, std::size_t>;
    using normal_type           = std::invoke_result_t<NormalMap, std::size_t>;
    using matrix_3d_type        = Eigen::Matrix<scalar_type, 3, 3>;
    using column_vector_3d_type = Eigen::Matrix<scalar_type, 3, 1>;
    using row_vector_3d_type    = Eigen::Matrix<scalar_type, 1, 3>;

    scalar_type constexpr two  = scalar_type{2.};
    scalar_type constexpr zero = scalar_type{0.};

    point_type const pcp_s = point_map(i);
    column_vector_3d_type const s{pcp_s.x(), pcp_s.y(), pcp_s.z()};
    auto const ci = coordinate_map(i);

    sphere_a<scalar_type> support_region{};
    support_region.position = ci;
    support_region.radius   = two * sigmaf;

    auto const neighbors = kdtree.range_search(support_region);

    /**
     * Jacobian of sum of:
     *
     * projection(s) * f(||s - p||) * g(||projection(s) - s||)
     */
    matrix_3d_type J_pi_f_g;
    J_pi_f_g.setZero();
    /**
     * sum of:
     *
     * projection(s) * f(||s - p||) * g(||projection(s) - s||)
     */
    column_vector_3d_type pi_f_g;
    pi_f_g.setZero();
    /**
     * Gradient of k(s)
     */
    row_vector_3d_type grad_k;
    grad_k.setZero();
    /**
     * k(s)
     */
    scalar_type k = zero;

    for (auto it = neighbors.begin(); it != neighbors.end(); ++it)
    {
        point_type const pcp_p           = point_map(*it);
        normal_type const pcp_np         = normal_map(*it);
        point_type const pcp_s_projected = projection(pcp_p, pcp_np, pcp_s);

        column_vector_3d_type const p{pcp_p.x(), pcp_p.y(), pcp_p.z()};
        column_vector_3d_type const np{pcp_np.nx(), pcp_np.ny(), pcp_np.nz()};
        column_vector_3d_type const s_projected{
            pcp_s_projected.x(),
            pcp_s_projected.y(),
            pcp_s_projected.z()};

        column_vector_3d_type const sp  = s - p;
        column_vector_3d_type const sps = s_projected - s;
        scalar_type const rf            = sp.norm();
        scalar_type const rg            = sps.norm();

        // f(||s - p||)
        scalar_type const wf = f(sigmaf, rf);
        // g(||projection(s) - s||)
        scalar_type const wg = g(sigmag, rg);

        scalar_type const w = wf * wg;

        // k(s) = sum f(||s - p||) * g(||projection(s) - s||)
        k += w;

        // projection(s) * f(||s - p||) * g(||projection(s) - s||)
        column_vector_3d_type translation{
            w * pcp_s_projected.x(),
            w * pcp_s_projected.y(),
            w * pcp_s_projected.z()};
        pi_f_g += translation;

        // derivative df/dr | r=||s-p||
        scalar_type const wdf            = df(sigmaf, rf);
        row_vector_3d_type const sp_unit = sp.normalized().transpose();
        // grad(f) = (s - p) / ||s - p|| * (df/dr | r=||s-p||)
        row_vector_3d_type const grad_f = sp_unit * wdf;

        // Jacobian of projection(s)
        matrix_3d_type Jpi;
        Jpi(0, 0) = 1 - (np.x() * np.x());
        Jpi(1, 1) = 1 - (np.y() * np.y());
        Jpi(2, 2) = 1 - (np.z() * np.z());
        Jpi(0, 1) = np.x() * np.y();
        Jpi(0, 2) = np.x() * np.z();
        Jpi(1, 2) = np.y() * np.z();
        Jpi(1, 0) = Jpi(0, 1);
        Jpi(2, 0) = Jpi(0, 2);
        Jpi(2, 1) = Jpi(1, 2);

        // derivative dg/dr | r = ||projection(s) - s||
        scalar_type const wdg             = dg(sigmag, rg);
        row_vector_3d_type const sps_unit = sps.normalized().transpose();
        /**
         * Let sps_unit = (projection(s) - s) / ||projection(s) - s||
         *
         * grad(g) =
         * (sps_unit * Jacobian(projection(s)) - sps_unit) *
         * (dg/dr | r = ||projection(s) - s||)
         */
        row_vector_3d_type const grad_g = (sps_unit * Jpi - sps_unit) * wdg;

        /**
         * Product rule grad(f(||s - p||) * g(||projection(s) - s||))
         */
        grad_k += (grad_f * wg) + (wf * grad_g);
        /**
         * Product rule grad(projection(s) * f(||s - p||) * g(||projection(s) - s||))
         */
        J_pi_f_g += (Jpi * wf * wg) + (sps * grad_f * wg) + (sps * wf * grad_g);
    }

    /**
     * Quotient rule grad(u/v) = (1/v^2) * (grad(u)*v - u*grad(v))
     */
    scalar_type const ks2_inv = scalar_type{1.} / (k * k);
    matrix_3d_type const J    = ks2_inv * (J_pi_f_g * k - pi_f_g * grad_k);

    normal_type const pcp_ns = normal_map(i);
    column_vector_3d_type const ns{pcp_ns.nx(), pcp_ns.ny(), pcp_ns.nz()};

    /**
     * In the original normal improvement paper, they use the inverse transpose
     * of the jacobian. In our case, we directly use the jacobian
     * of the filter F(s), because it is in line with the intuition
     * of using the local spatial deformation of the field F(s) to
     * adjust normals.
     */
    // ns' = J^(-T) * ns
    // matrix_3d_type const adj       = J.adjoint();
    column_vector_3d_type ns_prime = J * ns;
    ns_prime.normalize();
    return normal_type{ns_prime.x(), ns_prime.y(), ns_prime.z()};
}

} // namespace detail
} // namespace bilateral

/**
 * @ingroup smoothing-algorithm
 * @brief
 * Uses a bilateral filter to smooth an input point cloud.
 *
 * The bilateral filter is defined in the same way as in
 * 'Jones, Thouis R., Fredo Durand, and Matthias Zwicker. "Normal improvement for point rendering."
 * IEEE Computer Graphics and Applications 24.4 (2004): 53-56.'
 *
 * Internally, the filtering needs a temporary copy of the points p(k)
 * to compute p(k+1), the points at the next iteration. At each iteration,
 * a kd-tree must be built over the points p(k) to support optimized
 * range search queries.
 *
 * The transformation F (the bilateral filter) over our points P is done in parallel at each
 * iteration k.
 *
 * @tparam RandomAccessIter Iterator type satisfying Random Access requirements
 * @tparam OutputIter Iterator type dereferenceable to a type satisfying Point concept
 * @tparam PointMap Type satisfying PointMap concept
 * @tparam NormalMap Type satisfying NormalMap concept
 * @param begin Start iterator of input point cloud
 * @param end End iterator of input point cloud
 * @param out_begin Start iterator of output points
 * @param point_map The point map property map
 * @param normal_map The normal map property map
 * @param params The bilateral filter algorithm's parameters
 * @return End iterator of output sequence
 */
template <class RandomAccessIter, class OutputIter, class PointMap, class NormalMap>
OutputIter bilateral_filter_points(
    RandomAccessIter begin,
    RandomAccessIter end,
    OutputIter out_begin,
    PointMap const& point_map,
    NormalMap const& normal_map,
    bilateral::params_t const& params)
{
    using input_element_type = typename std::iterator_traits<RandomAccessIter>::value_type;
    using input_point_type   = std::invoke_result_t<PointMap, input_element_type>;
    using input_normal_type  = std::invoke_result_t<NormalMap, input_element_type>;
    using scalar_type        = typename input_point_type::coordinate_type;
    using output_point_type  = typename xstd::output_iterator_traits<OutputIter>::value_type;
    // using difference_type    = typename std::iterator_traits<RandomAccessIter>::difference_type;

    static_assert(
        traits::is_point_map_v<PointMap, input_element_type>,
        "point_map must satisfy PointMap concept");

    static_assert(
        traits::is_point_v<output_point_type>,
        "OutputIter must be dereferenceable to a type satisfying Point concept");

    std::size_t const N        = static_cast<std::size_t>(std::distance(begin, end));
    scalar_type const sigmaf   = static_cast<scalar_type>(params.sigmaf);
    scalar_type const sigmag   = static_cast<scalar_type>(params.sigmag);
    std::size_t const K        = params.K;
    scalar_type constexpr pi   = static_cast<scalar_type>(3.14159265358979323846);
    scalar_type constexpr zero = scalar_type{0.};

    assert(K > 0u);
    assert(N > 0u);
    assert(sigmaf > zero);
    assert(sigmag > zero);

    std::vector<std::size_t> indices(N);
    std::iota(indices.begin(), indices.end(), 0u);

    std::vector<input_point_type> points(N);
    std::transform(begin, end, points.begin(), [&](input_element_type const& e) {
        return point_map(e);
    });

    std::vector<input_point_type> temporary_points(N);

    std::vector<input_normal_type> normals(N);
    std::transform(begin, end, normals.begin(), [&](input_element_type const& e) {
        return normal_map(e);
    });

    auto const internal_point_map = [&](std::size_t const i) {
        return points[i];
    };
    auto const internal_normal_map = [&](std::size_t const i) {
        return normals[i];
    };

    auto const gaussian = [=](scalar_type const sigma, scalar_type const r) -> scalar_type {
        scalar_type const s2      = sigma * sigma;
        scalar_type const r2      = r * r;
        scalar_type const power   = -r2 / (2 * s2);
        scalar_type constexpr one = scalar_type{1.};
        scalar_type constexpr two = scalar_type{2.};
        scalar_type const coeff   = one / (sigma * std::sqrt(two * pi));
        return coeff * std::exp(power);
    };

    auto const projection = [&](input_point_type const& p,
                                input_normal_type const& np,
                                input_point_type const& s) -> input_point_type {
        auto const sp = p - s;
        common::basic_vector3d_t<scalar_type> const n{np.nx(), np.ny(), np.nz()};
        auto const d = common::inner_product(sp, n);
        return s + d * n;
    };

    using coordinates_type = std::array<scalar_type, 3u>;

    auto const coordinate_map = [&](std::size_t const pe) {
        auto const p = internal_point_map(pe);
        return coordinates_type{p.x(), p.y(), p.z()};
    };

    kdtree::construction_params_t kdtree_params;
    kdtree_params.compute_max_depth     = true;
    kdtree_params.construction          = kdtree::construction_t::nth_element;
    kdtree_params.max_elements_per_leaf = 64u;

    for (std::size_t k = 0u; k < K; ++k)
    {
        basic_linked_kdtree_t<input_element_type, 3u, decltype(coordinate_map)> kdtree{
            indices.begin(),
            indices.end(),
            coordinate_map,
            kdtree_params};

        std::transform(
            std::execution::par,
            indices.begin(),
            indices.end(),
            temporary_points.begin(),
            [&](std::size_t const i) {
                return bilateral::detail::compute_pi(
                    i,
                    sigmaf,
                    sigmag,
                    kdtree,
                    coordinate_map,
                    internal_point_map,
                    internal_normal_map,
                    gaussian,
                    gaussian,
                    projection);
            });

        std::copy(temporary_points.begin(), temporary_points.end(), points.begin());
    }

    return std::copy(points.begin(), points.end(), out_begin);
}

/**
 * @ingroup smoothing-algorithm
 * @brief
 * Deforms the normal field of an input point cloud using the local deformation field of the
 * 3d bilateral filter (in other words, its Jacobian at a point p).
 *
 * The bilateral filter is defined in the same way as in
 * 'Jones, Thouis R., Fredo Durand, and Matthias Zwicker. "Normal improvement for point rendering."
 * IEEE Computer Graphics and Applications 24.4 (2004): 53-56.'
 *
 * This method should be used only for point rendering. It does not smooth input normals as one
 * would expect for surface reconstruction. For normal smoothing, one should rather look at
 * techniques such as EAR (edge aware resampling):
 *
 * 'Huang, Hui, et al. "Edge-aware point set resampling." ACM transactions on graphics (TOG) 32.1
 * (2013): 1-12.'
 *
 * @tparam RandomAccessIter Iterator type satisfying Random Access requirements
 * @tparam OutputIter Iterator type dereferenceable to a type satisfying Normal concept
 * @tparam PointMap Type satisfying PointMap concept
 * @tparam NormalMap Type satisfying NormalMap concept
 * @param begin Start iterator of input point cloud
 * @param end End iterator of input point cloud
 * @param out_begin Start iterator of output points
 * @param point_map The point map property map
 * @param normal_map The normal map property map
 * @param params The bilateral filter algorithm's parameters
 * @return End iterator of output sequence
 */
template <class RandomAccessIter, class OutputIter, class PointMap, class NormalMap>
OutputIter bilateral_filter_normals(
    RandomAccessIter begin,
    RandomAccessIter end,
    OutputIter out_begin,
    PointMap const& point_map,
    NormalMap const& normal_map,
    bilateral::params_t const& params)
{
    using input_element_type = typename std::iterator_traits<RandomAccessIter>::value_type;
    using input_point_type   = std::invoke_result_t<PointMap, input_element_type>;
    using input_normal_type  = std::invoke_result_t<NormalMap, input_element_type>;
    using scalar_type        = typename input_point_type::coordinate_type;
    using output_normal_type = typename xstd::output_iterator_traits<OutputIter>::value_type;
    using difference_type    = typename std::iterator_traits<RandomAccessIter>::difference_type;

    static_assert(
        traits::is_point_map_v<PointMap, input_element_type>,
        "point_map must satisfy PointMap concept");

    std::size_t const N        = static_cast<std::size_t>(std::distance(begin, end));
    scalar_type const sigmaf   = static_cast<scalar_type>(params.sigmaf);
    scalar_type const sigmag   = static_cast<scalar_type>(params.sigmag);
    std::size_t const K        = params.K;
    scalar_type constexpr pi   = static_cast<scalar_type>(3.14159265358979323846);
    scalar_type constexpr zero = scalar_type{0.};

    assert(K > 0u);
    assert(N > 0u);
    assert(sigmaf > zero);
    assert(sigmag > zero);

    std::vector<std::size_t> indices(N);
    std::iota(indices.begin(), indices.end(), 0u);

    std::vector<input_normal_type> normals(N);
    std::transform(begin, end, normals.begin(), [&](input_element_type const& e) {
        return normal_map(e);
    });

    std::vector<input_normal_type> temporary_normals(N);

    auto const internal_point_map = [&](std::size_t const i) -> input_point_type {
        return point_map(*std::next(begin, static_cast<difference_type>(i)));
    };
    auto const internal_normal_map = [&](std::size_t const i) -> input_normal_type {
        return normals[i];
    };

    auto const gaussian = [=](scalar_type const sigma, scalar_type const r) {
        scalar_type const s2      = sigma * sigma;
        scalar_type const r2      = r * r;
        scalar_type const power   = -r2 / (2 * s2);
        scalar_type constexpr one = scalar_type{1.};
        scalar_type constexpr two = scalar_type{2.};
        scalar_type const coeff   = one / (sigma * std::sqrt(two * pi));
        return coeff * std::exp(power);
    };

    auto const dgaussian = [=](scalar_type const sigma, scalar_type const r) {
        scalar_type const s2      = sigma * sigma;
        scalar_type const s3      = sigma * s2;
        scalar_type const r2      = r * r;
        scalar_type const power   = -r2 / (2 * s2);
        scalar_type constexpr two = scalar_type{2.};
        scalar_type const coeff   = -r / (s3 * std::sqrt(two * pi));
        return coeff * std::exp(power);
    };

    auto const projection = [&](input_point_type const& p,
                                input_normal_type const& np,
                                input_point_type const& s) -> input_point_type {
        auto const sp = p - s;
        common::basic_vector3d_t<scalar_type> const n{np.nx(), np.ny(), np.nz()};
        auto const d = common::inner_product(sp, n);
        return s + d * n;
    };

    using coordinates_type = std::array<scalar_type, 3u>;

    auto const coordinate_map = [&](std::size_t const i) {
        auto const p = internal_point_map(i);
        return coordinates_type{p.x(), p.y(), p.z()};
    };

    kdtree::construction_params_t kdtree_params;
    kdtree_params.compute_max_depth     = true;
    kdtree_params.construction          = kdtree::construction_t::nth_element;
    kdtree_params.max_elements_per_leaf = 64u;

    basic_linked_kdtree_t<input_element_type, 3u, decltype(coordinate_map)> kdtree{
        indices.begin(),
        indices.end(),
        coordinate_map,
        kdtree_params};

    for (std::size_t k = 0u; k < K; ++k)
    {
        std::transform(
            std::execution::par,
            indices.begin(),
            indices.end(),
            temporary_normals.begin(),
            [&](std::size_t const i) {
                return bilateral::detail::compute_ni(
                    i,
                    sigmaf,
                    sigmag,
                    kdtree,
                    coordinate_map,
                    internal_point_map,
                    internal_normal_map,
                    gaussian,
                    dgaussian,
                    gaussian,
                    dgaussian,
                    projection);
            });

        std::copy(temporary_normals.begin(), temporary_normals.end(), normals.begin());
    }

    return std::copy(normals.begin(), normals.end(), out_begin);
}

} // namespace algorithm
} // namespace pcp

#endif // PCP_ALGORITHM_BILATERAL_FILTER_HPP


#ifndef ENKF_HPP
#define ENKF_HPP

#include <Eigen/Core>
#include <Eigen/Dense>
#include <Eigen/Cholesky>
#include <vector>
#include <functional>
#include <utility>
#include <cassert>
#include <random>
#include <cmath>
#include <iostream>

namespace fastmath
{
  using Vec = Eigen::VectorXd;
  using Mat = Eigen::MatrixXd;

  template<class State_t, class Obs_t, class ObsOp_t, class... ForecastArgs>
  class EnKF
  {
  public:

    constexpr EnKF() {};

    //filter with one observation
    /* virtual void filter(const Obs_t& y);*/

    virtual Obs_t observe(const State_t& x);

    virtual State_t forecast(State_t& x, ForecastArgs&... args);

    virtual ~EnKF() {};
  };

  //scalar EnKF with linear observations
  template<class... ForecastArgs>
  class VectorEnKF : public EnKF<Vec, Vec, Mat, ForecastArgs...>
  {

 private:

    Mat  m_H;

    Vec m_state;

    Vec m_obs = m_H * m_state;

    std::function<Vec(Vec&, ForecastArgs&...)> m_forecast;

    Vec m_ensemble_mean;

    Mat m_ensemble_covariance;

    Mat m_obs_covariance;

    int m_ensemble_size;

    int m_dimension;

    Mat m_ensemble = Mat::Zero(m_dimension, m_ensemble_size);

    bool m_initialized_ensemble = false;

    Mat m_background_covariance = Mat::Identity(m_dimension, m_dimension);

    Mat m_ensemble_transform = gen_covariance_transform(m_ensemble_covariance);
  public:

    VectorEnKF(const std::function<Vec(Vec&, ForecastArgs&...)>& n_forecast,
			 const Mat& n_H,
			 const Vec& n_initial_state,
			 const Mat& n_observe_err_covariance,
			 const Vec& n_initial_ensemble_mean,
			 const Mat&   n_initial_ensemble_covariance,
			 const int    n_ensemble_size
			 ) : EnKF<Vec, Vec, Mat, ForecastArgs...>(),
			     m_H(n_H),
			     m_state(n_initial_state),
			     m_forecast(n_forecast),
			     m_ensemble_mean(n_initial_ensemble_mean),
			     m_ensemble_covariance(n_initial_ensemble_covariance),
			     m_obs_covariance(n_observe_err_covariance),
			     m_ensemble_size(n_ensemble_size),
			     m_dimension(static_cast<int>(n_initial_state.size()))

    {};

    void set_state(const Vec& n_state) noexcept
    {
      m_state = n_state;
    }

    void set_ensemble_size(const int n_ensemble_size) noexcept
    {
      assert(n_ensemble_size > 0);
      m_ensemble_size = n_ensemble_size;
    }

    Vec state()
    {
      return m_state;
    }

    int ensemble_size()
    {
      return m_ensemble_size;
    }

    Vec ensemble_mean()
    {
      return m_ensemble_mean;
    }

    Mat ensemble_covariance()
    {
      return m_ensemble_covariance;
    }

    Vec observe(const Vec& x){
      return m_H * x;
    }

    Vec forecast(Vec& x, ForecastArgs&... args){
      return m_forecast(x, args...);
    }

  private:

    Mat gen_covariance_transform(const Mat& target_covariance) const
    {
      return target_covariance.ldlt().matrixL();
    }

    Mat gen_obs_perturbations()
    {
      auto rcovtransform = gen_covariance_transform(m_obs_covariance);

      std::random_device rd{};
      std::mt19937 gen{rd()};
      std::normal_distribution<> dis{0.0, 1.0};

      Mat A(m_H.rows(), m_ensemble_size);
      Vec epsilon(m_H.rows());

      for(auto i = 0; i < m_ensemble_size; ++i){
	for(auto e = 0; e < m_H.rows(); ++e){
	  epsilon[e] = dis(gen);
	}
	A.col(i) = epsilon;
      }

      return rcovtransform * A;
    }

    Mat gen_ensemble_perturbations(bool gen_cov_transform=true)
    {
      std::random_device rd{};
      std::mt19937 gen{rd()};
      std::normal_distribution<> dis{0.0, 1.0};

      Mat A(m_dimension,m_ensemble_size);
      Vec epsilon(m_dimension);

      if(gen_cov_transform){
	m_ensemble_transform = gen_covariance_transform(m_ensemble_covariance);
      }

      for(auto i = 0; i < m_ensemble_size; ++i){
	for(auto e = 0; e < m_dimension; ++e){
	  epsilon[e] = dis(gen);
	}
	A.col(i) = epsilon + m_ensemble_mean;
      }

      return  m_ensemble_transform * A;// * A;
    }

  /*std::pair<Vec,Mat> sample_mean_covariance(const std::vector<Vec>& sample)
    {
      Vec mu(sample[0].size());

      for(auto i = 0; i < m_ensemble_size; ++i){
	for(auto i = 0; i < sample[0].size(); ++i){
	  mu[i] += s[i];
	}
      }

      mu /= sample.size();

      Mat cov = Mat::Zero(sample.size(), sample.size());

      for(const auto& s : sample){
	auto delta = s - mu;
	cov.noalias() += delta * delta.transpose();
      }

      cov /= (sample.size() - 1);

      return std::make_pair(mu, cov);
      }*/

    //A is the result of a call to form_A
    Mat kalman_gain_A(const Mat& A)
    {
      auto V = m_H * A;

      auto vvr = V * V.transpose();// + m_obs_covariance;
      return A * V.transpose() * vvr.inverse();
    }

    //copy ensemble here
    Mat form_A(Mat ensemble)
    {
      for(auto i = 0; i < m_ensemble_size; ++i){
	ensemble.col(i) -= m_ensemble_mean;
      }

      return ensemble / std::sqrt(m_ensemble_size - 1);
    }

    Mat kalman_gain_ensemble(const Mat& ensemble)
    {
      auto A = form_A(ensemble);
      return kalman_gain_A(A);
    }

    //B matrix given an ensemble
    Mat ensemble_prior_covariance(const Mat& A)
    {
      return A * A.transpose();
    }

    Vec ensemble_member_update(const Mat& K, const Vec& perturbed_y, const Vec& xi)
    {
      auto err = perturbed_y - m_H * xi;

      return xi + K * err;
    }

    Mat ensemble_update(const Mat& K, const Mat& perturbed_yvals, const Mat& ensemble_xvals)
    {
      Mat updated_xvals(ensemble_xvals.rows(), ensemble_xvals.cols());
      for(auto i = 0; i < m_ensemble_size; ++i){
	updated_xvals.col(i) = ensemble_member_update(K, perturbed_yvals.col(i), ensemble_xvals.col(i));
      }
      return updated_xvals;
    }

    Mat ETKF_ensemble_update_mat(Mat& A)
    {
      auto V = m_H * A;
      auto Id = Mat::Identity(m_ensemble_size);
      auto emat = Id + V.transpose() * m_obs_covariance.inverse() * V;

      Eigen::SelfAdjointEigenSolver<Mat> esolver(emat);

      assert(esolver.info() == Eigen::Success);//, "ETKF eigendecomposition failed.");

      Mat Gamma = esolver.eigenvalues().asDiagonal();/* I + Gamma in the notes*/
      auto Q = esolver.eigenvectors();

      for(auto i = 0; i < Gamma.cols(); ++i){
	Gamma(i,i) = std::sqrt(Gamma(i,i));
      }

      return Q * Gamma * Q.transpose();
    }

    Mat ETKF_posterior_ensemble(Mat& A)
    {
      // A is prior ensemble
      auto X = ETKF_ensemble_update_mat(A);
      auto one = Vec::Ones(m_ensemble_size);
      auto aplus = A * one / m_ensemble_size;
      return A - aplus * one.transpose();
    }

    Mat ensemble_posterior_covariance(const Mat& B, const Mat& K)
    {
      return B - K * m_H * B;
    }

  public:

    void ETKF_filter(Vec& obs, ForecastArgs&... args)
    {
      for(auto i = 0; i < m_ensemble_size; ++i){
	  m_ensemble.col(i) = m_state;
	}

      m_ensemble += gen_ensemble_perturbations();

      m_ensemble = ETKF_posterior_ensemble(m_ensemble);

      m_state = m_ensemble.rowwise().mean();

      m_ensemble_mean = m_state;

      m_ensemble_covariance = m_ensemble * m_ensemble.transpose();
    }

    void filter(Vec& obs, ForecastArgs&... args)
    {
      // generate ensemble
      if(!m_initialized_ensemble){
	for(auto i = 0; i < m_ensemble_size; ++i){
	  m_ensemble.col(i) = m_state;
	}
	//m_initialized_ensemble = true;
      }
      m_ensemble += gen_ensemble_perturbations();
      /* forecast ensemble */
      for(auto i = 0; i < m_ensemble_size; ++i){
	Vec xstate = m_ensemble.col(i);
	m_ensemble.col(i) = m_forecast(xstate, args...);
      }

      m_ensemble_mean = m_ensemble.rowwise().mean();

      auto A = form_A(m_ensemble);
      m_background_covariance = A * A.transpose();

      Mat obs_perturbed = gen_obs_perturbations();

      for(auto i = 0; i < m_ensemble_size; ++i){
	obs_perturbed.col(i) += obs;
      }
      auto K = kalman_gain_A(A);
      A = ensemble_update(K, obs_perturbed / std::sqrt(m_ensemble_size-1), A);

      m_ensemble_covariance = A * A.transpose();

      /* get ensemble back */
      A *= std::sqrt(m_ensemble_size - 1);
      //add prior mean back and re-scale
      m_ensemble_mean = std::sqrt(m_ensemble_size - 1) * (A.rowwise().mean() + m_ensemble_mean);
      m_state = m_ensemble_mean;
    }

  };

}//namespace fastmath
#endif


#include <pybindings.h>
#include <container_pybindings.h>
#include <maps/pointing.h>
#include <G3Map.h>
#include <G3Units.h>

#include <boost/math/constants/constants.hpp>

#include <vector>
#include <math.h>
#include <iostream>
#include <iomanip>

#include <stdlib.h>
#include <time.h>

using namespace boost::math;

const double PI = constants::pi<double>();

#define ASIN asin
#define ATAN2 atan2

//#define CHECK_QUAT_INVERSE

/*
 * Quaternions cannot represent parity flips.  Since celestial coordinates
 * and az-el coordinates by construction have a different parity, we can't use
 * the general alpha delta angle to x-y-z mapping for one of the coordiate
 * systems.
 *
 * For the Euclidean quaternion representation at the pole,
 * the z coordinate = -sin(elevation) = sin(declination)
 */

static quat
project_on_plane(quat plane_normal, quat point)
{
	// Projects the quaternion onto a plane with unit normal plane_normal
	//   The plane is defined as going through the origin
	//   with normal = plane_normal

	quat out_q(point);
	//ensure unit vec
	plane_normal /= sqrt(dot3(plane_normal,plane_normal));
	out_q -= plane_normal * dot3(plane_normal, point);
	return out_q;
}

static bool
sloppy_eq(quat a, quat b, double slop = 1e-6)
{
	// Fuzzy quaternion equality comparison
	return ((fabs(a.R_component_1() - (b.R_component_1())) < slop ) &&
		(fabs(a.R_component_2() - (b.R_component_2())) < slop ) &&
		(fabs(a.R_component_3() - (b.R_component_3())) < slop ) &&
		(fabs(a.R_component_4() - (b.R_component_4())) < slop ));
}

quat
ang_to_quat(double alpha, double delta)
{
	double c_delta = cos(delta / G3Units::rad);
	return quat(0,
		    c_delta * cos(alpha/G3Units::rad),
		    c_delta * sin(alpha/G3Units::rad),
		    sin(delta / G3Units::rad));
}

void
quat_to_ang(quat q, double &alpha, double &delta)
{
	double d = dot3(q,q);
	if (fabs(d - 1.0) > 1e-6){
		q /= sqrt(d);
	}
	delta = ASIN(q.R_component_4()) * G3Units::rad;
	alpha = ATAN2(q.R_component_3(), q.R_component_2())*G3Units::rad;
}

static boost::python::tuple
py_quat_to_ang(quat q)
{
	double a,d;
	quat_to_ang(q, a, d);

	return boost::python::make_tuple(a, d);
}

quat
coord_quat_to_delta_hat(quat q)
{
	// computes the delta hat vector for a given point on the unit sphere
	// specified by q
	//
	// (The delta hat is equal to -alpha hat)

	q /= sqrt(dot3(q,q));
	double st = sqrt(1 - (q.R_component_4()*q.R_component_4()));
	quat u= quat(0,
		     -1 * (q.R_component_2() * q.R_component_4())/st,
		     -1 * (q.R_component_3() * q.R_component_4())/st,
		     st);
	u /= sqrt(dot3(u,u));
	return u;
}

double
get_rot_ang(quat start_q, quat end_q, quat trans)
{
	// delta is the physicist spherical coordinates delta
	// Computes delta hat for the start q applies trans to it
	// and then computes the angle between that and end_q's delta hat.

	quat t = trans * coord_quat_to_delta_hat(start_q) / trans;
	quat t_p = coord_quat_to_delta_hat(end_q);

	t /= sqrt(dot3(t,t));
	t_p /= sqrt(dot3(t_p,t_p));
	double d  = dot3(t,t_p);
	double sf = (dot3(end_q, cross3(t, t_p)) < 0) ? -1 : 1;
	if (d > 1) {
		g3_assert(d < 1.01);
		return 0;
	} else if (d < -1) {
		g3_assert(d > -1.01);
		return PI * G3Units::rad;
	} else {
		return sf * acos(d) * G3Units::rad;
	}
}

quat
get_transform_quat(double as_0, double ds_0, double ae_0, double de_0,
    double as_1, double ds_1, double ae_1, double de_1)
{
	/*
	 * as = alpha start
	 * ds = delta start
	 * ae = alpha end
	 * de = delta end
	 *
	 * The numeral postscripts are for which set of points.
	 *
	 * Computes a rotation that will take: (as_0,ds_0) to (ae_0, de_0) and
	 * (as_1, ds_1) to (ae_1, de_1)
	 *
	 */

	quat asds_0 = ang_to_quat(as_0, ds_0);
	quat asds_1 = ang_to_quat(as_1, ds_1);
	quat aede_0 = ang_to_quat(ae_0, de_0);
	quat aede_1 = ang_to_quat(ae_1, de_1);

	quat tquat = cross3(asds_0, aede_0);
	double mag = sqrt(dot3(tquat, tquat));
	double ang = acos(dot3(asds_0, aede_0));
	tquat *= sin(ang/2.0) / mag;
	tquat += quat(cos(ang/2.0),0,0,0);

	// trans_asds_1 and aede_1 should now be the same up to a rotation
	// around aede_0
	quat trans_asds_1 = tquat * asds_1 / tquat;

	// Project them on to a plane and find the angle between the two vectors
	// using (ae_0, de_0) as the normal since we are rotating around that
	// vector.
	quat p_asds1 = project_on_plane(aede_0, trans_asds_1);
	quat p_aede1 = project_on_plane(aede_0, aede_1);
	p_asds1 /= sqrt(dot3(p_asds1,p_asds1));
	p_aede1 /= sqrt(dot3(p_aede1,p_aede1));

	double rot_ang = acos(dot3(p_asds1, p_aede1));
	double sf = (dot3(aede_0, cross3(p_asds1, p_aede1)) < 0) ? -1 : 1;
	rot_ang *= sf;

	double sin_rot_ang_ov_2 = sin(rot_ang/2.0);
	quat rot_quat = quat(cos(rot_ang/2.0),
			     sin_rot_ang_ov_2 * aede_0.R_component_2(),
			     sin_rot_ang_ov_2 * aede_0.R_component_3(),
			     sin_rot_ang_ov_2 * aede_0.R_component_4());
	quat final_trans = rot_quat * tquat;

	return final_trans;
}

static std::vector<double>
test_trans(double az_0, double el_0, double ra_0, double dec_0,
    double az_1, double el_1, double ra_1, double dec_1,
    double az_t, double el_t)
{
	// computes the transform from the first 4 variables
	// returns that transform to the last 2 variables

	double ra_t, dec_t;
	quat q = get_transform_quat(az_0, -el_0,
				    ra_0, dec_0,
				    az_1, -el_1,
				    ra_1, dec_1);
	quat azel = ang_to_quat(az_t, -el_t);
	quat rad = q * azel / q;
	quat_to_ang(rad, ra_t, dec_t);
	std::vector<double> r(2,0);
	r[0] = ra_t;
	r[1] = dec_t;
	return r;
}

static std::vector<double>
test_gal_trans(double az_0, double el_0, double ra_0, double dec_0,
    double az_1, double el_1, double ra_1, double dec_1,
    double az_t, double el_t)
{
	double l_t, b_t;
	quat q = get_fk5_j2000_to_gal_quat() *get_transform_quat(az_0, -el_0,
								 ra_0, dec_0,
								 az_1, -el_1,
								 ra_1, dec_1);
	quat azel = ang_to_quat(az_t, -el_t);
	quat rad = q * azel / q;
	quat_to_ang(rad, l_t, b_t);
	std::vector<double> r(2,0);
	r[0] = l_t;
	r[1] = b_t;
	return r;
}

static double
test_gal_trans_rot(double ra, double dec)
{
	quat start = ang_to_quat(ra,dec);
	quat trans = get_fk5_j2000_to_gal_quat();

	return get_rot_ang(start, trans*start/trans, trans);
}

quat
offsets_to_quat(double x_offset, double y_offset)
{
	// Rotates the point (1,0,0) by the rotation matrix for the y_offset
	// and then the rotation matrix for the x_offset
	// quat t = (quat(cos(x_offset/(2.0*G3Units::rad)),0,0,sin(x_offset/(2.0*G3Units::rad))) *
	// 	  quat(cos(y_offset/(2.0*G3Units::rad)),0,sin(y_offset/(2.0*G3Units::rad)),0));
	// return t*quat(0,1.0,0,0)/t;
	// The above is exactly equal to:
	return ang_to_quat(x_offset, -y_offset);
}

quat
get_origin_rotator(double alpha, double delta)
{
	// Rotates the point (1,0,0) to the point specified by alpha and
	// delta via a rotation about the y axis and then the z axis
        return (quat(cos(alpha/2.0), 0, 0, sin(alpha/2.0)) *
                quat(cos(delta/2.0), 0, -sin(delta/2.0), 0));
}

static void
print_fk5_j2000_to_gal_quat()
{
	// uhh, so, this code was a super lazy way to get the quaternion
	// that takes fk5 j2000 to galactic j2000
	std::cout << std::setprecision(10) << std::endl;
	std::cout << get_transform_quat(
		0,0, 1.6814025470759737, -1.050488399695429,
		0,-0.7853981633974483, 5.750520098164818, -1.2109809382060603)
	          << std::endl;
}

quat
get_fk5_j2000_to_gal_quat()
{
	// returns the quaternion that rotates fk5j2000 to galactic J2000
	// coordinates
	return quat(0.4889475076,-0.483210684,0.1962537583,0.699229742);
}

static void
create_det_az_el_trans(const G3Timestream &az, const G3Timestream &el,
    G3VectorQuat &trans_quats) // XXX: switch to G3TimestreamQuat?
{
	// Creates the transform that takes (1,0,0) to az, -el
	// for why it's -el see the comment at the top of this document

	g3_assert(az.size() == el.size());
	trans_quats = G3VectorQuat(az.size(), quat(1,0,0,0));
	for (size_t i = 0; i < az.size(); i++)
		trans_quats[i] = get_origin_rotator(az[i], -el[i]);
}

static void
create_lazy_det_ra_dec_trans(const G3Timestream &ra, const G3Timestream &dec,
    G3VectorQuat &trans_quats)
{
	// Creates the transform that takes (1,0,0) to ra,dec
	g3_assert(ra.size() == dec.size());
	trans_quats = G3VectorQuat(ra.size(), quat(1,0,0,0));
	for (size_t i = 0; i < ra.size(); i++)
		trans_quats[i] = get_origin_rotator(ra[i], dec[i]);
}

static void
create_det_ra_dec_trans(const G3Timestream &az_0, const G3Timestream &el_0,
     const G3Timestream &ra_0, const G3Timestream &dec_0,
     const G3Timestream &az_1, const G3Timestream &el_1,
     const G3Timestream &ra_1, const G3Timestream &dec_1,
     G3VectorQuat & trans_quats)
{
	// Computes the transform that takes (1,0,0) to the point (ra_0, dec_0)
	// and properly handles rotation about the (ra_0, dec_0) point with the
	// inclusion of the second set of points.
	//
	// Stores the output in trans_quats.

	g3_assert(az_0.size() == el_0.size());
	g3_assert(az_0.size() == el_1.size());
	g3_assert(az_0.size() == az_1.size());
	g3_assert(az_0.size() == dec_0.size());
	g3_assert(az_0.size() == dec_1.size());
	g3_assert(az_0.size() == ra_0.size());
	g3_assert(az_0.size() == ra_1.size());
	trans_quats = G3VectorQuat(ra_0.size(), quat(1,0,0,0));

	for (size_t i = 0; i < ra_0.size(); i++) {
		trans_quats[i] = get_transform_quat(
		    az_0[i], -el_0[i],
		    ra_0[i], dec_0[i],
		    az_1[i], -el_1[i],
		    ra_1[i], dec_1[i]
		    )*get_origin_rotator(az_0[i], -el_0[i]);
	}
}

static void
convert_ra_dec_trans_to_gal(const G3VectorQuat &radec_trans,
    G3VectorQuat &gal_trans)
{
	// Converts a rotation from (1,0,0) to fk5 j2000 into a rotation that
	// takes (1,0,0) to our galactic (l,b)

	gal_trans = G3VectorQuat(radec_trans.size(), quat(1,0,0,0));
	quat gt = get_fk5_j2000_to_gal_quat();
	for (size_t i = 0; i < radec_trans.size(); i++)
		gal_trans[i] = gt*radec_trans[i];
}

void
get_detector_pointing(double x_offset, double y_offset,
    const G3VectorQuat &trans_quat, MapCoordReference coord_sys,
    std::vector<double> &alpha, std::vector<double> &delta)
{
	// For a detector x/y offset and a boresight position specified by
	// trans_quat with a given coordinate system coord_sys,
	// computes the individual detector pointing coordinates.

	quat det_pos = offsets_to_quat(x_offset, y_offset);
	delta.resize(trans_quat.size());
	alpha.resize(trans_quat.size());

	if ((!std::isfinite(x_offset)) || (!std::isfinite(y_offset))){
		log_debug("Found non-finite (inf or nan) offsets");
		for (size_t i=0; i<alpha.size(); i++){
			alpha[i] = nan("");
			delta[i] = nan("");
		}
		return;
	}

	for (size_t i = 0; i < alpha.size(); i++) {
		//using boost inverse
		//quat q=trans_quat[i]*det_pos/trans_quat[i];

		//uses an inverse that assumes we are on the unit sphere
		const quat & t = trans_quat[i];
		quat q=trans_quat[i]*det_pos * quat( t.R_component_1(),
		    -t.R_component_2(), -t.R_component_3(), -t.R_component_4());

		quat_to_ang(q, alpha[i], delta[i]);

		#ifdef CHECK_QUAT_INVERSE
		double a,d;
		quat u = trans_quat[i]*det_pos/trans_quat[i];
		quat_to_ang(u, a, d);
		if( fabs(a - alpha[i]) > 1e-5 || fabs(d - delta[i]) > 1e-5){
			log_fatal("Failed trans %lf %lf %lf %lf\n", a, alpha[i], d, delta[i]);
		}
		#endif
	}
	if (coord_sys == Local) {
		for (size_t i = 0; i < delta.size(); i++)
			delta[i] *= -1;
	}

}

void
get_detector_rotation(double x_offset, double y_offset,
    const G3VectorQuat &trans_quat, std::vector<double> &rot)
{
	// Computes the polarization angle rotation that occurs under the
	// transform trans_quat and stores it in rot.

	rot = std::vector<double>(trans_quat.size(), 0);
	quat det_pos = offsets_to_quat(x_offset, y_offset);
	for (size_t i = 0; i < rot.size(); i++) {
		quat q = trans_quat[i]*det_pos/trans_quat[i];
		rot[i] = get_rot_ang(det_pos, q, trans_quat[i]);
	}
}

static std::vector<double>
convert_celestial_offsets_to_local_offsets(G3VectorQuat trans_vec,
    double x_offset_celest, double y_offset_celest)
{
	double alpha, delta;
 	double x_offset_local, y_offset_local;

	// written at a time when we only have python bindings for the vector quat
	g3_assert( trans_vec.size() == 1);
	quat trans = trans_vec[0];

	// First we compute the celestial position of boresight.
	quat bs_cel = trans * quat(0,1,0,0) / trans;
	quat_to_ang(bs_cel, alpha, delta); //alpha, delta are celestial boresight

	// Next, with our celestial position offsets we compute the position
	// of our detector in celestial coordinates with the celestial offsets.
	// because the offset is in celestial coordinates we do that with the transform
	// that takes (1,0,0) in celestial coordinates to the boresight  while ignoring
	//boresight rotation

	quat celestial_trans = get_origin_rotator(alpha, delta);

	quat det_cel = celestial_trans * offsets_to_quat(x_offset_celest, y_offset_celest) /
	    celestial_trans;

	// We then transform those positions to local coordinates using the full transform
	quat inv_trans = boost::math::conj(trans);
	quat det_local = inv_trans * det_cel / inv_trans;

	#ifdef CHECK_QUAT_INVERSE
	//just to check I'm not insane can be dropped in the future
	g3_assert(sloppy_eq(inv_trans * bs_cel / inv_trans, quat(0,1,0,0)));
	#endif

	// we read off the position of the detector.
	quat_to_ang(det_local, x_offset_local, y_offset_local);

	// C++11 is magic.
	return {x_offset_local, -y_offset_local};
}

static double
angle_d2(double ra_0, double dec_0, double ra_1, double dec_1)
{
	double delta_ra_abs = fabs(ra_0 - ra_1);
	double delta_dec = dec_0 - dec_1;
	if (delta_ra_abs > PI){
		delta_ra_abs = 2 * PI - delta_ra_abs;
	}
	delta_ra_abs /= cos(dec_0);
	return delta_ra_abs*delta_ra_abs + delta_dec*delta_dec;

}

static G3VectorQuat
get_closest_transform(double ra, double dec,
    const G3VectorDouble &ras, const G3VectorDouble &decs,
    const G3VectorQuat &trans)
{
	// Returns quaternion transformation nearest the requested ra/dec

	double dist = angle_d2(ras[0]/G3Units::rad, decs[0]/G3Units::rad,
			       ra/G3Units::rad, dec/G3Units::rad);
	size_t ind = 0;
	for (size_t i = 0; i < ras.size(); i++) {
		double td = angle_d2(ras[i]/G3Units::rad, decs[i]/G3Units::rad,
				     ra/G3Units::rad, dec/G3Units::rad);
		if (td < dist) {
			dist = td;
			ind = i;
		}
	}
	G3VectorQuat v;
	v.push_back(trans[ind]);
	return v;
}

PYBINDINGS("maps")
{
	using namespace boost::python;

	def("test_trans_", test_trans);
	def("test_gal_trans_", test_gal_trans);
	def("test_gal_trans_rot_", test_gal_trans_rot);
	def("print_fk5_j2000_to_gal_quat_", print_fk5_j2000_to_gal_quat);
	def("c_quat_to_ang_", py_quat_to_ang);
	def("c_ang_to_quat_", ang_to_quat);

        def("get_origin_rotator", get_origin_rotator, (arg("alpha"), arg("delta")),
            "Compute the transformation quaternion that would rotate the "
            "vector (1, 0, 0) to point in the given direction.");
        def("offsets_to_quat", offsets_to_quat, (arg("x"), arg("y")),
            "Returns the vector quaternion (0,1,0,0) rotated by the given "
            "x and y offsets.  Equivalent to ``t * quat(0,1,0,0) / t``, where "
            "``t = get_origin_rotator(x, -y)``");
	def("create_det_az_el_trans", create_det_az_el_trans,
	    "Construct a quaternion vector from timestreams of detector "
	    "azimuth and elevation. Equivalent to ``R_z(az) * R_y(-el)``.");
	def("create_lazy_det_ra_dec_trans", create_lazy_det_ra_dec_trans,
	    "Construct a quaternion vector from timestreams of detector "
	    "RA and declination.  Equivalent to ``R_z(ra) * R_y(dec)``");
	def("create_det_ra_dec_trans", create_det_ra_dec_trans,
	    "Construct a quaternion vector from timestreams of detector "
	    "coordinates.  Computes the transformation from local "
	    "(az_0, el_0) coordinates to celestial (ra_0, dec_0), "
	    "accounting for rotation about the boresight by including "
	    "the second set of points.");
	def("convert_ra_dec_trans_to_gal", convert_ra_dec_trans_to_gal,
	    "Rotate a vector of quaternions from Equatorial to Galactic "
	    "coordinates");
}


// Copyright 2015-2022 The ALMA Project Developers
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//   http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
// implied. See the License for the specific language governing
// permissions and limitations under the License.

#pragma once

/// @file
/// This file contains some auxiliary code used for computing 1D
/// Green's functions and related quantities using piecewise cubic
/// interpolation.

#include <cmath>
#include <boost/math/special_functions/pow.hpp>

namespace alma {
namespace aux_cubic {
/// Auxiliary function used when computing rational integrals.
inline double I1(double x0) {
    return std::log(std::abs((1. - x0) / x0));
}

/// Auxiliary function used when computing rational integrals.
inline double I2(double dd, double ee) {
    double den{std::sqrt(4. * ee - boost::math::pow<2>(dd))};

    return -2. * (std::atan(dd / den) - std::atan((dd + 2. * ee) / den)) / den;
}

/// Auxiliary function used when computing rational integrals.
inline double I3(double dd, double ee) {
    double den = std::sqrt(4. * ee - boost::math::pow<2>(dd));

    return (2. * dd * (std::atan(dd / den) - std::atan((dd + 2. * ee) / den)) +
            den * std::log1p(dd + ee)) /
           (2 * ee * den);
}

/// Objects of this class perform rational integrals of the kind
/// found when approximating bands by piecewise cubic polynomials
/// to obtain Green's functions in 1D.
class Cubic_segment {
private:
    /// Coefficients of the cubic polynomial in the denominator.
    const std::array<double, 4> coeff;

public:
    /// Constructor.
    ///
    /// @param[in] a - value of the energy at the left
    /// @param[in] b - value of the energy at the right
    /// @param[in] ap - derivative of the energy at the left
    /// @param[in] bp - derivative of the energy at the right
    Cubic_segment(double a, double b, double ap, double bp)
        : coeff({{-(2. * a + ap - 2. * b + bp),
                  -(-3. * a - 2. * ap + 3. * b - bp),
                  -ap,
                  -a}}) {
    }
    /// Obtain the values of the four integrals needed for the
    /// calculation of the Green's function at a particular
    /// energy.
    ///
    /// @param[in] energ - value of the energy
    /// @return a vector with the four integrals in this order:
    /// ReG, ReGx, ImG, ImGx
    std::array<double, 4> calc_integrals(double ener) const {
        // Coefficients of the full polynomial in the
        // denominator.
        std::array<double, 4> p(this->coeff);
        p[3] += ener;
        // Get the number of roots based on the value of the
        // discriminant for the third-order equation.
        double a{p[1] / p[0]};
        double b{p[2] / p[0]};
        double c{p[3] / p[0]};
        double a2{boost::math::pow<2>(a)};
        double q{(a2 - 3. * b) / 9.};
        double r{(a * (2. * a2 - 9. * b) + 27. * c) / 54.};
        double r2{boost::math::pow<2>(r)};
        double q3{boost::math::pow<3>(q)};
        std::array<double, 4> nruter({{0., 0., 0., 0.}});

        // And use the formulae specific to each case.
        if (r2 < q3) {
            // Three real roots.
            double th{std::acos(r / std::sqrt(q3))};
            double q12{std::sqrt(q)};
            double x0{-2. * q12 * std::cos(th / 3.) - a / 3.};
            double x1{-2. * q12 * std::cos((th + 2. * constants::pi) / 3.) -
                      a / 3.};
            double x2{-2. * q12 * std::cos((th - 2. * constants::pi) / 3.) -
                      a / 3.};
            double pref{-x0 * x1 * x2 / p[3]};
            double aa{pref / ((x0 - x1) * (x0 - x2))};
            double bb{pref / ((x1 - x0) * (x1 - x2))};
            double cc{pref / ((x2 - x0) * (x2 - x1))};
            nruter[0] = aa * I1(x0) + bb * I1(x1) + cc * I1(x2);
            double aax{x0 * aa};
            double bbx{x1 * bb};
            double ccx{x2 * cc};
            nruter[1] = aax * I1(x0) + bbx * I1(x1) + ccx * I1(x2);

            if ((0. < x0) && (x0 < 1.)) {
                nruter[2] += std::abs(aa);
                nruter[3] += std::abs(aax);
            }

            if ((0. < x1) && (x1 < 1.)) {
                nruter[2] += std::abs(bb);
                nruter[3] += std::abs(bbx);
            }

            if ((0. < x2) && (x2 < 1.)) {
                nruter[2] += std::abs(cc);
                nruter[3] += std::abs(ccx);
            }
            nruter[2] *= constants::pi;
            nruter[3] *= constants::pi;
        }
        else {
            // Only one real root.
            double A{-signum(r) * std::cbrt(std::abs(r) + std::sqrt(r2 - q3))};
            double B{(A == 0. ? 0. : q / A)};
            double x0{A + B - a / 3.};
            double dd{A + B + 2. * a / 3.};
            double ee{boost::math::pow<2>(dd / 2.) +
                      .75 * boost::math::pow<2>(A - B)};
            double pref{-x0 * ee / p[3]};
            double aa{pref / (x0 * (x0 + dd) + ee)};
            double bb{-(x0 + dd) * aa};
            double cc{-aa};
            double aax{x0 * aa};
            double bbx{ee * aa};
            double ccx{-aax};
            nruter[0] = aa * I1(x0) + (bb / ee) * I2(dd / ee, 1. / ee) +
                        (cc / ee) * I3(dd / ee, 1. / ee);
            nruter[1] = aax * I1(x0) + (bbx / ee) * I2(dd / ee, 1. / ee) +
                        (ccx / ee) * I3(dd / ee, 1. / ee);

            if ((0. < x0) && (x0 < 1.)) {
                nruter[2] = constants::pi * std::abs(aa);
                nruter[3] = constants::pi * std::abs(aax);
            }
        }
        return nruter;
    }
};
} // namespace aux_cubic
} // namespace alma


/*************************************************************************
	> File Name: lqr_steer_control.cpp
	> Author: TAI Lei
	> Mail: ltai@ust.hk
	> Created Time: Wed Apr 17 11:48:46 2019
 ************************************************************************/

#include <iostream>
#include <limits>
#include <vector>
#include <opencv2/opencv.hpp>
#include <opencv2/core/core.hpp>
#include <opencv2/highgui/highgui.hpp>
#include <sys/time.h>
#include <Eigen/Eigen>
#include "cubic_spline.h"
#include "motion_model.h"
#include "cpprobotics_types.h"

#define DT 0.1
#define L 0.5
#define KP 1.0
#define MAX_STEER 45.0/180*M_PI

using namespace cpprobotics;
using Matrix5f = Eigen::Matrix<float, 5, 5>;
using Matrix52f = Eigen::Matrix<float, 5, 2>;
using Matrix25f = Eigen::Matrix<float, 2, 5>;
using RowVector5f = Eigen::Matrix<float, 1, 5>;
using Vector5f = Eigen::Matrix<float, 5, 1>;

cv::Point2i cv_offset(
    float x, float y, int image_width=2000, int image_height=2000){
  cv::Point2i output;
  output.x = int(x * 100) + 300;
  output.y = image_height - int(y * 100) - image_height/2;
  return output;
};

Vec_f calc_speed_profile(Vec_f rx, Vec_f ry, Vec_f ryaw, float target_speed){
    Vec_f speed_profile(ryaw.size(), target_speed);

  float direction = 1.0;
  for(unsigned int i=0; i < ryaw.size()-1; i++){
    float dyaw = std::abs(ryaw[i+1] - ryaw[i]);
    float switch_point = (M_PI/4.0< dyaw) && (dyaw<M_PI/2.0);

    if (switch_point) direction = direction * -1;
    if (direction != 1.0) speed_profile[i]= target_speed * -1;
    else speed_profile[i]= target_speed;

    if (switch_point) speed_profile[i] = 0.0;
  }

  for(int k=0; k< 40; k++){
    *(speed_profile.end()-k) = target_speed / (50 - k);
    if (*(speed_profile.end()-k) <= 1.0 / 3.6){
      *(speed_profile.end()-k) = 1.0 / 3.6;
    }
  }
  return speed_profile;
};

float calc_nearest_index(State state, Vec_f cx, Vec_f cy, Vec_f cyaw, int &ind){
  float mind = std::numeric_limits<float>::max();
  for(unsigned int i=0; i<cx.size(); i++){
    float idx = cx[i] - state.x;
    float idy = cy[i] - state.y;
    float d_e = idx*idx + idy*idy;

    if (d_e<mind){
      mind = d_e;
      ind = i;
    }
  }
  float dxl = cx[ind] - state.x;
  float dyl = cy[ind] - state.y;
  float angle = YAW_P2P(cyaw[ind] - std::atan2(dyl, dxl));
  if (angle < 0) mind = mind * -1;

  return mind;
};

Matrix5f solve_DARE(Matrix5f A, Matrix52f B, Matrix5f Q, Eigen::Matrix2f R){
  Matrix5f X = Q;
  int maxiter = 150;
  float eps = 0.01;

  for(int i=0; i<maxiter; i++){
    Matrix5f Xn = A.transpose()*X*A-A.transpose()*X*B*(R+B.transpose()*X*B).inverse() * B.transpose()*X*A+Q;
    Matrix5f error = Xn - X;
    if (error.cwiseAbs().maxCoeff()<eps){
      return Xn;
    }
    X = Xn;
  }

  return X;
};

Matrix25f dlqr(Matrix5f A, Matrix52f B, Matrix5f Q, Eigen::Matrix2f R){
  Matrix5f X = solve_DARE(A, B ,Q, R);
  Matrix25f K = (B.transpose()*X*B + R).inverse() * (B.transpose()*X*A);
  return K;
};

Vec_f lqr_steering_control(State state, Vec_f cx, Vec_f cy, Vec_f cyaw, Vec_f ck, Vec_f sp, float& pe, float& pth_e){
  int ind = 0;
  float e = calc_nearest_index(state, cx, cy, cyaw, ind);

  float k = ck[ind];
  float th_e = YAW_P2P(state.yaw - cyaw[ind]);
  float tv = sp[ind];

  Matrix5f A = Matrix5f::Zero();
  A(0, 0) = 1.0;
  A(0 ,1) = DT;
  A(1 ,2) = state.v;
  A(2 ,2) = 1.0;
  A(2 ,3) = DT;
  A(4 ,4) = 1.0;

  Matrix52f B = Matrix52f::Zero();
  B(3, 0) = state.v/L;
  B(4, 1) = DT;

  Matrix5f Q = Matrix5f::Identity();
  Eigen::Matrix2f R = Eigen::Matrix2f::Identity();

  // gain of lqr
  Matrix25f K = dlqr(A, B, Q, R);

  Vector5f x = Vector5f::Zero();
  x(0) = e;
  x(1) = (e-pe)/DT;
  x(2) = th_e;
  x(3) = (th_e-pth_e)/DT;
  x(4) = state.v - tv;

  Eigen::Vector2f ustar = -K * x;

  float ff = std::atan2((L*k), (double)1.0);
  float fb = YAW_P2P(ustar(0));
  float delta = ff+fb;
  float ai = ustar(1);

  pe = e;
  pth_e = th_e;
  return {ai, delta};
};

void update (State& state, float a, float delta){

  if (delta >= MAX_STEER) delta = MAX_STEER;
  if (delta <= - MAX_STEER) delta = - MAX_STEER;

  state.x = state.x + state.v * std::cos(state.yaw) * DT;
  state.y = state.y + state.v * std::sin(state.yaw) * DT;
  state.yaw = state.yaw + state.v / L * std::tan(delta) * DT;
  state.v = state.v + a * DT;

};

void closed_loop_prediction(Vec_f cx, Vec_f cy, Vec_f cyaw, Vec_f ck, Vec_f speed_profile, Poi_f goal){
  float T = 500.0;
  float goal_dis = 0.3;
  float stop_speed = 0.05;

  State state(-0.0, -0.0, 0.0, 0.0);

  float time_ = 0.0;
  Vec_f x;
  x.push_back(state.x);
  Vec_f y;
  y.push_back(state.y);
  Vec_f yaw;
  yaw.push_back(state.yaw);
  Vec_f v;
  v.push_back(state.v);
  Vec_f t;
  t.push_back(0.0);

  float e = 0;
  float e_th = 0;

  cv::namedWindow("lqr_full", cv::WINDOW_NORMAL);
  int count = 0;
  Vec_f x_h;
  Vec_f y_h;

  while (T >= time_){
    Vec_f control = lqr_steering_control(state, cx, cy, cyaw, ck, speed_profile, e, e_th);
    // float ai = KP * (speed_profile[ind]-state.v);
    update(state, control[0], control[1]);
    // if (std::abs(state.v) <= stop_speed) ind += 1;

    float dx = state.x - goal[0];
    float dy = state.y - goal[1];
    if (std::sqrt(dx*dx + dy*dy) <= goal_dis) {
      std::cout<<("Goal")<<std::endl;
      break;
    }

    x_h.push_back(state.x);
    y_h.push_back(state.y);

    // visualization
    cv::Mat bg(2000, 3000, CV_8UC3, cv::Scalar(255, 255, 255));
    for(unsigned int i=1; i<cx.size(); i++){
      cv::line(
        bg,
        cv_offset(cx[i-1], cy[i-1], bg.cols, bg.rows),
        cv_offset(cx[i], cy[i], bg.cols, bg.rows),
        cv::Scalar(0, 0, 0),
        10);
    }

    for(unsigned int j=0; j< x_h.size(); j++){
      cv::circle(
        bg,
        cv_offset(x_h[j], y_h[j], bg.cols, bg.rows),
        10, cv::Scalar(0, 0, 255), -1);
    }

    cv::putText(
      bg,
      "Speed: " + std::to_string(state.v*3.6).substr(0, 4) + "km/h",
      cv::Point2i((int)bg.cols*0.5, (int)bg.rows*0.1),
      cv::FONT_HERSHEY_SIMPLEX,
      3,
      cv::Scalar(0, 0, 0),
      10);

    // save image in build/bin/pngs
    struct timeval tp;
    gettimeofday(&tp, NULL);
    long int ms = tp.tv_sec * 1000 + tp.tv_usec / 1000;
    std::string int_count = std::to_string(ms);
    cv::imwrite("./pngs/"+int_count+".png", bg);
    // cv::imshow("lqr_full", bg);
    // cv::waitKey(5);
  }
};

int main(){
  Vec_f wx({0.0, 6.0,  12.5, 10.0, 17.5, 20.0, 25.0});
  Vec_f wy({0.0, -3.0, -5.0,  6.5, 3.0, 0.0, 0.0});

  Spline2D csp_obj(wx, wy);
  Vec_f r_x;
  Vec_f r_y;
  Vec_f ryaw;
  Vec_f rcurvature;
  Vec_f rs;
  for(float i=0; i<csp_obj.s.back(); i+=0.1){
    std::array<float, 2> point_ = csp_obj.calc_postion(i);
    r_x.push_back(point_[0]);
    r_y.push_back(point_[1]);
    ryaw.push_back(csp_obj.calc_yaw(i));
    rcurvature.push_back(csp_obj.calc_curvature(i));
    rs.push_back(i);
  }
  float target_speed = 10.0 / 3.6;
  Vec_f speed_profile = calc_speed_profile(r_x, r_y, ryaw, target_speed);
  closed_loop_prediction(r_x, r_y, ryaw, rcurvature, speed_profile, {{wx.back(), wy.back()}});
}