/*
 * NativeCommonOps.cpp
 *
 *  Created on: Nov 27, 2018
 *      Author: Georg Wiedebach
 */

#include <jni.h>
#include <Eigen/Dense>
#include <iostream>
#include "us_ihmc_matrixlib_NativeCommonOpsWrapper.h"

using Eigen::MatrixXd;

JNIEXPORT void JNICALL Java_us_ihmc_matrixlib_NativeCommonOpsWrapper_mult(JNIEnv *env, jobject thisObj,
		jdoubleArray result, jdoubleArray aData, jdoubleArray bData, jint aRows, jint aCols, jint bCols)
{
	jdouble *aDataArray = env->GetDoubleArrayElements(aData, NULL);
	jdouble *bDataArray = env->GetDoubleArrayElements(bData, NULL);
	MatrixXd A = Eigen::Map<Eigen::Matrix<double, Eigen::Dynamic, Eigen::Dynamic, Eigen::RowMajor>>(aDataArray, aRows, aCols);
	MatrixXd B = Eigen::Map<Eigen::Matrix<double, Eigen::Dynamic, Eigen::Dynamic, Eigen::RowMajor>>(bDataArray, aCols, bCols);

	MatrixXd AB = A * B;

	jdouble *resultDataArray = new jdouble[aRows * bCols];
	Eigen::Map<Eigen::Matrix<double, Eigen::Dynamic, Eigen::Dynamic, Eigen::RowMajor>>(resultDataArray, aRows, bCols) = AB;
	env->SetDoubleArrayRegion(result, 0, aRows * bCols, resultDataArray);

	env->ReleaseDoubleArrayElements(aData, aDataArray, 0);
	env->ReleaseDoubleArrayElements(bData, bDataArray, 0);
	delete resultDataArray;
}

JNIEXPORT void JNICALL Java_us_ihmc_matrixlib_NativeCommonOpsWrapper_multQuad(JNIEnv *env, jobject thisObj,
		jdoubleArray result, jdoubleArray aData, jdoubleArray bData, jint aRows, jint aCols)
{
	jdouble *aDataArray = env->GetDoubleArrayElements(aData, NULL);
	jdouble *bDataArray = env->GetDoubleArrayElements(bData, NULL);
	MatrixXd A = Eigen::Map<Eigen::Matrix<double, Eigen::Dynamic, Eigen::Dynamic, Eigen::RowMajor>>(aDataArray, aRows, aCols);
	MatrixXd B = Eigen::Map<Eigen::Matrix<double, Eigen::Dynamic, Eigen::Dynamic, Eigen::RowMajor>>(bDataArray, aRows, aRows);

	MatrixXd AtBA = A.transpose() * B * A;

	jdouble *resultDataArray = new jdouble[aCols * aCols];
	Eigen::Map<Eigen::Matrix<double, Eigen::Dynamic, Eigen::Dynamic, Eigen::RowMajor>>(resultDataArray, aCols, aCols) = AtBA;
	env->SetDoubleArrayRegion(result, 0, aCols * aCols, resultDataArray);

	env->ReleaseDoubleArrayElements(aData, aDataArray, 0);
	env->ReleaseDoubleArrayElements(bData, bDataArray, 0);
	delete resultDataArray;
}

JNIEXPORT void JNICALL Java_us_ihmc_matrixlib_NativeCommonOpsWrapper_invert(JNIEnv *env, jobject thisObj,
		jdoubleArray result, jdoubleArray aData, jint aRows)
{
	jdouble *aDataArray = env->GetDoubleArrayElements(aData, NULL);
	MatrixXd A = Eigen::Map<Eigen::Matrix<double, Eigen::Dynamic, Eigen::Dynamic, Eigen::RowMajor>>(aDataArray, aRows, aRows);

	MatrixXd x = A.lu().inverse();

	jdouble *resultDataArray = new jdouble[aRows * aRows];
	Eigen::Map<Eigen::Matrix<double, Eigen::Dynamic, Eigen::Dynamic, Eigen::RowMajor>>(resultDataArray, aRows, aRows) = x;
	env->SetDoubleArrayRegion(result, 0, aRows * aRows, resultDataArray);

	env->ReleaseDoubleArrayElements(aData, aDataArray, 0);
	delete resultDataArray;
}

JNIEXPORT void JNICALL Java_us_ihmc_matrixlib_NativeCommonOpsWrapper_solve(JNIEnv *env, jobject thisObj,
		jdoubleArray result, jdoubleArray aData, jdoubleArray bData, jint aRows)
{
	jdouble *aDataArray = env->GetDoubleArrayElements(aData, NULL);
	jdouble *bDataArray = env->GetDoubleArrayElements(bData, NULL);
	MatrixXd A = Eigen::Map<Eigen::Matrix<double, Eigen::Dynamic, Eigen::Dynamic, Eigen::RowMajor>>(aDataArray, aRows, aRows);
	MatrixXd B = Eigen::Map<Eigen::Matrix<double, Eigen::Dynamic, Eigen::Dynamic, Eigen::RowMajor>>(bDataArray, aRows, 1);

	MatrixXd x = A.lu().solve(B);

	jdouble *resultDataArray = new jdouble[aRows];
	Eigen::Map<Eigen::Matrix<double, Eigen::Dynamic, Eigen::Dynamic, Eigen::RowMajor>>(resultDataArray, aRows, 1) = x;
	env->SetDoubleArrayRegion(result, 0, aRows, resultDataArray);

	env->ReleaseDoubleArrayElements(aData, aDataArray, 0);
	env->ReleaseDoubleArrayElements(bData, bDataArray, 0);
}

JNIEXPORT jboolean JNICALL Java_us_ihmc_matrixlib_NativeCommonOpsWrapper_solveCheck(JNIEnv *env, jobject thisObj,
		jdoubleArray result, jdoubleArray aData, jdoubleArray bData, jint aRows)
{
	jdouble *aDataArray = env->GetDoubleArrayElements(aData, NULL);
	jdouble *bDataArray = env->GetDoubleArrayElements(bData, NULL);
	MatrixXd A = Eigen::Map<Eigen::Matrix<double, Eigen::Dynamic, Eigen::Dynamic, Eigen::RowMajor>>(aDataArray, aRows, aRows);
	MatrixXd B = Eigen::Map<Eigen::Matrix<double, Eigen::Dynamic, Eigen::Dynamic, Eigen::RowMajor>>(bDataArray, aRows, 1);

	const Eigen::FullPivLU<Eigen::Matrix<double, Eigen::Dynamic, Eigen::Dynamic> > fullPivLu = A.fullPivLu();
	if (fullPivLu.isInvertible())
	{
		MatrixXd x = fullPivLu.solve(B);

		jdouble *resultDataArray = new jdouble[aRows];
		Eigen::Map<Eigen::Matrix<double, Eigen::Dynamic, Eigen::Dynamic, Eigen::RowMajor>>(resultDataArray, aRows, 1) = x;
		env->SetDoubleArrayRegion(result, 0, aRows, resultDataArray);

		delete resultDataArray;
		env->ReleaseDoubleArrayElements(aData, aDataArray, 0);
		env->ReleaseDoubleArrayElements(bData, bDataArray, 0);
		return true;
	}
	else
	{
		env->ReleaseDoubleArrayElements(aData, aDataArray, 0);
		env->ReleaseDoubleArrayElements(bData, bDataArray, 0);
		return false;
	}
}

JNIEXPORT void JNICALL Java_us_ihmc_matrixlib_NativeCommonOpsWrapper_solveRobust(JNIEnv *env, jobject thisObj,
		jdoubleArray result, jdoubleArray aData, jdoubleArray bData, jint aRows, jint aCols)
{
	jdouble *aDataArray = env->GetDoubleArrayElements(aData, NULL);
	jdouble *bDataArray = env->GetDoubleArrayElements(bData, NULL);
	MatrixXd A = Eigen::Map<Eigen::Matrix<double, Eigen::Dynamic, Eigen::Dynamic, Eigen::RowMajor>>(aDataArray, aRows, aCols);
	MatrixXd B = Eigen::Map<Eigen::Matrix<double, Eigen::Dynamic, Eigen::Dynamic, Eigen::RowMajor>>(bDataArray, aRows, 1);

	MatrixXd x = A.householderQr().solve(B);

	jdouble *resultDataArray = new jdouble[aCols];
	Eigen::Map<Eigen::Matrix<double, Eigen::Dynamic, Eigen::Dynamic, Eigen::RowMajor>>(resultDataArray, aCols, 1) = x;
	env->SetDoubleArrayRegion(result, 0, aCols, resultDataArray);

	env->ReleaseDoubleArrayElements(aData, aDataArray, 0);
	env->ReleaseDoubleArrayElements(bData, bDataArray, 0);
	delete resultDataArray;
}

JNIEXPORT void JNICALL Java_us_ihmc_matrixlib_NativeCommonOpsWrapper_solveDamped(JNIEnv *env, jobject thisObj,
		jdoubleArray result, jdoubleArray aData, jdoubleArray bData, jint aRows, jint aCols, jdouble alpha)
{
	jdouble *aDataArray = env->GetDoubleArrayElements(aData, NULL);
	jdouble *bDataArray = env->GetDoubleArrayElements(bData, NULL);
	MatrixXd A = Eigen::Map<Eigen::Matrix<double, Eigen::Dynamic, Eigen::Dynamic, Eigen::RowMajor>>(aDataArray, aRows, aCols);
	MatrixXd B = Eigen::Map<Eigen::Matrix<double, Eigen::Dynamic, Eigen::Dynamic, Eigen::RowMajor>>(bDataArray, aRows, 1);

	MatrixXd outer = A * A.transpose() + MatrixXd::Identity(aRows, aRows) * alpha * alpha;
	MatrixXd x = A.transpose() * outer.llt().solve(B);

	jdouble *resultDataArray = new jdouble[aCols];
	Eigen::Map<Eigen::Matrix<double, Eigen::Dynamic, Eigen::Dynamic, Eigen::RowMajor>>(resultDataArray, aCols, 1) = x;
	env->SetDoubleArrayRegion(result, 0, aCols, resultDataArray);

	env->ReleaseDoubleArrayElements(aData, aDataArray, 0);
	env->ReleaseDoubleArrayElements(bData, bDataArray, 0);
	delete resultDataArray;
}

JNIEXPORT void JNICALL Java_us_ihmc_matrixlib_NativeCommonOpsWrapper_projectOnNullspace(JNIEnv *env, jobject thisObj,
		jdoubleArray result, jdoubleArray aData, jdoubleArray bData, jint aRows, jint aCols, jint bRows, jdouble alpha)
{
	jdouble *aDataArray = env->GetDoubleArrayElements(aData, NULL);
	jdouble *bDataArray = env->GetDoubleArrayElements(bData, NULL);
	MatrixXd A = Eigen::Map<Eigen::Matrix<double, Eigen::Dynamic, Eigen::Dynamic, Eigen::RowMajor>>(aDataArray, aRows, aCols);
	MatrixXd B = Eigen::Map<Eigen::Matrix<double, Eigen::Dynamic, Eigen::Dynamic, Eigen::RowMajor>>(bDataArray, bRows, aCols);

	MatrixXd BtB = B.transpose() * B;
	MatrixXd outer = BtB + MatrixXd::Identity(aCols, aCols) * alpha * alpha;
	MatrixXd x = A * (MatrixXd::Identity(aCols, aCols) - outer.llt().solve(BtB));

	jdouble *resultDataArray = new jdouble[aRows * aCols];
	Eigen::Map<Eigen::Matrix<double, Eigen::Dynamic, Eigen::Dynamic, Eigen::RowMajor>>(resultDataArray, aRows, aCols) = x;
	env->SetDoubleArrayRegion(result, 0, aRows * aCols, resultDataArray);

	env->ReleaseDoubleArrayElements(aData, aDataArray, 0);
	env->ReleaseDoubleArrayElements(bData, bDataArray, 0);
	delete resultDataArray;
}


#include <cstdlib>
#include <iostream>
#include <complex>
#include <boost/numeric/ublas/vector.hpp>
#include <boost/numeric/ublas/matrix.hpp>
#include <Eigen/Core>
#include <boost/numeric/bindings/trans.hpp>
#include <boost/numeric/bindings/ublas/vector.hpp>
#include <boost/numeric/bindings/ublas/matrix.hpp>
#include <boost/numeric/bindings/eigen/vector.hpp>
#include <boost/numeric/bindings/eigen/matrix.hpp>
#include <boost/numeric/bindings/blas/level2.hpp>
#include "print.hpp"
#include "random.hpp"

namespace ublas=boost::numeric::ublas;
namespace blas=boost::numeric::bindings::blas;

int main(int argc, char *argv[]) {
  {
    typedef std::complex<double> complex;
    typedef ublas::vector<complex> vector;
    typedef ublas::matrix<complex, ublas::column_major> matrix;
    typedef vector::size_type size_type;
    rand_normal<complex>::reset();
    size_type m=6, n=8;
    matrix A(m, n);
    for (size_type j=0; j<n; ++j)
      for (size_type i=0; i<m; ++i)
 	A(i, j)=rand_normal<complex>::get();
    matrix A_t(ublas::trans(A));
    vector x(n);
    for (size_type i=0; i<n; ++i)
      x(i)=rand_normal<complex>::get();
    vector y(m);
    for (size_type i=0; i<m; ++i)
      y(i)=rand_normal<complex>::get();
    complex alpha(rand_normal<complex>::get());
    complex beta(rand_normal<complex>::get());
    vector y1(alpha*ublas::prod(A, x)+beta*y);
    vector y2(y);
    blas::gemv(alpha, A, x, beta, y2);
    vector y3(y);
    blas::gemv(alpha, blas::trans(A_t), x, beta, y3);
    std::cout << "testing boost::ublas containers\n"
	      << "using ublas           : " << print_vec(y1) << '\n'
	      << "using blas            : " << print_vec(y2) << '\n'
	      << "using blas (tranposed): " << print_vec(y3) << '\n'
	      << '\n';
  }
  {
    typedef std::complex<double> complex;
    typedef Eigen::Matrix<complex, Eigen::Dynamic, 1> vector;
    typedef Eigen::Matrix<complex, Eigen::Dynamic, Eigen::Dynamic> matrix;
    typedef int size_type;
    size_type m=6, n=8;
    rand_normal<complex>::reset();
    matrix A(m, n);
    for (size_type j=0; j<n; ++j)
      for (size_type i=0; i<m; ++i)
  	A(i, j)=rand_normal<complex>::get();
    matrix A_t(A.transpose());
    vector x(n);
    for (size_type i=0; i<n; ++i)
      x(i)=rand_normal<complex>::get();
    vector y(m);
    for (size_type i=0; i<m; ++i)
      y(i)=rand_normal<complex>::get();
    complex alpha(rand_normal<complex>::get());
    complex beta(rand_normal<complex>::get());
    vector y1(alpha*A*x+beta*y);
    vector y2(y);
    blas::gemv(alpha, A, x, beta, y2);
    vector y3(y);
    blas::gemv(alpha, blas::trans(A_t), x, beta, y3);
    std::cout << "testing eigen++ containers\n"
	      << "using eigen++         : " << print_vec(y1) << '\n'
	      << "using blas            : " << print_vec(y2) << '\n'
	      << "using blas (tranposed): " << print_vec(y3) << '\n'
	      << '\n';
  }
  return EXIT_SUCCESS;
}


/* -*- mode: c++; tab-width: 4; indent-tabs-mode: nil; c-basic-offset: 4 -*- */

/*
 Copyright (C) 2015 Johannes Goettker-Schnetmann
 Copyright (C) 2015 Klaus Spanderen

 This file is part of QuantLib, a free-software/open-source library
 for financial quantitative analysts and developers - http://quantlib.org/

 QuantLib is free software: you can redistribute it and/or modify it
 under the terms of the QuantLib license.  You should have received a
 copy of the license along with this program; if not, please email
 <quantlib-dev@lists.sf.net>. The license is also available online at
 <http://quantlib.org/license.shtml>.

 This program is distributed in the hope that it will be useful, but WITHOUT
 ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or FITNESS
 FOR A PARTICULAR PURPOSE.  See the license for more details.
*/

#include <ql/experimental/finitedifferences/squarerootprocessrndcalculator.hpp>

#include <boost/math/distributions/non_central_chi_squared.hpp>

namespace QuantLib {

    SquareRootProcessRNDCalculator::SquareRootProcessRNDCalculator(
        Real v0, Real kappa, Real theta, Real sigma)
    : v0_(v0), kappa_(kappa), theta_(theta),
      d_(4*kappa/(sigma*sigma)), df_(d_*theta) {    }

    Real SquareRootProcessRNDCalculator::pdf(Real v, Time t) const {
        const Real e   = std::exp(-kappa_*t);
        const Real k   = d_/(1-e);
        const Real ncp = k*v0_*e;

        const boost::math::non_central_chi_squared_distribution<Real>
            dist(df_, ncp);

        return boost::math::pdf(dist, v*k) * k;
    }

    Real SquareRootProcessRNDCalculator::cdf(Real v, Time t) const {
        const Real e   = std::exp(-kappa_*t);
        const Real k   = d_/(1-e);
        const Real ncp = k*v0_*e;

        const boost::math::non_central_chi_squared_distribution<Real>
            dist(df_, ncp);

        return boost::math::cdf(dist, v*k);
    }

    Real SquareRootProcessRNDCalculator::invcdf(Real q, Time t) const {
        const Real e   = std::exp(-kappa_*t);
        const Real k   = d_/(1-e);
        const Real ncp = k*v0_*e;

        const boost::math::non_central_chi_squared_distribution<Real>
            dist(df_, ncp);

        return boost::math::quantile(dist, q) / k;
    }

    Real SquareRootProcessRNDCalculator::stationary_pdf(Real v) const {
        const Real alpha = 0.5*df_;
        const Real beta = alpha/theta_;

        return std::pow(beta, alpha)*std::pow(v, alpha-1)
                *std::exp(-beta*v-boost::math::lgamma(alpha));
    }

    Real SquareRootProcessRNDCalculator::stationary_cdf(Real v) const {
        const Real alpha = 0.5*df_;
        const Real beta = alpha/theta_;

        return boost::math::gamma_p(alpha, beta*v);
    }

    Real SquareRootProcessRNDCalculator::stationary_invcdf(Real q) const {
        const Real alpha = 0.5*df_;
        const Real beta = alpha/theta_;

        return boost::math::gamma_p_inv(alpha, q)/beta;
    }
}


#ifndef _DEPTH_ESTIMATOR_H_
#define _DEPTH_ESTIMATOR_H_

#include <vector>
#include <Eigen/Dense>
#include <memory>
#include <ceres/ceres.h>

namespace depth_estimator {

  class CostFunction {
  public:
    explicit inline CostFunction(const Eigen::Vector2d& p1, const Eigen::Vector2d& p2,
                                 int idx1, int idx2,
                                 double target_dist)
      : x1_(p1[0]), y1_(p1[1]),
        x2_(p2[0]), y2_(p2[1]),
        idx1_(idx1), idx2_(idx2),
        dist_(target_dist) {}

    template <typename T>
    inline bool operator() (const T* const d, T* residual) const {
      residual[0] = T(dist_) - ceres::sqrt(ceres::pow(d[idx1_] * T(x1_) - d[idx2_] * T(x2_), 2) +
                                           ceres::pow(d[idx1_] * T(y1_) - d[idx2_] * T(y2_), 2) +
                                           ceres::pow(d[idx1_]          - d[idx2_],          2));
      return true;
    }

  private:
    const double x1_;
    const double y1_;
    const double x2_;
    const double y2_;
    const int idx1_;
    const int idx2_;
    const double dist_;
  };

  template <int N>
  class DepthEstimator {
  public:
    inline DepthEstimator(const std::vector<Eigen::Vector2d>& points_2d,
                          const std::vector<Eigen::Vector3d>& points_3d)
      : points_2d_(points_2d), points_3d_(points_3d) {}

    inline Eigen::VectorXd Estimate() {
      std::vector<double> depths(N, 1.0);

      ceres::Problem problem;

      for (int i = N - 1; i >= 1; --i) {
        for (int j = i - 1; j >= 0; --j) {
          problem.AddResidualBlock(
            new ceres::AutoDiffCostFunction<CostFunction, 1, N>(
              new CostFunction(points_2d_[i], points_2d_[j],
                               i, j,
                               (points_3d_[i] - points_3d_[j]).norm())),
            nullptr, depths.data());
        }
      }

      ceres::Solver::Options options;
      options.linear_solver_type = ceres::DENSE_QR;
      options.minimizer_progress_to_stdout = false;
      options.parameter_tolerance = 0.01;

      ceres::Solver::Summary summary;
      ceres::Solve(options, &problem, &summary);

      Eigen::VectorXd result(N);
      for (int i = 0; i < N; ++i) {
        result[i] = depths[i];
      }
      return result;
    }

  private:
    const std::vector<Eigen::Vector2d>& points_2d_;
    const std::vector<Eigen::Vector3d>& points_3d_;
  };

}

#endif


#ifndef USE_CUDA

#include <iostream>

#include <kernel.hpp>
#include <Eigen/Core>
#include <Eigen/Dense>

// C++ Version
namespace Kernel
{
    double dot(const std::vector<Eigen::Vector3d> & v1, const std::vector<Eigen::Vector3d> & v2)
    {
        double x=0;
        for (int i=0; i<v1.size(); ++i)
        {
            x += v1[i].dot(v2[i]);
        }
        return x;
    }

    void run_eigen_solver(const std::vector<Eigen::Matrix3f> &m)
    {
        for (int i = 0; i < m.size(); ++i){
            Eigen::SelfAdjointEigenSolver<Eigen::Matrix3f> es(m[i]);
            std::cout << "Matrix " << i << ":" << std::endl << m[i] << std::endl;
            std::cout << "The eigenvalues :" << std::endl << es.eigenvalues() << std::endl;
            std::cout << "The eigenvectors :" << std::endl << es.eigenvectors() << std::endl;
            std::cout << "==================================================================" << std::endl;
        }
    }
}

#endif


#include <boost/lexical_cast.hpp>
#include <boost/math/constants/constants.hpp>
#include <cmath>
#include <fstream>
#include <iomanip>
#include <iostream>
#include <string>

#include "quadrature/qmaxwell.hpp"

using namespace std;
using namespace boltzmann;

// template<typename MAP>
// void print(const MAP& m, std::ofstream& fout, string title)
// {
//   fout << "----- " << title << endl;
//   for (auto it = m.begin(); it != m.end(); ++it) {
//     fout << it->first.first << "\t"
//          << it->first.second
//          << "\t"
//          << setprecision(16) << it->second
//          << endl;
//   }
// }

// template<typename CONT>
// void print_basis(const CONT& cont, std::ofstream& fout)
// {
//   for (int i = 0; i < cont.size(); ++i) {
//     fout  << cont[i].get_id() << endl;
//   }
// }

// ----------------------------------------------------------------------
int main(int argc, char *argv[])
{
  const double beta = 2;

  if (argc < 2) {
    cerr << "info: " << argv[0] << " N" << endl << "q: No. quad. points\n";
    return 1;
  }
  int N = atoi(argv[1]);

  int digits = 128;
  QMaxwell qmaxwell(1, N, digits);

  std::ofstream fout("quadrule_order" + boost::lexical_cast<string>(N) + "_" +
                     boost::lexical_cast<string>(digits) + ".dat");
  for (int i = 0; i < qmaxwell.size(); ++i) {
    fout << setprecision(30) << qmaxwell.pts(i) << "\t" << setprecision(30) << qmaxwell.wts(i)
         << endl;
  }
  fout.close();

  // test integration
  cout << "evaluate integral: \\int_0^{2pi} \\int_0^\\infty e^{-r^} r \\dd r"
       << "\n";
  double sum = 0;
  for (int i = 0; i < N; ++i) {
    sum += qmaxwell.wts(i);
  }
  const double pi = boost::math::constants::pi<double>();
  sum *= 2 * pi;
  cout << "\terror:" << std::abs(sum - pi) << "\n";

  return 0;
}


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

#ifndef ITL_QUASI_NEWTON_INCLUDE
#define ITL_QUASI_NEWTON_INCLUDE

#include <boost/numeric/mtl/concept/collection.hpp>
#include <boost/numeric/mtl/matrix/dense2D.hpp>
#include <boost/numeric/mtl/matrix/operators.hpp>
#include <boost/numeric/mtl/operation/operators.hpp>
#include <boost/numeric/mtl/operation/trans.hpp>
#include <boost/numeric/mtl/utility/gradient.hpp>

// #include <iostream>

namespace itl {

/// Quasi-Newton method
template <typename Matrix, typename Vector, typename F, typename Grad,
	  typename Step, typename Update, typename Iter>
Vector quasi_newton(Vector& x, F f, Grad grad_f, Step step, Update update, Iter& iter)
{
    typedef typename mtl::Collection<Vector>::value_type value_type;
    Vector         d, y, x_k, s;
    Matrix         H(size(x), size(x));

    H= 1;
    for (; !iter.finished(two_norm(grad_f(x))); ++iter) {
	d= H * -grad_f(x);                                                  // std::cout << "d is " << d << '\n';
	value_type alpha= step(x, d, f, grad_f); assert(alpha == alpha);
	x_k= x + alpha * d;                                                 // std::cout << "x_k is " << x_k << '\n';
	s= alpha * d;                                                       // std::cout << "alpha is " << alpha << '\n';
	y= grad_f(x_k) - grad_f(x);
	update(H, y, s);
	x= x_k;
    }
    return x;
}

/// Quasi-Newton method
template <typename Vector, typename F, typename Grad, typename Step, typename Update, typename Iter>
Vector inline quasi_newton(Vector& x, F f, Grad grad_f, Step step, Update update, Iter& iter)
{
    typedef typename mtl::traits::gradient<Vector>::type hessian_type;
    // typedef typename mtl::Collection<Vector>::value_type value_type;
    return quasi_newton<hessian_type>(x, f, grad_f, step, update, iter);
}

} // namespace itl

#endif // ITL_QUASI_NEWTON_INCLUDE


#include <iostream>
#include <vector>

#include <boost/range/adaptors.hpp>
namespace ba = boost::adaptors;

#include <dionysus/simplex.h>
#include <dionysus/filtration.h>
#include <dionysus/omni-field-persistence.h>
#include <dionysus/diagram.h>

namespace d = dionysus;

#include <format.h>

using Simplex       = d::Simplex<>;
using Filtration    = d::Filtration<Simplex>;
using Persistence   = d::OmniFieldPersistence<>;

int main()
{
    // Klein bottle
    Filtration filtration
    {
      {0}, {1}, {2}, {3}, {4}, {5}, {6}, {7}, {8},
      {0,1}, {1,2}, {2,0}, {0,3}, {3,4}, {4,0},
      {1,5}, {5,6}, {6,2}, {2,7}, {7,8}, {8,1},
      {3,5}, {5,7}, {7,3}, {4,6}, {6,8}, {8,4},
      {0,5}, {1,7}, {2,3}, {3,6}, {5,8}, {7,4},
      {4,2}, {6,1}, {8,0},
      {0,3,5}, {0,1,5}, {1,5,7}, {1,2,7}, {2,7,3}, {2,3,0},
      {3,4,6}, {3,5,6}, {5,6,8}, {5,7,8}, {7,8,4}, {7,3,4},
      {4,0,2}, {4,6,2}, {6,2,1}, {6,8,1}, {8,1,0}, {8,4,0}
    };

    fmt::print("Boundary matrix over Q\n");
    d::Q<> q;
    for (auto& s : filtration)
    {
        fmt::print("{} at {}\n", s, filtration.index(s));
        for (auto sb : s.boundary(q))
            fmt::print("   {} * {} at {}\n", sb.element(), sb.index(), filtration.index(sb.index()));
    }

    Persistence     persistence;
    for(auto& s : filtration)
    {
        using SimplexChainEntry = d::ChainEntry<Persistence::Field, Simplex>;
        using ChainEntry        = d::ChainEntry<Persistence::Field, Persistence::Index>;
        persistence.add(s.boundary(persistence.field()) |
                                                 ba::transformed([&filtration](const SimplexChainEntry& e)
                                                 { return ChainEntry(e.element(), filtration.index(e.index())); }));
    }
    fmt::print("Reduction finished\n");

    fmt::print("Special primes:");
    for (auto x : persistence.primes())
        fmt::print(" {}", x);
    fmt::print("\n");

    unsigned i = 0;
    fmt::print("Q chains finished\n");
    for (auto& c : persistence.q_chains())
    {
        fmt::print("{}: ", i);
        for (auto& ce : c)
            fmt::print(" + {} * {}", ce.element(), ce.index());
        fmt::print("\n");
        ++i;
    }

    fmt::print("Zp chains finished\n");
    for (auto& x : persistence.zp_chains())
    {
        unsigned i = x.first;
        fmt::print("{}:\n", i);
        for (auto& ec : x.second)
        {
            auto& e = ec.first;
            auto& c = ec.second;

            fmt::print("  mod {}:", e);

            for (auto& ce : c)
                fmt::print(" + {} * {}", ce.element(), ce.index());
        }
        fmt::print("\n");
    }

    auto primes = persistence.primes();
    primes.emplace(primes.begin(), 1);
    for (auto& p : primes)
    {
        if (p == 1)
            fmt::print("Over Z_p (for all p, except those specified explicitly)\n");
        else
            fmt::print("Over Z_{}:\n", p);
        auto diagrams = init_diagrams(prime_adapter(persistence, p), filtration,
                                      [&](const Simplex& s) -> float  { return filtration.index(s); },        // inefficient, but works
                                      [](Persistence::Index i)        { return i; });
        i = 0;
        for (auto& dgm : diagrams)
        {
            fmt::print("  Dimension {}:\n", i++);
            for (auto& pt : dgm)
                fmt::print("    {} {}\n", pt.birth(), pt.death());
        }
    }
}


// distribution_construction.cpp

// Copyright Paul A. Bristow 2007, 2010.

// Use, modification and distribution are 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)

// Caution: this file contains Quickbook markup as well as code
// and comments, don't change any of the special comment markups!

//[distribution_construction1

/*`

The structure of distributions is rather different from some other statistical libraries,
for example in less object-oriented language like FORTRAN and C,
that provide a few arguments to each free function.
This library provides each distribution as a template C++ class.
A distribution is constructed with a few arguments, and then
member and non-member functions are used to find values of the
distribution, often a function of a random variate.

First we need some includes to access the negative binomial distribution
(and the binomial, beta and gamma too).

*/

#include <boost/math/distributions/negative_binomial.hpp> // for negative_binomial_distribution
  using boost::math::negative_binomial_distribution; // default type is double.
  using boost::math::negative_binomial; // typedef provides default type is double.
#include <boost/math/distributions/binomial.hpp> // for binomial_distribution.
#include <boost/math/distributions/beta.hpp> // for beta_distribution.
#include <boost/math/distributions/gamma.hpp> // for gamma_distribution.
#include <boost/math/distributions/normal.hpp> // for normal_distribution.
/*`
Several examples of constructing distributions follow:
*/
//] [/distribution_construction1 end of Quickbook in C++ markup]

int main()
{
//[distribution_construction2
/*`
First, a negative binomial distribution with 8 successes
and a success fraction 0.25, 25% or 1 in 4, is constructed like this:
*/
  boost::math::negative_binomial_distribution<double> mydist0(8., 0.25);
  /*`
  But this is inconveniently long, so we might be tempted to write
  */
  using namespace boost::math;
  /*`
  but this might risk ambiguity with names in std random so
  *much better is explicit `using boost::math:: ` * ... statements like
  */
  using boost::math::negative_binomial_distribution;
  /*`
  and we can still reduce typing.

  Since the vast majority of applications use will be using double precision,
  the template argument to the distribution (RealType) defaults
  to type double, so we can also write:
  */

  negative_binomial_distribution<> mydist9(8., 0.25); // Uses default RealType = double.

  /*`
  But the name "negative_binomial_distribution" is still inconveniently long,
  so for most distributions, a convenience typedef is provided, for example:

     typedef negative_binomial_distribution<double> negative_binomial; // Reserved name of type double.

  [caution
  This convenience typedef is /not/ provided if a clash would occur
  with the name of a function: currently only "beta" and "gamma"
  fall into this category.
  ]

  So, after a using statement,
  */

  using boost::math::negative_binomial;

  /*`
  we have a convenient typedef to `negative_binomial_distribution<double>`:
  */
  negative_binomial mydist(8., 0.25);

  /*`
  Some more examples using the convenience typedef:
  */
  negative_binomial mydist10(5., 0.4); // Both arguments double.
  /*`
  And automatic conversion takes place, so you can use integers and floats:
  */
  negative_binomial mydist11(5, 0.4); // Using provided typedef double, int and double arguments.
  /*`
  This is probably the most common usage.
  */
  negative_binomial mydist12(5., 0.4F); // Double and float arguments.
  negative_binomial mydist13(5, 1); // Both arguments integer.

  /*`
  Similarly for most other distributions like the binomial.
  */
  binomial mybinomial(1, 0.5); // is more concise than
  binomial_distribution<> mybinomd1(1, 0.5);

  /*`
  For cases when the typdef distribution name would clash with a math special function
  (currently only beta and gamma)
  the typedef is deliberately not provided, and the longer version of the name
  must be used.  For example do not use:

     using boost::math::beta;
     beta mybetad0(1, 0.5); // Error beta is a math FUNCTION!

  Which produces the error messages:

  [pre
  error C2146: syntax error : missing ';' before identifier 'mybetad0'
  warning C4551: function call missing argument list
  error C3861: 'mybetad0': identifier not found
  ]

  Instead you should use:
  */
  using boost::math::beta_distribution;
  beta_distribution<> mybetad1(1, 0.5);
  /*`
  or for the gamma distribution:
  */
  gamma_distribution<> mygammad1(1, 0.5);

  /*`
  We can, of course, still provide the type explicitly thus:
  */

  // Explicit double precision:
  negative_binomial_distribution<double>        mydist1(8., 0.25);

  // Explicit float precision, double arguments are truncated to float:
  negative_binomial_distribution<float>         mydist2(8., 0.25);

  // Explicit float precision, integer & double arguments converted to float.
  negative_binomial_distribution<float>         mydist3(8, 0.25);

  // Explicit float precision, float arguments, so no conversion:
  negative_binomial_distribution<float>         mydist4(8.F, 0.25F);

  // Explicit float precision, integer arguments promoted to float.
  negative_binomial_distribution<float>         mydist5(8, 1);

  // Explicit double precision:
  negative_binomial_distribution<double>        mydist6(8., 0.25);

  // Explicit long double precision:
  negative_binomial_distribution<long double>   mydist7(8., 0.25);

  /*`
  And if you have your own RealType called MyFPType,
  for example NTL RR (an arbitrary precision type), then we can write:

     negative_binomial_distribution<MyFPType>  mydist6(8, 1); // Integer arguments -> MyFPType.

  [heading Default arguments to distribution constructors.]

  Note that default constructor arguments are only provided for some distributions.
  So if you wrongly assume a default argument you will get an error message, for example:

     negative_binomial_distribution<> mydist8;

  [pre error C2512 no appropriate default constructor available.]

  No default constructors are provided for the negative binomial,
  because it is difficult to chose any sensible default values for this distribution.
  For other distributions, like the normal distribution,
  it is obviously very useful to provide 'standard'
  defaults for the mean and standard deviation thus:

      normal_distribution(RealType mean = 0, RealType sd = 1);

  So in this case we can write:
  */
  using boost::math::normal;

  normal norm1;       // Standard normal distribution.
  normal norm2(2);    // Mean = 2, std deviation = 1.
  normal norm3(2, 3); // Mean = 2, std deviation = 3.

  return 0;
}  // int main()

/*`There is no useful output from this program, of course. */

//] [/end of distribution_construction2]


#include <iostream>

#include <elemental.hpp>
#include <boost/mpi.hpp>
#include <boost/format.hpp>
#include <skylark.hpp>

/*******************************************/
namespace bmpi =  boost::mpi;
namespace skybase = skylark::base;
namespace skysk =  skylark::sketch;
namespace skynla = skylark::nla;
namespace skyalg = skylark::algorithms;
namespace skyutil = skylark::utility;
/*******************************************/

const int m = 50000;
const int n = 500;

typedef elem::DistMatrix<double, elem::VC, elem::STAR> matrix_type;
typedef elem::DistMatrix<double, elem::VC, elem::STAR> rhs_type;
typedef elem::DistMatrix<double, elem::STAR, elem::STAR> sol_type;

template<typename MatrixType, typename RhsType, typename SolType>
void check_solution(const MatrixType &A, const RhsType &b, const SolType &x,
    const RhsType &r0,
    double &res, double &resAtr, double &resFac) {
    RhsType r(b);
    skybase::Gemv(elem::NORMAL, -1.0, A, x, 1.0, r);
    res = skybase::Nrm2(r);

    SolType Atr(x.Height(), x.Width(), x.Grid());
    skybase::Gemv(elem::TRANSPOSE, 1.0, A, r, 0.0, Atr);
    resAtr = skybase::Nrm2(Atr);

    skybase::Axpy(-1.0, r0, r);
    RhsType dr(b);
    skybase::Axpy(-1.0, r0, dr);
    resFac = skybase::Nrm2(r) / skybase::Nrm2(dr);
}

int main(int argc, char** argv) {
    double res, resAtr, resFac;

    elem::Initialize(argc, argv);

    bmpi::communicator world;
    int rank = world.rank();

    skybase::context_t context(23234);

    // Setup problem and righthand side
    // Using Skylark's uniform generator (as opposed to Elemental's)
    // will insure the same A and b are generated regardless of the number
    // of processors.
    matrix_type A =
        skyutil::uniform_matrix_t<matrix_type>::generate(m,
            n, elem::DefaultGrid(), context);
    matrix_type b =
        skyutil::uniform_matrix_t<matrix_type>::generate(m,
            1, elem::DefaultGrid(), context);

    sol_type x(n,1);
    rhs_type r(b);

    boost::mpi::timer timer;
    double telp;

    // Solve using Elemental. Note: Elemental only supports [MC,MR]...
    elem::DistMatrix<double> A1 = A, b1 = b, x1;
    timer.restart();
    elem::LeastSquares(elem::NORMAL, A1, b1, x1);
    telp = timer.elapsed();
    x = x1;
    check_solution(A, b, x, r, res, resAtr, resFac);
    if (rank == 0)
        std::cout << "Elemental:\t\t\t||r||_2 =  "
                  << boost::format("%.2f") % res
                  << "\t\t\t\t\t\t\t||A' * r||_2 = " << boost::format("%.2e") % resAtr
                  << "\t\tTime: " << boost::format("%.2e") % telp << " sec"
                  << std::endl;
    double res_opt = res;

    skybase::Gemv(elem::NORMAL, -1.0, A, x, 1.0, r);

    // Solve using Sylark
    timer.restart();
    skynla::FastLeastSquares(elem::NORMAL, A, b, x, context);
    telp = timer.elapsed();
    check_solution(A, b, x, r, res, resAtr, resFac);
    if (rank == 0)
        std::cout << "Skylark:\t\t\t||r||_2 =  "
                  << boost::format("%.2f") % res
                  << " (x " << boost::format("%.5f") % (res / res_opt) << ")"
                  << "\t||r - r*||_2 / ||b - r*||_2 = " << boost::format("%.2e") % resFac
                  << "\t||A' * r||_2 = " << boost::format("%.2e") % resAtr
                  << "\t\tTime: " << boost::format("%.2e") % telp << " sec"
                  << std::endl;

    // Approximately solve using Sylark
    timer.restart();
    skynla::ApproximateLeastSquares(elem::NORMAL, A, b, x, context);
    telp = timer.elapsed();
    check_solution(A, b, x, r, res, resAtr, resFac);
    if (rank == 0)
        std::cout << "Skylark (approximate):\t\t||r||_2 =  "
                  << boost::format("%.2f") % res
                  << " (x " << boost::format("%.5f") % (res / res_opt) << ")"
                  << "\t||r - r*||_2 / ||b - r*||_2 = " << boost::format("%.2e") % resFac
                  << "\t||A' * r||_2 = " << boost::format("%.2e") % resAtr
                  << "\t\tTime: " << boost::format("%.2e") % telp << " sec"
                  << std::endl;

    return 0;
}


#ifndef NCA_HPP
#define NCA_HPP

#include <iostream>
#include <limits>
#include <vector>

#include <Eigen/Core>

void nearest_neighbors(const std::vector<Eigen::VectorXd>& input, const std::vector<std::string>& label) {
    unsigned int correct = 0;

    for(unsigned int i = 0; i < input.size(); ++i) {
        double min_norm = std::numeric_limits<double>::infinity();
        std::string min_norm_label;
        for(unsigned int j = 0; j < input.size(); ++j) {
            if(i == j) continue;
            double norm = (input[i] - input[j]).norm();
            if(norm < min_norm) {
                min_norm = norm;
                min_norm_label = label[j];
            }
        }

        if(label[i] == min_norm_label) {
            ++correct;
        }
    }

    std::cout << "Got " << correct << " correct out of " << input.size() << std::endl;
}

Eigen::MatrixXd scaling_matrix(const std::vector<Eigen::VectorXd>& input) {
    Eigen::MatrixXd A;
    if(input.size() == 0) return A;

    int size = input[0].size();

    std::vector< std::pair<double, double> > minmax(size, std::make_pair(std::numeric_limits<double>::infinity(), -std::numeric_limits<double>::infinity()));
    for(std::vector<Eigen::VectorXd>::const_iterator i = input.begin(); i != input.end(); ++i) {
        for(int j = 0; j < i->size(); ++j) {
            double val = (*i)[j];
            if(val < minmax[j].first) {
                minmax[j] = std::make_pair(val, minmax[j].second);
            }
            if(val > minmax[j].second) {
                minmax[j] = std::make_pair(minmax[j].first, val);
            }
        }
    }

    A = Eigen::MatrixXd::Identity(size, size);
    for(unsigned int i = 0; i < minmax.size(); ++i) {
        A(i, i) = 1.0/(minmax[i].second - minmax[i].first);
    }

    return A;
}

std::vector<Eigen::VectorXd> scale(const Eigen::MatrixXd& ScaleA, const std::vector<Eigen::VectorXd>& input) {
    std::vector<Eigen::VectorXd> scaled_input;
    for(std::vector<Eigen::VectorXd>::const_iterator i = input.begin(); i != input.end(); ++i) {
        scaled_input.push_back(ScaleA * (*i));
    }

    return scaled_input;
}

Eigen::MatrixXd neighborhood_components_analysis(const std::vector<Eigen::VectorXd>& input, const std::vector<std::string>& label, const Eigen::MatrixXd& init, unsigned int iterations, double learning_rate) {
    Eigen::MatrixXd A = init;
    for(unsigned int it = 0; it < iterations; ++it) {
        unsigned int i = it % input.size();

        double softmax_normalization = 0.0;
        for(unsigned int k = 0; k < input.size(); ++k) {
            if(k == i) continue;
            softmax_normalization += std::exp(-(A*input[i] - A*input[k]).squaredNorm());
        }

        std::vector<double> softmax;
        for(unsigned int k = 0; k < input.size(); ++k) {
            if(k == i) softmax.push_back(0.0);
            else {
                softmax.push_back(std::exp(-(A*input[i] - A*input[k]).squaredNorm()) / softmax_normalization);
            }
        }

        double p = 0.0;
        for(unsigned int k = 0; k < softmax.size(); ++k) {
            if(label[k] == label[i]) p += softmax[k];
        }

        Eigen::MatrixXd first_term = Eigen::MatrixXd::Zero(input[0].size(), input[0].size());
        Eigen::MatrixXd second_term = Eigen::MatrixXd::Zero(input[0].size(), input[0].size());
        for(unsigned int k = 0; k < input.size(); ++k) {
            if(k == i) continue;
            Eigen::VectorXd xik = input[i] - input[k];
            Eigen::MatrixXd term = softmax[k] * (xik * xik.transpose());

            first_term += term;
            if(label[k] == label[i]) second_term += term;
        }
        first_term *= p;

        A += learning_rate*A*(first_term - second_term);
    }

    return A;
}

#endif // NCA_HPP


/* +------------------------------------------------------------------------+
   |                     Mobile Robot Programming Toolkit (MRPT)            |
   |                          https://www.mrpt.org/                         |
   |                                                                        |
   | Copyright (c) 2005-2022, Individual contributors, see AUTHORS file     |
   | See: https://www.mrpt.org/Authors - All rights reserved.               |
   | Released under BSD License. See: https://www.mrpt.org/License          |
   +------------------------------------------------------------------------+ */

#include "vision-precomp.h"	 // Precompiled headers
//
#include <Eigen/Dense>
#include <Eigen/SVD>
#include <iostream>

#include "posit.h"

mrpt::vision::pnp::posit::posit(
	Eigen::MatrixXd obj_pts_, Eigen::MatrixXd img_pts_,
	Eigen::MatrixXd camera_intrinsic_, int n0)
{
	obj_pts = obj_pts_;
	img_pts = img_pts_.block(0, 0, n0, 2);
	cam_intrinsic = camera_intrinsic_;
	R = Eigen::MatrixXd::Identity(3, 3);
	t = Eigen::VectorXd::Zero(3);
	f = (cam_intrinsic(0, 0) + cam_intrinsic(1, 1)) / 2;

	obj_matrix =
		(obj_pts.transpose() * obj_pts).inverse() * obj_pts.transpose();

	n = n0;

	obj_vecs = Eigen::MatrixXd::Zero(n0, 3);

	for (int i = 0; i < n; i++)
		obj_vecs.row(i) = obj_pts.row(i) - obj_pts.row(0);

	img_vecs = Eigen::MatrixXd::Zero(n0, 2);
	img_vecs_old = img_vecs;

	epsilons = Eigen::VectorXd::Zero(n);
}

void mrpt::vision::pnp::posit::POS()
{
	Eigen::Vector3d I0, J0, r1, r2, r3;
	double I0_norm, J0_norm;

	int i;
	double scale;

	for (i = 0; i < 3; i++)
	{
		I0(i) = obj_matrix.row(i).dot(img_vecs.col(0));
		J0(i) = obj_matrix.row(i).dot(img_vecs.col(1));
	}

	I0_norm = I0.norm();
	J0_norm = J0.norm();

	scale = (I0_norm + J0_norm) / 2;

	/*Computing TRANSLATION */
	t(0) = img_pts(0, 0) / scale;
	t(1) = img_pts(0, 1) / scale;
	t(2) = f / scale;

	/* Computing ROTATION */
	r1 = I0 / I0_norm;
	r2 = J0 / J0_norm;
	r3 = r1.cross(r2);

	R.row(0) = r1;
	R.row(1) = r2;
	R.row(2) = r3;
}

/**
Iterate over results obtained by the POS function;
see paper "Model-Based Object Pose in 25 Lines of Code", IJCV 15, pp. 123-141,
1995.
*/
bool mrpt::vision::pnp::posit::compute_pose(
	Eigen::Ref<Eigen::Matrix3d> R_, Eigen::Ref<Eigen::Vector3d> t_)
{
	Eigen::FullPivLU<Eigen::MatrixXd> lu(obj_pts);
	if (lu.rank() < 3) return false;

	int i, iCount;
	long imageDiff = 1000;

	for (iCount = 0; iCount < LOOP_MAX_COUNT; iCount++)
	{
		if (iCount == 0)
		{
			for (i = 0; i < img_vecs.rows(); i++)
				img_vecs.row(i) = img_pts.row(i) - img_pts.row(0);
		}

		else
		{
			// Compute new image vectors
			epsilons.setZero();
			for (i = 0; i < n; i++)
			{
				epsilons(i) += obj_vecs.row(i).dot(R.row(2));
			}
			epsilons /= t(2);

			// Corrected image vectors
			for (i = 0; i < n; i++)
			{
				img_vecs.row(i) =
					img_pts.row(i) * (1 + epsilons(i)) - img_pts.row(0);
			}

			imageDiff = this->get_img_diff();
		}

		img_vecs_old = img_vecs;

		this->POS();

		if (iCount > 0 && imageDiff == 0) break;

		if (iCount == LOOP_MAX_COUNT)
		{
			std::cout << "Solution Not converged" << std::endl << std::endl;
			break;
		}
	}
	R_ = R;
	t_ = t;

	return true;
}

long mrpt::vision::pnp::posit::get_img_diff()
{
	double sumOfDiffs = 0;

	for (int i = 0; i < n; i++)
	{
		for (int j = 0; j < 2; j++)
		{
			sumOfDiffs += std::abs(
				floor(0.5 + img_vecs(i, j)) - floor(0.5 + img_vecs_old(i, j)));
		}
	}
	return static_cast<long>(sumOfDiffs);
}


#define GLM_ENABLE_EXPERIMENTAL 1
#include <Eigen/Dense>
#include <Eigen/Geometry>
#include <Halide.h>
#include <fstream>
#include <glm/ext.hpp>
#include <glm/glm.hpp>
#include <halide_image_io.h>
#include <stdio.h>
#include <iostream>
#include <ImfRgbaFile.h>
#include <ImfStringAttribute.h>
#include <ImfMatrixAttribute.h>
#include <ImfArray.h>
#include <algorithm>
#include <ImfNamespace.h>
#include <tuple>
#include <math.h>

namespace IMF = OPENEXR_IMF_NAMESPACE;

using namespace IMF;
using namespace IMATH_NAMESPACE;

using namespace Halide;
using namespace Halide::Tools;
using namespace Eigen;
using namespace std;

using Vector4h = Matrix<Halide::Expr, 4, 1>;
using Matrix4h = Matrix<Halide::Expr, 4, 4>;

Var x, y, c, i, ii, xo, yo, xi, yi, img;

ostream &operator<<(ostream &os, const Matrix4h &m)
{
    os << m(0, 0) << " " << m(0, 1) << " " << m(0, 2) << " " << m(0, 3) << "\n"
       << m(1, 0) << " " << m(1, 1) << " " << m(1, 2) << " " << m(1, 3) << "\n"
       << m(2, 0) << " " << m(2, 1) << " " << m(2, 2) << " " << m(2, 3) << "\n"
       << m(3, 0) << " " << m(3, 1) << " " << m(3, 2) << " " << m(3, 3) << endl;
    return os;
}

ostream &operator<<(ostream &os, const Vector4h &v)
{
    os << v(0) << "\n"
       << v(1) << "\n"
       << v(2) << "\n"
       << v(3) << endl;
    return os;
}

Matrix4h getCameraMatrix(double focalLen, double pxDim, int width, int height)
{

    focalLen = focalLen * 10e-3;
    double u0 = width / 2.0;
    double v0 = height / 2.0;
    Matrix4h camMat;
    double diagEntry = focalLen / pxDim;
    camMat << diagEntry, 0.0d, u0, 0.0d,
        0.0d, diagEntry, v0, 0.0d,
        0.0d, 0.0d, 1.0d, 0.0d,
        0.0d, 0.0d, 0.0d, 1.0d;

    return camMat;
}

Matrix4h getInvCameraMat(double focalLen, double pxDim, int width, int height)
{
    focalLen = focalLen * 10e-3;
    double u0 = width / 2.0;
    double v0 = height / 2.0;
    Matrix4h camMat;
    camMat << pxDim / focalLen, 0.0d, -pxDim * u0 / focalLen, 0.0d,
        0.0d, pxDim / focalLen, -pxDim * v0 / focalLen, 0.0d,
        0.0d, 0.0d, 1.0d, 0.0d,
        0.0d, 0.0d, 0.0d, 1.0d;

    return camMat;
}

Matrix4h getTransMatProjToCam()
{
    Matrix4h transMat;
    Expr f64 = cast<double>(0.1d);
    transMat << cast<double>(0.9945219), 0.0d, -0.10452846d, -0.2d,
        0.0d, 1.0d, 0.0d, 0.0d,
        0.10452846d, 0.0d, 0.9945219d, 0.0d,
        0.0d, 0.0d, 0.0d, 1.0d;
    return transMat;
}

tuple<int, int> readOpenEXR(const char filename[], Array2D<Rgba> &pixels)
{
    RgbaInputFile file(filename);
    Box2i dw = file.dataWindow();

    int width = dw.max.x - dw.min.x + 1;
    int height = dw.max.y - dw.min.y + 1;
    cout << "Width: " << width << "    Height: " << height << endl;
    tuple<int, int> dim(width, height);
    pixels.resizeErase(height, width);

    file.setFrameBuffer(&pixels[0][0] - dw.min.x - dw.min.y * width, 1, width);
    file.readPixels(dw.min.y, dw.max.y);
    return dim;
}

void exrArrayToHalideBuffer(Array2D<Rgba> &pixels, Buffer<double> &halideBuffer, int width, int height)
{
    for (int row{0}; row < height; row++)
    {
        for (int col{0}; col < width; col++)
        {
            isinf(pixels[row][col].r) ? pixels[row][col].r = 0.0f : pixels[row][col].r = pixels[row][col].r;
            halideBuffer(col, row) = pixels[row][col].r;
        }
    }
}

void saveImage(Expr result, size_t width, size_t height, const string &basename)
{
    Target target = get_host_target();
    Func byteResult;
    byteResult(x, y) = cast<uint8_t>(clamp(result, 0.0f, 1.0f) * 255.0f);
    byteResult.compile_jit(target);
    Buffer<uint8_t> output(width, height);
    byteResult.realize(output);
    stringstream filename;
    filename << basename << ".png";
    save_image(output, filename.str());
}

void saveImageEXR(Expr result, int width, int height, const char fileName[])
{
    Target target = get_host_target();
    Func byteResult;
    byteResult(x, y) = result;
    Buffer<double> output(width, height);
    byteResult.compile_jit(target);
    byteResult.realize(output);

    Array2D<Rgba> pixels;
    pixels.resizeErase(height, width);
    for (int row{0}; row < height; row++)
    {
        for (int col{0}; col < width; col++)
        {
            pixels[row][col].r = output(col, row);
            pixels[row][col].g = output(col, row);
            pixels[row][col].b = output(col, row);
        }
    }
    RgbaOutputFile file(fileName, width, height, WRITE_RGBA);
    file.setFrameBuffer(&pixels[0][0], 1, width);
    file.writePixels(height);
}

void debugImageEXR(Expr channel1Expr, Expr channel2Expr, Expr channel3Expr, int width, int height, const char fileName[])
{
    Target target = get_host_target();
    Func byteResult;
    byteResult(x, y, c) = 0.0f;
    byteResult(x, y, 0) = channel1Expr;
    byteResult(x, y, 1) = channel2Expr;
    byteResult(x, y, 2) = channel3Expr;
    Buffer<double> output(width, height, 3);
    byteResult.compile_jit(target);
    byteResult.realize(output);

    Array2D<Rgba> pixels;
    pixels.resizeErase(height, width);
    for (int row{0}; row < height; row++)
    {
        for (int col{0}; col < width; col++)
        {
            pixels[row][col].r = output(col, row, 0);
            pixels[row][col].g = output(col, row, 1);
            pixels[row][col].b = output(col, row, 2);
        }
    }
    RgbaOutputFile file(fileName, width, height, WRITE_RGBA);
    file.setFrameBuffer(&pixels[0][0], 1, width);
    file.writePixels(height);
}

int main(int argc, char **argv)
{
    const string INPUTFILE = "depth-image-gt.exr";
    const double FOCAL_LEN = 36.1;
    const double PX_DIM = 10 * 10e-5;

    Array2D<Rgba> pixels;
    int width, height;
    tie(width, height) = readOpenEXR(INPUTFILE.c_str(), pixels);
    Buffer<double> input(width, height);
    exrArrayToHalideBuffer(pixels, input, width, height);
    cout << "Width: " << width << "  Height: " << height << endl;
    Matrix4h camMat = getCameraMatrix(FOCAL_LEN, PX_DIM, width, height);
    Matrix4h camMatInv = getInvCameraMat(FOCAL_LEN, PX_DIM, width, height);
    Matrix4h transfMatProjToCam = getTransMatProjToCam();
    cout << "Camera matrix" << endl;
    cout << camMat << endl;
    cout << "Inverse camera matrix" << endl;
    cout << camMatInv << endl;
    cout << camMatInv(0, 0) << endl;
    cout << "Transformation matrix projector to camera" << endl;
    cout << transfMatProjToCam << endl;

    Expr zDepthCam = input(x, y);
    Vector4h pxCam{x, y, 1.0f, 0.0f};
    Vector4h normCam = camMatInv * pxCam;
    debugImageEXR(normCam(0), normCam(1), normCam(2), width, height, "normCam.exr");
    Vector4h ptCam = normCam / normCam(2) * zDepthCam;
    ptCam(3) = 1.0f;
    debugImageEXR(ptCam(0), ptCam(1), ptCam(2), width, height, "ptCam.exr");
    Vector4h ptProj = transfMatProjToCam * ptCam;
    debugImageEXR(ptProj(0), ptProj(1), ptProj(2), width, height, "ptProj.exr");
    Vector4h normProj = ptProj / ptProj(2);
    ptProj(3) = 0.0f;
    Vector4h pxProj = camMat * normProj;
    //normProj = normProj / normProj(2);
    debugImageEXR(pxProj(0), pxProj(1), pxProj(2), width, height, "normProj.exr");
    Expr xPxProj = pxProj(0);
    //saveImage(xValProj, width, height, "x-val-proj-gt");
    saveImageEXR(xPxProj, width, height, "x-val-proj-gt.exr");

    return 0;
}


#ifndef _COCONUT_PULP_MATH_MATRIX_HPP_
#define _COCONUT_PULP_MATH_MATRIX_HPP_

#include <array>
#include <cmath>
#include <type_traits>
#include <iosfwd>
#include <algorithm>
#include <functional>

#include <boost/operators.hpp>

#include "coconut-tools/utils/InfixOstreamIterator.hpp"

#include "Angle.hpp"
#include "Handedness.hpp"
#include "Vector.hpp"
#include "ScalarEqual.hpp"

namespace coconut {
namespace pulp {
namespace math {

template <class NextViewType, class GetElementFunc>
class MatrixView {
public:

	using MatrixType = typename GetElementFunc::MatrixType;

	using Scalar = typename MatrixType::Scalar;

	static const auto ROWS = GetElementFunc::ROWS;

	static const auto COLUMNS = GetElementFunc::COLUMNS;

	constexpr MatrixView(NextViewType nextView, GetElementFunc getElementFunc = GetElementFunc()) :
		next_(nextView), // TODO: can't use std::move here when NextViewType is a reference. Figure out this view type better.
		getElementFunc_(std::move(getElementFunc))
	{
	}

	constexpr auto get(size_t row, size_t column) const noexcept -> decltype(auto) {
		return getElementFunc_(next_, row, column);
	}

	template <size_t ROWS_ = ROWS>
	auto determinant() const noexcept {
		static_assert(ROWS_ == ROWS, "Rows count changed");
		static_assert(ROWS == COLUMNS, "Determinant only available for square matrices");

		auto result = MatrixType::Scalar(0);
		for (size_t columnIndex = 0; columnIndex < COLUMNS; ++columnIndex) {
			const auto absElement = get(0, columnIndex) * submatrix(*this, 0, columnIndex).determinant();
			if (columnIndex % 2 == 0) {
				result += absElement;
			} else {
				result -= absElement;
			}
		}
		return result;
	}

	template <>
	auto determinant<1>() const noexcept {
		static_assert(ROWS == COLUMNS, "Determinant only available for square matrices");
		return get(0, 0);
	}

	auto cofactor(size_t rowIndex, size_t columnIndex) const noexcept {
		const auto det = submatrix(*this, rowIndex, columnIndex).determinant();
		return ((rowIndex + columnIndex) % 2 == 0) ? det : -det;
	}

	template <size_t ROWS_ = ROWS>
	auto inverse() const noexcept {
		static_assert(ROWS_ == ROWS, "Rows count changed");
		static_assert(ROWS == COLUMNS, "Inverse only available for square matrices");

		const auto det = determinant();
		assert(det != Scalar(0));
		const auto detInverse = Scalar(1) / det;

		auto result = MatrixType();

		for (size_t rowIndex = 0; rowIndex < COLUMNS; ++rowIndex) {
			for (size_t columnIndex = 0; columnIndex < COLUMNS; ++columnIndex) {
				result[rowIndex][columnIndex] = detInverse * cofactor(columnIndex, rowIndex);
			}
		}

		return result;
	}

private:

	NextViewType next_;

	GetElementFunc getElementFunc_;

};

template <class MatrixT>
class MatrixViewFunc {
public:

	using MatrixType = MatrixT;

	static const auto ROWS = MatrixType::ROWS;

	static const auto COLUMNS = MatrixType::COLUMNS;

	constexpr auto operator()(MatrixType& matrix, size_t row, size_t column) const noexcept
		-> decltype(auto)
	{
		return matrix[row][column];
	}

};

template <class MatrixType>
constexpr auto viewMatrix(MatrixType& matrix) noexcept {
	return MatrixView<MatrixType&, MatrixViewFunc<MatrixType>>(matrix);
}

template <class NextViewType>
class TransposedViewFunc {
public:

	using MatrixType = typename NextViewType::MatrixType;

	static const auto ROWS = NextViewType::ROWS;

	static const auto COLUMNS = NextViewType::COLUMNS;

	constexpr auto operator()(NextViewType nextView, size_t row, size_t column) const noexcept
		-> decltype(auto)
	{
		return nextView.get(column, row);
	}

};

template <class NextViewType>
constexpr auto transpose(NextViewType nextView) noexcept {
	return MatrixView<NextViewType, TransposedViewFunc<NextViewType>>(std::move(nextView));
}

template <class MatrixType>
constexpr auto viewMatrixTransposed(MatrixType& matrix) noexcept {
	return transpose(viewMatrix(matrix));
}

template <class NextViewType>
class SubmatrixViewFunc {
public:

	using MatrixType = typename NextViewType::MatrixType;

	static const auto ROWS = NextViewType::ROWS - 1;

	static const auto COLUMNS = NextViewType::COLUMNS - 1;

	constexpr SubmatrixViewFunc(size_t noRow, size_t noColumn) :
		noRow_(noRow),
		noColumn_(noColumn)
	{
		assert(noRow < NextViewType::ROWS);
		assert(noColumn < NextViewType::COLUMNS);
	}

	constexpr auto operator()(NextViewType nextView, size_t row, size_t column) const noexcept
		-> decltype(auto)
	{
		return nextView.get((row < noRow_) ? row : row + 1, (column < noColumn_) ? column : column + 1);
	}

private:

	size_t noRow_;

	size_t noColumn_;

};

template <class NextViewType>
constexpr auto submatrix(NextViewType nextView, size_t noRow, size_t noColumn) noexcept {
	return MatrixView<NextViewType, SubmatrixViewFunc<NextViewType>>(
		std::move(nextView), SubmatrixViewFunc<NextViewType>(noRow, noColumn));
}

template <class MatrixType>
constexpr auto viewSubmatrix(MatrixType& matrix, size_t noRow, size_t noColumn) noexcept {
	return submatrix(viewMatrix(matrix), noRow, noColumn);
}

template <
	class ScalarType,
	size_t ROWS_PARAM,
	size_t COLUMNS_PARAM,
	class ScalarEqualityFunc = ScalarEqual<ScalarType>
	>
class Matrix :
	boost::equality_comparable<Matrix<ScalarType, ROWS_PARAM, COLUMNS_PARAM, ScalarEqualityFunc>,
	boost::additive<Matrix<ScalarType, ROWS_PARAM, COLUMNS_PARAM, ScalarEqualityFunc>,
	boost::multiplicative<Matrix<ScalarType, ROWS_PARAM, COLUMNS_PARAM, ScalarEqualityFunc>, ScalarType
	>>>
{
public:

	using Scalar = ScalarType;

	static const auto ROWS = ROWS_PARAM;

	static const auto COLUMNS = COLUMNS_PARAM;

	static const Matrix IDENTITY;

	using Row = Vector<ScalarType, COLUMNS, ScalarEqualityFunc>;

	using Column = Vector<ScalarType, ROWS, ScalarEqualityFunc>;

	static constexpr auto IS_ROW_MAJOR = true;

	static constexpr auto IS_COLUMN_MAJOR = !IS_ROW_MAJOR;

	static constexpr auto VECTOR_IS_SINGLE_ROW_MATRIX = false;

	static constexpr auto VECTOR_IS_SINGLE_COLUMN_MATRIX = !VECTOR_IS_SINGLE_ROW_MATRIX;

	// --- CONSTRUCTORS AND OPERATORS

	template <class... CompatibleVectorType>
	explicit constexpr Matrix(CompatibleVectorType&&... rows) noexcept :
		elements_{ std::forward<CompatibleVectorType>(rows)... }
	{
		static_assert(sizeof...(rows) == ROWS, "Bad number of arguments");
	}

	template <class... CompatibleScalarType>
	explicit constexpr Matrix(std::initializer_list<CompatibleScalarType>... rows) noexcept :
		elements_{ rows... }
	{
		static_assert(sizeof...(rows) == ROWS || sizeof...(rows) == 0, "Bad number of arguments");
	}

	template <class NVT, class GEF>
	Matrix(const MatrixView<NVT, GEF>& view) {
		for (size_t rowIndex = 0; rowIndex < ROWS; ++rowIndex) {
			for (size_t columnIndex = 0; columnIndex < COLUMNS; ++columnIndex) {
				(*this)[rowIndex][columnIndex] = view.get(rowIndex, columnIndex);
			}
		}
	}

	friend std::ostream& operator<<(std::ostream& os, const Matrix& matrix) {
		os << '<';
		std::copy(matrix.elements_.begin(), matrix.elements_.end(),
			coconut_tools::InfixOstreamIterator<Row>(os, ", "));
		os << '>';
		return os;
	}

	friend bool operator==(const Matrix& lhs, const Matrix& rhs) noexcept {
		return std::equal(lhs.elements_.begin(), lhs.elements_.end(), rhs.elements_.begin());
	}

	Matrix& operator+=(const Matrix& other) noexcept {
		std::transform(elements_.begin(), elements_.end(), other.elements_.begin(), elements_.begin(), std::plus<>());
		return *this;
	}

	Matrix& operator-=(const Matrix& other) noexcept {
		std::transform(elements_.begin(), elements_.end(), other.elements_.begin(), elements_.begin(), std::minus<>());
		return *this;
	}

	Matrix operator-() noexcept {
		auto result = Matrix();
		std::transform(elements_.begin(), elements_.end(), result.elements_.begin(), std::negate<>());
		return result;
	}

	template <class CompatibleMatrixType>
	friend std::enable_if_t<
		COLUMNS == CompatibleMatrixType::ROWS,
		Matrix<ScalarType, ROWS, CompatibleMatrixType::COLUMNS, ScalarEqualityFunc>
		> operator*(const Matrix& lhs, const CompatibleMatrixType& rhs) noexcept
	{
		auto result = Matrix<ScalarType, ROWS, CompatibleMatrixType::COLUMNS, ScalarEqualityFunc>();
		for (auto rowIndex = 0u; rowIndex < ROWS; ++rowIndex) {
			for (auto columnIndex = 0u; columnIndex < CompatibleMatrixType::COLUMNS; ++columnIndex) {
				result[rowIndex][columnIndex] = dot(lhs[rowIndex], rhs.column(columnIndex));
			}
		}
		return result;
	}

	template <class CompatibleMatrixType>
	std::enable_if_t<
		COLUMNS == ROWS &&
		CompatibleMatrixType::COLUMNS == CompatibleMatrixType::ROWS &&
		ROWS == CompatibleMatrixType::ROWS,
		Matrix&
		> operator*=(const CompatibleMatrixType& rhs) noexcept
	{
		const auto result = *this * rhs;
		*this = result;
		return *this;
	}

	Matrix& operator*=(Scalar scalar) noexcept {
		std::transform(elements_.begin(), elements_.end(), elements_.begin(), [scalar](auto element) {
			return element * scalar;
		});
		return *this;
	}

	friend Vector<Scalar, ROWS, ScalarEqualityFunc> operator*(
		const Matrix& matrix,
		const Vector<Scalar, ROWS, ScalarEqualityFunc>& vector
		) noexcept
	{
		auto result = Vector<Scalar, ROWS, ScalarEqualityFunc>();
		for (auto columnIndex = 0u; columnIndex < COLUMNS; ++columnIndex) {
			result[columnIndex] = dot(matrix.column(columnIndex), vector);
		}
		return result;
	}

	Matrix& operator/=(Scalar scalar) noexcept {
		std::transform(elements_.begin(), elements_.end(), elements_.begin(), [scalar](auto element) {
			return element / scalar;
		});
		return *this;
	}

	// --- MATRIX-SPECIFIC OPERATIONS

	constexpr Matrix<Scalar, COLUMNS, ROWS, ScalarEqualityFunc> transpose() const noexcept {
		return viewMatrixTransposed(*this);
	}

	Scalar determinant() const noexcept {
		return viewMatrix(*this).determinant();
	}

	Scalar cofactor(size_t rowIndex, size_t columnIndex) const noexcept {
		return viewMatrix(*this).cofactor(rowIndex, columnIndex);
	}

	Matrix inverse() const noexcept {
		return viewMatrix(*this).inverse();
	}

	// --- ACCESSORS

	constexpr const Row& operator[](size_t rowIndex) const noexcept {
		assert(rowIndex < ROWS);
		return elements_[rowIndex];
	}

	Row& operator[](size_t rowIndex) noexcept {
		assert(rowIndex < ROWS);
		return elements_[rowIndex];
	}

	template <size_t ROW, size_t COLUMN>
	constexpr std::enable_if_t<(ROW < ROWS && COLUMN < COLUMNS), const Scalar&> get() const noexcept {
		return elements_[ROW].get<COLUMN>();
	}

	template <size_t ROW, size_t COLUMN>
	std::enable_if_t<(ROW < ROWS && COLUMN < COLUMNS), Scalar&> get() noexcept {
		return elements_[ROW].get<COLUMN>();
	}

	constexpr const Row& row(size_t rowIndex) const noexcept {
		assert(rowIndex < ROWS);
		return elements_[rowIndex];
	}

	Row& row(size_t rowIndex) noexcept {
		assert(rowIndex < ROWS);
		return elements_[rowIndex];
	}

	Column column(size_t columnIndex) const noexcept {
		assert(columnIndex < COLUMNS);
		auto column = Column();
		getColumn_<>(column, columnIndex);
		return column;
	}

	auto view() const noexcept {
		return viewMatrix(*this);
	}

	auto view() noexcept {
		return viewMatrix(*this);
	}

private:

	std::array<Row, ROWS> elements_;

	template <size_t ROW = 0>
	void getColumn_(Column& column, size_t columnIndex) const {
		column.get<ROW>() = elements_[ROW][columnIndex];
		getColumn_<ROW + 1>(column, columnIndex);
	}

	template <>
	void getColumn_<ROWS>(Column&, size_t) const {
	}

};

using Matrix4x4 = Matrix<float, 4, 4>;
static_assert(sizeof(Matrix4x4) == sizeof(float) * 16, "Empty base optimisation didn't work");
static_assert(std::is_trivially_copyable<Matrix4x4>::value, "Matrix is not trivially copyable");

namespace detail {

template <class ST, size_t R, size_t C, class SEF>
Matrix<ST, R, C, SEF> makeIdentity() {
	auto identity = Matrix<ST, R, C, SEF>();
	for (auto row = 0u; row < R; ++row) {
		identity[row][row] = ST(1);
	}
	return identity;
}

} // namespace detail

template <class ST, size_t R, size_t C, class SEF>
const Matrix<ST, R, C, SEF>	Matrix<ST, R, C, SEF>::IDENTITY =
	detail::makeIdentity<ST, R, C, SEF>();

} // namespace math

using math::Matrix;
using math::Matrix4x4;

} // namespace pulp
} // namespace coconut

#endif /* _COCONUT_PULP_MATH_MATRIX_HPP_ */


#ifndef RSVD_GRAMSCHMIDT_HPP_
#define RSVD_GRAMSCHMIDT_HPP_

#include <Eigen/Dense>
#include <algorithm>
#include <limits>

namespace Rsvd {

namespace Internal {

/// \brief Orthonormalize matrix columns inplace. Deflate if needed.
///
/// \warning This is a self-implemented algorithm, use it with caution. Use
/// #Rsvd::SubspaceIterationConditioner::Lu or #Rsvd::SubspaceIterationConditioner::Qr if numerical
/// stability and implementation quality is important.
///
/// \long This function implements the modified Gram--Schmidt process with deflation.
///
/// The deflation is implemented as follows: Let \f$i\f$ be the index of the current matrix column.
/// Let \f$n_i\f$ be its norm after projection onto a subspace orthogonal to the subspace spanned
/// by the previous columns \f$1, \ldots, n - 1\f$. Obviously, if \f$n_i\f$ is equal to zero, then
/// the \f$i\f$-th column is a linear combination of previous columns. Due to the limited precision
/// of the floating-point computations, \f$n_i\f$ will be small but greater than zero. Hence, we
/// use the following heuristic rule: If \f[ n_i < \max\{n_1, \ldots, n_{i - 1}\} \cdot
/// \varepsilon_{\mathrm{mach}} \cdot 100, \f] then the \f$i\f$-th column is deemed to be linearly
/// dependent and is filled with zeros.
///
/// Here, \f$ \varepsilon_{\mathrm{mach}} \f$ is the machine epsilon.
///
/// \note Although this function deflates if needed, it asserts that the matrix has fewer columns
/// than rows. This might help during debugging. If compiled without assertions, this function will
/// silently deflate on column rank loss.
///
/// \tparam MatrixType Eigen matrix type.
///
/// \param a Matrix whose columns should be orthonormalized inplace. The matrix can be over a real
/// or complex field.
template <typename MatrixType> void modifiedGramSchmidt(MatrixType &a) {
  using RealType = typename Eigen::NumTraits<typename MatrixType::Scalar>::Real;

  RealType largestNormSeen{0};
  // 100 is just an educated guess...
  const RealType tol{100 * std::numeric_limits<RealType>::epsilon()};

  // If a matrix has fewer rows than columns then the columns are linearly dependent
  assert(a.cols() <= a.rows());

  Eigen::Index currCol;
  for (currCol = 0; currCol < a.cols(); ++currCol) {
    for (Eigen::Index prevCol{0}; prevCol < currCol; ++prevCol) {
      /// \note Implementation detail: The order in the dot product is important for vectors over
      /// complex fields!
      a.col(currCol) -= a.col(prevCol).dot(a.col(currCol)) * a.col(prevCol);
    }

    // If the current column has near zero norm, it is a linear combination of previous columns
    const auto currColNorm{a.col(currCol).norm()};
    if (currColNorm < tol * largestNormSeen) {
      // Deflate
      a.col(currCol).setZero();
    } else {
      // Normalize
      a.col(currCol) /= currColNorm;
      largestNormSeen = std::max(largestNormSeen, currColNorm);
    }
  }
}

} // namespace Internal

} // namespace Rsvd

#endif


#include <limits>
#include <iostream>
#include <iomanip>
#include <string>
#include <cmath>
#include <algorithm>
#include <vector>

// BGL includes
#include <boost/graph/adjacency_list.hpp>
#include <boost/graph/cycle_canceling.hpp>
#include <boost/graph/push_relabel_max_flow.hpp>
#include <boost/graph/successive_shortest_path_nonnegative_weights.hpp>
#include <boost/graph/find_flow_cost.hpp>

// Graph Type with nested interior edge properties for Cost Flow Algorithms
typedef boost::adjacency_list_traits<boost::vecS, boost::vecS, boost::directedS> traits;
typedef boost::adjacency_list<boost::vecS, boost::vecS, boost::directedS, boost::no_property,
    boost::property<boost::edge_capacity_t, long,
        boost::property<boost::edge_residual_capacity_t, long,
            boost::property<boost::edge_reverse_t, traits::edge_descriptor,
                boost::property <boost::edge_weight_t, long> > > > > graph; // new! weightmap corresponds to costs

typedef boost::graph_traits<graph>::edge_descriptor             edge_desc;
typedef boost::graph_traits<graph>::out_edge_iterator           out_edge_it; // Iterator

// Custom edge adder class
class edge_adder {
 graph &G;

 public:
  explicit edge_adder(graph &G) : G(G) {}
  void add_edge(int from, int to, long capacity, long cost) {
    auto c_map = boost::get(boost::edge_capacity, G);
    auto r_map = boost::get(boost::edge_reverse, G);
    auto w_map = boost::get(boost::edge_weight, G); // new!
    const edge_desc e = boost::add_edge(from, to, G).first;
    const edge_desc rev_e = boost::add_edge(to, from, G).first;
    c_map[e] = capacity;
    c_map[rev_e] = 0; // reverse edge has no capacity!
    r_map[e] = rev_e;
    r_map[rev_e] = e;
    w_map[e] = cost;   // new assign cost
    w_map[rev_e] = -cost;   // new negative cost
  }
};

struct Request {
    int si;
    int ti;
    int di;
    int ai;
    int pi;
};

void testcase() {
    long n, s;
    std::cin >> n >> s;
    std::vector<int> l(s);
    long total_cars = 0;
    for(int i = 0; i < s; i++) {
        std::cin >> l[i];
        total_cars += l[i];
    }

    std::vector<std::set<int>> timestamps_per_s(s);

    std::vector<Request> requests(n);
    long max_arrival = 100'000;
    long max_p = 100;
    for(int i = 0; i < n; i++) {
        int si, ti, di, ai, pi;
        std::cin >> si >> ti >> di >> ai >> pi;
        requests[i] = {si - 1, ti - 1, di, ai, pi};
        if(di != 0) timestamps_per_s[si - 1].insert(di);
        if(ai != max_arrival) timestamps_per_s[ti - 1].insert(ai);
        // std::cerr << s * di / 30 + si - 1 << " " << s * ai / 30 + ti - 1 <<  " " << pi << std::endl;
    }

    int sum_dist_nodes = 0;
    std::vector<std::unordered_map<int, int>> map_to_g(s); // map (si, ti) -> j in graph
    for(int i = 0; i < s; i++) {
        map_to_g[i].insert({0, sum_dist_nodes++});
        map_to_g[i].insert({max_arrival, sum_dist_nodes++});
        for(auto t : timestamps_per_s[i]) {
            map_to_g[i].insert({t, sum_dist_nodes++});
            // std::cerr << "i " << i << " " << t  << "\n";
        }
    }
    graph G(sum_dist_nodes);
    const int v_source = boost::add_vertex(G);
    const int v_sink = boost::add_vertex(G);
    edge_adder adder(G);

    for(auto r : requests) {
        int si, ti, di, ai, pi;
        si = r.si; ti = r.ti; di = r.di;
        ai = r.ai; pi = r.pi;
        // std::cerr << -pi + max_p * (ai - di) << std::endl;
        adder.add_edge(map_to_g[si].at(di) , map_to_g[ti].at(ai), 1, -pi + max_p * (ai - di));
    }

    for(int i = 0; i < s; i++) {
        // source to time 0
        adder.add_edge(v_source, map_to_g[i].at(0), l[i], 0);
        adder.add_edge(map_to_g[i].at(max_arrival), v_sink, total_cars, 0);
        // time 0 to first time
        if(timestamps_per_s[i].size() > 0) {
            auto t1 = timestamps_per_s[i].begin();
            // zero to first time
            adder.add_edge(map_to_g[i].at(0), map_to_g[i].at(*t1), total_cars, *t1 * max_p);
            auto t2 = timestamps_per_s[i].begin(); t2++;
            for(; t2 != timestamps_per_s[i].end(); t1++, t2++) {
                adder.add_edge(map_to_g[i].at(*t1), map_to_g[i].at(*t2), total_cars, (*t2 - *t1) * max_p);
            }
            // second to last to max time
            adder.add_edge(map_to_g[i].at(*t1), map_to_g[i].at(max_arrival), total_cars, (max_arrival - *t1) * max_p);
        } else {
            adder.add_edge(map_to_g[i].at(0), map_to_g[i].at(max_arrival), total_cars, (max_arrival) * max_p);
        }
    }

    boost::successive_shortest_path_nonnegative_weights(G, v_source, v_sink);
    long cost = boost::find_flow_cost(G);
    // std::cout << -(cost) << "\n";
    std::cout << -(cost - (total_cars * max_p * max_arrival)) << "\n";

//   // // Retrieve the capacity map and reverse capacity map
//   const auto c_map = boost::get(boost::edge_capacity, G);
//   const auto rc_map = boost::get(boost::edge_residual_capacity, G);

//   // Iterate over all the edges to print the flow along them
//   auto edge_iters = boost::edges(G);
//   std::cerr << v_source << " " << v_sink;
//   for (auto edge_it = edge_iters.first; edge_it != edge_iters.second; ++edge_it) {
//     const edge_desc edge = *edge_it;
//     const long flow_through_edge = c_map[edge] - rc_map[edge];
//     std::cerr << "edge from " << boost::source(edge, G) << " to " << boost::target(edge, G)
//               << " with capacity " << c_map[edge]
//               << " runs " << flow_through_edge
//               << " units of flow (negative for reverse direction). \n";
//   }
    return;
}

int main() {
    std::ios_base::sync_with_stdio(false);

    int t;
    std::cin >> t;
    for (int i = 0; i < t; ++i)
        testcase();
}


#ifndef CANNON_ML_RECURSIVE_PIECEWISE_LSTD_H
#define CANNON_ML_RECURSIVE_PIECEWISE_LSTD_H

/*!
 * \file cannon/ml/piecewise_recursive_lstd.hpp
 * \brief File containing PiecewiseRecursiveLSTDFilter class definition.
 */

#include <Eigen/Core>
#include <Eigen/Sparse>

using namespace Eigen;

namespace cannon {
  namespace ml {

    /*!
     * \brief Class representing a Least-Squares Temporal Difference (LSTD)
     * approximator using a feature set that produces a piecewise-affine value
     * function approximation. This is for use in reinforcement learning
     * algorithms. This version of the algorithm is similar to a recursive
     * least squares filter.
     */
    class PiecewiseRecursiveLSTDFilter {
      public:
        PiecewiseRecursiveLSTDFilter() = delete;

        /*!
         * \brief Constructor taking state space dimension, number of affine
         * regions, and discount factor.
         */
        PiecewiseRecursiveLSTDFilter(unsigned int in_dim, unsigned int num_refs, double
            discount_factor, double alpha=1.0) : in_dim_(in_dim + 1),
        param_dim_(in_dim_ * num_refs), num_refs_(num_refs),
        discount_factor_(discount_factor), alpha_(alpha) {
          a_inv_ = MatrixXd::Identity(param_dim_, param_dim_) * alpha_;
          theta_ = VectorXd::Zero(param_dim_);
        }

        /*!
         * \brief Update this approximation in light of a single data point,
         * which is a transition from one state to another with associated
         * reward.
         *
         * \param in_vec Feature vector for first state.
         * \param next_in_vec Feature vector for next state.
         * \param idx Region index for first state.
         * \param next_idx Region index for next state.
         * \param reward Reward associated with this state transition.
         */
        void process_datum(const VectorXd& in_vec, const VectorXd& next_in_vec,
            unsigned int idx, unsigned int next_idx, double reward);

        /*!
         * \brief Get the matrix representing the linear portion of the local,
         * affine approximation in the region with for the input index.
         *
         * \param idx Region index
         *
         * \returns Estimated local linear approximation parameter matrix.
         */
        VectorXd get_mat(unsigned int idx) const;

        /*!
         * \brief Predict the value of the input state using the estimated
         * piecewise-affine value function.
         *
         * \param in_vec Feature vector for the state.
         * \param idx Region index of the state.
         *
         * \returns Value function prediction.
         */
        double predict(const VectorXd& in_vec, unsigned int idx) const;

        /*!
         * \brief Reset this value function approximation.
         */
        void reset();

      private:
        /*!
         * \brief Make internal feature vector for the input state in the
         * region with the input index.
         *
         * \param in_vec Input features for state.
         * \param idx Region index for the state.
         *
         * \returns Internal feature representation leading to piecewise-affine
         * function.
         */
        SparseMatrix<double>
        make_feature_vec_(VectorXd in_vec, unsigned int idx) const;

        /*!
         * \brief Update approximation of the linear portion of the LSTD filter
         * given a particular state transition.
         *
         * \param a_inv_feat_t A^-1 * feat.transpose()
         * \param diff feat - (discount_factor_ * next_feat)
         * \param inv_denom 1.0 / (diff * A^-1 * feat^T)
         */
        void update_a_inv_(const Ref<const VectorXd> &a_inv_feat_t,
                           const Ref<const SparseMatrix<double>> &diff,
                           double inv_denom);

        // Parameters
        unsigned int in_dim_; //!< Dimension of input
        unsigned int param_dim_; //!< Dimension of internal feature space
        unsigned int num_refs_; //!< Number of linear regions
        double discount_factor_; //!< Discount factor for value function
        double alpha_; //!< L2 regularization parameter

        // Matrices
        MatrixXd a_inv_; //!< Linear portion of LSTD filter
        VectorXd theta_; //!< Parameters of value function approximation
    };

  } // namespace ml
} // namespace cannon
#endif /* ifndef CANNON_ML_RECURSIVE_PIECEWISE_LSTD_H */


#ifndef ESKF_IMPL
#define ESKF_IMPL

#include "ESKF.h"
#include "utilities.h"

#include <Eigen/LU>
#include <Wire.h>

#ifndef SGN
#define SGN(X) ((X > 0) - (X < 0))
#endif

#ifndef RAD_TO_DEG
#define RAD_TO_DEG (180.0 / M_PI)
#endif

#ifndef DEG_TO_RAD
#define DEG_TO_RAD (M_PI / 180.0)
#endif

#ifndef MS2_TO_G
#define MS2_TO_G (1.0 / 9.81)
#endif

#ifndef G_TO_MS2
#define G_TO_MS2 (9.81)
#endif

namespace IMU_EKF
{

template <typename precision>
ESKF<precision>::ESKF()
{
    init();
}

template <typename precision>
void ESKF<precision>::init()
{
    qref_ = Quaternion<precision>();
    x_.setZero();
    P_.setIdentity();
    Q_.setIdentity();
    // TODO don't hardcode this
    Q_(0, 0) = 0.05;
    Q_(1, 1) = 0.05;
    Q_(2, 2) = 0.05;
    Q_(3, 3) = 0.05;
    Q_(4, 4) = 0.05;
    Q_(5, 5) = 0.05;
    Q_(6, 6) = 0.025;
    Q_(7, 7) = 0.025;
    Q_(8, 8) = 0.025;
    Q_(9, 9) = 0.025;
    Q_(10, 10) = 0.025;
    Q_(11, 11) = 0.025;
    Q_(12, 12) = 0.01;
    Q_(13, 13) = 0.01;
    Q_(14, 14) = 0.01;

    R_Gyr_.setIdentity();
    R_Gyr_(0, 0) = 0.0000045494 * DEG_TO_RAD;
    R_Gyr_(1, 1) = 0.0000039704 * DEG_TO_RAD;
    R_Gyr_(2, 2) = 0.0000093844 * DEG_TO_RAD;

    R_Acc_.setIdentity();
    R_Acc_(0, 0) = 0.0141615383 * G_TO_MS2;
    R_Acc_(1, 1) = 0.0164647549 * G_TO_MS2;
    R_Acc_(2, 2) = 0.0100094303 * G_TO_MS2;

    R_Mag_.setIdentity();
    R_Mag_(0, 0) = 0.00004;
    R_Mag_(1, 1) = 0.00004;
    R_Mag_(2, 2) = 0.00004;
}

template <typename precision>
void ESKF<precision>::initWithAcc(const float ax, const float ay, const float az)
{
    init();

    // see https://cache.freescale.com/files/sensors/doc/app_note/AN3461.pdf
    // rotation sequence R = Rx * Ry * Rz
    // eq. 38
    precision roll = std::atan2(ay, SGN(-az) * std::sqrt(az * az + 0.01 * ax * ax));
    // eq. 37
    precision pitch = std::atan(-ax / std::sqrt(ay * ay + az * az));

    precision sr05 = std::sin(0.5 * roll);
    precision cr05 = std::cos(0.5 * roll);
    precision sp05 = std::sin(0.5 * pitch);
    precision cp05 = std::cos(0.5 * pitch);
    qref_ = Quaternion<precision>(sr05, 0, 0, cr05) * Quaternion<precision>(0, sp05, 0, cp05);
}

template <typename precision>
void ESKF<precision>::initWithAccAndMag(const float ax, const float ay, const float az, const float mx, const float my, const float mz, const Eigen::Matrix<precision, 3, 3> &Winv, const Eigen::Matrix<precision, 3, 1> &V)
{
    init();

    // see https://cache.freescale.com/files/sensors/doc/app_note/AN3461.pdf
    // rotation sequence R = Rx * Ry * Rz
    // eq. 38
    precision roll = std::atan2(ay, SGN(-az) * std::sqrt(az * az + 0.01 * ax * ax));
    // eq. 37
    precision pitch = std::atan(-ax / std::sqrt(ay * ay + az * az));

    // see https://www.nxp.com/docs/en/application-note/AN4246.pdf
    // eq. 6 - 10
    precision sr = std::sin(roll);
    precision cr = std::cos(roll);
    precision sp = std::sin(pitch);
    precision cp = std::cos(pitch);

    Eigen::Matrix<precision, 3, 3> RxT;
    RxT << 1, 0, 0,
        0, cr, sr,
        0, -sr, cr;
    Eigen::Matrix<precision, 3, 3> RyT;
    RyT << cp, 0, -sp,
        0, 1, 0,
        sp, 0, cp;

    Eigen::Matrix<precision, 3, 1> Bp;
    Bp << mx, my, mz;

    Eigen::Matrix<precision, 3, 1> Bf;
    Bf = RyT * RxT * Winv * (Bp - V);

    precision yaw = -atan2(-Bf(1), Bf(0));

    precision sr05 = std::sin(0.5 * roll);
    precision cr05 = std::cos(0.5 * roll);
    precision sp05 = std::sin(0.5 * pitch);
    precision cp05 = std::cos(0.5 * pitch);
    precision sy05 = std::sin(0.5 * yaw);
    precision cy05 = std::cos(0.5 * yaw);

    qref_ = Quaternion<precision>(sr05, 0, 0, cr05) * Quaternion<precision>(0, sp05, 0, cp05) * Quaternion<precision>(0, 0, sy05, cy05);
}

template <typename precision>
void ESKF<precision>::predict(precision dt)
{
    // eq. 23
    Quaternion<precision> angular_velocity_quat(x_[12], x_[13], x_[14], 0);
    qref_ += dt * ((0.5 * angular_velocity_quat) * qref_);
    qref_.normalize();

    Eigen::Matrix<precision, 9, 9> A;
    A.setIdentity();
    A(0, 3) = dt;
    A(0, 6) = dt * dt * 0.5;
    A(1, 4) = dt;
    A(1, 7) = dt * dt * 0.5;
    A(2, 5) = dt;
    A(2, 8) = dt * dt * 0.5;
    A(3, 6) = dt;
    A(4, 7) = dt;
    A(4, 8) = dt;

    x_.segment(0, 9) = A * x_.segment(0, 9);
    P_.topLeftCorner(9, 9) = A * P_.topLeftCorner(9, 9) * A.transpose() + dt * Q_.topLeftCorner(9, 9);

    // eq. 38
    Eigen::Matrix<precision, 3, 1> angular_velocity = x_.segment(12, 3);
    Eigen::Matrix<precision, 3, 1> error = x_.segment(9, 3);
    Eigen::Matrix<precision, 6, 6> Jac = Eigen::Matrix<precision, 6, 6>::Zero();
    Jac.topLeftCorner(3, 3) = toCrossMatrix<precision>(error);
    Jac.topRightCorner(3, 3) = -toCrossMatrix<precision>(angular_velocity);
    // eq. 39
    Eigen::Matrix<precision, 6, 6> G = Eigen::Matrix<precision, 6, 6>::Identity();
    G.bottomRightCorner(3, 3) *= -1;

    // eq. 33 adpated
    P_.bottomRightCorner(6, 6) = Jac * P_.bottomRightCorner(6, 6) * Jac.transpose() + G * dt * Q_.bottomRightCorner(6, 6) * G.transpose();
}

template <typename precision>
void ESKF<precision>::correctGyr(const float gx, const float gy, const float gz)
{
    // z
    Eigen::Matrix<precision, MEASSUREMENT_GYR_SIZE, 1> z;
    z << gx * DEG_TO_RAD, gy * DEG_TO_RAD, gz * DEG_TO_RAD;

    // Kalman Gain
    // K = P * H' (H * P * H' + V * R * V')^-
    // H = eye(3)
    Eigen::Matrix<precision, 3, MEASSUREMENT_GYR_SIZE> K;
    K = P_.bottomRightCorner(3, 3) * (P_.bottomRightCorner(3, 3) + R_Gyr_).inverse();

    // x = x + K * (z - H * x)
    x_.segment(12, 3) += K * (z - x_.segment(12, 3));

    // P = (I - KH)P
    Eigen::Matrix<precision, 3, 3> IKH = Eigen::Matrix<precision, 3, 3>::Identity();
    IKH -= K;

    P_.bottomRightCorner(3, 3) = IKH * P_.bottomRightCorner(3, 3);
}

template <typename precision>
void ESKF<precision>::correctAcc(const float ax, const float ay, const float az)
{
    // z
    Eigen::Matrix<precision, MEASSUREMENT_ACC_SIZE, 1> z;
    z << ax * G_TO_MS2, ay * G_TO_MS2, az * G_TO_MS2;

    Eigen::Matrix<precision, 3, 1> gravity;
    gravity << 0.0, 0.0, -9.81;

    if (std::abs(z.norm() - 9.81) < 0.7)
    {
        // set acc to zero
        // z
        Eigen::Matrix<precision, 3, 1> z_acc = Eigen::Matrix<precision, 3, 1>::Zero();

        // Kalman Gain
        // K = P * H' (H * P * H' + V * R * V')^-
        // H = eye(3)
        Eigen::Matrix<precision, 3, 3> K_acc;
        K_acc = P_.block(6, 6, 3, 3) * (P_.block(6, 6, 3, 3) + R_Acc_).inverse();

        // x = x + K * (z - H * x)
        x_.segment(6, 3) += K_acc * (z_acc - x_.segment(6, 3));

        // P = (I - KH)P
        Eigen::Matrix<precision, 3, 3> IKH = Eigen::Matrix<precision, 3, 3>::Identity();
        IKH -= K_acc;

        P_.block(6, 6, 3, 3) = IKH * P_.block(6, 6, 3, 3);

        // correct tilt
        Eigen::Matrix<precision, 3, 1> vi = gravity;

        // eq. 42
        Eigen::Matrix<precision, 3, 1> error = x_.segment(9, 3);
        // Eigen::Matrix<precision, 3, 3> Aa = toRotationMatrix<precision>(error);
        Eigen::Matrix<precision, 3, 3> Aq = qref_.toRotationMatrix();
        Eigen::Matrix<precision, 3, 1> vb_pred = Aq * vi;

        // H
        // eq. 44
        Eigen::Matrix<precision, 3, 3> Ha = toCrossMatrix<precision>(vb_pred);

        // eq. 46
        Eigen::Matrix<precision, 3, 3> K;
        K = P_.block(9, 9, 3, 3) * Ha.transpose() * (Ha * P_.block(9, 9, 3, 3) * Ha.transpose() + R_Acc_).inverse();

        // eq. 47
        // h = vb_pred
        x_.segment(9, 3) += K * (z - vb_pred - Ha * error);

        // eq. 48
        P_.block(9, 9, 3, 3) -= K * Ha * P_.block(9, 9, 3, 3);
    }
    else
    {
        Eigen::Matrix<precision, 3, 1> vi = gravity + x_.segment(6, 3);

        // eq. 42
        Eigen::Matrix<precision, 3, 1> error = x_.segment(9, 3);
        Eigen::Matrix<precision, 3, 3> Aa = toRotationMatrix<precision>(error);
        Eigen::Matrix<precision, 3, 3> Aq = qref_.toRotationMatrix();
        Eigen::Matrix<precision, 3, 1> vb_pred = Aq * vi;

        // H
        // eq. 44
        Eigen::Matrix<precision, 3, 6> H;
        Eigen::Matrix<precision, 3, 3> Ha = toCrossMatrix<precision>(vb_pred);
        H.topLeftCorner(3, 3) = Aa * Aq;
        H.bottomRightCorner(3, 3) = Ha;

        // eq. 46
        Eigen::Matrix<precision, 6, 3> K;
        K = P_.block(6, 6, 6, 6) * H.transpose() * (H * P_.block(6, 6, 6, 6) * H.transpose() + R_Acc_).inverse();

        // eq. 47
        // h = v_bred
        x_.segment(6, 6) += K * (z - vb_pred - Ha * error);

        // eq. 48
        P_.block(6, 6, 6, 6) -= K * H * P_.block(6, 6, 6, 6);
    }
}

template <typename precision>
void ESKF<precision>::correctMag(const float mx, const float my, const float mz, const float incl, const float B, const Eigen::Matrix<precision, 3, 3> &W, const Eigen::Matrix<precision, 3, 1> &V)
{
    // z
    Eigen::Matrix<precision, MEASSUREMENT_MAG_SIZE, 1> z;
    z << mx, my, mz;

    Eigen::Matrix<precision, 3, 1> vi;
    vi << std::cos(incl), 0, -std::sin(incl);
    vi = B * vi;

    // eq. 42
    Eigen::Matrix<precision, 3, 1> error = x_.segment(9, 3);
    // Eigen::Matrix<precision, 3, 3> Aa = toRotationMatrix<precision>(error);
    Eigen::Matrix<precision, 3, 3> Aq = qref_.toRotationMatrix();
    Eigen::Matrix<precision, 3, 1> vb_pred = Aq * vi;

    // h(v) = W * v + V
    Eigen::Matrix<precision, 3, 1> h = W * vb_pred + V;

    // H
    // eq. 44
    Eigen::Matrix<precision, 3, 3> Ha = W * toCrossMatrix<precision>(vb_pred);

    // eq. 46
    Eigen::Matrix<precision, 3, 3> K;
    K = P_.block(9, 9, 3, 3) * Ha.transpose() * (Ha * P_.block(9, 9, 3, 3) * Ha.transpose() + R_Mag_).inverse();

    // eq. 47
    x_.segment(9, 3) += K * (z - h - Ha * error);

    // eq. 48
    P_.block(9, 9, 3, 3) -= K * Ha * P_.block(9, 9, 3, 3);
}

template <typename precision>
void ESKF<precision>::reset()
{
    // eq. 21
    qref_ = Quaternion<precision>(x_(9), x_(10), x_(11), 2.0) * qref_;
    // eq. 22
    qref_.normalize();
    x_(9) = 0.0;
    x_(10) = 0.0;
    x_(11) = 0.0;
}

template <typename precision>
Eigen::Matrix<precision, STATE_SIZE, 1> ESKF<precision>::getState() const
{
    return x_;
}

template <typename precision>
void ESKF<precision>::getAttitude(float &roll, float &pitch, float &yaw) const
{
    Eigen::Matrix<precision, 3, 1> angles = qref_.toEulerAngles();
    roll = angles(0);
    pitch = angles(1);
    yaw = angles(2);
}

template <typename precision>
Quaternion<precision> ESKF<precision>::getAttitude() const
{
    return qref_;
}

template <typename precision>
void ESKF<precision>::getAcceleration(float &x, float &y, float &z) const
{
    x = x_(6);
    y = x_(7);
    z = x_(8);
}
} // namespace IMU_EKF

#endif // ESKF_IMPL


//
//  Family.hpp
//  test_rcpp
//
//  Created by Dongjie Wu on 18/11/2021.
//

#ifndef Family_hpp
#define Family_hpp

#include <stdio.h>
#include <RcppArmadillo.h>
//#include <boost/math/distributions/normal.hpp>
//#include <boost/multiprecision/cpp_bin_float.hpp>
//#include "distribution.hpp"
//using namespace boost::math;
using namespace Rcpp;

// [[Rcpp::depends(RcppArmadillo)]]

// !! need to decide provide the option of not log
template <class T>
class Family {
public:
  arma::mat logLik(const arma::vec& theta, const arma::mat& Y,
                   const arma::mat& X, const bool& lg = false) {
    return(static_cast<T*>(this) -> logLik(theta, Y, X, lg));
  }

};

class FamilyNormal : public Family<FamilyNormal> {
public:
  arma::mat logLik(const arma::vec& theta, const arma::mat& Y,
                   const arma::mat& X, const bool& lg = false)
  {
    arma::mat mean_mat = (X * theta.subvec(1,(theta.n_elem-1)));
    double sd = sqrt(theta[0] * theta[0]);
    arma::mat l = arma::zeros(Y.n_rows, Y.n_cols);
    for (int i=0; i<Y.n_rows;++i) {
      for (int j=0; j<Y.n_cols;++j) {
          //normal_distribution nd(mean_mat.at(i,j),sd);
          l.at(i,j) = R::dnorm4(Y.at(i,j), mean_mat.at(i,j), sd, lg);
      }
    }
    return l;
  }
};

class FamilyPoisson : public Family<FamilyPoisson> {
public:
  arma::mat logLik(const arma::vec& theta, const arma::mat& Y,
                   const arma::mat& X, const bool& lg = false)
  {
    arma::mat mean_mat = arma::exp(X * theta); // the precision problem
    arma::mat l = arma::zeros(Y.n_rows, Y.n_cols);

    for (int i=0; i<Y.n_rows;++i) {
      for (int j=0; j<Y.n_cols;++j) {
          l.at(i,j) = R::dpois(Y.at(i,j), mean_mat.at(i,j), lg);
      }
    }
    return l;
  }
};

class FamilyLogit : public Family<FamilyLogit> {
public:
  arma::mat sigmoid(arma::mat x) {
          return (1/(1+arma::exp(-x)));
      };
  arma::mat logLik(const arma::vec& theta, const arma::mat& Y,
                   const arma::mat& X, const bool& lg = false)
  {
//      arma::mat mean_t = (X * theta);
//      arma::mat mean_mat = log(sigmoid(mean_t));
//      arma::mat l = arma::zeros(Y.n_rows, Y.n_cols);
//      for (int i=0; i<Y.n_rows;i++) {
//        for (int j=0; j<Y.n_cols; j++) {
//           l.at(i,j) = R::dbinom(Y.at(i,j), 1, mean_mat.at(i,j), lg);
//        }
//      }
//      return l;
    arma::mat mean_t = (X * theta);
    arma::mat l = arma::zeros(Y.n_rows, Y.n_cols);
    l = (Y % mean_t) - arma::log1p(arma::exp(mean_t));
    return l;
  }
};

class FamilyMultiNomial : public Family<FamilyMultiNomial> {
public:
  arma::mat logLik(const arma::vec& theta, const arma::mat& Y,
                   const arma::mat& X, const bool& lg = false)
  {
    arma::mat theta_t = arma::mat(theta);
    theta_t.reshape(X.n_cols, Y.n_cols);
    arma::mat mean_mat = (X * theta_t);
    arma::vec r = arma::sum((Y % mean_mat), 1) - log(1 + arma::sum(arma::exp(mean_mat), 1));
    arma::mat l = arma::mat(r);
    return l;
  }
};

class FamilyConditionalLogit : public Family<FamilyConditionalLogit> {
public:
  arma::mat logLik(const arma::vec& theta, const arma::mat& Y,
                   const arma::mat& X, const bool& lg = false)
 // The group variable is the last variable of X.
  {
    //auto t1 = std::chrono::high_resolution_clock::now();
    arma::uvec gV = arma::conv_to<arma::uvec>::from(X.col(X.n_cols-1));
    arma::uvec g = unique(gV);
    arma::mat l = arma::zeros(g.n_elem, 1);
    arma::mat l1 = arma::zeros(g.n_elem, 1);
    arma::mat l2 = arma::zeros(g.n_elem, 1);
    //auto t2 = std::chrono::high_resolution_clock::now();
    for (arma::uword i=0, nr=Y.n_rows; i<nr; ++i) {
        arma::uword j = X.at(i,X.n_cols-1)-1;
        double mean_x = 0;
        for (arma::uword k=0, nc=X.n_cols; k<nc-1; ++k) {
            mean_x += theta.at(k) * X.at(i, k);
        }
        l1.at(j, 0) += Y.at(i, 0) * mean_x;
        l2.at(j, 0) += exp(mean_x);
    }
    //auto t3 = std::chrono::high_resolution_clock::now();
    l = l1 - log(l2);
    //auto t4 = std::chrono::high_resolution_clock::now();

    /*
    const arma::mat& Xd = X.cols(0, X.n_cols-2);
    auto t2 = std::chrono::high_resolution_clock::now();

    arma::cube Xc1(Y.n_rows, Y.n_cols, 1);
    Xc1.slice(0) = Y % (Xd * theta);
    Xc1.reshape(gV.n_elem/g.n_elem, Y.n_cols, g.n_elem);
    arma::cube Xc2(Y.n_rows, Y.n_cols, 1);
    Xc2.slice(0) = arma::exp(Xd * theta);
    Xc2.reshape(gV.n_elem/g.n_elem, Y.n_cols, g.n_elem);
    auto t3 = std::chrono::high_resolution_clock::now();

    for (arma::uword i=0; i<g.n_elem;++i) {
      l.at(i, 0) = arma::accu(Xc1.slice(i)) -     log(arma::accu(Xc2.slice(i)));
    }
    auto t4 = std::chrono::high_resolution_clock::now();

    std::chrono::duration<double, std::milli> d1 = t2 - t1;
    std::chrono::duration<double, std::milli> d2 = t3 - t2;
    std::chrono::duration<double, std::milli> d3 = t4 - t3;
    std::cout << d1.count() << "ms" << std::endl;
    std::cout << d2.count() << "ms" << std::endl;
    std::cout << d3.count() << "ms" << std::endl;
     */
    return l;
  }
};

class FamilyUnidiff : public Family<FamilyUnidiff> {
public:
  arma::mat logLik(const arma::vec& theta, const arma::mat& Y,
                   const arma::mat& X, const bool& lg = false)
 // The group variable is the last variable of X.
    {
        if (Y.n_cols != 1 ) {
            throw std::invalid_argument("Y should only have 1 column!");
        }
        if (X.n_cols != 2 ) {
            throw std::invalid_argument("X should have 2 columns!");
        }
        arma::vec uY = arma::unique(Y);
        arma::mat X_X = fast_dummy(X.col(0));
        arma::mat X_Z = fast_dummy(X.col(1));
        arma::mat Y_Y = fast_dummy(Y.col(0));
        arma::mat W = arma::ones(X_Z.n_rows, (X_Z.n_cols+1));
        W(arma::span::all,arma::span(1,W.n_cols-1)) = X_Z;
        arma::uword colX = X_X.n_cols;
        arma::uword colZ = X_Z.n_cols;
        arma::uword colY = Y_Y.n_cols;
        if (theta.n_elem != (colY*(colZ+1)+colY*colX+colZ)) {
            throw std::invalid_argument("Wrong size of theta!");
        }
        arma::uword end = colY*(colZ+1)-1;
        arma::uword end2 = colY*colX;
        arma::uword end3 = colZ;
        arma::mat theta_y = arma::mat(theta(arma::span(0, end)));
        theta_y.reshape(colZ+1, colY);
        arma::mat psi_y = arma::mat(theta(arma::span(end+1, end+end2)));
        psi_y.reshape(colX, colY);
        arma::mat phi = arma::mat(theta(arma::span(end+end2+1,end+end2+end3)));
        phi.reshape(colZ,1);
        arma::mat expZ = arma::exp(X_Z * phi);
        arma::mat phiX = X_X  * psi_y;
        arma::mat part2 = phiX.each_col() % expZ;
        arma::mat sum_part = W * theta_y + part2;
        arma::mat first = arma::sum(Y_Y % sum_part, 1);
        arma::mat second = log(1+arma::sum(arma::exp(sum_part), 1));
        arma::mat l = first - second;
        return l;
    }
    arma::mat fast_dummy(const arma::vec& x) {
        arma::vec vec_u = arma::unique(x);
        vec_u = vec_u(arma::span(1,vec_u.n_elem-1)); // Remove the baseline
        arma::mat mats = arma::zeros(x.n_elem, (vec_u.n_elem));
        for (arma::uword i=0; i<vec_u.n_elem; ++i) {
            mats.col(i) = x;
            double idx = vec_u.at(i);
            mats.col(i).for_each([idx](arma::vec::elem_type& val){
                (val == idx) ? (val = 1) : (val = 0);
            });
        }
        return mats;
    }
};

#endif /* Family_hpp */