#include "SDOT/SemiDiscreteOT.h"

#include "SDOT/Assert.h"
#include "SDOT/Distances/Distances.h"

#include <Eigen/Core>
#include <Eigen/SparseCholesky>
#include <Eigen/Dense>

#include <CGAL/Kernel/global_functions.h>

#include <algorithm>
#include <chrono>

using namespace sdot;

template<typename ConjugateFunctionType>
SemidiscreteOT<ConjugateFunctionType>::SemidiscreteOT(std::shared_ptr<Distribution2d> const& distIn,
                               Eigen::Matrix2Xd                const& discrPtsIn,
                               Eigen::VectorXd                 const& discrProbsIn,
                               double                                 unbalancedPenaltyIn) : dist(distIn),
                                                                                grid(distIn->Grid()),
                                                                                discrPts(discrPtsIn),
                                                                                discrProbs(discrProbsIn),
                                                                                unbalancedPenalty(unbalancedPenaltyIn)
{
  if(discrPtsIn.cols()!=discrProbsIn.size())
  SDOT_ASSERT(discrPtsIn.cols()==discrProbsIn.size());

  SDOT_ASSERT(unbalancedPenalty>0);

  CheckNormalization();

  // Check to make sure all the points are inside the grid domain
  for(unsigned int i=0; i<discrPts.cols(); ++i){
    SDOT_ASSERT(discrPts(0,i)>=grid->xMin);
    SDOT_ASSERT(discrPts(0,i)<=grid->xMax);
    SDOT_ASSERT(discrPts(1,i)>=grid->yMin);
    SDOT_ASSERT(discrPts(1,i)<=grid->yMax);
  }
}

template<>
void SemidiscreteOT<sdot::distances::Wasserstein2>::CheckNormalization()
{
  double discrSum = discrProbs.sum();
  double contSum= dist->TotalMass();
  SDOT_ASSERT(std::abs(discrSum-contSum)<1e-5);
}

template<typename ConjugateFunctionType>
void SemidiscreteOT<ConjugateFunctionType>::SetPoints(Eigen::Matrix2Xd const& newPts){
  SDOT_ASSERT(newPts.cols()==discrProbs.size());

  // Check to make sure all the points are inside the grid domain
  for(unsigned int i=0; i<newPts.cols(); ++i){
    SDOT_ASSERT(newPts(0,i)>=grid->xMin);
    SDOT_ASSERT(newPts(0,i)<=grid->xMax);
    SDOT_ASSERT(newPts(1,i)>=grid->yMin);
    SDOT_ASSERT(newPts(1,i)<=grid->yMax);
  }

  discrPts = newPts;
}

template<typename ConjugateFunctionType>
std::tuple<double,Eigen::VectorXd, Eigen::SparseMatrix<double>> SemidiscreteOT<ConjugateFunctionType>::Objective(Eigen::VectorXd const& prices) const
{
    // Notes:
    //   - The cost c(x,y) is the squared distance between x and y
    //   - See (17) of https://arxiv.org/pdf/1710.02634.pdf

    const int numCells = discrPts.cols();
    SDOT_ASSERT(numCells==prices.size());

    // Construct the Laguerre diagram
    LaguerreDiagram lagDiag(grid->xMin, grid->xMax, grid->yMin, grid->yMax, discrPts, prices);

    double obj;
    Eigen::VectorXd grad;
    Eigen::SparseMatrix<double> hess;

    std::tie(obj,grad) = ComputeGradient(prices, lagDiag);
    hess = ComputeHessian(prices, lagDiag);

    return std::make_tuple(obj,grad,hess);
}

template<typename ConjugateFunctionType>
Eigen::Matrix2Xd SemidiscreteOT<ConjugateFunctionType>::PointGradient() const
{
  return PointGradient(optPrices, *lagDiag);
}

template<typename ConjugateFunctionType>
Eigen::Matrix2Xd SemidiscreteOT<ConjugateFunctionType>::PointGradient(Eigen::VectorXd const& prices,
                                                                      LaguerreDiagram const& lagDiag) const
{
  int numCells =  lagDiag.NumCells();
  Eigen::Matrix2Xd grad(2,numCells);

  // Loop over all ofthe cells
  for(int i=0; i<numCells; ++i){

    auto triFunc = [&](Eigen::Vector2d const& pt1, Eigen::Vector2d const& pt2, Eigen::Vector2d const& pt3)
      {
        return ConjugateFunctionType::TriangularIntegralPointGrad(prices(i), discrPts.col(i), pt1,pt2,pt3,unbalancedPenalty);
      };
    auto rectFunc = [&](Eigen::Vector2d const& pt1, Eigen::Vector2d const& pt2)
      {
        return ConjugateFunctionType::RectangularIntegralPointGrad(prices(i), discrPts.col(i), pt1,pt2,unbalancedPenalty);
      };

    grad.col(i) = lagDiag.IntegrateOverCell(i, triFunc, rectFunc, dist);
  }

  return grad;
}

template<typename ConjugateFunctionType>
Eigen::Matrix2Xd SemidiscreteOT<ConjugateFunctionType>::LloydPointHessian() const
{
  return LloydPointHessian(optPrices, *lagDiag);
}

template<typename ConjugateFunctionType>
Eigen::Matrix2Xd SemidiscreteOT<ConjugateFunctionType>::LloydPointHessian(Eigen::VectorXd const& prices,
                                                                         LaguerreDiagram const& lagDiag) const
{
  const unsigned int numCells = discrPts.cols();
  Eigen::Matrix2Xd hessVals(2,numCells);
  Eigen::Matrix2d intVal(2,2);

  // For unbalanced transport, the Diagonal of the hessian has an additional term
  for(unsigned int cellInd=0; cellInd<numCells; ++cellInd) {

    auto triFunc = [&](Eigen::Vector2d const& pt1, Eigen::Vector2d const& pt2, Eigen::Vector2d const& pt3)
      {
        return ConjugateFunctionType::TriangularIntegralPointHessDiag(prices(cellInd), discrPts.col(cellInd), pt1,pt2,pt3, unbalancedPenalty);
      };
    auto rectFunc = [&](Eigen::Vector2d const& pt1, Eigen::Vector2d const& pt2)
      {
        return ConjugateFunctionType::RectangularIntegralPointHessDiag(prices(cellInd), discrPts.col(cellInd), pt1,pt2, unbalancedPenalty);
      };

    intVal = -1.0*lagDiag.IntegrateOverCell(cellInd, triFunc, rectFunc, dist);
    hessVals(0,cellInd)   =  intVal(0,0);
    hessVals(1,cellInd)   =  intVal(1,1);
  }

  return hessVals;
}

template<typename ConjugateFunctionType>
Eigen::SparseMatrix<double> SemidiscreteOT<ConjugateFunctionType>::PointHessian(Eigen::VectorXd const& prices,
                                                                                LaguerreDiagram const& lagDiag) const
{
  const unsigned int numCells = discrPts.cols();
  typedef Eigen::Triplet<double> T;
  std::vector<T> hessVals;

  Eigen::VectorXd diagVals = Eigen::VectorXd::Zero(2*numCells);
  Eigen::Matrix2d intVal;

  // For unbalanced transport, the Diagonal of the hessian has an additional term
  for(unsigned int cellInd=0; cellInd<numCells; ++cellInd) {

    auto triFunc = [&](Eigen::Vector2d const& pt1, Eigen::Vector2d const& pt2, Eigen::Vector2d const& pt3)
      {
        return ConjugateFunctionType::TriangularIntegralPointHessDiag(prices(cellInd), discrPts.col(cellInd), pt1,pt2,pt3, unbalancedPenalty);
      };
    auto rectFunc = [&](Eigen::Vector2d const& pt1, Eigen::Vector2d const& pt2)
      {
        return ConjugateFunctionType::RectangularIntegralPointHessDiag(prices(cellInd), discrPts.col(cellInd), pt1,pt2, unbalancedPenalty);
      };

    intVal = -1.0*lagDiag.IntegrateOverCell(cellInd, triFunc, rectFunc, dist);
    hessVals.push_back(T(2*cellInd,2*cellInd,intVal(0,0)));
    hessVals.push_back(T(2*cellInd,2*cellInd+1,intVal(0,1)));
    hessVals.push_back(T(2*cellInd+1,2*cellInd,intVal(1,0)));
    hessVals.push_back(T(2*cellInd+1,2*cellInd+1,intVal(1,1)));
  }

  ///////
  // Off-diagonal parts

  unsigned int cellInd2;
  LaguerreDiagram::Point_2 srcPt, tgtPt;

  for(unsigned int cellInd1=0; cellInd1<numCells; ++cellInd1){

    for(auto edgeTuple : lagDiag.InternalEdges(cellInd1)){
      std::tie(cellInd2, srcPt, tgtPt) = edgeTuple;

      // Compute the integral of the target density along the edge
      auto func = [&](Eigen::Vector2d const& pt1, Eigen::Vector2d const& pt2)
      {
        return ConjugateFunctionType::LineIntegralPointHess(prices(cellInd1), discrPts.col(cellInd1), discrPts.col(cellInd2), pt1, pt2, unbalancedPenalty);
      };
      intVal = LineIntegral(func, srcPt, tgtPt)/(discrPts.col(cellInd1)-discrPts.col(cellInd2)).norm();

      hessVals.push_back(T(2*cellInd1,2*cellInd1,-intVal(0,0)));
      hessVals.push_back(T(2*cellInd1+1,2*cellInd1,-intVal(0,1)));
      hessVals.push_back(T(2*cellInd1,2*cellInd1+1,-intVal(1,0)));
      hessVals.push_back(T(2*cellInd1+1,2*cellInd1+1,-intVal(1,1)));

      hessVals.push_back(T(2*cellInd1,2*cellInd2,intVal(0,0)));
      hessVals.push_back(T(2*cellInd1+1,2*cellInd2,intVal(0,1)));
      hessVals.push_back(T(2*cellInd1,2*cellInd2+1,intVal(1,0)));
      hessVals.push_back(T(2*cellInd1+1,2*cellInd2+1,intVal(1,1)));

    }
  }

  Eigen::SparseMatrix<double> hess(2*numCells,2*numCells);
  hess.setFromTriplets(hessVals.begin(), hessVals.end());

  return hess;

}

template<typename ConjugateFunctionType>
Eigen::SparseMatrix<double> SemidiscreteOT<ConjugateFunctionType>::PointHessian() const
{
  return PointHessian(optPrices, *lagDiag);
}

template<typename ConjugateFunctionType>
std::pair<double,Eigen::VectorXd> SemidiscreteOT<ConjugateFunctionType>::ComputeGradient(Eigen::VectorXd const& prices,
                                                                                LaguerreDiagram const& lagDiag) const
{
  const int numCells = prices.size();

  // Holds the part of the objective for each cell in the Laguerre diagram
  Eigen::VectorXd objParts = Eigen::VectorXd::Zero(numCells);
  Eigen::VectorXd gradient = Eigen::VectorXd::Zero(numCells);

  Eigen::VectorXd probs(numCells);

  //Eigen::MatrixXd cellAreas = Eigen::MatrixXd::Zero(grid->Nx, grid->Ny);
#if defined(_OPENMP)
  #pragma omp parallel for
#endif
  for(int cellInd=0; cellInd<numCells; ++cellInd){

    // auto area_integrand = std::make_shared<ConstantIntegrand>();
    //
    // auto trans_integrand = std::make_shared<TransportIntegrand>(discrPts.col(cellInd));

    // auto triArea = [](Eigen::Vector2d const& pt1, Eigen::Vector2d const& pt2, Eigen::Vector2d const& pt3)
    //   {
    //     return 0.5*std::abs((pt2[0]*pt1[1]-pt1[0]*pt2[1])+(pt3[0]*pt2[1]-pt2[0]*pt3[1])+(pt1[0]*pt3[1]-pt3[0]*pt1[1]));
    //   };
    //
    // auto rectArea = [](Eigen::Vector2d const& bottomLeft, Eigen::Vector2d const& topRight)
    //   {
    //     return std::abs((topRight[0]-bottomLeft[0])*(topRight[1]-bottomLeft[1]));
    //   };
    //
    //
    // double weightedArea  = lagDiag.IntegrateOverCell(cellInd, triArea, rectArea, dist);

    objParts(cellInd) = -ConjugateFunctionType::Evaluate(-prices(cellInd), unbalancedPenalty)*discrProbs(cellInd);// - prices(cellInd)*weightedArea;

    gradient(cellInd) = ConjugateFunctionType::Derivative(-prices(cellInd), unbalancedPenalty)*discrProbs(cellInd);

    // if(weightedArea>0){

      auto triFunc = [&](Eigen::Vector2d const& pt1, Eigen::Vector2d const& pt2, Eigen::Vector2d const& pt3)
        {
          return ConjugateFunctionType::TriangularIntegral(prices(cellInd), discrPts.col(cellInd), pt1,pt2,pt3,unbalancedPenalty);
        };
      auto rectFunc = [&](Eigen::Vector2d const& pt1, Eigen::Vector2d const& pt2)
        {
          return ConjugateFunctionType::RectangularIntegral(prices(cellInd), discrPts.col(cellInd), pt1,pt2,unbalancedPenalty);
        };

      objParts(cellInd) += -lagDiag.IntegrateOverCell(cellInd, triFunc, rectFunc, dist);

      auto triFuncDeriv = [&](Eigen::Vector2d const& pt1, Eigen::Vector2d const& pt2, Eigen::Vector2d const& pt3)
        {
          return ConjugateFunctionType::TriangularIntegralDeriv(prices(cellInd), discrPts.col(cellInd), pt1,pt2,pt3,unbalancedPenalty);
        };
      auto rectFuncDeriv = [&](Eigen::Vector2d const& pt1, Eigen::Vector2d const& pt2)
        {
          return ConjugateFunctionType::RectangularIntegralDeriv(prices(cellInd), discrPts.col(cellInd), pt1,pt2,unbalancedPenalty);
        };

      gradient(cellInd) += -lagDiag.IntegrateOverCell(cellInd, triFuncDeriv, rectFuncDeriv, dist); //-weightedArea;
    }

  //   probs(cellInd) = weightedArea;
  // }
  //
  // double totalProb = probs.sum();
  // if(std::abs(totalProb-1.0)>1e-10){
  //
  //   std::cout << "Warning:  Total probability has an error of " << totalProb-1.0 << std::endl;
  //   // std::cout << "Prices = " << prices.transpose() << std::endl;
  //   // for(unsigned int cellInd=0; cellInd<lagDiag.NumCells(); ++cellInd){
  //   //   std::cout << "Cell " << cellInd << " has points " << std::endl;
  //   //   Eigen::MatrixXd pts = lagDiag.GetCellVertices(cellInd);
  //   //   for(unsigned int ptInd=0; ptInd<pts.cols(); ++ptInd){
  //   //      std::cout << "[" << pts(0,ptInd) << "," << pts(1,ptInd) << "], ";
  //   //   }
  //   //   std::cout << std::endl;
  //   // }
  //   throw std::runtime_error("Error in total probability.");
  //
  //
  //
  //  //SDOT_ASSERT(std::abs(weightedArea-1.0)<1e-10);
  // }

  return std::make_pair(objParts.sum(), gradient);
}

template<typename ConjugateFunctionType>
Eigen::SparseMatrix<double> SemidiscreteOT<ConjugateFunctionType>::ComputeHessian(Eigen::VectorXd const& prices,
                                                                                LaguerreDiagram const& lagDiag) const
{
  const unsigned int numCells = discrPts.cols();
  typedef Eigen::Triplet<double> T;

  /* The diagonal entries of the Hessian are the negative sum of the off diagonals
     See equation 27 of https://arxiv.org/pdf/1710.02634.pdf
     This vector is used to keep track of this sum for each cell.
  */
  Eigen::VectorXd diagVals = Eigen::VectorXd::Zero(numCells);

  // For unbalanced transport, the Diagonal of the hessian has an additional term
  for(unsigned int cellInd=0; cellInd<numCells; ++cellInd) {

    diagVals(cellInd) -= ConjugateFunctionType::Derivative2(-prices(cellInd), unbalancedPenalty)*discrProbs(cellInd);

    auto triFunc = [&](Eigen::Vector2d const& pt1, Eigen::Vector2d const& pt2, Eigen::Vector2d const& pt3)
      {
        return ConjugateFunctionType::TriangularIntegralDeriv2(prices(cellInd), discrPts.col(cellInd), pt1,pt2,pt3, unbalancedPenalty);
      };
    auto rectFunc = [&](Eigen::Vector2d const& pt1, Eigen::Vector2d const& pt2)
      {
        return ConjugateFunctionType::RectangularIntegralDeriv2(prices(cellInd), discrPts.col(cellInd), pt1,pt2, unbalancedPenalty);
      };

    diagVals(cellInd) -= lagDiag.IntegrateOverCell(cellInd, triFunc, rectFunc, dist);
  }

  // Hold the i,j,val triplets defining the sparse Hessian
  std::vector<T> hessVals;

  double intVal;
  unsigned int cellInd2;
  LaguerreDiagram::Point_2 srcPt, tgtPt;

  for(unsigned int cellInd1=0; cellInd1<numCells; ++cellInd1){

    for(auto edgeTuple : lagDiag.InternalEdges(cellInd1)){
      std::tie(cellInd2, srcPt, tgtPt) = edgeTuple;

      // Compute the integral of the target density along the edge
      auto func = [&](Eigen::Vector2d const& pt1, Eigen::Vector2d const& pt2)
      {
        return ConjugateFunctionType::LineIntegralDeriv(prices(cellInd1), discrPts.col(cellInd1), pt1, pt2, unbalancedPenalty);
      };
      intVal = 0.5*LineIntegral(func, srcPt, tgtPt)/(discrPts.col(cellInd1)-discrPts.col(cellInd2)).norm();

      diagVals(cellInd1) -= intVal;
      hessVals.push_back(T(cellInd1,cellInd2,intVal));
    }
  }

  for(int i=0; i<numCells; ++i){
    hessVals.push_back(T(i,i, diagVals(i)));
  }

  Eigen::SparseMatrix<double> hess(numCells,numCells);
  hess.setFromTriplets(hessVals.begin(), hessVals.end());

  return hess;
}

// double SemidiscreteOT::SquareIntegral(double xmin, double xmax,
//                                       double ymin, double ymax,
//                                       double px,   double py)
// {
//  double rectInt = (0.5/3.0)*(ymax-ymin)*(std::pow(xmax-px,3.0)-std::pow(xmin-px,3.0))
//                 + (0.5/3.0)*(xmax-xmin)*(std::pow(ymax-py,3.0)-std::pow(ymin-py,3.0));
//
//  return rectInt;
// }

// double SemidiscreteOT::TriangleIntegral(double x1, double y1,
//                                        double x2, double y2,
//                                        double x3, double y3,
//                                        double px, double py)
// {
//  double triInt = (1.0/2.0)*std::pow(px, 2) - 1.0/3.0*px*x1 - 1.0/3.0*px*x2 - 1.0/3.0*px*x3 + (1.0/2.0)*std::pow(py, 2) - 1.0/3.0*py*y1 - 1.0/3.0*py*y2 - 1.0/3.0*py*y3 + (1.0/12.0)*std::pow(x1, 2) + (1.0/12.0)*x1*x2 + (1.0/12.0)*x1*x3 + (1.0/12.0)*std::pow(x2, 2) + (1.0/12.0)*x2*x3 + (1.0/12.0)*std::pow(x3, 2) + (1.0/12.0)*std::pow(y1, 2) + (1.0/12.0)*y1*y2 + (1.0/12.0)*y1*y3 + (1.0/12.0)*std::pow(y2, 2) + (1.0/12.0)*y2*y3 + (1.0/12.0)*std::pow(y3, 2);
//  triInt *= 0.5*((x2-x1)*(y3-y1) - (x3-x1)*(y2-y1));
//
//  return triInt;
// }

// template<typename ConjugateFunctionType>
// Eigen::VectorXd SemidiscreteOT<ConjugateFunctionType>::GetValidPrices(Eigen::VectorXd const& x0)
// {
//
// }

template<typename ConjugateFunctionType>
std::pair<Eigen::VectorXd, double> SemidiscreteOT<ConjugateFunctionType>::Solve(Eigen::VectorXd                  const& prices0,
                                                         OptionList                              options)
{
  SDOT_ASSERT(prices0.size()==discrPts.cols());
  const unsigned int dim = prices0.size();

  unsigned int printLevel = GetOpt("Print Level", options, 3);

  // Trust region approach with a double dogleg step
  double trustRadius = GetOpt("Trust Radius", options, 1.0);
  const unsigned int maxEvals = GetOpt("Max Steps", options, 100.0);

  const double xtol_abs = GetOpt("XTol Abs", options, 1e-13*std::sqrt(double(dim)));
  const double gtol_abs = GetOpt("GTol Abs", options, 2e-4*std::sqrt(double(dim)));
  const double ftol_abs = GetOpt("FTol Abs", options, 1e-11);

  const double acceptRatio = GetOpt("Accept Ratio", options, 0.1);//0.1;
  const double shrinkRatio = GetOpt("Shrink Ratio", options, 0.25);//0.1;
  const double growRatio = GetOpt("Grow Ratio", options, 0.75);
  const double growRate = GetOpt("Grow Rate", options, 2.0);
  const double shrinkRate = GetOpt("Shrink Rate", options, 0.25);
  const double maxRadius = GetOpt("Max Radius", options, 10);

  SDOT_ASSERT(shrinkRatio>=acceptRatio);

  double fval, newF, gradNorm, newGradNorm;
  Eigen::VectorXd grad, newGrad;
  Eigen::SparseMatrix<double> hess;

  Eigen::VectorXd x = prices0;
  Eigen::VectorXd newX(x);
  Eigen::VectorXd step = Eigen::VectorXd::Zero(dim);

  std::shared_ptr<LaguerreDiagram> newLagDiag;

  // Compute an initial gradient and Hessian
  lagDiag  = std::make_shared<LaguerreDiagram>(grid->xMin, grid->xMax, grid->yMin, grid->yMax, discrPts, x);
  SDOT_ASSERT(lagDiag!=nullptr);

  std::tie(fval, grad) = ComputeGradient(x, *lagDiag);
  hess = ComputeHessian(x, *lagDiag);

  fval *= -1.0;
  grad *= -1.0;
  hess *= -1.0;
  gradNorm = grad.norm();

  if(printLevel>0){
    std::cout << "Using NewtonTrust optimizer..." << std::endl;
    std::cout << "  Iteration, TrustRadius, Empty Cells,           rho, Dual Objective,        ||g||,   ||g||/sqrt(dim)" << std::endl;
  }

  // count the non-empty Cells
  int numEmpty = 0;
  for(int i=0; i<x.size(); ++i){
    if(lagDiag->GetCellVertices(i).cols()<3)
      ++numEmpty;
  }

  for(int it=0; it<maxEvals; ++it) {

    if((gradNorm < gtol_abs)&&(numEmpty==0)){
      if(printLevel>0){
        std::printf("Terminating because gradient norm (%4.2e) is smaller than gtol_abs (%4.2e).\n", gradNorm, gtol_abs);
      }
      optPrices = x;
      return std::make_pair(x,fval);
    }

    step = SolveSubProblem(fval, grad,  hess, trustRadius);

    newX = x+step;

    // Try constructing the new Laguerre diagram.
    newLagDiag  = std::make_shared<LaguerreDiagram>(grid->xMin, grid->xMax, grid->yMin, grid->yMax, discrPts, newX);

    std::tie(newF, newGrad) = ComputeGradient(newX, *newLagDiag);
    newF *= -1.0;
    newGrad *= -1.0;
    newGradNorm = newGrad.norm();

    // Use the quadratic submodel to predict the change in the objective
    double trueDelta = newF-fval;
    double modDelta = grad.dot(step) + 0.5*step.dot(hess.selfadjointView<Eigen::Lower>()*step);

    double rho = trueDelta/modDelta;
    //std::cout << "Model, Truth = " << modDelta << ", " << trueDelta << std::endl;
    //std::cout << "          step.dot(grad) = " << step.dot(grad) << std::endl;
    // std::cout << "          delta f = " << trueDelta << std::endl;
    // std::cout << "          modDelta = " << modDelta << std::endl;
    // std::cout << "          New prices = " << newX.transpose() << std::endl;
    //std::cout << "          rho = " << rho << std::endl;

    if(printLevel>0){
      std::printf("  %9d, %11.2e,  %10d,    % 5.3e, % 14.3e,    %5.3e,  %15.3e\n", it, trustRadius, numEmpty, rho, -1.0*fval, gradNorm, gradNorm/std::sqrt(double(dim)));
    }

    double stepNorm = step.norm();
    if((stepNorm < xtol_abs)&&(numEmpty==0)){
      if(printLevel>0){
        std::printf("Terminating because stepsize (%4.2e) is smaller than xtol_abs (%4.2e).\n", stepNorm, xtol_abs);
      }
      optPrices = newX;
      return std::make_pair(newX,newF);
    }

    // Update the position.  If the model is really bad, we'll just stay put
    if(rho>acceptRatio){

      if(((std::abs(fval-newF)<ftol_abs))&&(numEmpty==0)){
        if(printLevel>0){
          std::printf("Terminating because change in objective (%4.2e) is smaller than ftol_abs (%4.2e).\n", fval-newF, ftol_abs);
        }
        optPrices = newX;
        return std::make_pair(newX,newF);
      }

      x = newX;
      fval = newF;
      lagDiag = newLagDiag;
      grad = newGrad;
      gradNorm = newGradNorm;

      // Recompute the Hessian at the new point
      hess = ComputeHessian(x, *lagDiag);
      hess *= -1.0;

      // Update the number of empty cells in the diagram
      numEmpty = 0;
      for(int i=0; i<x.size(); ++i){
        if(lagDiag->GetCellVertices(i).cols()<3)
          ++numEmpty;
      }
    }

    // Update the trust region size
    if(rho<shrinkRatio){
      trustRadius = shrinkRate*trustRadius; // shrink trust region

      if(printLevel>1)
        std::cout << "            Shrinking trust region because of submodel accuracy." << std::endl;

    }else if((rho>growRatio)&&(std::abs(step.norm()-trustRadius)<1e-10)) {
      trustRadius = std::min(growRate*trustRadius, maxRadius);

      if(printLevel>1)
        std::cout << "            Growing trust region." << std::endl;

    }
  }

  if(printLevel>0){
    std::printf("Terminating because maximum number of iterations (%d) was reached.", maxEvals);
  }

  optPrices = x;
  return std::make_pair(x,fval);
}

template<typename ConjugateFunctionType>
Eigen::VectorXd SemidiscreteOT<ConjugateFunctionType>::SolveSubProblem(double obj,
                                                Eigen::Ref<const Eigen::VectorXd> const& grad,
                                                Eigen::Ref<const Eigen::SparseMatrix<double>> const& hess,
                                                double trustRadius) const
{
  const double trustTol = 1e-12;
  const unsigned int dim = grad.size();

  // Current estimate of the subproblem minimum
  Eigen::VectorXd z = Eigen::VectorXd::Zero(dim);

  // Related to the step direction
  Eigen::VectorXd r = grad;
  Eigen::VectorXd d = -r;

  // If the gradient is small enough where we're starting, then we're done
  if(r.norm()<trustTol){
    return z;
  }

  Eigen::VectorXd Bd; // the Hessian (B) applied to a vector d

  double alpha, beta, gradd, dBd, rr;

  for(int i=0; i<dim; ++i){
    Bd = hess.selfadjointView<Eigen::Lower>()*d;
    gradd = grad.dot(d);
    dBd = d.dot(Bd);
    rr = r.squaredNorm();

    // If the Hessian isn't positive definite in this direction, we can go all
    // the way to the trust region boundary
    if(dBd<=0){
      // do something

      double dz = d.dot(z);
      double dd = d.squaredNorm();
      double zz = z.squaredNorm();
      double r2 = trustRadius*trustRadius;

      double tau1 = (-dz + sqrt(dz*dz - dd*(zz-r2)))/dd;
      double tau2 = (-dz - sqrt(dz*dz - dd*(zz-r2)))/dd;

      double zBd = z.dot(Bd);
      double mval1 = tau1*gradd + tau1*zBd + tau1*tau1*dBd;
      double mval2 = tau2*gradd + tau2*zBd + tau2*tau2*dBd;

      return (mval1<mval2) ? (z+tau1*d) : (z+tau2*d);
    }

    alpha = rr / dBd;
    Eigen::VectorXd newZ = z + alpha * d;

    if(newZ.norm()>trustRadius){

      double dz = d.dot(z);
      double dd = d.squaredNorm();
      double zz = z.squaredNorm();
      double r2 = trustRadius*trustRadius;

      double tau = (-dz + sqrt(dz*dz - dd*(zz-r2)))/dd;
      return z + tau*d;
    }

    z = newZ;

    r += alpha*Bd;

    if(r.norm()<trustTol){
      return z;
    }

    beta = r.squaredNorm() / rr;
    d = (-r + beta*d).eval();
  }

  return z;
}

template<typename ConjugateFunctionType>
std::shared_ptr<LaguerreDiagram> SemidiscreteOT<ConjugateFunctionType>::BuildCentroidal(std::shared_ptr<Distribution2d> const& dist,
                                                                 Eigen::Matrix2Xd                const& initialPoints,
                                                                 Eigen::VectorXd                 const& pointProbs,
                                                                 OptionList                             opts)
{

  unsigned int maxIts =  GetOpt("Lloyd Steps", opts, 100);
  double xtol = GetOpt("Lloyd Tol", opts, 1e-8);
  double gtol = xtol;

  auto optIt = opts.find("Print Level");
  if(optIt == opts.end())
    opts["Print Level"] = 0;

  const unsigned int numPts = pointProbs.size();
  SDOT_ASSERT(numPts==initialPoints.cols());

  double resid = xtol + 1.0;

  Eigen::MatrixXd newPts;
  Eigen::VectorXd prices = Eigen::VectorXd::Ones(numPts);
  double dualObj;
  Eigen::MatrixXd pts = initialPoints;
  std::shared_ptr<SemidiscreteOT> ot = std::make_shared<SemidiscreteOT>(dist, initialPoints, pointProbs, GetOpt("Penalty", opts, 1.0));

  std::cout << "Computing constrained centroidal diagram..." << std::endl;
  std::cout << "  Iteration,  ||g||,   max(dx)" << std::endl;

  for(unsigned int i=0; i<maxIts; ++i){
    SDOT_ASSERT(ot);

    std::tie(prices, dualObj) = ot->Solve(prices, opts);

    Eigen::MatrixXd pointGrad = ot->PointGradient();
    Eigen::Map<Eigen::VectorXd> pointGradVec(pointGrad.data(),2*pointGrad.cols());

    double stepSize = 1.0;
    Eigen::Matrix2Xd dir = -pointGrad.array()*(ot->LloydPointHessian()+1e-12*Eigen::MatrixXd::Ones(2,numPts)).array().inverse();

    newPts = pts + stepSize*dir;

    // Backtrack until all the points are  in the domain
    while((newPts.row(0).minCoeff()<ot->grid->xMin)||(newPts.row(0).maxCoeff()>ot->grid->xMax)||(newPts.row(1).minCoeff()<ot->grid->yMin)||(newPts.row(1).maxCoeff()>ot->grid->yMax)){
      stepSize *= 0.75;
      newPts = pts + stepSize*dir;
    }

    resid = (newPts - pts).cwiseAbs().maxCoeff();;
    std::printf("  %9d, %5.3e, %5.3e\n", i, pointGrad.norm(), resid);

    if(resid<xtol){
      std::cout << "Converged due to small stepsize." << std::endl;
      return ot->Diagram();
    }

    if((pointGradVec.norm()/(2*numPts))<gtol){
      std::cout << "Converged due to small gradient." << std::endl;
      return ot->Diagram();
    }

    pts = newPts;
    ot->SetPoints(pts);
  }

  std::cout << "WARNING: Did not converge to constrained centroidal diagram." << std::endl;
  return ot->Diagram();
}

template<typename ConjugateFunctionType>
std::shared_ptr<LaguerreDiagram> SemidiscreteOT<ConjugateFunctionType>::BuildCentroidal(std::shared_ptr<Distribution2d> const& dist,
                                                                 Eigen::VectorXd                 const& probs,
                                                                 OptionList                             opts)
{
  BoundingBox bbox(dist->Grid()->xMin, dist->Grid()->xMax, dist->Grid()->yMin, dist->Grid()->yMax);
  Eigen::Matrix2Xd initialPts = LaguerreDiagram::LatinHypercubeSample(bbox, probs.size());
  return BuildCentroidal(dist, initialPts, probs, opts);
}

template<typename ConjugateFunctionType>
std::shared_ptr<LaguerreDiagram> SemidiscreteOT<ConjugateFunctionType>::BuildCentroidal(std::shared_ptr<Distribution2d> const& dist,
                                                                 unsigned int                           numPts,
                                                                 OptionList                             opts)
{
  Eigen::VectorXd probs = (1.0/numPts)*Eigen::VectorXd::Ones(numPts);
  return BuildCentroidal(dist, probs, opts);
}

template<typename ConjugateFunctionType>
Eigen::Matrix2Xd SemidiscreteOT<ConjugateFunctionType>::MarginalCentroids(Eigen::VectorXd const& prices,
                                                                          LaguerreDiagram const& lagDiag) const
{
  int numCells =  lagDiag.NumCells();
  Eigen::Matrix2Xd centroids(2,numCells);

  // Loop over all ofthe cells
  for(int i=0; i<numCells; ++i){

    auto triFunc = [&](Eigen::Vector2d const& pt1, Eigen::Vector2d const& pt2, Eigen::Vector2d const& pt3)
      {
        return ConjugateFunctionType::TriangularIntegralMarginalCentroid(prices(i), discrPts.col(i), pt1,pt2,pt3,unbalancedPenalty);
      };
    auto rectFunc = [&](Eigen::Vector2d const& pt1, Eigen::Vector2d const& pt2)
      {
        return ConjugateFunctionType::RectangularIntegralMarginalCentroid(prices(i), discrPts.col(i), pt1,pt2,unbalancedPenalty);
      };

    centroids.col(i) = lagDiag.IntegrateOverCell(i, triFunc, rectFunc, dist);

    auto triFunc2 = [&](Eigen::Vector2d const& pt1, Eigen::Vector2d const& pt2, Eigen::Vector2d const& pt3)
      {
        return ConjugateFunctionType::TriangularIntegralMarginalMass(prices(i), discrPts.col(i), pt1,pt2,pt3,unbalancedPenalty);
      };
    auto rectFunc2 = [&](Eigen::Vector2d const& pt1, Eigen::Vector2d const& pt2)
      {
        return ConjugateFunctionType::RectangularIntegralMarginalMass(prices(i), discrPts.col(i), pt1,pt2,unbalancedPenalty);
      };

    centroids.col(i) /= lagDiag.IntegrateOverCell(i, triFunc2, rectFunc2, dist);
  }

  return centroids;
}

template<>
Eigen::Matrix2Xd SemidiscreteOT<sdot::distances::Wasserstein2>::MarginalCentroids(Eigen::VectorXd const& prices,
                                                                                LaguerreDiagram const& lagDiag) const
{
  return lagDiag.Centroids(dist);
}

namespace sdot{
  template class SemidiscreteOT<sdot::distances::Wasserstein2>;
  template class SemidiscreteOT<sdot::distances::QuadraticRegularization>;
  template class SemidiscreteOT<sdot::distances::GHK>;
}


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

/*
 Copyright (C) 2015 Peter Caspers

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

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

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

#include <ql/experimental/models/lgm1.hpp>
#include <ql/quotes/simplequote.hpp>

#include <boost/make_shared.hpp>

namespace QuantLib {

Lgm1::Lgm1(const Handle<YieldTermStructure> &yts,
           const std::vector<Date> &volstepdates,
           const std::vector<Real> &alpha, const Real &kappa)
    : Lgm(yts), CalibratedModel(2), volstepdates_(volstepdates),
      volsteptimes_(volstepdates_.size()),
      volsteptimesArray_(volstepdates_.size()), alpha_(arguments_[0]),
      kappa_(arguments_[1]) {
    alphaQuotes_.resize(alpha.size());
    for (Size i = 0; i < alpha.size(); ++i) {
        alphaQuotes_[i] =
            Handle<Quote>(boost::make_shared<SimpleQuote>(alpha[i]));
    }

    kappaQuote_ = Handle<Quote>(boost::make_shared<SimpleQuote>(kappa));

    initialize();
}

Lgm1::Lgm1(const Handle<YieldTermStructure> &yts,
           const std::vector<Date> &volstepdates,
           const std::vector<Handle<Quote> > &alpha, const Handle<Quote> &kappa)
    : Lgm(yts), CalibratedModel(2), volstepdates_(volstepdates),
      volsteptimes_(volstepdates_.size()),
      volsteptimesArray_(volstepdates_.size()), alphaQuotes_(alpha),
      kappaQuote_(kappa), alpha_(arguments_[0]), kappa_(arguments_[1]) {

    initialize();
}

void Lgm1::updateTimes() const {
    volsteptimes_.clear();
    int j = 0;
    for (std::vector<Date>::const_iterator i = volstepdates_.begin();
         i != volstepdates_.end(); ++i, ++j) {
        volsteptimes_.push_back(termStructure()->timeFromReference(*i));
        volsteptimesArray_[j] = volsteptimes_[j];
        if (j == 0)
            QL_REQUIRE(volsteptimes_[0] > 0.0, "volsteptimes must be positive ("
                                                   << volsteptimes_[0] << ")");
        else
            QL_REQUIRE(volsteptimes_[j] > volsteptimes_[j - 1],
                       "volsteptimes must be strictly increasing ("
                           << volsteptimes_[j - 1] << "@" << (j - 1) << ", "
                           << volsteptimes_[j] << "@" << j << ")");
    }
}

void Lgm1::updateAlpha() {
    for (Size i = 0; i < alpha_.size(); ++i) {
        alpha_.setParam(i, alphaQuotes_[i]->value());
    }
    update();
}

void Lgm1::updateKappa() {
    kappa_.setParam(0, kappaQuote_->value());
    update();
}

void Lgm1::initialize() {
    QL_REQUIRE(volstepdates_.size() + 1 == alphaQuotes_.size(),
               "alphas (" << alphaQuotes_.size() << ") and step dates ("
                          << volstepdates_.size() << ") inconsistent.");
    updateTimes();
    alpha_ = PiecewiseConstantParameter(volsteptimes_, NoConstraint());
    kappa_ = ConstantParameter(kappaQuote_->value(), NoConstraint());
    updateAlpha();
    alphaObserver_ = boost::make_shared<AlphaObserver>(this);
    kappaObserver_ = boost::make_shared<KappaObserver>(this);
    for (Size i = 0; i < alpha_.size(); ++i)
        alphaObserver_->registerWith(alphaQuotes_[i]);
    kappaObserver_->registerWith(kappaQuote_);
    setParametrization(
        boost::make_shared<detail::LgmPiecewiseAlphaConstantKappa>(
            volsteptimesArray_, alpha_.params(), kappa_.params()));
    stateProcess_ = boost::make_shared<
        LgmStateProcess<detail::LgmPiecewiseAlphaConstantKappa> >(
        parametrization());
    registerWith(stateProcess_);
    parametrization()->update();
}

} // namespace QuantLib


#ifndef STAN_MATH_PRIM_FUN_GAMMA_Q_HPP
#define STAN_MATH_PRIM_FUN_GAMMA_Q_HPP

#include <stan/math/prim/meta.hpp>
#include <stan/math/prim/functor/apply_scalar_binary.hpp>
#include <boost/math/special_functions/gamma.hpp>

namespace stan {
namespace math {

/**
 *
   \f[
   \mbox{gamma\_q}(a, z) =
   \begin{cases}
     \textrm{error} & \mbox{if } a\leq 0 \textrm{ or } z < 0\\
     Q(a, z) & \mbox{if } a > 0, z \geq 0 \\[6pt]
     \textrm{NaN} & \mbox{if } a = \textrm{NaN or } z = \textrm{NaN}
   \end{cases}
   \f]

   \f[
   \frac{\partial\, \mbox{gamma\_q}(a, z)}{\partial a} =
   \begin{cases}
     \textrm{error} & \mbox{if } a\leq 0 \textrm{ or } z < 0\\
     \frac{\partial\, Q(a, z)}{\partial a} & \mbox{if } a > 0, z \geq 0 \\[6pt]
     \textrm{NaN} & \mbox{if } a = \textrm{NaN or } z = \textrm{NaN}
   \end{cases}
   \f]

   \f[
   \frac{\partial\, \mbox{gamma\_q}(a, z)}{\partial z} =
   \begin{cases}
     \textrm{error} & \mbox{if } a\leq 0 \textrm{ or } z < 0\\
     \frac{\partial\, Q(a, z)}{\partial z} & \mbox{if } a > 0, z \geq 0 \\[6pt]
     \textrm{NaN} & \mbox{if } a = \textrm{NaN or } z = \textrm{NaN}
   \end{cases}
   \f]

   \f[
   Q(a, z)=\frac{1}{\Gamma(a)}\int_z^\infty t^{a-1}e^{-t}dt
   \f]

   \f[
   \frac{\partial \, Q(a, z)}{\partial a} =
 -\frac{\Psi(a)}{\Gamma^2(a)}\int_z^\infty t^{a-1}e^{-t}dt
   + \frac{1}{\Gamma(a)}\int_z^\infty (a-1)t^{a-2}e^{-t}dt
   \f]

   \f[
   \frac{\partial \, Q(a, z)}{\partial z} = -\frac{z^{a-1}e^{-z}}{\Gamma(a)}
   \f]
   * @throws domain_error if x is at pole
 */
inline double gamma_q(double x, double a) { return boost::math::gamma_q(x, a); }

/**
 * Enables the vectorised application of the gamma_q function,
 * when the first and/or second arguments are containers.
 *
 * @tparam T1 type of first input
 * @tparam T2 type of second input
 * @param a First input
 * @param b Second input
 * @return gamma_q function applied to the two inputs.
 */
template <typename T1, typename T2, require_any_container_t<T1, T2>* = nullptr>
inline auto gamma_q(const T1& a, const T2& b) {
  return apply_scalar_binary(
      a, b, [&](const auto& c, const auto& d) { return gamma_q(c, d); });
}

}  // namespace math
}  // namespace stan
#endif


#include <NTL/ZZ_pX.h>
#include <NTL/FFT.h>

#include <cstdio>

NTL_CLIENT

double clean_data(double *t)
{
   double x, y, z;
   long i, ix, iy, n;

   x = t[0]; ix = 0;
   y = t[0]; iy = 0;

   for (i = 1; i < 5; i++) {
      if (t[i] < x) {
         x = t[i];
         ix = i;
      }
      if (t[i] > y) {
         y = t[i];
         iy = i;
      }
   }

   z = 0; n = 0;
   for (i = 0; i < 5; i++) {
      if (i != ix && i != iy) z+= t[i], n++;
   }

   z = z/n;

   return z;
}

void print_flag()
{

#if defined(NTL_FFT_LAZYMUL)
printf("FFT_LAZYMUL ");
#endif

#if defined(NTL_SPMM_ULL)
printf("SPMM_ULL ");
#endif

#if defined(NTL_AVOID_BRANCHING)
printf("AVOID_BRANCHING ");
#endif

#if defined(NTL_FFT_BIGTAB)
printf("FFT_BIGTAB ");
#endif

printf("\n");

}

int main()
{

   SetSeed(ZZ(0));

   long n, k;

   n = 200;
   k = 10*NTL_ZZ_NBITS;

   ZZ p;

   RandomLen(p, k);
   if (!IsOdd(p)) p++;

   ZZ_p::init(p);         // initialization

   ZZ_pX f, g, h, r1, r2, r3;

   random(g, n);    // g = random polynomial of degree < n
   random(h, n);    // h =             "   "
   random(f, n);    // f =             "   "

   SetCoeff(f, n);  // Sets coefficient of X^n to 1

   // For doing arithmetic mod f quickly, one must pre-compute
   // some information.

   ZZ_pXModulus F;
   build(F, f);

   PlainMul(r1, g, h);  // this uses classical arithmetic
   PlainRem(r1, r1, f);

   MulMod(r2, g, h, F);  // this uses the FFT

   MulMod(r3, g, h, f);  // uses FFT, but slower

   // compare the results...

   if (r1 != r2) {
      printf("999999999999999 ");
      print_flag();
      return 0;
   }
   else if (r1 != r3) {
      printf("999999999999999 ");
      print_flag();
      return 0;
   }

   double t;
   long i, j;
   long iter;

   const int nprimes = 30;
   const long L = 12;
   const long N = 1L << L;
   long r;

   for (r = 0; r < nprimes; r++) UseFFTPrime(r);

   vec_long A1[nprimes], A2[nprimes];
   vec_long B1[nprimes], B2[nprimes];

   for (r = 0; r < nprimes; r++) {
      A1[r].SetLength(N);
      A2[r].SetLength(N);
      B1[r].SetLength(N);
      B2[r].SetLength(N);

      for (i = 0; i < N; i++) {
         A1[r][i] = RandomBnd(GetFFTPrime(r));
         A2[r][i] = RandomBnd(GetFFTPrime(r));
      }
   }

   for (r = 0; r < nprimes; r++) {
      long *A1p = A1[r].elts();
      long *A2p = A2[r].elts();
      long *B1p = B1[r].elts();
      long *B2p = B2[r].elts();
      long q = GetFFTPrime(r);
      mulmod_t qinv = GetFFTPrimeInv(r);

      FFTFwd(B1p, A1p, L, r);
      FFTFwd(B2p, A2p, L, r);
      for (i = 0; i < N; i++) B1p[i] = NormalizedMulMod(B1p[i], B2p[i], q, qinv);
      FFTRev1(B1p, B1p, L, r);
   }

   iter = 1;

   do {
     t = GetTime();
     for (j = 0; j < iter; j++) {
        for (r = 0; r < nprimes; r++) {
           long *A1p = A1[r].elts();
           long *A2p = A2[r].elts();
           long *B1p = B1[r].elts();
           long *B2p = B2[r].elts();
           long q = GetFFTPrime(r);
           mulmod_t qinv = GetFFTPrimeInv(r);

           FFTFwd(B1p, A1p, L, r);
           FFTFwd(B2p, A2p, L, r);
           for (i = 0; i < N; i++) B1p[i] = NormalizedMulMod(B1p[i], B2p[i], q, qinv);
           FFTRev1(B1p, B1p, L, r);
        }
     }
     t = GetTime() - t;
     iter = 2*iter;
   } while(t < 1);

   iter = iter/2;

   iter = long((3/t)*iter) + 1;

   double tvec[5];
   long w;

   for (w = 0; w < 5; w++) {
     t = GetTime();
     for (j = 0; j < iter; j++) {
        for (r = 0; r < nprimes; r++) {
           long *A1p = A1[r].elts();
           long *A2p = A2[r].elts();
           long *B1p = B1[r].elts();
           long *B2p = B2[r].elts();
           long q = GetFFTPrime(r);
           mulmod_t qinv = GetFFTPrimeInv(r);

           FFTFwd(B1p, A1p, L, r);
           FFTFwd(B2p, A2p, L, r);
           for (i = 0; i < N; i++) B1p[i] = NormalizedMulMod(B1p[i], B2p[i], q, qinv);
           FFTRev1(B1p, B1p, L, r);
        }
     }
     t = GetTime() - t;
     tvec[w] = t;
   }

   t = clean_data(tvec);

   t = floor((t/iter)*1e13);

   if (t < 0 || t >= 1e15)
      printf("999999999999999 ");
   else
      printf("%015.0f ", t);

   printf(" [%ld] ", iter);

   print_flag();

   return 0;
}


//=======================================================================
// Copyright (c) 2013 Maciej Andrejczuk
//               2014 Piotr Wygocki
//
// Distributed under the Boost Software License, Version 1.0. (See
// accompanying file LICENSE_1_0.txt or copy at
// http://www.boost.org/LICENSE_1_0.txt)
//=======================================================================
/**
 * @file steiner_component.hpp
 * @brief
 * @author Maciej Andrejczuk, Piotr Wygocki
 * @version 1.0
 * @date 2013-08-01
 */
#ifndef PAAL_STEINER_COMPONENT_HPP
#define PAAL_STEINER_COMPONENT_HPP

#include "paal/data_structures/metric/basic_metrics.hpp"
#include "paal/steiner_tree/dreyfus_wagner.hpp"
#include "paal/utils/irange.hpp"

#include <boost/range/algorithm/transform.hpp>
#include <boost/range/join.hpp>

#include <iosfwd>
#include <set>

namespace paal {
namespace ir {

/**
 * @class steiner_component
 * @brief Class represents k-components of Steiner Tree.
 * Component is a subtree whose terminals coincide with leaves.
 */
template <typename Vertex, typename Dist>
class steiner_component {
public:
    using Edge = typename std::pair<Vertex, Vertex>;
    using Vertices = typename std::vector<Vertex>;

    ///constructor
    template<typename Metric, typename Terminals>
    steiner_component(const Metric & cost_map, Vertices terminals, const Terminals& steiner_vertices) :
        m_terminals(std::move(terminals)), m_size(m_terminals.size()) {

        auto all_elements = boost::join(m_terminals, steiner_vertices);
        data_structures::array_metric<typename data_structures::metric_traits<Metric>::DistanceType>
            fast_metric(cost_map, all_elements);
        auto term_nr = boost::distance(m_terminals);
        auto all_elements_nr = boost::distance(all_elements);
        auto dw = paal::make_dreyfus_wagner(fast_metric,
                    irange(term_nr),
                    irange(int(term_nr), int(all_elements_nr)));
        dw.solve();
        m_cost = dw.get_cost();
        auto &steiner = dw.get_steiner_elements();
        m_steiner_elements.resize(steiner.size());
        auto id_to_elem = [&](int i){
                if(i < term_nr) {
                    return m_terminals[i];
                } else {
                    return steiner_vertices[i - term_nr];
                }
        };
        boost::transform(steiner, m_steiner_elements.begin(), id_to_elem);
        m_edges.resize(dw.get_edges().size());
        boost::transform(dw.get_edges(), m_edges.begin(), [=](std::pair<int, int> e) {
            return std::make_pair(id_to_elem(e.first), id_to_elem(e.second));
        });
    }

    /**
     * @brief Each component has versions, where sink is chosen from its
     * terminals
     */
    Vertex get_sink(int version) const {
        assert(version < count_terminals());
        return m_terminals[version];
    }

    /**
     * Returns vector composed of component's terminals.
     */
    const Vertices &get_terminals() const { return m_terminals; }

    /**
     * Returns vector composed of component's nonterminals, i.e. Steiner
     * elements.
     */
    const Vertices &get_steiner_elements() const {
        return m_steiner_elements;
    }

    /**
     * Returns edges spanning the component.
     */
    const std::vector<Edge> &get_edges() const { return m_edges; }

    /**
     * Returns degree of component, i.e. number of terminals.
     */
    int count_terminals() const { return m_size; }

    /**
     * Returns minimal cost of spanning a component.
     */
    Dist get_cost() const { return m_cost; }

    /**
     * Prints the component.
     */
    friend std::ostream &operator<<(std::ostream &stream,
                                    const steiner_component &component) {
        for (int i = 0; i < component.m_size; i++) {
            stream << component.m_terminals[i] << " ";
        }
        stream << ": ";
        for (auto edge : component.m_edges) {
            stream << "(" << edge.first << "," << edge.second << ") ";
        }
        stream << component.m_cost;
        return stream;
    }

  private:
    const Vertices m_terminals; // terminals of the component
    int m_size;                           // m_terminals.size()
    Dist m_cost; // minimal cost of spanning the component
    Vertices m_steiner_elements; // non-terminals selected for
                                            // spanning tree
    std::vector<Edge> m_edges;              // edges spanning the component
};

} // ir
} // paal

#endif // PAAL_STEINER_COMPONENT_HPP


/**
 * @author pkambadu
 */

#include <iostream>
#include <fstream>
#include <functional>
#include <cstring>
#include <cstdlib>
#include <vector>
#include <utility>
#include <cstdio>
#include <map>

#if defined(_OPENMP)
#include <omp.h>
#endif

#include <boost/math/distributions/chi_squared.hpp>
#include <boost/filesystem.hpp>

#include "dyck_options.hpp"
#include "generate_Galton_Watson.hpp"
#include "dyck_path.hpp"

typedef generate_Galton_Watson_t::graph_type graph_type;
typedef generate_Galton_Watson_t::vertex_index_map_t vertex_index_map_t;
typedef generate_Galton_Watson_t::vertex_name_map_t vertex_name_map_t;
typedef generate_Galton_Watson_t::edge_weight_map_t edge_weight_map_t;
typedef generate_Galton_Watson_t::vertex_distance_map_t vertex_weight_map_t;
typedef boost::graph_traits<graph_type>::vertex_iterator vertex_iter_t;

static const double PI =
    3.1415926535897932384626433832795028841971693993751058209749445923078164;

/**
 * Thin wrapper around boost::uniform<> so that it conforms to random_shuffle.
 */
template <typename EngineType>
struct random_functor_t : std::unary_function<int, int> {
  typedef EngineType engine_type;
  typedef boost::uniform_int<int> uni_int_type;
  engine_type engine;

  random_functor_t (engine_type engine) : engine(engine) {}

  int operator()(const int& ULIMIT) {
    uni_int_type dist (0, ULIMIT-1);
    return dist(engine);
  }
};

/**
 * Print the tree as a DOT file
 */
template <typename Graph>
void print_tree (Graph& G, const dyck_options_t& options) {
#pragma omp critical
  if (options.print_tree) {
    /** Set the properties that you need to extract from the graph */
    boost::dynamic_properties dp;
    vertex_index_map_t index_map = boost::get(boost::vertex_index, G);
    vertex_name_map_t name_map = boost::get(boost::vertex_name, G);
    vertex_weight_map_t weight_map = boost::get(boost::vertex_distance, G);

    dp.property("index", index_map);
    dp.property("name", name_map);
    dp.property("weight", weight_map);

    /** Create the output */
    std::ostream* output;
    if (0==strcmp("stdout", options.dot_out.c_str())) output = &(std::cout);
    else output = new std::ofstream(options.dot_out.c_str());

    boost::write_graphviz_dp(*output, G, dp, "index");

    if (0!=strcmp("stdout", options.dot_out.c_str())) delete output;
  }
}

/**
 * A function that lists all the DOT files in a directory.
 */
std::vector<std::string> get_dot_files (const std::string& dir_name) {
  std::vector<std::string> dot_file_list;

  if (boost::filesystem::exists(dir_name) &&
      boost::filesystem::is_directory(dir_name)) {
    boost::filesystem::directory_iterator end_iter;
    for(boost::filesystem::directory_iterator dir_iter(dir_name);
        dir_iter != end_iter;
        ++dir_iter) {
      if (boost::filesystem::is_regular_file(dir_iter->status())) {
        const boost::filesystem::path current_file = dir_iter->path();
        const boost::filesystem::path current_file_ext =
                                             current_file.extension();
        if (current_file_ext.native() == ".dot") {
          dot_file_list.push_back(current_file.native());
        }
      }
    }
  }

  return dot_file_list;
}

/**
 * The one annoying problem is that OutputIterator's do not define
 * a value type, so it may not always work! Hrmp!
 */
int generate_seed () {
  int r_int;
#pragma omp critical
  {
    std::ifstream rf("/dev/random", std::ios::binary);
    if(rf.is_open())rf.read(reinterpret_cast<char*>(&r_int),sizeof(r_int));
    else r_int = std::time(NULL);
  }

  return r_int;
}

graph_type generate_one_graph (const dyck_options_t& options,
                               int& gen_seed) {
  graph_type G;

  gen_seed = (0>options.seed)?generate_seed():options.seed;
  if (0 == strcmp("Poisson", options.gen_method.c_str())) {
    G = generate_Galton_Watson_t::generate_Poisson(options.n,
                                                   gen_seed,
                                                   options.lambda,
                                                   options.use_random_weights,
                                                   options.verbosity,
                                                   options.use_stupid_tree_gen);
  } else if (0 == strcmp("Binomial", options.gen_method.c_str())) {
    G = generate_Galton_Watson_t::generate_Binomial (options.n,
                                                     gen_seed,
                                                     options.k,
                                                     options.p,
                                                     options.use_random_weights,
                                                     options.verbosity,
                                                 options.use_stupid_tree_gen);
  } else if (0 == strcmp("Geometric", options.gen_method.c_str())) {
    G = generate_Galton_Watson_t::generate_Geometric(options.n,
                                                     gen_seed,
                                                     options.p,
                                                     options.use_random_weights,
                                                     options.verbosity,
                                                 options.use_stupid_tree_gen);
  } else if (0 == strcmp("Binary-0-2", options.gen_method.c_str())) {
    if (false == (options.n & 0x1)) {
      std::cerr << "Binary-0-2 can only generate odd number of nodes\
 due to a shortcoming of Devroye's algorithm. Sorry" << std::endl;
      exit(-1);
    }
    G = generate_Galton_Watson_t::generate_Binary_0_2
                                              (options.n,
                                               gen_seed,
                                               options.p,
                                               options.use_random_weights,
                                               options.verbosity,
                                               options.use_stupid_tree_gen);
  } else if (0 == strcmp("Binary-0-1-2", options.gen_method.c_str())) {
    G = generate_Galton_Watson_t::generate_Binary_0_1_2
                                              (options.n,
                                               gen_seed,
                                               options.use_random_weights,
                                               options.verbosity,
                                               options.use_stupid_tree_gen);
  } else if (0 == strcmp("Graphviz", options.gen_method.c_str())) {
    /** Get the input file */
    std::ifstream input_dot_file(options.dot_in.c_str());

    /** Set the properties that you need to extract from the graph */
    boost::dynamic_properties dp;
    vertex_index_map_t index_map = boost::get(boost::vertex_index, G);
    vertex_name_map_t name_map = boost::get(boost::vertex_name, G);
    vertex_weight_map_t weight_map = boost::get(boost::vertex_distance, G);

    dp.property("index", index_map);
    dp.property("name", name_map);
    dp.property("weight", weight_map);

    /** Read the graph --- note that indices are relabeled, which is fine */
    boost::read_graphviz(input_dot_file, G, dp, "name");

    /** Label it */
    bfs_label(G);

  } else {
    std::cout << "Chosen method of graph generation not supported"
              << std::endl;
  }
  return G;
}

template <typename NameToIndexMap,
          typename WeightMap>
double path_length (const std::string& parent,
                    const std::string& child,
                    const NameToIndexMap& name_map,
                    const WeightMap& weight_map,
                    const dyck_options_t& options) {

  double total_path_length = 0.0;

  if (0<options.verbosity) {
    std::cout << "Path length " << name_map.find(parent)->second << " ---> "
              << name_map.find(child)->second << std::endl;
  }

  /** Get the first place where things mismatch */
  std::string::const_iterator position =
    std::mismatch (parent.begin(), parent.end(), child.begin()).second;

  /** If they don't, we get to go one by one */
  while (position != child.end()) {
    const std::string current_vertex(child.begin(), position+1);
    const int current_vertex_index = name_map.find(current_vertex)->second;
    const double weight = boost::get(weight_map, current_vertex_index);
    if (0<options.verbosity) {
      std::cout << "Adding " << current_vertex_index << " (code:"
              << current_vertex << ") "
              << " of weight " << weight << std::endl;
    }
    total_path_length += weight;
    ++position;
  }

  return total_path_length;
}

/** total path length and generate trees by cutting things off */

std::pair<int,double>
run_one_lca_experiment(dyck_options_t& options,
                       int& gen_seed,
                       int& lca_seed,
                       int& num_nodes_in_G) {
  graph_type G;
  vertex_name_map_t name_map;
  vertex_weight_map_t weight_map;

  /** Generate the seeds if needed */
  lca_seed = (0>options.lca_seed)?generate_seed():options.lca_seed;
  gen_seed = (0>options.seed)?generate_seed():options.seed;

  if (options.verbosity) std::cout << "Generating graph .... " << std::flush;

  do {
    /** Generate a graph */
    G = generate_one_graph (options, gen_seed);

    /** Get the number of nodes in the graph */
    vertex_iter_t v_begin, v_end;
    boost::tie(v_begin, v_end) = boost::vertices(G);
    num_nodes_in_G = std::distance(v_begin, v_end);

    if (false==options.dot_in.empty()) {
      options.n = num_nodes_in_G;
    }
  } while (num_nodes_in_G<options.n);
  if (options.verbosity) std::cout << "DONE" << std::endl;

  name_map = boost::get(boost::vertex_name, G);
  weight_map = boost::get(boost::vertex_distance, G);

  int num_lca_samples = num_nodes_in_G*options.percent_lca_samples;
  std::vector<int> lca_vertices(num_lca_samples);
  const std::string file_name =
           boost::filesystem::path(options.dot_in).filename().string();

  /**
   * We want to generate some vertices from the given set of vertices
   * randomly and form the LCA tree. The thing wrong here is that my
   * compiler doesn't implement std::random_sample and std::iota. Hence
   * that long stupid way to sample elements from the vertex.
   *
   */
  std::vector<int> all_vertices(num_nodes_in_G-1);
  for (size_t i=0; i<all_vertices.size(); ++i) all_vertices[i]=(i+1);
  boost::mt19937 engine(lca_seed);
  random_functor_t<boost::mt19937> shuffle_prng (engine);
  std::random_shuffle (all_vertices.begin(), all_vertices.end(), shuffle_prng);

  if (options.sample_leaves) {
    /**
     * We want to only sample leaf nodes. The logic behind this is that we want
     * to only sample "original" data points. In case of Arvind's data, the
     * intermediate nodes are formed by hierarchical clustering and hence are
     * not original. There are two things to do:
     *
     * (1) Iterate through the shuffled vertices and
     *    -- copy those vertices that are leaves (up to num_lca_samples)
     *    -- count the number of leaves
     * (2) If num_leaves>num_lca_samples, we are golden. However, if this is
     *     not the case (num_lca_samples>=num_leaves), then we want to cut
     *     down on the number of samples.
     */
    int num_leaves = 0;
    std::vector<int>::iterator lca_vertices_iter = lca_vertices.begin();
    for (size_t i=0; i<all_vertices.size(); ++i) {
      const int vertex_id = all_vertices[i];
      if (0 == boost::out_degree(vertex_id, G)) {
        if (num_leaves<num_lca_samples) {
          *lca_vertices_iter = vertex_id;
          ++lca_vertices_iter;
        }
        ++num_leaves;
      }
    }

    if (num_lca_samples>=num_leaves) {
      std::cerr << file_name << ": Sampling " << num_lca_samples
                << " from " << num_leaves;
      num_lca_samples = 0.5*num_leaves;
      std::cerr << ". Changing to sample " << num_lca_samples << std::endl;
      if (0==num_lca_samples) {
        std::cerr << "Error: sampling " << num_lca_samples << std::endl;
        exit(-3);
      }
      lca_vertices.resize(num_lca_samples);
    }
  } else {
    /**
     * Approach 1:
     * We don't care whether we are sampling leaf or non-leaf nodes. As long as
     * it's a non-root node, we are golden.  All we are doing here is creating
     * an array of elements 1...(n-1), randomly shuffling it and then picking
     * the top num_lca_samples.
     */
    std::copy (all_vertices.begin(),
               all_vertices.begin() + num_lca_samples,
               lca_vertices.begin());
  }

  /**
   * Now, build the new tree with intermediate vertices and figure out the
   * number of edges in the new tree using the LCA algorithm to create
   * intermediate nodes.
   */
  if (0<options.verbosity) {
    std::cout << "Sampled LCA vertices are: " << std::endl;
    for (size_t i=0; i<lca_vertices.size(); ++i)
      std::cout << lca_vertices[i] << " (code:"
              << boost::get(name_map, lca_vertices[i]) << ")" << std::endl;
  }

  std::map<std::string, int> all_vertices_name_map;

  /** Create a name map */
  for (size_t i=0; i<num_nodes_in_G; ++i)
    all_vertices_name_map[boost::get(name_map, i)] = i;

  /** 2. Join the vertices one by one to get the LCA tree */
  std::map<std::string,int> new_lca_vertices;
  typedef std::map<std::string, int>::const_iterator MapIterator;
  new_lca_vertices[boost::get(name_map, 0)] = 0;

  std::set<std::pair<int,int> > added_edges_set;
  double total_path_length = 0.0;
  int total_num_edges = 0;

  for (size_t i=0; i<(lca_vertices.size()-1); ++i) {
    for (size_t j=(i+1); j<lca_vertices.size(); ++j) {
      /** Figure out the vertices we are dealing with */
      const int me = lca_vertices[i];
      const std::string my_name = boost::get(name_map, me);
      const int you = lca_vertices[j];
      const std::string your_name = boost::get(name_map, you);

      /** Push these two vertices into new LCA tree --- duplicate elim by map */
      new_lca_vertices[my_name] = me;
      new_lca_vertices[your_name] = you;

      /** Find the LCA, which also happens to be the common prefix */
      const std::string our_common_prefix =
        std::string(my_name.begin(),
        std::mismatch(my_name.begin(), my_name.end(), your_name.begin()).first);

      /** Search for the common ancestor */
      const bool found_ancestor = (new_lca_vertices.end() !=
                            new_lca_vertices.find(our_common_prefix));

      if (false == found_ancestor) {
      const int our_common_ancestor = all_vertices_name_map[our_common_prefix];
        new_lca_vertices[our_common_prefix] = our_common_ancestor;
      }

      /** If the common prefix is not found, then insert it in */
      if (0<options.verbosity) {
        std::cout << "Ancestor of " << me << " (code:" << my_name << ") and "
                  << you << " (code:" << your_name << ") is "
                  << new_lca_vertices[our_common_prefix] << " (code:"
            << our_common_prefix << ")" << (found_ancestor?" (already present)":
                                                    " (had to be added)")
            << std::endl;
      }
    }
  }

  if (options.sample_leaves) {
    std::cerr << file_name << ": " << num_lca_samples
              << " lca-tree (|V|=" << new_lca_vertices.size()
              << ", |E|=" << 2*num_lca_samples-1 << ")" << std::endl;
  }

  if (0<options.verbosity) {
    std::cout << "Sampled (and built) LCA tree has " << new_lca_vertices.size()
              << " nodes and these vertices: " << std::endl;
    for (MapIterator iter = new_lca_vertices.begin();
         iter != new_lca_vertices.end();
         ++iter) {
      std::cout << iter->second << " (code:" << iter->first << ")" << std::endl;
    }
  }

  const MapIterator range_first = new_lca_vertices.begin();
  MapIterator range_last = range_first;
  ++range_last;
  MapIterator iter = range_last;
  while (iter != new_lca_vertices.end()) {
    /** find parent of *iter */
    MapIterator lowest_parent = range_first;
    MapIterator first = range_first;
    while (first != range_last) {
      if (0<options.verbosity) {
        std::cout << "Testing " << first->second << " for "
                    << iter->second << std::endl;
      }
      if (first->first.end()==
          std::mismatch(first->first.begin(),
                        first->first.end(),
                        iter->first.begin()).first) {
        if (0<options.verbosity) {
          std::cout << first->second << " is an ancestor of "
                    << iter->second << std::endl;
        }
        lowest_parent = first;
      }
      ++first;
    }

    if (0<options.verbosity) {
      std::cout << lowest_parent->second << " is the LCA of "
                << iter->second << std::endl;
    }

    /** compute distance from parent to *iter */
    total_path_length += path_length(lowest_parent->first, iter->first,
                                    all_vertices_name_map, weight_map, options);
    ++total_num_edges;

    /** increment range_last */
    ++range_last;
    ++iter;
  }

  print_tree(G, options);

  return std::pair<int,double>(total_num_edges, total_path_length);
}

void run_one_experiment (const dyck_options_t& options,
                         int& gen_seed,
                         int& mle_seed,
                         int& num_nodes_in_G,
                         int& vertex_to_find,
                         double& vertex_height,
                         double& mean_height) {
  graph_type G;
  vertex_name_map_t name_map;
  vertex_weight_map_t weight_map;

  do {
    /** Generate a graph */
    G = generate_one_graph (options, gen_seed);

    /** Get the number of nodes in the graph */
    vertex_iter_t v_begin, v_end;
    boost::tie(v_begin, v_end) = boost::vertices(G);
    num_nodes_in_G = std::distance(v_begin, v_end);
  } while (num_nodes_in_G < (options.n));

  name_map = boost::get(boost::vertex_name, G);
  weight_map = boost::get(boost::vertex_distance, G);

  /**
   * If we want to identify the height of a random (non-root) vertex in the
   * tree, then we want generate a random number between [1,n), where n is
   * the number of nodes in the generated tree.
   */
  vertex_to_find = -1;
  mle_seed = (0>options.mle_seed)?generate_seed():options.mle_seed;
  boost::mt19937 engine (mle_seed);
  boost::uniform_int<int> dist (1, num_nodes_in_G-1);
  vertex_to_find = dist(engine);

  /** Now, compute the Dyck path using depth-first search */
  std::vector<double> y_axis;
  vertex_height = dyck_path (G, name_map, weight_map,
                             std::back_inserter(y_axis),
                             vertex_to_find,
                             options.verbosity);

  /** Compute the mean height of the tree */
  assert (y_axis.size() == 2*num_nodes_in_G);
  mean_height = 0.0;
  for (size_t i=0; i<y_axis.size(); ++i) mean_height += y_axis[i];
  mean_height /= (2*pow(num_nodes_in_G,1.5));

  print_tree (G, options);

#pragma omp critical
  if (options.print_path) {
    /** Write stuff out where needed */
    /** We need to ensure that the x-axis is such that the slope is always 45*/
    std::vector<double> x_axis(y_axis.size());
    x_axis[0] = 0;
    for (size_t i=1; i<x_axis.size(); ++i)
      x_axis[i] = x_axis[i-1] + std::abs(y_axis[i] - y_axis[i-1]);

    /** Remember, DO NOT PRINT THE LAST INDEX -- this is an artifact of DFS */
    std::ostream* output;
    if (0==strcmp("stdout", options.dyck_out.c_str())) output = &(std::cout);
    else output = new std::ofstream(options.dyck_out.c_str());
    for (size_t i=0; i<(x_axis.size()-1); ++i) {
      *output << x_axis[i] << "  " << y_axis[i] << std::endl;
    }
    if (0!=strcmp("stdout", options.dyck_out.c_str())) delete output;
  }

}

int main (int argc, char** argv) {
  dyck_options_t options (argc, argv);
  if (options.exit_on_return) { return -1; }

  if (1<options.verbosity) options.pretty_print();

  std::vector<std::string> dot_file_names;
  if (options.dot_in_dir != "") {
    dot_file_names = get_dot_files(options.dot_in_dir);
    options.num_trials = dot_file_names.size();
    options.gen_method = "Graphviz";
  }

  std::vector<int> gen_seed_vec(options.num_trials);
  std::vector<int> mle_seed_vec(options.num_trials);
  std::vector<int> lca_seed_vec(options.num_trials);
  std::vector<int> num_nodes_vec(options.num_trials);
  std::vector<int> vertex_to_find_vec(options.num_trials);
  std::vector<double> vertex_height_vec(options.num_trials);
  std::vector<double> mean_height_vec(options.num_trials);

  if (options.measure_lca) {

    char header_string[1024];
    int h_count = sprintf (header_string,
            "%9s %9s %2s %10s %2s %7s %7s %7s %5s %12s %12s",
                "dist-type", "f-name", "k", "p", "lm", "n",
                "n-tru", "n-smpl", "lca-edges", "lca-path-ln");
    std::cout << header_string << std::endl;

#pragma omp parallel for num_threads(options.num_threads)
    for (int i=0; i<options.num_trials; ++i) {
      /**
       * TODO: This is a hack. What we are doing is that we are checking if
       * the input method is a DOT file. If so, we are checking if the input
       * mentioned is actually a directory listing of files. In this case,
       * we will populate options.dot_in to be a different file name every
       * single time.
       */
      if (false==options.dot_in_dir.empty()) {
        options.dot_in = dot_file_names[i];
      }
      std::pair<int,double> lca_result =
        run_one_lca_experiment (options, gen_seed_vec[i],
                                lca_seed_vec[i], num_nodes_vec[i]);

      char value_string[1024];
      const std::string file_name =
              boost::filesystem::path(options.dot_in).filename().string();
      int v_count = sprintf (value_string,
               "%9s %9s %2d %.4e %2d %7d %7d %6.2f %12d %.6e",
                               options.gen_method.c_str(),
                               file_name.c_str(),
                               options.k,
                               options.p,
                               options.lambda,
                               options.n,
                               num_nodes_vec[i],
                               options.percent_lca_samples,
                               lca_result.first,
                               lca_result.second);
      std::cout << value_string << std::endl;
    }

    goto END;
  }

#pragma omp parallel for num_threads(options.num_threads)
  for (int i=0; i<options.num_trials; ++i) {
    /**
     * TODO: This is a hack. What we are doing is that we are checking if
     * the input method is a DOT file. If so, we are checking if the input
     * mentioned is actually a directory listing of files. In this case,
     * we will populate options.dot_in to be a different file name every
     * single time.
     */
    if (false==options.dot_in_dir.empty()) {
      options.dot_in = dot_file_names[i];
    }
    run_one_experiment (options, gen_seed_vec[i], mle_seed_vec[i],
                        num_nodes_vec[i], vertex_to_find_vec[i],
                        vertex_height_vec[i], mean_height_vec[i]);
  }

  double true_var;
  if (0 == strcmp("Poisson", options.gen_method.c_str()))
    true_var=options.lambda;
  else if (0 == strcmp("Binomial", options.gen_method.c_str()))
    true_var=options.k*options.p*(1-options.p);
  else if (0 == strcmp("Geometric", options.gen_method.c_str()))
    true_var=(1-options.p)/(options.p*options.p);
  else if (0 == strcmp("Binary-0-2", options.gen_method.c_str()))
    true_var=1;
  else if (0 == strcmp("Binary-0-1-2", options.gen_method.c_str()))
    true_var=2./3.;

  if (options.dump_numbers) {
    for (int i=0; i<options.num_trials; ++i) {
      double x_i = (1.0/(sqrt((double)num_nodes_vec[i])))*vertex_height_vec[i];
      printf (" %-8.3f", x_i);
    }
    printf ("\n");
    for (int i=0;i<options.num_trials;++i) {
      printf(" %-8.3f", mean_height_vec[i]);
    }
    printf ("\n");
  }

  if (options.measure_mle) {
    if (options.verbosity) {
      char header_string[1024];
      int h_count = sprintf (header_string,
            "%9s %9s %2s %10s %2s %7s %7s %7s %7s %12s %12s %12s",
                "dist-type", "filename", "k", "p", "lm", "n",
                "n-tru", "node-#","height-raw","height-nrmzd","height-mean");
      if (1<options.verbosity)
   sprintf (&(header_string[h_count]), " %10s %10s", "gen-seed", "MLE-seed");
      std::cout << header_string << std::endl;
    }

    double sum_of_x_i_sqr = 0.0;
    double mean_of_x_i_sqr = 0.0;
    for (int i=0; i<options.num_trials; ++i) {
      double x_i = (1.0/(sqrt((double)num_nodes_vec[i])))*vertex_height_vec[i];
      sum_of_x_i_sqr += x_i*x_i;
      mean_of_x_i_sqr += x_i;

      if (options.verbosity) {
        char value_string[1024];
        const std::string file_name =
              boost::filesystem::path(dot_file_names[i]).filename().string();
        int v_count = sprintf (value_string,
               "%9s %9s %2d %.4e %2d %7d %7d %7d %.6e %.6e %.6e",
                               options.gen_method.c_str(),
                               file_name.c_str(),
                               options.k,
                               options.p,
                               options.lambda,
                               options.n,
                               num_nodes_vec[i],
                               vertex_to_find_vec[i],
                               vertex_height_vec[i],
                               x_i,
                               mean_height_vec[i]);
        if (1<options.verbosity)
          sprintf (&(value_string[v_count]), " %10u %10u",
                          gen_seed_vec[i], mle_seed_vec[i]);
        std::cout << value_string << std::endl;
      }
    }
    mean_of_x_i_sqr /= options.num_trials;
    mean_of_x_i_sqr *= mean_of_x_i_sqr;

    double mle_estimate_var = (2*options.num_trials)/sum_of_x_i_sqr;
    double ano_estimate_var = PI/(2*mean_of_x_i_sqr);

    std::cout << "True variance = " << true_var
              << " MLE = " << mle_estimate_var
              << " Another = " << ano_estimate_var << std::endl;
    std::cout << " Mean of x_i^2 =  " << mean_of_x_i_sqr << std::endl;
  }

  if (options.test_confidence) {
    int num_batches = options.num_trials / options.batch_size;
    int num_contained = 0;

    int num_nodes = num_nodes_vec[0];
    bool experimental_error = false;
    for (int i=1; i<options.num_trials; ++i) {
      if (num_nodes != num_nodes_vec[i]) {
        experimental_error = true;
        break;
      }
    }

    if (experimental_error) {
      std::cout << "Run this experiment again, please" << std::endl;
      goto END;
    }

    std::vector<double> other_variances;
    if (0 == strcmp("Poisson", options.gen_method.c_str())) {
      other_variances.push_back(1 - 1./2.);
      other_variances.push_back(1 - 1./3.);
      other_variances.push_back(1 - 1./4.);
      other_variances.push_back(1 - 1./5.);
      other_variances.push_back(2);
    } else if (0 == strcmp("Binomial", options.gen_method.c_str())) {
      other_variances.push_back(1);
      other_variances.push_back(2);
      other_variances.push_back(2.0/3.0);
    } else if (0 == strcmp("Geometric", options.gen_method.c_str())) {
      other_variances.push_back(1);
      other_variances.push_back(1 - 1./2.);
      other_variances.push_back(1 - 1./3.);
      other_variances.push_back(1 - 1./4.);
      other_variances.push_back(1 - 1./5.);
    } else if (0 == strcmp("Binary-0-2", options.gen_method.c_str())) {
      other_variances.push_back(1 - 1./2.);
      other_variances.push_back(1 - 1./4.);
      other_variances.push_back(1 - 1./5.);
      other_variances.push_back(2);
      other_variances.push_back(2.0/3.0);
    } else if (0 == strcmp("Binary-0-1-2", options.gen_method.c_str())) {
      other_variances.push_back(1 - 1./2.);
      other_variances.push_back(1 - 1./3.);
      other_variances.push_back(1 - 1./4.);
      other_variances.push_back(1 - 1./5.);
      other_variances.push_back(1);
      other_variances.push_back(2);
    }
    std::vector<double> power_against_others (other_variances.size(), 0.0);

    boost::math::chi_squared chi_dist (2*options.batch_size);
    const double a = boost::math::quantile(chi_dist, 0.025);
    const double b = boost::math::quantile(chi_dist, 0.975);

    for (int i=0; i<options.num_trials; i+=options.batch_size) {
      double sum_of_x_j_sqr = 0.0;
      for (int j=0; j<options.batch_size; ++j) {
        double x_j = (1.0/(sqrt((double)num_nodes_vec[i+j])))*
                                        vertex_height_vec[i+j];
        sum_of_x_j_sqr += x_j*x_j;
      }

      const double A = a/sum_of_x_j_sqr;
      const double B = b/sum_of_x_j_sqr;
      if (A <= true_var && B >= true_var) ++num_contained;
      for (int l=0; l<other_variances.size(); ++l)
        if (A <= other_variances[l] && B >= other_variances[l])
          power_against_others[l]+=1;

      if (1<options.verbosity) {
        std::cout << "Lower = " << A << " Upper = " << B
                  << " sum_of_x_j_sqr = " << sum_of_x_j_sqr << std::endl;
      }
    }
    const double contained =
       100*static_cast<double>(num_contained)/static_cast<double>(num_batches);

    printf ("%12s %5d %3d %7.3f%%(tru-var=%4.3f) ", options.gen_method.c_str(),
                                                    num_batches,
                                                    options.batch_size,
                                                    contained,
                                                    true_var);
    for (int l=0; l<other_variances.size(); ++l)
      printf (" %7.3f%%(var=%4.3f)", 100*(power_against_others[l]/num_batches),
                                     other_variances[l]);
    printf ("\n");
  }

END:

  return 0;
}


//
// Created by mknoe on 04.04.2021.
//
#include <windows.h>        // Must have for Windows platform builds

#include <cmath>
#include <cstdio>
#include <iostream>
#include <sstream>
#include <fstream>
#include <chrono>
#include <cstdlib>
#include <utility>
#include <vector>
#include <algorithm>

#include <iterator>
#include "../exercise01/drawLine.h"
#include "../basics.h"
#include "../bench.h"
#include "../exercise02/drawCircle.h"
#include "bezier.h"
#include <boost/math/interpolators/cardinal_cubic_b_spline.hpp>
#include <boost/random.hpp>
#include <boost/math/statistics/bivariate_statistics.hpp>
#include <random>
#include <boost/math/tools/roots.hpp>
#include "splines.cpp"
#include "splines.h"
#include "vec3.hpp"
#include <boost/config.hpp>
#include <cstddef>
#include <iosfwd>
#include <GL/glu.h>

//https://chi3x10.wordpress.com/2009/10/18/de-boor-algorithm-in-c/
class Point {
public:
    Point() {
        x = 0.;
        y = 0.;
        z = 0.;
    };

    // copy operator
    Point operator=(const Point &pt);

    Point operator+(const Point &pt) const;

    //Point operator-(const Point &pt) const;
    Point operator*(double m) const;

    Point operator/(double m) const;

    double x, y, z;
};

Point Point::operator=(const Point &pt) {
    x = pt.x;
    y = pt.y;
    z = pt.z;
    return *this;
}

Point Point::operator+(const Point &pt) const {
    Point temp;
    temp.x = x + pt.x;
    temp.y = y + pt.y;
    temp.z = z + pt.z;
    return temp;
}

Point Point::operator*(double m) const {
    Point temp;
    temp.x = x * m;
    temp.y = y * m;
    temp.z = z * m;
    return temp;
}

Point Point::operator/(double m) const {
    Point temp;
    temp.x = x / m;
    temp.y = y / m;
    temp.z = z / m;
    return temp;
}

Cordinate deBoorRecursiv(int k, int degree, int i, double x, vector<double> knots, vector<Cordinate> ctrlPoints) {
    // Please see wikipedia page for detail
    // note that the algorithm here kind of traverses in reverse order
    // comapred to that in the wikipedia page
    if (k == 0)
        return ctrlPoints[i];
    else {
        double alpha = (x - knots[i]) / (knots[i + degree + 1 - k] - knots[i]);
        return (deBoorRecursiv(k - 1, degree, i - 1, x, knots, ctrlPoints) * (1 - alpha) +
                deBoorRecursiv(k - 1, degree, i, x, knots, ctrlPoints) * alpha);
    }
}

Cordinate deBoor(int k, double x, vector<double> t, vector<Cordinate> c, int p) {
/*
k: Index of knot interval that contains x.
x: Position.
t: Array of knot positions, needs to be padded as described above.
c: Array of control points.
p: Degree of B-spline.
 */
    vector<Cordinate> d{};
    for (int j = 0; j < p + 1; ++j) {
        d.push_back(c[j + k - p]);
    }
    for (int r = 1; r < p + 1; ++r) {
        for (int j = p; j > r - 1; --j) {
            int alpha = (x - t[j + k - p]) / (t[j + 1 + k - r] - t[j + k - p]);
            d[j] = (1.0 - alpha) * d[j - 1] + alpha * d[j];
        }
    }
    return d[p];
}

int WhichInterval(double x, vector<double> knot, int ti) {
    for (int i = 1; i < ti - 1; i++) {
        if (x < knot[i])
            return (i - 1);
        else if (x == knot[ti - 1])
            return (ti - 1);
    }
    return -1;
}

Cordinate bSpline(vector<Cordinate> d, int n, vector<double> t, int i) {
    auto m = d.size() - 1;
    int ts;
    if (t[i + 1] - 1 >= t[i])
        ts = t[i + 1] - 1;
    else
        ts = t[i];

    // if (0 < i <= m) throw new exception();

    vector<vector<Cordinate>> d0{};
    for (int j = 0; j < n; ++j) {
        d0.emplace_back();
        for (int l = i - n + j; l < i; l++) {
            if (j == 0) {
                d0[0].push_back(d[l]);
            } else {
                double tl = (ts - t[l]) / t[l + n - j] - t[l];
                d0[j][l] = (1 - tl) * d0[j - 1][l - 1] + tl * d0[j - 1][l];
            }
        }
    }
    return d0[n][i];
}

Cordinate closeBSpline(vector<Cordinate> d, int n, vector<double> t, int i) {
    auto m = d.size() - 1;
    int ts;
    if (t[i + 1] - 1 >= t[i])
        ts = t[i + 1] - 1;
    else
        ts = t[i];

    //  if (0 < i <= m) throw new exception();

    vector<vector<Cordinate>> d0{};

    for (int j = 0; j < n; ++j) {
        int l = i - n + j - 1;
        do {
            l++;
            if (l < 0) {
                l = l + m + 1;
                ts = ts - t[0] + t[m + 1];
            } else {
                if (l >= m + 1) {
                    l = l - m - 1;
                    ts = ts - t[m + 1] + t[0];
                }
            }
            if (j == 0) {
                d0[0].push_back(d[l]);
            } else {
                double tl = (ts - t[l]) / t[l + n - j] - t[l];
                d0[j][l] = (1 - tl) * d0[j - 1][l - 1] + tl * d0[j - 1][l];
            }
        } while (l != i);
    }
    return d0[n][i];
}

int main(int argc, char *argv[]) {

    RGBPixel color = RGBPixel(40, 40, 40);
    RGBPixel color2 = RGBPixel(0, 255, 255);
    RGBPixel color3 = RGBPixel(0, 0, 255);

    Display::init(argc, argv, 256, 256);
    auto display = Display::getInstance();
    ////////// put your framebuffer drawing code here /////////////

    const Cordinate P1 = Cordinate(30, 45);
    const Cordinate P2 = Cordinate(35, 100);
    const Cordinate P3 = Cordinate(70, 95);
    const Cordinate P4 = Cordinate(80, 35);

    std::vector<std::vector<double>> vList = {};
    const vector<Cordinate> list = {P1, P2, P3, P4};

    int k{3};
    vector<double> t{0, 0, 0, 0.5, 0.5, 0.5, 1, 1, 1};
    vector<Cordinate> c{P1, P2, P3, P4};
    auto degree = 2;
    int n = 2;
    vector<Cordinate> points;
//    for (int i = 0; i < list.size() - 1; i++) {
//        auto point = bSpline(list, n, t, i);
//
//        points.push_back(point);
//        display->setPixel(point, color);
//
//    }

    for (int x = 0; x < display->RESOLUTION; x++) {
        Cordinate b = deBoor(k, x, t, c, degree);
        Cordinate b2 = deBoorRecursiv(k, degree, WhichInterval(x, t, t.size()), x, t, c);
        display->setPixel(b, color3);
        display->setPixel(b2, color2);
        cout << "b= " << b << "\n";
        cout << "b2= " << b2 << "\n";
    }
    drawList(list, color);

    /////////////////////////////////
    glutMainLoop();
    return 0;
}


// This file is part of OpenMVG, an Open Multiple View Geometry C++ library.

// Copyright (c) 2014 cDc and Pierre MOULON.

// This Source Code Form is subject to the terms of the Mozilla Public
// License, v. 2.0. If a copy of the MPL was not distributed with this
// file, You can obtain one at http://mozilla.org/MPL/2.0/.

#include "openMVG/multiview/rotation_averaging_l1.hpp"
#include "openMVG/numeric/l1_solver_admm.hpp"

#ifdef HAVE_BOOST
#include <boost/accumulators/accumulators.hpp>
#include <boost/accumulators/statistics/max.hpp>
#include <boost/accumulators/statistics/mean.hpp>
#include <boost/accumulators/statistics/min.hpp>
#include <boost/accumulators/statistics/stats.hpp>
#include <boost/foreach.hpp>
#include <boost/graph/adjacency_list.hpp>
#include <boost/graph/kruskal_min_spanning_tree.hpp>
#include <boost/graph/prim_minimum_spanning_tree.hpp>
using namespace boost;
#else
#include "lemon/adaptors.h"
#include "lemon/dfs.h"
#include "lemon/kruskal.h"
#include "lemon/list_graph.h"
#include "lemon/path.h"
using namespace lemon;
#endif

#include <ceres/ceres.h>
#include <ceres/rotation.h>

#include <Eigen/Cholesky>
#include <Eigen/SparseCholesky>

#include <queue>

namespace openMVG   {
namespace rotation_averaging  {
namespace l1  {

/////////////////////////

// given an array of values, compute the X84 threshold as in:
// Hampel FR, Rousseeuw PJ, Ronchetti EM, Stahel WA
// "Robust Statistics: the Approach Based on Influence Functions"
// Wiley Series in Probability and Mathematical Statistics, John Wiley & Sons, 1986
// returns the pair(median,trust_region)
// upper-bound threshold = median+trust_region
// lower-bound threshold = median-trust_region
template<typename TYPE>
inline std::pair<TYPE, TYPE>
ComputeX84Threshold(const TYPE* const values, uint32_t size, TYPE mul=TYPE(5.2))
{
  assert(size > 0);
  typename std::vector<TYPE> data(values, values+size);
  typename std::vector<TYPE>::iterator mid = data.begin() + size / 2;
  std::nth_element(data.begin(), mid, data.end());
  const TYPE median = *mid;
  // threshold = 5.2 * MEDIAN(ABS(values-median));
  for (size_t i=0; i<size; ++i)
    data[i] = std::abs(values[i]-median);
  std::nth_element(data.begin(), mid, data.end());
  return {median, mul*(*mid)};
} // ComputeX84Threshold

/////////////////////////

using Matrix3x3 = openMVG::Mat3;
using IndexArr = std::vector<uint32_t>;

// find the shortest cycle for the given graph and starting vertex
struct Node {
  using InternalType = IndexArr;
  InternalType edges; // array of vertex indices
};
using NodeArr = std::vector<Node>;

struct Link {
  uint32_t ID; // node index
  uint32_t parentID;// parent link
  inline Link(uint32_t ID_=0, uint32_t parentID_=0) : ID(ID_), parentID(parentID_) {}
};
using LinkQue = std::queue<Link>;

#ifdef HAVE_BOOST
using edge_property_t = boost::property<boost::edge_weight_t, float>;
using graph_t = boost::adjacency_list<boost::vecS, boost::vecS, boost::undirectedS, uint32_t, edge_property_t>;
using vertex_t = graph_t::vertex_descriptor;
using edge_t = graph_t::edge_descriptor;
using edge_iter = boost::graph_traits<graph_t>::edge_iterator;
#else
using graph_t = lemon::ListGraph;
using map_EdgeMap = graph_t::EdgeMap<double>;
#endif
using MapEdgeIJ2R = std::map<std::pair<uint32_t,uint32_t>, Matrix3x3>;

// Look for the maximum spanning tree along the graph of relative rotations
// since we look for the maximum spanning tree using a minimum spanning tree algorithm
// weight are negated.
uint32_t FindMaximumSpanningTree(const RelativeRotations& RelRs, graph_t& g, MapEdgeIJ2R& mapIJ2R, NodeArr& minGraph)
{
  assert(!RelRs.empty());
#ifdef HAVE_BOOST
  for (size_t p = 0; p < RelRs.size(); ++p) {
    const RelativeRotation& relR = RelRs[p];
    boost::add_edge(relR.i, relR.j, - relR.weight, g);
    mapIJ2R[{relR.i, relR.j}] = relR.Rij;
    mapIJ2R[{relR.j, relR.i}] = relR.Rij.transpose();
  }
  // find the minimum spanning tree
  const size_t nViews = boost::num_vertices(g);
  minGraph.resize(nViews);
  std::vector<edge_t> spanningTree;
  boost::kruskal_minimum_spanning_tree(g, std::back_inserter(spanningTree));
  for (std::vector<edge_t>::const_iterator ei=spanningTree.begin(); ei!=spanningTree.end(); ++ei) {
    const edge_t& edge = *ei;
    minGraph[edge.m_source].edges.push_back(edge.m_target);
    minGraph[edge.m_target].edges.push_back(edge.m_source);
  }
  const size_t nEdges = spanningTree.size();
  return nEdges;
#else

  //A-- Compute the number of node we need
  std::set<uint32_t> setNodes;
  for (size_t p = 0; p < RelRs.size(); ++p) {
    const RelativeRotation& relR = RelRs[p];
    setNodes.insert(relR.i);
    setNodes.insert(relR.j);
  }

  //B-- Create a node graph for each element of the set
  using map_NodeT = std::map<uint32_t, graph_t::Node>;
  map_NodeT map_index_to_node;
  for (const auto & iter : setNodes)
  {
    map_index_to_node[iter] = g.addNode();
  }

  //C-- Create a graph from RelRs with weighted edges
  map_EdgeMap map_edgeMap(g);
  for (size_t p = 0; p < RelRs.size(); ++p) {
    const RelativeRotation& relR = RelRs[p];
    mapIJ2R[{relR.i, relR.j}] = relR.Rij;
    mapIJ2R[{relR.j, relR.i}] = relR.Rij.transpose();

    // add edge to the graph
    graph_t::Edge edge =  g.addEdge(map_index_to_node[relR.i], map_index_to_node[relR.j]);
    map_edgeMap[ edge ] = - relR.weight;
  }

  //D-- Compute the MST of the graph
  std::vector<graph_t::Edge> tree_edge_vec;
  lemon::kruskal(g, map_edgeMap, std::back_inserter(tree_edge_vec));

  const size_t nViews = lemon::countNodes(g);
  minGraph.resize(nViews);

  //E-- Export compute MST
  for (size_t i= 0; i < tree_edge_vec.size(); i++)
  {
    minGraph[g.id(g.u(tree_edge_vec[i]))].edges.push_back(g.id(g.v(tree_edge_vec[i])));
    minGraph[g.id(g.v(tree_edge_vec[i]))].edges.push_back(g.id(g.u(tree_edge_vec[i])));
  }
  return tree_edge_vec.size();
#endif
}
//----------------------------------------------------------------

// Filter the given relative rotations using the known global rotations
// returns the number of inliers
unsigned int FilterRelativeRotations(
  const RelativeRotations& RelRs,
  const Matrix3x3Arr& Rs,
  float threshold,
  std::vector<bool> * vec_inliers)
{
  assert(!RelRs.empty() && !Rs.empty());
  assert(threshold >= 0);
  // compute errors for each relative rotation
  std::vector<float> errors(RelRs.size());
  for (size_t r= 0; r<RelRs.size(); ++r) {
    const RelativeRotation& relR = RelRs[r];
    const Matrix3x3& Ri = Rs[relR.i];
    const Matrix3x3& Rj = Rs[relR.j];
    const Matrix3x3& Rij = relR.Rij;
    const Mat3 eRij(Rj.transpose()*Rij*Ri);
    const openMVG::Vec3 erij;
    ceres::RotationMatrixToAngleAxis((const double*)eRij.data(), (double*)erij.data());
    errors[r] = (float)erij.norm();
  }
  if (threshold == 0) {
    // estimate threshold
    const std::pair<float,float> res = ComputeX84Threshold(&errors[0], errors.size());
    threshold = res.first+res.second;
  }
  if (vec_inliers)  {
    vec_inliers->resize(RelRs.size());
  }
  // mark outliers
  unsigned int nInliers = 0;
  for (size_t r=0; r<errors.size(); ++r) {
    const bool bInlier = errors[r] < threshold;
    if (vec_inliers)
      (*vec_inliers)[r] = bInlier;
    if (bInlier)
      ++nInliers;
  }
  return nInliers;
} // FilterRelativeRotations
//----------------------------------------------------------------

double RelRotationAvgError
(
  const RelativeRotations& RelRs,
  const Matrix3x3Arr& Rs,
  double* pMin=nullptr,
  double* pMax=nullptr
)
{
#ifdef HAVE_BOOST
  boost::accumulators::accumulator_set<double,
    boost::accumulators::stats<
      boost::accumulators::tag::min,
      boost::accumulators::tag::mean,
      boost::accumulators::tag::max>> acc;

  for (int i=0; i < RelRs.size(); ++i) {
    const RelativeRotation& relR = RelRs[i];
    acc(openMVG::FrobeniusNorm(relR.Rij  - (Rs[relR.j]*Rs[relR.i].transpose())));
  }
  if (pMin)
    *pMin = boost::accumulators::min(acc);
  if (pMax)
    *pMax = boost::accumulators::max(acc);
  return boost::accumulators::mean(acc);
#else
  std::vector<double> vec_err(RelRs.size(), 0.0);
  for (size_t i=0; i < RelRs.size(); ++i) {
    const RelativeRotation& relR = RelRs[i];
    vec_err[i] = openMVG::FrobeniusNorm(relR.Rij  - (Rs[relR.j]*Rs[relR.i].transpose()));
  }
  float min, max, mean, median;
  minMaxMeanMedian(vec_err.begin(), vec_err.end(), min, max, mean, median);
  if (pMin)
    *pMin = min;
  if (pMax)
    *pMax = max;
  return mean;
#endif
}
//----------------------------------------------------------------

void InitRotationsMST
(
  const RelativeRotations& RelRs,
  Matrix3x3Arr& Rs,
  const uint32_t nMainViewID
)
{
  assert(!Rs.empty());

  // -- Compute coarse global rotation estimates:
  //   - by finding the maximum spanning tree and linking the relative rotations
  //   - Initial solution is driven by relative rotations data confidence.
  graph_t g;
  MapEdgeIJ2R mapIJ2R;
  NodeArr minGraph;
  // find the Maximum Spanning Tree
  FindMaximumSpanningTree(RelRs, g, mapIJ2R, minGraph);
  g.clear();

  // start from the main view and link all views using the relative rotation estimates
  LinkQue stack;
  stack.push(Link(nMainViewID, uint32_t(0)));
  Rs[nMainViewID] = Matrix3x3::Identity();
  do {
    const Link& link = stack.front();
    const Node& node = minGraph[link.ID];

    for (Node::InternalType::const_iterator pEdge = node.edges.begin();
      pEdge != node.edges.end(); ++pEdge) {
        const size_t edge = *pEdge;
        if (edge == link.parentID) {
          // compute the global rotation for the current node
          assert(mapIJ2R.find({link.parentID, link.ID}) != mapIJ2R.end());
          const Matrix3x3& Rij = mapIJ2R[{link.parentID, link.ID}];
          Rs[link.ID] = Rij * Rs[link.parentID];
        } else {
          // add edge to the processing queue
          stack.push(Link(edge, link.ID));
        }
    }
    stack.pop();
  } while (!stack.empty());
}

// Robustly estimate global rotations from relative rotations as in:
// "Efficient and Robust Large-Scale Rotation Averaging", Chatterjee and Govindu, 2013
// and detect outliers relative rotations and return them with 0 in arrInliers
bool GlobalRotationsRobust(
  const RelativeRotations& RelRs,
  Matrix3x3Arr& Rs,
  const uint32_t nMainViewID,
  float threshold,
  std::vector<bool> * vec_Inliers)
{
  assert(!Rs.empty());

  // -- Compute coarse global rotation estimates:
  InitRotationsMST(RelRs, Rs, nMainViewID);

  // refine global rotations based on the relative rotations
  const bool bOk = RefineRotationsAvgL1IRLS(RelRs, Rs, nMainViewID);

  // find outlier relative rotations
  if (threshold>=0 && vec_Inliers)  {
    FilterRelativeRotations(RelRs, Rs, threshold, vec_Inliers);
  }

  return bOk;
} // GlobalRotationsRobust
//----------------------------------------------------------------

namespace internal
{

// build A in Ax=b
inline void FillMappingMatrix(
  const RelativeRotations& RelRs,
  const uint32_t nMainViewID,
  sMat& A)
{
  A.reserve(A.rows()*2); // estimate of the number of non-zeros (optional)
  sMat::Index i = 0, j = 0;
  for (size_t r=0; r<RelRs.size(); ++r) {
    const RelativeRotation& relR = RelRs[r];
    if (relR.i != nMainViewID) {
      j = 3*(relR.i<nMainViewID ? relR.i : relR.i-1);
      A.insert(i+0,j+0) = -1.0;
      A.insert(i+1,j+1) = -1.0;
      A.insert(i+2,j+2) = -1.0;
    }
    if (relR.j != nMainViewID) {
      j = 3*(relR.j<nMainViewID ? relR.j : relR.j-1);
      A.insert(i+0,j+0) = 1.0;
      A.insert(i+1,j+1) = 1.0;
      A.insert(i+2,j+2) = 1.0;
    }
    i+=3;
  }
  A.makeCompressed();
}

// compute errors for each relative rotation
inline void FillErrorMatrix(
  const RelativeRotations& RelRs,
  const Matrix3x3Arr& Rs,
  Vec & b)
{
  for (size_t r = 0; r < RelRs.size(); ++r) {
    const RelativeRotation& relR = RelRs[r];
    const Matrix3x3& Ri = Rs[relR.i];
    const Matrix3x3& Rj = Rs[relR.j];
    const Matrix3x3& Rij = relR.Rij;
    const Mat3 eRij(Rj.transpose()*Rij*Ri);
    const openMVG::Vec3 erij;
    ceres::RotationMatrixToAngleAxis((const double*)eRij.data(), (double*)erij.data());
    b.block<3,1>(3*r,0) = erij;
  }
}

// apply correction to global rotations
inline void CorrectMatrix(
  const Mat& x,
  const uint32_t nMainViewID,
  Matrix3x3Arr& Rs)
{
  for (size_t r = 0; r < Rs.size(); ++r) {
    if (r == nMainViewID)
      continue;
    Matrix3x3& Ri = Rs[r];
    const uint32_t i = (r<nMainViewID ? r : r-1);
    const openMVG::Vec3 eRid = openMVG::Vec3(x.block<3,1>(3*i,0));
    const Mat3 eRi;
    ceres::AngleAxisToRotationMatrix((const double*)eRid.data(), (double*)eRi.data());
    Ri = Ri*eRi;
  }
}

// L1RA -> L1 Rotation Averaging implementation
bool SolveL1RA
(
  const RelativeRotations& RelRs,
  Matrix3x3Arr& Rs,
  const sMat & A,
  const unsigned int nMainViewID
)
{
  const unsigned nObss = (unsigned)RelRs.size();
  const unsigned nVars = (unsigned)Rs.size()-1; // one view is kept constant
  const unsigned m = nObss*3;
  const unsigned n = nVars*3;

  // init x with 0 that corresponds to trusting completely the initial Ri guess
  Vec x(Vec::Zero(n)), b(m);

  // Current error and the previous one
  double e = std::numeric_limits<double>::max(), ep;
  unsigned iter = 0;
  // L1RA iterate optimization till the desired precision is reached
  do {
    // compute errors for each relative rotation
    FillErrorMatrix(RelRs, Rs, b);

    // solve the linear system using l1 norm
    L1Solver<sMat >::Options options;
    L1Solver<sMat > l1_solver(options, A);
    l1_solver.Solve(b, &x);

    ep = e; e = x.norm();
    if (ep < e)
      break;
    // apply correction to global rotations
    CorrectMatrix(x, nMainViewID, Rs);
  } while (++iter < 32 && e > 1e-5 && (ep-e)/e > 1e-2);

  std::cout << "L1RA Converged in " << iter << " iterations." << std::endl;

  return true;
}

// Iteratively Reweighted Least Squares (IRLS) implementation
bool SolveIRLS
(
  const RelativeRotations& RelRs,
  Matrix3x3Arr& Rs,
  const sMat & A,
  const unsigned int nMainViewID,
  const double sigma
)
{
  const unsigned nObss = (unsigned)RelRs.size();
  const unsigned nVars = (unsigned)Rs.size()-1; // one view is kept constant
  const unsigned m = nObss*3;
  const unsigned n = nVars*3;

  // init x with 0 that corresponds to trusting completely the initial Ri guess
  Vec x(Vec::Zero(n)), b(m);

  // Since the sparsity pattern will not change with each linear solve
  //  compute it once to speed up the solution time.
  using Linear_Solver_T = Eigen::SimplicialLDLT<sMat >;

  Linear_Solver_T linear_solver;
  linear_solver.analyzePattern(A.transpose() * A);
  if (linear_solver.info() != Eigen::Success) {
    std::cerr << "Cholesky decomposition failed." << std::endl;
    return false;
  }

  const double sigmaSq(Square(sigma));

  Eigen::ArrayXd errors, weights;
  Vec xp(n);
  // current error and the previous one
  double e = std::numeric_limits<double>::max(), ep;
  unsigned int iter = 0;
  do
  {
    xp = x;
    // compute errors for each relative rotation
    FillErrorMatrix(RelRs, Rs, b);

    // Compute the weights for each error term
    errors = (A * x - b).array();

    // compute robust errors using the Huber-like loss function
    weights = sigmaSq / (errors.square() + sigmaSq).square();

    // Update the factorization for the weighted values
    const sMat at_weight = A.transpose() * weights.matrix().asDiagonal();
    linear_solver.factorize(at_weight * A);
    if (linear_solver.info() != Eigen::Success) {
      std::cerr << "Failed to factorize the least squares system." << std::endl;
      return false;
    }

    // Solve the least squares problem
    x = linear_solver.solve(at_weight * b);
    if (linear_solver.info() != Eigen::Success) {
      std::cerr << "Failed to solve the least squares system." << std::endl;
      return false;
    }

    // apply correction to global rotations
    CorrectMatrix(x, nMainViewID, Rs);

    ep = e; e = (xp-x).norm();

  } while (++iter < 32 && e > 1e-5 && (ep-e)/e > 1e-2);

  std::cout << "IRLS Converged in " << iter << " iterations." << std::endl;

  return true;
}

} // namespace internal

// Refine the global rotations using to the given relative rotations, similar to:
// "Efficient and Robust Large-Scale Rotation Averaging", Chatterjee and Govindu, 2013
// L1 Rotation Averaging (L1RA) and Iteratively Reweighted Least Squares (IRLS) implementations combined
bool RefineRotationsAvgL1IRLS(
  const RelativeRotations& RelRs,
  Matrix3x3Arr& Rs,
  const uint32_t nMainViewID,
  const double sigma)
{
  assert(!RelRs.empty() && !Rs.empty());
  assert(Rs[nMainViewID] == Matrix3x3::Identity());

  double fMinBefore, fMaxBefore;
  const double fMeanBefore = RelRotationAvgError(RelRs, Rs, &fMinBefore, &fMaxBefore);

  const unsigned nObss = (unsigned)RelRs.size();
  const unsigned nVars = (unsigned)Rs.size()-1; // main view is kept constant
  const unsigned m = nObss*3;
  const unsigned n = nVars*3;

  // build mapping matrix A in Ax=b
  sMat A(m, n);
  internal::FillMappingMatrix(RelRs, nMainViewID, A);

  if (!internal::SolveL1RA(RelRs, Rs, A, nMainViewID))
  {
    std::cerr << "Could not solve the L1 regression step." << std::endl;
    return false;
  }

  if (!internal::SolveIRLS(RelRs, Rs, A, nMainViewID, sigma))
  {
    std::cerr << "Could not solve the ILRS step." << std::endl;
    return false;
  }

  double fMinAfter, fMaxAfter;
  const double fMeanAfter = RelRotationAvgError(RelRs, Rs, &fMinAfter, &fMaxAfter);

  std::cout << "Refine global rotations using L1RA-IRLS and " << nObss << " relative rotations:\n"
    << " error reduced from " << fMeanBefore << "(" <<fMinBefore << " min, " << fMaxBefore << " max)\n"
    << " to " << fMeanAfter << "(" << fMinAfter << "min,"<< fMaxAfter<< "max)" << std::endl;

  return true;
} // RefineRotationsAvgL1IRLS

} // namespace l1
} // namespace rotation_averaging
} // namespace openMVG