#include <iostream>
#include <cmath>
#include <vector>
#include <chrono>
#include <memory>
#include <random>
#include <boost/range/irange.hpp>
#include <Eigen/Dense>
#include "GPs.h"

// Include CppOptLib files
#include "./include/cppoptlib/meta.h"
#include "./include/cppoptlib/boundedproblem.h"
#include "./include/cppoptlib/solver/lbfgsbsolver.h"

// Retrieve aliases from GP namescope
using Matrix = GP::Matrix;
using Vector = GP::Vector;

// Compute pairwise distance between lists of points
void GP::pdist(Matrix & Dv, Matrix & X1, Matrix & X2)
{
  auto n = static_cast<int>(X1.rows());
  int k = 0;
  auto entryCount = static_cast<int>( (n*(n-1))/2);
  Dv.resize(entryCount, 1);
  for ( auto i : boost::irange(0,n-1) )
    {
      for ( auto j : boost::irange(i+1,n) )
          Dv(k++) = (static_cast<Vector>(X1.row(i) - X2.row(j))).squaredNorm();
    }
}

// Re-assemble pairwise distances into a dense matrix
void GP::squareForm(Matrix & D, Matrix & Dv, int n, double diagVal)
{

  D.resize(n,n);

  int k = 0;
  for ( auto i : boost::irange(0,n-1) )
    {
      for ( auto j : boost::irange(i+1,n) )
        D(i,j) = D(j,i) = Dv(k++,0);
    }

  D.diagonal() = diagVal * Eigen::MatrixXd::Ones(n,1);
}

// Compute covariance matrix (and gradients) from a vector of squared pairwise distances Dv
std::vector<Matrix> GP::RBF::computeCov(Matrix & K, Matrix & Dv, Vector & params, double jitter, bool evalGrad)
{
  auto n = static_cast<int>(K.rows());
  int lengthIndex;
  if ( params.size() > paramCount )
    {
      Kv.noalias() = ( (-0.5 / std::pow(params(1),2)) * Dv ).array().exp().matrix();
      squareForm(K, Kv, n, 1.0 + params(0) + jitter);
      lengthIndex = 1;
    }
  else
    {
      Kv.noalias() = ( (-0.5 / std::pow(params(0),2)) * Dv ).array().exp().matrix();
      squareForm(K, Kv, n, 1.0 + noiseLevel + jitter);
      lengthIndex = 0;
    }

  // Compute gradient list if "evalGrad=true"
  std::vector<Matrix> gradList;
  if ( evalGrad )
    {
      dK_iv.noalias() = 1/std::pow(params(lengthIndex),2) * ( Dv.array() * Kv.array() ).matrix();
      squareForm(dK_i, dK_iv, n); // Note: diagVal = 0.0
      gradList.push_back(dK_i);
    }

  return gradList;

};

// Compute cross covariance between two input vectors using kernel parameters params
void GP::RBF::computeCrossCov(Matrix & K, Matrix & X1, Matrix & X2, Vector & params)
{
  // Get prediction count
  auto m = static_cast<int>(X2.rows());

   // Define lambda function to create unary operator (by clamping kernelParams argument)
  auto lambda = [=,&params](double d)->double { return evalDistKernel(d, params, 0); };
  for ( auto j : boost::irange(0,m) )
    {
      K.col(j) = ((X1.rowwise() - X2.row(j)).rowwise().squaredNorm()).unaryExpr(lambda);
    }

};

// Precompute distance matrix to avoid repeated calculations during optimization procedure
void GP::GaussianProcess::computeDistMat()
{
  pdist(distMatrix, obsX, obsX);
}

// Evaluate NLML for specified kernel hyperparameters p
double GP::GaussianProcess::evalNLML(const Vector & p, Vector & g, bool evalGrad)
{
  // Get matrix input observation count
  auto n = static_cast<int>(obsX.rows());

  // ASSUME OPTIMIZATION OVER LOG VALUES
  auto params = static_cast<Vector>(p);
  params = params.array().exp().matrix();

  // Compute covariance matrix and store Cholesky factor
  K.resize(n,n);
  //time start = high_resolution_clock::now();
  auto gradList = (*kernel).computeCov(K, distMatrix, params, jitter, evalGrad);
  //time end = high_resolution_clock::now();
  //time_computecov += getTime(start, end);

  //start = high_resolution_clock::now();
  cholesky = K.llt();
  //end = high_resolution_clock::now();
  //time_cholesky_llt += getTime(start, end);

  // Store alpha for DNLML calculation
  _alpha.noalias() = cholesky.solve(obsY);

  // Compute NLML value
  //start = high_resolution_clock::now();
  double NLML_value = 0.5*(obsY.transpose()*_alpha)(0);
  NLML_value += static_cast<Matrix>(cholesky.matrixL()).diagonal().array().log().sum();
  NLML_value += 0.5*n*std::log(2*PI);
  //end = high_resolution_clock::now();
  //time_NLML += getTime(start, end);

  if ( evalGrad )
    {

      //
      // Precompute the multiplicative term in derivative expressions
      //
      // [ THIS APPEARS TO BE A COMPUTATIONAL BOTTLE-NECK ]
      //

      //start = high_resolution_clock::now();
      // Direct evaluation of inverse matrix
      term.noalias() = cholesky.solve(Matrix::Identity(n,n)) - _alpha*_alpha.transpose();

      /*
      //
      //          ---  PARALLEL IMPLEMENTATION  ---
      //
      //  [ Typically much slower; possibly not thread-safe... ]
      //
      term.resize(n,n);
      //Matrix ei = Eigen::MatrixXd::Zero(n,1);
      int i = 0;
      //#pragma omp parallel for private(i) shared(cholesky, term, n)
#pragma omp parallel for private(i) shared(cholesky, _alpha, term, n)
      for ( i=0 ; i<n ; i++ )
        {
          Matrix ei = Eigen::MatrixXd::Zero(n,1);
          ei(i) = 1.0;
          //term.col(i) = cholesky.solve(ei);
          term.col(i) = cholesky.solve(ei);
          term.col(i) -= _alpha(i)*_alpha;
        }
      */

      //term.noalias() -= _alpha*_alpha.transpose();
      //end = high_resolution_clock::now();
      //time_term += getTime(start, end);

      //start = high_resolution_clock::now();
      // Compute gradient for noise term if 'fixedNoise=false'
      int index = 0;
      if (!fixedNoise)
        {
          // Specify gradient of white noise kernel
          g(index++) = 0.5 * (term * params(0)).trace() ;
        }

      // Compute gradients with respect to kernel hyperparameters
      for (auto dK_i = gradList.begin(); dK_i != gradList.end(); ++dK_i)
        {
          // Compute trace of full matrix
          g(index++) = 0.5 * (term * (*dK_i) ).trace() ;
        }
      //end = high_resolution_clock::now();
      //time_grad += getTime(start, end);

    }

  return NLML_value;

}

// Define simplified interface for evaluating NLML without gradient calculation
double GP::GaussianProcess::evalNLML(const Vector & p)
{
  Vector nullGrad(0);
  return evalNLML(p,nullGrad,false);
}

// Define function for uniform sampling
Matrix GP::sampleUnif(double a, double b, int N, int dim)
{
  //return (b-a)*(Eigen::MatrixXd::Random(N,1) * 0.5 + 0.5*Eigen::MatrixXd::Ones(N,1)) + a*Eigen::MatrixXd::Ones(N,1);
  return (b-a)*(Eigen::MatrixXd::Random(N,dim) * 0.5 + 0.5*Eigen::MatrixXd::Ones(N,dim) ) + a*Eigen::MatrixXd::Ones(N,dim);
}

// Define function for uniform sampling [Vectors]
Vector GP::sampleUnifVector(Vector lbs, Vector ubs)
{
  auto n = static_cast<int>(lbs.rows());

  Vector sampleVector = ((ubs-lbs).array()*( 0.5*Eigen::MatrixXd::Random(n,1) + 0.5*Eigen::MatrixXd::Ones(n,1) ).array()).matrix() + (lbs.array()*Eigen::MatrixXd::Ones(n,1).array()).matrix();

  return sampleVector;
}

// Define function for sampling from standard normal distribution
Matrix GP::sampleNormal(int N)
{
  // Note: Boost random is currently throwing deprecated header warnings...
  //boost::random::mt19937 rng;
  //boost::random::normal_distribution<> normalDist;
  std::default_random_engine rng;
  std::normal_distribution<double> normalDist(0.0,1.0);
  Matrix sampleVals(N,1);
  for ( auto i : boost::irange(0,N) )
    sampleVals(i) = normalDist(rng);
  return sampleVals;
}

// Generate equally spaced points on an interval or square region
Matrix GP::linspace(double a, double b, int N, int dim)
{

  Matrix linspaceVals;
  if ( dim == 1 )
    {
      linspaceVals.resize(N,1);
      linspaceVals = Eigen::Array<double, Eigen::Dynamic, 1>::LinSpaced(N, a, b);
    }
  else if ( dim == 2 )
    {
      linspaceVals.resize(N*N,2);
      Matrix linspaceVals1D = Eigen::Array<double, Eigen::Dynamic, 1>::LinSpaced(N, a, b);
      int k = 0;
      for ( auto i : boost::irange(0,N) )
        {
          for ( auto j : boost::irange(0,N) )
            {
              linspaceVals(k,0) = linspaceVals1D(i);
              linspaceVals(k,1) = linspaceVals1D(j);
              k++;
            }
        }
    }
  else
      std::cout << "[*] GP::linspace has not been implemented for dim > 2\n";

  return linspaceVals;
}

// Define utility function for formatting hyperparameter bounds
void GP::GaussianProcess::parseBounds(Vector & lbs, Vector & ubs, int augParamCount)
{
  lbs.resize(augParamCount);
  ubs.resize(augParamCount);

  double defaultLowerBound = 0.00001;
  double defaultUpperBound = 10.0;

  if ( fixedBounds )
    {
      // Check if bounds for noise parameter were provided
      if ( lowerBounds.size() < augParamCount )
        {
          // Set noise bounds to defaults
          lbs(0) = std::log( defaultLowerBound );
          ubs(0) = std::log( defaultUpperBound );

          // Convert specified bounds to log-scale
          for ( auto bi : boost::irange(1,augParamCount) )
            {
              lbs(bi) = std::log(lowerBounds(bi-1));
              ubs(bi) = std::log(upperBounds(bi-1));
            }
        }
      else
        {
          // Convert specified bounds to log-scale
          lbs = (lowerBounds.array().log()).matrix();
          ubs = (upperBounds.array().log()).matrix();
        }
    }
  else
    {
      // Set noise and hyperparameter bounds to defaults
      lbs = ( defaultLowerBound * Eigen::MatrixXd::Ones(augParamCount,1) ).array().log().matrix();
      ubs = ( defaultUpperBound * Eigen::MatrixXd::Ones(augParamCount,1) ).array().log().matrix();
    }

  /*
  std::cout << "Bounds:\n";
  std::cout << lbs.array().exp().matrix().transpose() << std::endl;
  std::cout << ubs.array().exp().matrix().transpose() << std::endl;
  std::cout << "\nLog Bounds:\n";
  std::cout << lbs.transpose() << std::endl;
  std::cout << ubs.transpose() << std::endl << std::endl;
  */

}

// Fit model hyperparameters
void GP::GaussianProcess::fitModel()
{

  // Get combined parameter/noise vector size
  paramCount = (*kernel).getParamCount();
  augParamCount = (fixedNoise) ? static_cast<int>(paramCount) : static_cast<int>(paramCount) + 1 ;

  // Pass noise level to kernel when 'fixedNoise=true'
  if ( fixedNoise )
    (*kernel).setNoise(noiseLevel);

  // Precompute distance matrix
  computeDistMat();

  // Declare vector for storing gradient calculations
  Vector g(augParamCount);

  // Initialize gradient vector size
  cppOptLibgrad.resize(augParamCount);

  // Convert hyperparameter bounds to log-scale
  Vector lbs, ubs;
  parseBounds(lbs, ubs, augParamCount);

  Vector optParams = Eigen::MatrixXd::Zero(augParamCount,1);

  this->setLowerBound(lbs);
  this->setUpperBound(ubs);
  cppoptlib::LbfgsbSolver<GaussianProcess> solver;

  // Specify stopping criteria
  cppoptlib::Criteria<double> crit = cppoptlib::Criteria<double>::defaults();
  //crit.iterations = 5000;
  //crit.gradNorm = 10.0;   //!< Minimum norm of gradient vector
  //crit.xDelta = 0;      //!< Minimum change in parameter vector
  //crit.fDelta = 0;      //!< Minimum change in cost function
  //crit.condition = 0;
  solver.setStopCriteria(crit);

  solver.minimize(*this, optParams);

  // ASSUME OPTIMIZATION OVER LOG VALUES
  optParams = optParams.array().exp().matrix();

  //start = high_resolution_clock::now();
  ///* [ This is included in the SciKit Learn model.fit() call as well ]

  // Recompute covariance and Cholesky factor
  auto nullGradList = (*kernel).computeCov(K, distMatrix, optParams);
  cholesky = K.llt();
  _alpha.noalias() = cholesky.solve(obsY);

  //*/
  //end = high_resolution_clock::now();
  //time_final_chol += getTime(start, end);

  // Assign tuned parameters to model
  if (!fixedNoise)
    {
      noiseLevel = static_cast<double>((optParams.head(1))[0]);
      optParams = static_cast<Vector>(optParams.tail(augParamCount - 1));
    }

  (*kernel).setParams(optParams);

  // DISPLAY TIMING INFORMATION
  /*
  std::cout << "\n\n Time Diagnostics |\n";
  std::cout << "------------------\n";
  std::cout << "computeCov():\t  " << time_computecov  << std::endl;
  std::cout << "cholesky.llt():\t  " << time_cholesky_llt  << std::endl;
  std::cout << "NLML:\t  \t  " << time_NLML  << std::endl;
  std::cout << "Grad term:\t  " << time_term  << std::endl;
  std::cout << "Gradient:\t  " << time_grad  << std::endl;
  std::cout << "\n" << std::endl;
  std::cout << "paramSearch:\t  " << time_paramsearch  << std::endl;
  std::cout << "Minimize:\t  " << time_minimize  << std::endl;
  std::cout << "Final Cholesky:\t  " << time_final_chol  << std::endl;
  */

};

// Compute predicted values
void GP::GaussianProcess::predict()
{
  // Get matrix input observation count
  auto n = static_cast<int>(obsX.rows());
  auto m = static_cast<int>(predX.rows());

  // Get optimized kernel hyperparameters
  Vector params = (*kernel).getParams();

  // Compute cross covariance for test points
  Matrix kstar;
  kstar.resize(n,m);
  (*kernel).computeCrossCov(kstar, obsX, predX, params);

  // Compute covariance matrix for test points
  Matrix kstarmat;
  kstarmat.resize(m,m);
  (*kernel).computeCrossCov(kstarmat, predX, predX, params);

  // Possible redundant calculations; should simplify...
  Matrix cholMat(cholesky.matrixL());
  //Matrix alpha = cholesky.solve(obsY);
  Matrix v = kstar;
  cholMat.triangularView<Eigen::Lower>().solveInPlace(v);

  // Set predictive means/variances and compute negative log marginal likelihood
  //predMean = kstar.transpose() * alpha;
  predMean = kstar.transpose() * _alpha;
  predCov = kstarmat - v.transpose() * v;

}

// Draw sample paths from posterior distribution
Matrix GP::GaussianProcess::getSamples(int count)
{
  // Get number of target points
  auto n = static_cast<int>(predX.rows());

  // Construct simple random generator engine from a time-based seed
  unsigned seed = std::chrono::system_clock::now().time_since_epoch().count();
  std::default_random_engine generator (seed);
  std::normal_distribution<double> normal (0.0,1.0);

  // Assign i.i.d. random normal values to uVals
  Matrix uVals(n,count);
  for (auto i : boost::irange(0,n))
    {
      for (auto j : boost::irange(0,count))
          uVals(i,j) = normal(generator);
    }

  // Compute Cholesky factor L
  Matrix L = ( predCov + (noiseLevel+jitter)*Matrix::Identity(static_cast<int>(predCov.cols()), static_cast<int>(predCov.cols())) ).llt().matrixL();

  // Draw samples using the formula:  y = m + L*u
  Matrix samples = predMean.replicate(1,count) + L*uVals;

  return samples;
}

// Evaluate NLML [public interface]
double GP::GaussianProcess::computeNLML(const Vector & p, double noise)
{
  // Compute log-hyperparameters
  Vector logparams(augParamCount);

  if ( !fixedNoise )
    {
      logparams(0) = std::log(noise);
      for ( auto i : boost::irange(1,augParamCount) )
        logparams(i) = std::log(p(i-1));
    }
  else
    {
      logparams = logparams.array().log().matrix();
    }

  // Evaluate NLML using log-hyperparameters
  return evalNLML(logparams);
}

// Evaluate NLML with default noise level [public interface]
double GP::GaussianProcess::computeNLML(const Vector & p)
{
  return computeNLML(p, noiseLevel);
}

// Evaluate NLML with default noise level [public interface]
double GP::GaussianProcess::computeNLML()
{
  auto params = (*kernel).getParams();
  return computeNLML(params, noiseLevel);
}

// Define function for retrieving time from chrono
float GP::getTime(std::chrono::high_resolution_clock::time_point start, std::chrono::high_resolution_clock::time_point end)
{
  return static_cast<float>(std::chrono::duration_cast<std::chrono::microseconds>( end - start ).count() / 1000000.0);
};

// Define kernel function for RBF
double GP::RBF::evalKernel(Matrix & x, Matrix & y, Vector & params, int n)
{
  switch (n)
    {
    case 0: return std::exp( -(x-y).squaredNorm() / (2.0*std::pow(params(0),2)));
    case 1: return (x-y).squaredNorm() / std::pow(params(0),3) * std::exp( -(x-y).squaredNorm() / (2.0*std::pow(params(0),2)));
    default: std::cout << "\n[*] UNDEFINED DERIVATIVE\n"; return 0.0;
    }
};

// Define distance kernel function for RBF
//
// REVISION:  Optimize w.r.t. theta = log(l) for stability  ==>   / l^2  instead of  / l^3
//
double GP::RBF::evalDistKernel(double d, Vector & params, int n)
{
  switch (n)
    {
    case 0: return std::exp( -d / (2.0*std::pow(params(0),2)));
    case 1: return d / std::pow(params(0),2) * std::exp( -d / (2.0*std::pow(params(0),2)));
    default: std::cout << "\n[*] UNDEFINED DERIVATIVE\n"; return 0.0;
    }
};

/*
// POTENTIAL PARALLEL IMPLEMENTATION OF SQUARE FORM; SPEED-UP APPEARS NEGLIGIBLE
// Re-assemble pairwise distances into a dense matrix
void GP::squareFormParallel(Matrix & D, Matrix & Dv, int n, double diagVal)
{

  D.resize(n,n);
  int i;
  int j;
#pragma omp parallel for private(i,j) shared(D,Dv,n)
  for ( i = 0 ; i<n-1; i++ )
    {
      for ( j = i+1 ; j<n ; j++ )
        {
          D(i,j) = D(j,i) = Dv( static_cast<int>(i*n - (i*(i+1))/2 + j - i -1) ,0);
        }
    }
  D.diagonal() = diagVal * Eigen::MatrixXd::Ones(n,1);
}
*/

//
//   ORIGINAL MINIMIZATION IMPLEMENTATION USING RASMUSSEN'S CODE
//

//
//    NOTE:
//
//    Remember to include "utils/minimize.h" and derive the
//   'GaussianProcess' class from the 'GradientObj' class:
//
//    i.e.  class GaussianProcess : public minimize::GradientObj
//

/*
void minimize(...)
{

  //
  // Specify the parameters for the minimization algorithm
  //

  //   HIGH ACCURACY SETTINGS   //
  // max of MAX function evaluations per line search
  //int MAX = 30;
  // max number of line searches = length
  //int length = 20;
  // don't reevaluate within INT of the limit of the current bracket
  //double INT = 0.00001;
  // SIG is a constant controlling the Wolfe-Powell conditions
  //double SIG = 0.9;
  // extrapolate maximum EXT times the current step-size
  //double EXT = 5.0;

  //  EFFICIENT SETTINGS  //

  int MAX = 15;
  int length = 10;
  double INT = 0.00001;
  double SIG = 0.9;
  double EXT = 5.0;

  // Define number of exploratory NLML evaluations for specifying
  // a reasonable initial value for the optimization algorithm
  int initParamSearchCount = 30;

  // Define restart count for optimizer
  int restartCount = 0;

  // Convert hyperparameter bounds to log-scale
  Vector lbs, ubs;
  parseBounds(lbs, ubs, augParamCount);

  // Declare variables to store optimization loop results
  double currentVal;
  double optVal = 1e9;
  Vector theta(augParamCount);
  Vector optParams(augParamCount);

  //time start = high_resolution_clock::now();
  // First explore hyperparameter space to get a reasonable initializer for optimization
  for ( auto i : boost::irange(0,initParamSearchCount) )
    {
      if ( i == 0 )
          theta = Eigen::MatrixXd::Zero(augParamCount,1);
      else
          theta = sampleUnifVector(lbs, ubs);

      // Compute current NLML and store parameters if optimal
      currentVal = evalNLML(theta);
      if ( currentVal < optVal )
        {
          optVal = currentVal;
          optParams = theta;
        }
      //std::cout << "Theta Search:  " << theta.transpose() << "  [ NLML = " << currentVal << " ]"<< std::endl;
    }
  //time end = high_resolution_clock::now();
  //time_paramsearch += getTime(start, end);

  // NOTE: THIS NEEDS TO BE RE-WRITTEN TO USE THE PRELIMINARY PARAMETER SEARCH RESULTS
  // Evaluate optimizer with various different initializations
  for ( auto i : boost::irange(0,restartCount) )
    {
      // Avoid compiler warning for unused variable
      //(void)i;

      if ( i == 0 )
        {
          // Set initial guess (should make this user specifiable...)
          theta = Eigen::MatrixXd::Zero(augParamCount,1);
        }
      else
        {
          // Sample initial hyperparameter vector
          theta = sampleUnifVector(lbs, ubs);
        }

      // Optimize hyperparameters
      minimize::cg_minimize(theta, this, g, length, SIG, EXT, INT, MAX);

      // Compute current NLML and store parameters if optimal
      currentVal = evalNLML(theta);
      if ( currentVal < optVal )
        {
          optVal = currentVal;
          optParams = theta;
        }
    }

  // Perform one last optimization starting from best parameters so far
  if ( ( initParamSearchCount == 0 ) && ( restartCount == 0 ) )
    optParams = sampleUnifVector(lbs, ubs);
  //std::cout << "\n[*] FINAL - Initial Values (log):  " << optParams.transpose() << std::endl;
  //std::cout << "[*] FINAL - Initial Values (std):  " << optParams.transpose().array().exp().matrix() << std::endl;
  //start = high_resolution_clock::now();
  minimize::cg_minimize(optParams, this, g, length, SIG, EXT, INT, MAX);
  //end = high_resolution_clock::now();
  //time_minimize += getTime(start, end);

}
*/


#include <iostream>
#include <random>
#include <Eigen/Geometry>

#include "GeometryDerivatives.h"
#include "MeshConnectivity.h"

Eigen::Matrix3d crossMatrix(Eigen::Vector3d v)
{
    Eigen::Matrix3d ret;
    ret << 0, -v[2], v[1],
        v[2], 0, -v[0],
        -v[1], v[0], 0;
    return ret;
}

Eigen::Matrix2d adjugate(Eigen::Matrix2d M)
{
    Eigen::Matrix2d ret;
    ret << M(1, 1), -M(0, 1), -M(1, 0), M(0, 0);
    return ret;
}

double angle(const Eigen::Vector3d& v, const Eigen::Vector3d& w, const Eigen::Vector3d& axis,
    Eigen::Matrix<double, 1, 9>* derivative, // v, w, a
    Eigen::Matrix<double, 9, 9>* hessian
)
{
    double theta = 2.0 * atan2((v.cross(w).dot(axis) / axis.norm()), v.dot(w) + v.norm() * w.norm());

    if (derivative)
    {
        derivative->segment<3>(0) = -axis.cross(v) / v.squaredNorm() / axis.norm();
        derivative->segment<3>(3) = axis.cross(w) / w.squaredNorm() / axis.norm();
        derivative->segment<3>(6).setZero();
    }
    if (hessian)
    {
        hessian->setZero();
        hessian->block<3, 3>(0, 0) += 2.0 * (axis.cross(v)) * v.transpose() / v.squaredNorm() / v.squaredNorm() / axis.norm();
        hessian->block<3, 3>(3, 3) += -2.0 * (axis.cross(w)) * w.transpose() / w.squaredNorm() / w.squaredNorm() / axis.norm();
        hessian->block<3, 3>(0, 0) += -crossMatrix(axis) / v.squaredNorm() / axis.norm();
        hessian->block<3, 3>(3, 3) += crossMatrix(axis) / w.squaredNorm() / axis.norm();

        Eigen::Matrix3d dahat = (Eigen::Matrix3d::Identity() / axis.norm() - axis * axis.transpose() / axis.norm() / axis.norm() / axis.norm());

        hessian->block<3, 3>(0, 6) += crossMatrix(v) * dahat / v.squaredNorm();
        hessian->block<3, 3>(3, 6) += -crossMatrix(w) * dahat / w.squaredNorm();
    }

    return theta;
}

Eigen::Vector3d faceNormal(const MeshConnectivity &mesh,
    const Eigen::MatrixXd &curPos,
    int face, int startidx,
    Eigen::Matrix<double, 3, 9> *derivative,
    std::vector<Eigen::Matrix<double, 9, 9> > *hessian)
{
    if (derivative)
        derivative->setZero();

    if (hessian)
    {
        hessian->resize(3);
        for (int i = 0; i < 3; i++) (*hessian)[i].setZero();
    }

    int v0 = startidx % 3;
    int v1 = (startidx + 1) % 3;
    int v2 = (startidx + 2) % 3;
    Eigen::Vector3d qi0 = curPos.row(mesh.faceVertex(face, v0)).transpose();
    Eigen::Vector3d qi1 = curPos.row(mesh.faceVertex(face, v1)).transpose();
    Eigen::Vector3d qi2 = curPos.row(mesh.faceVertex(face, v2)).transpose();
    Eigen::Vector3d n = (qi1 - qi0).cross(qi2 - qi0);

    if (derivative)
    {
        derivative->block<3, 3>(0, 0) += crossMatrix(qi2 - qi1);
        derivative->block<3, 3>(0, 3) += crossMatrix(qi0 - qi2);
        derivative->block<3, 3>(0, 6) += crossMatrix(qi1 - qi0);
    }

    if (hessian)
    {
        for (int j = 0; j < 3; j++)
        {
            Eigen::Vector3d ej(0, 0, 0);
            ej[j] = 1.0;
            Eigen::Matrix3d ejc = crossMatrix(ej);
            (*hessian)[j].block<3, 3>(0, 3) -= ejc;
            (*hessian)[j].block<3, 3>(0, 6) += ejc;
            (*hessian)[j].block<3, 3>(3, 6) -= ejc;
            (*hessian)[j].block<3, 3>(3, 0) += ejc;
            (*hessian)[j].block<3, 3>(6, 0) -= ejc;
            (*hessian)[j].block<3, 3>(6, 3) += ejc;
        }
    }

    return n;
}

double triangleAltitude(const MeshConnectivity &mesh,
    const Eigen::MatrixXd &curPos,
    int face,
    int edgeidx,
    Eigen::Matrix<double, 1, 9> *derivative,
    Eigen::Matrix<double, 9, 9> *hessian)
{
    if (derivative)
        derivative->setZero();
    if (hessian)
        hessian->setZero();

    Eigen::Matrix<double, 3, 9> nderiv;
    std::vector<Eigen::Matrix<double, 9, 9> > nhess;
    Eigen::Vector3d n = faceNormal(mesh, curPos, face, edgeidx, (derivative || hessian ? &nderiv : NULL), hessian ? &nhess : NULL);

    int v2 = (edgeidx + 2) % 3;
    int v1 = (edgeidx + 1) % 3;
    Eigen::Vector3d q2 = curPos.row(mesh.faceVertex(face, v2)).transpose();
    Eigen::Vector3d q1 = curPos.row(mesh.faceVertex(face, v1)).transpose();

    Eigen::Vector3d e = q2 - q1;
    double nnorm = n.norm();
    double enorm = e.norm();
    double h = nnorm / enorm;

    if (derivative)
    {
        for (int i = 0; i < 3; i++)
        {
            *derivative += nderiv.row(i) * n[i] / nnorm / enorm;
        }
        derivative->block(0, 6, 1, 3) += -nnorm / enorm / enorm / enorm * e.transpose();
        derivative->block(0, 3, 1, 3) += nnorm / enorm / enorm / enorm * e.transpose();
    }

    if (hessian)
    {
        for (int i = 0; i < 3; i++)
        {
            *hessian += nhess[i] * n[i] / nnorm / enorm;
        }

        Eigen::Matrix3d idMat = Eigen::Matrix3d::Identity();

        Eigen::Matrix3d P = idMat / nnorm - n*n.transpose() / nnorm / nnorm / nnorm;
        *hessian += nderiv.transpose() * P * nderiv / enorm;
        hessian->block<3, 9>(6, 0) += -e * n.transpose() * nderiv / nnorm / enorm / enorm / enorm;
        hessian->block<3, 9>(3, 0) += e * n.transpose() * nderiv / nnorm / enorm / enorm / enorm;
        hessian->block<9, 3>(0, 6) += -nderiv.transpose() * n * e.transpose() / nnorm / enorm / enorm / enorm;
        hessian->block<9, 3>(0, 3) += nderiv.transpose() * n * e.transpose() / nnorm / enorm / enorm / enorm;
        hessian->block<3, 3>(6, 6) += -nnorm / enorm / enorm / enorm * idMat;
        hessian->block<3, 3>(6, 3) += nnorm / enorm / enorm / enorm * idMat;
        hessian->block<3, 3>(3, 6) += nnorm / enorm / enorm / enorm * idMat;
        hessian->block<3, 3>(3, 3) += -nnorm / enorm / enorm / enorm * idMat;

        Eigen::Matrix3d outer = e*e.transpose() * 3.0 * nnorm / enorm / enorm / enorm / enorm / enorm;
        hessian->block(6, 6, 3, 3) += outer;
        hessian->block(6, 3, 3, 3) += -outer;
        hessian->block(3, 6, 3, 3) += -outer;
        hessian->block(3, 3, 3, 3) += outer;
    }

    return h;
}

Eigen::Matrix2d firstFundamentalForm(
    const MeshConnectivity &mesh,
    const Eigen::MatrixXd &curPos,
    int face,
    Eigen::Matrix<double, 4, 9> *derivative, // F(face, i)
    std::vector <Eigen::Matrix<double, 9, 9> > *hessian)
{
    Eigen::Vector3d q0 = curPos.row(mesh.faceVertex(face, 0));
    Eigen::Vector3d q1 = curPos.row(mesh.faceVertex(face, 1));
    Eigen::Vector3d q2 = curPos.row(mesh.faceVertex(face, 2));
    Eigen::Matrix2d result;
    result << (q1 - q0).dot(q1 - q0), (q1 - q0).dot(q2 - q0),
        (q2 - q0).dot(q1 - q0), (q2 - q0).dot(q2 - q0);

    if (derivative)
    {
        derivative->setZero();
        derivative->block<1, 3>(0, 3) += 2.0 * (q1 - q0).transpose();
        derivative->block<1, 3>(0, 0) -= 2.0 * (q1 - q0).transpose();
        derivative->block<1, 3>(1, 6) += (q1 - q0).transpose();
        derivative->block<1, 3>(1, 3) += (q2 - q0).transpose();
        derivative->block<1, 3>(1, 0) += -(q1 - q0).transpose() - (q2 - q0).transpose();
        derivative->block<1, 3>(2, 6) += (q1 - q0).transpose();
        derivative->block<1, 3>(2, 3) += (q2 - q0).transpose();
        derivative->block<1, 3>(2, 0) += -(q1 - q0).transpose() - (q2 - q0).transpose();
        derivative->block<1, 3>(3, 6) += 2.0 * (q2 - q0).transpose();
        derivative->block<1, 3>(3, 0) -= 2.0 * (q2 - q0).transpose();
    }

    if (hessian)
    {
        hessian->resize(4);
        for (int i = 0; i < 4; i++)
        {
            (*hessian)[i].setZero();
        }
        Eigen::Matrix3d I = Eigen::Matrix3d::Identity();
        (*hessian)[0].block<3, 3>(0, 0) += 2.0*I;
        (*hessian)[0].block<3, 3>(3, 3) += 2.0*I;
        (*hessian)[0].block<3, 3>(0, 3) -= 2.0*I;
        (*hessian)[0].block<3, 3>(3, 0) -= 2.0*I;

        (*hessian)[1].block<3, 3>(3, 6) += I;
        (*hessian)[1].block<3, 3>(6, 3) += I;
        (*hessian)[1].block<3, 3>(0, 3) -= I;
        (*hessian)[1].block<3, 3>(0, 6) -= I;
        (*hessian)[1].block<3, 3>(3, 0) -= I;
        (*hessian)[1].block<3, 3>(6, 0) -= I;
        (*hessian)[1].block<3, 3>(0, 0) += 2.0*I;

        (*hessian)[2].block<3, 3>(3, 6) += I;
        (*hessian)[2].block<3, 3>(6, 3) += I;
        (*hessian)[2].block<3, 3>(0, 3) -= I;
        (*hessian)[2].block<3, 3>(0, 6) -= I;
        (*hessian)[2].block<3, 3>(3, 0) -= I;
        (*hessian)[2].block<3, 3>(6, 0) -= I;
        (*hessian)[2].block<3, 3>(0, 0) += 2.0*I;

        (*hessian)[3].block<3, 3>(0, 0) += 2.0*I;
        (*hessian)[3].block<3, 3>(6, 6) += 2.0*I;
        (*hessian)[3].block<3, 3>(0, 6) -= 2.0*I;
        (*hessian)[3].block<3, 3>(6, 0) -= 2.0*I;
    }

    return result;
}


/**
 * @file LossTests.cpp
 *
 * @breif Tests for our loss functions
 *
 * @date 1/06/17
 * @author Ben Caine
 */

#define BOOST_TEST_DYN_LINK
#define BOOST_TEST_MODULE LossTests

#include <boost/test/unit_test.hpp>
#include <iostream>
#include "loss/CrossEntropy.h"
#include "loss/HuberLoss.h"
#include "loss/MeanSquaredError.h"

BOOST_AUTO_TEST_CASE(test_cross_entropy_loss) {
    // TODO: Really just checking compilation right now. Need to use TF/Pytorch/Numpy to generate test cases
    nn::CrossEntropyLoss<float, 2> lossFunction;

    int batchSize = 4;
    int numClasses = 3;
    Eigen::Tensor<float, 2> predictions(batchSize, numClasses);
    predictions.setValues({
                          {0.1, 0.7, 0.2},
                          {0.9, 0.0, 0.1},
                          {0.0, 0.0, 1.0},
                          {0.3, 0.4, 0.3}
                          });

    Eigen::Tensor<float, 2> labels(batchSize, numClasses);
    labels.setValues({
                             {0, 1, 0},
                             {1, 0, 0},
                             {1, 0, 0},
                             {0, 1, 0}
                     });

    auto loss = lossFunction.loss(predictions, labels);
    auto accuracy = lossFunction.accuracy(predictions, labels);
    auto backwardsResult = lossFunction.backward(predictions, labels);
    BOOST_REQUIRE_MESSAGE(accuracy == 0.75, "Accuracy was not correct");
}

BOOST_AUTO_TEST_CASE(test_mse_loss) {
    nn::MeanSquaredError<float, 2> lossFunction;

    const int batchSize = 4;

    Eigen::Tensor<float, 2> predictions(batchSize, 1);
    predictions.setValues({{2}, {3}, {4}, {5}});

    Eigen::Tensor<float, 2> labels(batchSize, 1);
    labels.setValues({{2}, {1}, {3}, {0}});

    // Expected squared error of each should be:
    // 0^2 + 2^2 + 1^2 + 5^2 = 0 + 4 + 1 + 25 = 30 / 4 = 7.5

    auto loss = lossFunction.loss(predictions, labels);
    BOOST_REQUIRE_MESSAGE(loss == 7.5, "Loss not what was expected");

    auto backwardsResult = lossFunction.backward(predictions, labels);
    std::array<float, batchSize> expectedBackwardsResults = {0, 2, 1, 5};

    for (int ii = 0; ii < batchSize; ++ii) {
        BOOST_REQUIRE_CLOSE(backwardsResult(ii, 0), expectedBackwardsResults[ii], 1e-3);
    }
}

BOOST_AUTO_TEST_CASE(test_huber_loss) {
    float threshold = 1.5;
    nn::HuberLoss<float, 2> lossFunction(threshold);

    const int batchSize = 4;
    Eigen::Tensor<float, 2> predictions(batchSize, 1);
    predictions.setValues({{2}, {3}, {4}, {5}});

    Eigen::Tensor<float, 2> labels(batchSize, 1);
    labels.setValues({{2}, {1}, {3}, {0}});

    // Expected absolute error:
    // [0, 2, 1, 5]
    // If our thresh is 1.5, then two terms are squared loss and two are absolute
    // So, we expect:
    // 0.5 * 0^2 + (1.5 * 2 - 0.5 * 1.5^2) + 0.5 * 1^2 + (1.5 * 5 - 0.5 * 1.5^2)
    // Which is:
    // [0, 1.875, 0.5, 6.375] = 8.75 / 4 = 2.1875

    auto loss = lossFunction.loss(predictions, labels);
    BOOST_REQUIRE_CLOSE(loss, 2.1875, 1e-3);

    auto backwardsResult = lossFunction.backward(predictions, labels);

    std::array<float, batchSize> expectedBackwardsResults = {0, threshold, 1, threshold};

    for (int ii = 0; ii < batchSize; ++ii) {
        BOOST_REQUIRE_CLOSE(backwardsResult(ii, 0), expectedBackwardsResults[ii], 1e-3);
    }
}


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

#include <cmath>
#include <iterator>
#include <utility>
#include <algorithm>
#include <type_traits>
#include <initializer_list>
#include <boost/math/special_functions/logaddexp.hpp>

namespace boost { namespace math {

// https://nhigham.com/2021/01/05/what-is-the-log-sum-exp-function/
// See equation (#)
template <typename ForwardIterator, typename Real = typename std::iterator_traits<ForwardIterator>::value_type>
Real logsumexp(ForwardIterator first, ForwardIterator last)
{
    using std::exp;
    using std::log1p;

    const auto elem = std::max_element(first, last);
    const Real max_val = *elem;

    Real arg = 0;
    while (first != last)
    {
        if (first != elem)
        {
            arg += exp(*first - max_val);
        }

        ++first;
    }

    return max_val + log1p(arg);
}

template <typename Container, typename Real = typename Container::value_type>
inline Real logsumexp(const Container& c)
{
    return logsumexp(std::begin(c), std::end(c));
}

template <typename... Args, typename Real = typename std::common_type<Args...>::type,
          typename std::enable_if<std::is_floating_point<Real>::value, bool>::type = true>
inline Real logsumexp(Args&& ...args)
{
    std::initializer_list<Real> list {std::forward<Args>(args)...};

    if(list.size() == 2)
    {
        return logaddexp(*list.begin(), *std::next(list.begin()));
    }
    return logsumexp(list.begin(), list.end());
}

}} // Namespace boost::math


/*=================================================================================
 *	                    Copyleft! 2018 William Yu
 *          Some rights reserved：CC(creativecommons.org)BY-NC-SA
 *                      Copyleft! 2018 William Yu
 *      版权部分所有，遵循CC(creativecommons.org)BY-NC-SA协议授权方式使用
 *
 * Filename                :
 * Description             : 视觉SLAM十四讲/ch3/useGeometry/eigenGeometry.cpp 学习记录
								矩阵库
 * Reference               :
 * Programmer(s)           : William Yu, windmillyucong@163.com
 * Company                 : HUST, DMET国家重点实验室FOCUS团队
 * Modification History	   : ver1.0, 2018.03.26, William Yu

=================================================================================*/
/// Include Files
#include <iostream>
#include <cmath>
#include <Eigen/Core>
// Eigen 几何模块
#include <Eigen/Geometry>
using namespace std;

/// Global Variables

/// Function Definitions

/**
 * @function main
 * @author William Yu
 * @brief Eigen几何模块的使用方法
 * @param  None
 * @retval None
 */
int main ( int argc, char** argv )
{
    //---------------------------------------------------------------------------
    //  旋转平移表示
    //----------------------------------------------------------------
    // Eigen/Geometry 几何模块提供了各种旋转和平移的表示
    // 3D 旋转矩阵直接使用 Matrix3d 或 Matrix3f
    Eigen::Matrix3d rotation_matrix = Eigen::Matrix3d::Identity();
    // 旋转向量使用 AngleAxis, 它底层不直接是Matrix，但运算可以当作矩阵（因为重载了运算符）
    Eigen::AngleAxisd rotation_vector ( M_PI/4, Eigen::Vector3d ( 0,0,1 ) );     //沿 Z 轴旋转 45 度
    cout .precision(3);
    cout<<"rotation matrix =\n"<<rotation_vector.matrix() <<endl;                //用matrix()转换成矩阵
    // 也可以直接赋值
    rotation_matrix = rotation_vector.toRotationMatrix();
    // 用 AngleAxis 可以进行坐标变换
    Eigen::Vector3d v ( 1,0,0 );
    Eigen::Vector3d v_rotated = rotation_vector * v;
    cout<<"(1,0,0) after rotation = "<<v_rotated.transpose()<<endl;
    // 或者用旋转矩阵
    v_rotated = rotation_matrix * v;
    cout<<"(1,0,0) after rotation = "<<v_rotated.transpose()<<endl;

    //---------------------------------------------------------------------------
    //  欧拉角
    //----------------------------------------------------------------
    //-- 可以将旋转矩阵直接转换成欧拉角
    Eigen::Vector3d euler_angles = rotation_matrix.eulerAngles ( 2,1,0 ); // ZYX顺序，即roll pitch yaw顺序
    cout<<"yaw pitch roll = "<<euler_angles.transpose()<<endl;

    // 欧氏变换矩阵使用 Eigen::Isometry
    Eigen::Isometry3d T=Eigen::Isometry3d::Identity();                // 虽然称为3d，实质上是4＊4的矩阵
    T.rotate ( rotation_vector );                                     // 按照rotation_vector进行旋转
    T.pretranslate ( Eigen::Vector3d ( 1,3,4 ) );                     // 把平移向量设成(1,3,4)
    cout << "Transform matrix = \n" << T.matrix() <<endl;

    // 用变换矩阵进行坐标变换
    Eigen::Vector3d v_transformed = T*v;                              // 相当于R*v+t
    cout<<"v tranformed = "<<v_transformed.transpose()<<endl;

    // 对于仿射和射影变换，使用 Eigen::Affine3d 和 Eigen::Projective3d 即可，略

    //---------------------------------------------------------------------------
    //  四元数
    //----------------------------------------------------------------
    // 可以直接把AngleAxis赋值给四元数，反之亦然
    Eigen::Quaterniond q = Eigen::Quaterniond ( rotation_vector );
    cout<<"quaternion = \n"<<q.coeffs() <<endl;   // 请注意coeffs的顺序是(x,y,z,w),w为实部，前三者为虚部
    // 也可以把旋转矩阵赋给它
    q = Eigen::Quaterniond ( rotation_matrix );
    cout<<"quaternion = \n"<<q.coeffs() <<endl;
    // 使用四元数旋转一个向量，使用重载的乘法即可
    v_rotated = q*v; // 注意数学上是qvq^{-1}
    cout<<"(1,0,0) after rotation = "<<v_rotated.transpose()<<endl;

    return 0;
}