#include "filter.h"
#include <Eigen/Dense>
#include <Eigen/Core>
#include <iostream>
#include <fstream>
#include <math.h>

using namespace Eigen;
using namespace std;

inline double Gauss(double sigma, double x)
{
	double expVal = -1 * (pow(x, 2) / (2*pow(sigma, 2)));
	double divider = sqrt(2 * PI * pow(sigma, 2));
	return (1 / divider) * exp(expVal);
}

// a : signal , b : filter
inline MatrixXd Conv(MatrixXd a, MatrixXd b)
{
	int size = a.rows();
	int sizeb = b.rows();
	int window(0);
	MatrixXd out(size, 1);
	double sum(0);

	for (int i = 0; i < size; i++) {
		if (i >= sizeb) window = sizeb;
		else  window = i;

		for (int j = 0; j < window; j++) {
			sum += b(j, 0) * a(i - j);
		}
		out(i, 0) = sum;
		sum = 0;
	}

	return out;
}

Filter::Filter()
	: sum_kernel_(0.0), weight_(0.0)
{
	kernel_.resize(kkernelLevel, 1);
}

Filter::~Filter()
{}

void Filter::KernelMake(int samples, double sigma)
{
	kernel_.resize(samples, 1);
	bool doubleCenter = false;
	if (kkernelLevel % 2 == 0) {
		doubleCenter = true;
		samples--;
	}

	int steps = (samples - 1) / 2;
	double stepSize = (3 * sigma) / steps;

	for (int i = steps; i >= 1; i--) {
		kernel_(steps - i, 0) = Gauss(sigma, -1*i*stepSize);
		kernel_(kernel_.rows() + i - steps - 1, 0) = Gauss(sigma, i * stepSize);
	}
	kernel_(steps, 0) = Gauss(sigma, 0);
	if (doubleCenter) kernel_(steps + 1, 0) = Gauss(sigma, 0);

	sum_kernel_ = 0;
	for (int i = 0; i < kernel_.rows(); i++) {
		sum_kernel_ += kernel_(i, 0);
	}
	weight_ = 1 / sum_kernel_;
}

MatrixXd Filter::LanczosDiffLow(MatrixXd const& a, int order = 5, int circle = 0)
{
	int n = a.rows();

	if (n < 4) {
		cout << "size of matrix is wrong" << endl;
	}

	diff_.resize(n, 1);

	int m = (order - 1) / 2;
	double temp(0);
	int start(m);

	if (circle) start = 0;

	for (int i = start; i < n - start; i++) {
		temp = 0;
		for (int j = 1; j < m + 1; j++) {
			int front = i + j;
			int back = i - j;

			if (front > n - 1) front = i + j - n;
			if (back < 0) back = i - j + n;
			temp += j * (a(front, 0) - a(back, 0)) / (m * (m + 1) * (2 * m + 1));
		}
		diff_(i, 0) = 3 * temp;
	}

	if(!circle) {
		for (int i = 0; i < m; i++) {
			diff_(i, 0) = a(i + 1) - a(i, 0);
			diff_(n - i - 1, 0) = a(n - i - 1, 0) - a(n - i - 2, 0);
		}
	}

	return diff_;
}

void Filter::Gaussian(double sigma,
	MatrixXd const& in,
	MatrixXd& out,
	int circle)
{
	int samples = kkernelLevel;
	KernelMake(samples, sigma);
	int sampleSide = samples / 2;
	int valueIdx = samples / 2 + 1;
	int ubound = in.rows();
	out.resize(ubound, 1);
	MatrixXd temp(ubound, 1);
	temp = in;
	int start(sampleSide);

	if (circle) start = 0;

	else if (!circle) {
		for (int i = 0; i < sampleSide; i++) {
			out(i, 0) = in(i, 0);
			out(ubound - i - 1) = in(ubound - i - 1);
		}
	}

	for (int i = start; i < ubound - start; i++) {
		double sum = 0;
		int sampleCtr = 0;
		int init = i - sampleSide;
		int limit = i + sampleSide;

		for (int j = init ; j < limit ; j++) {
			int index(j);
			if (j < 0) index = j + ubound;
			else if (j > ubound - 1) index = j - ubound;

			int sampleWeightIndex = sampleSide + (j - i);
			sum += kernel_(sampleWeightIndex, 0) * temp(index);
			sampleCtr++;
		}
		double smoothed = sum * weight_;
		out(i, 0) = smoothed;
	}
}

void CalculateFeatureAxisless(Geom& shard, int axis_index)
{
	Filter filter;

	if (!shard.is_matching_) return;

	// Feature calculation
	int number_of_points = shard.edge_line_.point_.cols();
	MatrixXd Dist(number_of_points, 1), Thickness(number_of_points, 1);
	MatrixXd Height(number_of_points, 1), Theta(number_of_points, 1);

	for (int i = 0; i < number_of_points; i++) {
		Vector3d r = { shard.edge_line_.point_(0, i), shard.edge_line_.point_(1, i), 0 };
		Dist(i) = r.norm();
		Theta(i) = atan2(shard.edge_line_.point_(1, i), shard.edge_line_.point_(0, i));
		Height(i) = shard.edge_line_.point_(2, i);
	}
	for (int i = 1; i < Theta.rows(); i++) {
		if (Theta(i, 0) - Theta(i - 1, 0) > PI) {
			Theta(i, 0) = Theta(i, 0) - 2 * PI;
		}
		else if (Theta(i, 0) - Theta(i - 1, 0) < -PI) {
			Theta(i, 0) = Theta(i, 0) + 2 * PI;
		}
	}

	if ((axis_index == 0) && (!shard.is_thickness_)) {
		Thickness = GetThickness(shard);
	}
	else {
		Thickness = shard.edge_line_.feature_[0].row(6);
	}

	MatrixXd Dist_Diff = filter.LanczosDiffLow(Dist, 7, 0);
	MatrixXd Height_Diff = filter.LanczosDiffLow(Height, 7, 0);
	MatrixXd Theta_Diff = filter.LanczosDiffLow(Theta, 7, 0);

	//Calculate Dist*d(Theta)
	for (int i = 0; i < Theta_Diff.rows(); i++) {
		Theta_Diff(i) = Theta_Diff(i) * Dist(i);
	}

	filter.Gaussian(2, Dist_Diff, Dist_Diff, 1);
	filter.Gaussian(2, Height_Diff, Height_Diff, 1);
	filter.Gaussian(2, Theta_Diff, Theta_Diff, 1);

	int num_features = Dist_Diff.rows();

	shard.edge_line_.feature_[axis_index].resize(7, num_features);
	for (int i = 0; i < num_features; i++) {
		shard.edge_line_.feature_[axis_index].col(i) << Dist_Diff(i),
			Height_Diff(i),
			Theta_Diff(i),
			Dist(i),
			Height(i),
			Theta(i),
			Thickness(i);
	}
}

MatrixXd GetThickness(Geom& shard)
{
	int num_breakline = shard.edge_line_.point_.cols();
	int num_o_sur = shard.sur_out_.point_.cols();
	MatrixXd thickness(num_breakline, 1);

	int index_out(0);
	for (size_t i = 0; i < num_breakline; i++) {
		Eigen::Vector4f pt(1, 0, 0, 0), line_pt(0, 0, 0, 0), line_dir(1, 1, 0, 0);
		double point2line_disance = 99999999, dist_Tmp = 0;
		line_pt << shard.edge_line_.point_(0, i), shard.edge_line_.point_(1, i), shard.edge_line_.point_(2, i), 0;
		line_dir << shard.edge_line_.normal_(0, i), shard.edge_line_.normal_(1, i), shard.edge_line_.normal_(2, i), 0;

		for (size_t j = 0; j < num_o_sur; j++) {
			pt << shard.sur_out_.point_(0, j), shard.sur_out_.point_(1, j), shard.sur_out_.point_(2, j), 0;
			dist_Tmp = sqrt(pcl::sqrPointToLineDistance(pt, line_pt, line_dir));
			if (dist_Tmp < point2line_disance)
			{
				point2line_disance = dist_Tmp;
				index_out = j;
			}
		}

		Eigen::Vector4f a;
		a << shard.sur_out_.point_(0, index_out), shard.sur_out_.point_(1, index_out),
			shard.sur_out_.point_(2, index_out), 0;

		if ((point2line_disance > 1.0) || (line_dir.dot(a - line_pt) > 0)) {
			thickness(i) = -1.0;
		}
		else {
			thickness(i) = (shard.sur_out_.point_.col(index_out) - shard.edge_line_.point_.col(i)).norm();
		}

	}

	shard.is_thickness_ = true;

	return thickness;
}

// (r, theta, z) 3*n matrix
MatrixXd ToCylindricalInterpolation(const BreakLine& breakline,
	bool theta_sort,
	bool r_sort)
{
	MatrixXd b1 = breakline.point_;
	int number_of_points = b1.cols();
	vector<Vector3d> cy;
	double base_theta;
	for (int i = 0; i < number_of_points; i++) {
		Vector3d r = { b1(0, i), b1(1, i), 0 }, tmp;
		tmp(0) = r.norm();
		tmp(1) = atan2(b1(1, i), b1(0, i));
		tmp(2) = b1(2, i);
		cy.push_back(tmp);
	}

	for (int i = 1; i < number_of_points; i++) {
		double gap = abs(cy[i](1) - cy[i - 1](1));
		if ((gap > 0.04) && (gap < 3.14)) {
			int num_interpol = gap / 0.04;
			double r_step = (cy[i](0) - cy[i - 1](0)) / num_interpol;
			double theta_step = (cy[i](1) - cy[i - 1](1)) / num_interpol;
			double z_step = (cy[i](2) - cy[i - 1](2)) / num_interpol;
			for (int j = 0; j < num_interpol; j++) {
				Vector3d tmp;
				tmp(0) = cy[i - 1](0) + r_step * (j + 1);
				tmp(1) = cy[i - 1](1) + theta_step * (j + 1);
				tmp(2) = cy[i - 1](2) + z_step * (j + 1);
				cy.push_back(tmp);
			}
		}
	}

	if (theta_sort) {
		sort(cy.begin(), cy.end(), [](Vector3d a, Vector3d b) -> bool {
			return a(1) < b(1);
		});
	}
	else if (r_sort) {
		sort(cy.begin(), cy.end(), [](Vector3d a, Vector3d b) -> bool {
			return a(0) > b(0);
		});
	}
	number_of_points = cy.size();
	MatrixXd output(3, number_of_points);

	for(int i = 0; i < number_of_points; i++){
		output(0, i) = cy[i](0);
		output(1, i) = cy[i](1);
		output(2, i) = cy[i](2);
	}

	return output;
}

void ToCylindricalInterpolation(const BreakLine& breakline,
	vector<Vector3d>& out,
	bool interpol)
{
	MatrixXd b1 = breakline.point_;
	int number_of_points = b1.cols();
	double base_theta;
	for (int i = 0; i < number_of_points; i++) {
		Vector3d r = { b1(0, i), b1(1, i), 0 }, tmp;
		tmp(0) = r.norm();
		tmp(1) = atan2(b1(1, i), b1(0, i));
		tmp(2) = b1(2, i);
		out.push_back(tmp);
	}

	if (interpol) {
		for (int i = 1; i < number_of_points; i++) {
			double gap = abs(out[i](1) - out[i - 1](1));
			if ((gap > 0.04) && (gap < 3.14)) {
				int num_interpol = gap / 0.04;
				double r_step = (out[i](0) - out[i - 1](0)) / num_interpol;
				double theta_step = (out[i](1) - out[i - 1](1)) / num_interpol;
				double z_step = (out[i](2) - out[i - 1](2)) / num_interpol;
				for (int j = 0; j < num_interpol; j++) {
					Vector3d tmp;
					tmp(0) = out[i - 1](0) + r_step * (j + 1);
					tmp(1) = out[i - 1](1) + theta_step * (j + 1);
					tmp(2) = out[i - 1](2) + z_step * (j + 1);
					out.push_back(tmp);
				}
			}

			double r_gap = abs(out[i](0) - out[i - 1](0));
			if ((r_gap > 0.1) && (gap < 3.14)) {
				int num_interpol = r_gap / 0.1;
				double r_step = (out[i](0) - out[i - 1](0)) / num_interpol;
				double theta_step = (out[i](1) - out[i - 1](1)) / num_interpol;
				double z_step = (out[i](2) - out[i - 1](2)) / num_interpol;
				for (int j = 0; j < num_interpol; j++) {
					Vector3d tmp;
					tmp(0) = out[i - 1](0) + r_step * (j + 1);
					tmp(1) = out[i - 1](1) + theta_step * (j + 1);
					tmp(2) = out[i - 1](2) + z_step * (j + 1);
					out.push_back(tmp);
				}
			}
		}
	}
}


/* Copyright (C) 2010-2019, The Regents of The University of Michigan.
 All rights reserved.

 This software was developed as part of the The Vulcan project in the Intelligent Robotics Lab
 under the direction of Benjamin Kuipers, kuipers@umich.edu. Use of this code is governed by an
 MIT-style License that can be found at "https://github.com/h2ssh/Vulcan".
*/

/**
* \file     univariate_distribution_estimator_impl.cpp
* \author   Collin Johnson
*
* Definition of implementations of the UnivariateDistributionEstimator for the subclasses of UnivariateDistribution:
*
*   - UnivariateGaussianDistribution
*   - GammaDistribution
*   - BetaDistribution
*   - ExponentialDistribution
*   - TruncatedGaussianDistribution
*/

#include <math/univariate_distribution_estimator_impl.h>
#include <math/univariate_gaussian.h>
#include <math/discrete_gaussian.h>
#include <math/beta_distribution.h>
#include <math/exponential_distribution.h>
#include <math/gamma_distribution.h>
#include <math/truncated_gaussian_distribution.h>
#include <math/statistics.h>
#include <boost/math/special_functions/digamma.hpp>
#include <boost/math/special_functions/gamma.hpp>
#include <iostream>

namespace vulcan
{
namespace math
{

using DistPtr = std::unique_ptr<UnivariateDistribution>;

double log_likelihood_gamma_func(double k, double sumXi, double sumLogXi, std::size_t n)
{
    return (k-1)*sumXi - n*k - n*k*std::log(sumXi/(k*n)) - n*std::log(boost::math::tgamma(k));
}

double log_likelihood_gamma_deriv(double k, double sumXi, double sumLogXi, std::size_t n)
{
    return n*(std::log(k) - boost::math::digamma(k) - std::log(sumXi/n)) + sumLogXi;
}

DistPtr UnivariateGaussianDistributionEstimator::estimateDistribution(const std::vector<double>::const_iterator dataBegin,
                                                                      const std::vector<double>::const_iterator dataEnd) const
{
    return DistPtr{new UnivariateGaussianDistribution(estimate(dataBegin, dataEnd))};
}

UnivariateGaussianDistribution UnivariateGaussianDistributionEstimator::estimate(const std::vector<double>::const_iterator dataBegin,
                                                                                const std::vector<double>::const_iterator dataEnd) const
{
    double var = variance(dataBegin, dataEnd);

    if(var == 0.0)
    {
        std::cerr << "WARNING: UnivariateGaussianDistributionEstimator: Variance in the data was 0. Setting variance to 1e-4.\n";
        var = 1e-4;
    }

    return UnivariateGaussianDistribution(mean(dataBegin, dataEnd), var);
}

DistPtr DiscreteGaussianDistributionEstimator::estimateDistribution(const std::vector<double>::const_iterator dataBegin,
                                                                    const std::vector<double>::const_iterator dataEnd) const
{
    return DistPtr{new DiscreteGaussianDistribution(estimate(dataBegin, dataEnd))};
}

DiscreteGaussianDistribution DiscreteGaussianDistributionEstimator::estimate(const std::vector<double>::const_iterator dataBegin,
                                                                             const std::vector<double>::const_iterator dataEnd) const
{
    double var = variance(dataBegin, dataEnd);

    if(var == 0.0)
    {
        std::cerr << "WARNING: DiscreteGaussianDistributionEstimator: Variance in the data was 0. Setting variance to 1e-4.\n";
        var = 1e-4;
    }

    return DiscreteGaussianDistribution(mean(dataBegin, dataEnd), var);
}

DistPtr GammaDistributionEstimator::estimateDistribution(const std::vector<double>::const_iterator dataBegin,
                                                         const std::vector<double>::const_iterator dataEnd) const
{
    return DistPtr{ new GammaDistribution(estimate(dataBegin, dataEnd)) };
}

GammaDistribution GammaDistributionEstimator::estimate(const std::vector<double>::const_iterator dataBegin,
                                                       const std::vector<double>::const_iterator dataEnd) const
{
    std::vector<double> filtered(dataBegin, dataEnd);
    auto validDataEnd = std::remove_if(filtered.begin(), filtered.end(), [](double x) { return x <= 0.0; });

    int numValidData = std::distance(filtered.begin(), validDataEnd);
    assert(numValidData);

    using namespace std::placeholders;

    double sumXi    = std::accumulate(filtered.begin(), validDataEnd, 0.0);
    double sumLogXi = 0.0;
    for(auto val : boost::make_iterator_range(filtered.begin(), validDataEnd))
    {
        sumLogXi += std::log(val);
    }

    double s  = std::log(sumXi/numValidData) - sumLogXi/numValidData;
    double k0 = (3.0 - s + std::sqrt(std::pow(s-3.0, 2.0) + 24.0*s)) / (12.0 * s);

    //     NewtonRaphsonErrorFunc<double> newton(std::bind(log_likelihood_func,  _1, sumXi, sumLogXi, filtered.size()),
    //                                           std::bind(log_likelihood_deriv, _1, sumXi, sumLogXi, filtered.size()));
    //
    //     double k     = find_single_root(newton, k0, 1e-5);
    double k     = k0;
    double theta = sumXi / (numValidData * k);
    return GammaDistribution(k, theta);
}

DistPtr BetaDistributionEstimator::estimateDistribution(const std::vector<double>::const_iterator dataBegin,
                                                        const std::vector<double>::const_iterator dataEnd) const
{
    return DistPtr{ new BetaDistribution(estimate(dataBegin, dataEnd)) };
}

BetaDistribution BetaDistributionEstimator::estimate(const std::vector<double>::const_iterator dataBegin,
                                                     const std::vector<double>::const_iterator dataEnd) const
{
    double mean     = math::mean    (dataBegin, dataEnd);
    double variance = math::variance(dataBegin, dataEnd);

    if(variance == 0.0)
    {
        std::cerr << "WARNING: BetaDistributionEstimator: Variance of data was 0. Setting to 1e-4.\n";
        variance = 1e-4;
    }

    if(variance < mean*(1-mean) && (variance > 0.0))
    {
        double alpha = mean         * ((mean * (1.0-mean) / variance) - 1.0);
        double beta  = (1.0 - mean) * ((mean * (1.0-mean) / variance) - 1.0);
        return BetaDistribution(alpha, beta);
    }
    else
    {
        std::cerr << "ERROR: BetaDistributionEstimator: Could not use method-of-moments to find parameters. Mean:" << mean << " Variance:" << variance << '\n'
                  << " Variance should be less than " << (mean * (1-mean)) << '\n';
        return BetaDistribution();
    }
}

ExponentialDistributionEstimator::ExponentialDistributionEstimator(double maxValue)
: max_(maxValue)
{
}

DistPtr ExponentialDistributionEstimator::estimateDistribution(const std::vector<double>::const_iterator dataBegin,
                                                               const std::vector<double>::const_iterator dataEnd) const
{
    return DistPtr{ new ExponentialDistribution(estimate(dataBegin, dataEnd)) };
}

ExponentialDistribution ExponentialDistributionEstimator::estimate(const std::vector<double>::const_iterator dataBegin,
                                                                   const std::vector<double>::const_iterator dataEnd) const
{
    double mean = math::mean(dataBegin, dataEnd);
    if(mean > 0.0)
    {
        return ExponentialDistribution(1.0/mean, max_);
    }
    else
    {
        std::cerr << "ERROR: ExponentialDistributionEstimator: Invalid mean for the data:" << mean <<" Must be greater than 0.\n";
        return ExponentialDistribution(1.0, max_);
    }
}

TruncatedGaussianDistributionEstimator::TruncatedGaussianDistributionEstimator(double lower, double upper)
: lower_(lower)
, upper_(upper)
{
}

DistPtr TruncatedGaussianDistributionEstimator::estimateDistribution(const std::vector<double>::const_iterator dataBegin,
                                                                     const std::vector<double>::const_iterator dataEnd) const
{
    return DistPtr{ new TruncatedGaussianDistribution(estimate(dataBegin, dataEnd)) };
}

TruncatedGaussianDistribution TruncatedGaussianDistributionEstimator::estimate(const std::vector<double>::const_iterator dataBegin,
                                                                               const std::vector<double>::const_iterator dataEnd) const
{
    double var = variance(dataBegin, dataEnd);

    if(var == 0.0)
    {
        std::cerr << "WARNING: TruncatedGaussianDistributionEstimator: Variance in the data was 0. Setting variance to 1e-4.\n";
        var = 1e-4;
    }

    return TruncatedGaussianDistribution(mean(dataBegin, dataEnd), var, lower_, upper_);
}

} // namespace math
} // namespace vulcan


/* Copyright (C) 2010-2019, The Regents of The University of Michigan.
 All rights reserved.

 This software was developed as part of the The Vulcan project in the Intelligent Robotics Lab
 under the direction of Benjamin Kuipers, kuipers@umich.edu. Use of this code is governed by an
 MIT-style License that can be found at "https://github.com/h2ssh/Vulcan".
*/

/**
 * \file     em.cpp
 * \author   Collin Johnson
 *
 * Definition of em_1d_linear and em_2d_fixed.
 */

#include <math/clustering.h>
#include <core/multivariate_gaussian.h>
#include <math/statistics.h>
#include <math/univariate_gaussian.h>
#include <boost/range/iterator_range.hpp>
#include <algorithm>
#include <iomanip>
#include <iostream>

namespace vulcan
{
namespace math
{

template <class Data>
using DataIt = typename std::vector<Data>::const_iterator;

template<class Data, class Dist>
struct em_data_t
{
    int numClusters;
    std::vector<int> clusterSizes;
    std::vector<Dist> dists;
    std::vector<int> assignments;
    std::vector<Data> clusterData;  // temporary storage for computing distributions
    double totalProb;
};

template <class Data, class Dist, class ProbFunc, class Estimator>
em_data_t<Data, Dist> run_em(DataIt<Data> begin,
                             DataIt<Data> end,
                             int k,
                             int maxIterations,
                             ProbFunc prob,
                             Estimator estimator);

template <class Data, class Dist, class Estimator>
void initialize_dists(DataIt<Data> begin, DataIt<Data> end, int k, em_data_t<Data, Dist>& dists, Estimator estimator);

template <class Data, class Dist, class Estimator>
void calculate_dists(DataIt<Data> begin, DataIt<Data> end, em_data_t<Data, Dist>& dists, Estimator estimator);

template <class Data, class Dist, class ProbFunc>
int assign_dists(DataIt<Data> begin, DataIt<Data> end, em_data_t<Data, Dist>& dists, ProbFunc prob);

template <class Data, class Dist, class ProbFunc>
int most_probable_cluster(Data data, const em_data_t<Data, Dist>& dists, ProbFunc prob);

template <class Data, class Dist>
void print_dists(const em_data_t<Data, Dist>& dists);

struct UnivariateEstimator
{
    UnivariateGaussianDistribution operator()(DataIt<double> begin, DataIt<double> end)
    {
        double var = variance(begin, end);
        return UnivariateGaussianDistribution(mean(begin, end), std::max(var, 1e-10));
    }
};

struct MultivariateEstimator
{
    MultivariateGaussian operator()(DataIt<Point<float>> begin, DataIt<Point<float>> end)
    {
        Vector mean(2);
        mean.zeros();
        for(auto p : boost::make_iterator_range(begin, end))
        {
            mean[0] += p.x;
            mean[1] += p.y;
        }

        if(std::distance(begin, end) > 0)
        {
            mean /= std::distance(begin, end);
        }

        Matrix cov(2, 2);
        cov.zeros();
        Vector diff(2);
        for(auto p : boost::make_iterator_range(begin, end))
        {
            diff[0] = p.x - mean[0];
            diff[1] = p.y - mean[1];

            cov += diff * arma::trans(diff);
        }

        if(std::distance(begin, end) < 2)
        {
            cov.zeros();
            cov.diag().fill(1e-16);
        }
        else
        {
            cov /= std::distance(begin, end) - 1;
        }

        try
        {
            Matrix inv = arma::inv(cov);
        }
        catch(std::exception& e)
        {
            std::cerr << "Failed to take inverse: " << e.what() << " Matrix:\n" << cov;
            cov.zeros();
            cov.diag().fill(1e-16);
        }

        return MultivariateGaussian(mean, cov);
    }
};

inline double probability_univariate(double value, const UnivariateGaussianDistribution& mean)
{
    return mean.likelihood(value);
}

inline double probability_multivariate(const Point<float>& value , const MultivariateGaussian& mean)
{
    Vector vec(2);
    vec[0] = value.x;
    vec[1] = value.y;
    return mean.probability(vec);
}

clustering_result_t em_1d_linear(std::vector<double>::const_iterator begin,
                                 std::vector<double>::const_iterator end,
                                 int kMax,
                                 int maxIterations)
{
    int dataSize = std::distance(begin, end);
    kMax = std::min(kMax, dataSize);    // can't have more clusters than data!

    std::vector<em_data_t<double, UnivariateGaussianDistribution>> attemptedMeans;
    for(int k = 1; k <= kMax; ++k)
    {
        attemptedMeans.push_back(run_em<double, UnivariateGaussianDistribution>(begin,
                                                                                end,
                                                                                k,
                                                                                maxIterations,
                                                                                probability_univariate,
                                                                                UnivariateEstimator()));
    }

    auto maxProbIt = std::max_element(attemptedMeans.begin(), attemptedMeans.end(), [](const auto& lhs, const auto& rhs) {
        return lhs.totalProb < rhs.totalProb;
    });

    // Convert the em_data_t to a clustering_result_t
    clustering_result_t results;
    results.numClusters = maxProbIt->dists.size();
    results.clusterSizes = std::move(maxProbIt->clusterSizes);
    results.assignedCluster = std::move(maxProbIt->assignments);

#ifdef DEBUG_RESULTS
    std::cout << "INFO: em_1d_linear: Cluster results:\nNum clusters:" << results.numClusters << " Cluster sizes:\n";
    for(std::size_t n = 0; n < minMeanIt->dists.size(); ++n)
    {
        std::cout << std::setprecision(5) << std::setw(5) << minMeanIt->dists[n] << "->" << results.clusterSizes[n] << '\n';
    }
#endif

    return results;
}

clustering_result_t em_2d_linear(std::vector<Point<float>>::const_iterator begin,
                                 std::vector<Point<float>>::const_iterator end,
                                 const int kMax,
                                 const int maxIterations)
{
    std::vector<em_data_t<Point<float>, MultivariateGaussian>> clusters;

    for(int k = 1; k < kMax; ++k)
    {
        clusters.push_back(run_em<Point<float>, MultivariateGaussian>(begin,
                                                                      end,
                                                                      k,
                                                                      maxIterations,
                                                                      probability_multivariate,
                                                                      MultivariateEstimator()));
    }

    auto maxProbIt = std::max_element(clusters.begin(), clusters.end(), [](const auto& lhs, const auto& rhs) {
        return lhs.totalProb < rhs.totalProb;
    });

    // Convert the em_data_t to a clustering_result_t
    clustering_result_t results;
    results.numClusters = maxProbIt->dists.size();
    results.clusterSizes = std::move(maxProbIt->clusterSizes);
    results.assignedCluster = std::move(maxProbIt->assignments);

#ifdef DEBUG_RESULTS
    std::cout << "INFO: em_2d_fixed: Cluster results:\nNum clusters:" << results.numClusters << " Cluster sizes:\n";
    for(std::size_t n = 0; n < dists.dists.size(); ++n)
    {
        std::cout << std::setprecision(5) << std::setw(5) << dists.dists[n] << "->" << results.clusterSizes[n] << '\n';
    }
#endif

    return results;
}

template <class Data, class Dist, class ProbFunc, class Estimator>
em_data_t<Data, Dist> run_em(DataIt<Data> begin,
                             DataIt<Data> end,
                             int k,
                             int maxIterations,
                             ProbFunc prob,
                             Estimator estimator)
{
    em_data_t<Data, Dist> dists;
    dists.numClusters = k;
    initialize_dists(begin, end, k, dists, estimator);
    calculate_dists(begin, end, dists, estimator);

    int numChanges = 0;
    int numIterations = 0;

    do
    {
        numChanges = assign_dists(begin, end, dists, prob);
        calculate_dists(begin, end, dists, estimator);
        ++numIterations;

#ifdef DEBUG_KMEANS
        std::cout << "Iteration:" << numIterations << " Changes:" << numChanges << '\n';
        print_dists(dists);
#endif

    } while((numChanges > 0) && (numIterations <= maxIterations));

    // Once complete, sum the error amongst all the dists
    for(std::size_t n = 0, size = std::distance(begin, end); n < size; ++n)
    {
        dists.totalProb += prob(*(begin + n), dists.dists[dists.assignments[n]]);
    }

    // If any cluster is empty, zero probability of correctness
    for(std::size_t n = 0; n < dists.clusterSizes.size(); ++n)
    {
        if(dists.clusterSizes[n] == 0)
        {
            dists.totalProb = 0.0;
        }
    }

#ifdef DEBUG_KMEANS
    std::cout << "INFO: em_1d: k:" << k << " Error:" << dists.totalProb << " Num iterations:" << numIterations
        << '\n';
    print_dists(dists);
#endif

    return dists;
}

template <class Data, class Dist, class Estimator>
void initialize_dists(DataIt<Data> begin, DataIt<Data> end, int k, em_data_t<Data, Dist>& dists, Estimator estimator)
{
    // Allocate the buffers
    dists.clusterSizes.resize(k);
    dists.dists.resize(k);
    dists.assignments.resize(std::distance(begin, end));
    std::iota(dists.assignments.begin(), dists.assignments.end(), 0);
    dists.totalProb = 0.0;

    std::transform(dists.assignments.begin(), dists.assignments.end(), dists.assignments.begin(), [k](int c) {
        return c % k;
    });
}

template <class Data, class Dist, class Estimator>
void calculate_dists(DataIt<Data> begin, DataIt<Data> end, em_data_t<Data, Dist>& dists, Estimator estimator)
{
    for(int cluster = 0; cluster < dists.numClusters; ++cluster)
    {
        dists.clusterData.clear();

        for(std::size_t n = 0; n < dists.assignments.size(); ++n)
        {
            if(dists.assignments[n] == cluster)
            {
                dists.clusterData.push_back(*(begin + n));
            }
        }

        dists.dists[cluster] = estimator(dists.clusterData.begin(), dists.clusterData.end());
    }
}

template <class Data, class Dist, class ProbFunc>
int assign_dists(DataIt<Data> begin, DataIt<Data> end, em_data_t<Data, Dist>& dists, ProbFunc prob)
{
    std::fill(dists.clusterSizes.begin(), dists.clusterSizes.end(), 0);

    int numChanged = 0;

    for(std::size_t n = 0, size = std::distance(begin, end); n < size; ++n)
    {
        int newAssignment = most_probable_cluster(*(begin + n), dists, prob);
        if(newAssignment != dists.assignments[n])
        {
            dists.assignments[n] = newAssignment;
            ++numChanged;
        }

        ++dists.clusterSizes[newAssignment];
    }

    return numChanged;
}

template <class Data, class Dist, class ProbFunc>
int most_probable_cluster(Data data, const em_data_t<Data, Dist>& dists, ProbFunc prob)
{
    auto maxIt = std::max_element(dists.dists.begin(), dists.dists.end(), [&](auto& lhs, auto& rhs) {
        return prob(data, lhs) < prob(data, rhs);
    });

    return std::distance(dists.dists.begin(), maxIt);
}

template <class Data, class Dist>
void print_dists(const em_data_t<Data, Dist>& dists)
{
    for(std::size_t n = 0; n < dists.dists.size(); ++n)
    {
        std::cout << std::setprecision(5) << std::setw(5) << dists.dists[n] << "->" << dists.clusterSizes[n] << '\n';
    }
}

} // namespace utils
} // namespace vulcan


// SPDX-FileCopyrightText: 2015 - 2021 Marcin Łoś <marcin.los.91@gmail.com>
// SPDX-License-Identifier: MIT

#ifndef ADS_SIMULATION_BASIC_SIMULATION_3D_HPP
#define ADS_SIMULATION_BASIC_SIMULATION_3D_HPP

#include <array>
#include <cstddef>

#include <boost/range/counting_range.hpp>

#include "ads/lin/tensor.hpp"
#include "ads/simulation/boundary.hpp"
#include "ads/simulation/dimension.hpp"
#include "ads/util/function_value.hpp"
#include "ads/util/iter/product.hpp"

namespace ads {

class basic_simulation_3d {
public:
    virtual ~basic_simulation_3d() = default;

    basic_simulation_3d() = default;
    basic_simulation_3d(const basic_simulation_3d&) = delete;
    basic_simulation_3d& operator=(const basic_simulation_3d&) = delete;
    basic_simulation_3d(basic_simulation_3d&&) = delete;
    basic_simulation_3d& operator=(basic_simulation_3d&&) = delete;

protected:
    using vector_type = lin::tensor<double, 3>;
    using vector_view = lin::tensor_view<double, 3>;
    using value_type = function_value_3d;

    using index_type = std::array<int, 3>;
    using index_1d_iter_type = boost::counting_iterator<int>;
    using index_iter_type = util::iter_product3<index_1d_iter_type, index_type>;
    using index_range = boost::iterator_range<index_iter_type>;

    using point_type = std::array<double, 3>;

    struct L2 {
        double operator()(value_type a) const { return a.val * a.val; }
    };

    struct H10 {
        double operator()(value_type a) const { return a.dx * a.dx + a.dy * a.dy + a.dz * a.dz; }
    };

    struct H1 {
        double operator()(value_type a) const {
            return a.val * a.val + a.dx * a.dx + a.dy * a.dy + a.dz * a.dz;
        }
    };

    value_type eval_basis(index_type e, index_type q, index_type a, const dimension& x,
                          const dimension& y, const dimension& z) const {
        auto loc = dof_global_to_local(e, a, x, y, z);

        const auto& bx = x.basis;
        const auto& by = y.basis;
        const auto& bz = z.basis;

        double B1 = bx.b[e[0]][q[0]][0][loc[0]];
        double B2 = by.b[e[1]][q[1]][0][loc[1]];
        double B3 = bz.b[e[2]][q[2]][0][loc[2]];
        double dB1 = bx.b[e[0]][q[0]][1][loc[0]];
        double dB2 = by.b[e[1]][q[1]][1][loc[1]];
        double dB3 = bz.b[e[2]][q[2]][1][loc[2]];

        double v = B1 * B2 * B3;
        double dxv = dB1 * B2 * B3;
        double dyv = B1 * dB2 * B3;
        double dzv = B1 * B2 * dB3;

        return {v, dxv, dyv, dzv};
    }

    double laplacian(index_type e, index_type q, index_type a, const dimension& x,
                     const dimension& y, const dimension& z) const {
        auto loc = dof_global_to_local(e, a, x, y, z);

        const auto& bx = x.basis;
        const auto& by = y.basis;
        const auto& bz = z.basis;

        double B1 = bx.b[e[0]][q[0]][0][loc[0]];
        double B2 = by.b[e[1]][q[1]][0][loc[1]];
        double B3 = bz.b[e[2]][q[2]][0][loc[2]];

        double ddB1 = bx.b[e[0]][q[0]][2][loc[0]];
        double ddB2 = by.b[e[1]][q[1]][2][loc[1]];
        double ddB3 = by.b[e[2]][q[2]][2][loc[2]];

        return ddB1 * B2 * B3 + B1 * ddB2 * B3 + B1 * B2 * ddB3;
    }

    template <typename Sol>
    value_type eval(const Sol& v, index_type e, index_type q, const dimension& x,
                    const dimension& y, const dimension& z) const {
        value_type u{};
        for (auto b : dofs_on_element(e, x, y, z)) {
            double c = v(b[0], b[1], b[2]);
            value_type B = eval_basis(e, q, b, x, y, z);
            u += c * B;
        }
        return u;
    }

    index_range elements(const dimension& x, const dimension& y, const dimension& z) const {
        return util::product_range<index_type>(x.element_indices(), y.element_indices(),
                                               z.element_indices());
    }

    index_range quad_points(const dimension& x, const dimension& y, const dimension& z) const {
        auto rx = boost::counting_range(0, x.basis.quad_order);
        auto ry = boost::counting_range(0, y.basis.quad_order);
        auto rz = boost::counting_range(0, z.basis.quad_order);
        return util::product_range<index_type>(rx, ry, rz);
    }

    index_range dofs_on_element(index_type e, const dimension& x, const dimension& y,
                                const dimension& z) const {
        auto rx = x.basis.dof_range(e[0]);
        auto ry = y.basis.dof_range(e[1]);
        auto rz = z.basis.dof_range(e[2]);
        return util::product_range<index_type>(rx, ry, rz);
    }

    index_range elements_supporting_dof(index_type dof, const dimension& x, const dimension& y,
                                        const dimension& z) const {
        auto rx = x.basis.element_range(dof[0]);
        auto ry = y.basis.element_range(dof[1]);
        auto rz = z.basis.element_range(dof[2]);
        return util::product_range<index_type>(rx, ry, rz);
    }

    bool supported_in(index_type dof, index_type e, const dimension& x, const dimension& y,
                      const dimension& z) const {
        auto xrange = x.basis.element_ranges[dof[0]];
        auto yrange = y.basis.element_ranges[dof[1]];
        auto zrange = z.basis.element_ranges[dof[2]];

        return e[0] >= xrange.first && e[0] <= xrange.second && e[1] >= yrange.first
            && e[1] <= yrange.second && e[2] >= zrange.first && e[2] <= zrange.second;
    }

    index_type dof_global_to_local(index_type e, index_type a, const dimension& x,
                                   const dimension& y, const dimension& z) const {
        const auto& bx = x.basis;
        const auto& by = y.basis;
        const auto& bz = z.basis;
        return {{a[0] - bx.first_dof(e[0]), a[1] - by.first_dof(e[1]), a[2] - bz.first_dof(e[2])}};
    }

    template <typename RHS>
    void update_global_rhs(RHS& global, const vector_type& local, index_type e, const dimension& x,
                           const dimension& y, const dimension& z) const {
        for (auto a : dofs_on_element(e, x, y, z)) {
            auto loc = dof_global_to_local(e, a, x, y, z);
            global(a[0], a[1], a[2]) += local(loc[0], loc[1], loc[2]);
        }
    }

    index_range dofs(const dimension& x, const dimension& y, const dimension& z) const {
        auto rx = boost::counting_range(0, x.dofs());
        auto ry = boost::counting_range(0, y.dofs());
        auto rz = boost::counting_range(0, z.dofs());
        return util::product_range<index_type>(rx, ry, rz);
    }

    index_range internal_dofs(const dimension& x, const dimension& y, const dimension& z) const {
        auto rx = boost::counting_range(1, x.dofs() - 1);
        auto ry = boost::counting_range(1, y.dofs() - 1);
        auto rz = boost::counting_range(1, z.dofs() - 1);
        return util::product_range<index_type>(rx, ry, rz);
    }

    double jacobian(index_type e, const dimension& x, const dimension& y,
                    const dimension& z) const {
        return x.basis.J[e[0]] * y.basis.J[e[1]] * z.basis.J[e[2]];
    }

    double weight(index_type q, const dimension& x, const dimension& y, const dimension& z) const {
        return x.basis.w[q[0]] * y.basis.w[q[1]] * z.basis.w[q[2]];
    }

    point_type point(index_type e, index_type q, const dimension& x, const dimension& y,
                     const dimension& z) const {
        double px = x.basis.x[e[0]][q[0]];
        double py = y.basis.x[e[1]][q[1]];
        double pz = z.basis.x[e[2]][q[2]];
        return {px, py, pz};
    }

    auto overlapping_dofs(int dof, int begin, int end, const dimension& x) const {
        using std::max;
        using std::min;

        auto minx = max(begin, dof - x.B.degree);
        auto maxx = min(end, dof + x.B.degree + 1);

        return boost::counting_range(minx, maxx);
    }

    index_range overlapping_dofs(index_type dof, const dimension& x, const dimension& y,
                                 const dimension& z) const {
        auto rx = overlapping_dofs(dof[0], 0, x.dofs(), x);
        auto ry = overlapping_dofs(dof[1], 0, y.dofs(), y);
        auto rz = overlapping_dofs(dof[1], 0, z.dofs(), z);
        return util::product_range<index_type>(rx, ry, rz);
    }

    index_range overlapping_dofs(index_type dof, const dimension& Ux, const dimension& Uy,
                                 const dimension& Uz, const dimension& Vx, const dimension& Vy,
                                 const dimension& Vz) const {
        auto xrange = Ux.basis.element_ranges[dof[0]];
        auto yrange = Uy.basis.element_ranges[dof[1]];
        auto zrange = Uz.basis.element_ranges[dof[2]];

        auto x0 = Vx.basis.first_dof(xrange.first);
        auto x1 = Vx.basis.last_dof(xrange.second) + 1;

        auto y0 = Vy.basis.first_dof(yrange.first);
        auto y1 = Vy.basis.last_dof(yrange.second) + 1;

        auto z0 = Vz.basis.first_dof(zrange.first);
        auto z1 = Vz.basis.last_dof(zrange.second) + 1;

        auto rx = boost::counting_range(x0, x1);
        auto ry = boost::counting_range(y0, y1);
        auto rz = boost::counting_range(z0, z1);

        return util::product_range<index_type>(rx, ry, rz);
    }

    index_range overlapping_internal_dofs(index_type dof, const dimension& x, const dimension& y,
                                          const dimension& z) const {
        auto rx = overlapping_dofs(dof[0], 1, x.dofs() - 1, x);
        auto ry = overlapping_dofs(dof[1], 1, y.dofs() - 1, y);
        auto rz = overlapping_dofs(dof[2], 1, z.dofs() - 1, z);
        return util::product_range<index_type>(rx, ry, rz);
    }

    int linear_index(index_type dof, const dimension& x, const dimension& y,
                     const dimension& z) const {
        auto order = reverse_ordering<3>({x.dofs(), y.dofs(), z.dofs()});
        return order.linear_index(dof[0], dof[1], dof[2]);
    }

    template <typename Fun>
    void for_boundary_dofs(const dimension& x, const dimension& y, const dimension& z,
                           Fun&& fun) const {
        for (auto jx = 0; jx < x.dofs(); ++jx) {
            for (auto jy = 0; jy < y.dofs(); ++jy) {
                fun({jx, jy, 0});
                fun({jx, jy, z.dofs() - 1});
            }
        }
        for (auto jx = 0; jx < x.dofs(); ++jx) {
            for (auto jz = 1; jz < z.dofs() - 1; ++jz) {
                fun({jx, 0, jz});
                fun({jx, y.dofs() - 1, jz});
            }
        }
        for (auto jy = 1; jy < y.dofs() - 1; ++jy) {
            for (auto jz = 1; jz < z.dofs() - 1; ++jz) {
                fun({0, jy, jz});
                fun({x.dofs() - 1, jy, jz});
            }
        }
    }

    bool is_boundary(int dof, const dimension& x) const { return dof == 0 || dof == x.dofs() - 1; }

    bool is_boundary(index_type dof, const dimension& x, const dimension& y,
                     const dimension& z) const {
        return is_boundary(dof[0], x) || is_boundary(dof[1], y) || is_boundary(dof[2], z);
    }

    template <typename Norm, typename Fun>
    double norm(const dimension& Ux, const dimension& Uy, const dimension& Uz, Norm&& norm,
                Fun&& fun) const {
        double val = 0;

        for (auto e : elements(Ux, Uy, Uz)) {
            double J = jacobian(e, Ux, Uy, Uz);
            for (auto q : quad_points(Ux, Uy, Uz)) {
                double w = weight(q, Ux, Uy, Uz);
                auto x = point(e, q, Ux, Uy, Uz);
                auto d = fun(x);
                val += norm(d) * w * J;
            }
        }
        return std::sqrt(val);
    }

    template <typename Fun>
    double normL2(const dimension& Ux, const dimension& Uy, const dimension& Uz, Fun&& fun) const {
        return norm(Ux, Uy, Uz, L2{}, fun);
    }

    template <typename Fun>
    double normH1(const dimension& Ux, const dimension& Uy, const dimension& Uz, Fun&& fun) const {
        return norm(Ux, Uy, Uz, H1{}, fun);
    }

    template <typename Sol, typename Norm>
    double norm(const Sol& u, const dimension& Ux, const dimension& Uy, const dimension& Uz,
                Norm&& norm) const {
        double val = 0;

        for (auto e : elements(Ux, Uy, Uz)) {
            double J = jacobian(e, Ux, Uy, Uz);
            for (auto q : quad_points(Ux, Uy, Uz)) {
                double w = weight(q, Ux, Uy, Uz);
                value_type uu = eval(u, e, q, Ux, Uy, Uz);
                val += norm(uu) * w * J;
            }
        }
        return std::sqrt(val);
    }

    template <typename Sol>
    double normL2(const Sol& u, const dimension& Ux, const dimension& Uy,
                  const dimension& Uz) const {
        return norm(u, Ux, Uy, Uz, L2{});
    }

    template <typename Sol>
    double normH1(const Sol& u, const dimension& Ux, const dimension& Uy,
                  const dimension& Uz) const {
        return norm(u, Ux, Uy, Uz, H1{});
    }

    template <typename Sol, typename Fun, typename Norm>
    double error(const Sol& u, const dimension& Ux, const dimension& Uy, const dimension& Uz,
                 Norm&& norm, Fun&& fun) const {
        double error = 0;

        for (auto e : elements(Ux, Uy, Uz)) {
            double J = jacobian(e, Ux, Uy, Uz);
            for (auto q : quad_points(Ux, Uy, Uz)) {
                double w = weight(q, Ux, Uy, Uz);
                auto x = point(e, q, Ux, Uy, Uz);
                value_type uu = eval(u, e, q, Ux, Uy, Uz);

                auto d = uu - fun(x);
                error += norm(d) * w * J;
            }
        }
        return std::sqrt(error);
    }

    template <typename Sol, typename Fun, typename Norm>
    double error_relative(const Sol& u, const dimension& Ux, const dimension& Uy,
                          const dimension& Uz, Norm&& norm, Fun&& fun) const {
        double error = 0;
        double ref_norm = 0;

        for (auto e : elements(Ux, Uy, Uz)) {
            double J = jacobian(e, Ux, Uy, Uz);
            for (auto q : quad_points(Ux, Uy, Uz)) {
                double w = weight(q, Ux, Uy, Uz);
                auto x = point(e, q, Ux, Uy, Uz);
                value_type uu = eval(u, e, q, Ux, Uy, Uz);
                auto fx = fun(x);

                error += norm(uu - fx) * w * J;
                ref_norm += norm(fx) * w * J;
            }
        }
        return std::sqrt(error / ref_norm);
    }

    template <typename Sol, typename Fun>
    double errorL2(const Sol& u, const dimension& Ux, const dimension& Uy, const dimension& Uz,
                   Fun&& fun) const {
        return error(u, Ux, Uy, Uz, L2{}, fun);
    }

    template <typename Sol, typename Fun>
    double error_relative_L2(const Sol& u, const dimension& Ux, const dimension& Uy,
                             const dimension& Uz, Fun&& fun) const {
        return error_relative(u, Ux, Uy, Uz, L2{}, fun);
    }

    template <typename Sol, typename Fun>
    double errorH1(const Sol& u, const dimension& Ux, const dimension& Uy, const dimension& Uz,
                   Fun&& fun) const {
        return error(u, Ux, Uy, Uz, H1{}, fun);
    }

    template <typename Sol, typename Fun>
    double error_relative_H1(const Sol& u, const dimension& Ux, const dimension& Uy,
                             const dimension& Uz, Fun&& fun) const {
        return error_relative(u, Ux, Uy, Uz, H1{}, fun);
    }

    template <typename Sol>
    double norm_rot(const Sol& X, const Sol& Y, const Sol& Z,                          //
                    const dimension& U1x, const dimension& U1y, const dimension& U1z,  //
                    const dimension& U2x, const dimension& U2y, const dimension& U2z,  //
                    const dimension& U3x, const dimension& U3y, const dimension& U3z) const {
        double val = 0;

        for (auto e : elements(U1x, U1y, U1z)) {
            double J = jacobian(e, U1x, U1y, U1z);
            for (auto q : quad_points(U1x, U1y, U1z)) {
                double w = weight(q, U1x, U1y, U1z);
                auto x = eval(X, e, q, U1x, U1y, U1z);
                auto y = eval(Y, e, q, U2x, U2y, U2z);
                auto z = eval(Z, e, q, U3x, U3y, U3z);

                auto rx = z.dy - y.dz;
                auto ry = x.dz - z.dx;
                auto rz = y.dx - x.dy;

                val += (rx * rx + ry * ry + rz * rz) * w * J;
            }
        }
        return std::sqrt(val);
    }

    template <typename Sol, typename FunX, typename FunY, typename FunZ>
    double error_rot(const Sol& X, const Sol& Y, const Sol& Z,                          //
                     const dimension& U1x, const dimension& U1y, const dimension& U1z,  //
                     const dimension& U2x, const dimension& U2y, const dimension& U2z,  //
                     const dimension& U3x, const dimension& U3y, const dimension& U3z,  //
                     FunX&& fx, FunY&& fy, FunZ&& fz) const {
        double val = 0;

        for (auto e : elements(U1x, U1y, U1z)) {
            double J = jacobian(e, U1x, U1y, U1z);
            for (auto q : quad_points(U1x, U1y, U1z)) {
                auto p = point(e, q, U1x, U2y, U3z);
                double w = weight(q, U1x, U2y, U3z);
                auto x = eval(X, e, q, U1x, U1y, U1z) - fx(p);
                auto y = eval(Y, e, q, U2x, U2y, U2z) - fy(p);
                auto z = eval(Z, e, q, U3x, U3y, U3z) - fz(p);

                auto rx = z.dy - y.dz;
                auto ry = x.dz - z.dx;
                auto rz = y.dx - x.dy;

                val += (rx * rx + ry * ry + rz * rz) * w * J;
            }
        }
        return std::sqrt(val);
    }

    template <typename Sol>
    double norm_div(const Sol& X, const Sol& Y, const Sol& Z,                          //
                    const dimension& U1x, const dimension& U1y, const dimension& U1z,  //
                    const dimension& U2x, const dimension& U2y, const dimension& U2z,  //
                    const dimension& U3x, const dimension& U3y, const dimension& U3z) const {
        double val = 0;

        for (auto e : elements(U1x, U1y, U1z)) {
            double J = jacobian(e, U1x, U1y, U1z);
            for (auto q : quad_points(U1x, U1y, U1z)) {
                double w = weight(q, U1x, U1y, U1z);
                auto x = eval(X, e, q, U1x, U1y, U1z);
                auto y = eval(Y, e, q, U2x, U2y, U2z);
                auto z = eval(Z, e, q, U3x, U3y, U3z);

                auto v = x.dx + y.dy + z.dz;
                val += v * v * w * J;
            }
        }
        return std::sqrt(val);
    }
};

}  // namespace ads

#endif  // ADS_SIMULATION_BASIC_SIMULATION_3D_HPP


/*Copyright (c) 2021 James Gayvert

Permission is hereby granted, free of charge, to any person obtaining a copy
of this software and associated documentation files (the "Software"), to deal
in the Software without restriction, including without limitation the rights
to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
copies of the Software, and to permit persons to whom the Software is
furnished to do so, subject to the following conditions:

The above copyright notice and this permission notice shall be included in all
copies or substantial portions of the Software.

THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
SOFTWARE.
*/

/*
 * transforms.cpp
 *
 */

#include <math.h>
#include <algorithm>
#include <array>
#include <Eigen/Dense>
#include <vector>

#include "BasisSet.h"
#include "gto_ordering.h"
#include "Shell.h"
#include "utils.h"

void uniform_cart_norm(Eigen::MatrixXd &my_mat, BasisSet &bs)
{
	unsigned int bf_idx = 0;
	for(auto&shell:bs.basis)
	{
		std::vector<std::array<size_t,3>> order = opencap_carts_ordering(shell.l);
		for(unsigned int i=0;i<shell.num_carts();i++)
		{
			std::array<size_t,3> cart = order[i];
			double scale = sqrt(fact2(2*shell.l-1)/fact2(2*cart[0]-1)
					/fact2(2*cart[1]-1)/fact2(2*cart[2]-1));
			my_mat.col(bf_idx+i) = my_mat.col(bf_idx+i)*scale;
			my_mat.row(bf_idx+i) = my_mat.row(bf_idx+i)*scale;
		}
		bf_idx += shell.num_carts();
	}
}

//DOI: 10.1002/qua.560540202
//Last part of eqn 15, for M<0 we want imaginary part, for M>0 we want real part
double term4(int M, int exp_num)
{
	if(M<0)
	{
		//integer powers will be real
		if(exp_num%2==0)
			return 0;
		else
			return (exp_num-1)%4==0 ? 1:-1;
	}
	else
	{
		//half-integer powers will be imaginary
		if(exp_num%2!=0)
			return 0;
		else
			return exp_num%4==0 ? 1:-1;
	}
}

//DOI: 10.1002/qua.560540202
//Equation 15, currently assumes fully normalized cartesians. Will add non-fully normalized carts soon enough...
double get_coeff(int L, int m, int lx, int ly, int lz)
{
      auto abs_m = std::abs(m);
      if ((lx + ly - abs_m)%2)
        return 0.0;
      auto j = (lx + ly - abs_m)/2;
      if (j < 0)
        return 0.0;
      auto term1 = sqrt( fact(2*lx)*fact(2*ly)*fact(2*lz)*fact(L)*fact(L-abs_m)
    		  	  	  	  /fact(2*L)/fact(lx)/fact(ly)/fact(lz)/fact(L+abs_m));
      term1/=fact(L);
      term1/=pow(2,L);
      auto term2 = 0;
      for(int i=0;i<=(L-abs_m)/2;i++)
      {
    	  term2+=binom(L,i)*binom(i,j)
    			  *parity(i) * fact(2*(L-i))
				  /fact(L-abs_m-2*i);
      }
      double term3=0;
      for (int k=0;k<=j;k++)
    	  term3+=binom(j,k)*binom(abs_m,lx-2*k)*term4(m,abs_m - lx +2*k);
      // for m!=0, real solid harmonics are linear combinations of complex ones
      // R+(l,m) = ( Y(l,m) + Y(l,-m) )/ sqrt(2) ;  R-(l,m) = ( Y(l,m) - Y(l,-m) )/ sqrt(-2)
      double result = (m == 0) ? term1*term2*term3 : M_SQRT2*term1*term2*term3;
      return result;
}

Eigen::MatrixXd get_trans_mat(Shell &shell)
{
	std::vector<std::array<size_t,3>> cart_order = opencap_carts_ordering(shell.l);
	std::vector<int> sph_order = opencap_harmonic_ordering(shell.l);
	Eigen::MatrixXd trans_mat(shell.num_bf(),shell.num_carts());
	for(size_t i=0;i<shell.num_bf();i++)
	{
		int M = sph_order[i];
		for(size_t j=0;j<shell.num_carts();j++)
		{
			std::array<size_t,3> cart = cart_order[j];
			trans_mat(i,j) = get_coeff(shell.l,M,cart[0],cart[1],cart[2]);
		}
	}
	return trans_mat;
}

Eigen::MatrixXd transform_block(Shell &shell1, Shell &shell2, Eigen::MatrixXd cart_block)
{

	if(!shell1.pure && !shell2.pure)
		return cart_block;
	else if(shell1.pure && !shell2.pure)
		return get_trans_mat(shell1)*cart_block;
	else if(shell1.pure && shell2.pure)
		return get_trans_mat(shell1)*cart_block*get_trans_mat(shell2).transpose();
	else
		return cart_block*get_trans_mat(shell2).transpose();
}

void cart2spherical(Eigen::MatrixXd &cart_ints, Eigen::MatrixXd &spherical_ints, BasisSet &bs)
{
	//indices for first basis function for cart and spherical matrices
	unsigned int cart_row_idx = 0;
	unsigned int sph_row_idx = 0;
	for(auto shell1:bs.basis)
	{
		//indices for 2nd basis function for cart and spherical matrices
		unsigned int cart_col_idx = 0;
		unsigned int sph_col_idx = 0;
		for(auto shell2:bs.basis)
		{
			Eigen::MatrixXd cart_block = cart_ints.block(cart_row_idx,cart_col_idx, shell1.num_carts(),shell2.num_carts());
			spherical_ints.block(sph_row_idx,sph_col_idx,shell1.num_bf(),shell2.num_bf())
							= transform_block(shell1,shell2,cart_block);
			cart_col_idx+=shell2.num_carts();
			sph_col_idx+=shell2.num_bf();
		}
		cart_row_idx+=shell1.num_carts();
		sph_row_idx+=shell1.num_bf();
	}
}


// This file is part of libigl, a simple c++ geometry processing library.
//
// Copyright (C) 2016 Alec Jacobson <alecjacobson@gmail.com>
//
// This Source Code Form is subject to the terms of the Mozilla Public License
// v. 2.0. If a copy of the MPL was not distributed with this file, You can
// obtain one at http://mozilla.org/MPL/2.0/.
#include "bbw.h"
#include "mosek_quadprog.h"
#include "../harmonic.h"
#include "../slice_into.h"
#include <Eigen/Sparse>
#include <iostream>
#include <cstdio>

template <
  typename DerivedV,
  typename DerivedEle,
  typename Derivedb,
  typename Derivedbc,
  typename DerivedW>
IGL_INLINE bool igl::mosek::bbw(
  const Eigen::PlainObjectBase<DerivedV> & V,
  const Eigen::PlainObjectBase<DerivedEle> & Ele,
  const Eigen::PlainObjectBase<Derivedb> & b,
  const Eigen::PlainObjectBase<Derivedbc> & bc,
  igl::BBWData & data,
  igl::mosek::MosekData & mosek_data,
  Eigen::PlainObjectBase<DerivedW> & W
  )
{
  using namespace std;
  using namespace Eigen;
  assert(!data.partition_unity && "partition_unity not implemented yet");
  // number of domain vertices
  int n = V.rows();
  // number of handles
  int m = bc.cols();
  // Build biharmonic operator
  Eigen::SparseMatrix<typename DerivedV::Scalar> Q;
  harmonic(V,Ele,2,Q);
  W.derived().resize(n,m);
  // No linear terms
  VectorXd c = VectorXd::Zero(n);
  // No linear constraints
  SparseMatrix<typename DerivedW::Scalar> A(0,n);
  VectorXd uc(0,1),lc(0,1);
  // Upper and lower box constraints (Constant bounds)
  VectorXd ux = VectorXd::Ones(n);
  VectorXd lx = VectorXd::Zero(n);
  // Loop over handles
  for(int i = 0;i<m;i++)
  {
    if(data.verbosity >= 1)
    {
      cout<<"BBW: Computing weight for handle "<<i+1<<" out of "<<m<<
        "."<<endl;
    }
    VectorXd bci = bc.col(i);
    VectorXd Wi;
    // impose boundary conditions via bounds
    slice_into(bci,b,ux);
    slice_into(bci,b,lx);
    bool r = mosek_quadprog(Q,c,0,A,lc,uc,lx,ux,mosek_data,Wi);
    if(!r)
    {
      return false;
    }
    W.col(i) = Wi;
  }
#ifndef NDEBUG
    const double min_rowsum = W.rowwise().sum().array().abs().minCoeff();
    if(min_rowsum < 0.1)
    {
      cerr<<"bbw.cpp: Warning, minimum row sum is very low. Consider more "
        "active set iterations or enforcing partition of unity."<<endl;
    }
#endif

  return true;
}

#ifdef IGL_STATIC_LIBRARY
// Explicit template instantiation
template bool igl::mosek::bbw<Eigen::Matrix<double, -1, -1, 0, -1, -1>, Eigen::Matrix<int, -1, -1, 0, -1, -1>, Eigen::Matrix<int, -1, 1, 0, -1, 1>, Eigen::Matrix<double, -1, -1, 0, -1, -1>, Eigen::Matrix<double, -1, -1, 0, -1, -1> >(Eigen::PlainObjectBase<Eigen::Matrix<double, -1, -1, 0, -1, -1> > const&, Eigen::PlainObjectBase<Eigen::Matrix<int, -1, -1, 0, -1, -1> > const&, Eigen::PlainObjectBase<Eigen::Matrix<int, -1, 1, 0, -1, 1> > const&, Eigen::PlainObjectBase<Eigen::Matrix<double, -1, -1, 0, -1, -1> > const&, igl::BBWData&, igl::mosek::MosekData&, Eigen::PlainObjectBase<Eigen::Matrix<double, -1, -1, 0, -1, -1> >&);
#endif


#ifdef _MAIN_

#else
#include <armadillo>
#include "../core/UAmoeba.hpp"
#ifdef _ALGEBRA_TOOLS_
#else
#include "../core/algebraTools.hpp"
#endif
#ifdef _AMOEBA_INIT_
#else
#include "../core/amoebaParam.cpp"
#endif

#endif

using namespace arma;

template<typename T>
void serialSolver(const AmoebaParam<T> amoebaParam)
{

	arma_rng::set_seed_random();

	vector<vec> newGuess;

	newGuess.resize(amoebaParam.nGridPoints);
	for(int i=0; i<amoebaParam.nGridPoints; ++i)
	{
		newGuess[i] = randu<vec>(amoebaParam.lieDimension);
	}

	UAmoeba* amoeba = new UAmoeba(amoebaParam.maxAmoebaIters,
																amoebaParam.nGridPoints,
																amoebaParam.precision,
																amoebaParam.matSize,
																amoebaParam.lieDimension,
																amoebaParam.basis);

	amoeba->solver(newGuess, amoebaParam.startBoundary, amoebaParam.endBoundary);
	amoeba->curvePrint();

}


// Boost.Geometry (aka GGL, Generic Geometry Library)

// Copyright (c) 2019 Tinko Bartels, Berlin, Germany.

// Use, modification and distribution is subject to the Boost Software License,
// Version 1.0. (See accompanying file LICENSE_1_0.txt or copy at
// http://www.boost.org/LICENSE_1_0.txt)

#ifndef BOOST_GEOMETRY_EXTENSIONS_RANDOM_STRATEGIES_SPHERICAL_EDWILLIAMS_AVFORM_INTERMEDIATE_HPP
#define BOOST_GEOMETRY_EXTENSIONS_RANDOM_STRATEGIES_SPHERICAL_EDWILLIAMS_AVFORM_INTERMEDIATE_HPP

#include <random>
#include <cmath>

#include <boost/geometry/algorithms/equals.hpp>
#include <boost/geometry/arithmetic/arithmetic.hpp>

#include <boost/geometry/extensions/random/strategies/uniform_point_distribution.hpp>

namespace boost { namespace geometry
{

namespace strategy { namespace uniform_point_distribution {

template
<
    typename Point,
    typename DomainGeometry
>
struct edwilliams_avform_intermediate
{
    edwilliams_avform_intermediate(DomainGeometry const& g) {}
    bool equals(DomainGeometry const& l_domain,
                DomainGeometry const& r_domain,
                edwilliams_avform_intermediate const& r_strategy) const
    {
        return boost::geometry::equals(l_domain.domain(), r_domain.domain());
    }

    // The following implementation is adapted from
    // https://www.edwilliams.org/avform.htm#Intermediate
    template
    <
        typename LengthType,
        typename PointIn
    >
    static Point map(PointIn const& p1, PointIn const& p2, LengthType const& f)
    {
        Point out;
        const auto lat1 = get_as_radian<1>(p1);
        const auto lon1 = get_as_radian<0>(p1);
        const auto lat2 = get_as_radian<1>(p2);
        const auto lon2 = get_as_radian<0>(p2);
        LengthType const d = std::acos(
              std::sin(lat1) * std::sin(lat2)
            + std::cos(lat1) * std::cos(lat2) * std::cos(lon1 - lon2));
        LengthType const A = std::sin( ( 1 - f ) * d ) / std::sin(d);
        LengthType const B = std::sin( f * d ) / std::sin( d );
        LengthType const x = A * std::cos(lat1) * std::cos(lon1)
                           + B * std::cos(lat2) * std::cos(lon2);
        LengthType const y = A * std::cos(lat1) * std::sin(lon1)
                           + B * std::cos(lat2) * std::sin(lon2);
        LengthType const z = A * std::sin(lat1) + B * std::sin(lat2);
        LengthType const lat = std::atan2(z, std::sqrt(x * x + y * y));
        LengthType const lon = std::atan2(y, x);
        set_from_radian<1>(out, lat);
        set_from_radian<0>(out, lon);
        return out;
    }

    template<typename Gen>
    Point apply(Gen& g, DomainGeometry const& d)
    {
        typedef typename select_most_precise
            <
                typename coordinate_type<DomainGeometry>::type,
                double
            >::type sample_type;
        std::uniform_real_distribution<sample_type> real_dist(0, 1);
        return map<sample_type>(d, real_dist(g));
    }
    void reset(DomainGeometry const&) {};
};

namespace services {

template
<
    typename Point,
    typename DomainGeometry
>
struct default_strategy
<
    Point,
    DomainGeometry,
    segment_tag,
    single_tag,
    2,
    spherical_tag
> : public edwilliams_avform_intermediate<Point, DomainGeometry> {
    typedef edwilliams_avform_intermediate<Point, DomainGeometry> base;
    using base::base;
};

} // namespace services

}} // namespace strategy::uniform_point_distribution

}} // namespace boost::geometry

#endif // BOOST_GEOMETRY_EXTENSIONS_RANDOM_STRATEGIES_SPHERICAL_EDWILLIAMS_AVFORM_INTERMEDIATE_HPP


#define _USE_MATH_DEFINES

#include <fstream>
#include <iostream>
#include <omp.h>
#include <cmath>
#include <complex>
#include <cstdio>
#include <chrono>
#include <string>
#include <armadillo>
#include <vector>
#include "wignerSymbols.h"
//#include "wigner/gaunt.hpp"

//using namespace std::chrono;
using namespace arma;

static const double kB = 1.3806504e-23;         // J/K
static const double NA = 6.02214179e23;         // 1/mol
static const double EHARTREE = 4.35974434e-18;  // J/Hartree
static const double AMU = 1.660538921e-27;      // kg/amu
static const double HBAR = 1.054571726e-34;     // J.s
static const double HBAR1 = HBAR / EHARTREE;    // Hartree.s
static const double HBAR2 = HBAR * 1e20 / AMU;  // amu.Å^2/s
static const double SCH4 = 186.25; // J/mol.K

Mat<complex<double>> getHamiltonian(int lmax, double Ix, double Iy, double Iz, std::vector<double>& a);
SpMat<complex<double>> getSparseHam(int lmax, double Ix, double Iy, double Iz, std::vector<double>& a);
//Col<double> getCoefficients(std::string sysname);
std::vector<double> getWignerCoeffs(std::string sysname, bool freeRotor);
std::vector<double> getMomentOfInertia(std::string sysname="METH-CHA");
std::string getDirectory(std::string sysname);
double getPartitionFunction(double T, Mat<complex<double>>& H, int sym=1);
double getSparseQ(double T, SpMat<complex<double>>& H, int sym=1);

int main(int argc, char** argv) {
    /* argv[0]  Program name
     * argv[1]  System name
     * argv[2]  Lmax (start)
     * argv[3]  free-rotor
     * argv[4]  temperature / K
     */
    if (argc != 5) {
        throw std::runtime_error("Enter systemName as it appears in data directory, Lmax, free-rotor flag, temperature in Kelvin");
    }

    auto start = std::chrono::high_resolution_clock::now();

    std::string sysname = argv[1];
    bool freeRotor = std::stoi(argv[3]);
    double T = std::stod(argv[4]);
    std::string dirname = getDirectory(sysname);
    std::cout << "Directory containing data is: " << dirname << std::endl;
    std::vector<double> a = getWignerCoeffs(sysname, freeRotor);
    int sigma = 1;

    std::vector<double> Ivec = getMomentOfInertia(sysname);
    std::cout << "Moments of Inertia:" << std::endl;
    for (double i : Ivec) {
        std::cout << i << '\t';
    }

    /*
     *  Rotational Constants + Classical Partition Function
     */
    double B = HBAR1*HBAR2/(2.0*Ivec[2]);
    double A = HBAR1*HBAR2/(2.0*Ivec[1]);
    double C = HBAR1*HBAR2/(2.0*Ivec[0]);
    //std::cout << "kB T / Hartree:\n" << kB * 298 / EHARTREE << std::endl;
    std::cout << "Qapprox = " << sqrt(M_PI)/sigma * sqrt(pow(kB*300/EHARTREE, 3) / (A*B*C)) << std::endl;

    /*
     *  Dense Matrix Implementation
     */
    bool converge = false;
    bool dense = true;
    int lmax = atoi(argv[2]);
    double Q;
    double Qprev = 0;
    std::cout << "Lmax\t\tQ\t\tTime\t\tdQ" << std::endl;
    do {
        auto qstart = std::chrono::high_resolution_clock::now();
        if (dense) {
            /****  Dense Matrix Implementation  ****/
            unsigned long long size = (lmax+1)*(2*lmax+1)*(2*lmax+3)/3.0;
            Mat<complex<double>> H = getHamiltonian(lmax, Ivec[0], Ivec[1], Ivec[2], a);
            if (!H.is_hermitian(1e-3)) {
                H.brief_print("H = ");
            }
            try {
                Q = getPartitionFunction(298.15, H, sigma);
            } catch (const std::logic_error& e) {
                //std::cout << "Required Memory is " << size*size*8*1e-9 << " Gb." << std::endl;
                //std::cout << "Memory problem caused failure. Switching to sparse implementation." << std::endl;
                //dense = false;
                throw e;
                break;
            }
            /* *********************************** */
        } else {
            break;
            /****  Sparse Matrix Implementation  ****/
            SpMat<complex<double>> H = getSparseHam(lmax, Ivec[0], Ivec[1], Ivec[2], a);
            Q = getSparseQ(298.15, H, sigma);
            /****************************************/
        }
        auto qend = std::chrono::high_resolution_clock::now();
        auto qduration = std::chrono::duration_cast<std::chrono::microseconds>(qend - qstart);
        double dQ = fabs(Q-Qprev);
        std::cout << lmax << "\t\t" << Q << "\t\t" << qduration.count()/1e6 << " sec." << "\t\t" << dQ << std::endl;
        if (dQ < 1e-4) {
            converge = true;
            std::cout << "DeltaQ = " << fabs(Q-Qprev) << std::endl;
            std::cout << "Convergence criterion met!" << std::endl;
            std::cout << "Lmax = " << lmax << std::endl;
        }
        lmax++;
        Qprev = Q;
    } while (!converge);
    std::cout << std::endl;
    std::cout << std::endl;
    auto stop = std::chrono::high_resolution_clock::now();
    auto duration = std::chrono::duration_cast<std::chrono::microseconds>(stop - start);
    std::cout << std::endl << duration.count()/1e6 << " secs" << std::endl;
    return 0;
}

double getPartitionFunction(double T, Mat<complex<double>>& H, int sym) {
    /* Solve the Eigenvalues
     * Inputs:  T;  the temperature [=] K
     *          H;  the Hamiltonian matrix [=] Hartree
     */
    double b = pow(kB * T, -1) * EHARTREE;
    double tr = 0;

    //std::cout << "Solving Matrix exponential" << std::endl;
    //Mat<complex<double>> expbH = expmat_sym(-b*H);
    vec eigval = eig_sym(H);
    for (double e : eigval) {
        tr += exp(-b*e);
    }
    //std::cout  <<  "Solved Matrix exponential" << std::endl;
    //double tr = trace(expmat_sym(-b*H));
    //std::cout << std::endl << "Q predicted by eig_sym: " << tr/sym << std::endl;
    return double(tr/sym);
    //return double(Q/sym);
}

Mat<complex<double>> getHamiltonian(int lmax, double Ix, double Iy, double Iz, std::vector<double>& a) {
    /* Construct the Hamiltonian Matrix
     * Inputs:  lmax;   the maximum quantum number forthe spherical basis
     *             I;   the gas-phase moments of inertia [=] amu*Å^2
     *             a;   coefficients for potential in the Wigner D Matrix basis
     * Outputs:    H;   the Hamiltonian matrix (Hermitian) [=] Hartree
     */
    long double size = (lmax+1)*(2*lmax+1)*(2*lmax+3)/3.0;
    Mat<complex<double>> H = zeros<mat>(size, size);

    // Define rotational constants:
    double B = HBAR1*HBAR2/(2.0*Iz);
    double A = HBAR1*HBAR2/(2.0*Iy);
    double C = HBAR1*HBAR2/(2.0*Ix);
    double kap = (A == B && A == C) ? 0 : (2.0*B - (A+C)) / (A-C);
    #pragma omp parallel
    {
        unsigned long long i = 0;
        #pragma omp for
        for (int el = 0; el < lmax+1; el++) {
            for (int m = -el; m <= el; m++) {
                for (int k = -el; k <= el; k++) {
                    //unsigned long long i = (4*el*el*el/3.0) + 2*el*el + (5*el/3.0) + 2*m*el + m + k;
                    //std::cout << j << '\t';
                    unsigned long long j = 0;
                    for (int ell = 0; ell < lmax+1; ell++) {
                        for (int mm = -ell; mm <= ell; mm++) {
                            for (int kk = -ell; kk <= ell; kk++) {
                                //unsigned long long j = (4*ell*ell*ell/3.0) + 2*ell*ell + (5*ell/3.0) + 2*mm*ell + mm + kk;
                                if (j > i) continue;
                                if (i == j) {
                                    try {
                                        H(i,j) += 0.5*(A+C)*el*(el+1) + 0.5*(A-C)*kap*k*k;
                                        if (k+2 <= el) {
                                            double val = 0.25*(C-A)*sqrt(el*(el+1)-k*(k+1))*sqrt(el*(el+1)-(k+1)*(k+2));
                                            H(i+2,j) += val;
                                            H(j,i+2) += val;
                                        }
                                    } catch (const std::exception& e) {
                                        std::cout << "Failure at index: " <<
                                            i << "\t(" << el << ',' << m << ',' <<
                                            k << ')' << std::endl;
                                    }
                                }
                                //if (a.size() == 1) continue;
                                //double Vij = 0;
                                //for (int L = 0; L < 8; L++) {
                                //    for (int M = -L; M <= L; M++) {
                                //        for (int K = -L; K <= L; K++) {
                                //            unsigned long long ind = (4*L*L*L/3.0) + 2*L*L + (5*L/3.0) + 2*M*L + M + K;
                                //            if (ind > a.size()-1 || a[ind] == 0) continue;
                                //            //double Clm = WignerSymbols::clebschGordan(L,ell,el,M,mm,m);
                                //            //double Clk = WignerSymbols::clebschGordan(L,ell,el,K,kk,k);
                                //            //double val = 8*M_PI*M_PI*Clm*Clk/(2.0*el+1.0);
                                //            double Wlm = WignerSymbols::wigner3j(L,ell,el,M,mm,-m);
                                //            double Wlk = WignerSymbols::wigner3j(L,ell,el,K,kk,-k);
                                //            double val = 8*M_PI*M_PI*pow(-1.0, -m-k)*Wlm*Wlk;
                                //            Vij += a[ind] * val;
                                //        }
                                //    }
                                //}
                                //H(i,j) += Vij;
                                //H(j,i) += Vij;
                                j++;
                            }
                        }
                    }
                    i++;
                }
            }
        }
    } // end parallel
    //H.print("H = ");
    //std::cout << "Constructed matrix with LMAX " << lmax << '.' << std::endl;
    return H;
}

std::vector<double> getWignerCoeffs(std::string sysname, bool freeRotor) {
    if (freeRotor) {
        std::vector<double> a = { 0.0 };
        return a;
    }
    std::string dirname = getDirectory(sysname);
    std::string filename = dirname+"/a.txt";
    std::ifstream is(filename);
    double val;
    std::vector<double> a;
    while (is) {
        if(!(is >> val)) {
            break;
        }
        a.push_back(val);
    }
    return a;
}

std::vector<double> getMomentOfInertia(std::string sysname) {
    std::string dirname = getDirectory(sysname);
    std::string filename = dirname+'/'+"I.txt";
    std::ifstream is(filename);
    double val;
    std::vector<double> Ivec;
    while (is) {
        if (!(is >> val)) {
            break;
        }
        Ivec.push_back(val);
    }
    return Ivec;
}

std::string getDirectory(std::string sysname) {
    //std::string dirname = "/Users/lancebettinson/Thesis/umrr/code/hamiltonian-cpp/data";
    std::string dirname = "/global/scratch/lbettins/rotational-hamiltonian/data";
    if (sysname == "") {
        std::string sysname;
        std::cout << "Enter system name:" << std::endl;
        std::cin >> sysname;
    }
    return dirname+'/'+sysname;
}

SpMat<complex<double>> getSparseHam(int lmax, double Ix, double Iy, double Iz, std::vector<double>& a) {
    /* Construct the Hamiltonian Matrix
     * Inputs:  lmax;   the maximum quantum number forthe spherical basis
     *             I;   the gas-phase moments of inertia [=] amu*Å^2
     * Outputs:    H;   the Hamiltonian matrix (Hermitian) [=] Hartree
     */
    unsigned long long size = (lmax+1)*(2*lmax+1)*(2*lmax+3)/3.0;
    SpMat<complex<double>> H = sp_mat(size, size);

    // Define rotational constants:
    double B = HBAR1*HBAR2/(2.0*Iz);
    double A = HBAR1*HBAR2/(2.0*Iy);
    double C = HBAR1*HBAR2/(2.0*Ix);

    double kap = (A == B && A == C) ? 0 : (2.0*B - (A+C)) / (A-C);
    #pragma omp parallel
    {
        #pragma omp for
        for (int el = 0; el < lmax+1; el++) {
            for (int m = -el; m <= el; m++) {
                for (int k = -el; k <= el; k++) {
                    unsigned long long i = (4*el*el*el/3.0) + 2*el*el + (5*el/3.0) + 2*m*el + m + k;
                    for (int ell = 0; ell < lmax+1; ell++) {
                        for (int mm = -ell; mm <= ell; mm++) {
                            for (int kk = -ell; kk <= ell; kk++) {
                                unsigned long long j = (4*ell*ell*ell/3.0) + 2*ell*ell + (5*ell/3.0)
                                    + 2*mm*ell + mm + kk;
                                if (j > i) continue;
                                if (i == j) {
                                    try {
                                        H(i,j) += 0.5*(A+C)*el*(el+1) + 0.5*(A-C)*kap*k*k;
                                        if (k+2 <= el) {
                                            double val = 0.25*(C-A)*sqrt(el*(el+1)-k*(k+1))*sqrt(el*(el+1)-(k+1)*(k+2));
                                            H(i+2,j) += val;
                                            H(j,i+2) += val;
                                        }
                                    } catch (const std::exception& e) {
                                        std::cout << "Failure at index: " <<
                                            i << "\t(" << el << ',' << m << ',' <<
                                            k << ')' << std::endl;
                                    }
                                }
                            }
                        }
                    }
                }
            }
        }
    } // end parallel
    return H;
}

double getSparseQ(double T, SpMat<complex<double>>& H, int sym) {
    /* Solve the Eigenvalues forSparse Matrix
     * Inputs:  T;  the temperature [=] K
     *          H;  the (sparse) Hamiltonian matrix [=] Hartree
     */
    double b = pow(kB * T, -1) * EHARTREE;
    vec eigval;
    mat eigvec;
    //std::cout << "Number of rows: " << H.n_rows << std::endl;
    eigs_sym(eigval, eigvec, H, H.n_rows-1);
    double Q = 0;
    //std::cout << "Eigenvalues: " << std::endl;
    for (double e : eigval) {
        //std::cout << e << '\t';
        Q += exp(-b * e);
    }
    //std::cout << "Q predicted by eigs_sym: " << Q/sym << std::endl;
    return double(Q/sym);
}

Col<double> getCoefficients(std::string sysname) {
    std::string dirname = getDirectory(sysname);
    std::string filename = dirname+'/'+"vdat.txt";
    std::ifstream is(filename);
    if (is.fail())
    {
        std::cout << "cannot open file " << filename;
    }
    double theta, phi, v;
    std::vector<double> my_vec;
    while (is) {
        if (!(is >> theta >> phi >> v)) {
            break;
        }
        //std::cout << theta << '\t' << phi << '\t' << v << std::endl;
        my_vec.push_back(v);
    }
    Col<double> cvec = conv_to<vec>::from(my_vec);
    //cvec.print();
    is.close();
    return cvec;
}


// Software License for MTL
//
// Copyright (c) 2007 The Trustees of Indiana University.
//               2008 Dresden University of Technology and the Trustees of Indiana University.
//               2010 SimuNova UG (haftungsbeschrÃ¤nkt), www.simunova.com.
// All rights reserved.
// Authors: Peter Gottschling and Andrew Lumsdaine
//
// This file is part of the Matrix Template Library
//
// See also license.mtl.txt in the distribution.

#include <iostream>
#include <boost/numeric/mtl/mtl.hpp>

using namespace std;
using namespace mtl;

template <typename Matrix>
void init_data(Matrix& T)
{
    for (unsigned i= 0; i < num_rows(T); i++)
	for (unsigned j= 0; j < num_cols(T); j++)
	    T[i][j]= i + j + 2; // add 2 because i and j are zero-based
}

template <typename Matrix, typename Vector>
void w_semblance(const Matrix& T, Vector& semb, int nwin, int lwin, int linc)
{
    int nsamp= num_rows(T), ntr= num_cols(T);
    dense_vector<double> sumsqr(nsamp), sqrsum(nsamp);

    for (int i= 0; i < nsamp; i++) {
	sumsqr[i]= unary_dot(T[i][iall]);
	sqrsum[i]= square(sum(T[i][iall]));
    }

    for (int i= 1, ll= 0; i <= nwin; i++, ll+= linc) {
	int kkend= ll + kkend > nsamp ? nsamp - ll : lwin;
	irange r(ll, kkend);
	double sumsq= sum(sumsqr[r]),
	       sumampsq= sum(sqrsum[r]),
	       value= ntr * sumsq;
	semb[i]= value != 0.0 ? sumampsq / value : 0.0;
    }
}

int main()
{
    //int nsamples= 30001, ntrc= 2000;
    int nsamples= 8, ntrc= 6;

    dense2D<float> traces(nsamples, ntrc);
    init_data(traces);
    // cout << "T is\n" << traces;

    dense_vector<double> semb(nsamples/2 + 1, 0.0);
    int lwin= 5, linc= 2, nwin= size(semb) - linc;
    w_semblance(traces, semb, nwin, lwin, linc);

    cout << "Semblance is " << semb << '\n';

    return 0;
}


#include "Polynomial.h"

#include <complex>
#include <iostream>

#define BOOST_TEST_MODULE Polynomial
#include <boost/test/included/unit_test.hpp>

using namespace std;
using namespace Storage_B::Polynomials;

// 0 degree polynomial
BOOST_AUTO_TEST_CASE(zero_degree)
{
  Polynomial<double> p0(42.0);

  BOOST_CHECK(p0.degree() == 0);
  BOOST_CHECK(p0[0] == 42.0);
  BOOST_CHECK(p0(0.123) == 42.0);
}

// first degree polynomial
BOOST_AUTO_TEST_CASE(first_degree)
{
  Polynomial<double> p1(0.0, 1);
  p1[0] = -1.0;
  p1[1] = 1.0;

  BOOST_CHECK(p1.degree() == 1);
  BOOST_CHECK(p1(0.0) == -1.0);
  BOOST_CHECK(p1(1.0) == 0.0);
  BOOST_CHECK(p1(2.0) == 1.0);
  BOOST_CHECK(p1(-1.0) == -2.0);
}

// second degree polynomial
BOOST_AUTO_TEST_CASE(second_degree)
{
  double c[3] ;

  c[0] = 2.0 ;
  c[1] = -1.0;
  c[2] = 3.0 ;

  Polynomial<double> p2(c, 2) ;

  BOOST_CHECK(p2.degree() == 2);
  BOOST_CHECK(p2(0.0) == 2.0);
  BOOST_CHECK(p2(1.0) == 4.0);
  BOOST_CHECK(p2(2.0) == 12.0);
}

// third degree polynomial
BOOST_AUTO_TEST_CASE(third_degree)
{
  vector<float> c;

  c.push_back(-1.0);
  c.push_back(1.0);
  c.push_back(-2.0);
  c.push_back(3.0);

  Polynomial<float> p3(c);

  BOOST_CHECK(p3.degree() == 3);
  BOOST_CHECK(p3(0.0) == -1.0);
  BOOST_CHECK(p3(1.0) == 1.0);
  BOOST_CHECK(p3(-1.0) == -7.0);
}

// increment degree
BOOST_AUTO_TEST_CASE(increment_degree)
{
  double c[3] ;

  c[0] = 2.0 ;
  c[1] = -1.0;
  c[2] = 3.0 ;

  Polynomial<double> p(c, 2) ;

  BOOST_CHECK(p.degree() == 2);

  p.IncrementDegree(12.0);

  BOOST_CHECK(p.degree() == 3);
  BOOST_CHECK(p[0] == 2);
  BOOST_CHECK(p[1] == -1.0);
  BOOST_CHECK(p[2] == 3.0);
  BOOST_CHECK(p[3] == 12.0);
}

// decrement degree
BOOST_AUTO_TEST_CASE(decrement_degree)
{
  double c[3] ;

  c[0] = 2.0 ;
  c[1] = -1.0;
  c[2] = 3.0 ;

  Polynomial<double> p(c, 2);

  BOOST_CHECK(p.degree() == 2);

  auto val = p.DecrementDegree();

  BOOST_CHECK(p.degree() == 1);
  BOOST_CHECK(p[0] == 2);
  BOOST_CHECK(p[1] == -1.0);
  BOOST_CHECK(val == 3.0);
}

// complex polynomial
BOOST_AUTO_TEST_CASE(complex_poly)
{
  Polynomial<complex<double> > p;
  p[0] = 1.0;
  p.IncrementDegree(0.0);
  p.IncrementDegree(1.0);

  BOOST_CHECK(p.degree() == 2);
  BOOST_CHECK(p(1.0i) == 0.0);

  p.IncrementDegree();

  BOOST_CHECK(p.degree() == 3);
  BOOST_CHECK(p[3] == 0.0);

  p[3] = 2.0 - 3.0i;
  complex<double> result = p(2.0 + 1.0i);

  BOOST_CHECK(result.real() == 41);
  BOOST_CHECK(result.imag() == 20);
}

// Evaluate
BOOST_AUTO_TEST_CASE(derivative)
{
  vector<double> c;

  c.push_back(2.0) ;
  c.push_back(-1.0);
  c.push_back(3.0) ;
  Polynomial<double> f(c);

  c.clear();
  c.push_back(f[1]);
  c.push_back(2.0 * f[2]);
  Polynomial<double> fprime(c);

  double x = 42.0;

  auto res = f.eval(x);

  BOOST_CHECK(res.first == f(x)); // y = f(x)
  BOOST_CHECK(res.second == fprime(x)); // dy/dx = f'(x)
}


#include <kv/Heine.hpp>
#include <kv/qAiry.hpp>
#include <kv/qBessel.hpp>
#include <cmath>
#include <iostream>
#include <boost/numeric/ublas/vector.hpp>
#include <boost/numeric/ublas/io.hpp>
typedef kv::interval<double> itv;
typedef kv::complex< kv::interval<double> > cp;
using namespace std;
namespace ub = boost::numeric::ublas;
int main()
{
  cout.precision(17);
  int n=20;
  itv nu,q;
  ub::vector< itv > x(100);
  q="0.7";
  nu=1.5;
  x(0)=4.5;
  x(1)=4.45;
  for(int i=1;i<=n;i++){
    x(i+1)=x(i)-kv::Hahn_Exton(itv(x(i)),itv(nu),itv(q))*(x(i)-x(i-1))
      /(kv::Hahn_Exton(itv(x(i)),itv(nu),itv(q))-kv::Hahn_Exton(itv(x(i-1)),itv(nu),itv(q)));
  cout<<x(i+1)<<endl;
  cout<<"value of HE inf"<<kv::Hahn_Exton(itv(x(i+1).lower()),itv(nu),itv(q))<<endl;
  cout<<"value of HE sup"<<kv::Hahn_Exton(itv(x(i+1).upper()),itv(nu),itv(q))<<endl;
  cout<<"value of HE mid"<<kv::Hahn_Exton(itv(mid(x(i+1))),itv(nu),itv(q))<<endl;
  }
}


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

  @copyright 2016 NumScale SAS

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

#if defined(DOXYGEN_ONLY)
namespace boost { namespace simd
{
  /*!

    @ingroup group-arithmetic
    Function object function implementing correct_fma

    Computes the "correct" fused multiply/add of its parameter: this is perhaps not that you need
    if speed is needed more than accuracy (this remark is hardware dependent).

    @par semantic:
    For any given value @c x, @c y, @c z of type @c T:

    @code
    T r = correct_fma(x, y, z);
    @endcode

    is similar to:

    @code
    T r = x*y+z;
    @endcode

    but is only rounded once and with no intermediate overflow.

    @par Notes:
    - For integers x*y+z is always performed in 2-complement wraping.

    - For floating points numbers, the correct fused multiply add is computed,
    meaning the computation of x*y+z with only one rounding operation.
    This operation may be very expensive on architectures where this
    functionality is not hardware provided.

    - @c correct_fma is in fact a transitory function which ensures
    strict @ref fma capabilities, i.e. only one rounding operation and no undue
    overflow in intermediate computations.

    - If you are using this function on an architecture without @ref fma
    and are sure that overflow is not a problem
    you may define BOOST_SIMD_DONT_CARE_CORRECT_FMA_OVERFLOW for better
    performance.

    - @c correct_fma is never used internally in boost.simd.

    @par Decorators
      std_ for floating entries calls std::fma but does not guarantee performances...

    @see  fma, fms fnma, fnms
**/
     Value correct_fma(Value const& v0, Value const& v1, Value const& v2);
} }
#endif

#include <boost/simd/function/scalar/correct_fma.hpp>
#include <boost/simd/function/simd/correct_fma.hpp>

#endif


#ifndef GICP_COST_HPP
#define GICP_COST_HPP

#include <ceres/ceres.h>
#include <ceres/rotation.h>
#include <Eigen/Core>
#include <Eigen/Geometry>
#include <Eigen/Dense>
#include "litamin2/ceres_cost/PoseSE3Parameterization.hpp"
// L^T*error
struct GICP_FACTOR {
  GICP_FACTOR(Eigen::Vector3d p_mean, Eigen::Vector3d q_mean, Eigen::Matrix3d p_cov, Eigen::Matrix3d q_cov) : p_mean_(p_mean), q_mean_(q_mean), p_cov_(p_cov), q_cov_(q_cov) {}

  template <typename T>
  bool operator()(const T* const q, const T* const t, T* residuals) const {
    Eigen::Map<Eigen::Matrix<T, 3, 1>> residuals_map(residuals);
    Eigen::Matrix<T, 3, 1> p_m(p_mean_.cast<T>());
    Eigen::Matrix<T, 3, 1> q_m(q_mean_.cast<T>());
    Eigen::Matrix<T, 3, 3> p_c = p_cov_.cast<T>();
    Eigen::Matrix<T, 3, 3> q_c = q_cov_.cast<T>();

    Eigen::Quaternion<T> quat(q);
    Eigen::Matrix<T, 3, 1> translation(t);

    Eigen::Matrix<T, 3, 3> mahalanobis = (q_c + quat * p_c * quat.inverse()).inverse();
    Eigen::Matrix<T, 3, 3> LT = mahalanobis.llt().matrixL().transpose();
    residuals_map = LT * (q_m - (quat * p_m + translation));

    return true;
  }

  static ceres::CostFunction* Create(Eigen::Vector3d p_mean_, Eigen::Vector3d q_mean_, Eigen::Matrix3d p_cov_, Eigen::Matrix3d q_cov_) {
    // 分别是残差，q，t的维度
    return (new ceres::AutoDiffCostFunction<GICP_FACTOR, 3, 4, 3>(new GICP_FACTOR(p_mean_, q_mean_, p_cov_, q_cov_)));
  }

  Eigen::Vector3d p_mean_, q_mean_;
  Eigen::Matrix3d p_cov_, q_cov_;
};

class GICPAnalyticCostFunction : public ceres::SizedCostFunction<3, 7> {
public:
  GICPAnalyticCostFunction(Eigen::Vector3d p_mean, Eigen::Vector3d q_mean, Eigen::Matrix3d p_cov, Eigen::Matrix3d q_cov)
  : p_mean_(p_mean),
    q_mean_(q_mean),
    p_cov_(p_cov),
    q_cov_(q_cov) {}
  virtual ~GICPAnalyticCostFunction() {}
  // parameters是[0,0,0,1,x,y,z]
  virtual bool Evaluate(double const* const* parameters, double* residuals, double** jacobians) const {
    // 默认已经归一化
    Eigen::Quaterniond q_last_curr(parameters[0]);
    Eigen::Vector3d t_last_curr(parameters[0] + 4);
    Eigen::Matrix3d LT = (q_cov_ + q_last_curr * q_cov_ * q_last_curr.inverse()).inverse().llt().matrixL().transpose();
    Eigen::Map<Eigen::Vector3d> residuals_map(residuals);
    Eigen::Vector3d p_mean_trans = q_last_curr * p_mean_ + t_last_curr;
    residuals_map = LT * (q_mean_ - p_mean_trans);

    if (jacobians != NULL) {
      if (jacobians[0] != NULL) {
        Eigen::Map<Eigen::Matrix<double, 3, 7, Eigen::RowMajor>> J_se3(jacobians[0]);
        J_se3.setZero();
        Eigen::Matrix<double, 3, 6> dp_by_se3;
        dp_by_se3.block<3, 3>(0, 0) = skew(p_mean_trans);
        dp_by_se3.block<3, 3>(0, 3) = -Eigen::Matrix3d::Identity();
        J_se3.block<3, 6>(0, 0) = LT * dp_by_se3;
      }
    }
    return true;
  }

  Eigen::Vector3d p_mean_, q_mean_;
  Eigen::Matrix3d p_cov_, q_cov_;
};

// 弃用
// 精度不高，而且耗时
// Z是R*Sigma_i*RT对李代数的雅可比矩阵，来自于d2d-ndt，具体如何计算，我不清楚
class GICPDoubleAnalyticCostFunction : public ceres::SizedCostFunction<3, 7> {
public:
  GICPDoubleAnalyticCostFunction(Eigen::Vector3d p_mean, Eigen::Vector3d q_mean, Eigen::Matrix3d p_cov, Eigen::Matrix3d q_cov)
  : p_mean_(p_mean),
    q_mean_(q_mean),
    p_cov_(p_cov),
    q_cov_(q_cov) {}
  virtual ~GICPDoubleAnalyticCostFunction() {}
  // parameters是[0,0,0,1,x,y,z]
  virtual bool Evaluate(double const* const* parameters, double* residuals, double** jacobians) const {
    // 默认已经归一化
    Eigen::Quaterniond q_last_curr(parameters[0]);
    Eigen::Vector3d t_last_curr(parameters[0] + 4);
    Eigen::Matrix3d mahalanobis = (q_cov_ + q_last_curr * q_cov_ * q_last_curr.inverse()).inverse();
    // Eigen::Matrix3d LT = mahalanobis.llt().matrixL().transpose();
    Eigen::Map<Eigen::Vector3d> residuals_map(residuals);
    Eigen::Vector3d p_mean_trans = q_last_curr * p_mean_ + t_last_curr;
    residuals_map = mahalanobis * (q_mean_ - p_mean_trans);

    if (jacobians != NULL) {
      if (jacobians[0] != NULL) {
        Eigen::Map<Eigen::Matrix<double, 3, 7, Eigen::RowMajor>> J_se3(jacobians[0]);
        J_se3.setZero();
        Eigen::Matrix<double, 3, 6> dp_by_se3;
        Eigen::Matrix<double, 3, 6> BZBU(Eigen::Matrix<double,3,6>::Zero());
        dp_by_se3.block<3, 3>(0, 0) = skew(p_mean_trans);
        dp_by_se3.block<3, 3>(0, 3) = -Eigen::Matrix3d::Identity();
        Eigen::Matrix3d Z1, Z2, Z3;
        Z1 << 0, -p_cov_(0, 2), p_cov_(0, 1), -p_cov_(0, 2), -2 * p_cov_(1, 2), -p_cov_(2, 2) + p_cov_(1, 1), p_cov_(0, 1), -p_cov_(2, 2) + p_cov_(1, 1), 2 * p_cov_(1, 2);
        Z2 << 2 * p_cov_(0, 2), p_cov_(1, 2), -p_cov_(0, 0) + p_cov_(2, 2), p_cov_(1, 2), 0, -p_cov_(0, 1), -p_cov_(0, 0) + p_cov_(2, 2), -p_cov_(0, 1), -2 * p_cov_(0, 2);
        Z3 << -2 * p_cov_(0, 1), -p_cov_(1, 1) + p_cov_(0, 0), -p_cov_(1, 2), -p_cov_(1, 1) + p_cov_(0, 0), 2 * p_cov_(0, 1), p_cov_(0, 2), -p_cov_(1, 2), p_cov_(0, 2), 0;
        BZBU.block<3,1>(0,0) = mahalanobis * Z1 * mahalanobis * p_mean_trans;
        BZBU.block<3,1>(0,1) = mahalanobis * Z2 * mahalanobis * p_mean_trans;
        BZBU.block<3,1>(0,2) = mahalanobis * Z3 * mahalanobis * p_mean_trans;
        J_se3.block<3, 6>(0, 0) = -BZBU + mahalanobis * dp_by_se3;
      }
    }
    return true;
  }

  Eigen::Vector3d p_mean_, q_mean_;
  Eigen::Matrix3d p_cov_, q_cov_;
};

// 正常工作，但精度并不是特别的高
// TODO 删除无效代码
class ICPAnalyticCostFunction : public ceres::SizedCostFunction<3, 7> {
public:
  EIGEN_MAKE_ALIGNED_OPERATOR_NEW
  ICPAnalyticCostFunction(Eigen::Vector3d p_mean, Eigen::Vector3d q_mean, Eigen::Matrix3d p_cov, Eigen::Matrix3d q_cov)
  : p_mean_(p_mean),
    q_mean_(q_mean),
    p_cov_(p_cov),
    q_cov_(q_cov) {}
  virtual ~ICPAnalyticCostFunction() {}
  // parameters是[0,0,0,1,x,y,z]
  virtual bool Evaluate(double const* const* parameters, double* residuals, double** jacobians) const {
    // 默认已经归一化
    Eigen::Quaterniond q_last_curr(parameters[0]);
    Eigen::Vector3d t_last_curr(parameters[0] + 4);
    // Eigen::Matrix3d mahalanobis = q_cov_ + q_last_curr * q_cov_ * q_last_curr.inverse();
    // Eigen::Matrix3d LT = mahalanobis.llt().matrixL().transpose();
    Eigen::Map<Eigen::Vector3d> residuals_map(residuals);
    Eigen::Vector3d p_mean_trans = q_last_curr * p_mean_ + t_last_curr;
    residuals_map = q_mean_ - p_mean_trans;

    if (jacobians != NULL) {
      if (jacobians[0] != NULL) {
        Eigen::Map<Eigen::Matrix<double, 3, 7, Eigen::RowMajor>> J_se3(jacobians[0]);
        J_se3.setZero();
        Eigen::Matrix<double, 3, 6> dp_by_se3;
        dp_by_se3.block<3, 3>(0, 0) = skew(p_mean_trans);
        dp_by_se3.block<3, 3>(0, 3) = -Eigen::Matrix3d::Identity();
        J_se3.block<3, 6>(0, 0) = dp_by_se3;
      }
    }
    return true;
  }

  Eigen::Vector3d p_mean_, q_mean_;
  Eigen::Matrix3d p_cov_, q_cov_;
};
#endif