/**
 * \file boost/numeric/ublasx/operation/lsq.hpp
 *
 * \brief Least Square problem solvers.
 *
 * Copyright (c) 2010, Marco Guazzone
 *
 * Distributed under 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)
 *
 * \author Marco Guazzone, marco.guazzone@gmail.com
 */

#ifndef BOOST_NUMERIC_UBLASX_OPERATION_LSQ_HPP
#define BOOST_NUMERIC_UBLASX_OPERATION_LSQ_HPP

#include <algorithm>
#include <boost/numeric/bindings/lapack/driver/gels.hpp>
#include <boost/numeric/bindings/lapack/driver/gelss.hpp>
#include <boost/numeric/bindings/ublas.hpp>
#include <boost/numeric/ublas/detail/temporary.hpp>
#include <boost/numeric/ublas/expression_types.hpp>
#include <boost/numeric/ublas/fwd.hpp>
#include <boost/numeric/ublas/matrix.hpp>
#include <boost/numeric/ublas/vector.hpp>
#include <boost/numeric/ublas/traits.hpp>
#include <boost/numeric/ublasx/operation/num_columns.hpp>
#include <boost/numeric/ublasx/operation/num_rows.hpp>
#include <boost/numeric/ublasx/operation/rcond.hpp>
#include <boost/numeric/ublasx/operation/size.hpp>

namespace boost { namespace numeric { namespace ublasx {

using namespace ::boost::numeric::ublas;

namespace detail { namespace /*<unnamed>*/ {

template <typename MatrixT, typename VectorT>
void llsq_qr_impl(MatrixT& A, VectorT& b, column_major_tag)
{
	typedef typename promote_traits<
				typename matrix_traits<MatrixT>::size_type,
				typename vector_traits<VectorT>::size_type
		>::promote_type size_type;

//	size_type m = num_rows(A);
	size_type n = num_columns(A);

	::boost::numeric::bindings::lapack::gels(A, b);

	b.resize(n, true);
}

template <typename MatrixT, typename VectorT>
void llsq_qr_impl(MatrixT& A, VectorT& b, row_major_tag)
{
	typedef typename matrix_traits<MatrixT>::value_type value_type;
	typedef matrix<value_type, column_major> colmaj_matrix_type;

	colmaj_matrix_type tmp_A(A);

	llsq_qr_impl(tmp_A, b, column_major_tag());
}

template <typename MatrixT, typename VectorT>
void llsq_qr_impl(matrix_expression<MatrixT> const& A, VectorT& b, column_major_tag)
{
	typedef typename matrix_traits<MatrixT>::value_type value_type;
	typedef matrix<value_type, column_major> colmaj_matrix_type;

	colmaj_matrix_type tmp_A(A);

	llsq_qr_impl(tmp_A, b, column_major_tag());
}

template <typename MatrixT, typename VectorT>
void llsq_qr_impl(matrix_expression<MatrixT> const& A, VectorT& b, row_major_tag)
{
	typedef typename matrix_traits<MatrixT>::value_type value_type;
	typedef matrix<value_type, column_major> colmaj_matrix_type;

	colmaj_matrix_type tmp_A(A);

	llsq_qr_impl(tmp_A, b, column_major_tag());
}

template <typename MatrixT, typename VectorT>
void llsq_svd_impl(MatrixT& A, VectorT& b, column_major_tag)
{
	typedef typename matrix_traits<MatrixT>::value_type value_type;
	typedef typename type_traits<value_type>::real_type real_type;
	typedef typename promote_traits<
				typename matrix_traits<MatrixT>::size_type,
				typename vector_traits<VectorT>::size_type
		>::promote_type size_type;
	typedef vector<real_type> work_vector_type;

	size_type m = num_rows(A);
	size_type n = num_columns(A);
	size_type k = ::std::min(m, n);
	real_type rc = rcond(A);
	::fortran_int_t r;
	//work_vector_type tmp_b(b);//TODO: should we do this in order to avoid problem with other types of vector (like sparse vector)
	work_vector_type dummy_s(k);

	::boost::numeric::bindings::lapack::gelss(A, b, dummy_s, rc, r);

	b.resize(n, true);
}

template <typename MatrixT, typename VectorT>
void llsq_svd_impl(MatrixT& A, VectorT& b, row_major_tag)
{
	typedef typename matrix_traits<MatrixT>::value_type value_type;
	typedef matrix<value_type, column_major> colmaj_matrix_type;

	colmaj_matrix_type tmp_A(A);

	llsq_svd_impl(tmp_A, b, column_major_tag());
}

template <typename MatrixT, typename VectorT>
void llsq_svd_impl(matrix_expression<MatrixT> const& A, VectorT& b, column_major_tag)
{
	typedef typename matrix_traits<MatrixT>::value_type value_type;
	typedef matrix<value_type, column_major> colmaj_matrix_type;

	colmaj_matrix_type tmp_A(A);

	llsq_svd_impl(tmp_A, b, column_major_tag());
}

template <typename MatrixT, typename VectorT>
void llsq_svd_impl(matrix_expression<MatrixT> const& A, VectorT& b, row_major_tag)
{
	typedef typename matrix_traits<MatrixT>::value_type value_type;
	typedef matrix<value_type, column_major> colmaj_matrix_type;

	colmaj_matrix_type tmp_A(A);

	llsq_svd_impl(tmp_A, b, column_major_tag());
}

}} // Namespace detail::<unnamed>

/**
 * \brief Solve the linear (ordinary) least square problem by using the QR
 *  decomposition.
 * \tparam MatrixExprT Type of the input matrix expression.
 * \tparam VectorExprT Type of the input/output vector.
 * \param A The input matrix expression (i.e., the design matrix).
 * \param b On entry, the input vector (i.e., the observations vector); on exit,
 *  the least square solution.
 *
 * Orthogonal decomposition methods of solving the least squares problem are
 * slower than directly solving the normal equations but are more numerically
 * stable.
 */
template <typename MatrixExprT, typename VectorT>
BOOST_UBLAS_INLINE
void llsq_qr_inplace(matrix_expression<MatrixExprT> const& A, VectorT& b)
{
	typedef typename matrix_traits<MatrixExprT>::orientation_category orientation_category;

	detail::llsq_qr_impl(A, b, orientation_category());
}

/**
 * \brief Solve the linear (ordinary) least square problem by using the QR
 *  decomposition.
 * \tparam MatrixExprT Type of the input matrix expression.
 * \tparam VectorExprT Type of the input/output vector.
 * \param A The input matrix expression (i.e., the design matrix).
 * \param b The input vector (i.e., the observations vector).
 * \return The least square solution.
 *
 * Orthogonal decomposition methods of solving the least squares problem are
 * slower than directly solving the normal equations but are more numerically
 * stable.
 */
template <typename MatrixExprT, typename VectorExprT>
BOOST_UBLAS_INLINE
typename vector_temporary_traits<VectorExprT>::type llsq_qr(matrix_expression<MatrixExprT> const& A, vector_expression<VectorExprT> const& b)
{
	typedef typename vector_temporary_traits<VectorExprT>::type out_vector_type;

	out_vector_type x(b);

	llsq_qr_inplace(A, x);

	return x;
}

/**
 * \brief Solve the linear (ordinary) least square problem by using the Singular
 * Value  Decomposition (SVD) method.
 * \tparam MatrixExprT Type of the input matrix expression.
 * \tparam VectorExprT Type of the input/output vector.
 * \param A The input matrix expression (i.e., the design matrix).
 * \param b On entry, the input vector (i.e., the observations vector); on exit,
 *  the least square solution.
 *
 * This method is the most computationally intensive, but is particularly useful
 * if the normal equations matrix is very ill-conditioned (i.e. if its condition
 * number multiplied by the machine's relative round-off error is appreciably
 * large).
 */
template <typename MatrixExprT, typename VectorT>
BOOST_UBLAS_INLINE
void llsq_svd_inplace(matrix_expression<MatrixExprT> const& A, VectorT& b)
{
	typedef typename matrix_traits<MatrixExprT>::orientation_category orientation_category;

	detail::llsq_svd_impl(A, b, orientation_category());
}

/**
 * \brief Solve the linear (ordinary) least square problem by using the Singular
 *  Value Decomposition (SVD) method.
 * \tparam MatrixExprT Type of the input matrix expression.
 * \tparam VectorExprT Type of the input/output vector.
 * \param A The input matrix expression (i.e., the design matrix).
 * \param b The input vector (i.e., the observations vector).
 * \return The least square solution.
 *
 * This method is the most computationally intensive, but is particularly useful
 * if the normal equations matrix is very ill-conditioned (i.e. if its condition
 * number multiplied by the machine's relative round-off error is appreciably
 * large).
 */
template <typename MatrixExprT, typename VectorExprT>
BOOST_UBLAS_INLINE
typename vector_temporary_traits<VectorExprT>::type llsq_svd(matrix_expression<MatrixExprT> const& A, vector_expression<VectorExprT> const& b)
{
	typedef typename vector_temporary_traits<VectorExprT>::type out_vector_type;

	out_vector_type x(b);

	llsq_svd_inplace(A, x);

	return x;
}

/**
 * \brief Solve the linear (ordinary) least square problem.
 * \tparam MatrixExprT Type of the input matrix expression.
 * \tparam VectorExprT Type of the input/output vector.
 * \param A The input matrix expression (i.e., the design matrix).
 * \param b On entry, the input vector (i.e., the observations vector); on exit,
 *  the least square solution.
 */
template <typename MatrixExprT, typename VectorT>
BOOST_UBLAS_INLINE
void llsq_inplace(matrix_expression<MatrixExprT> const& A, VectorT& b)
{
	typedef typename matrix_traits<MatrixExprT>::orientation_category orientation_category;

	detail::llsq_svd_impl(A, b, orientation_category());
}

/**
 * \brief Solve the linear (ordinary) least square problem.
 * \tparam MatrixExprT Type of the input matrix expression.
 * \tparam VectorExprT Type of the input/output vector.
 * \param A The input matrix expression (i.e., the design matrix).
 * \param b The input vector (i.e., the observations vector).
 * \return The least square solution.
 */
template <typename MatrixExprT, typename VectorExprT>
BOOST_UBLAS_INLINE
typename vector_temporary_traits<VectorExprT>::type llsq(matrix_expression<MatrixExprT> const& A, vector_expression<VectorExprT> const& b)
{
	typedef typename vector_temporary_traits<VectorExprT>::type out_vector_type;

	out_vector_type x(b);

	llsq_inplace(A, x);

	return x;
}

}}} // Namespace boost::numeric::ublasx

#endif // BOOST_NUMERIC_UBLASX_OPERATION_LSQ_HPP


// Copyright (c) Microsoft Corporation. All rights reserved.
// Licensed under the MIT license.

#include "gemm.h"
#include "common.h"
#include <Eigen/Core>
#include <Eigen/Dense>

// Wasm interop method
void gemm_f32(void *data) {
  uint32_t *dataIndex = static_cast<uint32_t *>(data);
  uint32_t const argc = dataIndex[0];

  gemm_f32_imp(
      PARAM_BOOL(data, dataIndex[1]), PARAM_BOOL(data, dataIndex[2]),
      PARAM_INT32(data, dataIndex[3]), PARAM_INT32(data, dataIndex[4]),
      PARAM_INT32(data, dataIndex[5]), PARAM_FLOAT(data, dataIndex[6]),
      PARAM_FLOAT_PTR(data, dataIndex[7]), PARAM_FLOAT_PTR(data, dataIndex[8]),
      PARAM_FLOAT(data, dataIndex[9]), PARAM_FLOAT_PTR(data, dataIndex[10]));
}

// Core operator implementation
void gemm_f32_imp(const bool TransA, const bool TransB, const int M,
                  const int N, const int K, const float alpha, const float *A,
                  const float *B, const float beta, float *C) {
  auto C_mat = Eigen::Map<Eigen::Matrix<float, Eigen::Dynamic, Eigen::Dynamic>>(
      C, static_cast<int64_t>(N), static_cast<int64_t>(M));
  if (beta == 0) {
    C_mat.setZero();
  } else {
    C_mat *= beta;
  }

  if (!TransA) {
    if (!TransB) {
      C_mat.noalias() +=
          alpha *
          (Eigen::Map<
               const Eigen::Matrix<float, Eigen::Dynamic, Eigen::Dynamic>>(
               B, static_cast<int64_t>(N), static_cast<int64_t>(K)) *
           Eigen::Map<
               const Eigen::Matrix<float, Eigen::Dynamic, Eigen::Dynamic>>(
               A, static_cast<int64_t>(K), static_cast<int64_t>(M)));
    } else {
      C_mat.noalias() +=
          alpha *
          (Eigen::Map<
               const Eigen::Matrix<float, Eigen::Dynamic, Eigen::Dynamic>>(
               B, static_cast<int64_t>(K), static_cast<int64_t>(N))
               .transpose() *
           Eigen::Map<
               const Eigen::Matrix<float, Eigen::Dynamic, Eigen::Dynamic>>(
               A, static_cast<int64_t>(K), static_cast<int64_t>(M)));
    }
  } else {
    if (!TransB) {
      C_mat.noalias() +=
          alpha *
          (Eigen::Map<
               const Eigen::Matrix<float, Eigen::Dynamic, Eigen::Dynamic>>(
               B, static_cast<int64_t>(N), static_cast<int64_t>(K)) *
           Eigen::Map<
               const Eigen::Matrix<float, Eigen::Dynamic, Eigen::Dynamic>>(
               A, static_cast<int64_t>(M), static_cast<int64_t>(K))
               .transpose());
    } else {
      C_mat.noalias() +=
          alpha *
          (Eigen::Map<
               const Eigen::Matrix<float, Eigen::Dynamic, Eigen::Dynamic>>(
               B, static_cast<int64_t>(K), static_cast<int64_t>(N))
               .transpose() *
           Eigen::Map<
               const Eigen::Matrix<float, Eigen::Dynamic, Eigen::Dynamic>>(
               A, static_cast<int64_t>(M), static_cast<int64_t>(K))
               .transpose());
    }
  }
}


#include <HElib/FHE.h>
#include <HElib/FHEContext.h>
#include <HElib/EncryptedArray.h>
#include <HElib/PAlgebra.h>
#include <NTL/ZZX.h>
#include "SMP/Timer.hpp"
#include "SMP/literal.hpp"
#include "SMP/HElib.hpp"
bool check(std::vector<ZZX> const& factors, long d) {
    std::cout << "factors: \n";
    for (auto const &f : factors) {
        std::cout << f << "\n";
        for (long i = 1; i < d; i++) {
            if (NTL::coeff(f, i) != 0)
                return false;
        }
    }
    return true;
}

std::ostream& print_poly(NTL::ZZX const& poly, long d) {
    std::cout << "[";
    for (; d > 0; d--)
        std::cout << NTL::coeff(poly, d) << ",";
    std::cout << NTL::coeff(poly, 0) << "]";
    return std::cout;
}

long inner_product(NTL::zz_pX const& a,
                   NTL::zz_pX const& b)
{
    NTL::zz_p ip;
    long deg = NTL::deg(a);
    for (long i = 0; i <= deg; i++) {
        ip += (NTL::coeff(a, i) * NTL::coeff(b, deg - i));
    }
    return ip._zz_p__rep;
}

/// Return the leading coeff of the rem polynomial f % h
/// h = X^d + a
NTL::zz_p ModGetLeadingCoeff(NTL::zz_pX const& f,
                             NTL::zz_pX const& h)
{
    long m = NTL::deg(f) + 1;
    long d = NTL::deg(h);
    long n = m / d;

    // printf("|f| = %ld, |h| = %ld\n", m, d);
    // std::cout << f << "\n" << h << std::endl;
    std::cout << h << std::endl;
    NTL::zz_p beta = NTL::coeff(h, 0);
    NTL::zz_p beta_power(1);
    NTL::zz_p ret(NTL::coeff(f, d - 1));
    std::cout << ret;
    for (long i = 2; i <= n; i++) {
        auto coeff = NTL::coeff(f, i * d - 1);
        beta_power *= beta;
        coeff *= beta_power;
        if ((i & 1) == 0) {
            ret -= coeff;
        std::cout << " - " << beta_power << " * " << NTL::coeff(f, i * d - 1) ;
        } else {
        std::cout << " + " << beta_power << " * " << NTL::coeff(f, i * d - 1) ;
            ret += coeff;
        }
    }
    std::cout << "\n";
    auto be = beta;
    for (long i = 1; i <= n; i++) {
        std::cout << be << " ";
        be *= beta;
    }
    std::cout << "\n";
    // std::cout << "result " << ret << std::endl;
    return ret;
}

std::vector<long> precompute_power(long beta, long p, long l)
{
    std::vector<long> beta_power_l(l);
    /// (-beta)^k mod p for 0 <= k < l
    for (long i = 0; i < l; i++)
        beta_power_l[i] = NTL::PowerMod(i & 1 ? p - beta : beta, i, p);
    return beta_power_l;
}

long mod_with_precomputed_table(NTL::ZZX const& poly,
                                std::vector<long> const& tbl,
                                FHEcontext const &context)
{
    long d = context.ea->getDegree();
    long l = context.ea->size();
    long p = context.alMod.getPPowR();
    long phim = context.zMStar.getPhiM();
    long ret = 0;
    auto inv_p = NTL::PrepMulMod(p);
    for (long i = 0; i < l; i++) {
        long coeff_loc = (i + 1) * d - 1;
        assert(coeff_loc < phim);
        long coeff = NTL::to_long(NTL::coeff(poly, coeff_loc));
        coeff = NTL::MulMod(coeff, tbl[i], p, inv_p);
        ret = NTL::AddMod(ret, coeff, p);
    }
    return ret;
}

void faster() {
    NTL::SetSeed(NTL::to_ZZ(132));
    long m = 8192;
    long p = 70913;
    NTL::zz_p::init(p);
    FHEcontext context(m, p, 1);
    context.bitsPerLevel = 60;
    buildModChain(context, 2);
    FHESecKey sk(context);
    sk.GenSecKey(64);

    const EncryptedArray *ea = context.ea;
    const long l = ea->size();
    const long d = ea->getDegree();
    std::cout << "l = " << l << ", d = " << d << std::endl;

    std::vector<GMMPrecompTable> tables = precompute_gmm_tables(context);
    const auto &encoder = context.alMod.getDerived(PA_zz_p());

    std::vector<NTL::zz_pX> vec_A(l);
    std::vector<NTL::zz_pX> vec_B(l);
    for (long i = 0; i < l; i++) {
        NTL::random(vec_A[i], d);
        NTL::random(vec_B[i], d);
        std::cout << inner_product(vec_A[i], vec_B[i]) << " ";
    }
    std::cout << "\n";

    NTL::zz_pX encoded_A, encoded_B;
    encoder.CRT_reconstruct(encoded_A, vec_A);
    encoder.CRT_reconstruct(encoded_B, vec_B);

    Ctxt ctx_vec_A(sk), ctx_vec_B(sk);
    sk.Encrypt(ctx_vec_A, NTL::conv<NTL::ZZX>(encoded_A));
    sk.Encrypt(ctx_vec_B, NTL::conv<NTL::ZZX>(encoded_B));
    ctx_vec_A.multiplyBy(ctx_vec_B);

    std::vector<NTL::zz_pX> slots;
    NTL::ZZX decrypted;
    sk.Decrypt(decrypted, ctx_vec_A);
    double raw_dec_time = 0.;
    {
        AutoTimer timer(&raw_dec_time);
        rawDecode(slots, decrypted, context);
    }

    double table_dec_time = 0.;
    std::vector<long> inner_products;
    {
        AutoTimer timer(&table_dec_time);
        extract_inner_products(inner_products, decrypted, tables, context);
    }
    std::cout << inner_products << std::endl;

    std::cout << raw_dec_time << " : " << table_dec_time << std::endl;
}

void test_poly_mod()
{
    NTL::zz_p::init(769);
    NTL::zz_pX ply;
    ply.SetLength(3);
    ply[0] = 3; ply[2] = 1; // X^2 + 3
    NTL::zz_pXModulus mod(ply);

    NTL::zz_pX rnd;
    NTL::random(rnd, 6);
    std::cout << rnd << std::endl;

    NTL::zz_pX rm;
    NTL::rem(rm, rnd, mod);
    std::cout << rm << std::endl;
}

void normal() {
    long m = 8192;
    long p = 769;
    FHEcontext context(m, p, 1);
    context.bitsPerLevel = 59;
    buildModChain(context, 2);
    FHESecKey sk(context);
    sk.GenSecKey(64);

    const auto &factors = context.alMod.getFactorsOverZZ();
    EncryptedArray *ea = new EncryptedArray(context, factors[0]);
    const long l = ea->size();
    const long d = ea->getDegree();

    NTL::zz_p::init(p);
    std::vector<NTL::zz_pX> vec_A(l);
    std::vector<NTL::zz_pX> vec_B(l);
    std::vector<NTL::ZZX> Vec_A(l);
    std::vector<NTL::ZZX> Vec_B(l);
    for (long i = 0; i < l; i++) {
        NTL::random(vec_A[i], d);
        NTL::random(vec_B[i], d);
        NTL::conv(Vec_A[i], vec_A[i]);
        NTL::conv(Vec_B[i], vec_B[i]);
        std::cout << inner_product(vec_A[i], vec_B[i]) << " ";
    }
    std::cout << "\n";

    NTL::ZZX encoded_A, encoded_B;
    for (long i = 0; i < 100; i++) {
        FHE_NTIMER_START(EAEncode);
        ea->encode(encoded_A, Vec_A);
        ea->encode(encoded_B, Vec_B);
        FHE_NTIMER_STOP(EAEncode);

        FHE_NTIMER_START(rawEncode);
        rawEncode(encoded_A, vec_A, context);
        rawEncode(encoded_B, vec_B, context);
        FHE_NTIMER_STOP(rawEncode);
    }

    Ctxt ctx_vec_A(sk), ctx_vec_B(sk);
    sk.Encrypt(ctx_vec_A, encoded_A);
    sk.Encrypt(ctx_vec_B, encoded_B);
    ctx_vec_A.multiplyBy(ctx_vec_B);

    NTL::ZZX decrypted;
    sk.Decrypt(decrypted, ctx_vec_A);
    std::vector<NTL::zz_pX> results;
    std::vector<NTL::ZZX> Results;
    std::vector<double> raw_decodes;
    for (long i = 0; i < 100; i++) {
        FHE_NTIMER_START(EADecode);
        ea->decode(Results, decrypted);
        FHE_NTIMER_STOP(EADecode);

        FHE_NTIMER_START(rawDecode);
        do {
            raw_decodes.push_back(0.);
            AutoTimer timer(&(raw_decodes.back()));
            rawDecode(results, decrypted, context);
        } while (0);
        FHE_NTIMER_STOP(rawDecode);
    }
    //for (auto &s : results) {
    //    std::cout << NTL::coeff(s, d - 1) << " ";
    //}
    //std::cout << "\n";

    printNamedTimer(std::cout, "EAEncode");
    printNamedTimer(std::cout, "EADecode");
    printNamedTimer(std::cout, "rawEncode");
    printNamedTimer(std::cout, "rawDecode");
    auto decode = mean_std(raw_decodes);
    std::cout << decode.first / 1000. << " +- " << decode.second << std::endl;
}

int main() {
    //normal();
    faster();
    //test_poly_mod();
    return 0;
}


/**
 * @file prefinement.cc
 * @brief Creates convergence plots for experiment 3.2.3.11
 * @author Tobias Rohner
 * @date April 2020
 * @copyright MIT License
 */

#define _USE_MATH_DEFINES

#include <lf/fe/fe.h>
#include <lf/io/io.h>
#include <lf/mesh/hybrid2d/hybrid2d.h>
#include <lf/mesh/mesh.h>
#include <lf/mesh/utils/utils.h>
#include <lf/refinement/mesh_function_transfer.h>
#include <lf/refinement/refinement.h>

#include <Eigen/Dense>
#include <Eigen/Sparse>
#include <boost/program_options.hpp>
#include <cmath>
#include <cstdlib>
#include <filesystem>
#include <fstream>
#include <iostream>
#include <memory>
#include <string>
#include <tuple>
#include <vector>

namespace po = boost::program_options;

/**
 * @brief Builds a mesh on [0, 1]^2 from two triangles
 * @returns A shared pointer to a mesh
 */
std::shared_ptr<lf::mesh::Mesh> getSquareDomain() {
  lf::mesh::hybrid2d::MeshFactory factory(2);
  // Add the vertices
  std::vector<lf::mesh::MeshFactory::size_type> vertices;
  Eigen::Vector2d vertex_coord;
  vertex_coord << 0, 0;
  vertices.push_back(factory.AddPoint(vertex_coord));
  vertex_coord << 1, 0;
  vertices.push_back(factory.AddPoint(vertex_coord));
  vertex_coord << 1, 1;
  vertices.push_back(factory.AddPoint(vertex_coord));
  vertex_coord << 0, 1;
  vertices.push_back(factory.AddPoint(vertex_coord));
  // Add the triangles
  Eigen::Matrix<double, Eigen::Dynamic, 3> coords(2, 3);
  lf::mesh::MeshFactory::size_type nodes[3];  // NOLINT
  coords << 0, 1, 0, 0, 0, 1;
  nodes[0] = vertices[0];
  nodes[1] = vertices[1];
  nodes[2] = vertices[3];
  auto geom_tria1 = std::make_unique<lf::geometry::TriaO1>(coords);
  factory.AddEntity(lf::base::RefEl::kTria(), nodes, std::move(geom_tria1));
  coords << 1, 1, 0, 0, 1, 1;
  nodes[0] = vertices[1];
  nodes[1] = vertices[2];
  nodes[2] = vertices[3];
  auto geom_tria2 = std::make_unique<lf::geometry::TriaO1>(coords);
  factory.AddEntity(lf::base::RefEl::kTria(), nodes, std::move(geom_tria2));
  // Build the mesh
  return factory.Build();
}

/**
 * @brief Builds a mesh on [-1, 1]^2 \ (]0, 1[x]-1, 0[) from four triangles
 * @returns A shared pointer to a mesh
 */
std::shared_ptr<lf::mesh::Mesh> getLDomain() {
  lf::mesh::hybrid2d::MeshFactory factory(2);
  // Add the vertices
  std::vector<lf::mesh::MeshFactory::size_type> vertices;
  Eigen::Vector2d vertex_coord;
  vertex_coord << -1, -1;
  vertices.push_back(factory.AddPoint(vertex_coord));
  vertex_coord << 0, -1;
  vertices.push_back(factory.AddPoint(vertex_coord));
  vertex_coord << 0, 0;
  vertices.push_back(factory.AddPoint(vertex_coord));
  vertex_coord << 1, 0;
  vertices.push_back(factory.AddPoint(vertex_coord));
  vertex_coord << 1, 1;
  vertices.push_back(factory.AddPoint(vertex_coord));
  vertex_coord << -1, 1;
  vertices.push_back(factory.AddPoint(vertex_coord));
  // Add the triangles
  Eigen::Matrix<double, Eigen::Dynamic, 3> coords(2, 3);
  lf::mesh::MeshFactory::size_type nodes[3];  // NOLINT
  coords << -1, 0, -1, -1, 0, 1;
  nodes[0] = vertices[0];
  nodes[1] = vertices[2];
  nodes[2] = vertices[5];
  auto geom_tria1 = std::make_unique<lf::geometry::TriaO1>(coords);
  factory.AddEntity(lf::base::RefEl::kTria(), nodes, std::move(geom_tria1));
  coords << 0, 1, -1, 0, 1, 1;
  nodes[0] = vertices[2];
  nodes[1] = vertices[4];
  nodes[2] = vertices[5];
  auto geom_tria2 = std::make_unique<lf::geometry::TriaO1>(coords);
  factory.AddEntity(lf::base::RefEl::kTria(), nodes, std::move(geom_tria2));
  coords << -1, 0, 0, -1, -1, 0;
  nodes[0] = vertices[0];
  nodes[1] = vertices[1];
  nodes[2] = vertices[2];
  auto geom_tria3 = std::make_unique<lf::geometry::TriaO1>(coords);
  factory.AddEntity(lf::base::RefEl::kTria(), nodes, std::move(geom_tria3));
  coords << 0, 1, 1, 0, 0, 1;
  nodes[0] = vertices[2];
  nodes[1] = vertices[3];
  nodes[2] = vertices[4];
  auto geom_tria4 = std::make_unique<lf::geometry::TriaO1>(coords);
  factory.AddEntity(lf::base::RefEl::kTria(), nodes, std::move(geom_tria4));
  // Build the mesh
  return factory.Build();
}

/**
 * @brief Get the error of a test problem for the given mesh and FE space
 * @param degree The polynomial degree to use
 * @param mesh The mesh on which to solve the PDE
 * @param fe_space The fe space to use
 *
 * The test problem is formulated on [0, 1]^2 with Dirichlet boundary conditions
 and has the analytic solution
 * \f[
        u(x) = \sin(\pi x_1)\sin(\pi x_2)
   \f]
 */
std::tuple<double, double> computeErrorsSquareDomain(
    unsigned degree, const std::shared_ptr<lf::mesh::Mesh> &mesh,
    const std::shared_ptr<lf::fe::ScalarFESpace<double>> &fe_space) {
  // The analytic solution
  const auto u = [](const Eigen::VectorXd &x) -> double {
    return std::sin(M_PI * x[0]) * std::sin(M_PI * x[1]);
  };
  const lf::mesh::utils::MeshFunctionGlobal mf_u(u);
  // The gradient of the analytic solution
  const auto u_grad = [](const Eigen::Vector2d &x) -> Eigen::Vector2d {
    Eigen::Vector2d grad;
    grad[0] = M_PI * std::cos(M_PI * x[0]) * std::sin(M_PI * x[1]);
    grad[1] = M_PI * std::sin(M_PI * x[0]) * std::cos(M_PI * x[1]);
    return grad;
  };
  const lf::mesh::utils::MeshFunctionGlobal mf_u_grad(u_grad);

  // Define the load function of the manufactured solution
  const auto load = [](const Eigen::Vector2d &x) -> double {
    return 2 * M_PI * M_PI * std::sin(M_PI * x[0]) * std::sin(M_PI * x[1]);
  };
  const lf::mesh::utils::MeshFunctionGlobal mf_load(load);

  // Assemble the system matrix and right hand side
  const lf::assemble::DofHandler &dofh = fe_space->LocGlobMap();
  lf::assemble::COOMatrix<double> A_COO(dofh.NumDofs(), dofh.NumDofs());
  Eigen::VectorXd rhs = Eigen::VectorXd::Zero(dofh.NumDofs());
  std::cout << "\t\t> Assembling System Matrix" << std::endl;
  const lf::mesh::utils::MeshFunctionConstant<double> mf_alpha(1);
  lf::fe::DiffusionElementMatrixProvider element_matrix_provider(fe_space,
                                                                 mf_alpha);
  lf::assemble::AssembleMatrixLocally(0, dofh, dofh, element_matrix_provider,
                                      A_COO);
  std::cout << "\t\t> Assembling right Hand Side" << std::endl;
  lf::fe::ScalarLoadElementVectorProvider element_vector_provider(fe_space,
                                                                  mf_load);
  lf::assemble::AssembleVectorLocally(0, dofh, element_vector_provider, rhs);

  // Enforce zero dirichlet boundary conditions
  std::cout << "\t\t> Enforcing Boundary Conditions" << std::endl;
  const auto boundary = lf::mesh::utils::flagEntitiesOnBoundary(mesh);
  const auto selector = [&](unsigned int idx) -> std::pair<bool, double> {
    const auto &entity = dofh.Entity(idx);
    return {entity.Codim() > 0 && boundary(entity), 0};
  };
  lf::assemble::FixFlaggedSolutionComponents(selector, A_COO, rhs);

  // Solve the LSE using the cholesky decomposition
  std::cout << "\t\t> Solving LSE" << std::endl;
  Eigen::SparseMatrix<double> A = A_COO.makeSparse();
  Eigen::SimplicialLDLT<Eigen::SparseMatrix<double>> solver(A);
  const Eigen::VectorXd solution = solver.solve(rhs);
  const lf::fe::MeshFunctionFE<double, double> mf_numeric(fe_space, solution);
  const lf::fe::MeshFunctionGradFE<double, double> mf_numeric_grad(fe_space,
                                                                   solution);

  // Compute the H1 and L2 errors
  std::cout << "\t\t> Computing Error Norms" << std::endl;
  auto qr_segment =
      lf::quad::make_QuadRule(lf::base::RefEl::kSegment(), 2 * degree - 1);
  auto qr_tria =
      lf::quad::make_QuadRule(lf::base::RefEl::kTria(), 2 * degree - 1);
  auto qr_quad =
      lf::quad::make_QuadRule(lf::base::RefEl::kQuad(), 2 * degree - 1);
  const auto quadrule_provider = [&](const lf::mesh::Entity &entity) {
    const lf::base::RefEl refel = entity.RefEl();
    switch (refel) {
      case lf::base::RefEl::kTria():
        return qr_tria;
      case lf::base::RefEl::kSegment():
        return qr_segment;
      case lf::base::RefEl::kQuad():
        return qr_quad;
      default:
        return lf::quad::make_QuadRule(refel, 2 * degree - 1);
    }
  };
  const double H1_err = std::sqrt(lf::fe::IntegrateMeshFunction(
      *mesh, lf::mesh::utils::squaredNorm(mf_u_grad - mf_numeric_grad),
      quadrule_provider));
  const double L2_err = std::sqrt(lf::fe::IntegrateMeshFunction(
      *mesh, lf::mesh::utils::squaredNorm(mf_u - mf_numeric),
      quadrule_provider));

  // Return the errors
  return {H1_err, L2_err};
}

/**
 * @brief Get the error of a test problem for the given mesh and FE space
 * @param degree The polynomial degree of the basis functions
 * @param mesh The mesh on which to solve the PDE
 * @param fe_space The Finite Element Space to use for the computation
 *
 * The test problem is formulated on [-1, 1]^2 \ (]0, 1[x]-1, 0[) with Dirichlet
 boundary conditions and has the analytic solution
 * \f[
        u(r, \phi) = r^{\frac{2}{3}}\sin(\frac{2}{3}\phi)
   \f]
 */
std::tuple<double, double> computeErrorsLDomain(
    unsigned degree, const std::shared_ptr<lf::mesh::Mesh> &mesh,
    const std::shared_ptr<const lf::fe::ScalarFESpace<double>> &fe_space) {
  // The analytic solution
  const auto u = [](const Eigen::Vector2d &x) -> double {
    const double r = x.norm();
    double phi = std::atan2(x[1], x[0]);
    if (phi < 0) {
      phi += 2 * M_PI;
    }
    return std::pow(r, 2. / 3) * std::sin(2. / 3 * phi);
  };
  lf::mesh::utils::MeshFunctionGlobal mf_u(u);
  // The gradient of the analytic solution
  const auto u_grad = [](const Eigen::Vector2d &x) -> Eigen::Vector2d {
    const double r = x.norm();
    double phi = std::atan2(x[1], x[0]);
    if (phi < 0) {
      phi += 2 * M_PI;
    }
    Eigen::Vector2d grad;
    grad[0] = 2. / 3 * std::pow(r, -4. / 3) *
              (x[0] * std::sin(2. / 3 * phi) - x[1] * std::cos(2. / 3 * phi));
    grad[1] = 2. / 3 * std::pow(r, -4. / 3) *
              (x[1] * std::sin(2. / 3 * phi) + x[0] * std::cos(2. / 3 * phi));
    return grad;
  };
  lf::mesh::utils::MeshFunctionGlobal mf_u_grad(u_grad);

  // Get a few useful variables
  const lf::assemble::DofHandler &dofh = fe_space->LocGlobMap();

  // Assemble the system matrix
  std::cout << "\t\t> Assembling System Matrix" << std::endl;
  const lf::mesh::utils::MeshFunctionConstant<double> mf_alpha(1);
  lf::fe::DiffusionElementMatrixProvider element_matrix_provider(fe_space,
                                                                 mf_alpha);
  lf::assemble::COOMatrix<double> A_COO(dofh.NumDofs(), dofh.NumDofs());
  lf::assemble::AssembleMatrixLocally(0, dofh, dofh, element_matrix_provider,
                                      A_COO);

  // The right hand side is zero because we have no load
  Eigen::VectorXd rhs = Eigen::VectorXd::Zero(dofh.NumDofs());

  // Enforce the dirichlet boundary conditions
  std::cout << "\t\t> Enforcing Boundary Conditions" << std::endl;
  const auto boundary = lf::mesh::utils::flagEntitiesOnBoundary(mesh);
  Eigen::VectorXd boundary_dofs = Eigen::VectorXd::Zero(dofh.NumDofs());
  for (const auto *const edge : mesh->Entities(1)) {
    if (boundary(*edge)) {
      const auto *const sfl = fe_space->ShapeFunctionLayout(*edge);
      const auto eval_nodes = sfl->EvaluationNodes();
      const Eigen::RowVectorXd nodal_values = Eigen::Map<Eigen::RowVectorXd>(
          mf_u(*edge, eval_nodes).data(), eval_nodes.cols());
      const Eigen::VectorXd locdofs = sfl->NodalValuesToDofs(nodal_values);
      const auto dofidxs = dofh.GlobalDofIndices(*edge);
      for (long i = 0; i < dofidxs.size(); ++i) {
        boundary_dofs[dofidxs[i]] = locdofs[i];
      }
    }
  }
  const auto selector = [&](unsigned int idx) -> std::pair<bool, double> {
    const lf::mesh::Entity &entity = dofh.Entity(idx);
    return {entity.Codim() > 0 && boundary(entity), boundary_dofs[idx]};
  };
  lf::assemble::FixFlaggedSolutionComponents(selector, A_COO, rhs);

  // Solve the LSE using Cholesky decomposition
  std::cout << "\t\t> Solving LSE" << std::endl;
  Eigen::SparseMatrix<double> A = A_COO.makeSparse();
  Eigen::SimplicialLDLT<Eigen::SparseMatrix<double>> solver(A);
  const Eigen::VectorXd solution = solver.solve(rhs);
  const lf::fe::MeshFunctionFE<double, double> mf_numeric(fe_space, solution);
  const lf::fe::MeshFunctionGradFE<double, double> mf_numeric_grad(fe_space,
                                                                   solution);

  // Compute the H1 and L2 errors
  std::cout << "\t\t> Computing Error Norms" << std::endl;
  auto qr_segment =
      lf::quad::make_QuadRule(lf::base::RefEl::kSegment(), 2 * degree - 1);
  auto qr_tria =
      lf::quad::make_QuadRule(lf::base::RefEl::kTria(), 2 * degree - 1);
  auto qr_quad =
      lf::quad::make_QuadRule(lf::base::RefEl::kQuad(), 2 * degree - 1);
  const auto quadrule_provider = [&](const lf::mesh::Entity &entity) {
    const lf::base::RefEl refel = entity.RefEl();
    switch (refel) {
      case lf::base::RefEl::kTria():
        return qr_tria;
      case lf::base::RefEl::kSegment():
        return qr_segment;
      case lf::base::RefEl::kQuad():
        return qr_quad;
      default:
        return lf::quad::make_QuadRule(refel, 2 * degree - 1);
    }
  };
  const double H1_err = std::sqrt(lf::fe::IntegrateMeshFunction(
      *mesh, lf::mesh::utils::squaredNorm(mf_u_grad - mf_numeric_grad),
      quadrule_provider));
  const double L2_err = std::sqrt(lf::fe::IntegrateMeshFunction(
      *mesh, lf::mesh::utils::squaredNorm(mf_u - mf_numeric),
      quadrule_provider));

  // Return the errors
  return {H1_err, L2_err};
}

int main(int argc, char *argv[]) {
  po::options_description desc("Allowed options");
  desc.add_options()("output,o", po::value<std::string>(),
                     "Name of the output file")(
      "max_p,p", po::value<unsigned>(), "Maximum polynomial degree");
  po::variables_map vm;
  po::store(po::parse_command_line(argc, argv, desc), vm);
  if (vm.count("output") == 0 || vm.count("max_p") == 0) {
    std::cout << desc << std::endl;
    exit(1);
  }
  const std::string output_file = vm["output"].as<std::string>();
  const unsigned max_p = vm["max_p"].as<unsigned>();

  const std::filesystem::path here = __FILE__;
  // Load the unit square mesh
  const auto square_mesh = getSquareDomain();
  // Load the L-shaped domain mesh
  const auto L_mesh = getLDomain();

  // Compute the errors for different polynomial degrees
  Eigen::MatrixXd results(max_p, 7);
  for (unsigned p = 1; p <= max_p; ++p) {
    std::cout << "> Polynomial Degree: " << p << std::endl;

    // Solve the problem on the unit square domain
    std::cout << "\t> Unit Square Domain";
    const auto fe_space_square =
        std::make_shared<lf::fe::HierarchicScalarFESpace<double>>(square_mesh,
                                                                  p);
    std::cout << " (" << fe_space_square->LocGlobMap().NumDofs() << " DOFs)"
              << std::endl;
    const auto [H1_square, L2_square] =
        computeErrorsSquareDomain(p, square_mesh, fe_space_square);

    // Solve the problem on the L-shaped domain
    std::cout << "\t> L-shaped Domain";
    const auto fe_space_L =
        std::make_shared<lf::fe::HierarchicScalarFESpace<double>>(L_mesh, p);
    std::cout << " (" << fe_space_L->LocGlobMap().NumDofs() << " DOFs)"
              << std::endl;
    const auto [H1_L, L2_L] = computeErrorsLDomain(p, L_mesh, fe_space_L);

    // Store the computed quantities in the results matrix
    results(p - 1, 0) = p;
    results(p - 1, 1) = fe_space_square->LocGlobMap().NumDofs();
    results(p - 1, 2) = fe_space_L->LocGlobMap().NumDofs();
    results(p - 1, 3) = H1_square;
    results(p - 1, 4) = L2_square;
    results(p - 1, 5) = H1_L;
    results(p - 1, 6) = L2_L;
  }

  // Output the resulting errors to a file
  const static Eigen::IOFormat CSVFormat(Eigen::StreamPrecision,
                                         Eigen::DontAlignCols, ", ", "\n");
  std::ofstream file;
  file.open(output_file);
  file << results.format(CSVFormat);
  file.close();

  return 0;
}


/*
The MIT License (MIT)

Copyright (c) 2016 Alvaro Collet

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.

Neither name of this software nor the names of its contributors may be used to
endorse or promote products derived from this software without specific
prior written permission.

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.
*/

// gmm_example.cpp : Defines the entry point for the console application.
//

#include "stdafx.h"
#include "gmm/gmm.h" // Remember, you need to put the $(SolutionDir) in "Additional Include Directories"
#include "gmm/kmeans.h"
#include "gmm/math_utils.h"
#include <Eigen/Core> // Remember, you need to put the path to Eigen in "Additional Include Directories"
#include <Eigen/Geometry>

///////////////////////////////////////////////////////////////////////////////
// Basic test to show how to use KMeans
bool TestKMeans3D()
{
    using namespace AC;

    // Create some centroids
    int numCentroids = 5;
    std::vector<Vec3> centroids(numCentroids);
    float scale = 10;
    centroids[0] = Vec3(-scale, 0, 0);
    centroids[1] = Vec3(0, -scale, 0);
    centroids[2] = Vec3(0, 0, -scale);
    centroids[3] = Vec3(scale, 0, 0);
    centroids[4] = Vec3(0, scale, 0);

    // Create some observations
    int numObservationsPerCentroid = 100;
    int numObservations = numObservationsPerCentroid * numCentroids;
    std::default_random_engine generator;
    std::uniform_real_distribution<double> distribution(-scale / 4, scale / 4);
    std::vector<Vec3> observations(numObservations);
    std::vector<int> assignmentsGT(numObservations);
    auto noise = std::bind(distribution, generator);
    Vec3 observation;
    for (int k = 0; k < numCentroids; ++k) {
        for (int i = 0; i < numObservationsPerCentroid; ++i) {
            observations[numObservationsPerCentroid*k + i] = Vec3(centroids[k][0] + float(noise()), centroids[k][1] + float(noise()), centroids[k][2] + float(noise())); assignmentsGT[numObservationsPerCentroid*k + i] = k;
        }
    }

    // Run kmeans
    KMeans<Vec3> kmeans(numCentroids);
    std::vector<int> assignmentsKMeans(numObservations);
    kmeans.setMaxIterations(100);
    kmeans.Process(observations, 30 /* restart 30 times with different initializations */, assignmentsKMeans);

    // Compute average best distance between centroids and KMeans
    double avgBestDistance = 0;
    int assignment;
    for (auto& centroidGT : centroids)
        avgBestDistance += kmeans.ClosestCentroid(centroidGT, assignment);
    avgBestDistance /= (int)centroids.size();

    // Compute assignment differences between the ground truth and KMeans
    int badAssignments = 0;
    for (int i = 0; i < (int)assignmentsGT.size(); ++i) {
        for (int j = i; j < (int)assignmentsGT.size(); ++j) {
            if ((assignmentsGT[i] == assignmentsGT[j] && assignmentsKMeans[i] != assignmentsKMeans[j]) ||
                (assignmentsGT[i] != assignmentsGT[j] && assignmentsKMeans[i] == assignmentsKMeans[j])) {
                badAssignments++;
            }
        }
    }
    // Assignments are pairwise, so there is a total of N*(N-1)/2 possible bad assignments;
    int maxAssignments = (int)assignmentsGT.size() * ((int)assignmentsGT.size() - 1) / 2;
    // If centroids are within 10% distance of GT, and there are less than 10% assignment errors, declare success.
    if (avgBestDistance < scale && badAssignments < (int)maxAssignments / 10)
        return true;
    else
        return false;
}

///////////////////////////////////////////////////////////////////////////////
// Basic test for GMM. Create a noisy set of observations from a (known) multivariate gaussian distribution, fit GMM to it, and compare the differences in labeling.
bool TestGMM3D()
{
    using namespace AC;

    // Create some centroids
    int numModes = 5;
    std::vector<Vec3> centroids(numModes);
    double scale = 10;
    centroids[0] = Vec3(-scale, 0, 0);
    centroids[1] = Vec3(0, -scale, 0);
    centroids[2] = Vec3(0, 0, -scale);
    centroids[3] = Vec3(scale, 0, 0);
    centroids[4] = Vec3(0, scale, 0);

    // Create some covariance matrices, and rotate and scale them
    std::vector<Mat3> covariances(numModes);
    covariances[0] = Vec3(1.0, 1.0, 1.0).asDiagonal() * Eigen::AngleAxisd(0.0, Vec3::UnitX());
    covariances[1] = Vec3(2.0, 1.0, 0.5).asDiagonal() * Eigen::AngleAxisd(0.2, Vec3::UnitY());
    covariances[2] = Vec3(1.0, 2.0, 1.0).asDiagonal() * Eigen::AngleAxisd(0.4, Vec3::UnitZ());
    covariances[3] = Vec3(1.0, 1.0, 2.0).asDiagonal() * Eigen::AngleAxisd(0.6, Vec3::UnitX());
    covariances[4] = Vec3(0.75, 1.0, 0.75).asDiagonal() * Eigen::AngleAxisd(0.8, Vec3::UnitY());

    // Create some observations
    int numObservations = 100;
    std::default_random_engine generator;
    std::uniform_real_distribution<double> distribution(-scale / 4, scale / 4);
    std::vector<Vec3> observations(numObservations * numModes);
    std::vector<int> assignmentsGT(numObservations * numModes);
    auto noise = std::bind(distribution, generator);
    Vec3 observation;
    for (int k = 0; k < numModes; ++k) {
        for (int i = 0; i < numObservations; ++i) {
            observations[numObservations*k + i] = covariances[k] * Vec3(centroids[k][0] + noise(), centroids[k][1] + noise(), centroids[k][2] + noise());
            assignmentsGT[numObservations*k + i] = k;
        }
    }

    // Run GMM
    GMM::GMM3D gmm(numModes);
    gmm.Process(observations);

    // Compute hard assignments and probabilities
    std::vector<int> assignmentsGMM(numObservations * numModes);
    double avgProbability = 0;
    AC::OnlineMean<double> avgLogProb;
    for (int i = 0; i < (int)observations.size(); ++i) {
        gmm.ClosestMode(observations[i], assignmentsGMM[i] /* output value */);
        avgLogProb.Push(gmm.LogResponsibility(observations[i], assignmentsGMM[i]));
    }
    avgProbability = exp(avgLogProb.Mean());

    // Compute assignment differences between the ground truth and KMeans
    int badAssignments = 0;
    for (int i = 0; i < (int)assignmentsGT.size(); ++i) {
        for (int j = i; j < (int)assignmentsGT.size(); ++j) {
            if ((assignmentsGT[i] == assignmentsGT[j] && assignmentsGMM[i] != assignmentsGMM[j]) ||
                (assignmentsGT[i] != assignmentsGT[j] && assignmentsGMM[i] == assignmentsGMM[j])) {
                badAssignments++;
            }
        }
    }
    // Assignments are pairwise, so there is a total of N*(N-1)/2 possible bad assignments;
    int maxAssignments = (int)assignmentsGT.size() * ((int)assignmentsGT.size() - 1) / 2;
    // If average point probability is larger than 0.9, and there are less than 10% assignment errors, declare success.
    if (avgProbability > 0.9 && badAssignments < (int)maxAssignments / 10)
        return true;
    else
        return false;
}

///////////////////////////////////////////////////////////////////////////////
int main()
{
    bool isKMeansOK = TestKMeans3D();
    bool isGMMOK = TestGMM3D();
    return isKMeansOK && isGMMOK;
}


//==================================================================================================
/*!
  @file
  @copyright 2016 NumScale SAS

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

#include <boost/simd/function/std.hpp>
#include <boost/simd/arch/common/detail/generic/gammaln_kernel.hpp>
#include <boost/simd/constant/inf.hpp>
#include <boost/simd/constant/invpi.hpp>
#include <boost/simd/constant/logpi.hpp>
#include <boost/simd/constant/logsqrt2pi.hpp>
#include <boost/simd/constant/nan.hpp>
#include <boost/simd/constant/one.hpp>
#include <boost/simd/constant/pi.hpp>
#include <boost/simd/constant/two.hpp>
#include <boost/simd/constant/zero.hpp>
#include <boost/simd/function/abs.hpp>
#include <boost/simd/function/dec.hpp>
#include <boost/simd/function/floor.hpp>
#include <boost/simd/function/fma.hpp>
#include <boost/simd/function/inc.hpp>
#include <boost/simd/function/is_eqz.hpp>
#include <boost/simd/function/log.hpp>
#include <boost/simd/function/rec.hpp>
#include <boost/simd/function/sinpi.hpp>
#include <boost/simd/function/sqr.hpp>
#include <boost/simd/detail/dispatch/function/overload.hpp>
#include <boost/config.hpp>
#include <cmath>

namespace boost { namespace simd { namespace ext
{
  namespace bd = boost::dispatch;
  namespace bs = boost::simd;
  BOOST_DISPATCH_OVERLOAD ( gammaln_
                          , (typename A0)
                          , bd::cpu_
                          , bd::scalar_< bd::single_<A0> >
                          )
  {
    BOOST_FORCEINLINE A0 operator() (A0 a0) const BOOST_NOEXCEPT
    {
      A0 Maxgammaln = Constant<A0, 0x7bc3f8eaUL>(); //2.035093e36f
      if ((a0 > Maxgammaln) || bs::is_eqz(a0) ) return bs::Inf<A0>();
      A0 x = a0;
      A0 q = bs::abs(x);
      if( x < 0.0f )
      {
        if(q > Maxgammaln) return Nan<A0>();
        A0 w = gammaln_pos(q);
        A0 p =  bs::floor(q);
        if (p == q) return bs::Inf<A0>();
        A0 z = q - p;
        if( z > bs::Half<A0>() )
        {
          p += bs::One<A0>();
          z = p-q;
        }
        z = q*bs::sinpi(z);
        if( bs::is_eqz(z) ) return bs::Inf<A0>();
        return -bs::log(Invpi<A0>()*bs::abs(z))-w;
      }
      else
      {
        return gammaln_pos(x);
      }
    }
  private:
    static /*BOOST_FORCEINLINE*/ A0 gammaln_pos(A0 x) BOOST_NOEXCEPT
    {
      if( x < 6.5f )
      {
        A0 z = One<A0>();
        A0 tx = x;
        A0 nx = Zero<A0>();
        if( x >= 1.5f )
        {
          while( tx > 2.5f )
          {
            nx = dec(nx);
            tx = x + nx;
            z *=tx;
          }
          x += nx - Two<A0>();
          A0 p = x * detail::gammaln_kernel<A0>::gammalnB(x);
          return p+bs::log(z);

        }
        if( x >= 1.25f )
        {
          z *= x;
          x =  dec(x);
          A0 p = x *  detail::gammaln_kernel<A0>::gammalnB(x);
          return p-bs::log(z);
        }
        if( x >= 0.75f )
        {
          x = dec(x); //-= 1.0f;
          return x * detail::gammaln_kernel<A0>::gammalnC(x);
        }
        while( tx < 1.5f )
        {
          if(is_eqz(tx) ) return Inf<A0>();
          z *=tx;
          nx = inc(nx);
          tx = x + nx;
        }
        x += nx - Two<A0>();
        A0 p = x *  detail::gammaln_kernel<A0>::gammalnB(x);
        return p-log(z);
      }
      A0 q = fma(( x - 0.5f ), bs::log(x), Logsqrt2pi<A0>() - x);
      if( x <= 1.0e4f )
      {
        A0 z = rec(x);
        A0 p = sqr(z);
        q = fma(z, detail::gammaln_kernel<A0>::gammaln2(p), q);
      }
      return q;
    }
  };
  BOOST_DISPATCH_OVERLOAD ( gammaln_
                          , (typename A0)
                          , bd::cpu_
                          , bd::scalar_< bd::double_<A0> >
                          )
  {
    BOOST_FORCEINLINE A0 operator() (A0 a0) const BOOST_NOEXCEPT
    {
      A0 Maxgammaln = Constant<A0, 0x7f574c5dd06d2516ULL>();
      if ((a0 == bs::Inf<A0>()) || bs::is_eqz(a0) ) return bs::Inf<A0>(); //2.556348e305
      A0 x = a0;
      A0 q = bs::abs(x);
      if(x > Maxgammaln) return Inf<A0>();
      if( x < -34.0 )
      {
        if(q > Maxgammaln) return Nan<A0>();
        A0 w = gammaln_pos(q);
        A0 p =  bs::floor(q);
        if (p == q) return bs::Inf<A0>();
        A0 z = q - p;
        if( z > bs::Half<A0>() )
        {
          p += bs::One<A0>();
          z = p-q;
        }
        z = q*bs::sinpi(z);
        if( bs::is_eqz(z) ) return bs::Inf<A0>();
        return Logpi<A0>()-bs::log(z)-w;
      }
      else
      {
        return gammaln_pos(x);
      }
    }
  private:
    static /*BOOST_FORCEINLINE*/ A0 gammaln_pos(A0 x) BOOST_NOEXCEPT
    {
      if( x < 13.0 )
      {
        A0 z = One<A0>();
        A0 p = Zero<A0>();
        A0 u = x;
        while( u >= 3.0 )
        {
          p -= 1.0;
          u = x + p;
          z *= u;
        }
        while( u < 2.0 )
        {
          if( u == 0.0 )  return Inf<A0>();
          z /= u;
          p += 1.0;
          u = x + p;
        }
        z = bs::abs(z);
        if( u == 2.0 ) return( bs::log(z) );
        p -= 2.0;
        x = x + p;
        p = x * detail::gammaln_kernel<A0>::gammaln1(x);
        return bs::log(z) + p ;
      }
      A0 q = fma(( x - 0.5 ), bs::log(x), Logsqrt2pi<A0>()-x);
      if( x > 1.0e8 ) return q;

      A0 p = rec(sqr(x));
      q += detail::gammaln_kernel<A0>::gammalnA(p)/x;
      return q;
    }
  };
  BOOST_DISPATCH_OVERLOAD ( gammaln_
                          , (typename A0)
                          , bd::cpu_
                          , bd::scalar_< bd::floating_<A0> >
                          , bd::scalar_< bd::floating_<A0> >
                          )
  {
    BOOST_FORCEINLINE A0 operator() (A0 a0, A0&sgn) const BOOST_NOEXCEPT
    {
      sgn = signgam(a0);
      return std::lgamma(a0);
    }
  };
  BOOST_DISPATCH_OVERLOAD ( gammaln_
                          , (typename A0)
                          , bd::cpu_
                          , bs::std_tag
                          , bd::scalar_< bd::floating_<A0> >
                          )
  {
    BOOST_FORCEINLINE A0 operator() (const std_tag &, A0 a0) const BOOST_NOEXCEPT
    {
      return std::lgamma(a0);
    }
  };
} } }

#endif