/*
   For more information, please see: http://software.sci.utah.edu

   The MIT License

   Copyright (c) 2020 Scientific Computing and Imaging Institute,
   University of Utah.

   Permission is hereby granted, free of charge, to any person obtaining a
   copy of this software and associated documentation files (the "Software"),
   to deal in the Software without restriction, including without limitation
   the rights to use, copy, modify, merge, publish, distribute, sublicense,
   and/or sell copies of the Software, and to permit persons to whom the
   Software is furnished to do so, subject to the following conditions:

   The above copyright notice and this permission notice shall be included
   in all copies or substantial portions of the Software.

   THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS
   OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
   FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL
   THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
   LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING
   FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER
   DEALINGS IN THE SOFTWARE.
*/

#include <Core/Algorithms/Legacy/FiniteElements/BuildMatrix/BuildFEMatrix.h>

#include <Core/Datatypes/DenseMatrix.h>
#include <Core/Datatypes/SparseRowMatrix.h>

#include <Core/Thread/Barrier.h>
#include <Core/Thread/Parallel.h>

#include <Core/Datatypes/Legacy/Field/Mesh.h>
#include <Core/Datatypes/Legacy/Field/VMesh.h>
#include <Core/Datatypes/Legacy/Field/Field.h>
#include <Core/Datatypes/Legacy/Field/VField.h>
#include <Core/GeometryPrimitives/Tensor.h>
#include <Core/Algorithms/Base/AlgorithmPreconditions.h>
#include <Core/Algorithms/Base/AlgorithmVariableNames.h>
#include <Core/Logging/Log.h>

#include <string>
#include <vector>
#include <algorithm>
#include <boost/shared_array.hpp>

using namespace SCIRun;
using namespace SCIRun::Core::Geometry;
using namespace SCIRun::Core::Datatypes;
using namespace SCIRun::Core::Thread;
using namespace SCIRun::Core::Algorithms;
using namespace SCIRun::Core::Algorithms::FiniteElements;
using namespace SCIRun::Core::Logging;

namespace SCIRun {
	namespace Core {
		namespace Algorithms {
			namespace FiniteElements {

        template <typename T>
        using matrix_type = Datatypes::SparseRowMatrixGeneric<T>;
        template <typename T>
        using matrix_pointer_type = SharedPointer<matrix_type<T>>;

template <typename T>
class BuildFEMatrixAlgoImpl
{
public:
  explicit BuildFEMatrixAlgoImpl(const AlgorithmBase* algo) : algo_(algo) {}
  bool run(FieldHandle input, Datatypes::DenseMatrixHandle ctable, matrix_pointer_type<T>& output) const;
private:
  const AlgorithmBase* algo_;
  mutable int generation_ = 0;
  mutable std::vector<std::vector<T>> basis_values_;
  mutable matrix_pointer_type<T> basis_fematrix_;
};

// Helper class
template <typename T>
class FEMBuilder
{
public:
  explicit FEMBuilder(const AlgorithmBase* algo) :
    algo_(algo), numprocessors_(Parallel::NumCores()),
    barrier_("FEMBuilder Barrier", numprocessors_),
    mesh_(nullptr), field_(nullptr),
    domain_dimension(0), local_dimension_nodes(0),
    local_dimension_add_nodes(0),
    local_dimension_derivatives(0),
    local_dimension(0),
    global_dimension_nodes(0),
    global_dimension_add_nodes(0),
    global_dimension_derivatives(0),
    global_dimension(0)
  {
  }

  bool build_matrix(FieldHandle input,
                    DenseMatrixHandle ctable,
                    matrix_pointer_type<T>& output);

private:
  const AlgorithmBase* algo_;
  int numprocessors_;
  Barrier barrier_;

  VMesh* mesh_;
  VField *field_;

  matrix_pointer_type<T> fematrix_;

  std::vector<bool> success_;

  boost::shared_array<index_type> rows_;
  boost::shared_array<index_type> allcols_;
  std::vector<index_type> colidx_;

  index_type domain_dimension;

  index_type local_dimension_nodes;
  index_type local_dimension_add_nodes;
  index_type local_dimension_derivatives;
  index_type local_dimension;

  index_type global_dimension_nodes;
  index_type global_dimension_add_nodes;
  index_type global_dimension_derivatives;
  index_type global_dimension;

  // A copy of the tensors list that was generated by SetConductivities
  std::vector<std::pair<std::string, Tensor> > tensors_;
  std::vector<std::pair<std::string, T> > scalars_;

  // Entry point for the parallel version
  void parallel(int proc);

  void add_lcl_gbl(index_type row, const std::vector<index_type> &cols, const std::vector<T> &lcl_a)
  {
    for (size_t i = 0; i < lcl_a.size(); i++)
      fematrix_->coeffRef(row, cols[i]) += lcl_a[i];
  }

  void create_numerical_integration(std::vector<VMesh::coords_type>& p,
                                    std::vector<double>& w,
                                    std::vector<std::vector<double>>& d);
  bool build_local_matrix(VMesh::Elem::index_type c_ind,
                          index_type row,
                          std::vector<T>& l_stiff,
                          std::vector<VMesh::coords_type>& p,
                          std::vector<double>& w,
                          std::vector<std::vector<double>>& d);
  bool build_local_matrix_regular(VMesh::Elem::index_type c_ind,
                                  index_type row,
                                  std::vector<T>& l_stiff,
                                  std::vector<VMesh::coords_type>& p,
                                  std::vector<double>& w,
                                  std::vector<std::vector<double>>& d,
                                  std::vector<std::vector<T>>& precompute);
  bool setup();

};
}}}}

template <typename T>
bool
FEMBuilder<T>::build_matrix(FieldHandle input,
                         DenseMatrixHandle ctable,
                         matrix_pointer_type<T>& output)
{
  // Get virtual interface to data
  field_ = input->vfield();
  mesh_  = input->vmesh();

#ifdef SCIRUN4_CODE_TO_BE_ENABLED_LATER
  // If we have the Conductivity property use it, if not we assume the values on
  // the data to be the actual tensors.
  field_->get_property("conductivity_table",tensors_);
#endif

  // We added a second system of adding a conductivity table, using a matrix
  // Convert that matrix into the conductivity table
  if (ctable)
  {
    tensors_.clear();
    auto mat = ctable;
    // Only if we can convert it into a dense matrix, otherwise skip it
    if (mat)
    {
      auto data = mat->data();
      size_type m = mat->nrows();
      size_type n = mat->ncols();
      Tensor tensor;

      // Case the table has isotropic conductivities
      if (mat->ncols() == 1)
      {
        for (size_type p=0; p<m;p++)
        {
          // Set the diagonals to the proper version.
          tensor.val(0,0) = data[p*n+0];
          tensor.val(1,0) = 0.0;
          tensor.val(2,0) = 0.0;
          tensor.val(0,1) = 0.0;
          tensor.val(1,1) = data[p*n+0];
          tensor.val(2,1) = 0.0;
          tensor.val(0,2) = 0.0;
          tensor.val(1,2) = 0.0;
          tensor.val(2,2) = data[p*n+0];
          tensors_.push_back(std::make_pair("",tensor));
        }
      }

      // Use our compressed way of storing tensors
      if (mat->ncols() == 6)
      {
        for (size_type p=0; p<m;p++)
        {
          tensor.val(0,0) = data[0+p*n];
          tensor.val(1,0) = data[1+p*n];
          tensor.val(2,0) = data[2+p*n];
          tensor.val(0,1) = data[1+p*n];
          tensor.val(1,1) = data[3+p*n];
          tensor.val(2,1) = data[4+p*n];
          tensor.val(0,2) = data[2+p*n];
          tensor.val(1,2) = data[4+p*n];
          tensor.val(2,2) = data[5+p*n];
          tensors_.push_back(std::make_pair("",tensor));
        }
      }

      // Use the full symmetric tensor. We will make the tensor symmetric here.
      if (mat->ncols() == 9)
      {
        for (size_type p=0; p<m;p++)
        {
          tensor.val(0,0) = data[0+p*n];
          tensor.val(1,0) = data[1+p*n];
          tensor.val(2,0) = data[2+p*n];
          tensor.val(0,1) = data[1+p*n];
          tensor.val(1,1) = data[4+p*n];
          tensor.val(2,1) = data[5+p*n];
          tensor.val(0,2) = data[2+p*n];
          tensor.val(1,2) = data[5+p*n];
          tensor.val(2,2) = data[8+p*n];
          tensors_.push_back(std::make_pair("",tensor));
        }
      }
    }
  }

  success_.resize(numprocessors_,true);

  // Start the multi threaded FE matrix builder.
  Parallel::RunTasks([this](int i) { parallel(i); }, numprocessors_);
  for (size_t j=0; j<success_.size(); j++)
  {
    if (!success_[j])
    {
      std::ostringstream oss;
      oss << "Algorithm failed in thread " << j;
      algo_->error(oss.str());
      return false;
    }
  }

  // Make sure it is symmetric
  if (algo_->get(BuildFEMatrixAlgo::ForceSymmetry).toBool())
  {
    // Make sure the matrix is fully symmetric, this compensates for round off
    // errors
    matrix_type<T> transpose = fematrix_->transpose();
    output.reset(new matrix_type<T>(0.5*(transpose + *fematrix_)));
  }
  else
  {
    // Assume that the builder did a good job and the matrix is numerically almost
    // symmetric
    output = fematrix_;
  }
  return true;
}

template <typename T>
void
FEMBuilder<T>::create_numerical_integration(std::vector<VMesh::coords_type> &p,
                                         std::vector<double> &w,
                                         std::vector<std::vector<double> > &d)
{
  //ScopedTimeLogger s1("FEMBuilder::create_numerical_integration");
  int int_basis = 1;
  if (mesh_->is_quad_element() ||
      mesh_->is_hex_element() ||
      mesh_->is_prism_element())
  {
    int_basis = 2;
  }

  mesh_->get_gaussian_scheme(p,w,int_basis);
  d.resize(p.size());
  for (size_t j=0; j<p.size();j++)
  {
    mesh_->get_derivate_weights(p[j],d[j],1);
    size_t pad_size = ( 3 - p[ j ].size() ) * d[ j ].size();

    if (pad_size > 0)
      d[j].resize(pad_size + d[j].size(), 0.0);
  }
}

/// build line of the local stiffness matrix
template <typename T>
bool
FEMBuilder<T>::build_local_matrix(VMesh::Elem::index_type c_ind,
                               index_type row,
                               std::vector<T> &l_stiff,
                               std::vector<VMesh::coords_type> &p,
                               std::vector<double> &w,
                               std::vector<std::vector<double>>  &d)
{
  Tensor tensor;

  if (tensors_.empty())
  {
    field_->get_value(tensor,c_ind);
  }
  else
  {
    int tensor_index;
    field_->get_value(tensor_index,c_ind);
    tensor = tensors_[tensor_index].second;
  }

  auto Ca = tensor.val(0,0);
  auto Cb = tensor.val(0,1);
  auto Cc = tensor.val(0,2);
  auto Cd = tensor.val(1,1);
  auto Ce = tensor.val(1,2);
  auto Cf = tensor.val(2,2);

  if ( (Ca==0) && (Cb==0) && (Cc==0) && (Cd==0) && (Ce==0) && (Cf==0) )
  {
    for (int j = 0; j<local_dimension; j++)
    {
      l_stiff[j] = 0.0;
    }
  }
  else
  {
    /// @todo: replace with std::fill
    for (int i=0; i<local_dimension; i++)
      l_stiff[i] = 0.0;

    auto local_dimension2=2*local_dimension;

    // These calls are direct lookups in the base of the VMesh
    // The compiler should optimize these well
    auto vol = mesh_->get_element_size();
    const int dim = mesh_->dimensionality();

    if (dim < 1 || dim > 3)
    {
      algo_->error("Mesh dimension is 0 or larger than 3, for which no FE implementation is available");
      return false;
    }
    for (size_t i = 0; i < d.size(); i++)
    {
      double Ji[9];
      auto detJ = mesh_->inverse_jacobian(p[i],c_ind,Ji);

      // If Jacobian is negative there is a problem with the mesh
      if (detJ <= 0.0)
      {
        algo_->error("Mesh has elements with negative jacobians, check the order of the nodes that define an element");
        return false;
      }

      // Volume associated with the local Gaussian Quadrature point:
      // weightfactor * Volume Unit element * Volume ratio (real element/unit element)
      detJ *= w[i] * vol;

      // Build local stiffness matrix
      // Get the local derivatives of the basis functions in the basis element
      // They are all the same and are thus precomputed in matrix d
      auto Nxi = &d[i][0];
      auto Nyi = &d[i][local_dimension];
      auto Nzi = &d[i][local_dimension2];
      // Gradients associated with the node we are calculating
      const auto& Nxip = Nxi[row];
      const auto &Nyip = Nyi[row];
      const auto &Nzip = Nzi[row];
      // Calculating gradient shape function * inverse Jacobian * volume scaling factor
      const auto uxp = detJ*(Nxip*Ji[0] + Nyip*Ji[1] + Nzip*Ji[2]);
      const auto uyp = detJ*(Nxip*Ji[3] + Nyip*Ji[4] + Nzip*Ji[5]);
      const auto uzp = detJ*(Nxip*Ji[6] + Nyip*Ji[7] + Nzip*Ji[8]);
      // Matrix multiplication with conductivity tensor :
      const auto uxyzpabc = uxp*Ca + uyp*Cb + uzp*Cc;
      const auto uxyzpbde = uxp*Cb + uyp*Cd + uzp*Ce;
      const auto uxyzpcef = uxp*Cc + uyp*Ce + uzp*Cf;

      // The above is constant for this node. Now multiply with the weight function
      // We assume the weight factors are the same as the local gradients
      // Galerkin approximation:

      for (int j = 0; j<local_dimension; j++)
      {
        const auto &Nxj = Nxi[j];
        const auto &Nyj = Nyi[j];
        const auto &Nzj = Nzi[j];

        // Matrix multiplication Gradient with inverse Jacobian:
        const auto ux = Nxj*Ji[0] + Nyj*Ji[1] + Nzj*Ji[2];
        const auto uy = Nxj*Ji[3] + Nyj*Ji[4] + Nzj*Ji[5];
        const auto uz = Nxj*Ji[6] + Nyj*Ji[7] + Nzj*Ji[8];

        // Add everything together into one coefficient of the matrix
        l_stiff[j] += ux*uxyzpabc+uy*uxyzpbde+uz*uxyzpcef;
      }
    }
  }
  return true;
}

template <typename T>
bool
FEMBuilder<T>::build_local_matrix_regular(VMesh::Elem::index_type c_ind,
                                       index_type row,
                                       std::vector<T> &l_stiff,
                                       std::vector<VMesh::coords_type> &p,
                                       std::vector<double> &w,
                                       std::vector<std::vector<double>> &d,
                                       std::vector<std::vector<T>> &precompute)
{
  Tensor tensor;

  if (tensors_.empty())
  {
    // Call to virtual interface. Get the tensor value. Actually this call relies
    // on the automatic casting feature of the virtual interface to convert scalar
    // values into a tensor.
    field_->get_value(tensor,c_ind);
  }
  else
  {
    int tensor_index;
    field_->get_value(tensor_index,c_ind);
    tensor = tensors_[tensor_index].second;
  }

  auto Ca = tensor.val(0,0);
  auto Cb = tensor.val(0,1);
  auto Cc = tensor.val(0,2);
  auto Cd = tensor.val(1,1);
  auto Ce = tensor.val(1,2);
  auto Cf = tensor.val(2,2);

  if ( (Ca==0) && (Cb==0) && (Cc==0) && (Cd==0) && (Ce==0) && (Cf==0) )
  {
    for (int j = 0; j<local_dimension; j++)
    {
      l_stiff[j] = 0.0;
    }
  }
  else
  {

    if (precompute.empty())
    {
      precompute.resize(d.size());
      for (int m=0; m < static_cast<int>(d.size()); m++)
      {
        precompute[m].resize(10);
      }

      for(int i=0; i<local_dimension; i++)
        l_stiff[i] = 0.0;

      auto local_dimension2=2*local_dimension;

      auto vol = mesh_->get_element_size();

      for (size_t i = 0; i < d.size(); i++)
      {
        auto& pc = precompute[i];

        double Ji[9];
        auto detJ = mesh_->inverse_jacobian(p[i], c_ind, Ji);

        // Volume elements can return negative determinants if the order of elements
        // is put in a different order
        /// @todo: It seems to be that a negative determinant is not necessarily bad,
        // we should be more flexible on this point
        if (detJ <= 0.0)
        {
          algo_->error("Mesh has elements with negative jacobians, check the order of the nodes that define an element");
          return false;
        }
        // Volume associated with the local Gaussian Quadrature point:
        // weightfactor * Volume Unit element * Volume ratio (real element/unit element)
        detJ*=w[i]*vol;

        pc[0] = Ji[0];
        pc[1] = Ji[1];
        pc[2] = Ji[2];
        pc[3] = Ji[3];
        pc[4] = Ji[4];
        pc[5] = Ji[5];
        pc[6] = Ji[6];
        pc[7] = Ji[7];
        pc[8] = Ji[8];
        pc[9] = detJ;

        // Build local stiffness matrix
        // Get the local derivatives of the basis functions in the basis element
        // They are all the same and are thus precomputed in matrix d
        auto Nxi = &d[i][0];
        auto Nyi = &d[i][local_dimension];
        auto Nzi = &d[i][local_dimension2];
        // Gradients associated with the node we are calculating
        const auto &Nxip = Nxi[row];
        const auto &Nyip = Nyi[row];
        const auto &Nzip = Nzi[row];
        // Calculating gradient shape function * inverse Jacobian * volume scaling factor
        const auto uxp = pc[9]*(Nxip*pc[0]+Nyip*pc[1]+Nzip*pc[2]);
        const auto uyp = pc[9]*(Nxip*pc[3]+Nyip*pc[4]+Nzip*pc[5]);
        const auto uzp = pc[9]*(Nxip*pc[6]+Nyip*pc[7]+Nzip*pc[8]);
        // Matrix multiplication with conductivity tensor :
        const auto uxyzpabc = uxp*Ca+uyp*Cb+uzp*Cc;
        const auto uxyzpbde = uxp*Cb+uyp*Cd+uzp*Ce;
        const auto uxyzpcef = uxp*Cc+uyp*Ce+uzp*Cf;

        // The above is constant for this node. Now multiply with the weight function
        // We assume the weight factors are the same as the local gradients
        // Galerkin approximation:

        for (int j = 0; j<local_dimension; j++)
        {
          const auto &Nxj = Nxi[j];
          const auto &Nyj = Nyi[j];
          const auto &Nzj = Nzi[j];

          // Matrix multiplication Gradient with inverse Jacobian:
          const auto ux = Nxj*pc[0]+Nyj*pc[1]+Nzj*pc[2];
          const auto uy = Nxj*pc[3]+Nyj*pc[4]+Nzj*pc[5];
          const auto uz = Nxj*pc[6]+Nyj*pc[7]+Nzj*pc[8];

          // Add everything together into one coefficient of the matrix
          l_stiff[j] += ux*uxyzpabc+uy*uxyzpbde+uz*uxyzpcef;
        }
      }
    }
    else
    {
      for(int i=0; i<local_dimension; i++)
        l_stiff[i] = 0.0;

      auto local_dimension2=2*local_dimension;

      for (size_t i = 0; i < d.size(); i++)
      {
        auto& pc = precompute[i];

        // Build local stiffness matrix
        // Get the local derivatives of the basis functions in the basis element
        // They are all the same and are thus precomputed in matrix d
        auto Nxi = &d[i][0];
        auto Nyi = &d[i][local_dimension];
        auto Nzi = &d[i][local_dimension2];
        // Gradients associated with the node we are calculating
        const auto &Nxip = Nxi[row];
        const auto &Nyip = Nyi[row];
        const auto &Nzip = Nzi[row];
        // Calculating gradient shape function * inverse Jacobian * volume scaling factor
        const auto uxp = pc[9]*(Nxip*pc[0]+Nyip*pc[1]+Nzip*pc[2]);
        const auto uyp = pc[9]*(Nxip*pc[3]+Nyip*pc[4]+Nzip*pc[5]);
        const auto uzp = pc[9]*(Nxip*pc[6]+Nyip*pc[7]+Nzip*pc[8]);
        // Matrix multiplication with conductivity tensor :
        const auto uxyzpabc = uxp*Ca+uyp*Cb+uzp*Cc;
        const auto uxyzpbde = uxp*Cb+uyp*Cd+uzp*Ce;
        const auto uxyzpcef = uxp*Cc+uyp*Ce+uzp*Cf;

        // The above is constant for this node. Now multiply with the weight function
        // We assume the weight factors are the same as the local gradients
        // Galerkin approximation:

        for (int j = 0; j<local_dimension; j++)
        {
          const auto &Nxj = Nxi[j];
          const auto &Nyj = Nyi[j];
          const auto &Nzj = Nzi[j];

          // Matrix multiplication Gradient with inverse Jacobian:
          const auto ux = Nxj*pc[0]+Nyj*pc[1]+Nzj*pc[2];
          const auto uy = Nxj*pc[3]+Nyj*pc[4]+Nzj*pc[5];
          const auto uz = Nxj*pc[6]+Nyj*pc[7]+Nzj*pc[8];

          // Add everything together into one coefficient of the matrix
          l_stiff[j] += ux*uxyzpabc+uy*uxyzpbde+uz*uxyzpcef;
        }
      }
    }
  }
  return true;
}

template <typename T>
bool
FEMBuilder<T>::setup()
{
  // The domain dimension
  domain_dimension = mesh_->dimensionality();
  if (domain_dimension < 1)
  {
    algo_->error("This mesh type cannot be used for FE computations");
    return false;
  }

  local_dimension_nodes = mesh_->num_nodes_per_elem();
  if (field_->basis_order() == 2)
  {
    local_dimension_add_nodes = mesh_->num_enodes_per_elem();
  }
  else
  {
    local_dimension_add_nodes = 0;
  }

  local_dimension_derivatives = 0;

  // Local degrees of freedom per element
  local_dimension = local_dimension_nodes +
  local_dimension_add_nodes +
  local_dimension_derivatives; ///< degrees of freedom (dofs) of system

  VMesh::Node::size_type mns;
  mesh_->size(mns);

  // Number of mesh points (not necessarily number of nodes)
  global_dimension_nodes = mns;
  if (field_->basis_order() == 2) // quadratic
  {
    mesh_->synchronize(Mesh::ENODES_E);
    global_dimension_add_nodes = mesh_->num_enodes();
  }
  else
  {
    global_dimension_add_nodes = 0;
  }

  global_dimension_derivatives = 0;
  global_dimension = global_dimension_nodes+
  global_dimension_add_nodes+
  global_dimension_derivatives;

  if (mns > 0)
  {
    // We only need edges for the higher order basis in case of quartic Lagrangian
    // Hence we should only synchronize it for this case
    if (global_dimension_add_nodes > 0)
      mesh_->synchronize(Mesh::EDGES_E|Mesh::NODE_NEIGHBORS_E);
    else
      mesh_->synchronize(Mesh::NODE_NEIGHBORS_E);
  }
  else
  {
    algo_->error("Mesh size < 0");
    success_[0] = false;
  }
  LOG_DEBUG("Allocating buffer for nonzero row indices of size: {}", global_dimension+1);
  rows_.reset(new index_type[global_dimension+1]);

  colidx_.resize(numprocessors_+1);
  return true;
}

// -- callback routine to execute in parallel
template <typename T>
void
FEMBuilder<T>::parallel(int proc_num)
{
  success_[proc_num] = true;

  if (proc_num == 0)
  {
    try
    {
      success_[proc_num] = setup();
    }
    catch (...)
    {
      algo_->error("BuildFEMatrix could not setup FE Stiffness computation");
      success_[proc_num] = false;
    }
  }

  barrier_.wait();

  // In case one of the threads fails, we should have them fail all
  for (int q = 0; q < numprocessors_; q++)
  {
    if (!success_[q])
    {
      std::ostringstream oss;
      oss << "FEMBuilder::setup failed in thread " << q;
      algo_->error(oss.str());
      return;
    }
  }

  /// distributing dofs among processors
  const index_type start_gd = (global_dimension * proc_num)/numprocessors_;
  const index_type end_gd  = (global_dimension * (proc_num+1))/numprocessors_;

  /// creating sparse matrix structure
  std::vector<index_type> mycols;

  VMesh::Elem::array_type ca;
  VMesh::Node::array_type na;
  VMesh::Edge::array_type ea;
  std::vector<index_type> neib_dofs;

  /// loop over system dofs for this thread
  int cnt = 0;
  size_type size_gd = end_gd-start_gd;
  auto updateFrequency = 2*size_gd / 100;
  try
  {
    mycols.reserve((end_gd - start_gd)*local_dimension*8);  //<! rough estimate

    for (VMesh::Node::index_type i = start_gd; i<end_gd; ++i)
    {
      rows_[i] = mycols.size();

      neib_dofs.clear();
      /// check for nodes
      if (i < global_dimension_nodes)
      {
        /// get neighboring cells for node
        mesh_->get_elems(ca, i);
      }
      else if (i < global_dimension_nodes+global_dimension_add_nodes)
      {
        /// check for additional nodes at edges
        /// get neighboring cells for node
        VMesh::Edge::index_type ii(i-global_dimension_nodes);
        mesh_->get_elems(ca,ii);
      }
      else
      {
        // There is some functionality implemented for higher order basis functions,
        // but it seems not to be accessible, entirely implemented nor validated.
        algo_->warning("BuildFEMatrix only supports linear basis functions.");
      }

      for(size_t j = 0; j < ca.size(); j++)
      {
        /// get neighboring nodes
        mesh_->get_nodes(na, ca[j]);

        for(size_t k = 0; k < na.size(); k++)
        {
          neib_dofs.push_back(static_cast<index_type>(na[k]));
        }

        /// check for additional nodes at edges
        if (global_dimension_add_nodes)
        {
          /// get neighboring edges
          mesh_->get_edges(ea, ca[j]);

          for(size_t k = 0; k < ea.size(); k++)
            neib_dofs.push_back(global_dimension + ea[k]);
        }
      }

      std::sort(neib_dofs.begin(), neib_dofs.end());

      for (size_t j=0; j<neib_dofs.size(); j++)
      {
        if (j == 0 || neib_dofs[j] != mycols.back())
        {
          mycols.push_back(neib_dofs[j]);
        }
      }
      if (proc_num == 0)
      {
        cnt++;
        if (cnt == updateFrequency)
        {
          cnt = 0;
          algo_->update_progress_max(i,2*size_gd);
        }
      }
    }

    colidx_[proc_num] = mycols.size();
    success_[proc_num] = true;
  }
  catch (...)
  {
    algo_->error("BuildFEMatrix crashed mapping out stiffness matrix");
    success_[proc_num] = false;
  }

  /// check point
  barrier_.wait();

  // Bail out if one of the processes failed
  for (int q=0; q<numprocessors_;q++)
  {
    if (!success_[q])
    {
      return;
    }
  }

  std::vector<std::vector<T>> precompute;
  index_type st = 0;

  if (proc_num == 0)
    allcols_.reset();

  try
  {
    if (proc_num == 0)
    {
      for(int i=0; i<numprocessors_; i++)
      {
        const index_type ns = colidx_[i];
        colidx_[i] = st;
        st += ns;
      }

      colidx_[numprocessors_] = st;
      allcols_.reset(new index_type[st]);
    }
    success_[proc_num] = true;
  }
  catch (...)
  {
    if (proc_num == 0)
      allcols_.reset();

    algo_->error("Could not allocate enough memory");
    success_[proc_num] = false;
  }

  /// check point
  barrier_.wait();

  // Bail out if one of the processes failed
  for (int q=0; q<numprocessors_;q++)
  {
    if (! success_[q])
      return;
  }

  try
  {
    /// updating global column by each of the processors
    const index_type s = colidx_[proc_num];
    const size_t n = mycols.size();

    for(size_t i=0; i<n; i++)
      allcols_[i+s] = mycols[i];

    for(index_type i = start_gd; i<end_gd; i++)
      rows_[i] += s;

    success_[proc_num] = true;
  }
  catch (...)
  {
    algo_->error("BuildFEMatrix crashed while setting up row compression");
    success_[proc_num] = false;
  }

  /// check point
  barrier_.wait();

  // Bail out if one of the processes failed
  for (auto q=0; q<numprocessors_; q++)
  {
    if (!success_[q])
      return;
  }

  try
  {
    /// the main thread makes the matrix
    if (proc_num == 0)
    {
      rows_[global_dimension] = st;
      algo_->remark("Creating fematrix on main thread.");
      fematrix_ = makeShared<matrix_type<T>>(global_dimension, global_dimension, rows_.get(), allcols_.get(), st);
      rows_.reset();
      allcols_.reset();
    }
    success_[proc_num] = true;
  }
  catch (...)
  {
    algo_->error("BuildFEMatrix crashed while creating final stiffness matrix");
    success_[proc_num] = false;
  }

  /// check point
  barrier_.wait();

  // Bail out if one of the processes failed
  for (auto q=0; q<numprocessors_;q++)
  {
    if (!success_[q])
      return;
  }

  try
  {
    /// zeroing in parallel
    const auto ns = colidx_[proc_num];
    const auto ne = colidx_[proc_num+1];
    auto a = &(fematrix_->valuePtr()[ns]), ae=&(fematrix_->valuePtr()[ne]);
    while (a<ae) *a++=0.0;

    std::vector<VMesh::coords_type> ni_points;
    std::vector<double> ni_weights;
    std::vector<std::vector<double>> ni_derivatives;

    create_numerical_integration(ni_points, ni_weights, ni_derivatives);

    std::vector<T> lsml; ///< line of local stiffnes matrix
    lsml.resize(local_dimension);

    /// loop over system dofs for this thread
    cnt = 0;
    size_gd = end_gd-start_gd;
    for (VMesh::Node::index_type i = start_gd; i<end_gd; ++i)
    {
      if (i < global_dimension_nodes)
      {
        /// check for nodes
        /// get neighboring cells for node
        mesh_->get_elems(ca,i);
      }
      else if (i < global_dimension_nodes + global_dimension_add_nodes)
      {
        /// check for additional nodes at edges
        /// get neighboring cells for additional nodes
        VMesh::Edge::index_type ii(i-global_dimension_nodes);
        mesh_->get_elems(ca,ii);
      }
      else
      {
        // There is some functionality implemented for higher order basis functions,
        // but it seems not to be accessible, entirely implemented nor validated.
        algo_->warning("BuildFEMatrix only supports linear basis functions.");
      }

      /// loop over elements attributed elements

      if (mesh_->is_regularmesh())
      {
        for (size_t j = 0; j < ca.size(); j++)
        {
          mesh_->get_nodes(na, ca[j]); ///< get neighboring nodes
          neib_dofs.resize(na.size());
          for(size_t k = 0; k < na.size(); k++)
          {
            neib_dofs[k] = na[k]; // Must cast to (int) for SGI compiler :-(
          }

          for(size_t k = 0; k < na.size(); k++)
          {
            if (na[k] == i)
            {
              auto successLocal = build_local_matrix_regular(ca[j], k , lsml, ni_points, ni_weights, ni_derivatives,precompute);
							if (!successLocal)
							{
								success_[proc_num] = false;
								return;
							}
              add_lcl_gbl(i, neib_dofs, lsml);
            }
          }
        }
      }
      else
      {
        for (size_t j = 0; j < ca.size(); j++)
        {
          neib_dofs.clear();
          mesh_->get_nodes(na, ca[j]); ///< get neighboring nodes
          for(size_t k = 0; k < na.size(); k++)
          {
            neib_dofs.push_back(na[k]); // Must cast to (int) for SGI compiler :-(
          }
          /// check for additional nodes at edges
          if (global_dimension_add_nodes)
          {
            mesh_->get_edges(ea, ca[j]); ///< get neighboring edges
            for(size_t k = 0; k < ea.size(); k++)
            {
              neib_dofs.push_back(global_dimension + ea[k]);
            }
          }

          ASSERT(static_cast<int>(neib_dofs.size()) == local_dimension);

          for(size_t k = 0; k < na.size(); k++)
          {
            if (na[k] == i)
            {
              auto successLocal = build_local_matrix(ca[j], k , lsml, ni_points, ni_weights, ni_derivatives);
							if (!successLocal)
							{
								success_[proc_num] = false;
								return;
							}
              add_lcl_gbl(i, neib_dofs, lsml);
            }
          }

          if (global_dimension_add_nodes)
          {
            for (size_t k = 0; k < ea.size(); k++)
            {
              if (global_dimension + static_cast<int>(ea[k]) == i)
              {
                auto successLocal = build_local_matrix(ca[j], k+na.size(), lsml, ni_points, ni_weights, ni_derivatives);
								if (!successLocal)
								{
									success_[proc_num] = false;
									return;
								}
                add_lcl_gbl(i, neib_dofs, lsml);
              }
            }
          }
        }
      }

      if (proc_num == 0)
      {
        cnt++;
        if (cnt == updateFrequency)
        {
          cnt = 0;
          algo_->update_progress_max(i+size_gd,2*size_gd);
        }
      }
    }
    success_[proc_num] = true;
  }
  catch (...)
  {
    algo_->error("BuildFEMatrix crashed while filling out stiffness matrix");
    success_[proc_num] = false;
  }

  barrier_.wait();

  // Bail out if one of the processes failed
  for (int q=0; q<numprocessors_; q++)
  {
    if (!success_[q])
      return;
  }
}

const AlgorithmParameterName BuildFEMatrixAlgo::ForceSymmetry("ForceSymmetry");
const AlgorithmParameterName BuildFEMatrixAlgo::GenerateBasis("GenerateBasis");

template <typename T>
bool
BuildFEMatrixAlgoImpl<T>::run(FieldHandle input, DenseMatrixHandle ctable, matrix_pointer_type<T>& output) const
{
  ScopedAlgorithmStatusReporter s(algo_, "BuildFEMatrix");

  if (!input)
  {
    algo_->error("Could not obtain input field");
    return false;
  }

  if (input->vfield()->is_vector())
  {
    algo_->error("This function has not yet been defined for elements with vector data");
    return false;
  }

  if (input->vfield()->basis_order()!=0)
  {
    algo_->error("This function has only been defined for data that is located at the elements");
    return false;
  }

  if (ctable)
  {
    if ((ctable->ncols() != 1)&&(ctable->ncols() != 6)&&(ctable->ncols() != 9))
    {
      algo_->error("Conductivity table needs to have 1, 6, or 9 columns");
      return false;
    }
    if (ctable->nrows() == 0)
    {
      algo_->error("ConductivityTable is empty");
      return false;
    }
  }

  FEMBuilder<T> builder(algo_);

  if (algo_->get(BuildFEMatrixAlgo::GenerateBasis).toBool())
  {
    if (!ctable)
    {
      std::vector<std::pair<std::string,Tensor> > tens;

      input->properties().get_property("conductivity_table",tens);

      if (!tens.empty())
      {
        ctable.reset(new DenseMatrix(tens.size(), 1));
        auto data = ctable->data();
        for (size_t i=0; i<tens.size();i++)
        {
          auto t = tens[i].second.val(0,0);
          data[i] = t;
        }
      }
    }

    if (ctable)
    {
      auto nconds = ctable->nrows();
      if ( (input->vmesh()->generation() != generation_) ||
          (!basis_fematrix_) )
      {
        auto con = makeShared<DenseMatrix>(nconds, 1, 0.0);
        auto data = con->data();

        if (!builder.build_matrix(input, con, basis_fematrix_) )
        {
          algo_->error("Build matrix method failed when building FEMatrix structure");
          return false;
        }

        if (!basis_fematrix_)
        {
          algo_->error("Failed to build FEMatrix structure");
          return false;
        }

        basis_values_.resize(nconds);
        for (size_type i=0; i < nconds; i++)
        {
          matrix_pointer_type<T> stiffness;
          /// @todo: can initialize array using std::fill
          data[i] = 1.0;

          if (!builder.build_matrix(input, con, stiffness) )
          {
            algo_->error("Build matrix method failed for one of the tissue types");
            return false;
          }

          if (!stiffness)
          {
            algo_->error("Failed to build FEMatrix component for one of the tissue types");
            return false;
          }

          basis_values_[i].resize(stiffness->nonZeros());
          for (size_type p=0; p< stiffness->nonZeros(); p++)
          {
            basis_values_[i][p] = stiffness->valuePtr()[p];
          }
          data[i] = 0.0;
        }

        generation_ = input->vmesh()->generation();
      }

      output.reset(basis_fematrix_->clone());

      auto sum = output->valuePtr();
      auto cdata = ctable->data();
      auto n = ctable->ncols();

      if (!basis_values_.empty())
      {
        for (size_t p=0; p < basis_values_[0].size(); p++)
          sum[p] = 0.0;
      }

      for (auto i=0; i<nconds; i++)
      {
        auto weight = cdata[i*n];
        for (size_t p=0; p < basis_values_[i].size(); p++)
        {
          sum[p] += weight * basis_values_[i][p];
        }
      }

    }
    else
    {
      algo_->error("No conductivity table present: The generate_basis option only works for indexed conductivities");
      return false;
    }
  }

  if (!builder.build_matrix(input,ctable,output) )
  {
    algo_->error("Build matrix method failed to build output matrix");
    return false;
  }

  if (!output)
  {
    algo_->error("Could not build output matrix");
    return false;
  }

  return true;
}

const AlgorithmInputName BuildFEMatrixAlgo::Conductivity_Table("Conductivity_Table");
const AlgorithmOutputName BuildFEMatrixAlgo::Stiffness_Matrix("Stiffness_Matrix");
const AlgorithmOutputName BuildFEMatrixAlgo::Stiffness_Matrix_Complex("Stiffness_Matrix_Complex");

AlgorithmOutput BuildFEMatrixAlgo::run(const AlgorithmInput& input) const
{
  auto field = input.get<Field>(Variables::InputField);
  auto ctable = input.get<DenseMatrix>(Conductivity_Table);

	AlgorithmOutput output;
  if (field && field->vfield() && field->vfield()->is_complex_double())
	{
		matrix_pointer_type<complex> stiffness;
	  BuildFEMatrixAlgoImpl<complex> impl(this);
	  if (!impl.run(field, ctable, stiffness))
	    THROW_ALGORITHM_PROCESSING_ERROR("False returned on legacy run call.--complex detected	");
		output[Stiffness_Matrix_Complex] = stiffness;
	}
	else
	{
		matrix_pointer_type<double> stiffness;
	  BuildFEMatrixAlgoImpl<double> impl(this);
	  if (!impl.run(field, ctable, stiffness))
	    THROW_ALGORITHM_PROCESSING_ERROR("False returned on legacy run call.");
		output[Stiffness_Matrix] = stiffness;
	}

  return output;
}


#ifndef AIKIDO_CONSTRAINT_DART_TSR_HPP_
#define AIKIDO_CONSTRAINT_DART_TSR_HPP_

#include <Eigen/Dense>
#include <dart/math/MathTypes.hpp>

#include "aikido/constraint/Differentiable.hpp"
#include "aikido/constraint/Projectable.hpp"
#include "aikido/constraint/Sampleable.hpp"
#include "aikido/constraint/Testable.hpp"
#include "aikido/statespace/SE3.hpp"

namespace aikido {
namespace constraint {
namespace dart {

AIKIDO_DECLARE_POINTERS(TSR)

/// TSRs describe end-effector constraint sets as subsets of SE(3).
/// A TSR consists of three parts:
///     T0_w: transform from the origin to the TSR frame w
///     B_w: 6 × 2 matrix of bounds in the coordinates of w.
///     Tw_e: end-effector offset transform in the coordinates of w
/// See:
/// Berenson, Dmitry, Siddhartha S. Srinivasa, and James Kuffner.
/// "Task space regions: A framework for pose-constrained manipulation
/// planning." IJRR 2001:
/// http://repository.cmu.edu/cgi/viewcontent.cgi?article=2024&context=robotics
class TSR
  : public Sampleable
  , public Differentiable
  , public Testable
  , public Projectable
{
public:
  EIGEN_MAKE_ALIGNED_OPERATOR_NEW

  using Differentiable::getValueAndJacobian;

  /// Constructor.
  /// \param _rng Random number generator used by SampleGenerators for this TSR.
  /// \param _T0_w transform from the origin to the TSR frame w
  /// \param _Bw 6 × 2 matrix of bounds in the coordinates of w.
  ///        Top three rows bound translation, and bottom three rows
  ///        bound rotation following Roll-Pitch-Yaw convention.
  ///        _Bw(i, 0) should be less than or equal to _Bw(i, 1).
  /// \param _Tw_e end-effector offset transform in the coordinates of w
  /// \param _testableTolerance tolerance used in isSatisfiable as testable
  TSR(std::unique_ptr<common::RNG> _rng,
      const Eigen::Isometry3d& _T0_w = Eigen::Isometry3d::Identity(),
      const Eigen::Matrix<double, 6, 2>& _Bw
      = Eigen::Matrix<double, 6, 2>::Zero(),
      const Eigen::Isometry3d& _Tw_e = Eigen::Isometry3d::Identity(),
      double _testableTolerance = 1e-6);

  /// Constructor with default random seed generator.
  /// \param _T0_w transform from the origin to the TSR frame w
  /// \param _Bw 6 × 2 matrix of bounds in the coordinates of w.
  ///        Top three rows bound translation, and bottom three rows
  ///        bound rotation following Roll-Pitch-Yaw convention.
  ///        _Bw(i, 0) should be less than or equal to _Bw(i, 1).
  /// \param _Tw_e end-effector offset transform in the coordinates of w
  /// \param _testableTolerance tolerance used in isSatisfiable as testable
  TSR(const Eigen::Isometry3d& _T0_w = Eigen::Isometry3d::Identity(),
      const Eigen::Matrix<double, 6, 2>& _Bw
      = Eigen::Matrix<double, 6, 2>::Zero(),
      const Eigen::Isometry3d& _Tw_e = Eigen::Isometry3d::Identity(),
      double _testableTolerance = 1e-6);

  TSR(const TSR& other);
  TSR(TSR&& other);

  TSR& operator=(const TSR& other);
  TSR& operator=(TSR&& other);

  virtual ~TSR() = default;

  // Documentation inherited.
  statespace::ConstStateSpacePtr getStateSpace() const override;

  /// Returns the SE3 which this TSR operates in.
  std::shared_ptr<statespace::SE3> getSE3() const;

  // Documentation inherited.
  std::unique_ptr<SampleGenerator> createSampleGenerator() const override;

  // Documentation inherited.
  bool isSatisfied(
      const statespace::StateSpace::State* _s,
      TestableOutcome* outcome = nullptr) const override;

  /// Return an instance of DefaultTestableOutcome, since this class doesn't
  /// have a more specialized TestableOutcome derivative assigned to it.
  std::unique_ptr<TestableOutcome> createOutcome() const override;

  /// Throws an invalid_argument exception if this TSR is invalid.
  /// For a TSR to be valid, mBw(i, 0) <= mBw(i, 1).
  void validate() const;

  /// Set the random number generator used by SampleGenerators for this TSR.
  void setRNG(std::unique_ptr<common::RNG> rng);

  // Documentation inherited.
  std::size_t getConstraintDimension() const override;

  // Documentation inherited.
  void getValue(const statespace::StateSpace::State* _s, Eigen::VectorXd& _out)
      const override;

  /// Jacobian of TSR with respect to the se(3) tangent vector of _s.
  /// The jacobian is w.r.t. the origin frame.
  /// se(3) tangent vector follows dart convention:
  ///   top 3 rows is the angle-axis representation of _s's rotation.
  ///   bottom 3 rows represent the translation.
  /// \param _s State to be evaluated at.
  /// \param[out] _out Jacobian, 6 x 6 matrix.
  void getJacobian(
      const statespace::StateSpace::State* _s,
      Eigen::MatrixXd& _out) const override;

  // Documentation inherited.
  std::vector<ConstraintType> getConstraintTypes() const override;

  // Documentation inherited.
  bool project(
      const statespace::StateSpace::State* _s,
      statespace::StateSpace::State* _out) const override;

  /// Get the testable tolerance used in isSatisfiable.
  double getTestableTolerance();

  /// Set the testable tolerance used in isSatisfiable.
  /// \param _testableTolerance Testable tolerance to set.
  void setTestableTolerance(double _testableTolerance);

  /// Transformation from origin frame into the TSR frame "w".
  /// "w" is usually centered at the origin of an object held by the hand
  ///  or at a location on an object that is useful for grasping.
  Eigen::Isometry3d mT0_w;

  /// Bounds on "wiggling" in `x, y, z, roll, pitch, yaw`.
  Eigen::Matrix<double, 6, 2> mBw;

  /// Transformation from "w" frame into end frame.
  /// This often represent an offset from "w" to the origin of the end-effector.
  Eigen::Isometry3d mTw_e;

private:
  /// Tolerance used in isSatisfied as a testable
  double mTestableTolerance;
  std::unique_ptr<common::RNG> mRng;
  std::shared_ptr<statespace::SE3> mStateSpace;
};

} // namespace dart
} // namespace constraint
} // namespace aikido

#endif // AIKIDO_CONSTRAINT_DART_TSR_HPP_


/*
 Copyright (C) 2016 Quaternion Risk Management Ltd
 All rights reserved.

 This file is part of ORE, a free-software/open-source library
 for transparent pricing and risk analysis - http://opensourcerisk.org

 ORE is free software: you can redistribute it and/or modify it
 under the terms of the Modified BSD License.  You should have received a
 copy of the license along with this program.
 The license is also available online at <http://opensourcerisk.org>

 This program is distributed on the basis that it will form a useful
 contribution to risk analytics and model standardisation, but WITHOUT
 ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or
 FITNESS FOR A PARTICULAR PURPOSE. See the license for more details.
*/

#include <qle/methods/multipathgeneratorbase.hpp>

#include <boost/make_shared.hpp>

using namespace QuantLib;

namespace QuantExt {

MultiPathGeneratorMersenneTwister::MultiPathGeneratorMersenneTwister(
    const boost::shared_ptr<StochasticProcess>& process, const TimeGrid& grid, BigNatural seed, bool antitheticSampling)
    : process_(process), grid_(grid), seed_(seed), antitheticSampling_(antitheticSampling), antitheticVariate_(true) {
    reset();
}

void MultiPathGeneratorMersenneTwister::reset() {
    PseudoRandom::rsg_type rsg = PseudoRandom::make_sequence_generator(process_->size() * (grid_.size() - 1), seed_);
    pg_ = boost::make_shared<MultiPathGenerator<PseudoRandom::rsg_type> >(process_, grid_, rsg, false);
    antitheticVariate_ = true;
}

MultiPathGeneratorSobol::MultiPathGeneratorSobol(const boost::shared_ptr<StochasticProcess>& process,
                                                 const TimeGrid& grid, BigNatural seed,
                                                 SobolRsg::DirectionIntegers directionIntegers)
    : process_(process), grid_(grid), seed_(seed), directionIntegers_(directionIntegers) {
    reset();
}

void MultiPathGeneratorSobol::reset() {
    pg_ = boost::make_shared<MultiPathGenerator<InverseCumulativeRsg<SobolRsg, InverseCumulativeNormal> > >(
        process_, grid_,
        InverseCumulativeRsg<SobolRsg, InverseCumulativeNormal>(
            SobolRsg(process_->size() * (grid_.size() - 1), seed_, directionIntegers_)));
}

MultiPathGeneratorSobolBrownianBridge::MultiPathGeneratorSobolBrownianBridge(
    const boost::shared_ptr<StochasticProcess>& process, const TimeGrid& grid,
    SobolBrownianGenerator::Ordering ordering, BigNatural seed, SobolRsg::DirectionIntegers directionIntegers)
    : process_(process), grid_(grid), ordering_(ordering), seed_(seed), directionIntegers_(directionIntegers),
      next_(MultiPath(process->size(), grid), 1.0) {
    reset();
}

void MultiPathGeneratorSobolBrownianBridge::reset() {
    gen_ = boost::make_shared<SobolBrownianGenerator>(process_->size(), grid_.size() - 1, ordering_, seed_,
                                                      directionIntegers_);
}

const Sample<MultiPath>& MultiPathGeneratorSobolBrownianBridge::next() const {
    Array asset = process_->initialValues();
    MultiPath& path = next_.value;
    for (Size j = 0; j < asset.size(); ++j) {
        path[j].front() = asset[j];
    }
    next_.weight = gen_->nextPath();
    std::vector<Real> output(asset.size());
    for (Size i = 1; i < grid_.size(); ++i) {
        Real t = grid_[i - 1];
        Real dt = grid_.dt(i - 1);
        gen_->nextStep(output);
        Array tmp(output.begin(), output.end());
        asset = process_->evolve(t, asset, dt, tmp);
        for (Size j = 0; j < asset.size(); ++j) {
            path[j][i] = asset[j];
        }
    }
    return next_;
}

boost::shared_ptr<MultiPathGeneratorBase> makeMultiPathGenerator(const SequenceType s,
                                                                 const boost::shared_ptr<StochasticProcess>& process,
                                                                 const TimeGrid& timeGrid, const BigNatural seed,
                                                                 const SobolBrownianGenerator::Ordering ordering,
                                                                 const SobolRsg::DirectionIntegers directionIntegers) {
    switch (s) {
    case MersenneTwister:
        return boost::make_shared<QuantExt::MultiPathGeneratorMersenneTwister>(process, timeGrid, seed, false);
    case MersenneTwisterAntithetic:
        return boost::make_shared<QuantExt::MultiPathGeneratorMersenneTwister>(process, timeGrid, seed, true);
    case Sobol:
        return boost::make_shared<QuantExt::MultiPathGeneratorSobol>(process, timeGrid, seed, directionIntegers);
    case SobolBrownianBridge:
        return boost::make_shared<QuantExt::MultiPathGeneratorSobolBrownianBridge>(process, timeGrid, ordering, seed,
                                                                                directionIntegers);
    default:
        QL_FAIL("Unknown sequence type");
    }
}

std::ostream& operator<<(std::ostream& out, const SequenceType s) {
    switch (s) {
    case MersenneTwister:
        return out << "MersenneTwister";
    case MersenneTwisterAntithetic:
        return out << "MersenneTwisterAntithetic";
    case Sobol:
        return out << "Sobol";
    case SobolBrownianBridge:
        return out << "SobolBrownianBridge";
    default:
        return out << "Unknown sequence type";
    }
}

} // namespace QuantExt


//---------------------------------------------------------------------------//
// Copyright (c) 2018-2021 Mikhail Komarov <nemo@nil.foundation>
// Copyright (c) 2020-2021 Nikita Kaskov <nbering@nil.foundation>
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//    http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
//---------------------------------------------------------------------------//

#include <iostream>

#include <boost/filesystem.hpp>
#include <boost/program_options.hpp>

#include "detail/r1cs_examples.hpp"
#include "detail/sha256_component.hpp"

#include <nil/crypto3/algebra/curves/bls12.hpp>
#include <nil/crypto3/algebra/fields/bls12/base_field.hpp>
#include <nil/crypto3/algebra/fields/bls12/scalar_field.hpp>
#include <nil/crypto3/algebra/fields/arithmetic_params/bls12.hpp>
#include <nil/crypto3/algebra/curves/params/multiexp/bls12.hpp>
#include <nil/crypto3/algebra/curves/params/wnaf/bls12.hpp>

#include <nil/crypto3/zk/components/blueprint.hpp>
#include <nil/crypto3/zk/components/blueprint_variable.hpp>

#include <nil/crypto3/zk/snark/schemes/ppzksnark/r1cs_gg_ppzksnark.hpp>
#include <nil/crypto3/zk/snark/schemes/ppzksnark/r1cs_gg_ppzksnark/marshalling.hpp>

#include <nil/crypto3/zk/snark/algorithms/generate.hpp>
#include <nil/crypto3/zk/snark/algorithms/verify.hpp>
#include <nil/crypto3/zk/snark/algorithms/prove.hpp>

#include <nil/marshalling/status_type.hpp>

using namespace nil::crypto3;
using namespace nil::crypto3::zk;

typedef algebra::curves::bls12<381> curve_type;
typedef typename curve_type::scalar_field_type field_type;

typedef zk::snark::r1cs_gg_ppzksnark<curve_type> scheme_type;

int main(int argc, char *argv[]) {
    boost::filesystem::path pout, pkout, vkout, piout, viout;
    boost::program_options::options_description options(
        "R1CS Generic Group PreProcessing Zero-Knowledge Succinct Non-interactive ARgument of Knowledge "
        "(https://eprint.iacr.org/2016/260.pdf) CLI Proof Generator");
    // clang-format off
    options.add_options()("help,h", "Display help message")
    ("version,v", "Display version")
    ("generate", "Generate proofs and/or keys")
    ("verify", "verify proofs and/or keys")
    ("proof-output,po", boost::program_options::value<boost::filesystem::path>(&pout)->default_value("proof"))
    ("primary-input-output,pio", boost::program_options::value<boost::filesystem::path>(&piout)->default_value
("pinput"))
    ("proving-key-output,pko", boost::program_options::value<boost::filesystem::path>(&pkout)->default_value("pkey"))
    ("verifying-key-output,vko", boost::program_options::value<boost::filesystem::path>(&viout)->default_value("vkey"))
    ("verifier-input-output,vio", boost::program_options::value<boost::filesystem::path>(&vkout)->default_value("vio"));
    // clang-format on

    boost::program_options::variables_map vm;
    boost::program_options::store(boost::program_options::command_line_parser(argc, argv).options(options).run(), vm);
    boost::program_options::notify(vm);

    if (vm.count("help") || argc < 2) {
        std::cout << options << std::endl;
        return 0;
    }

    std::cout << "SHA2-256 blueprint generation started." << std::endl;

    components::blueprint<field_type> bp = sha2_two_to_one_bp<field_type>();

    std::cout << "SHA2-256 blueprint generation finished." << std::endl;

    std::cout << "R1CS generation started." << std::endl;

    r1cs_example<field_type> example =
        r1cs_example<field_type>(bp.get_constraint_system(), bp.primary_input(), bp.auxiliary_input());

    std::cout << "R1CS generation finished." << std::endl;

    // const bool bit = run_r1cs_gg_ppzksnark<curve_type>(example);

    // zk::snark::detail::r1cs_example<field_type> example =
    //     zk::snark::detail::r1cs_example<field_type>(bp.get_constraint_system(), bp.primary_input(),
    //     bp.auxiliary_input());

    // zk::snark::r1cs_constraint_system<field_type> constraint_system = bp.get_constraint_system();

    std::cout << "Starting generator" << std::endl;

    typename scheme_type::keypair_type keypair = zk::snark::generate<scheme_type>(example.constraint_system);

    std::cout << "Starting prover" << std::endl;

    const typename scheme_type::proof_type proof =
        prove<scheme_type>(keypair.first, example.primary_input, example.auxiliary_input);

    // std::cout << "Starting verifier" << std::endl;

    // const bool ans = verify<basic_proof_system>(keypair.second, example.primary_input, proof);

    // std::cout << "Verifier finished, result: " << ans << std::endl;

    std::vector<std::uint8_t> proving_key_byteblob =
        nil::marshalling::verifier_input_serializer_tvm<scheme_type>::process(keypair.first);
    std::vector<std::uint8_t> verification_key_byteblob =
        nil::marshalling::verifier_input_serializer_tvm<scheme_type>::process(keypair.second);
    std::vector<std::uint8_t> proof_byteblob =
        nil::marshalling::verifier_input_serializer_tvm<scheme_type>::process(proof);
    std::vector<std::uint8_t> primary_input_byteblob =
        nil::marshalling::verifier_input_serializer_tvm<scheme_type>::process(example.primary_input);

    if (vm.count("proving-key-output")) {
        boost::filesystem::ofstream out(pkout);
        for (const auto &v : proving_key_byteblob) {
            out << v;
        }
        out.close();
    }

    if (vm.count("verifying-key-output")) {
        boost::filesystem::ofstream out(vkout);
        for (const auto &v : verification_key_byteblob) {
            out << v;
        }
        out.close();
    }

    if (vm.count("proof-output")) {
        boost::filesystem::ofstream out(pout);
        for (const auto &v : proof_byteblob) {
            out << v;
        }
        out.close();
    }

    if (vm.count("primary-input-output")) {
        boost::filesystem::ofstream out(piout);
        for (const auto &v : primary_input_byteblob) {
            out << v;
        }
        out.close();
    }

    // nil::marshalling::status_type provingProcessingStatus = nil::marshalling::status_type::success;
    // typename scheme_type::proving_key_type other =
    //             nil::marshalling::verifier_input_deserializer_tvm<scheme_type>::proving_key_process(
    //                 proving_key_byteblob.cbegin(),
    //                 proving_key_byteblob.cend(),
    //                 provingProcessingStatus);

    // assert(keypair.first == other);

    if (vm.count("verifier-input-output")) {
        std::vector<std::uint8_t> verifier_input_output_byteblob(proof_byteblob.begin(), proof_byteblob.end());

        verifier_input_output_byteblob.insert(verifier_input_output_byteblob.end(), primary_input_byteblob.begin(),
                                              primary_input_byteblob.end());
        verifier_input_output_byteblob.insert(verifier_input_output_byteblob.end(), verification_key_byteblob.begin(),
                                              verification_key_byteblob.end());

        boost::filesystem::ofstream poutf(pout);
        for (const auto &v : verifier_input_output_byteblob) {
            poutf << v;
        }
        poutf.close();
    }

    return 0;
}


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

// ---
// Author: Zachary Kann

// A set of functions for calculation of electric field / frequency maps
// as outlined in the papers of Skinner et al. See, for example, S.M. Gruenbaum,
// et al. J. Chem. Theory Comput. 2013 9 (7), 3109. These maps are necessary
// for calculation of vibrational spectra using the methodology  outline in the
// aforementioned papers.

// Abbreviations:
//   IR = infrared
//   SFG = sum-frequency generation

#include <complex>
#include <cmath>
#include <armadillo>
#include "z_histogram.hpp"
#include "z_atom_group.hpp"
#include "z_molecule_group.hpp"
#include "z_cx_tcf.hpp"
#include "z_conversions.hpp"
#ifndef _Z_FREQUENCY_MAP_HPP_
#define _Z_FREQUENCY_MAP_HPP_

enum Chromophore {kOH, kOD};
enum ModelType {kThreeSite, kFourSite};

class FrequencyMap {
 public:
  inline FrequencyMap(const MoleculeGroup& water_group,
                      const Chromophore chromophore, const double timestep,
                      const int correlation_length = 200)
      : step_(0), correlation_length_(correlation_length),
        num_chromophores_(2*water_group.num_molecules()),
        chromophore_(chromophore), timestep_(timestep), frequency_bins_(4000),
        min_frequency_(chromophore == kOH ? 1650.0 : 850.0),
        max_frequency_(chromophore == kOH ? 4950.0 : 4150.0),
        frequency_distribution_(frequency_bins_, min_frequency_,
                                max_frequency_),
        spectral_density_(frequency_bins_, min_frequency_, max_frequency_),
        spectra_tcf_(correlation_length,timestep_, 1,
                    (4096+2)/2-correlation_length) {
    steps_guess_ = 1000;
    omega_01_ = arma::zeros<arma::mat>(num_chromophores_, steps_guess_);

    switch(chromophore) {
      case kOH:
        SetupOH();
        break;
      case kOD:
        SetupOD();
        break;
      default:
        assert(false && "Unrecognized chromophore option.");
        break;
    }
    switch(water_group.size()/water_group.num_molecules()) {
      case 3:
        model_type_ = kThreeSite;
        break;
      case 4:
        model_type_ = kFourSite;
        break;
      default:
        assert(false && "Number of atoms in water model is not one 3 or 4.");
        break;
    }
  }

  inline void CalculateFrequencies(const MoleculeGroup& hydrogen_group,
                                   const AtomGroup& oxygen_group,
                                   const arma::rowvec& box) {
    if(step_ == steps_guess_) {
      steps_guess_ += kStepsGuessIncrement_;
      ResizeArrays();
    }
    for (int i_chrom = 0; i_chrom < hydrogen_group.size(); ++i_chrom) {
      int i_oxygen = i_chrom/2;
      oh_vector_ =
          hydrogen_group.position(i_chrom) - oxygen_group.position(i_oxygen);
      oh_unit_vector_ = arma::normalise(oh_vector_);
      field_projection_ =
          arma::dot(oh_vector_, hydrogen_group.electric_field(i_chrom))*
          AU_TO_NM*AU_TO_NM;
      transition_dipole_position_ = oxygen_group.position(i_oxygen) +
                                    transition_dipole_shift_*oh_vector_;
      UseMapping(i_chrom, box);
      frequency_distribution_.Add(omega_01_(i_chrom));
      double spectral_density_weight = SpectralDensityWeight(i_chrom);
      spectral_density_.Add(omega_01_(i_chrom), spectral_density_weight);
    }
    step_++;
  }

  inline void CalculateSpectra() {
    // Two mean functions needed because of armadillo convention.
    avg_frequency_ = arma::mean(arma::mean(omega_01_));
    omega_01_ -= avg_frequency_;
    int num_corr = spectra_tcf_.CalculateNumCorr(step_);
    std::complex<double> exp_argument, exp_omega;
    for (int i_chrom = 0; i_chrom < num_chromophores_; ++i_chrom) {
      for (int i_num = 0; i_num < num_corr; ++i_num) {
        int corr_start = i_num*spectra_tcf_.interval();
        double omega_integral = 0.0;
        exp_argument = -1.0i*omega_integral;
        exp_omega = std::exp(exp_argument);
        double mu_or_alpha_product =
            MuOrAlphaProduct(i_chrom, corr_start, 0);
        spectra_tcf_.AddTCF(0,mu_or_alpha_product*exp_omega);
        for (int i_corr = 1; i_corr < correlation_length_; ++i_corr) {
          omega_integral += timestep_/2.0*TWO_PI_C *
                            (omega_01_(i_chrom,corr_start+i_corr)
                             +omega_01_(i_chrom,corr_start));
          exp_argument = -1.0i*omega_integral;
          exp_omega = std::exp(exp_argument);
          mu_or_alpha_product =
              MuOrAlphaProduct(i_chrom, corr_start, i_corr);
          spectra_tcf_.AddTCF(i_corr,mu_or_alpha_product*exp_omega);
        }
      }
    }
    for (int i_corr = 0; i_corr < correlation_length_; ++i_corr)
      spectra_tcf_.MultiplyTCF(std::exp(-i_corr*timestep_/2.0/lifetime_));
    spectra_tcf_.FourierPlus();
  }

  inline void PrintSpectra(const std::string& filename) {
    spectra_tcf_.PrintFTReal(filename, avg_frequency_);
  }

  //accessors
  inline int num_chromophores() const { return num_chromophores_; }

 protected:
  int step_;
  static const double mu_prime_0_ = 0.1646;
  static const double mu_prime_1_ = 11.39;
  static const double mu_prime_2_ = 63.41;
  int steps_guess_;
  static const int kStepsGuessIncrement_ = 1000;
  arma::mat omega_01_;
  arma::cube mu_01_;
  arma::rowvec transition_dipole_position_;
  arma::rowvec oh_unit_vector_;
  double field_projection_;
  double omega_01_0_;
  double omega_01_1_;
  double omega_01_2_;
  double chi_01_0_;
  double chi_01_1_;

  inline double SpectralDensityWeight(const int chromophore) {
    return MuOrAlphaProduct(chromophore, step_, 0);
  }

 private:
  double avg_frequency_;
  double lifetime_;
  int correlation_length_;
  double num_chromophores_;
  static const double transition_dipole_shift_ = 0.067; // Formerly dipPos
  Chromophore chromophore_;
  double timestep_;
  int frequency_bins_;
  double min_frequency_;
  double max_frequency_;
  ModelType model_type_;
  arma::mat alpha_;
  arma::mat switcher_;
  //arma::cube u; // needed for coupled spectra
  arma::rowvec oh_vector_;
  Histogram frequency_distribution_;
  Histogram spectral_density_;
  CxTCF spectra_tcf_;

  virtual void ResizeArrays() = 0;

  virtual double MuOrAlphaProduct(const int chromophore, const int corr_start,
                                  const int i_corr) = 0;

  inline void SetupOH() {
    lifetime_ = 0.700; //TODO(Zak): change this for coupled spectra
    omega_01_0_ = 3760.2;
    omega_01_1_ = -3541.7;
    omega_01_2_ = -152677.0;
    chi_01_0_ = 0.19285;
    chi_01_1_ = -1.7261E-5;
  }

  inline void SetupOD() {
    lifetime_ = 1.800; //TODO(Zak): change this for coupled spectra
    omega_01_0_ = 2767.8;
    omega_01_1_ = -2630.3;
    omega_01_2_ = -102601.0;
    chi_01_0_ = 0.16593;
    chi_01_1_ = -2.0632E-5;
  }

  // Calculates omega and mu/alpha using the frequency maps.
  virtual void UseMapping(const int chromophore,
                  const arma::rowvec& box) = 0;

  // TODO(Zak): Move this to a coupled spectra class
  // Calculates the intramolecular coupling between OH groups for inclusion
  // in the system Hamiltonian needed for coupled spectra calculations
  //extern double IntraCouple();

  // TODO(Zak): Move this to a coupled spectra class
  // Calculates the intermolecular coupling between OH groups for inclusion
  // in the system Hamiltonian needed for coupled spectra calculations
  //extern void InterCouple(arma::rowvec& omega, const arma::mat& xdipole,
  //                        const arma::rowvec& chi01, const arma::rowvec& muprime,
  //                        const arma::mat& u, const arma::rowvec& box,
  //                        const arma::icube& shift, const int numMols,
  //                        const int numChromos);
};
#endif