#include <iostream>
#include <g2o/core/g2o_core_api.h>
#include <g2o/core/base_vertex.h>
#include <g2o/core/base_unary_edge.h>
#include <g2o/core/block_solver.h>
#include <g2o/core/optimization_algorithm_levenberg.h>
#include <g2o/core/optimization_algorithm_gauss_newton.h>
#include <g2o/core/optimization_algorithm_dogleg.h>
#include <g2o/solvers/dense/linear_solver_dense.h>
#include <Eigen/Core>
#include <opencv2/core/core.hpp>
#include <cmath>
#include <chrono>

using namespace std;

// Vertex of curve model, template parameters: optimize variable dimension and data type
class CurveFittingVertex : public g2o::BaseVertex<3, Eigen::Vector3d> {
public:
  EIGEN_MAKE_ALIGNED_OPERATOR_NEW

  // Reset
  virtual void setToOriginImpl() override {
    _estimate << 0, 0, 0;
  }

  // update
  virtual void oplusImpl(const double *update) override {
    _estimate += Eigen::Vector3d(update);
  }

  // Save and read: leave blank
  virtual bool read(istream &in) {}

  virtual bool write(ostream &out) const {}
};

// Error model Template parameters: observation dimension, type, connected vertex type
class CurveFittingEdge : public g2o::BaseUnaryEdge<1, double, CurveFittingVertex> {
public:
  EIGEN_MAKE_ALIGNED_OPERATOR_NEW

  CurveFittingEdge(double x) : BaseUnaryEdge(), _x(x) {}

  // Calculation curve model error
  virtual void computeError() override {
    const CurveFittingVertex *v = static_cast<const CurveFittingVertex *> (_vertices[0]);
    const Eigen::Vector3d abc = v->estimate();
    _error(0, 0) = _measurement - std::exp(abc(0, 0) * _x * _x + abc(1, 0) * _x + abc(2, 0));
  }

  // Calculate the Jacobian matrix
  virtual void linearizeOplus() override {
    const CurveFittingVertex *v = static_cast<const CurveFittingVertex *> (_vertices[0]);
    const Eigen::Vector3d abc = v->estimate();
    double y = exp(abc[0] * _x * _x + abc[1] * _x + abc[2]);
    _jacobianOplusXi[0] = -_x * _x * y;
    _jacobianOplusXi[1] = -_x * y;
    _jacobianOplusXi[2] = -y;
  }

  virtual bool read(istream &in) {}

  virtual bool write(ostream &out) const {}

public:
  double _x;  // x value, y value is _measurement
};

int main(int argc, char **argv) {
  double ar = 1.0, br = 2.0, cr = 1.0;         // True parameter value
  double ae = 2.0, be = -1.0, ce = 5.0;        // Estimate parameter values
  int N = 100;                                 // data point
  double w_sigma = 1.0;                        // noise Sigma value
  double inv_sigma = 1.0 / w_sigma;
  cv::RNG rng;                                 // OpenCV random number generator

  vector<double> x_data, y_data;      // data
  for (int i = 0; i < N; i++) {
    double x = i / 100.0;
    x_data.push_back(x);
    y_data.push_back(exp(ar * x * x + br * x + cr) + rng.gaussian(w_sigma * w_sigma));
  }

  //Build graph optimization, first set g2o
  typedef g2o::BlockSolver<g2o::BlockSolverTraits<3, 1>> BlockSolverType;  // The optimization variable dimension of each error term is 3, and the error value dimension is 1
  typedef g2o::LinearSolverDense<BlockSolverType::PoseMatrixType> LinearSolverType; // Linear solver type

  // Gradient descent method, choose from GN, LM, DogLe
  auto solver = new g2o::OptimizationAlgorithmGaussNewton(
    g2o::make_unique<BlockSolverType>(g2o::make_unique<LinearSolverType>()));
  g2o::SparseOptimizer optimizer;     // graph model
  optimizer.setAlgorithm(solver);   // Set the solver
  optimizer.setVerbose(true);       // Turn on debug output

  // 往图中增加顶点
  CurveFittingVertex *v = new CurveFittingVertex();
  v->setEstimate(Eigen::Vector3d(ae, be, ce));
  v->setId(0);
  optimizer.addVertex(v);

  // Add a side to the picture
  for (int i = 0; i < N; i++) {
    CurveFittingEdge *edge = new CurveFittingEdge(x_data[i]);
    edge->setId(i);
    edge->setVertex(0, v);                //  Set the vertices of the connection
    edge->setMeasurement(y_data[i]);      // Observe the value
    edge->setInformation(Eigen::Matrix<double, 1, 1>::Identity() * 1 / (w_sigma * w_sigma)); //Information matrix: the inverse of the covariance matrix
    optimizer.addEdge(edge);
  }

  // Perform optimization
  cout << "start optimization" << endl;
  chrono::steady_clock::time_point t1 = chrono::steady_clock::now();
  optimizer.initializeOptimization();
  optimizer.optimize(20);
  chrono::steady_clock::time_point t2 = chrono::steady_clock::now();
  chrono::duration<double> time_used = chrono::duration_cast<chrono::duration<double>>(t2 - t1);
  cout << "solve time cost = " << time_used.count() << " seconds. " << endl;

  // output optimized value
  Eigen::Vector3d abc_estimate = v->estimate();
  cout << "estimated model: " << abc_estimate.transpose() << endl;

  return 0;
}


// adapted from lassopack, see:
// https://github.com/statalasso/lassopack/blob/master/lassoutils.ado

#include "../eigen-3.4.0/Eigen/Eigen"
#include <iostream>
#include <math.h>
#include <pybind11/eigen.h>
#include <pybind11/pybind11.h>
#include <Eigen/Eigen>

namespace py = pybind11;
using namespace Eigen;
using namespace std;
// using MatrixXdRef = Eigen::Ref<Eigen::MatrixXd>;
// using VectorXdRef = Eigen::Ref<Eigen::VectorXd>;
using MatrixXdRef = Eigen::MatrixXd;
using VectorXdRef = Eigen::VectorXd;

VectorXd coordinateDescent(MatrixXdRef X, VectorXdRef y, MatrixXdRef XX, VectorXdRef Xy,
                       double lambd, VectorXdRef psi, VectorXdRef startingValues,
                       bool sqrtLasso = false, bool fitIntercept = true,
                       double optTol = 1e-10, int maxIter = 1000,
                       double zeroTol = 1e-4) {

  int n = X.rows(), p = X.cols();

  VectorXd beta = startingValues;

  VectorXd sdVec = X.colwise().norm();
  double ySd = y.norm();
  // normal lasso shooting
  if (!sqrtLasso) {

    XX *= 2, Xy *= 2;
    // loop over max iter
    for (int iter = 0; iter < maxIter; iter++) {
      // copy beta to beta_old
      VectorXd betaOld = beta;
      // loop over p
      for (int j = 0; j < p; j++) {
        // calculate s0 and do shooting
        double S0 = (XX.row(j) * beta).sum() - XX(j, j) * beta(j) - Xy(j);

        if (S0 > lambd * psi(j)) {
          beta(j) = (lambd * psi(j) - S0) / XX(j, j);

        } else if (S0 < -lambd * psi(j)) {
          beta(j) = (-lambd * psi(j) - S0) / XX(j, j);

        } else {
          beta(j) = 0;
        }
      }
      // check for convergence
      double diff = ((beta - betaOld).cwiseAbs() * sdVec / ySd).sum();
      // double diff = (beta - betaOld).cwiseAbs().sum();
      if (diff < optTol) {
        break;
      }
    }
    // sqrt-lasso shooting algorithm
  } else {
    // rescale XX and Xy
    XX /= n;
    Xy /= n;
    double MaxErrorNorm = 1.0e-10;

    // demean X and y
    // if (fitIntercept) {
    //   MatrixXd meanX = X.colwise().mean().replicate(n, 1);
    //   VectorXd meanY = VectorXd::Ones(n) * y.mean();
    //   X = X - meanX;
    //   y = y - meanY;
    // }

    // get error
    VectorXd error = y - X * beta;
    double qhat = error.squaredNorm() / n;

    for (int iter = 0; iter < maxIter; iter++) {
      VectorXd betaOld = beta;

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

        if (fabs(beta(j)) > 0) {
          error += X.col(j) * beta(j);
          qhat = error.squaredNorm() / n;
        }

        double S0 = XX.row(j).dot(beta) - XX(j, j) * beta(j) - Xy(j);
        double qqhat = max(qhat - (pow(S0, 2) / XX(j, j)), 0.0);

        if (pow(n, 2) < pow(lambd * psi(j), 2) / XX(j, j)) {
          beta(j) = 0;
        }

        else if (S0 > lambd / n * psi(j) * sqrt(qhat)) {
          beta[j] = ((lambd * psi(j) /
                      sqrt(pow(n, 2) - pow(lambd * psi(j), 2) / XX(j, j))) *
                         sqrt(qqhat) -
                     S0) /
                    XX(j, j);
          error -= X.col(j) * beta[j];
        }

        else if (S0 < -lambd / n * psi(j) * sqrt(qhat)) { // Optimal beta(j) > 0
          beta[j] = (-(lambd * psi(j) /
                       sqrt(pow(n, 2) - pow(lambd * psi(j), 2) / XX(j, j))) *
                         sqrt(qqhat) -
                     S0) /
                    XX(j, j);
          error -= X.col(j) * beta(j);
        }

        else {
          beta(j) = 0;
        }
      } // end loop beta^(i)_j

      // Update primal and dual value
      double errorNorm = (y - X * beta).norm();
      double fobj =
          errorNorm / sqrt(n) + (lambd / n) * (psi * beta.cwiseAbs()).sum();

      double dual;
      if (errorNorm > MaxErrorNorm) {
        VectorXd aaa = sqrt(n) * (error / errorNorm);
        double bbb = ((lambd / n) * psi - (X.transpose() * aaa / n).cwiseAbs()).cwiseAbs().transpose() * beta.cwiseAbs();
        dual = aaa.transpose() * (y / n) - bbb;
      } else {
        dual = (lambd / n) * (psi * beta.cwiseAbs()).sum();
      }

      // check for convergence
      double diff = (beta - betaOld).cwiseAbs().sum();
      // double diff = ((beta - betaOld).cwiseAbs() * (sdVec / ySd)).sum();
      if (diff < optTol) {
        if ((fobj - dual)  < 1e-6) {
        // if ((fobj - dual) / ySd < optTol) {
          break;
        }
      }
    }
  }
  // set beta to zero below threshold
  for (int j = 0; j < p; j++) {
    if (abs(beta(j)) < zeroTol) {
      beta(j) = 0;
    }
  }
  return beta;
}

// docstring
string docstring = R"mydelimiter(
    "Lasso Shooting algorithm for and sqrt lasso."

    Parameters
    ----------
    X : numpy array
        design matrix
    y : numpy array
        response vector
    XX : ndarray
      cross product matrix of X
    Xy : ndarray
      cross product vector of X and y
    lambd : float
        Regularization parameter.
    psi : float
        Penalty loadings.
    starting_values : ndarray, optional, default: None
        Initial beta estimate.
    sqrt : bool, optional, default False
        If True, use sqrt lasso.
        beta = min ||(y - X @ beta)||_2^2 + lambd ||psi @ beta||_1
    fit_intercept : bool, optional, default true
        If True, fit intercept. Only relevant for sqrt lasso.
    max_iter : int, optional, default: 1000
        Maximum number of iterations.
    opt_tol : float, optional, default: 1e-10
        Optimality tolerance.
    zero_tol : float, optional, default: 1e-4
        Zero tolerance. If beta(j) is smaller than zero_tol,
        set beta(j) = 0.

    Returns
    -------
    beta : ndarray
        Estimated beta.
)mydelimiter";

PYBIND11_MODULE(_solver_fast, m) {
  py::options options;
  options.disable_function_signatures();
  m.doc() = "Coordinate descent solver for lasso and sqrt-lasso";
  m.def("_cd_solver", &coordinateDescent, docstring.c_str(),
        py::arg("X").noconvert() = NULL,
        py::arg("y").noconvert() = NULL,
        py::arg("XX").noconvert() = NULL,
        py::arg("Xy").noconvert() = NULL,
        py::arg("lambd").noconvert() = NULL,
        py::arg("psi").noconvert() = NULL,
        py::arg("starting_values").noconvert() = NULL,
        py::arg("sqrt") = false,
        py::arg("fit_intercept") = true,
        py::arg("opt_tol") = 1e-10,
        py::arg("max_iter") = 1000,
        py::arg("zero_tol") = 1e-4);
};


<<<<<<< HEAD
/*    Copyright (c) 2010-2018, Delft University of Technology
=======
/*    Copyright (c) 2010-2019, Delft University of Technology
>>>>>>> origin/master
 *    All rigths reserved
 *
 *    This file is part of the Tudat. Redistribution and use in source and
 *    binary forms, with or without modification, are permitted exclusively
 *    under the terms of the Modified BSD license. You should have received
 *    a copy of the license with this file. If not, please or visit:
 *    http://tudat.tudelft.nl/LICENSE.
 */

#include <Eigen/Geometry>

#include "Tudat/Astrodynamics/Relativity/relativisticAccelerationCorrection.h"

namespace tudat
{

namespace relativity
{

//! Function to compute a term common to several relativistic acceleration terms
double calculateRelativisticAccelerationCorrectionsCommonterm(
        const double centralBodyGravitationalParameter,
        const double relativeDistance )
{
    return centralBodyGravitationalParameter /
            ( physical_constants::SPEED_OF_LIGHT * physical_constants::SPEED_OF_LIGHT *
              relativeDistance * relativeDistance * relativeDistance );
}

//! Function to compute the Schwarzschild term of the relativistic acceleration correction.
Eigen::Vector3d calculateScharzschildGravitationalAccelerationCorrection(
        const double centralBodyGravitationalParameter,
        const Eigen::Vector3d& relativePosition,
        const Eigen::Vector3d& relativeVelocity,
        const double relativeDistance,
        const double commonCorrectionTerm,
        const double ppnParameterGamma,
        const double ppnParameterBeta )
{
    Eigen::Vector3d acceleration = ( 2.0 * ( ppnParameterGamma + ppnParameterBeta ) *
                      centralBodyGravitationalParameter / relativeDistance - ppnParameterGamma *
                      relativeVelocity.dot( relativeVelocity ) ) * relativePosition +
            2.0 * ( 1.0 + ppnParameterGamma ) * ( relativePosition.dot( relativeVelocity ) ) * relativeVelocity;
    return commonCorrectionTerm * acceleration;
}

//! Function to compute the Schwarzschild term of the relativistic acceleration correction.
Eigen::Vector3d calculateScharzschildGravitationalAccelerationCorrection(
        double centralBodyGravitationalParameter,
        const Eigen::Vector6d& relativeState,
        double ppnParameterGamma,
        double ppnParameterBeta )
{
    return calculateScharzschildGravitationalAccelerationCorrection(
                centralBodyGravitationalParameter, relativeState.segment( 0, 3 ),
                relativeState.segment( 3, 3 ), relativeState.segment( 0, 3 ).norm( ),
                calculateRelativisticAccelerationCorrectionsCommonterm(
                    centralBodyGravitationalParameter, relativeState.segment( 0, 3 ).norm( ) ),
                ppnParameterGamma, ppnParameterBeta );
}

//! Function to compute the Lense-Thirring term of the relativistic acceleration correction.
Eigen::Vector3d calculateLenseThirringCorrectionAcceleration(
        const Eigen::Vector3d& relativePosition,
        const Eigen::Vector3d& relativeVelocity,
        const double relativeDistance,
        const double commonCorrectionTerm,
        const Eigen::Vector3d& centralBodyAngularMomentum,
        const double ppnParameterGamma )
{
    Eigen::Vector3d acceleration = 3.0 / (
                relativeDistance * relativeDistance ) *
            relativePosition.cross( relativeVelocity ) *
            ( relativePosition.dot( centralBodyAngularMomentum ) ) +
            relativeVelocity.cross( centralBodyAngularMomentum );
    return acceleration * ( 1.0 + ppnParameterGamma ) * commonCorrectionTerm;
}

//! Function to compute the Lense-Thirring term of the relativistic acceleration correction.
Eigen::Vector3d calculateLenseThirringCorrectionAcceleration(
        const double centralBodyGravitationalParameter,
        const Eigen::Vector6d& relativeState,
        const Eigen::Vector3d& centralBodyAngularMomentum,
        const double ppnParameterGamma )
{
    return calculateLenseThirringCorrectionAcceleration(
                relativeState.segment( 0, 3 ), relativeState.segment( 3, 3 ), relativeState.segment( 0, 3 ).norm( ),
                calculateRelativisticAccelerationCorrectionsCommonterm(
                    centralBodyGravitationalParameter, relativeState.segment( 0, 3 ).norm( ) ),
                centralBodyAngularMomentum, ppnParameterGamma );

}

//! Function to compute the de Sitter term of the relativistic acceleration correction.
Eigen::Vector3d calculateDeSitterCorrectionAcceleration(
        const Eigen::Vector3d& orbiterRelativeVelocity,
        const Eigen::Vector3d& orbitedBodyPositionWrtLargerBody,
        const Eigen::Vector3d& orbitedBodyVelocityWrtLargerBody,
        const double commonCorrectionTermOfLargerBody,
        const double ppnParameterGamma )
{
    return - commonCorrectionTermOfLargerBody * ( 1.0 + 2.0 * ppnParameterGamma ) *
            ( orbitedBodyVelocityWrtLargerBody.cross( orbitedBodyPositionWrtLargerBody ) ).cross( orbiterRelativeVelocity );
}

//! Function to compute the de Sitter term of the relativistic acceleration correction.
Eigen::Vector3d calculateDeSitterCorrectionAcceleration(
        const double largerBodyGravitationalParameter,
        const Eigen::Vector6d& orbiterRelativeState,
        const Eigen::Vector6d&orbitedBodyStateWrtLargerBody,
        const double ppnParameterGamma )
{
    return calculateDeSitterCorrectionAcceleration(
                orbiterRelativeState.segment( 3, 3 ),
                orbitedBodyStateWrtLargerBody.segment( 0, 3 ),
                orbitedBodyStateWrtLargerBody.segment( 3, 3 ),
                calculateRelativisticAccelerationCorrectionsCommonterm(
                    largerBodyGravitationalParameter,
                    orbitedBodyStateWrtLargerBody.segment( 0, 3 ).norm( ) ),
                ppnParameterGamma );
}

}

}


// The MIT License (MIT)
//
// Copyright (c) 2014 Julian Gehring
//
// Permission is hereby granted, free of charge, to any person obtaining a copy
// of this software and associated documentation files (the "Software"), to deal
// in the Software without restriction, including without limitation the rights
// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
// copies of the Software, and to permit persons to whom the Software is
// furnished to do so, subject to the following conditions:
//
// The above copyright notice and this permission notice shall be included in all
// copies or substantial portions of the Software.
//
// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
// SOFTWARE.

#include <boost/math/distributions/hypergeometric.hpp>
#include <algorithm>

double fisher_test(int a, int b, int c, int d) {
  unsigned N = a + b + c + d;
  unsigned r = a + c;
  unsigned n = c + d;
  unsigned max_for_k = std::min(r, n);
  unsigned min_for_k = (unsigned)std::max(0, int(r + n - N));
  boost::math::hypergeometric_distribution<> hgd(r, n, N);
  double cutoff = pdf(hgd, c);
  double tmp_p = 0.0;
  for(int k = min_for_k;k < max_for_k + 1;k++) {
    double p = pdf(hgd, k);
    if(p <= cutoff) tmp_p += p;
  }
  return tmp_p;
}


/**
 * @file stableevaluationatapoint_main.cc
 * @brief NPDE homework StableEvaluationAtAPoint
 * @author Amélie Loher, Erick Schulz & Philippe Peter
 * @date 29.11.2021
 * @copyright Developed at SAM, ETH Zurich
 */

#include <lf/assemble/assemble.h>
#include <lf/io/io.h>
#include <lf/mesh/hybrid2d/hybrid2d.h>
#include <lf/mesh/mesh.h>
#include <lf/uscalfe/uscalfe.h>

#include <Eigen/Core>
#include <cmath>
#include <cstdlib>
#include <fstream>
#include <iostream>
#include <memory>
#include <string>
#include <utility>

#include "stableevaluationatapoint.h"

int main(int /*argc*/, const char ** /*argv*/) {
  // exact solution
  auto uExact = [](Eigen::Vector2d x) -> double {
    Eigen::Vector2d one(1.0, 0.0);
    return std::log((x + one).norm());
  };
  // Fixed evaluation point
  Eigen::Vector2d x(0.3, 0.4);
  std::cout << "Exact evaluation at (0.3,0.4) : " << uExact(x) << std::endl;

  // Number of meshes used in the error analysis:
  int N_meshes = 6;

  // Error analysis vectors:
  Eigen::VectorXd mesh_sizes(N_meshes);
  mesh_sizes.setZero();
  Eigen::VectorXd dofs(N_meshes);
  dofs.setZero();

  // Error vector used for the error analysis in exercise b)
  Eigen::VectorXd errors_potential(N_meshes);
  errors_potential.setZero();

  // Error vectors used for the error analysis in exercise h
  Eigen::VectorXd errors_direct(N_meshes);
  errors_direct.setZero();
  Eigen::VectorXd errors_stable(N_meshes);
  errors_stable.setZero();

  // iterate over meshes:
  for (int k = 0; k < N_meshes; k++) {
    // read mesh::
    std::string idx = std::to_string(k + 1);
    auto mesh_factory = std::make_unique<lf::mesh::hybrid2d::MeshFactory>(2);
    lf::io::GmshReader reader(std::move(mesh_factory), CURRENT_SOURCE_DIR
                                                           "/../meshes/square" +
                                                           idx + ".msh");
    auto mesh_p = reader.mesh();

    // Initialize fe-space and dofh
    auto fe_space =
        std::make_shared<lf::uscalfe::FeSpaceLagrangeO1<double>>(mesh_p);
    const lf::assemble::DofHandler &dofh = fe_space->LocGlobMap();
    dofs(k) = dofh.NumDofs();

    // Printing mesh statistics
    mesh_sizes(k) = StableEvaluationAtAPoint::MeshSize(mesh_p);
    std::cout << "square" + idx + ".msh: "
              << "N_dofs = " << dofs(k) << ", h=" << mesh_sizes(k) << std::endl;

    // Error anlysis part b) (Potentials)
    errors_potential(k) = StableEvaluationAtAPoint::PointEval(mesh_p);

    // error analysis part g/h: Compare direct vs stable point evaluation:
    auto [direct_eval, stable_eval] =
        StableEvaluationAtAPoint::ComparePointEval(fe_space, uExact, x);
    errors_direct(k) = std::abs(uExact(x) - direct_eval);
    errors_stable(k) = std::abs(uExact(x) - stable_eval);
  }

  // Compute rates of convergence:
  Eigen::VectorXd rates_potential(N_meshes - 1);
  Eigen::VectorXd rates_direct(N_meshes - 1);
  Eigen::VectorXd rates_stable(N_meshes - 1);

#if SOLUTION
  for (int k = 0; k < N_meshes - 1; ++k) {
    double log_denum = std::log(mesh_sizes(k) / mesh_sizes(k + 1));
    rates_potential(k) =
        std::log(errors_potential(k) / errors_potential(k + 1)) / log_denum;
    rates_direct(k) =
        std::log(errors_direct(k) / errors_direct(k + 1)) / log_denum;
    rates_stable(k) =
        std::log(errors_stable(k) / errors_stable(k + 1)) / log_denum;
  }
#else
  //====================
  // Your code goes here
  //====================
#endif

  // Report computed errors and rates:
  std::cout << "Subtask b) Evaluation based on Potentials \n";
  std::cout << "Errors: \n" << errors_potential << "\n";
  std::cout << "Rates: \n" << rates_potential << "\n";
  std::cout << "Subtask h) Comparison of direct and stable evaluation: \n";
  std::cout << "Errors direct: \n" << errors_direct << "\n";
  std::cout << "Rates direct: \n" << rates_direct << "\n";
  std::cout << "Errors stable: \n" << errors_stable << "\n";
  std::cout << "Rates stable: \n" << rates_stable << "\n";

  // Output
  const static Eigen::IOFormat CSVFormat(Eigen::StreamPrecision,
                                         Eigen::DontAlignCols, ", ", "\n");

  Eigen::MatrixXd convergence_potential(N_meshes, 2);
  convergence_potential << mesh_sizes, errors_potential;

  Eigen::MatrixXd convergence_stable(N_meshes, 3);
  convergence_stable << mesh_sizes, errors_direct, errors_stable;

  std::ofstream file;
  file.open("convergence_potential.csv");
  file << "h, Error u(x) (Potential) \n";
  file << convergence_potential.format(CSVFormat);
  file.close();
  std::cout << "Generated " CURRENT_BINARY_DIR "/convergence_potential.csv"
            << std::endl;

  file.open("convergence_stable.csv");
  file << "h, Error u(x) (Direct), Error u(x) (Stable) \n";
  file << convergence_stable.format(CSVFormat);
  file.close();
  std::cout << "Generated " CURRENT_BINARY_DIR "/convergence_stable.csv"
            << std::endl;

  // Plot
  std::system("python3 " CURRENT_SOURCE_DIR
              "/plot_convergence_potential.py " CURRENT_BINARY_DIR);
  std::system("python3 " CURRENT_SOURCE_DIR
              "/plot_convergence_stable.py " CURRENT_BINARY_DIR);

  return 0;
}


#pragma once

#include <Eigen/Dense>
#include <Eigen/Geometry>
#include <numeric>
#include <vector>
#include <memory>
#include <map>
#include <tuple>
#include <exception>
#include <iostream>
#include <functional>
#include <mutex>          // std::mutex
#include "ThreadPool.h"
#include "cubature.h"
#if !defined(NO_CUBA)
#include "cuba.h"
#endif
#include "MultiComplex/MultiComplex.hpp"

#include "nlohmann/json.hpp"

double factorial(double n) {
    return std::tgamma(n + 1);
}

std::mutex mtx;  // mutex for cout

#if !defined(M_PI)
constexpr auto M_PI = 3.14159265358979323846;
#endif

auto geomspace(double xmin, double xmax, int N) {
    std::vector<double> vec;
    double dT = (log(xmax) - log(xmin)) / (N - 1);
    for (auto i = 0; i < N; ++i) {
        vec.push_back(exp(log(xmin) + dT * i));
    }
    return vec;
}

double gr = (sqrt(5) + 1) / 2;

auto gss(std::function<double(double)> f, double a, double b, const double tol = 1e-5) {
    /*
    Golden section search
    C++ translation of https://en.wikipedia.org/wiki/Golden-section_search#Algorithm
    Text is available under the Creative Commons Attribution-ShareAlike License
    */
    auto c = b - (b - a) / gr;
    auto d = a + (b - a) / gr;
    while (abs(c - d) > tol) {
        if (f(c) < f(d)) {
            b = d;
        }
        else {
            a = c;
        }
        // We recompute both c and d here to avoid loss of precision which may lead to incorrect results or infinite loop
        c = b - (b - a) / gr;
        d = a + (b - a) / gr;
    }
    return std::make_tuple((b + a) / 2, f((b + a) / 2));
}

/*
Trapezoidal integration -- the workhorse numerical integration routine
*/
template<typename TYPEX, typename TYPEY>
TYPEY trapz(const Eigen::Array<TYPEX, Eigen::Dynamic, 1>& x,
           const Eigen::Array<TYPEY, Eigen::Dynamic, 1>& y) {
    TYPEY out = 0;
    for (auto i = 0; i < x.size()-1; ++i) {
        auto ymean = (y[i+1]+y[i])/2.0;
        out = out + ymean*(x[i+1] - x[i]);
    }
    return out;
}
/*
Simpson's integration rule
*/
template<typename TYPEX, typename TYPEY>
TYPEY simps(const Eigen::Array<TYPEX, Eigen::Dynamic, 1>& x,
           const Eigen::Array<TYPEY, Eigen::Dynamic, 1>& f) {
    // C++ translation of https://en.wikipedia.org/wiki/Simpson%27s_rule#Composite_Simpson's_rule_for_irregularly_spaced_data
    auto N = x.size() - 1;
    auto h = x.tail(N) - x.head(N);
    auto cube = [](TYPEX x) { return x * x * x; };
    auto sq = [](TYPEX x) { return x * x; };
    TYPEY result = 0.0;
    for (auto i = 1; i < N; i += 2) {
        auto hph = h[i] + h[i - 1];
        result += f[i] * (cube(h[i]) + cube(h[i - 1]) + 3.0 * h[i] * h[i - 1] * hph) / (6 * h[i] * h[i - 1]);
        result += f[i - 1] * (2.0 * cube(h[i - 1.0]) - cube(h[i]) + 3.0 * h[i] * sq(h[i - 1])) / (6 * h[i - 1] * hph);
        result += f[i + 1] * (2.0 * cube(h[i]) - cube(h[i - 1]) + 3.0 * h[i - 1] * sq(h[i])) / (6 * h[i] * hph);
    }
    if ((N + 1) % 2 == 0) {
        result += f[N] * (2 * sq(h[N - 1]) + 3.0 * h[N - 2] * h[N - 1]) / (6 * (h[N - 2] + h[N - 1]));
        result += f[N - 1] * (sq(h[N - 1]) + 3.0 * h[N - 1] * h[N - 2]) / (6 * h[N - 2]);
        result -= f[N - 2] * cube(h[N - 1]) / (6 * h[N - 2] * (h[N - 2] + h[N - 1]));
    }
    return result;
}

template <typename TYPE>
class Molecule {

public:
    using EColArray = Eigen::Array<TYPE, Eigen::Dynamic, 1>;
    using CoordMatrix = Eigen::Array<TYPE, 3, Eigen::Dynamic>;
    CoordMatrix coords, coords_initial;

    Molecule(const std::vector<std::vector<TYPE>>& pts) {
        coords.resize(3, pts.size());
        for (auto i = 0; i < pts.size(); ++i) {
            auto &pt = pts[i];
            for (auto j = 0; j < pt.size(); ++j) {
                coords(j, i) = pt[j];
            }
        }
        coords_initial = coords;
    };
    void reset() {
        coords = coords_initial;
    };
    CoordMatrix rotZ3(TYPE theta) const {
        // Rotation matrix for rotation around the z axis in 3D
        // See https://en.wikipedia.org/wiki/Rotation_matrix#Basic_rotations
        return (Eigen::ArrayXXd(3,3) <<
            cos(theta), -sin(theta), 0,
            sin(theta), cos(theta),  0,
                0,           0,      1 ).finished();
    };
    CoordMatrix rotY3(const TYPE theta) const {
        // Rotation matrix for rotation around the y axis in 3D
        // See https://en.wikipedia.org/wiki/Rotation_matrix#Basic_rotations
        const TYPE c = cos(theta), s = sin(theta);
        return (Eigen::ArrayXXd(3,3) <<
            c,  0, s,
            0,  1, 0,
            -s, 0, c).finished();
    }
    CoordMatrix rotX3(TYPE theta) const {
        // Rotation matrix for rotation around the x axis in 3D
        // See https://en.wikipedia.org/wiki/Rotation_matrix#Basic_rotations
        return (Eigen::ArrayXXd(3,3) <<
            1,      0,          0,
            0, cos(theta), -sin(theta),
            0, sin(theta), cos(theta)).finished();
    }
    void rotate_plusx(TYPE angle) {
        Eigen::Transform<double, 3, Eigen::Affine> rot(Eigen::AngleAxisd(angle, Eigen::Vector3d::UnitX()));
        coords = rot.linear() * coords.matrix();
        //coords = rotX3(angle).matrix() * coords.matrix(); // Old method
    }
    void rotate_negativex(TYPE angle) {
        rotate_plusx(-angle);
    }
    void rotate_plusy(TYPE angle) {
        Eigen::Transform<double, 3, Eigen::Affine> rot(Eigen::AngleAxisd(angle, Eigen::Vector3d::UnitY()));
        coords = rot.linear() * coords.matrix();
        //coords = rotY3(angle).matrix() * coords.matrix(); // Old method
    }
    void rotate_negativey(TYPE angle) {
        rotate_plusy(-angle);
    }
    void rotate_plusz(TYPE angle) {
        Eigen::Transform<double, 3, Eigen::Affine> rot(Eigen::AngleAxisd(angle, Eigen::Vector3d::UnitZ()));
        coords = rot.linear()*coords.matrix();
    }
    void rotate_negativez(TYPE angle) {
        rotate_plusz(-angle);
    }
    void translatex(TYPE dx) {
        coords.row(0) += dx;
    }
    void translatey(TYPE dy) {
        coords.row(1) += dy;
    }
    template<typename ArrayType>
    TYPE get_dist(const ArrayType&rA, const ArrayType&rB) const{
        return sqrt((rA - rB).square().sum());
    }
    Eigen::Index get_Natoms() const {
        return coords.cols();
    }
    auto get_xyz_atom(const Eigen::Index i) const {
        if (i > coords.cols()) {
            throw std::invalid_argument("Bad atom index");
        }
        return coords.col(i);
    }
};

template<typename TYPE>
class PotentialEvaluator {
public:
    std::map<std::tuple<std::size_t, std::size_t>, std::function<double(double)> > potential_map;

    ///< Additional contributions, handled in a very generic way via a callback
    std::vector<std::function<double(const Molecule<TYPE> &, const Molecule<TYPE>&)>> generic_contributions;

    TYPE eval_pot(const Molecule<TYPE>& molA, const Molecule<TYPE>& molB) const {
        TYPE u = 0.0;
        // Sum up site-site contributions to total potential energy
        for (auto iatom = 0; iatom < molA.get_Natoms(); ++iatom) {
            auto xyzA = molA.get_xyz_atom(iatom);
            for (auto jatom = 0; jatom < molB.get_Natoms(); ++jatom) {
                auto xyzB = molB.get_xyz_atom(jatom);
                TYPE distij = molA.get_dist(xyzA, xyzB);
                auto f_potential = get_potential(iatom, jatom);
                u += f_potential(distij);
            }
        }
        // Add also contributions for other kinds of interactions (e.g., dipole, quadrupole, etc.)
        for (auto contrib : generic_contributions) {
            u += contrib(molA, molB);
        }
        return u;
    }
    /*
    Connect up all the site-site potentials, all molecules have the same number of sites
    */
    void connect_potentials(std::function<double(double)>& f,std::size_t Natoms) {
        for (auto iatom = 0; iatom < Natoms; ++iatom) {
            for (auto jatom = 0; jatom < Natoms; ++jatom) {
                potential_map[std::make_tuple(iatom, jatom)] = f;
            }
        }
    }
    /*
    Add a potential for a particular site-site interaction

    @param iatom The index of site on molecule A
    @param jatom The index of site on molecule B
    */
    void add_potential(std::size_t iatom, std::size_t jatom, std::function<double(double)>& f) {
        potential_map[std::make_tuple(iatom, jatom)] = f;
    }
    /*
    Add a generalized contribution to the overall potential energy.  Callback function takes two molecules
    */
    void add_generic_contribution(const std::function<double(const Molecule<TYPE>&, const Molecule<TYPE>&)>& f) {
        generic_contributions.push_back(f);
    }
    /*
    Get a reference to the potential function for a particular site-site interaction

    @param i The index of site on molecule A
    @param j The index of site on molecule B

    */
    auto& get_potential(std::size_t i, std::size_t j) const {
        auto itf = potential_map.find(std::make_tuple(i, j));
        if (itf != potential_map.end()) {
            return itf->second;
        }
        else {
            throw std::invalid_argument("Bad potential");
        }
    }
};
/// A helper class
template<typename TYPE, typename TEMPTYPE>
class SharedDataBase {
public:
    TEMPTYPE Tstar;
    TYPE rstar;
    Molecule<TYPE> molA, molB;
    std::valarray<TYPE> xmin, xmax;
    const PotentialEvaluator<TYPE>& evaltr;
    SharedDataBase(TEMPTYPE Tstar, TYPE rstar,
        Molecule<TYPE> molA, Molecule<TYPE> molB,
        const PotentialEvaluator<TYPE>& evaltr,
        const std::valarray<TYPE>& xmin = {},
        const std::valarray<TYPE>& xmax = {}
    )
        : Tstar(Tstar), rstar(rstar), molA(molA), molB(molB), evaltr(evaltr), xmin(xmin), xmax(xmax) {};

    TYPE eval_pot(const Molecule<TYPE>& molA, const Molecule<TYPE>& molB) {
        return evaltr.eval_pot(molA, molB);
    };

    /*
    Given the orientational angles, calculate the integrand
    */
    void orient_integrand(double theta1, double theta2, double phi, double *fval)
    {
        // Rotate molecule #1
        molA.reset(); // Back to COM at origin
        molA.rotate_negativey(theta1); // First rotate around -y axis

        // Rotate and move molecule #2
        molB.reset(); // Back to COM at origin
        molB.rotate_negativey(theta2); // First rotate around -y axis
        molB.rotate_negativex(phi); // Then rotate around +x
        molB.translatex(rstar); // Then translate

        auto V = eval_pot(molA, molB); // And finally evaluate the potential
        auto a = (exp(-V/Tstar)-1.0)*sin(theta1)*sin(theta2)*pow(rstar, 2);

        if constexpr (std::is_same<decltype(Tstar), double>::value) {
            // If T is double (real)
            fval[0] = a;
        }
        else if constexpr (std::is_same<decltype(Tstar), std::complex<double>>::value) {
            // If T is a complex number (perhaps for complex step derivatives)
            fval[0] = a.real();
            fval[1] = a.imag();
        }
        else if constexpr (std::is_same<decltype(Tstar), MultiComplex<double>>::value) {
            // If T is a multicomplex number
            auto &c = a.get_coef();
            for (auto i = 0; i < c.size(); ++i) {
                fval[i] = c[i];
            }
        }
    }
    /*
    Given the separations and angle, calculate the integrand for B_3 for a spherically-symmetric potential
    for an atomic fluid
    */
    void atomic_B3_integrand(const double r12, const double r13, const double eta_angle, double* fval)
    {
        // Get the potential function V(r) that we should use
        auto &pot = this->evaltr.get_potential(0, 0);
        TEMPTYPE Tstar = this->Tstar; // Local reference just for sharing with the lambda function f
        auto f = [pot, Tstar](double r) -> TEMPTYPE { return 1.0 - exp(-pot(r)/Tstar); };
        auto SQUARE = [](double x) { return x*x; };
        auto rangle = sqrt(SQUARE(r12) + SQUARE(r13) - 2*r12*r13*eta_angle);
        auto a = SQUARE(r12)*f(r12)*SQUARE(r13)*f(r13)*f(rangle);

        if constexpr (std::is_same<decltype(Tstar), double>::value) {
            // If T is double (real)
            fval[0] = a;
        }
        else if constexpr (std::is_same<decltype(Tstar), std::complex<double>>::value) {
            // If T is a complex number (perhaps for complex step derivatives)
            fval[0] = a.real();
            fval[1] = a.imag();
        }
        else if constexpr (std::is_same<decltype(Tstar), MultiComplex<double>>::value) {
            // If T is a multicomplex number
            auto& c = a.get_coef();
            for (auto i = 0; i < c.size(); ++i) {
                fval[i] = c[i];
            }
        }
    }

	/*
	Given the separations and angles, calculate the integrand for B4_1 for a spherically-symmetric potential
	for an atomic fluid
	*/
	void atomic_B4_1_integrand(const double r14, const double r13, const double gamma_angle, const double r12, const double eta_angle, double* fval)
	{
		// Get the potential function V(r) that we should use
		auto &pot = this->evaltr.get_potential(0, 0);
		TEMPTYPE Tstar = this->Tstar; // Local reference just for sharing with the lambda function f
		auto f = [pot, Tstar](double r) -> TEMPTYPE { return 1.0 - exp(-pot(r) / Tstar); };
		auto SQUARE = [](double x) { return x * x; };
		auto sq_r12 = SQUARE(r12);
		auto sq_r13 = SQUARE(r13);
		auto sq_r14 = SQUARE(r14);
		auto rangle_12_13 = sqrt(sq_r12 + sq_r13 - 2 * r12*r13*eta_angle);
		auto rangle_13_14 = sqrt(sq_r14 + sq_r13 - 2 * r14*r13*gamma_angle);

		auto a = sq_r12 * f(r12)*sq_r13*sq_r14*f(r14)*f(rangle_12_13)*f(rangle_13_14);

		if constexpr (std::is_same<decltype(Tstar), double>::value) {
			// If T is double (real)
			fval[0] = a;
		}
		else if constexpr (std::is_same<decltype(Tstar), std::complex<double>>::value) {
			// If T is a complex number (perhaps for complex step derivatives)
			fval[0] = a.real();
			fval[1] = a.imag();
		}
		else if constexpr (std::is_same<decltype(Tstar), MultiComplex<double>>::value) {
			// If T is a multicomplex number
			auto& c = a.get_coef();
			for (auto i = 0; i < c.size(); ++i) {
				fval[i] = c[i];
			}
		}
	}

	/*
Given the separations and angles, calculate the integrand for B4_2 for a spherically-symmetric potential
for an atomic fluid
*/
	void atomic_B4_2_integrand(const double eta_angle, const double r12, const double r13, const double gamma_angle, const double r14, double* fval)
	{
		// Get the potential function V(r) that we should use
		auto &pot = this->evaltr.get_potential(0, 0);
		TEMPTYPE Tstar = this->Tstar; // Local reference just for sharing with the lambda function f
		auto f = [pot, Tstar](double r) -> TEMPTYPE { return 1.0 - exp(-pot(r) / Tstar); };
		auto SQUARE = [](double x) { return x * x; };
		auto sq_r12 = SQUARE(r12);
		auto sq_r13 = SQUARE(r13);
		auto sq_r14 = SQUARE(r14);
		auto rangle_12_13 = sqrt(sq_r12 + sq_r13 - 2 * r12*r13*eta_angle);
		auto rangle_13_14 = sqrt(sq_r14 + sq_r13 - 2 * r14*r13*gamma_angle);

		auto a = sq_r12 * f(r12)*sq_r13*f(r13)*sq_r14*f(r14)*f(rangle_12_13)*f(rangle_13_14);

		if constexpr (std::is_same<decltype(Tstar), double>::value) {
			// If T is double (real)
			fval[0] = a;
		}
		else if constexpr (std::is_same<decltype(Tstar), std::complex<double>>::value) {
			// If T is a complex number (perhaps for complex step derivatives)
			fval[0] = a.real();
			fval[1] = a.imag();
		}
		else if constexpr (std::is_same<decltype(Tstar), MultiComplex<double>>::value) {
			// If T is a multicomplex number
			auto& c = a.get_coef();
			for (auto i = 0; i < c.size(); ++i) {
				fval[i] = c[i];
			}
		}
	}

	/*
Given the separations and angles, calculate the integrand for B4_3 for a spherically-symmetric potential
for an atomic fluid
*/
	void atomic_B4_3_integrand(const double eta_angle, const double zeta_angle, const double gamma_angle, const double r12, const double r13, const double r14, double* fval)
	{
		// Get the potential function V(r) that we should use
		auto &pot = this->evaltr.get_potential(0, 0);
		TEMPTYPE Tstar = this->Tstar; // Local reference just for sharing with the lambda function f
		auto f = [pot, Tstar](double r) -> TEMPTYPE { return 1.0 - exp(-pot(r) / Tstar); };
		auto SQUARE = [](double x) { return x * x; };
		auto sq_r12 = SQUARE(r12);
		auto sq_r13 = SQUARE(r13);
		auto sq_r14 = SQUARE(r14);
		auto rangle_12_13 = sqrt(sq_r12 + sq_r14 - 2 * r12*r14*eta_angle);
		auto rangle_13_14 = sqrt(sq_r13 + sq_r14 - 2 * r13*r14*gamma_angle);
		auto rangle_12_14 = sqrt(sq_r12 + sq_r13 - 2 * r12*r13*(eta_angle * gamma_angle + sqrt(1.0 - SQUARE(eta_angle))*sqrt(1.0 - SQUARE(gamma_angle))*cos(zeta_angle)));
		auto a = sq_r12 * f(r12)*sq_r13*f(r13)*sq_r14*f(r14)*f(rangle_12_13)*f(rangle_13_14)*f(rangle_12_14);

		if constexpr (std::is_same<decltype(Tstar), double>::value) {
			// If T is double (real)
			fval[0] = a;
		}
		else if constexpr (std::is_same<decltype(Tstar), std::complex<double>>::value) {
			// If T is a complex number (perhaps for complex step derivatives)
			fval[0] = a.real();
			fval[1] = a.imag();
		}
		else if constexpr (std::is_same<decltype(Tstar), MultiComplex<double>>::value) {
			// If T is a multicomplex number
			auto& c = a.get_coef();
			for (auto i = 0; i < c.size(); ++i) {
				fval[i] = c[i];
			}
		}
	}
};

template<typename TYPE>
class Integrator {
private:
    std::unique_ptr<ThreadPool> m_pool;
    nlohmann::json m_conf;

public:
    using EColArray = Eigen::Array<TYPE, Eigen::Dynamic, 1>;
    const Molecule<TYPE> mol1, mol2;
    PotentialEvaluator<TYPE> potcls;

    Integrator(const Molecule<TYPE>& mol1, const Molecule<TYPE>& mol2) : mol1(mol1), mol2(mol2) {};

    auto& get_conf_view() {
        return m_conf;
    }

    /* For a one-dimensional integration for B_2, use trapezoidal integration to calculate B_2 */
    template <typename TEMPTYPE>
    TEMPTYPE radial_integrate_B2(TEMPTYPE Tstar, TYPE rstart, TYPE rend, int N) {
        using arr = Eigen::Array<TYPE, Eigen::Dynamic, 1>;
        using arrT = Eigen::Array<TEMPTYPE, Eigen::Dynamic, 1>;
        arr rv = exp(arr::LinSpaced(N, log(rstart), log(rend)));
        arrT integrand;
        integrand.resize(rv.size());
        Molecule<TYPE> mol1 = this->mol1, mol2 = this->mol2;
        for (auto ir = 0; ir < rv.size(); ++ir) {
            auto r = rv[ir];
            mol2.reset();
            mol2.translatex(r);
            auto V = potcls.eval_pot(mol1, mol2);
            integrand[ir] = (exp(-V/Tstar)-1.0)*r*r;
        }
        return -2*M_PI*trapz(rv, integrand);
    };
    /*
    Get a reference to the evaluator class, giving access to matrix of site-site potential functions, for instance
    */
    auto &get_evaluator() {
        return potcls;
    }

    void init_thread_pool(short Nthreads) {
        if (!m_pool || m_pool->GetThreads().size() != Nthreads) {
            // Make a thread pool for the workers
            m_pool = std::unique_ptr<ThreadPool>(new ThreadPool(Nthreads));
        }
    }
    /*
    A helper function to evaluate the potential given COM separation r and the orientation angles
    */
    double potential(double r, double theta1, double theta2, double phi){
        Molecule<TYPE> molA = mol1, molB = mol2;
        // Rotate molecule #1
        molA.reset(); // Back to COM at origin
        molA.rotate_negativey(theta1); // First rotate around -y axis

        // Rotate and move molecule #2
        molB.reset(); // Back to COM at origin
        molB.rotate_negativey(theta2); // First rotate around -y axis
        molB.rotate_negativex(phi); // Then rotate around +x
        molB.translatex(r); // Then translate

        auto V = potcls.eval_pot(molA, molB); // And finally evaluate the potential in the form V/epsilon
        return V;
    }
    /*
    Calculate the orientationally-averaged potential
    */
    TYPE orient_averaged_potential(TYPE rstar, Molecule<TYPE> mol1, Molecule<TYPE> mol2) const {
        using SharedData = SharedDataBase<TYPE, double>;
        SharedData shared(0.0, 0.0, mol1, mol2, potcls, {0,0,0}, { M_PI, M_PI, 2 * M_PI });
        //typedef int (*integrand) (unsigned ndim, const double *x, void *, unsigned fdim, double* fval);
        auto f_integrand = [](unsigned ndim, const double* x, void* p_shared_data, unsigned fdim, double* fval) {
            auto& shared = *((class SharedDataBase<double, double>*)(p_shared_data));
            auto& molA = shared.molA;
            auto& molB = shared.molB;
            double theta1 = x[0], theta2 = x[1], phi = x[2];
            shared.orient_integrand(theta1, theta2, phi, fval);
            return 0; // success
        };
        shared.rstar = rstar;
        int ndim = 1;
        std::valarray<double> val(0.0, 4), err(0.0, 4);
        hcubature(ndim, f_integrand, &shared, 3, &(shared.xmin[0]), &(shared.xmax[0]), 100000, 0, 1e-13, ERROR_INDIVIDUAL, &(val[0]), &(err[0]));
        return val[0]/(8*M_PI);
    }

    /**
    Do the calculations for one temperature

    @param order The order of the virial coefficient (2=B_2, 3=B_3, etc.)
    @param Tstar The temperature
    @param rstart The initial value of r to be considered in integration
    @param rend The final value of r to be considered in integration
    @param mol1 The first molecule
    @param mol2 The second molecule
    @returns Tuple of (value, estimated error in value)
    @note The return numerical type maybe be one of double, std::complex<double> or MultiComplex<double>
    */
    template <typename TEMPTYPE>
    std::tuple<TEMPTYPE,TEMPTYPE> one_temperature(int order, TEMPTYPE Tstar, TYPE rstart, TYPE rend, Molecule<TYPE> mol1, Molecule<TYPE> mol2) const
    {

        // Some local typedefs to avoid typing
        using SharedData = SharedDataBase<TYPE, TEMPTYPE>;

        std::vector<std::valarray<double>> xmins, xmaxs;
        switch (order)
        {
        case 2:
            xmins = {{ 0, 0, 0, rstart }}, xmaxs = {{ M_PI, M_PI, 2 * M_PI, rend }}; // Limits on theta1, theta2, phi, r
            break;
        case 3:{
            double rbreak = 1.3;
            xmins = {{ rstart, rstart, -1 },{ rbreak, rstart, -1 } }, xmaxs = {{ rbreak, rend, 1 },{ rend, rend, 1 } }; // Limits on r12, r13, eta
            break;
            }
        case 4: {
            // Limits on r14, r13, gamma, r12, eta
            // Limits on eta, r12, r13, eta, r14
            // Limits on eta,zeta,gamma, r12, r13, r14
            xmins = { {rstart, rstart, -1 , rstart, -1 }, { -1, rstart, rstart, -1 , rstart }, { -1,  rstart, -1 , rstart, rstart , rstart } };
            xmaxs = { { rend, rend, 1 , rend, 1 }, { 1, rend, rend, 1 , rend }, { 1, 2 * M_PI, 1, rend, rend , rend } };
            break;
        }
        default:
            break;
        }
        SharedData shared(Tstar, 0.0, mol1, mol2, potcls);

        int ndim = 1; // If T is a floating point number (default)
        std::valarray<double> outval, outerr;
        std::vector<std::valarray<double >> vals(10), errs(10);
        if constexpr (std::is_same<decltype(shared.Tstar), std::complex<double>>::value) {
            ndim = 2;
        }
        else if constexpr (std::is_same<decltype(shared.Tstar), MultiComplex<double>>::value) {
            ndim = static_cast<int>(shared.Tstar.get_coef().size());
        }
        // Fill with zero
        for (auto& val : vals) {
            val = std::valarray<double>(0.0, ndim);
        }
        for (auto& err : errs) {
            err = std::valarray<double>(0.0, ndim);
        }
        outval = std::valarray<double>(0.0, ndim);
        outerr = std::valarray<double>(0.0, ndim);

        int feval_max = 0;
        if (m_conf.contains("feval_max")) {
            feval_max = static_cast<int>(m_conf["feval_max"]);
        }
        else {
            throw std::invalid_argument("Key \"feval_max\" must be specified in the configuration JSON");
        }

#if !defined(NO_CUBA)
        auto Cuba_integrand = [](const int *pndim, const cubareal x[], const int *pncomp, cubareal fval[], void *p_shared_data) {
            auto& shared = *((class SharedDataBase<double, TEMPTYPE>*)(p_shared_data));

            double theta1, theta2, phi, r;
            double jacobian = 1.0;
            for (auto i  = 0; i < *pndim; ++i){
                auto range = shared.xmax[i] - shared.xmin[i];
                jacobian *= range;
            }
            theta1 = shared.xmin[0] + x[0]*(shared.xmax[0]-shared.xmin[0]);
            theta2 = shared.xmin[1] + x[1]*(shared.xmax[1]-shared.xmin[1]);
            phi =    shared.xmin[2] + x[2]*(shared.xmax[2]-shared.xmin[2]);
            r   =    shared.xmin[3] + x[3]*(shared.xmax[3]-shared.xmin[3]);
            shared.rstar = r;
            shared.orient_integrand(theta1, theta2, phi, fval);
            for (auto i = 0; i < *pncomp; ++i){
                fval[i] *= jacobian;
            }
            return 0; // success
        };

        int NVEC = 1;
        int EPSREL = 1e-8;
        int EPSABS = 1e-12;
        int VERBOSE = 0;
        int LAST = 4;
        int MINEVAL = 0;
        int MAXEVAL = feval_max;
        int NSTART = 1000;
        int NINCREASE = 500;
        int NBATCH = 1000;
        int GRIDNO = 0;
        const char *STATEFILE = nullptr;
        void *SPIN = nullptr;
        int neval, fail;
        cubareal integral[ndim], error[ndim], prob[ndim];

        int nregions;
        int KEY = 0;
        auto startTimeC = std::chrono::high_resolution_clock::now();
        Cuhre(4, ndim, Cuba_integrand, &shared, NVEC,
        EPSREL, EPSABS, VERBOSE | LAST,
        MINEVAL, MAXEVAL, KEY,
        STATEFILE, SPIN,
        &nregions, &neval, &fail, integral, error, prob);
        auto endTimeC = std::chrono::high_resolution_clock::now();
        auto timeC = std::chrono::duration<double>(endTimeC - startTimeC).count();

        // The quadruple integral needs to be divided by 8*pi, but the leading term in the
        // expression for B_2 is -2\pi, so factor becomes -1/4, or -0.25
        for (auto i = 0; i < outval.size(); ++i) {
            outval[i] = -0.25 * integral[i];
            outerr[i] = -0.25 * error[i];
        }
#else
        // Use cubature to do the integration...

        switch (order) {
        case 2:
        {
            // The integrand function
            //typedef int (*integrand) (unsigned ndim, const double *x, void *, unsigned fdim, double* fval);
            auto cubature_integrand = [](unsigned ndim, const double* x, void* p_shared_data, unsigned fdim, double* fval) {
                auto& shared = *((class SharedDataBase<double, TEMPTYPE>*)(p_shared_data));
                double theta1 = x[0], theta2 = x[1], phi = x[2], r = x[3];
                shared.rstar = r;
                shared.orient_integrand(theta1, theta2, phi, fval);
                return 0; // success
            };

            int naxes = 4; // How many dimensions the integral is taken over (theta, phi1, phi2, r)
            hcubature(ndim, cubature_integrand, &shared, naxes, &(xmins[0][0]), &(xmaxs[0][0]), feval_max, 0, 1e-13, ERROR_INDIVIDUAL, &(vals[0][0]), &(errs[0][0]));

            // Copy into output
            // ....
            // The quadruple integral needs to be divided by 8*pi, but the leading term in the
            // expression for B_2 is -2\pi, so factor becomes -1/4, or -0.25
            outval = -0.25*vals[0]; outerr = -0.25*errs[0];
            break;
        }
        case 3:
        {
            // The integrand function
            //typedef int (*integrand) (unsigned ndim, const double *x, void *, unsigned fdim, double* fval);
            auto cubature_integrand = [](unsigned ndim, const double* x, void* p_shared_data, unsigned fdim, double* fval) {
                auto& shared = *((class SharedDataBase<double, TEMPTYPE>*)(p_shared_data));
                shared.atomic_B3_integrand(x[0], x[1], x[2], fval);
                return 0; // success
            };

            int naxes = 3; // How many dimensions the integral is taken over (r12, r13, eta)
            for (auto i = 0; i < xmins.size(); ++i){

                hcubature(ndim, cubature_integrand, &shared, naxes, &(xmins[i][0]), &(xmaxs[i][0]), feval_max, 0, 1e-13, ERROR_INDIVIDUAL, &(vals[i][0]), &(errs[i][0]));

                // Copy into output
                outval += vals[i]; outerr += std::abs(errs[i]);
            }
            // Rescale with the leading factor
            outval *= 8*M_PI*M_PI/3; outerr *= 8*M_PI*M_PI/3;

            break;
        }
		case 4:
		{
			// Fourth virial coefficient: three parts B_4_1,B_4_2,B_4_3
			// The integrand functions
			//typedef int (*integrand) (unsigned ndim, const double *x, void *, unsigned fdim, double* fval);
			auto cubature_integrand_1 = [](unsigned ndim, const double* x, void* p_shared_data, unsigned fdim, double* fval) {
				auto& shared = *((class SharedDataBase<double, TEMPTYPE>*)(p_shared_data));
				shared.atomic_B4_1_integrand(x[0], x[1], x[2], x[3], x[4], fval);
				return 0; // success
			};

			auto cubature_integrand_2 = [](unsigned ndim, const double* x, void* p_shared_data, unsigned fdim, double* fval) {
				auto& shared = *((class SharedDataBase<double, TEMPTYPE>*)(p_shared_data));
				shared.atomic_B4_2_integrand(x[0], x[1], x[2], x[3], x[4], fval);
				return 0; // success
			};

			auto cubature_integrand_3 = [](unsigned ndim, const double* x, void* p_shared_data, unsigned fdim, double* fval) {
				auto& shared = *((class SharedDataBase<double, TEMPTYPE>*)(p_shared_data));
				shared.atomic_B4_3_integrand(x[0], x[1], x[2], x[3], x[4], x[5], fval);
				return 0; // success
			};

			// prefactors for each contribution
			std::valarray<double> pre_factors = {-3.0*(27.0/4.0), 3.0*(27.0/2.0), -27.0/(8.0*M_PI)};

			int naxes = 5; // How many dimensions the integral is taken over (r14, r13, gamma, r12, eta)
			hcubature(ndim, cubature_integrand_1, &shared, naxes, &(xmins[0][0]), &(xmaxs[0][0]), feval_max, 0, 1e-4, ERROR_INDIVIDUAL, &(vals[0][0]), &(errs[0][0]));
			hcubature(ndim, cubature_integrand_2, &shared, naxes, &(xmins[1][0]), &(xmaxs[1][0]), feval_max, 0, 1e-4, ERROR_INDIVIDUAL, &(vals[1][0]), &(errs[1][0]));

			naxes = 6;
			hcubature(ndim, cubature_integrand_3, &shared, naxes, &(xmins[2][0]), &(xmaxs[2][0]), feval_max, 0, 1e-4, ERROR_INDIVIDUAL, &(vals[2][0]), &(errs[2][0]));

			for (auto i = 0; i < pre_factors.size(); ++i) {
                vals[i] *= pre_factors[i];
                errs[i] *= pre_factors[i];
			}

            // Copy into output
			outval = vals[0] + vals[1] + vals[2];
			outerr = std::abs(errs[0]) + std::abs(errs[1]) + std::abs(errs[2]);

			break;
		}
        default:
            throw -1;
        }
#endif
        if constexpr (std::is_same<decltype(Tstar), double>::value) {
            // If T is double (real)
            return std::make_tuple(outval[0],outerr[0]);
        }
        else if constexpr (std::is_same<decltype(Tstar), std::complex<double>>::value) {
            // If T is a complex number (perhaps for complex step derivatives)
            return std::make_tuple(decltype(Tstar)(outval[0], outval[1]), decltype(Tstar)(outerr[0], outerr[1]));
        }
        else if constexpr (std::is_same<decltype(Tstar), MultiComplex<double>>::value) {
            // If T is a multicomplex number
            return std::make_tuple(decltype(Tstar)(outval), decltype(Tstar)(outerr));
        }
    }

    std::map<std::string, double> B_and_derivs(int order, int Nderivs, double T, double rstart, double rend, Molecule<TYPE> mol1, Molecule<TYPE> mol2){

        if (Nderivs == 0) {
            auto [val,esterr] = this->one_temperature(order, T, rstart, rend, mol1, mol2);
            return {
                {"T", T},
                {"B", val},
                {"error(B)", esterr}
            };
        }
        if (Nderivs == 1) {
            double h = 1e-100;
            auto [val,esterr] = this->one_temperature(order, std::complex<double>(T,h), rstart, rend, mol1, mol2);
            return {
                {"T", T},
                {"B", val.real()},
                {"error(B)", esterr.real()},
                {"dBdT", val.imag()/h},
                {"error(dBdT)", esterr.imag()/h},
            };
        }
        else {
            std::function<std::tuple<MultiComplex<double>,MultiComplex<double>>(const MultiComplex<double>&)> f(
                [this, order, rstart, rend, mol1, mol2](const MultiComplex<double>& T) {
                    return this->one_temperature(order, T, rstart, rend, mol1, mol2);
                });
            bool and_val = true;
            auto [val,esterr] = diff_mcx1(f, T, Nderivs, and_val);
            std::map<std::string, double> o = { {"T", T} };
            for (auto i = 0; i <= Nderivs; ++i) {
                switch (i) {
                case 0:
                    o["B"] = val[0];
                    o["error(B)"] = esterr[0];
                    break;
                case 1:
                    o["dBdT"] = val[1];
                    o["error(dBdT)"] = esterr[1];
                    break;
                default:
                    auto n = std::to_string(i);
                    o["d" + n + "BdT" + n] = val[i];
                    o["error(d" + n + "BdT" + n + ")"] = esterr[i];
                }
            }
            return o;
        }
    }
    auto parallel_B_and_derivs(int order, int Nthreads, int Nderivs, std::vector<double> Tvec, double rstart, double rend, Molecule<TYPE> mol1, Molecule<TYPE> mol2)
    {
        init_thread_pool(Nthreads);
        std::vector<double> times(Tvec.size());
        std::vector<std::map<std::string, double>> outputs(Tvec.size());
        std::size_t i = 0;
        for (auto T : Tvec) {
            auto& result = outputs[i];
            auto& time = times[i];
            std::function<void(void)> one_Temp = [this, order, Nderivs, T, rstart, rend, mol1, mol2, &result, &time]() {
                auto startTime = std::chrono::high_resolution_clock::now();
                result = this->B_and_derivs(order, Nderivs, T, rstart, rend, mol1, mol2);
                auto endTime = std::chrono::high_resolution_clock::now();
                time = std::chrono::duration<double>(endTime - startTime).count();
                {
                    mtx.lock();
                    std::cout << "Done " << T << " in " << time << " seconds\n" ;
                    mtx.unlock();
                }
                result["elapsed / s"] = time;
            };
            m_pool->AddJob(one_Temp);
            i++;
        }
        // Wait until all the threads finish...
        m_pool->WaitAll();
        return outputs;
    }
    TYPE orientation_averaged_integrate(const TYPE Tstar, const TYPE rstart, const TYPE rend) {
        using arr = Eigen::Array<TYPE, Eigen::Dynamic, 1>;
        bool parallel = false;
        if (parallel){
            // Parallel
            TYPE result, time;
            Molecule<TYPE> mol1 = this->mol1, mol2 = this->mol2;
            auto one_temperature_job = [this, Tstar, rstart, rend, mol1, mol2, &result, &time]() {
                auto startTime = std::chrono::high_resolution_clock::now();
                result = one_temperature<double>(Tstar, rstart, rend, mol1, mol2);
                auto endTime = std::chrono::high_resolution_clock::now();
                time = std::chrono::duration<double>(endTime - startTime).count();
            };
            m_pool->AddJob(one_temperature_job);
            // Wait until all the threads finish...
            m_pool->WaitAll();
            return result;
        }
        else {
            // Serial
            return one_temperature(Tstar, rstart, rend, mol1, mol2);
        }
    }
};


// math functions

//local headers
#include "math.h"

//third party headers
#include <boost/math/special_functions/binomial.hpp>

//standard headers
#include <cmath>
#include <cstdint>

std::size_t bin_coeff_get_max_k(std::size_t size)
{
    // what is the largest value of (n/2) such that
    // n choose (n/2) will fit in an integral T{} without overflow?

    // note: results found by incrementing 'n' until failure
    switch (size)
    {
        case 1:
            return 10/2;
        case 2:
            return 18/2;
        case 4:
            return 34/2;
        case 8:
            return 67/2;
        case 16:
            return 131/2;
        case 32:
            return 260/2;
        case 64:
            return 516/2;
        case 128:
            return 1029/2;
        break;
    };

    // unknown size, approximate the answer
    return (10 * size);
}

std::vector<std::uint16_t> get_primes_up_to(const std::uint16_t n)
{
    // Eratosthenes sieve
    // get all primes in range [2, n]
    // '1' and '0' not considered primes
    std::vector<bool> sieve;
    std::vector<std::uint16_t> result;
    sieve.resize(n + 1, false);
    result.reserve(n/2);

    for (std::size_t i{2}; i <= n; ++i)
    {
        // If false, then it's a prime.
        if (sieve[i] == false)
        {
            // Set all multiples of 'i' to true (i.e. non-prime)
            if (i * i <= n)
            {
                for (std::size_t m{i * i}; m <= n; m += i)
                    sieve[m] = true;
            }

            // save result
            result.push_back(i);
        }
    }

    return result;
}

std::int32_t n_choose_k_bwrap(const std::uint32_t n, const std::uint32_t k)
{
    if (k > n)
        return 0;

    // note: result is 'double::max()' aka 'infinity' if any error occurs
    using namespace boost::math::policies;
    double fp_result = boost::math::binomial_coefficient<double>(n, k,
      make_policy(overflow_error<ignore_error>(),
        evaluation_error<ignore_error>(),
        domain_error<ignore_error>(),
        pole_error<ignore_error>()));

    if (fp_result < 0 || fp_result == std::numeric_limits<double>::infinity())
        return 0;

    if (fp_result > std::numeric_limits<std::int32_t>::max())
        return 0;

    return std::round(fp_result);
}


// Boost.Geometry

// Copyright (c) 2015 Oracle and/or its affiliates.

// Contributed and/or modified by Adam Wulkiewicz, on behalf of Oracle

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

#ifndef BOOST_GEOMETRY_ALGORITHMS_DETAIL_THOMAS_INVERSE_HPP
#define BOOST_GEOMETRY_ALGORITHMS_DETAIL_THOMAS_INVERSE_HPP

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

#include <boost/geometry/core/radius.hpp>
#include <boost/geometry/core/srs.hpp>

#include <boost/geometry/util/math.hpp>

#include <boost/geometry/algorithms/detail/flattening.hpp>

namespace boost { namespace geometry { namespace detail
{

/*!
\brief The solution of the inverse problem of geodesics on latlong coordinates,
       Forsyth-Andoyer-Lambert type approximation with second order terms.
\author See
    - Technical Report: PAUL D. THOMAS, MATHEMATICAL MODELS FOR NAVIGATION SYSTEMS, 1965
      http://www.dtic.mil/docs/citations/AD0627893
    - Technical Report: PAUL D. THOMAS, SPHEROIDAL GEODESICS, REFERENCE SYSTEMS, AND LOCAL GEOMETRY, 1970
      http://www.dtic.mil/docs/citations/AD703541
*/
template <typename CT>
class thomas_inverse
{
public:
    template <typename T1, typename T2, typename Spheroid>
    thomas_inverse(T1 const& lon1,
                   T1 const& lat1,
                   T2 const& lon2,
                   T2 const& lat2,
                   Spheroid const& spheroid)
        : m_a(get_radius<0>(spheroid))
        , m_f(detail::flattening<CT>(spheroid))
        , m_is_result_zero(false)
    {
        // coordinates in radians

        if ( math::equals(lon1, lon2)
          && math::equals(lat1, lat2) )
        {
            m_is_result_zero = true;
            return;
        }

        CT const one_minus_f = CT(1) - m_f;

//        CT const tan_theta1 = one_minus_f * tan(lat1);
//        CT const tan_theta2 = one_minus_f * tan(lat2);
//        CT const theta1 = atan(tan_theta1);
//        CT const theta2 = atan(tan_theta2);

        CT const pi_half = math::pi<CT>() / CT(2);
        CT const theta1 = math::equals(lat1, pi_half) ? lat1 :
                          math::equals(lat1, -pi_half) ? lat1 :
                          atan(one_minus_f * tan(lat1));
        CT const theta2 = math::equals(lat2, pi_half) ? lat2 :
                          math::equals(lat2, -pi_half) ? lat2 :
                          atan(one_minus_f * tan(lat2));

        CT const theta_m = (theta1 + theta2) / CT(2);
        CT const d_theta_m = (theta2 - theta1) / CT(2);
        m_d_lambda = lon2 - lon1;
        CT const d_lambda_m = m_d_lambda / CT(2);

        m_sin_theta_m = sin(theta_m);
        m_cos_theta_m = cos(theta_m);
        m_sin_d_theta_m = sin(d_theta_m);
        m_cos_d_theta_m = cos(d_theta_m);
        CT const sin2_theta_m = math::sqr(m_sin_theta_m);
        CT const cos2_theta_m = math::sqr(m_cos_theta_m);
        CT const sin2_d_theta_m = math::sqr(m_sin_d_theta_m);
        CT const cos2_d_theta_m = math::sqr(m_cos_d_theta_m);
        CT const sin_d_lambda_m = sin(d_lambda_m);
        CT const sin2_d_lambda_m = math::sqr(sin_d_lambda_m);

        CT const H = cos2_theta_m - sin2_d_theta_m;
        CT const L = sin2_d_theta_m + H * sin2_d_lambda_m;
        m_cos_d = CT(1) - CT(2) * L;
        CT const d = acos(m_cos_d);
        m_sin_d = sin(d);

        CT const one_minus_L = CT(1) - L;

        if ( math::equals(m_sin_d, CT(0))
          || math::equals(L, CT(0))
          || math::equals(one_minus_L, CT(0)) )
        {
            m_is_result_zero = true;
            return;
        }

        CT const U = CT(2) * sin2_theta_m * cos2_d_theta_m / one_minus_L;
        CT const V = CT(2) * sin2_d_theta_m * cos2_theta_m / L;
        m_X = U + V;
        m_Y = U - V;
        m_T = d / m_sin_d;
        //CT const D = CT(4) * math::sqr(T);
        //CT const E = CT(2) * cos_d;
        //CT const A = D * E;
        //CT const B = CT(2) * D;
        //CT const C = T - (A - E) / CT(2);
    }

    inline CT distance() const
    {
        if ( m_is_result_zero )
        {
            // TODO return some approximated value
            return CT(0);
        }

        //CT const n1 = X * (A + C*X);
        //CT const n2 = Y * (B + E*Y);
        //CT const n3 = D*X*Y;

        //CT const f_sqr = math::sqr(f);
        //CT const f_sqr_per_64 = f_sqr / CT(64);

        CT const delta1d = m_f * (m_T*m_X-m_Y) / CT(4);
        //CT const delta2d = f_sqr_per_64 * (n1 - n2 + n3);

        return m_a * m_sin_d * (m_T - delta1d);
        //double S2 = a * sin_d * (T - delta1d + delta2d);
    }

    inline CT azimuth() const
    {
        // NOTE: if both cos_latX == 0 then below we'd have 0 * INF
        // it's a situation when the endpoints are on the poles +-90 deg
        // in this case the azimuth could either be 0 or +-pi
        if ( m_is_result_zero )
        {
            return CT(0);
        }

        // may also be used to calculate distance21
        //CT const D = CT(4) * math::sqr(T);
        CT const E = CT(2) * m_cos_d;
        //CT const A = D * E;
        //CT const B = CT(2) * D;
        // may also be used to calculate distance21
        CT const f_sqr = math::sqr(m_f);
        CT const f_sqr_per_64 = f_sqr / CT(64);

        CT const F = CT(2)*m_Y-E*(CT(4)-m_X);
        //CT const M = CT(32)*T-(CT(20)*T-A)*X-(B+CT(4))*Y;
        CT const G = m_f*m_T/CT(2) + f_sqr_per_64;
        CT const tan_d_lambda = tan(m_d_lambda);
        CT const Q = -(F*G*tan_d_lambda) / CT(4);

        CT const d_lambda_p = (m_d_lambda + Q) / CT(2);
        CT const tan_d_lambda_p = tan(d_lambda_p);

        CT const v = atan2(m_cos_d_theta_m, m_sin_theta_m * tan_d_lambda_p);
        CT const u = atan2(-m_sin_d_theta_m, m_cos_theta_m * tan_d_lambda_p);

        CT const pi = math::pi<CT>();
        CT alpha1 = v + u;
        if ( alpha1 > pi )
        {
            alpha1 -= CT(2) * pi;
        }

        return alpha1;
    }

private:
    CT const m_a;
    CT const m_f;

    CT m_d_lambda;
    CT m_cos_d;
    CT m_sin_d;
    CT m_X;
    CT m_Y;
    CT m_T;
    CT m_sin_theta_m;
    CT m_cos_theta_m;
    CT m_sin_d_theta_m;
    CT m_cos_d_theta_m;

    bool m_is_result_zero;
};

}}} // namespace boost::geometry::detail

#endif // BOOST_GEOMETRY_ALGORITHMS_DETAIL_THOMAS_INVERSE_HPP


//<%
//cfg['compiler_args'] = ['-std=c++11']
//cfg['include_dirs'] = ['../eigen']
//setup_pybind11(cfg)
//%>

#include <pybind11/pybind11.h>
#include <pybind11/eigen.h>
#include <Eigen/Dense>
#include <iostream>
#include <cmath>
#include <vector>
#include <numeric>

constexpr auto pi = 3.141592653589793;

double log_marginal_probability_cpp(Eigen::VectorXd m, Eigen::MatrixXd S,
    double r, double v, Eigen::MatrixXd X) {

    int N = X.rows();
    int p = X.cols();
    Eigen::VectorXd xsum = X.colwise().sum().transpose();

    Eigen::MatrixXd Sprime = S + X.transpose() * X + r * N / (N + r) * (m * m.transpose()) - 1 / (N + r)
    * (xsum * xsum.transpose()) - r / (N + r) * (m * xsum.transpose() + xsum * m.transpose());

    double vprime = v + N;

    std::vector<double> gamma1(p, 0);
    std::vector<double> gamma2(p, 0);
    for (int i = 0; i < p; ++i) {
        gamma1[i] = lgamma((v - i) / 2);
        gamma2[i] = lgamma((vprime - i) / 2);
    }
    double gamma3 = std::accumulate(gamma1.begin(), gamma1.end(), 0.0);
	double gamma4 = std::accumulate(gamma2.begin(), gamma2.end(), 0.0);

	double log_prob = -N * p / 2 * log(2 * pi) + p / 2 * log(r / (N + r)) + v / 2 * log(S.determinant()) - vprime / 2
		* log(Sprime.determinant()) + vprime * p / 2 * log(2) + gamma4 - v * p / 2 * log(2) - gamma3;

	return log_prob;
}

PYBIND11_PLUGIN(log_marginal_prob) {
    pybind11::module m("log_marginal_prob", "auto-compiled c++ extension");
    m.def("log_marginal_probability_cpp", &log_marginal_probability_cpp);
    return m.ptr();
}


#include "FKDerivative.h"
#include <Eigen/Dense>
#include <totalmodel.h>
#include <cmath>
#include <cassert>
#include <chrono>

using namespace Eigen;

void AngleAxisToRotationMatrix_Derivative(const double* pose, double* dR_data, const int idj, const int numberColumns)
{
    Eigen::Map< Eigen::Matrix<double, 9, Eigen::Dynamic, Eigen::RowMajor> > dR(dR_data, 9, numberColumns);
    std::fill(dR_data, dR_data + 9 * numberColumns, 0.0);
    const double theta2 = pose[0] * pose[0] + pose[1] * pose[1] + pose[2] * pose[2];
    if (theta2 > std::numeric_limits<double>::epsilon())
    {
        const double theta = sqrt(theta2);
        const double s = sin(theta);
        const double c = cos(theta);
        const Eigen::Map< const Eigen::Matrix<double, 3, 1> > u(pose);
        Eigen::VectorXd e(3);
        e[0] = pose[0] / theta; e[1] = pose[1] / theta; e[2] = pose[2] / theta;

        // dR / dtheta
        Eigen::Matrix<double, 9, 1> dRdth(9, 1);
        Eigen::Map< Eigen::Matrix<double, 3, 3, Eigen::RowMajor> > dRdth_(dRdth.data());
        // skew symmetric
        dRdth_ << 0.0, -e[2], e[1],
                  e[2], 0.0, -e[0],
                  -e[1], e[0], 0.0;
        // dRdth_ = dRdth_ * c - Matrix<double, 3, 3>::Identity() * s + s * e * e.transpose();
        dRdth_ = - dRdth_ * c - Matrix<double, 3, 3>::Identity() * s + s * e * e.transpose();

        // dR / de
        Eigen::Matrix<double, 9, 3, RowMajor> dRde(9, 3);
        // d(ee^T) / de
        dRde <<
            2 * e[0], 0., 0.,
            e[1], e[0], 0.,
            e[2], 0., e[0],
            e[1], e[0], 0.,
            0., 2 * e[1], 0.,
            0., e[2], e[1],
            e[2], 0., e[0],
            0., e[2], e[1],
            0., 0., 2 * e[2];
        Eigen::Matrix<double, 9, 3, RowMajor> dexde(9, 3);
        dexde <<
            0, 0, 0,
            0, 0, -1,
            0, 1, 0,
            0, 0, 1,
            0, 0, 0,
            -1, 0, 0,
            0, -1, 0,
            1, 0, 0,
            0, 0, 0;
        // dRde = dRde * (1. - c) + c * dexde;
        dRde = dRde * (1. - c) - s * dexde;
        Eigen::Matrix<double, Dynamic, Dynamic, RowMajor> dedu = Matrix<double, 3, 3>::Identity() / theta - u * u.transpose() / theta2 / theta;

        dR.block(0, 3 * idj, 9, 3) = dRdth * e.transpose() + dRde * dedu;
    }
    else
    {
        dR(1, 3 * idj + 2) = 1;
        dR(2, 3 * idj + 1) = -1;
        dR(3, 3 * idj + 2) = -1;
        dR(5, 3 * idj) = 1;
        dR(6, 3 * idj + 1) = 1;
        dR(7, 3 * idj) = -1;
    }
}

void EulerAnglesToRotationMatrix_Derivative(const double* pose, double* dR_data, const int idj, const int numberColumns)
{
    const double degrees_to_radians = 3.14159265358979323846 / 180.0;
    const double pitch(pose[0] * degrees_to_radians);
    const double roll(pose[1] * degrees_to_radians);
    const double yaw(pose[2] * degrees_to_radians);

    const double c1 = cos(yaw);
    const double s1 = sin(yaw);
    const double c2 = cos(roll);
    const double s2 = sin(roll);
    const double c3 = cos(pitch);
    const double s3 = sin(pitch);

    Eigen::Matrix<double, 9, 3, Eigen::RowMajor> dRdp(9, 3);
    // dRdp <<
    //  -s1 * c2, -s2 * c1, 0.,
    //  -c1 * c3 - s1 * s2 * s3, c1 * c2 * s3, s1 * s3 + c1 * s2 * c3,
    //  c1 * s3 - s1 * s2 * c3, c1 * c2 * c3, s1 * c3 - c1 * s2 * s3,
    //  c1 * c2, -s1 * s2, 0.,
    //  -s1 * c3 + c1 * s2 * s3, s1 * c2 * s3, -c1 * s3 + s1 * s2 * c3,
    //  s1 * s3 + c1 * s2 * c3, s1 * c2 * c3, -c1 * c3 - s1 * s2 * s3,
    //  0, -c2, 0,
    //  0, -s2 * s3, c2 * c3,
    //  0, -s2 * c3, -c2 * s3;

    dRdp <<
        0.,                       -s2 * c1,      -s1 * c2,
        s1 * s3 + c1 * s2 * c3,    c1 * c2 * s3, -c1 * c3 - s1 * s2 * s3,
        s1 * c3 - c1 * s2 * s3,    c1 * c2 * c3, c1 * s3 - s1 * s2 * c3,
        0.,                       -s1 * s2,      c1 * c2,
        -c1 * s3 + s1 * s2 * c3,   s1 * c2 * s3, -s1 * c3 + c1 * s2 * s3,
        -c1 * c3 - s1 * s2 * s3,   s1 * c2 * c3, s1 * s3 + c1 * s2 * c3,
        0,                        -c2,           0,
        c2 * c3,                  -s2 * s3,      0,
        -c2 * s3,                 -s2 * c3,      0;

    Eigen::Map< Eigen::Matrix<double, 9, Eigen::Dynamic, Eigen::RowMajor> > dR(dR_data, 9, numberColumns);
    std::fill(dR_data, dR_data + 9 * numberColumns, 0.0);
    dR.block(0, 3 * idj, 9, 3) = degrees_to_radians * dRdp;
}

void Product_Derivative(const double* const A_data, const double* const dA_data, const double* const B_data,
                        const double* const dB_data, double* dAB_data, const int B_col)
{
    assert(dA_data != NULL || dB_data != NULL);
    assert(B_col == 3 || B_col == 1);  // matrix multiplication or matrix-vector multiplication
    if (dA_data != NULL && dB_data != NULL)
    {
        const Eigen::Map<const Eigen::Matrix<double, 9, TotalModel::NUM_JOINTS * 3, Eigen::RowMajor> > dA(dA_data);
        if (B_col == 1)
        {
            // B_col == 1
            // d(AB) = AdB + (dA)B
            const Eigen::Map<const Eigen::Matrix<double, 3, 3, Eigen::RowMajor> > A(A_data);
            const Eigen::Map<const Eigen::Matrix<double, 3, TotalModel::NUM_JOINTS * 3, Eigen::RowMajor> > dB(dB_data);
            Eigen::Map< Eigen::Matrix<double, 3, TotalModel::NUM_JOINTS * 3, Eigen::RowMajor> > dAB(dAB_data);
            for (int r = 0; r < 3; r++)
            {
                const int baseIndex = 3*r;
                dAB.row(r) = A(r, 0) * dB.row(0) + A(r, 1) * dB.row(1) + A(r, 2) * dB.row(2) +
                    B_data[0] * dA.row(baseIndex) + B_data[1] * dA.row(baseIndex + 1) + B_data[2] * dA.row(baseIndex + 2);
            }
        }
        else
        {
            // B_col == 3
            // d(AB) = AdB + (dA)B
            const Eigen::Map<const Eigen::Matrix<double, 9, TotalModel::NUM_JOINTS * 3, Eigen::RowMajor> > dB(dB_data);
            Eigen::Map< Eigen::Matrix<double, 9, TotalModel::NUM_JOINTS * 3, Eigen::RowMajor> > dAB(dAB_data);
            for (int r = 0; r < 3; r++)
            {
                const int baseIndex = 3*r;
                for (int c = 0; c < 3; c++)
                {
                    dAB.row(baseIndex + c) = A_data[baseIndex] * dB.row(c) + A_data[baseIndex+1] * dB.row(3 + c) + A_data[baseIndex+2] * dB.row(6 + c) +
                        B_data[c] * dA.row(baseIndex) + B_data[3 + c] * dA.row(baseIndex + 1) + B_data[6 + c] * dA.row(baseIndex + 2);
                }
            }
        }
    }
    else if (dA_data != NULL && dB_data == NULL)  // B is a constant matrix / vector, no derivative
    {
        const Eigen::Map<const Eigen::Matrix<double, 9, TotalModel::NUM_JOINTS * 3, Eigen::RowMajor> > dA(dA_data);
        if (B_col == 1)
        {
            // d(AB) = AdB + (dA)B
            Eigen::Map< Eigen::Matrix<double, 3, TotalModel::NUM_JOINTS * 3, Eigen::RowMajor> > dAB(dAB_data);
            // // Matrix form (slower)
            // for (int r = 0; r < 3; r++)
            //     dABAux.row(r) = B * dA.block<3, TotalModel::NUM_JOINTS * 3>(r, 0);
            // For loop form
            for (int r = 0; r < 3; r++)
            {
                const int baseIndex = 3*r;
                dAB.row(r) = B_data[0] * dA.row(baseIndex) + B_data[1] * dA.row(baseIndex + 1) + B_data[2] * dA.row(baseIndex + 2);
            }
        }
        else
        {
            // B_col == 3
            Eigen::Map< Eigen::Matrix<double, 9, TotalModel::NUM_JOINTS * 3, Eigen::RowMajor> > dAB(dAB_data);
            for (int r = 0; r < 3; r++)
                for (int c = 0; c < 3; c++)
                    dAB.row(3 * r + c) = B_data[c] * dA.row(3 * r) + B_data[3 + c] * dA.row(3 * r + 1) + B_data[6 + c] * dA.row(3 * r + 2);  // d(AB) = AdB + (dA)B
        }
    }
    else // A is a constant matrix, no derivative
    {
        const Eigen::Map<const Eigen::Matrix<double, 3, 3, Eigen::RowMajor> > A(A_data);
        // dA_data == NULL && dB_data != NULL
        if (B_col == 1)
        {
            const Eigen::Map<const Eigen::Matrix<double, 3, TotalModel::NUM_JOINTS * 3, Eigen::RowMajor> > dB(dB_data);
            Eigen::Map< Eigen::Matrix<double, 3, TotalModel::NUM_JOINTS * 3, Eigen::RowMajor> > dAB(dAB_data);
            dAB.setZero();
            for (int r = 0; r < 3; r++)
                dAB.row(r) = A(r, 0) * dB.row(0) + A(r, 1) * dB.row(1) + A(r, 2) * dB.row(2);
        }
        else
        {
            // B_col == 3
            const Eigen::Map<const Eigen::Matrix<double, 9, TotalModel::NUM_JOINTS * 3, Eigen::RowMajor> > dB(dB_data);
            Eigen::Map< Eigen::Matrix<double, 9, TotalModel::NUM_JOINTS * 3, Eigen::RowMajor> > dAB(dAB_data);
            dAB.setZero();
            for (int r = 0; r < 3; r++)
                for (int c = 0; c < 3; c++)
                    dAB.row(3 * r + c) = A(r, 0) * dB.row(0 + c) + A(r, 1) * dB.row(3 + c) + A(r, 2) * dB.row(6 + c);
        }
    }
}

void SparseProductDerivative(const double* const A_data, const double* const dA_data, const double* const B_data,
                             const double* const dB_data, const int colIndex, const std::vector<int>& parentIndexes, double* dAB_data, const int numberColumns)
{
    // d(AB) = AdB + (dA)B
    Eigen::Map< Eigen::Matrix<double, 9, Eigen::Dynamic, Eigen::RowMajor> > dAB(dAB_data, 9, numberColumns);

    std::fill(dAB_data, dAB_data + 9 * numberColumns, 0.0);
    // // Dense dAB (sparse dB) version
    // const Eigen::Map<const Eigen::Matrix<double, 9, numberColumns, Eigen::RowMajor> > dA(dA_data);
    // const Eigen::Map<const Eigen::Matrix<double, 9, numberColumns, Eigen::RowMajor> > dB(dB_data);
    // dAB.row(baseIndex + c) = B_data[c] * dA.row(baseIndex) + B_data[3 + c] * dA.row(baseIndex + 1) + B_data[6 + c] * dA.row(baseIndex + 2);
    // dAB.block<1,3>(baseIndex + c, 3*colIndex) += A_data[baseIndex] * dB.block<1,3>(c, 3*colIndex)
    //                                            + A_data[baseIndex+1] * dB.block<1,3>(3+c, 3*colIndex)
    //                                            + A_data[baseIndex+2] * dB.block<1,3>(6+c, 3*colIndex);
    // Sparse sped up equivalent
    const auto colOffset = 3*colIndex;
    for (int r = 0; r < 3; r++)
    {
        const int baseIndex = 3*r;
        for (int c = 0; c < 3; c++)
        {
            // AdB
            for (int subIndex = 0; subIndex < 3; subIndex++)
            {
                const auto finalOffset = colOffset + subIndex;
                dAB_data[numberColumns*(baseIndex + c) + finalOffset] +=
                    A_data[baseIndex] * dB_data[numberColumns*c + finalOffset]
                    + A_data[baseIndex+1] * dB_data[numberColumns*(3+c) + finalOffset]
                    + A_data[baseIndex+2] * dB_data[numberColumns*(6+c) + finalOffset];
            }
            // // AdB - Slower equivalent
            // dAB.block<1,3>(baseIndex + c, colOffset) += A_data[baseIndex] * dB.block<1,3>(c, colOffset)
            //                                            + A_data[baseIndex+1] * dB.block<1,3>(3+c, colOffset)
            //                                            + A_data[baseIndex+2] * dB.block<1,3>(6+c, colOffset);
            // (dA)B
            for (const auto& parentIndex : parentIndexes)
            {
                const auto parentOffset = 3*parentIndex;
                for (int subIndex = 0; subIndex < 3; subIndex++)
                {
                    const auto finalOffset = parentOffset + subIndex;
                    dAB_data[numberColumns*(baseIndex + c) + finalOffset] +=
                        B_data[c] * dA_data[numberColumns*baseIndex + finalOffset]
                        + B_data[3 + c] * dA_data[numberColumns*(baseIndex+1) + finalOffset]
                        + B_data[6 + c] * dA_data[numberColumns*(baseIndex+2) + finalOffset];
                }
            }
            // // (dA)B - Slower equivalent
            // for (const auto& parentIndex : parentIndexes)
            // {
            //     const auto parentOffset = 3*parentIndex;
            //     dAB.block<1,3>(baseIndex + c, parentOffset) += B_data[c] * dA.block<1,3>(baseIndex, parentOffset)
            //                                                  + B_data[3 + c] * dA.block<1,3>(baseIndex+1, parentOffset)
            //                                                  + B_data[6 + c] * dA.block<1,3>(baseIndex+2, parentOffset);
            // }
        }
    }
}

void SparseProductDerivative(const double* const dA_data, const double* const B_data,
                             const std::vector<int>& parentIndexes, double* dAB_data, const int numberColumns)
{
    // d(AB) = AdB + (dA)B
    // Sparse for loop form
    std::fill(dAB_data, dAB_data + 3 * numberColumns, 0.0);
    for (int r = 0; r < 3; r++)
    {
        const int baseIndex = 3*r;
        for (const auto& parentIndex : parentIndexes)
        {
            const auto parentOffset = 3*parentIndex;
            for (int subIndex = 0; subIndex < 3; subIndex++)
            {
                const auto finalOffset = parentOffset + subIndex;
                dAB_data[numberColumns*r + finalOffset] +=
                    B_data[0] * dA_data[numberColumns*baseIndex + finalOffset]
                    + B_data[1] * dA_data[numberColumns*(baseIndex+1) + finalOffset]
                    + B_data[2] * dA_data[numberColumns*(baseIndex+2) + finalOffset];
            }
        }
    }
    // // Dense Matrix form (slower)
    // Eigen::Map< Eigen::Matrix<double, 3, numberColumns, Eigen::RowMajor> > dAB(dAB_data);
    // const Eigen::Map<const Eigen::Matrix<double, 9, numberColumns, Eigen::RowMajor> > dA(dA_data);
    // for (int r = 0; r < 3; r++)
    //     dABAux.row(r) = B * dA.block<3, numberColumns>(r, 0);
    // // Dense for loop form
    // for (int r = 0; r < 3; r++)
    // {
    //     const int baseIndex = 3*r;
    //     dAB.row(r) = B_data[0] * dA.row(baseIndex) + B_data[1] * dA.row(baseIndex + 1) + B_data[2] * dA.row(baseIndex + 2);
    // }
}

void SparseProductDerivativeConstA(const double* const A_data, const double* const dB_data,
                             const std::vector<int>& parentIndexes, double* dAB_data, const int numberColumns)
{
	// d(AB) = AdB (A is a constant.)
    // Sparse for loop form
    std::fill(dAB_data, dAB_data + 3 * numberColumns, 0.0);
    for (int r = 0; r < 3; r++)
    {
		for (const auto& parentIndex : parentIndexes)
	    {
	    	const auto parentOffset = 3*parentIndex;
	    	for (int subIndex = 0; subIndex < 3; subIndex++)
            {
            	const auto finalOffset = parentOffset + subIndex;
            	dAB_data[numberColumns * r + finalOffset] = A_data[3 * r + 0] * dB_data[finalOffset] + A_data[3 * r + 1] * dB_data[finalOffset + numberColumns] +
            		A_data[3 * r + 2] * dB_data[finalOffset + numberColumns + numberColumns];
            }
		}
	}
}

void SparseAdd(const double* const B_data, const std::vector<int>& parentIndexes, double* A_data, const int numberColumns)
{
    // d(AB) += d(AB)_parent
    Eigen::Map< Eigen::Matrix<double, 3, Eigen::Dynamic, Eigen::RowMajor>> A(A_data, 3, numberColumns);
    const Eigen::Map<const Eigen::Matrix<double, 3, Eigen::Dynamic, Eigen::RowMajor>> B(B_data, 3, numberColumns);
    // Sparse for loop
    for (int r = 0; r < 3; r++)
    {
        for (const auto& parentIndex : parentIndexes)
        {
            const auto parentOffset = 3*parentIndex;
            for (int subIndex = 0; subIndex < 3; subIndex++)
            {
                const auto finalOffset = parentOffset + subIndex;
                A_data[numberColumns*r + finalOffset] += B_data[numberColumns*r + finalOffset];
            }
        }
    }
    // // Dense equivalent
    // dMtdPIdj += dJdP.block<3, numberColumns>(3 * ipar, 0);
    // A += B;
}

void SparseSubtract(const double* const B_data, const std::vector<int>& parentIndexes, double* A_data, const int numberColumns)
{
    // d(AB) += d(AB)_parent
    Eigen::Map< Eigen::Matrix<double, 3, Eigen::Dynamic, Eigen::RowMajor>> A(A_data, 3, numberColumns);
    const Eigen::Map<const Eigen::Matrix<double, 3, Eigen::Dynamic, Eigen::RowMajor>> B(B_data, 3, numberColumns);
    // Sparse for loop
    for (int r = 0; r < 3; r++)
    {
        for (const auto& parentIndex : parentIndexes)
        {
            const auto parentOffset = 3*parentIndex;
            for (int subIndex = 0; subIndex < 3; subIndex++)
            {
                const auto finalOffset = parentOffset + subIndex;
                A_data[numberColumns*r + finalOffset] -= B_data[numberColumns*r + finalOffset];
            }
        }
    }
    // // Dense equivalent
    // dMtdPIdj -= dJdP.block<3, numberColumns>(3 * ipar, 0);
    // A -= B;
}

void projection_Derivative(double* dPdI_data, const double* dJdI_data, const int ncol, double* XYZ, const double* pK_, int offsetP, int offsetJ, float weight)
{
	// Dx/Dt = dx/dX * dX/dt + dx/dY * dY/dt + dx/dZ * dZ/dt
	const double X = XYZ[0], Y = XYZ[1], Z = XYZ[2];
	double* P_row0 = dPdI_data + offsetP * ncol;
	double* P_row1 = dPdI_data + (offsetP + 1) * ncol;
	const double* J_row0 = dJdI_data + offsetJ * ncol;
	const double* J_row1 = dJdI_data + (offsetJ + 1) * ncol;
	const double* J_row2 = dJdI_data + (offsetJ + 2) * ncol;
	for (int i = 0; i < ncol; i++)
		P_row0[i] = weight * ( pK_[0] * J_row0[i] + pK_[1] * J_row1[i] - (pK_[0] * X + pK_[1] * Y) * J_row2[i] / Z ) / Z;
	for (int i = 0; i < ncol; i++)
		P_row1[i] = weight * pK_[4] * ( J_row1[i] - Y / Z * J_row2[i] ) / Z;
	// equivalent to
	// dPdI.row(offsetP + 0) = weight * (pK_[0] / Z * dJdI.row(offsetJ + 0)
	// 					 + pK_[1] / Z * dJdI.row(offsetJ + 1)
	// 					 - (pK_[0] * X + pK_[1] * Y) / Z / Z * dJdI.row(offsetJ + 2));
	// dPdI.row(offsetP + 1) = weight * (pK_[4] / Z * dJdI.row(offsetJ + 1)
	// 					 - pK_[4] * Y / Z / Z * dJdI.row(offsetJ + 2));
}


// ===========================================================================
// Imagine++ Libraries
// Copyright (C) Imagine
// For detailed information: http://imagine.enpc.fr/software
// ===========================================================================

#include "Imagine/LinAlg/MyEigen.h"
#include <Eigen/Dense>

using namespace Eigen;

//Existe dans Eigen Blas
//Experimental et pas de decomposition uplow='U' dans eigen et seulement pour les matrices def pos
//info important car utilise dans Matrix.h

// Cholesky decomposition (replaces xpotf2_)
template <typename T>
void subCholeskyDecomposition(const char* job, int *n, T* a, int *info)
{
    Map<Matrix<T, Dynamic, Dynamic> > AMap (a,*n,*n);
    Eigen::LLT<Matrix<T, Dynamic, Dynamic> > llt (AMap);
    AMap = (*job == 'L') ? llt.matrixLLT() : llt.matrixU();
    *info = int(llt.info());
};

void choleskyDecomposition(const char* job, int *n, double* a, int *info) { subCholeskyDecomposition<double>(job, n, a, info); }
void choleskyDecomposition(const char* job, int *n, float* a, int *info) { subCholeskyDecomposition<float>(job, n, a, info); }

// Norms
template <typename T>
T subMatrixNorm(char* norm, int *m, int *n, T *a, int *lda)
{
    Map<Matrix<T, Dynamic, Dynamic>, 0, OuterStride<> > AMap (a,*m,*n,OuterStride<>(*lda));

    T to_return = 0;
    switch (*norm)
    {
        case 'M':
            to_return = AMap.template lpNorm<Infinity>();
            break;
        case ('1'):
            to_return = (AMap.colwise().sum()).maxCoeff();
            break;
        case ('I'):
            to_return = (AMap.rowwise().sum()).maxCoeff();
            break;
        case ('F'):
            to_return = AMap.norm();
            break;
        default: assert(false);
    }

    return to_return;
}

double matrixNorm(char* norm, int *m, int *n, double *a, int *lda) { return subMatrixNorm<double>(norm, m, n, a, lda); }
float matrixNorm(char* norm, int *m, int *n, float *a, int *lda) { return subMatrixNorm<float>(norm, m, n, a, lda); }

// System solving by LU decomposition (replaces xgetrs_)
template <typename T>
void subLUSystemSolve(int *n,T *a,int *ipiv, T *b,int *info)
{
  Map<Matrix<T, Dynamic, Dynamic>, 0, OuterStride<> > AMap (a,*n,*n,OuterStride<>(*n));
  Map<Matrix<T, Dynamic, 1>, 0, InnerStride<1> > BMap (b,*n);
  Matrix<T, Dynamic, Dynamic> B = BMap;
  T nb = B.norm();
  // Eigen::FullPivLU<Matrix<T, Dynamic, Dynamic> > lu(AMap);

  // BMap = lu.solve(B);
  //  BMap = AMap.fullPivLu().solve(B);
  BMap = AMap.colPivHouseholderQr().solve(B);

  /* // test
  std::cout << "-- LU system solve : ||AX - B||_2 = " << (AMap * BMap - B).norm() << std::endl;
  std::cout << "A " << std::endl;
  printM(a, *n, *n);
  std::cout << "solution" << std::endl;
  printV(b, *n);
  std::cout << "AB " << std::endl;
  T* c = new T[*n];
  affinity(*n, *n, 1, a, b, 0, c, 'N');
  printV(c, *n);
  std::cout << "-- LU system solve : ||AX - B||_2 = " << (AMap * BMap - B).norm() << std::endl;
  delete[] c;
  */

  *info = ((AMap * BMap - B).norm() <= nb * 1.0e-3f) ? 0 : 1;

}

void LUSystemSolve(int *n,double *a,int *ipiv,double *b,int *info) { subLUSystemSolve<double>(n, a, ipiv, b, info); }
void LUSystemSolve(int *n,float *a,int *ipiv,float *b,int *info) { subLUSystemSolve<float>(n, a, ipiv, b, info); }

// System solving by QR decomposition (replaces xgels_)
template <typename T>
void subQRSystemSolve(int *m,int *n,T *a,T *b,int *info)
{
    Map<Matrix<T, Dynamic, Dynamic>, 0, OuterStride<> > AMap (a, *m, *n, OuterStride<>(*m));
    Map<Matrix<T, Dynamic, Dynamic>, 0, OuterStride<> > BMap (b, std::max(*m,*n), 1, OuterStride<>(std::max(*m, *n)));
    *info = 0;
    Eigen::HouseholderQR<Matrix<T, Dynamic, Dynamic> > qr(AMap);
    Matrix<T, Dynamic, Dynamic> x = qr.solve(BMap.block(0, 0, *m, 1));
    if(*m <= *n) *info = 1 - (AMap * x).isApprox(BMap.block(0, 0, *m, 1));
    BMap.block(0, 0, *n, 1) = x;
}

void QRSystemSolve(int *m, int *n, double *a, double *b, int *info) { subQRSystemSolve<double>(m, n, a, b, info); }
void QRSystemSolve(int *m, int *n, float *a, float *b, int *info) { subQRSystemSolve<float>(m, n, a, b, info); }

// QR factorization
template <typename T>
void subQRFactorization(int *m,int *n,int *k,T *a,T *r,int *info)
{
    Map<Matrix<T, Dynamic, Dynamic> > AMap (a, *m, *k);
    Map<Matrix<T, Dynamic, Dynamic> > RMap (r, *n, *k);
    Eigen::HouseholderQR<Matrix<T, Dynamic, Dynamic> > qr(AMap);
    Matrix<T, Dynamic, Dynamic> R1 = Matrix<T, Dynamic, Dynamic>::Zero(*m, *k);
    R1= qr.matrixQR().template triangularView<Upper>();
#ifdef _WIN32
    for(int i=0; i<*n; i++)
        for(int j=0; j<*k; j++)
            RMap(i,j) = R1(i,j);
#else
    RMap= R1.block(0, 0, *n, *k);
#endif
    Matrix<T, Dynamic, Dynamic> q = qr.householderQ();
    Matrix<T, Dynamic, Dynamic> original = q.block(0, 0, *m, *n) * RMap;
    *info=AMap.isApprox(original);
    if (*n == *k)
#ifdef _WIN32
    for(int i=0; i<*m; i++)
        for(int j=0; j<*k; j++)
            AMap(i,j) = q(i,j);
#else
        AMap=q.block(0, 0, *m, *k);
#endif
    else
    {
        Map<Matrix<T, Dynamic, Dynamic> > QMap (a, *m, *m);
        QMap = q;
    }
}

void QRFactorization(int *m, int *n, int *k, double *a, double *r, int *info) { subQRFactorization<double>(m, n, k, a, r, info); }
void QRFactorization(int *m, int *n, int *k, float *a, float *r, int *info) { subQRFactorization<float>(m, n, k, a, r, info); }

// Reciprocal condition number of triangular matrix
template <typename T>
void subConditionNumber(int *n,T *a,int *lda,T *rcond,int *info)
{
    char x = '1';
    T l = (matrixNorm(&x, n, n, a, lda));
    if (l != 0)
    {
        *rcond = 1 / l;
        *info = 0;
    }
    else *info = 1;
}

void conditionNumber(int *n, double *a, int *lda, double *rcond, int *info) { subConditionNumber<double>(n, a, lda, rcond, info); }
void conditionNumber(int *n, float *a, int *lda, float *rcond, int *info) { subConditionNumber<float>(n, a, lda, rcond, info); }


/*
 * sharpqballthread.cpp
 *
 * Created on: 27.12.2012
 * @author Ralph Schurade
 */

#include "sharpqballthread.h"

#include "fmath.h"
#include "../data/datasets/datasetdwi.h"

#include "../gui/gl/glfunctions.h"

#include <qmath.h>

#include <boost/math/special_functions/spherical_harmonic.hpp>

#include "math.h"

SharpQBallThread::SharpQBallThread( DatasetDWI* ds, int order, int id ) :
    m_ds( ds ),
    m_order( order ),
    m_id( id )
{
}

SharpQBallThread::~SharpQBallThread()
{
    m_qBallVector.clear();
}

QVector<ColumnVector> SharpQBallThread::getQBallVector()
{
    return m_qBallVector;
}

void SharpQBallThread::run()
{
    QVector<QVector3D> bvecs = m_ds->getBvecs();

    Matrix gradients( bvecs.size(), 3 );
    for ( int i = 0; i < bvecs.size(); ++i )
    {
        gradients( i + 1, 1 ) = bvecs.at( i ).x();
        gradients( i + 1, 2 ) = bvecs.at( i ).y();
        gradients( i + 1, 3 ) = bvecs.at( i ).z();
    }

    QVector<ColumnVector>* data = m_ds->getData();
    QVector<float>* b0Data = m_ds->getB0Data();

    m_qBallVector.clear();

    // inverse direction matrix for calculation:
    //const matrixT A( pseudoinverse (sh_base( gradients, order ) ) );
    Matrix B = FMath::sh_base( gradients, m_order );
    Matrix A = ( B.t() * B ).i() * B.t();

    int numThreads = GLFunctions::idealThreadCount;

    int chunkSize = data->size() / numThreads;

    int begin = m_id * chunkSize;
    int end = m_id * chunkSize + chunkSize;

    if ( m_id == numThreads - 1 )
    {
        end = data->size();
    }

    // for all voxels:
    for ( int i = begin; i < end; ++i )
    {
        ColumnVector voxel( data->at( i ) / b0Data->at( i ) );

        // regularize data data:
        regularize_sqball( 0.15, 0.15, voxel );
        voxel = FMath::vlog( FMath::vlog( voxel ) * ( -1 ) );
        ColumnVector coeff = A * voxel;

        for ( int k( 2 ); k <= m_order; k += 2 )
        {
            double frt_val = 2.0 * M_PI * boost::math::legendre_p<double>( k, 0 );
            double lbt_val = -k * ( k + 1 );

            for ( int degree = -k; degree <= k; degree++ )
            {
                int l = k * ( k + 1 ) / 2 + degree + 1;
                coeff( l ) *= frt_val * lbt_val;
            }
        }

        if ( b0Data->at( i ) > 0.0 )
        {
            coeff( 1 ) = 1.0 / sqrt( 4. * M_PI );
        }

        m_qBallVector.push_back( coeff );
    }
}

void SharpQBallThread::regularize_sqball( const double par_1, const double par_2, ColumnVector& data )
{
    for ( int i = 1; i <= data.Nrows(); ++i )
    {
        if ( data( i ) < 0 )
        {
            data( i ) = 0.5 * par_1;
        }
        else if ( data( i ) < par_1 )
        {
            data( i ) = 0.5 * par_1 + 0.5 * FMath::pow2( data( i ) ) / par_1;
        }
        else if ( data( i ) < 1.0 - par_2 )
        {
            // do nothing with data
        }
        else if ( data( i ) < 1.0 )
        {
            data( i ) = 1.0 - 0.5 * par_2 - 0.5 * FMath::pow2( 1.0 - data( i ) ) / par_2;
        }
        else
        {
            data( i ) = 1.0 - 0.5 * par_2;
        }
    }
}


/**
 *  @file eigen_map.cpp
 *  @author Maximilian Harr <maximilian.harr@daimler.com>
 *  @date 29.06.2017
 *
 *  @brief Eigen map computations
 *
 *
 *
 *          Coding Standard:
 *          wiki.ros.org/CppStyleGuide
 *          https://google.github.io/styleguide/cppguide.html
 *
 *
 *  @bug
 *
 *
 *  @todo
 *
 *
 */

// PRAGMA

// SYSTEM INCLUDES
#include <iostream>  /* Header that defines the standard input/output stream objects */
#include <cstdlib>  /* Header that defines several general purpose functions */
#include <string>
#include <Eigen/Dense>  /* template library for linear algebra http://eigen.tuxfamily.org */
#include <math.h> /* Defines M_PI for pi value */
#include <Eigen/Eigenvalues>
#include <Eigen/Geometry>

// PROJECT INCLUDES

// LOCAL INCLUDES

// FORWARD REFERENCES

// FUNCTION PROTOTYPES

/** @brief Standard command line parameter processing.
 *  @param Pass command line parameters
 *  @return 0 if -h flag is set
 */
int cmd_check(int, char*[]);

// GLOBAL VARIABLES

using namespace Eigen;
using namespace std;

void outputAsMatrix(const Eigen::Quaterniond& q)
{
    std::cout << "R=" << std::endl << q.normalized().toRotationMatrix() << std::endl;
}

//// MAIN //////////////////////////////////////////////////////////////////////////////////////////
int main(int argc, char* argv[])
{
  /* Check command line options. Stop execution if -h is set. */
  if(!cmd_check(argc,argv)) return 0;

  /* Map array of doubles in Eigen::Matrix3d */
  double array[9];
  for(int i = 0; i < 9; ++i) array[i] = i;
  Eigen::Map< Eigen::Matrix3d> residual_0 = Eigen::Map<Eigen::Matrix3d>(array);
  Eigen::Map< Eigen::Matrix<double, 9, 1> > residual_1 = Eigen::Map< Eigen::Matrix<double, 9, 1> >(array);
  std::cout << "residual_0: " << std::endl << residual_0 << std::endl;
  std::cout << "res 0+0: " << std::endl << residual_0 + residual_0 << std::endl;
  std::cout << "residual_1: " << std::endl << residual_1 << std::endl;

	/* Map Eigen::Matrix3d to array of doubles */
	double array2[9];
	Eigen::Map<Matrix<double,3,3,RowMajor> >(array2,3,3) = residual_0;
	std::cout << "Array2:" << std::endl;
	for(int i = 0; i < 9; ++i){
		std::cout << array2[i] << " ";
	}
	std::cout << std::endl;

  /* Cast Eigen VectorXd to Matrix of float */
  Eigen::Vector3d vector_0;
  vector_0 << 0,1,2;
  Eigen::Matrix<float, 3, 1> p_m = vector_0.template cast<float>();
  std::cout << "p_m: " << std::endl << p_m << std::endl;

}

//// FUNCTION DEFINITIONS //////////////////////////////////////////////////////////////////////////
int cmd_check(int argc, char* argv[])
{
  int option;
  /* third argument of getopt specifies valid options and whether they need input(:) */
  while((option = getopt(argc,argv,"hp:"))>=0)
  {
    switch (option)
    {
      case 'h': std::cout
                << "Usage: <filename> [options] \n\n"
                << "<desription> \n\n"
                << "Options: \n"
                << " -h                    show this help message and exit \n"
                << " -p <PARAMETER>        <description> \n"
                << " \n";
                return 0; /* do not execute main with this option */
      case 'p': std::cout << "-p = " << optarg << "\n"; /* optarg is option argument */
                break;
    }
  }
  return 1;
}