// -----------------------------------------------------------------------
//
// Copyright (C) 2020  - David Fernández Castellanos
//
// This file is part of the MEPLS software. You can use it, redistribute
// it, and/or modify it under the terms of the Creative Commons Attribution
// 4.0 International Public License. The full text of the license can be
// found in the file LICENSE at the top level of the MEPLS distribution.
//
// -----------------------------------------------------------------------

#include <example.h>
#include <mepls/utils.h>
#include <mepls/solver.h>
#include <mepls/system.h>
#include <mepls/event.h>
#include <mepls/history.h>
#include <mepls/dynamics.h>
#include <cmdparser.hpp>
#include <deal.II/base/parameter_handler.h>
#include <deal.II/base/path_search.h>
#include <boost/property_tree/json_parser.hpp>

// new headers
#include <omp.h>
#include <deal.II/base/conditional_ostream.h>

struct Parameters
{
   unsigned int seed = 1234567;
   unsigned int n_rep = 1;
   unsigned int Nx = 32;
   unsigned int Ny = 32;
   double G = 30.;
   double nu = 0.3;
   double gamma = 0.05;
   double k = 6.;
   double strain_limit_aqs = 0.05;
   double time_limit_creep = 1e5;
   double ext_stress_creep = 0.;
   double lambda = 1.;
   double temperature = 1.;
   bool verbose = true;

   void declare_entries(dealii::ParameterHandler &prm)
   {
      prm.enter_subsection("Section1");

      prm.declare_entry("seed", mepls::utils::str::to_string(seed), dealii::Patterns::Integer(0), "");
      prm.declare_entry("n_rep", mepls::utils::str::to_string(n_rep), dealii::Patterns::Integer(0), "");
      prm.declare_entry("Nx", mepls::utils::str::to_string(Nx), dealii::Patterns::Integer(0), "");
      prm.declare_entry("Ny", mepls::utils::str::to_string(Ny), dealii::Patterns::Integer(0), "");
      prm.declare_entry("G", mepls::utils::str::to_string(G), dealii::Patterns::Double(0.0), "");
      prm.declare_entry("nu", mepls::utils::str::to_string(nu), dealii::Patterns::Double(0.0), "");
      prm.declare_entry("gamma", mepls::utils::str::to_string(gamma), dealii::Patterns::Double(0.0), "");
      prm.declare_entry("strain_limit_aqs", mepls::utils::str::to_string(strain_limit_aqs), dealii::Patterns::Double(0.0), "");
      prm.declare_entry("time_limit_creep", mepls::utils::str::to_string(time_limit_creep), dealii::Patterns::Double(0.0), "");
      prm.declare_entry("ext_stress_creep", mepls::utils::str::to_string(ext_stress_creep), dealii::Patterns::Double(0.0), "");
      prm.declare_entry("temperature", mepls::utils::str::to_string(temperature), dealii::Patterns::Double(0.0), "");
      prm.declare_entry("lambda", mepls::utils::str::to_string(lambda), dealii::Patterns::Double(0.0), "");
      prm.declare_entry("k", mepls::utils::str::to_string(k), dealii::Patterns::Double(0.0), "");
	  prm.declare_entry("verbose", mepls::utils::str::to_string(verbose), dealii::Patterns::Bool(), "");

      prm.leave_subsection();
   }

   void load_entries(dealii::ParameterHandler &prm)
   {
      prm.enter_subsection("Section1");

      seed = prm.get_integer("seed");
      n_rep = prm.get_integer("n_rep");
      Nx = prm.get_integer("Nx");
      Ny = prm.get_integer("Ny");
      G = prm.get_double("G");
      nu = prm.get_double("nu");
      gamma = prm.get_double("gamma");
      k = prm.get_double("k");
      strain_limit_aqs = prm.get_double("strain_limit_aqs");
      time_limit_creep = prm.get_double("time_limit_creep");
      ext_stress_creep = prm.get_double("ext_stress_creep");
      temperature = prm.get_double("temperature");
      lambda = prm.get_double("lambda");
      verbose = prm.get_bool("verbose");

      prm.leave_subsection();
   }

   void load_file(const std::string &filename)
   {
      dealii::ParameterHandler prm;
      declare_entries(prm);
      prm.parse_input(filename);
      load_entries(prm);
   }

   void generate_file(const std::string &filename)
   {
      std::ofstream outfile(filename);
      dealii::ParameterHandler prm;
      declare_entries(prm);
      prm.print_parameters(outfile, dealii::ParameterHandler::OutputStyle::Text);
   }

};

template<int dim>
void write_data(const mepls::history::History<dim> &creep_history,
                const mepls::history::History<dim> &aqs_history,
                const Parameters &p)
{
   boost::property_tree::ptree data_tree;

   data_tree.put("Name", "Step5");
   data_tree.put("Description", "System undergoing creep deformation, and then driven in "
								"athermal quasistatic shear");

   data_tree.put("Parameters.dim", 2);
   data_tree.put("Parameters.seed", p.seed);
   data_tree.put("Parameters.Nx", p.Nx);
   data_tree.put("Parameters.Ny", p.Ny);
   data_tree.put("Parameters.G", p.G);
   data_tree.put("Parameters.nu", p.nu);
   data_tree.put("Parameters.gamma", p.gamma);
   data_tree.put("Parameters.lambda", p.lambda);
   data_tree.put("Parameters.temperature", p.temperature);
   data_tree.put("Parameters.k", p.k);
   data_tree.put("Parameters.ext_stress_creep", p.ext_stress_creep);
   data_tree.put("Parameters.time_limit_creep", p.time_limit_creep);
   data_tree.put("Parameters.strain_limit_aqs", p.strain_limit_aqs);

   	// -------- creep history ----------
	{
	   std::ostringstream plastic_events_csv;
	   plastic_events_csv << "index,element,eigenstrain_00,eigenstrain_11,eigenstrain_01\n";
	   for(auto &row : creep_history.plastic)
		  plastic_events_csv << row.index << "," << row.element << "," << row.eigenstrain_00 << ","
							 << row.eigenstrain_11 << "," << row.eigenstrain_01 << "\n";

	   data_tree.put("Data.creep.plastic_events", plastic_events_csv.str());

	   std::ostringstream driving_events_csv;
	   driving_events_csv << "index,dtime,dext_stress,dtotal_strain\n";
	   for(auto &row : creep_history.driving)
		  driving_events_csv << row.index << "," << row.dtime << ","
							 << row.dext_stress << "," << row.dtotal_strain << "\n";

	   data_tree.put("Data.creep.driving_events", driving_events_csv.str());

	   std::ostringstream macro_evolution_csv;
	   macro_evolution_csv << "index,ext_stress,total_strain,time,av_vm_stress,av_vm_plastic_strain\n";
	   for(auto &row : creep_history.macro_evolution)
		  macro_evolution_csv << row.index << "," << row.ext_stress << ","
							 << row.total_strain << "," << row.time << ","
							 << row.av_vm_stress << "," << row.av_vm_plastic_strain << "\n";

	   data_tree.put("Data.creep.macro_evolution", macro_evolution_csv.str());
	}

	// -------- aqs history ----------
	{
	   std::ostringstream plastic_events_csv;
	   plastic_events_csv << "index,element,eigenstrain_00,eigenstrain_11,eigenstrain_01\n";
	   for(auto &row : aqs_history.plastic)
		  plastic_events_csv << row.index << "," << row.element << "," << row.eigenstrain_00 << ","
							 << row.eigenstrain_11 << "," << row.eigenstrain_01 << "\n";

	   data_tree.put("Data.AQS.plastic_events", plastic_events_csv.str());

	   std::ostringstream driving_events_csv;
	   driving_events_csv << "index,dtime,dext_stress,dtotal_strain\n";
	   for(auto &row : aqs_history.driving)
		  driving_events_csv << row.index << "," << row.dtime << ","
							 << row.dext_stress << "," << row.dtotal_strain << "\n";

	   data_tree.put("Data.AQS.driving_events", driving_events_csv.str());

	   std::ostringstream macro_evolution_csv;
	   macro_evolution_csv << "index,ext_stress,total_strain,time,av_vm_stress,av_vm_plastic_strain\n";
	   for(auto &row : aqs_history.macro_evolution)
		  macro_evolution_csv << row.index << "," << row.ext_stress << ","
							 << row.total_strain << "," << row.time << ","
							 << row.av_vm_stress << "," << row.av_vm_plastic_strain << "\n";

	   data_tree.put("Data.AQS.macro_evolution", macro_evolution_csv.str());
	}

	// create a descriptive filename
   	std::ostringstream filename;

	filename << "Nx_" << p.Nx
	         << "+gamma_" << std::fixed << std::setprecision(2) << p.gamma
	         << "+lambda_" << p.lambda
	         << "+k_" << p.k
		     << "+G_" << p.G
		     << "+T_" << p.temperature
		     << "+ext_stress_" << p.ext_stress_creep
			 << "+seed_" << p.seed
			 << ".json";

   std::ofstream output_file( filename.str() );
   boost::property_tree::json_parser::write_json(output_file, data_tree);
   output_file.close();
}

void run(const Parameters &p, dealii::ConditionalOStream & cout)
{
   //----- SET UP ------

   constexpr unsigned dim = 2;
   std::mt19937 generator(p.seed);

   dealii::SymmetricTensor<4, dim> C = mepls::utils::tensor::make_isotropic_stiffness<dim>(p.G, p.nu);

   mepls::element::Vector<dim> elements;

   for(double n = 0; n < p.Nx * p.Ny; ++n)
   {
      example::element::Scalar<dim>::Config conf;
      conf.number = n;
      conf.gamma = p.gamma;
      conf.lambda = p.lambda;
      conf.k = p.k;
      conf.temperature = p.temperature;

      auto element = new example::element::Scalar<dim>(conf, generator);
      element->C(C);

      elements.push_back(element);
   }

   mepls::elasticity_solver::LeesEdwards<dim> solver(p.Nx, p.Ny, mepls::elasticity_solver::ControlMode::stress);
   for(auto &element : elements)
      solver.set_elastic_properties(element->number(), element->C());
   solver.setup_and_assembly();

   mepls::element::calculate_local_stress_coefficients_central(elements, solver);
   mepls::element::calculate_ext_stress_coefficients(elements, solver);

   mepls::system::Standard<dim> system(elements, solver, generator);

   //----- CREEP DEFORMATION ------

   mepls::history::History<dim> creep_history("creep_history");
   system.set_history(creep_history);
   creep_history.add_macro(system);

   // apply an external (stress) load of amplitude p.ext_stress_creep
   mepls::dynamics::fixed_load_increment(p.ext_stress_creep, system);
   creep_history.add_macro(system);

   mepls::dynamics::KMC<dim> kmc;

   mepls::utils::ContinueSimulation continue_creep;

	while( continue_creep() )
	{
		cout << system.macrostate["time"] << " " << system.macrostate["total_strain"] <<" "
			  << system.macrostate["ext_stress"] << std::endl;

		kmc(system);
		creep_history.add_macro(system);

		mepls::dynamics::relaxation(system, continue_creep);
		creep_history.add_macro(system);

	  continue_creep(system.macrostate["time"] < p.time_limit_creep, "creep time limit reached");
	}

	 cout << continue_creep << std::endl;

   //----- TRANSITION TO AQS ------

	// remove the external load
    mepls::dynamics::fixed_load_increment(-p.ext_stress_creep, system);
	creep_history.add_macro(system);

	// relax possible unstable slip system after the load change
	mepls::dynamics::relaxation(system, continue_creep);
	creep_history.add_macro(system);

    // in the AQS, we start measuring the macroscale strain from zero
	system.macrostate.clear();

	for(auto & element : elements)
		element->state_to_prestress();

	// we switch the driving mode to strain-controlled during AQS
    solver.set_control_mode(mepls::elasticity_solver::ControlMode::strain);

    // since the type of driving conditions have changed, the local stress change induced by a unit
    // load increment is different. We need to re-compute the ext_stress_coefficients
    mepls::element::calculate_ext_stress_coefficients(elements, solver);

    //----- ATHERMAL QUASISTATIC SHEAR ------

   	mepls::history::History<dim> aqs_history("aqs_history");
  	system.set_history(aqs_history);
   	aqs_history.add_macro(system);

   mepls::utils::ContinueSimulation continue_AQS_simulation;
   while(continue_AQS_simulation())
   {
		cout << system.macrostate["total_strain"] << " " << system.macrostate["ext_stress"] << std::endl;

      mepls::dynamics::extremal_dynamics_step(system);
      aqs_history.add_macro(system);

      mepls::dynamics::relaxation(system, continue_AQS_simulation);
      aqs_history.add_macro(system);

      continue_AQS_simulation(system.macrostate["total_strain"] < p.strain_limit_aqs, "AQS strain limit reached");
   }

	   cout << continue_AQS_simulation << std::endl;

   for(auto &element : elements)
      delete element;

	// write the simulation data, using its own dedicated function
	write_data(creep_history, aqs_history, p);
}

int main(int argc, char *argv[])
{
   // Read the command line arguments. We define the -f 'filename' to pass the
   // path to the parameters file
   cli::Parser parser(argc, argv);
   parser.set_optional<std::string>("f", "file", "./default.prm", "Name of the input configuration file");
   parser.run_and_exit_if_error();
   // you can check https://github.com/FlorianRappl/CmdParser for a further documentation of cli::Parser

   // Create the parameters object
   Parameters p;

   // We try to load the parameters file, but if it doesn't exist, we generate a new one with the name
   // default.prm and default values
   try
   {
      p.load_file(parser.get<std::string>("f"));
   }
   catch(dealii::PathSearch::ExcFileNotFound &)
   {
      p.generate_file(parser.get<std::string>("f"));
      std::cout << "Configuration file " << parser.get<std::string>("f") << " created" << std::endl;
      return 1;
   }

	unsigned int n_rep = p.n_rep;
	if(n_rep < omp_get_max_threads())
		n_rep = omp_get_max_threads();

	// initialize the master engine with the master seed
	std::srand(p.seed);

	#pragma omp parallel
	{
		unsigned int n_threads = omp_get_max_threads();
		unsigned int id = omp_get_thread_num();
		unsigned int rep_per_thread = int( n_rep / n_threads );

		dealii::ConditionalOStream cout(std::cout, id==0 and p.verbose);

		for(unsigned int n = 0; n < rep_per_thread; ++n)
		{
			Parameters p_thread = p;
			#pragma critical
			{
				// generate a seed for a simulation run
				p_thread.seed = std::rand();
			};

			run(p_thread, cout);
		}

	}

   return 0;
}


// Boost.Geometry (aka GGL, Generic Geometry Library)

// Copyright (c) 2007-2012 Barend Gehrels, Amsterdam, the Netherlands.
// Copyright (c) 2008-2012 Bruno Lalande, Paris, France.
// Copyright (c) 2009-2012 Mateusz Loskot, London, UK.
// Copyright (c) 2013 Adam Wulkiewicz, Lodz, Poland.

// Parts of Boost.Geometry are redesigned from Geodan's Geographic Library
// (geolib/GGL), copyright (c) 1995-2010 Geodan, Amsterdam, the Netherlands.

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

#ifndef BOOST_GEOMETRY_EXTENSIONS_ALGEBRA_GEOMETRIES_MATRIX_HPP
#define BOOST_GEOMETRY_EXTENSIONS_ALGEBRA_GEOMETRIES_MATRIX_HPP

#include <cstddef>

#include <boost/geometry/extensions/algebra/core/tags.hpp>
#include <boost/geometry/extensions/algebra/geometries/concepts/matrix_concept.hpp>

namespace boost { namespace geometry {

namespace model {

template <typename T, std::size_t Rows, std::size_t Cols>
class matrix
{
    BOOST_CONCEPT_ASSERT( (concept::Matrix<matrix>) );

public:

    /// @brief Default constructor, no initialization
    inline matrix()
    {}

    /// @brief Get a coordinate
    /// @tparam I row index
    /// @tparam J col index
    /// @return the cell value
    template <std::size_t I, std::size_t J>
    inline T const& get() const
    {
        BOOST_STATIC_ASSERT(I < Rows);
        BOOST_STATIC_ASSERT(J < Cols);
        return m_values[I + Rows * J];
    }

    /// @brief Set a coordinate
    /// @tparam I row index
    /// @tparam J col index
    /// @param value value to set
    template <std::size_t I, std::size_t J>
    inline void set(T const& value)
    {
        BOOST_STATIC_ASSERT(I < Rows);
        BOOST_STATIC_ASSERT(J < Cols);
        m_values[I + Rows * J] = value;
    }

private:

    T m_values[Rows * Cols];
};

} // namespace model

#ifndef DOXYGEN_NO_TRAITS_SPECIALIZATIONS
namespace traits
{

template <typename CoordinateType, std::size_t Rows, std::size_t Cols>
struct tag<model::matrix<CoordinateType, Rows, Cols> >
{
    typedef matrix_tag type;
};

template <typename CoordinateType, std::size_t Rows, std::size_t Cols>
struct coordinate_type<model::matrix<CoordinateType, Rows, Cols> >
{
    typedef CoordinateType type;
};

//template <typename CoordinateType, std::size_t Rows, std::size_t Cols>
//struct coordinate_system<model::matrix<CoordinateType, Rows, Cols> >
//{
//    typedef cs::cartesian type;
//};

// TODO - move this class to traits.hpp
template <typename Geometry, std::size_t Index>
struct indexed_dimension
{
     BOOST_MPL_ASSERT_MSG(false,
                          NOT_IMPLEMENTED_FOR_THIS_GEOMETRY_OR_INDEX,
                          (Geometry, boost::integral_constant<std::size_t, Index>));
};

template <typename CoordinateType, std::size_t Rows, std::size_t Cols>
struct indexed_dimension<model::matrix<CoordinateType, Rows, Cols>, 0>
    : boost::integral_constant<std::size_t, Rows>
{};

template <typename CoordinateType, std::size_t Rows, std::size_t Cols>
struct indexed_dimension<model::matrix<CoordinateType, Rows, Cols>, 1>
    : boost::integral_constant<std::size_t, Cols>
{};

template <typename CoordinateType, std::size_t Rows, std::size_t Cols, std::size_t I, std::size_t J>
struct indexed_access<model::matrix<CoordinateType, Dimension>, I, J>
{
    typedef CoordinateType coordinate_type;

    static inline coordinate_type get(model::matrix<CoordinateType, Rows, Cols> const& m)
    {
        return m.template get<I, J>();
    }

    static inline void set(model::matrix<CoordinateType, Rows, Cols> & m, coordinate_type const& value)
    {
        m.template set<I, J>(value);
    }
};

} // namespace traits
#endif // DOXYGEN_NO_TRAITS_SPECIALIZATIONS

}} // namespace boost::geometry

#endif // BOOST_GEOMETRY_EXTENSIONS_ALGEBRA_GEOMETRIES_MATRIX_HPP


//Bsplines file
#include <Eigen/Sparse>
#include <Eigen/Eigenvalues>
#include <Eigen/StdVector>
#include <cstdlib>
#include <iostream>
#include <mpi.h>
#include <fstream>
#include <iterator>
#include <vector>
#include <algorithm>
#include <complex>
#include <sstream>
#include <math.h>
#include <omp.h>
#include <numeric>
#include <chrono>
#include <ctime>
#include <limits>
#include "bsplines.h"

#define bp(x) cout << "breakpoint" << (#x) << endl

using namespace Eigen;
using namespace std::complex_literals;
using namespace std::chrono;
using std::cout;
using std::endl;

void generateBsplines(InputParameter Inp, double dX,ArrayXd x,ArrayXXd &BB,ArrayXXd &d_BB, ArrayXXd &diff_BB){
    std::vector<Array<double,Dynamic,Dynamic>
        ,aligned_allocator<Array<double,Dynamic,Dynamic> > > B(Inp.n+1);
    for(int i = 0 ; i < Inp.n+1; i++){
        B[i].setZero(x.size(),Inp.N_b-1+2*Inp.n);
    }

    ArrayXd t = ArrayXd::Zero(Inp.N_b+2*Inp.n);

    for(int i = Inp.n+1; i< Inp.N_b +Inp.n;i++){
        B[0].col(i-1).segment((i-Inp.n-1) * Inp.n,Inp.n) = 1;
        t(i-1) = (i-Inp.n-1) * dX;

    }
    t.segment(Inp.N_b + Inp.n - 1,Inp.n+1) = Inp.b;
    for( int k = 1; k < Inp.n+1;k++){
        for(int i = 1; i < Inp.N_b-1+2*Inp.n-k+1; i++){
            double tmp1 = t(i+k-1) - t(i-1);
            double tmp2 = t(i+k+0) - t(i-0);

            if(tmp1 == 0)
                tmp1 =1;
            if(tmp2 == 0)
                tmp2 =1;

            B[k].col(i-1) = ((x - t(i-1))/tmp1) * B[k-1].col(i-1)
                           + (t(i+k) - x)/tmp2 * B[k-1].col(i);
       }
    }
    bp("first order derivaitve of bsplines");
    //d_B is the first derivative
    std::vector<Array<double,Dynamic,Dynamic>
        ,aligned_allocator<Array<double,Dynamic,Dynamic> > > d_B(Inp.n+1);
    for(int i = 0 ; i < Inp.n+1; i++){
        d_B[i].setZero(x.size(),Inp.N_b-1+2*Inp.n);
    }

    for( int k = 2; k < Inp.n+1; k++){
        for(int i = 1; i < Inp.N_b-1+2*Inp.n-k+1; i++){
            double tmp1 = t(i+k-1) - t(i-1);
            double tmp2 = t(i+k+0) - t(i+0);

            if(tmp1 == 0)
                tmp1 =1;
            if(tmp2 == 0)
                tmp2 =1;

            d_B[k].col(i-1) = k * (B[k-1].col(i-1)/tmp1 - B[k-1].col(i)/tmp2);

        }
    }
    bp("scond order derivative of bsplines ");
    //diff of B is the second derivative
    std::vector<Array<double,Dynamic,Dynamic>
        ,aligned_allocator<Array<double,Dynamic,Dynamic> > > diff_B(Inp.n+1);
    for(int i = 0 ; i < Inp.n+1; i++){
        diff_B[i].setZero(x.size(),Inp.N_b-1+2*Inp.n);
    }

    for( int k = 2; k < Inp.n+1; k++){
        for(int i = 1; i < Inp.N_b-1+2*Inp.n-k+1; i++){
            double tmp1 = t(i+k-1) - t(i-1);
            double tmp2 = t(i+k+0) - t(i+0);
            double tmp3 = t(i+k-2) - t(i-1);
            double tmp4 = t(i+k-1) - t(i+0);
            double tmp5 = t(i+k+0) - t(i+1);

            if(tmp1 == 0)
                tmp1 =1;
            if(tmp2 == 0)
                tmp2 =1;
            if(tmp3 == 0)
                tmp3 =1;
            if(tmp4 == 0)
                tmp4 =1;
            if(tmp5 == 0)
                tmp5 =1;

        diff_B[k].col(i-1) = k * (k-1)/tmp1 * (B[k-2].col(i-1)/tmp3 - B[k-2].col(i)/tmp4)
                           - k * (k-1)/tmp2 * (B[k-2].col(i)/tmp4 - B[k-2].col(i+1)/tmp5);
        }
    }
    BB=B[Inp.n];
    d_BB=d_B[Inp.n];
    diff_BB=diff_B[Inp.n];

    cout <<"Bsplines generated"<<endl;

}

void gausLegendreSetup(int k, int nmbr_breakpoints,
        ArrayXd x_break, ArrayXd &x, ArrayXd &weights){
    //Return evenly spaced values within a given interval.
    std::vector<int> I(k);
    std::iota(I.begin(),I.end(),1);

    MatrixXd J = MatrixXd::Zero(k,k);
    for(int i = 0; i < k-1; i++){
        J(i,i+1) = 0.5 / sqrt(1-pow(2*I[i],-2));
        J(i+1,i) = 0.5 / sqrt(1-pow(2*I[i],-2));
    }
    //get the eigenvalues and eigenvectors

    SelfAdjointEigenSolver<MatrixXd> es;
    es.compute(J);
    VectorXd eigval = es.eigenvalues();
    MatrixXd eigvect = es.eigenvectors();
    for(int i = 2; i < nmbr_breakpoints+1;i++){
        x.segment((i-2)*k,k) = (x_break(i-1)-x_break(i-2))*0.5*  eigval.array()
            +(x_break[i-1]+x_break[i-2])*0.5;
        weights.segment((i-2)*k,k) = 2*eigvect.row(0).array().square();

    }

}


#ifndef SKYLARK_BLOCKADMM_HPP
#define SKYLARK_BLOCKADMM_HPP

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

#ifdef SKYLARK_HAVE_OPENMP
#include <omp.h>
#endif

#include "../utility/timer.hpp"
#include "hilbert.hpp"

// Columns are examples, rows are features
typedef elem::DistMatrix<double, elem::STAR, elem::VC> DistInputMatrixType;

// Rows are examples, columns are target values
typedef elem::DistMatrix<double, elem::VC, elem::STAR> DistTargetMatrixType;

typedef elem::Matrix<double> LocalMatrixType;
typedef skylark::base::sparse_matrix_t<double> sparse_matrix_t;

template <class T>
class BlockADMMSolver
{
public:

    typedef skylark::sketch::sketch_transform_t<T, LocalMatrixType>
    feature_transform_t;
    typedef std::vector<const feature_transform_t *> feature_transform_array_t;

    // No feature transdeforms (aka just linear regression).
    BlockADMMSolver(const lossfunction* loss,
        const regularization* regularizer,
        double lambda, // regularization parameter
        int NumFeatures,
        int NumFeaturePartitions = 1);

    // Easy interface, aka kernel based.
    template<typename Kernel, typename MapTypeTag>
    BlockADMMSolver<T>(skylark::base::context_t& context,
        const lossfunction* loss,
        const regularization* regularizer,
        double lambda, // regularization parameter
        int NumFeatures,
        Kernel kernel,
        MapTypeTag tag,
        int NumFeaturePartitions = 1);

    // Easy interface, aka kernel based, with quasi-random features.
    template<typename Kernel>
    BlockADMMSolver<T>(skylark::base::context_t& context,
        const lossfunction* loss,
        const regularization* regularizer,
        double lambda, // regularization parameter
        int NumFeatures,
        Kernel kernel,
        skylark::ml::quasi_feature_transform_tag tag,
        int NumFeaturePartitions);

    // Guru interface.
    BlockADMMSolver<T>(const lossfunction* loss,
        const regularization* regularizer,
        const feature_transform_array_t& featureMaps,
        double lambda, // regularization parameter
        bool ScaleFeatureMaps = true);

    void set_nthreads(int NumThreads) { this->NumThreads = NumThreads; }
    void set_rho(double RHO) { this->RHO = RHO; }
    void set_maxiter(double MAXITER) { this->MAXITER = MAXITER; }
    void set_tol(double TOL) { this->TOL = TOL; }
    void set_cache_transform(bool CacheTransforms) {this->CacheTransforms = CacheTransforms;}

    ~BlockADMMSolver();

    void InitializeFactorizationCache();
    void InitializeTransformCache(int n);

    skylark::ml::model_t<T, LocalMatrixType>* train(T& X,
        LocalMatrixType& Y, T& Xv, LocalMatrixType& Yv,
        const boost::mpi::communicator& comm);

    int get_numfeatures() {return NumFeatures;}

    feature_transform_array_t& get_feature_maps() {return featureMaps;}

private:

    feature_transform_array_t featureMaps;
    int NumFeatures;
    int NumFeaturePartitions;
    lossfunction* loss;
    regularization* regularizer;
    std::vector<int> starts, finishes;
    bool ScaleFeatureMaps;
    bool OwnFeatureMaps;
    LocalMatrixType **Cache;
    LocalMatrixType **TransformCache;
    int NumThreads;

    double lambda;
    double RHO;
    int MAXITER;
    double TOL;

    bool CacheTransforms;
};

template <class T>
void BlockADMMSolver<T>::InitializeFactorizationCache() {
    Cache = new LocalMatrixType* [NumFeaturePartitions];
    for(int j=0; j<NumFeaturePartitions; j++) {
        int start = starts[j];
        int finish = finishes[j];
        int sj = finish - start  + 1;
        Cache[j]  = new elem::Matrix<double>(sj, sj);
    }
}

template <class T>
void BlockADMMSolver<T>::InitializeTransformCache(int n) {
    TransformCache = new LocalMatrixType* [NumFeaturePartitions];
    for(int j=0; j<NumFeaturePartitions; j++) {
        int start = starts[j];
        int finish = finishes[j];
        int sj = finish - start  + 1;
        TransformCache[j]  = new elem::Matrix<double>(sj, n);
    }
}

// No feature transforms (aka just linear regression).
template <class T>
BlockADMMSolver<T>::BlockADMMSolver(
        const lossfunction* loss,
        const regularization* regularizer,
        double lambda, // regularization parameter
        int NumFeatures,
        int NumFeaturePartitions) :
        NumFeatures(NumFeatures),
            NumFeaturePartitions(NumFeaturePartitions),
            starts(NumFeaturePartitions), finishes(NumFeaturePartitions),
            NumThreads(1), RHO(1.0), MAXITER(1000), TOL(0.1) {

    this->loss = const_cast<lossfunction *> (loss);
    this->regularizer = const_cast<regularization *> (regularizer);
    this->lambda = lambda;
    this->NumFeaturePartitions = NumFeaturePartitions;
    int blksize = int(ceil(double(NumFeatures) / NumFeaturePartitions));
    for(int i = 0; i < NumFeaturePartitions; i++) {
        starts[i] = i * blksize;
        finishes[i] = std::min((i + 1) * blksize, NumFeatures) - 1;
    }
    this->ScaleFeatureMaps = false;
    OwnFeatureMaps = false;
    InitializeFactorizationCache();
    CacheTransforms = false;
}

// Easy interface, aka kernel based.
template<class T>
template<typename Kernel, typename MapTypeTag>
BlockADMMSolver<T>::BlockADMMSolver(skylark::base::context_t& context,
    const lossfunction* loss,
    const regularization* regularizer,
    double lambda, // regularization parameter
    int NumFeatures,
    Kernel kernel,
    MapTypeTag tag,
    int NumFeaturePartitions) :
    featureMaps(NumFeaturePartitions),
    NumFeatures(NumFeatures), NumFeaturePartitions(NumFeaturePartitions),
    starts(NumFeaturePartitions), finishes(NumFeaturePartitions),
    NumThreads(1), RHO(1.0), MAXITER(1000), TOL(0.1) {

    this->loss = const_cast<lossfunction *> (loss);
    this->regularizer = const_cast<regularization *> (regularizer);
    this->lambda = lambda;
    int blksize = int(ceil(double(NumFeatures) / NumFeaturePartitions));
    for(int i = 0; i < NumFeaturePartitions; i++) {
        starts[i] = i * blksize;
        finishes[i] = std::min((i + 1) * blksize, NumFeatures) - 1;
        int sj = finishes[i] - starts[i] + 1;
        featureMaps[i] =
            kernel.template create_rft< T, LocalMatrixType >(sj, tag, context);
    }
    this->ScaleFeatureMaps = true;
    OwnFeatureMaps = true;
    InitializeFactorizationCache();
    CacheTransforms = false;
}

// Easy interface, aka kernel based, with quasi-random features.
template<class T>
template<typename Kernel>
BlockADMMSolver<T>::BlockADMMSolver(skylark::base::context_t& context,
    const lossfunction* loss,
    const regularization* regularizer,
    double lambda, // regularization parameter
    int NumFeatures,
    Kernel kernel,
    skylark::ml::quasi_feature_transform_tag tag,
    int NumFeaturePartitions) :
    featureMaps(NumFeaturePartitions),
    NumFeatures(NumFeatures), NumFeaturePartitions(NumFeaturePartitions),
    starts(NumFeaturePartitions), finishes(NumFeaturePartitions),
    NumThreads(1), RHO(1.0), MAXITER(1000), TOL(0.1) {

    this->loss = const_cast<lossfunction *> (loss);
    this->regularizer = const_cast<regularization *> (regularizer);
    this->lambda = lambda;
    int blksize = int(ceil(double(NumFeatures) / NumFeaturePartitions));
    skylark::utility::leaped_halton_sequence_t<double>
        qmcseq(kernel.qrft_sequence_dim()); // TODO size
    for(int i = 0; i < NumFeaturePartitions; i++) {
        starts[i] = i * blksize;
        finishes[i] = std::min((i + 1) * blksize, NumFeatures) - 1;
        int sj = finishes[i] - starts[i] + 1;
        featureMaps[i] =
            kernel.template create_qrft< T, LocalMatrixType,
              skylark::utility::leaped_halton_sequence_t>(sj, qmcseq,
                  starts[i], context);
    }
    this->ScaleFeatureMaps = true;
    OwnFeatureMaps = true;
    InitializeFactorizationCache();
    CacheTransforms = false;
}

// Guru interface
template <class T>
BlockADMMSolver<T>::BlockADMMSolver(const lossfunction* loss,
    const regularization* regularizer,
    const feature_transform_array_t &featureMaps,
    double lambda,
    bool ScaleFeatureMaps) :
    featureMaps(featureMaps),
    NumFeaturePartitions(featureMaps.size()),
    starts(NumFeaturePartitions), finishes(NumFeaturePartitions),
    NumThreads(1), RHO(1.0), MAXITER(1000), TOL(0.1)  {

    this->loss = const_cast<lossfunction *> (loss);
    this->regularizer = const_cast<regularization *> (regularizer);
    this->lambda = lambda;
    NumFeaturePartitions = featureMaps.size();
    NumFeatures = 0;
    for(int i = 0; i < NumFeaturePartitions; i++) {
        starts[i] = NumFeatures;
        finishes[i] = NumFeatures + featureMaps[i]->get_S() - 1;
        NumFeatures += featureMaps[i]->get_S();
    }
    this->ScaleFeatureMaps = ScaleFeatureMaps;
    OwnFeatureMaps = false;
    InitializeFactorizationCache();
    CacheTransforms = false;
}

template <class T>
BlockADMMSolver<T>::~BlockADMMSolver() {
    for(int i=0; i  < NumFeaturePartitions; i++) {
        delete Cache[i];
        if (OwnFeatureMaps)
            delete featureMaps[i];
    }
    delete[] Cache;
}

template <class T>
skylark::ml::model_t<T, LocalMatrixType>* BlockADMMSolver<T>::train(T& X, LocalMatrixType& Y, T& Xv, LocalMatrixType& Yv,
    const boost::mpi::communicator& comm) {

       int rank = comm.rank();
       int size = comm.size();

       int P = size;

       int ni = skylark::base::Width(X);
       int d = skylark::base::Height(X);
       int targets = GetNumTargets(comm, Y);

       skylark::ml::model_t<T, LocalMatrixType>* model =
           new skylark::ml::model_t<T, LocalMatrixType>(featureMaps,
               ScaleFeatureMaps, NumFeatures, targets);

       elem::Matrix<double> Wbar;
       elem::View(Wbar, model->get_coef());

       int k = Wbar.Width();

       // number of classes, targets - to generalize

       int D = NumFeatures;

       // exception: check if D = Wbar.Height();

       LocalMatrixType O(k, ni); //uses default Grid
       elem::MakeZeros(O);

       LocalMatrixType Obar(k, ni); //uses default Grid
       elem::MakeZeros(Obar);

       LocalMatrixType nu(k, ni); //uses default Grid
       elem::MakeZeros(nu);

       LocalMatrixType W, mu, Wi, mu_ij, ZtObar_ij;

       if(rank==0) {
           elem::Zeros(W,  D, k);
           elem::Zeros(mu, D, k);
       }
       elem::Zeros(Wi, D, k);
       elem::Zeros(mu_ij, D, k);
       elem::Zeros(ZtObar_ij, D, k);

       int iter = 0;

       // int ni = O.LocalWidth();

       //elem::Matrix<double> x = X.Matrix();
       //elem::Matrix<double> y = Y.Matrix();

       double localloss = loss->evaluate(O, Y);
       double totalloss, accuracy, obj;

       int Dk = D*k;
       int nik  = ni*k;
       int start, finish, sj;

       boost::mpi::timer timer;

       LocalMatrixType sum_o, del_o, wbar_output;
       elem::Zeros(del_o, k, ni);
       LocalMatrixType Yp(Yv.Height(), k);
       LocalMatrixType Yp_labels(Yv.Height(), 1);

       /*LocalMatrixType wbar_tmp;
       //if (NumThreads > 1)

       elem::Zeros(wbar_tmp, k, ni);*/

       if (CacheTransforms)
                   InitializeTransformCache(ni);

       SKYLARK_TIMER_INITIALIZE(ITERATIONS_PROFILE);
       SKYLARK_TIMER_INITIALIZE(COMMUNICATION_PROFILE);
       SKYLARK_TIMER_INITIALIZE(TRANSFORM_PROFILE);
       SKYLARK_TIMER_INITIALIZE(ZTRANSFORM_PROFILE);
       SKYLARK_TIMER_INITIALIZE(ZMULT_PROFILE);
       SKYLARK_TIMER_INITIALIZE(PROXLOSS_PROFILE);
       SKYLARK_TIMER_INITIALIZE(BARRIER_PROFILE);
       SKYLARK_TIMER_INITIALIZE(PREDICTION_PROFILE);

       while(iter<MAXITER) {

           SKYLARK_TIMER_RESTART(ITERATIONS_PROFILE);

           iter++;

           SKYLARK_TIMER_RESTART(COMMUNICATION_PROFILE);
           broadcast(comm, Wbar.Buffer(), Dk, 0);

           SKYLARK_TIMER_ACCUMULATE(COMMUNICATION_PROFILE)

           // mu_ij = mu_ij - Wbar
           elem::Axpy(-1.0, Wbar, mu_ij);

           // Obar = Obar - nu
           elem::Axpy(-1.0, nu, Obar);

           SKYLARK_TIMER_RESTART(PROXLOSS_PROFILE);
           loss->proxoperator(Obar, 1.0/RHO, Y, O);
           SKYLARK_TIMER_ACCUMULATE(PROXLOSS_PROFILE);

           if(rank==0) {
               regularizer->proxoperator(Wbar, lambda/RHO, mu, W);
           }

           elem::Zeros(sum_o, k, ni);
           elem::Zeros(wbar_output, k, ni);

           int j;
           const feature_transform_t* featureMap;

           SKYLARK_TIMER_RESTART(TRANSFORM_PROFILE);

   #       ifdef SKYLARK_HAVE_OPENMP
   #       pragma omp parallel for if(NumThreads > 1) private(j, start, finish, sj, featureMap) num_threads(NumThreads)
   #       endif
           for(j = 0; j < NumFeaturePartitions; j++) {
               start = starts[j];
               finish = finishes[j];
               sj = finish - start  + 1;

               elem::Matrix<double> z(sj, ni);

               if (CacheTransforms && (iter > 1))
               {
                    elem::View(z,  *TransformCache[j], 0, 0, sj, ni);
               }
               else {
                   if (featureMaps.size() > 0) {
                       featureMap = featureMaps[j];

                       SKYLARK_TIMER_RESTART(ZTRANSFORM_PROFILE);
                       featureMap->apply(X, z, skylark::sketch::columnwise_tag());
                       SKYLARK_TIMER_ACCUMULATE(ZTRANSFORM_PROFILE)

                       if (ScaleFeatureMaps)
                           elem::Scal(sqrt(double(sj) / d), z);
                       } else {
                          // for linear case just use Z = X no slicing business.
                          // skylark::base::ColumnView<double>(z, x, );
                          // ;// VIEWS on SPARSE MATRICES: elem::View(z, x, start, 0, sj, ni);
                       }
               }

               elem::Matrix<double> tmp(sj, k);
               elem::Matrix<double> rhs(sj, k);
               elem::Matrix<double> o(k, ni);

               if(iter==1) {

                   elem::Matrix<double> Ones;
                   elem::Ones(Ones, sj, 1);
                   elem::Gemm(elem::NORMAL, elem::TRANSPOSE, 1.0, z, z, 0.0, *Cache[j]);
                   Cache[j]->UpdateDiagonal(Ones);
                   elem::Inverse(*Cache[j]);

                   if (CacheTransforms) {
                       *TransformCache[j] = z;
                       //DEBUG
                        std::cout << "CACHING TRANSFORMS..." << std::endl;
                        elem::Write(*TransformCache[0], "FeatureMatrix.asc", elem::ASCII, "");
                   }
               }

               elem::View(tmp, Wbar, start, 0, sj, k); //tmp = Wbar[J,:]

               LocalMatrixType wbar_tmp;
               elem::Zeros(wbar_tmp, k, ni);

               if (NumThreads > 1) {
                   elem::Gemm(elem::TRANSPOSE, elem::NORMAL, 1.0, tmp, z, 0.0, wbar_tmp);

   #               ifdef SKYLARK_HAVE_OPENMP
   #               pragma omp critical
   #               endif
                   elem::Axpy(1.0, wbar_tmp, wbar_output);
               } else
                   elem::Gemm(elem::TRANSPOSE, elem::NORMAL, 1.0, tmp, z, 1.0, wbar_output);

               rhs = tmp; //rhs = Wbar[J,:]
               elem::View(tmp, mu_ij, start, 0, sj, k); //tmp = mu_ij[J,:]
               elem::Axpy(-1.0, tmp, rhs); // rhs = rhs - mu_ij[J,:] = Wbar[J,:] - mu_ij[J,:]
               elem::View(tmp, ZtObar_ij, start, 0, sj, k);
               elem::Axpy(+1.0, tmp, rhs); // rhs = rhs + ZtObar_ij[J,:]

               SKYLARK_TIMER_RESTART(ZMULT_PROFILE);
               elem::Matrix<double> dsum = del_o;
               elem::Axpy(NumFeaturePartitions + 1.0, nu, dsum);
               elem::Gemm(elem::NORMAL, elem::TRANSPOSE, 1.0/(NumFeaturePartitions + 1.0), z, dsum, 1.0, rhs); // rhs = rhs + z'*(1/(n+1) * del_o + nu)
               SKYLARK_TIMER_ACCUMULATE(ZMULT_PROFILE);

               elem::View(tmp, Wi, start, 0, sj, k);
               elem::Gemm(elem::NORMAL, elem::NORMAL, 1.0, *Cache[j], rhs, 0.0, tmp); // ]tmp = Wi[J,:] = Cache[j]*rhs

               SKYLARK_TIMER_RESTART(ZMULT_PROFILE);
               elem::Gemm(elem::TRANSPOSE, elem::NORMAL, 1.0, tmp, z, 0.0, o); // o = (z*tmp)' = (z*Wi[J,:])'
               SKYLARK_TIMER_ACCUMULATE(ZMULT_PROFILE);

               // mu_ij[JJ,:] = mu_ij[JJ,:] + Wi[JJ,:];
               elem::View(tmp, mu_ij, start, 0, sj, k); //tmp = mu_ij[J,:]
               elem::View(rhs, Wi, start, 0, sj, k);
               elem::Axpy(+1.0, rhs, tmp);

               //ZtObar_ij[JJ,:] = numpy.dot(Z.T, o);
               elem::View(tmp, ZtObar_ij, start, 0, sj, k);
               elem::Gemm(elem::NORMAL, elem::TRANSPOSE, 1.0, z, o, 0.0, tmp);

               //  sum_o += o
               if (NumThreads > 1) {
   #               ifdef SKYLARK_HAVE_OPENMP
   #               pragma omp critical
   #               endif
                   elem::Axpy(1.0, o, sum_o);
               } else
                   elem::Axpy(1.0, o, sum_o);

               z.Empty();
           }

           SKYLARK_TIMER_ACCUMULATE(TRANSFORM_PROFILE);

           localloss = 0.0 ;
           //  elem::Zeros(o, ni, k);
           elem::Matrix<double> o(k, ni);
           elem::MakeZeros(o);
           elem::Scal(-1.0, sum_o);
           elem::Axpy(+1.0, O, sum_o); // sum_o = O.Matrix - sum_o
           del_o = sum_o;

           SKYLARK_TIMER_RESTART(PREDICTION_PROFILE);
           if (skylark::base::Width(Xv) > 0) {
               elem::MakeZeros(Yp);
               elem::MakeZeros(Yp_labels);
               model->predict(Xv, Yp_labels, Yp, NumThreads);
               accuracy = model->evaluate(Yv, Yp, comm);
           }
           SKYLARK_TIMER_ACCUMULATE(PREDICTION_PROFILE);

           localloss += loss->evaluate(wbar_output, Y);

           SKYLARK_TIMER_RESTART(COMMUNICATION_PROFILE);
           reduce(comm, localloss, totalloss, std::plus<double>(), 0);
           SKYLARK_TIMER_ACCUMULATE(COMMUNICATION_PROFILE);

           if(rank == 0) {
               obj = totalloss + lambda*regularizer->evaluate(Wbar);
               if (skylark::base::Width(Xv) <=0) {
                   std::cout << "iteration " << iter << " objective " << obj << " time " << timer.elapsed() << " seconds" << std::endl;
               }
               else {
                   std::cout << "iteration " << iter << " objective " << obj << " accuracy " << accuracy << " time " << timer.elapsed() << " seconds" << std::endl;
               }
           }

           elem::Copy(O, Obar);
           elem::Scal(1.0/(NumFeaturePartitions+1.0), sum_o);
           elem::Axpy(-1.0, sum_o, Obar);

           elem::Axpy(+1.0, O, nu);
           elem::Axpy(-1.0, Obar, nu);

           //Wbar = comm.reduce(Wi)
           SKYLARK_TIMER_RESTART(COMMUNICATION_PROFILE);
           boost::mpi::reduce (comm,
                                   Wi.LockedBuffer(),
                                   Wi.MemorySize(),
                                   Wbar.Buffer(),
                                   std::plus<double>(),
                                   0);
           SKYLARK_TIMER_ACCUMULATE(COMMUNICATION_PROFILE);

           if(rank==0) {
               //Wbar = (Wisum + W)/(P+1)
               elem::Axpy(1.0, W, Wbar);
               elem::Scal(1.0/(P+1), Wbar);

               // mu = mu + W - Wbar;
               elem::Axpy(+1.0, W, mu);
               elem::Axpy(-1.0, Wbar, mu);
           }

           SKYLARK_TIMER_RESTART(BARRIER_PROFILE);
           comm.barrier();
           SKYLARK_TIMER_ACCUMULATE(BARRIER_PROFILE);

           SKYLARK_TIMER_ACCUMULATE(ITERATIONS_PROFILE);
       }

       SKYLARK_TIMER_PRINT(ITERATIONS_PROFILE, comm);
       SKYLARK_TIMER_PRINT(COMMUNICATION_PROFILE, comm);
       SKYLARK_TIMER_PRINT(TRANSFORM_PROFILE, comm);
       SKYLARK_TIMER_PRINT(ZTRANSFORM_PROFILE, comm);
       SKYLARK_TIMER_PRINT(ZMULT_PROFILE, comm);
       SKYLARK_TIMER_PRINT(PROXLOSS_PROFILE, comm);
       SKYLARK_TIMER_PRINT(BARRIER_PROFILE, comm);
       SKYLARK_TIMER_PRINT(PREDICTION_PROFILE, comm);

       return model;
}

#endif /* SKYLARK_BLOCKADDM_HPP */


#include "surface_flow.h"
#include "fractional_laplacian.h"
#include "helpers.h"

#include "sobolev/h1.h"
#include "sobolev/h1_lbfgs.h"
#include "sobolev/h2.h"
#include "sobolev/l2_lbfgs.h"
#include "sobolev/bqn_lbfgs.h"
#include "sobolev/hs.h"
#include "sobolev/hs_schur.h"
#include "sobolev/hs_iterative.h"
#include "sobolev/constraints.h"
#include "spatial/convolution.h"

#include <Eigen/SparseCholesky>

namespace rsurfaces
{

    SurfaceFlow::SurfaceFlow(SurfaceEnergy *energy_)
    {
        energies.push_back(energy_);

        mesh = energy_->GetMesh();
        geom = energy_->GetGeom();
        stepCount = 0;

        origBarycenter = meshBarycenter(geom, mesh);
        std::cout << "Original barycenter = " << origBarycenter << std::endl;
        RecenterMesh();
        secretBarycenter = 0;
        obstacleEnergy = 0;

        verticesMutated = false;
        lbfgs = 0;
        bqn_B = 0;
    }

    void SurfaceFlow::AddAdditionalEnergy(SurfaceEnergy *extraEnergy)
    {
        energies.push_back(extraEnergy);
    }

    void SurfaceFlow::AddObstacleEnergy(SurfaceEnergy *obsEnergy)
    {
        obstacleEnergy = obsEnergy;
        AddAdditionalEnergy(obstacleEnergy);
    }

    void SurfaceFlow::UpdateEnergies()
    {
        for (SurfaceEnergy *energy : energies)
        {
            energy->Update();
        }
    }

    inline double guessStepSize(double gProjNorm)
    {
        // double initGuess = (gProjNorm < 1) ? 1.0 / sqrt(gProjNorm) : 1.0 / gProjNorm;
        double initGuess = 1.0 / gProjNorm;
        initGuess *= 2;
        return initGuess;
    }

    void SurfaceFlow::StepL2Unconstrained()
    {
        stepCount++;
        UpdateEnergies();

        Eigen::MatrixXd l2diff;
        l2diff.setZero(mesh->nVertices(), 3);
        AssembleGradients(l2diff);

        double initGuess = guessStepSize(l2diff.norm());
        LineSearch search(mesh, geom, energies, maxStepSize);
        search.BacktrackingLineSearch(l2diff, initGuess, 1);
    }

    void SurfaceFlow::StepL2Projected()
    {
        stepCount++;
        UpdateEnergies();

        Eigen::MatrixXd l2diff;
        l2diff.setZero(mesh->nVertices(), 3);
        AssembleGradients(l2diff);

        // Make a saddle matrix with just the identity in the corner
        std::vector<Triplet> triplets;
        for (size_t i = 0; i < mesh->nVertices(); i++)
        {
            size_t i3 = 3 * i;
            triplets.push_back(Triplet(i3, i3, 1));
            triplets.push_back(Triplet(i3 + 1, i3 + 1, 1));
            triplets.push_back(Triplet(i3 + 2, i3 + 2, 1));
        }

        size_t nConstraintRows = addConstraintTriplets(triplets, true);
        size_t dims = 3 * mesh->nVertices() + nConstraintRows;

        Eigen::SparseMatrix<double> A(dims, dims);
        A.setFromTriplets(triplets.begin(), triplets.end());

        Eigen::VectorXd l2col;
        l2col.setZero(dims);
        MatrixUtils::MatrixIntoColumn(l2diff, l2col);

        SparseFactorization factorizedA;
        factorizedA.Compute(A);

        l2col = factorizedA.Solve(l2col);
        MatrixUtils::ColumnIntoMatrix(l2col, l2diff);

        double initGuess = guessStepSize(l2diff.norm());
        LineSearch search(mesh, geom, energies, maxStepSize);
        search.BacktrackingLineSearch(l2diff, initGuess, 1);

        // Constraint projection
        l2col.setZero();
        size_t curRow = 3 * mesh->nVertices();

        for (Constraints::SimpleProjectorConstraint* spc : simpleConstraints)
        {
            spc->addErrorValues(l2col, mesh, geom, curRow);
            curRow += spc->nRows();
        }

        for (ConstraintPack pack : schurConstraints)
        {
            pack.constraint->addErrorValues(l2col, mesh, geom, curRow);
            curRow += pack.constraint->nRows();
        }

        l2col = factorizedA.Solve(l2col);

        VertexIndices inds = mesh->getVertexIndices();
        for (GCVertex v : mesh->vertices())
        {
            size_t base = inds[v] * 3;
            Vector3 corr{l2col(base), l2col(base + 1), l2col(base + 2)};
            geom->inputVertexPositions[v] -= corr;
        }

        geom->refreshQuantities();
    }

    double SurfaceFlow::evaluateEnergy()
    {
        return GetEnergyValue(energies);
    }

    void SurfaceFlow::AssembleGradients(Eigen::MatrixXd &dest)
    {
        AddGradientsToMatrix(energies, dest);
    }

    std::unique_ptr<Hs::HsMetric> SurfaceFlow::GetHsMetric()
    {
        std::unique_ptr<Hs::HsMetric> hs(new Hs::HsMetric(energies, obstacleEnergy, simpleConstraints, schurConstraints));
        hs->disableNearField = disableNearField;
        return hs;
    }

    void SurfaceFlow::StepProjectedGradientExact()
    {
        long timeStart = currentTimeMilliseconds();
        stepCount++;
        std::cout << "=== Iteration " << stepCount << " ===" << std::endl;
        std::cout << "Using Hs projected gradient method..." << std::endl;
        UpdateEnergies();

        // Assemble sum of L2 differentials of all energies involved
        // (including tangent-point energy)
        Eigen::MatrixXd l2diff, gradientProj;
        l2diff.setZero(mesh->nVertices(), 3);
        gradientProj.setZero(mesh->nVertices(), 3);

        AssembleGradients(l2diff);
        double gNorm = l2diff.norm();

        std::unique_ptr<Hs::HsMetric> hs = GetHsMetric();

        Eigen::MatrixXd M = hs->GetHsMatrixConstrained();

        // Flatten the gradient into a single column
        Eigen::VectorXd gradientCol;
        gradientCol.setZero(M.rows());

        // Solve the dense system
        MatrixUtils::MatrixIntoColumn(l2diff, gradientCol);
        Eigen::PartialPivLU<Eigen::MatrixXd> solver = M.partialPivLu();
        gradientCol = solver.solve(gradientCol);
        MatrixUtils::ColumnIntoMatrix(gradientCol, gradientProj);

        VertexIndices inds = mesh->getVertexIndices();

        double gProjNorm = gradientProj.norm();
        // Measure dot product of search direction with original gradient direction
        double gradDot = (l2diff.transpose() * gradientProj).trace() / (gNorm * gProjNorm);

        // Guess a step size
        // double initGuess = prevStep * 1.25;
        double initGuess = guessStepSize(gProjNorm);

        std::cout << "  * Initial step size guess = " << initGuess << std::endl;

        // Take the step using line search
        LineSearch search(mesh, geom, energies, maxStepSize);
        search.BacktrackingLineSearch(gradientProj, initGuess, gradDot);
        geom->refreshQuantities();

        // Reuse factorized matrix for constraint projection
        gradientCol.setZero();
        size_t curRow = 3 * mesh->nVertices();

        for (Constraints::SimpleProjectorConstraint *cons : simpleConstraints)
        {
            cons->addErrorValues(gradientCol, mesh, geom, curRow);
            curRow += cons->nRows();
        }

        for (const ConstraintPack &c : schurConstraints)
        {
            c.constraint->addErrorValues(gradientCol, mesh, geom, curRow);
            curRow += c.constraint->nRows();
        }

        gradientCol = solver.solve(gradientCol);

        VertexIndices verts = mesh->getVertexIndices();
        for (GCVertex v : mesh->vertices())
        {
            int base = 3 * verts[v];
            Vector3 vertCorr{gradientCol(base), gradientCol(base + 1), gradientCol(base + 2)};
            geom->inputVertexPositions[v] -= vertCorr;
        }

        long timeEnd = currentTimeMilliseconds();
        std::cout << "  Total time for gradient step = " << (timeEnd - timeStart) << " ms" << std::endl;
    }

    inline void printSolveInfo(size_t numNewton)
    {
        if (numNewton > 0)
        {
            std::cout << "  * With " << numNewton << " Newton constraint(s), Hs projection will require "
                      << (numNewton + 1) << " linear solves" << std::endl;
        }
        else
        {
            std::cout << "  * With no Newton constraints, Hs projection will require 1 linear solve" << std::endl;
        }
    }

    void SurfaceFlow::StepProjectedGradient()
    {
        ptic("SurfaceFlow::StepProjectedGradient");

        long timeStart = currentTimeMilliseconds();
        stepCount++;
        std::cout << "=== Iteration " << stepCount << " ===" << std::endl;
        std::cout << "Using Hs projected gradient method..." << std::endl;
        UpdateEnergies();

        // Assemble sum of L2 differentials of all energies involved
        // (including tangent-point energy)
        Eigen::MatrixXd l2diff, gradientProj;
        l2diff.setZero(mesh->nVertices(), 3);
        gradientProj.setZero(mesh->nVertices(), 3);

        AssembleGradients(l2diff);
        double gNorm = l2diff.norm();

        std::unique_ptr<Hs::HsMetric> hs = GetHsMetric();
        hs->allowBarycenterShift = allowBarycenterShift;
        printSolveInfo(hs->newtonConstraints.size());

        Vector3 shift{0, 0, 0};
        if (allowBarycenterShift)
        {
            shift = averageOfMatrixRows(geom, mesh, l2diff);
            std::cout << "Average shift of L2 diff = " << shift << std::endl;
        }

        Hs::ProjectViaSchur<Hs::SparseInverse>(*hs, l2diff, gradientProj);

        if (allowBarycenterShift)
        {
            addShiftToMatrixRows(gradientProj, mesh->nVertices(), shift);
        }

        VertexIndices inds = mesh->getVertexIndices();

        double gProjNorm = gradientProj.norm();
        // Measure dot product of search direction with original gradient direction
        double gradDot = (l2diff.transpose() * gradientProj).trace() / (gNorm * gProjNorm);

        // Guess a step size
        // double initGuess = prevStep * 1.25;
        double initGuess = guessStepSize(gProjNorm);

        std::cout << "  * Initial step size guess = " << initGuess << std::endl;

        // Take the step using line search
        LineSearch search(mesh, geom, energies, maxStepSize);
        double delta = search.BacktrackingLineSearch(gradientProj, initGuess, gradDot);

        if (schurConstraints.size() > 0)
        {
            Hs::ProjectSchurConstraints<Hs::SparseInverse>(*hs, 1);
        }

        if (allowBarycenterShift)
        {
            // The barycenter goes wherever it wants.
            // Assumes no pin constraints; free barycenter isn't meant to be used with pins.
        }
        else
        {
            hs->ProjectSimpleConstraints();
        }

        incrementSchurConstraints();
        geom->refreshQuantities();

        long timeEnd = currentTimeMilliseconds();
        std::cout << "  Total time for gradient step = " << (timeEnd - timeStart) << " ms" << std::endl;

        ptoc("SurfaceFlow::StepProjectedGradient");
    }

    void SurfaceFlow::StepProjectedGradientIterative()
    {
        ptic("SurfaceFlow::StepProjectedGradientIterative");

        long timeStart = currentTimeMilliseconds();
        stepCount++;
        std::cout << "=== Iteration " << stepCount << " ===" << std::endl;
        std::cout << "Using iterative Hs projected gradient method..." << std::endl;
        UpdateEnergies();

        // Assemble sum of L2 differentials of all energies involved
        // (including tangent-point energy)
        Eigen::MatrixXd l2diff, gradientProj;
        l2diff.setZero(mesh->nVertices(), 3);
        gradientProj.setZero(mesh->nVertices(), 3);
        AssembleGradients(l2diff);
        double gNorm = l2diff.norm();

        std::unique_ptr<Hs::HsMetric> hs = GetHsMetric();
        printSolveInfo(hs->newtonConstraints.size());

        Vector3 shift{0, 0, 0};
        if (allowBarycenterShift)
        {
            shift = averageOfMatrixRows(geom, mesh, l2diff);
            std::cout << "Average shift of L2 diff = " << shift << std::endl;
        }

        Hs::ProjectConstrainedHsIterativeMat(*hs, l2diff, gradientProj);

        if (allowBarycenterShift)
        {
            addShiftToMatrixRows(gradientProj, mesh->nVertices(), shift);
        }

        VertexIndices inds = mesh->getVertexIndices();

        double gProjNorm = gradientProj.norm();
        // Measure dot product of search direction with original gradient direction
        double gradDot = (l2diff.transpose() * gradientProj).trace() / (gNorm * gProjNorm);

        // Guess a step size
        // double initGuess = prevStep * 1.25;
        double initGuess = guessStepSize(gProjNorm);

        std::cout << "  * Initial step size guess = " << initGuess << std::endl;

        // Take the step using line search
        LineSearch search(mesh, geom, energies, maxStepSize);
        double delta = search.BacktrackingLineSearch(gradientProj, initGuess, gradDot);

        // Constraint projection
        if (schurConstraints.size() > 0)
        {
            // hs->ResetSchurComplement();
            std::cout << "  Projecting Newton constraints..." << std::endl;
            Hs::ProjectSchurConstraints<Hs::IterativeInverse>(*hs, 1);
        }
        if (allowBarycenterShift)
        {
            // The barycenter goes wherever it wants.
            // Assumes no pin constraints; free barycenter isn't meant to be used with pins.
        }
        else
        {
            hs->ProjectSimpleConstraints();
        }

        incrementSchurConstraints();
        geom->refreshQuantities();

        long timeEnd = currentTimeMilliseconds();
        std::cout << "  Total time for gradient step = " << (timeEnd - timeStart) << " ms" << std::endl;

        ptoc("SurfaceFlow::StepProjectedGradientIterative");
    }

    size_t SurfaceFlow::addConstraintTriplets(std::vector<Triplet> &triplets, bool includeSchur)
    {
        size_t curRow = 3 * mesh->nVertices();
        size_t nConstraintRows = 0;
        for (Constraints::SimpleProjectorConstraint *spc : simpleConstraints)
        {
            Constraints::addTripletsToSymmetric(*spc, triplets, mesh, geom, curRow);
            size_t nr = spc->nRows();
            curRow += nr;
            nConstraintRows += nr;
        }

        if (includeSchur)
        {
            for (ConstraintPack &c : schurConstraints)
            {
                Constraints::addTripletsToSymmetric(*c.constraint, triplets, mesh, geom, curRow);
                size_t nr = c.constraint->nRows();
                curRow += nr;
                nConstraintRows += nr;
            }
        }
        return nConstraintRows;
    }

    void SurfaceFlow::prefactorConstrainedLaplacian(SparseFactorization &factored, bool includeSchur)
    {
        Eigen::SparseMatrix<double> L;
        prefactorConstrainedLaplacian(L, factored, includeSchur);
    }

    void SurfaceFlow::prefactorConstrainedLaplacian(Eigen::SparseMatrix<double> &L, SparseFactorization &factored, bool includeSchur)
    {
        // Assemble triplets for the Laplacian
        std::vector<Triplet> h1Triplets, h1Triplets3x;
        H1::getTriplets(h1Triplets, mesh, geom, 1e-10);
        MatrixUtils::TripleTriplets(h1Triplets, h1Triplets3x);
        // Add constraint rows at bottom
        size_t nConstraintRows = addConstraintTriplets(h1Triplets3x, includeSchur);
        size_t dims = 3 * mesh->nVertices() + nConstraintRows;

        // Builds and factorize the Laplacian
        L.resize(dims, dims);
        L.setFromTriplets(h1Triplets3x.begin(), h1Triplets3x.end());
        factored.Compute(L);
    }

    void SurfaceFlow::StepH1ProjGrad()
    {
        long timeStart = currentTimeMilliseconds();
        stepCount++;
        std::cout << "=== Iteration " << stepCount << " ===" << std::endl;
        std::cout << "Using H1 projected gradient method..." << std::endl;
        UpdateEnergies();

        // Assemble sum of L2 differentials of all energies involved
        // (including tangent-point energy)
        Eigen::MatrixXd l2diff;
        l2diff.setZero(mesh->nVertices(), 3);
        AssembleGradients(l2diff);
        double gNorm = l2diff.norm();

        Vector3 shift{0, 0, 0};
        if (allowBarycenterShift)
        {
            shift = averageOfMatrixRows(geom, mesh, l2diff);
            std::cout << "Average shift of L2 diff = " << shift << std::endl;
        }

        SparseFactorization factorizedL;
        prefactorConstrainedLaplacian(factorizedL, true);
        size_t dims = factorizedL.nRows;
        std::cout << "Prefactorized" << std::endl;

        // Project the H1 gradient
        Eigen::VectorXd gradientVec;
        gradientVec.setZero(dims);
        MatrixUtils::MatrixIntoColumn(l2diff, gradientVec);

        gradientVec = factorizedL.Solve(gradientVec);
        Eigen::MatrixXd gradientProj;
        gradientProj.setZero(l2diff.rows(), l2diff.cols());
        MatrixUtils::ColumnIntoMatrix(gradientVec, gradientProj);

        if (allowBarycenterShift)
        {
            addShiftToMatrixRows(gradientProj, mesh->nVertices(), shift);
        }

        double gProjNorm = gradientProj.norm();
        // Measure dot product of search direction with original gradient direction
        double gradDot = (l2diff.transpose() * gradientProj).trace() / (gNorm * gProjNorm);
        // Guess a step size
        double initGuess = guessStepSize(gProjNorm);
        std::cout << "  * Initial step size guess = " << initGuess << std::endl;
        // Take the step using line search
        LineSearch search(mesh, geom, energies, maxStepSize);
        double delta = search.BacktrackingLineSearch(gradientProj, initGuess, gradDot);

        // Do corrective constraint projection by reusing the H1 metric
        H1::ProjectConstraints(mesh, geom, simpleConstraints, schurConstraints, factorizedL, 1);

        incrementSchurConstraints();
        geom->refreshQuantities();

        long timeEnd = currentTimeMilliseconds();
        std::cout << "  Total time for gradient step = " << (timeEnd - timeStart) << " ms" << std::endl;
    }

    void savePositions(MeshPtr &mesh, GeomPtr &geom, Eigen::MatrixXd &positions)
    {
        positions.setZero(mesh->nVertices(), 3);

        VertexIndices inds = mesh->getVertexIndices();

        for (GCVertex v : mesh->vertices())
        {
            Vector3 pos = geom->inputVertexPositions[v];
            size_t i = inds[v];
            positions(i, 0) = pos.x;
            positions(i, 1) = pos.y;
            positions(i, 2) = pos.z;
        }
    }

    void SurfaceFlow::StepAQP(double invKappa)
    {
        long timeStart = currentTimeMilliseconds();
        if (stepCount == 0 || verticesMutated)
        {
            // Reset Nesterov memory of previous step
            savePositions(mesh, geom, prevPositions1);
            savePositions(mesh, geom, prevPositions2);
        }
        else
        {
            double theta = (1 - sqrt(invKappa)) / (1 + sqrt(invKappa));
            // 1. Nesterov step
            Eigen::MatrixXd y_n = (1 + theta) * prevPositions1 - theta * prevPositions2;
            for (GCVertex v : mesh->vertices())
            {
                Vector3 pos = MatrixUtils::GetRowAsVector3(y_n, v.getIndex());
                geom->inputVertexPositions[v] = pos;
            }
            geom->refreshQuantities();
        }
        stepCount++;

        // 2. H1 gradient step following the Nesterov step
        std::cout << "=== Iteration " << stepCount << " ===" << std::endl;
        std::cout << "Using AQP..." << std::endl;
        UpdateEnergies();

        // Assemble sum of L2 differentials of all energies involved
        // (including tangent-point energy)
        Eigen::MatrixXd l2diff;
        l2diff.setZero(mesh->nVertices(), 3);
        AssembleGradients(l2diff);
        double gNorm = l2diff.norm();

        Vector3 shift{0, 0, 0};
        if (allowBarycenterShift)
        {
            shift = averageOfMatrixRows(geom, mesh, l2diff);
            std::cout << "Average shift of L2 diff = " << shift << std::endl;
        }

        SparseFactorization factorizedL;
        // Only use "simple" positional constraints (Nesterov would break hard constraints anyway)
        prefactorConstrainedLaplacian(factorizedL, false);
        size_t dims = factorizedL.nRows;

        // Project the H1 gradient
        Eigen::VectorXd gradientVec;
        gradientVec.setZero(dims);
        MatrixUtils::MatrixIntoColumn(l2diff, gradientVec);

        gradientVec = factorizedL.Solve(gradientVec);
        Eigen::MatrixXd gradientProj;
        gradientProj.setZero(l2diff.rows(), l2diff.cols());
        MatrixUtils::ColumnIntoMatrix(gradientVec, gradientProj);

        if (allowBarycenterShift)
        {
            addShiftToMatrixRows(gradientProj, mesh->nVertices(), shift);
        }

        double gProjNorm = gradientProj.norm();
        // Measure dot product of search direction with original gradient direction
        double gradDot = (l2diff.transpose() * gradientProj).trace() / (gNorm * gProjNorm);
        // Guess a step size
        double initGuess = guessStepSize(gProjNorm);
        std::cout << "  * Initial step size guess = " << initGuess << std::endl;
        // Take the step using line search
        LineSearch search(mesh, geom, energies, maxStepSize);
        double delta = search.BacktrackingLineSearch(gradientProj, initGuess, gradDot);

        // Make sure pins don't drift
        for (Constraints::SimpleProjectorConstraint *spc : simpleConstraints)
        {
            spc->ProjectConstraint(mesh, geom);
        }
        // Save previous positions
        prevPositions2 = prevPositions1;
        savePositions(mesh, geom, prevPositions1);
        geom->refreshQuantities();

        long timeEnd = currentTimeMilliseconds();
        std::cout << "  Total time for gradient step = " << (timeEnd - timeStart) << " ms" << std::endl;
    }

    void SurfaceFlow::StepH1LBFGS()
    {
        long timeStart = currentTimeMilliseconds();

        stepCount++;
        std::cout << "=== Iteration " << stepCount << " ===" << std::endl;
        std::cout << "Using H1 L-BFGS..." << std::endl;

        if (!lbfgs)
        {
            lbfgs = new H1_LBFGS(20, simpleConstraints);
        }

        if (verticesMutated)
        {
            std::cout << "  * Vertices were mutated; resetting memory" << std::endl;
            lbfgs->ResetMemory();
        }

        UpdateEnergies();
        // Assemble sum of L2 differentials of all energies involved
        // (including tangent-point energy)
        Eigen::MatrixXd l2diff, projected, positions;
        l2diff.setZero(mesh->nVertices(), 3);
        positions.setZero(mesh->nVertices(), 3);
        projected.setZero(mesh->nVertices(), 3);
        AssembleGradients(l2diff);
        double gNorm = l2diff.norm();

        Eigen::VectorXd l2diffvec(3 * mesh->nVertices());
        MatrixUtils::MatrixIntoColumn(l2diff, l2diffvec);

        savePositions(mesh, geom, positions);
        Eigen::VectorXd posvec(3 * mesh->nVertices());
        MatrixUtils::MatrixIntoColumn(positions, posvec);

        // Do the L-BFGS update with current position and gradient
        lbfgs->SetUpInnerProduct(mesh, geom);
        lbfgs->UpdateDirection(posvec, l2diffvec);
        double gProjNorm = lbfgs->direction().norm();

        MatrixUtils::ColumnIntoMatrix(lbfgs->direction(), projected);
        double gradDot = (l2diffvec.dot(lbfgs->direction())) / (gNorm * gProjNorm);
        std::cout << "  * Dot product = " << gradDot << std::endl;

        LineSearch search(mesh, geom, energies, maxStepSize);
        // Take the step using line search
        double initGuess = guessStepSize(gProjNorm);
        double delta = search.BacktrackingLineSearch(projected, initGuess, fmax(0, gradDot));

        if (gradDot < 0)
        {
            std::cout << "  * Negative dot product; resetting memory" << std::endl;
            lbfgs->ResetMemory();
        }

        // Make sure pins don't drift
        for (Constraints::SimpleProjectorConstraint *spc : simpleConstraints)
        {
            spc->ProjectConstraint(mesh, geom);
        }
        geom->refreshQuantities();

        long timeEnd = currentTimeMilliseconds();
        std::cout << "  Total time for gradient step = " << (timeEnd - timeStart) << " ms" << std::endl;
    }

    void SurfaceFlow::StepBQN()
    {
        long timeStart = currentTimeMilliseconds();

        stepCount++;
        std::cout << "=== Iteration " << stepCount << " ===" << std::endl;
        std::cout << "Using blended L-BFGS (BQN)..." << std::endl;

        if (!lbfgs)
        {
            // B = (total area) ^ (2(d-1) / d) for d = 3
            // (equation 13 of BCQN / Zhu et al.)
            double b = pow(totalArea(geom, mesh), 4.0 / 3.0);
            lbfgs = new BQN_LBFGS(20, simpleConstraints, b);
        }

        if (verticesMutated)
        {
            std::cout << "  * Vertices were mutated; resetting memory" << std::endl;
            lbfgs->ResetMemory();
        }

        UpdateEnergies();
        // Assemble sum of L2 differentials of all energies involved
        // (including tangent-point energy)
        Eigen::MatrixXd l2diff, projected, positions;
        l2diff.setZero(mesh->nVertices(), 3);
        positions.setZero(mesh->nVertices(), 3);
        projected.setZero(mesh->nVertices(), 3);
        AssembleGradients(l2diff);
        double gNorm = l2diff.norm();

        Eigen::VectorXd l2diffvec(3 * mesh->nVertices());
        MatrixUtils::MatrixIntoColumn(l2diff, l2diffvec);

        savePositions(mesh, geom, positions);
        Eigen::VectorXd posvec(3 * mesh->nVertices());
        MatrixUtils::MatrixIntoColumn(positions, posvec);

        // Do the L-BFGS update with current position and gradient
        lbfgs->SetUpInnerProduct(mesh, geom);
        lbfgs->UpdateDirection(posvec, l2diffvec);
        double gProjNorm = lbfgs->direction().norm();

        MatrixUtils::ColumnIntoMatrix(lbfgs->direction(), projected);
        double gradDot = (l2diffvec.dot(lbfgs->direction())) / (gNorm * gProjNorm);
        std::cout << "  * Dot product = " << gradDot << std::endl;

        LineSearch search(mesh, geom, energies, maxStepSize);
        // Take the step using line search
        double initGuess = guessStepSize(gProjNorm);
        double delta = search.BacktrackingLineSearch(projected, initGuess, fmax(0, gradDot));

        if (gradDot < 0)
        {
            std::cout << "  * Negative dot product; resetting memory" << std::endl;
            lbfgs->ResetMemory();
        }

        // Make sure pins don't drift
        for (Constraints::SimpleProjectorConstraint *spc : simpleConstraints)
        {
            spc->ProjectConstraint(mesh, geom);
        }
        geom->refreshQuantities();

        long timeEnd = currentTimeMilliseconds();
        std::cout << "  Total time for gradient step = " << (timeEnd - timeStart) << " ms" << std::endl;
    }

    void SurfaceFlow::StepH2Projected()
    {
        long timeStart = currentTimeMilliseconds();
        stepCount++;
        std::cout << "=== Iteration " << stepCount << " ===" << std::endl;
        std::cout << "Using H2 projected gradient..." << std::endl;
        UpdateEnergies();

        // Assemble sum of L2 differentials of all energies involved
        // (including tangent-point energy)
        Eigen::MatrixXd l2diff;
        l2diff.setZero(mesh->nVertices(), 3);
        AssembleGradients(l2diff);
        double gNorm = l2diff.norm();

        // Get bi-Laplacian triplets
        std::vector<Triplet> biTriplets, biTriplets3x;
        H2::getTriplets(biTriplets, mesh, geom, 0);
        MatrixUtils::TripleTriplets(biTriplets, biTriplets3x);
        size_t nConstraintRows = addConstraintTriplets(biTriplets3x, true);

        // Build tripled matrix
        size_t dims = 3 * mesh->nVertices() + nConstraintRows;
        Eigen::SparseMatrix<double> biLaplacian(dims, dims);
        biLaplacian.setFromTriplets(biTriplets3x.begin(), biTriplets3x.end());

        // Pre-factorize it
        SparseFactorization factorizedL;
        factorizedL.Compute(biLaplacian);

        // Reshape the gradient into a column
        Eigen::VectorXd gradientVec;
        gradientVec.setZero(dims);
        MatrixUtils::MatrixIntoColumn(l2diff, gradientVec);

        // Solve for the H2 gradient
        gradientVec = factorizedL.Solve(gradientVec);
        Eigen::MatrixXd gradientProj;
        gradientProj.setZero(l2diff.rows(), l2diff.cols());
        MatrixUtils::ColumnIntoMatrix(gradientVec, gradientProj);

        // Take the step using line search
        double gProjNorm = gradientProj.norm();
        double gradDot = (l2diff.transpose() * gradientProj).trace() / (gNorm * gProjNorm);
        double initGuess = guessStepSize(gProjNorm);
        std::cout << "  * Initial step size guess = " << initGuess << std::endl;
        LineSearch search(mesh, geom, energies, maxStepSize);
        double delta = search.BacktrackingLineSearch(gradientProj, initGuess, gradDot);

        // Do corrective constraint projection by reusing the metric
        H1::ProjectConstraints(mesh, geom, simpleConstraints, schurConstraints, factorizedL, 1);

        incrementSchurConstraints();
        geom->refreshQuantities();

        long timeEnd = currentTimeMilliseconds();
        std::cout << "  Total time for gradient step = " << (timeEnd - timeStart) << " ms" << std::endl;
    }

    void SurfaceFlow::RecenterMesh()
    {
        Vector3 center = meshBarycenter(geom, mesh);
        translateMesh(geom, mesh, origBarycenter - center);
    }

    void SurfaceFlow::ResetAllConstraints()
    {
        for (ConstraintPack &p : schurConstraints)
        {
            p.constraint->ResetFunction(mesh, geom);
        }
        for (Constraints::SimpleProjectorConstraint *c : simpleConstraints)
        {
            c->ResetFunction(mesh, geom);
        }
    }

    void SurfaceFlow::ResetAllPotentials()
    {
        for (SurfaceEnergy *energy : energies)
        {
            energy->ResetTargets();
        }
    }

    SurfaceEnergy *SurfaceFlow::BaseEnergy()
    {
        return energies[0];
    }

} // namespace rsurfaces


// Std includes
#include <iostream> // cout, endl
#include <vector>
#include <memory> // shared_ptr
#include <map>
// Thirdparties includes
#include <Eigen/Dense>
// Lib includes
#include "v0l/bin/file_data.h"
#include "m0sh/uniform.h"
#include "m0sh/structured_sub.h"
#include "fl0p/unstationary.h"

const unsigned int DIM = 3;

using TypeScalar = double;
using TypeVectorScalar = Eigen::Matrix<TypeScalar, 1, 1>;
using TypeVector = Eigen::Matrix<TypeScalar, DIM, 1>;
using TypeMatrix = Eigen::Matrix<TypeScalar, DIM, DIM>;
template<typename... Args>
using TypeRef = Eigen::Ref<Args...>;

template<typename ...Args>
using TypeContainer = std::vector<Args...>;
using TypeMesh = m0sh::Uniform<TypeVector, TypeRef, TypeContainer>;
using TypeMeshSub = m0sh::StructuredSub<TypeVector, TypeRef, TypeContainer>;
using TypeTimeMesh = m0sh::Uniform<TypeVectorScalar, TypeRef, TypeContainer>;
using TypeTimeMeshSub = m0sh::StructuredSub<TypeVectorScalar, TypeRef, TypeContainer>;
using TypeFlow = fl0w::fl0p::Unstationary<TypeVector, TypeMatrix, TypeRef, TypeMesh, TypeContainer, TypeMeshSub, TypeTimeMesh, TypeTimeMeshSub, TypeVectorScalar, v0l::FileData>;

void print(const TypeFlow& flow, const TypeVector& x, const TypeScalar& t) {
    std::cout << std::endl;
    std::cout << "flow.getVelocity(" << x.transpose() << ", " << t << ") = \n" << flow.getVelocity(x, t).transpose() << std::endl;
    //std::cout << "flow.getJacobian(" << x.transpose() << ", " << t << ") = \n" << flow.getJacobian(x, t) << std::endl;
    //std::cout << "flow.getVorticity(" << x.transpose() << ", " << t << ") = \n" << flow.getVorticity(x, t).transpose() << std::endl;
    //std::cout << "flow.getAcceleration(" << x.transpose() << ", " << t << ") = \n" << flow.getAcceleration(x, t).transpose() << std::endl;
    std::cout << std::endl;
}

int main () {
    std::vector<std::vector<v0l::FileData<float>>> velocity(2);
    for(std::size_t i = 0; i < velocity.size(); i++) {
        std::string fileName = std::string("../data/v050") + std::to_string(i) + ".vtk";
        velocity[i].emplace_back(fileName, 0);
        velocity[i].emplace_back(fileName, 1);
        velocity[i].emplace_back(fileName, 2);
    }
    // mesh data
    // // create lengths
    TypeVector origin;
    std::vector<double> lengths;
    for(std::size_t i = 0; i < velocity[0][0].meta.dimensions.size(); i++) {
        lengths.push_back(velocity[0][0].meta.dimensions[i] * velocity[0][0].meta.spacing[i]);
        origin[i] = velocity[0][0].meta.origin[i];
    }
    // // flow
    std::cout << "building flow..." << std::endl;
    TypeFlow flow(std::make_shared<TypeMesh>(velocity[0][0].meta.dimensions, lengths, origin, TypeContainer<bool>(DIM, true)), velocity, 1, std::make_shared<TypeTimeMesh>(std::vector<std::size_t>(1, 1), std::vector<double>(1, 0.02), TypeVectorScalar(0.0), std::vector<bool>(1, false)), 1);
    std::cout << "flow built !" << std::endl;
    // print
    print(flow, TypeVector({-0.5, -0.5, -0.5}), 0.0);
    print(flow, TypeVector({0.0, 0.0, 0.0}), 0.0);
    print(flow, TypeVector({0.5, 0.5, 0.5}), 0.0);
    print(flow, TypeVector({-0.5, -0.5, -0.5}), 0.01);
    print(flow, TypeVector({0.0, 0.0, 0.0}), 0.01);
    print(flow, TypeVector({0.5, 0.5, 0.5}), 0.01);
}


/* ---------------------------------------------------------------------
 *
 * Copyright (C) 1999 - 2021 by the deal.II authors
 *
 * This file is part of the deal.II library.
 *
 * The deal.II library is free software; you can use it, redistribute
 * it, and/or modify it under the terms of the GNU Lesser General
 * Public License as published by the Free Software Foundation; either
 * version 2.1 of the License, or (at your option) any later version.
 * The full text of the license can be found in the file LICENSE.md at
 * the top level directory of deal.II.
 *
 * ---------------------------------------------------------------------
 *
 *

 *
 *
 * This program was contributed by Peter Munch. This work and the required
 * generalizations of the internal data structures of deal.II form part of the
 * project "Virtual Materials Design" funded by the Helmholtz Association of
 * German Research Centres.
 *
 *
 *
 *
 *
Introduction

 *
 *
Motivation

 *
 * The motivation for using simplex meshes (as done in step-3simplex) is
 * straightforward: many freely available mesh-generation tools are very good in
 * creating good-quality meshes in such a format, while they struggle with
 * hex-only meshes. Hex-only meshes, on the other hand, are characterized with
 * better numerical properties (e.g., less degrees of freedom for the same
 * degree of accuracy and possibly better performance since the tensor-product
 * structure can be exploited) and are, as a consequence, the natural choice for
 * rather simple geometries and for meshes described by a coarse mesh with a few
 * cells (like a hyperball) and obtained in their final form through iterative
 * local refinement.
 *
 * Mixed meshes try to combine the best of both worlds by partitioning the
 * geometry in parts that can be easily meshed by hypercube cells
 * (quadrilaterals in 2D, hexahedrons in 3D) and in parts that can not be meshed
 * easily, requiring simplices (triangles in 2D, tetrahedrons in 3D). Since one
 * assumes that the region requiring simplices is rather small compared to the
 * rest of the domain where more efficient and accurate methods can be applied,
 * one can expect that the overall efficiency is hardly impacted by such an
 * approach.
 *
 * One should note that in 3D, one also needs a transition region between
 * hypercube and simplex regions. Here, one can use wedges/prisms and/or
 * pyramids.
 *
 *
 *
Working with mixed meshes

 *
 *
 * In the following, we concentrate, for the sake of simplicity, on 2D meshes:
 * they can only contain triangles and quadrilaterals. However, as detailed in
 * the outlook, an extension of the presented approach to the 3D case is
 * straightforward.
 *
 *
 * The complexity of working with mixed meshes in 2D results from the fact
 * that it contains of two
 * types of geometrical objects: quadrilaterals and triangles. How to deal with
 * quadrilaterals, we have discussed in step-3: we selected an appropriate
 * finite element, quadrature rule and mapping object, e.g., FE_Q, QGauss, and
 * MappingFE (initialized with FE_Q). For simplex meshes, we selected in
 * step-3simplex FE_SimplexP, QGaussSimplex, and MappingFE (intialized with
 * FE_SimplexP).
 *
 * For mixed meshes, we need multiple finite elements, quadrature rules, and
 * mapping objects (one set for triangles and one set for quadrilaterals) in the
 * same program. To ease the work with the multitude of objects (in particular
 * in 3D, we need at least four of each), you can collect the objects and group
 * them together in hp::FECollection, hp::QCollection, and
 * hp::MappingCollection.
 *
 * Just like in the context of finite elements, quadrature rules, and mapping
 * objects, we need multiple FEValues objects: the collection of FEValues is
 * called hp::FEValues. It returns the FEValues object needed for the current
 * cell via the method hp::FEValues::get_present_fe_values().
 *
 * For hp::FEValues, to be able to select the right finite element/quadrature
 * rule/ mapping object set, it queries the active_fe_index of the given cell
 * during hp::FEValues::reinit(). The indices have to be set - as shown below -
 * before calling DoFHandler::distribute_dofs() by the user.
 *
 *
 * The namespace name of hp::FECollection, hp::QCollection,
 * hp::MappingCollection, and hp::FEValues indicates that these classes have not
 * been written for mixed meshes in the first place, but for problems where each
 * (hypercube) cell could have a different type of finite element assigned - in
 * the simplest case, all cells have the same element type but different
 * polynomial degrees p (the reason for the letter "p" in "hp"). An extension of
 * this infrastructure to work not only on different element types but also on
 * different geometrical objects was a natural choice. For further details on
 * hp-methods, see step-27.
 *
 *
 *
Mesh generation

 *
 * Just like in step-3simplex, we read an externally generated mesh. For this
 * tutorial, we have created the mesh (square with width and height of one;
 * quadrilaterals on the left half and triangles on the right half) with Gmsh
 * with the following journal file "box_2D_mixed.geo":
 *
 * @code
 * Rectangle(1) = {0.0, 0, 0, 0.5, 1, 0};
 * Rectangle(2) = {0.5, 0, 0, 0.5, 1, 0};
 * Recombine Surface{1};
 * Physical Surface("All") = {1, 2};
 * Mesh 2;
 * Coherence Mesh;
 * Save "box_2D_mixed.msh";
 * @endcode
 *
 * The journal file can be processed by Gmsh generating the actual mesh with the
 * ending ".msh":
 *
 * @code
 * gmsh box_2D_mixed.geo
 * @endcode
 *
 * We have included in the tutorial folder both the journal file and the mesh
 * file in the event that one does not have access to Gmsh.
 *
 */

// @sect3{Include files}

// Include files, as used in step-3:
#include <deal.II/base/function.h>

#include <deal.II/dofs/dof_handler.h>
#include <deal.II/dofs/dof_tools.h>

#include <deal.II/fe/fe_q.h>
#include <deal.II/fe/fe_values.h>

#include <deal.II/grid/tria.h>

#include <deal.II/lac/dynamic_sparsity_pattern.h>
#include <deal.II/lac/full_matrix.h>
#include <deal.II/lac/precondition.h>
#include <deal.II/lac/solver_cg.h>
#include <deal.II/lac/sparse_matrix.h>
#include <deal.II/lac/vector.h>

#include <deal.II/numerics/data_out.h>
#include <deal.II/numerics/matrix_tools.h>
#include <deal.II/numerics/vector_tools.h>

#include <fstream>
#include <iostream>

// Include files, as added in step-3simplex:
#include <deal.II/base/quadrature_lib.h>

#include <deal.II/fe/fe_simplex_p.h>
#include <deal.II/fe/mapping_fe.h>

#include <deal.II/grid/grid_in.h>

// Include files that we need in this tutorial to be able to deal with
// collections of finite elements, quadrature rules, mapping objects, and
// FEValues.
#include <deal.II/hp/fe_collection.h>
#include <deal.II/hp/fe_values.h>
#include <deal.II/hp/mapping_collection.h>
#include <deal.II/hp/q_collection.h>

using namespace dealii;

// @sect3{The Step3 class}
//
// This is the main class of the tutorial. Since it is very similar to the
// version from step-3 and step-3simplex, we will only point out and explain
// the relevant differences that allow to perform simulations on mixed meshes.

class Step3
{
public:
  Step3();

  void run();

private:
  void make_grid();
  void setup_system();
  void assemble_system();
  void solve();
  void output_results() const;

  Triangulation<2> triangulation;

  // As already explained, we are not working with mapping objects, finite
  // elements, and quadrature rules directly but with collections of them.
  const hp::MappingCollection<2> mapping;
  const hp::FECollection<2>      fe;
  const hp::QCollection<2>       quadrature_formula;

  DoFHandler<2> dof_handler;

  SparsityPattern      sparsity_pattern;
  SparseMatrix<double> system_matrix;

  Vector<double> solution;
  Vector<double> system_rhs;
};

// @sect4{Step3::Step3}
//
// In the constructor of the Step3 class, we fill the collections. Here, we
// position the objects related to triangles in the first place (index 0) and
// the ones related to quadrilaterals in the second place (index 1).
Step3::Step3()
  : mapping(MappingFE<2>(FE_SimplexP<2>(1)), MappingFE<2>(FE_Q<2>(1)))
  , fe(FE_SimplexP<2>(2), FE_Q<2>(2))
  , quadrature_formula(QGaussSimplex<2>(3), QGauss<2>(3))
  , dof_handler(triangulation)
{}

// @sect4{Step3::make_grid}
//
// Read the external mesh file "box_2D_mixed.msh" as in step-3simplex.
void Step3::make_grid()
{
  GridIn<2>(triangulation).read("box_2D_mixed.msh");

  std::cout << "Number of active cells: " << triangulation.n_active_cells()
            << std::endl;
}

// @sect4{Step3::setup_system}
//
// In contrast to step-3 and step-3simplex, we need here a preprocessing step
// that assigns an active_fe_index to each cell consistently according to the
// indices in the collections and the cell type.
void Step3::setup_system()
{
  for (const auto &cell : dof_handler.active_cell_iterators())
    {
      if (cell->reference_cell() == ReferenceCells::Triangle)
        cell->set_active_fe_index(0);
      else if (cell->reference_cell() == ReferenceCells::Quadrilateral)
        cell->set_active_fe_index(1);
      else
        Assert(false, ExcNotImplemented());
    }

  dof_handler.distribute_dofs(fe);
  std::cout << "Number of degrees of freedom: " << dof_handler.n_dofs()
            << std::endl;
  DynamicSparsityPattern dsp(dof_handler.n_dofs());
  DoFTools::make_sparsity_pattern(dof_handler, dsp);
  sparsity_pattern.copy_from(dsp);

  system_matrix.reinit(sparsity_pattern);

  solution.reinit(dof_handler.n_dofs());
  system_rhs.reinit(dof_handler.n_dofs());
}

// @sect4{Step3::assemble_system}
//
// The following function looks similar to the version in step-3 and
// step-3simplex with the following two differences:
//  - We do not work with FEValues directly but with the collection class
//    hp::FEValues. It gives us - after it has been initialized with the current
//    cell - a reference to the right FEValues object (constructed
//    with the correct mapping object, finite element, and quadrature rule),
//    which can be used as usual to compute the cell integrals.
//  - The cell-local stiffness matrix and the right-hand-side vector have
//    different sizes depending on the cell type (6 DoFs vs. 9 DoFs) so that
//    they might need to be resized for each cell.
//
// Apart from these two changes, the code has not changed. Not even, the
// cell integrals have been changed depending on whether one operates on
// hypercube, simplex, or mixed meshes.
void Step3::assemble_system()
{
  hp::FEValues<2> hp_fe_values(mapping,
                               fe,
                               quadrature_formula,
                               update_values | update_gradients |
                                 update_JxW_values);

  FullMatrix<double>                   cell_matrix;
  Vector<double>                       cell_rhs;
  std::vector<types::global_dof_index> local_dof_indices;

  for (const auto &cell : dof_handler.active_cell_iterators())
    {
      hp_fe_values.reinit(cell);

      const auto &fe_values = hp_fe_values.get_present_fe_values();

      const unsigned int dofs_per_cell = cell->get_fe().n_dofs_per_cell();
      cell_matrix.reinit(dofs_per_cell, dofs_per_cell);
      cell_rhs.reinit(dofs_per_cell);
      local_dof_indices.resize(dofs_per_cell);

      cell_matrix = 0;
      cell_rhs    = 0;

      for (const unsigned int q_index : fe_values.quadrature_point_indices())
        {
          for (const unsigned int i : fe_values.dof_indices())
            for (const unsigned int j : fe_values.dof_indices())
              cell_matrix(i, j) +=
                (fe_values.shape_grad(i, q_index) * // grad phi_i(x_q)
                 fe_values.shape_grad(j, q_index) * // grad phi_j(x_q)
                 fe_values.JxW(q_index));           // dx

          for (const unsigned int i : fe_values.dof_indices())
            cell_rhs(i) += (fe_values.shape_value(i, q_index) * // phi_i(x_q)
                            1. *                                // f(x_q)
                            fe_values.JxW(q_index));            // dx
        }
      cell->get_dof_indices(local_dof_indices);

      for (const unsigned int i : fe_values.dof_indices())
        for (const unsigned int j : fe_values.dof_indices())
          system_matrix.add(local_dof_indices[i],
                            local_dof_indices[j],
                            cell_matrix(i, j));

      for (const unsigned int i : fe_values.dof_indices())
        system_rhs(local_dof_indices[i]) += cell_rhs(i);
    }

  std::map<types::global_dof_index, double> boundary_values;
  VectorTools::interpolate_boundary_values(
    mapping, dof_handler, 0, Functions::ZeroFunction<2>(), boundary_values);
  MatrixTools::apply_boundary_values(boundary_values,
                                     system_matrix,
                                     solution,
                                     system_rhs);
}

// @sect4{Step3::solve}
//
// Nothing has changed here.
void Step3::solve()
{
  SolverControl            solver_control(1000, 1e-12);
  SolverCG<Vector<double>> solver(solver_control);
  solver.solve(system_matrix, solution, system_rhs, PreconditionIdentity());
}

// @sect4{Step3::output_results}
//
// Nothing has changed here.
void Step3::output_results() const
{
  DataOut<2> data_out;

  DataOutBase::VtkFlags flags;
  flags.write_higher_order_cells = true;
  data_out.set_flags(flags);

  data_out.attach_dof_handler(dof_handler);
  data_out.add_data_vector(solution, "solution");
  data_out.build_patches(mapping, 2);
  std::ofstream output("solution.vtk");
  data_out.write_vtk(output);
}

// @sect4{Step3::run}
//
// Nothing has changed here.
void Step3::run()
{
  make_grid();
  setup_system();
  assemble_system();
  solve();
  output_results();
}

// @sect3{The main function}
//
// Nothing has changed here.
int main()
{
  deallog.depth_console(2);

  Step3 laplace_problem;
  laplace_problem.run();

  return 0;
}

/**
 *
Results

 *
 * The following figures show the mesh and the result obtained by executing this
 * program:
 *
 *

 *

 *
 *         @image html step_3_mixed_0.png
 *
 *
 *         @image html step_3_mixed_1.png
 *
 *

 *

 *
 * Not surprisingly, the result looks as expected.
 *
 *
 *
Possibilities for extensions

 *
 * In this tutorial, we presented how to use the deal.II simplex infrastructure
 * to solve a simple Poisson problem on a mixed mesh in 2D. In this scope, we
 * could only present a small section of the capabilities. In the following, we
 * point out further capabilities briefly.
 *
 *
 *
Pure hypercube and simplex meshes

 *
 * In this tutorial, we worked on a mesh consisting both of quadrilaterals and
 * triangles. However, the program and the underlying concepts also work if the
 * mesh only contains either quadrilaterals or triangles. Interested users can
 * try this out: we have provided appropriate journal files and meshes for such
 * cases.
 *
 *
 *
3D meshes

 *
 * In 3D, meshes might also consist of wedges/prisms and pyramids. Therefore,
 * the above introduced collections might consist of four components.
 *
 * For wedge/prism and pyramid cell types, following finite-element and
 * quadrature-rule classes are available:
 *  - wedge: FE_WedgeP, FE_WedgeDGP, QGaussWedge, MappingFE
 *  - pyramid: FE_PyramidP, FE_PyramidDGP, QGaussPyramid, MappingFE
 *
 *
Parallelization, face integrals, discontinuous Galerkin methods, and
 * matrix-free operator evaluation

 *
 * Regarding these aspects, the same comments are valid that are described in
 * step-3simplex for pure simplex meshes.
 *
 */


/**
 * io.cpp
 *
 * 2021 Gabriel Moreira
 *
 * https://github.com/gabmoreira/maks
 *
 * This software and the related documents  are provided as  is,  with no express
 * or implied  warranties,  other  than those  that are  expressly stated  in the
 * License.
 *
 * Copyright © 2021 Gabriel Moreira. All rights reserved.
 */

#include <iostream>
#include <string>
#include <vector>
#include <fstream>

#include <Eigen/Geometry>

#include "io.hpp"

using std::vector;
using std::string;
using std::stod;
using std::stoi;

using Eigen::Quaterniond;
using Eigen::Vector3d;
using Eigen::VectorXi;
using Eigen::Matrix3d;
using Eigen::MatrixXd;
using Eigen::Ref;

#define G2O_PRECISION 15

/**
 * Splits a string based on a delimiter.
 *
 * @param str Input string to be split.
 * @param delim Input string delimiter.
 * @return vector of tokens.
 */
vector<string> strSplit(const string& str, const string& delim) {
    vector<string> tokens;
    size_t prev = 0, pos = 0;
    do {
        pos = str.find(delim, prev);
        if (pos == string::npos)
            pos = str.length();
        string token = str.substr(prev, pos-prev);
        if (!token.empty())
            tokens.push_back(token);
        prev = pos + delim.length();
    } while (pos < str.length() && prev < str.length());

    return tokens;
};

/**
 * Reads g2o file.
 *
 * @param path (input) char*  - g2o file name.
 * @param node_i (output) Eigen::VectorXi - node indices.
 * @param node_r (output) Eigen::MatrixXd - SO(3) node rotations.
 * @param node_t (output) Eigen::MatrixXd - node translations.
 * @param edge_i (output) Eigen::VectorXi - edge indices (from).
 * @param edge_j (output) Eigen::VectorXi - edge indices (to).
 * @param edge_r (output) Eigen::MatrixXd - SO(3) edge rotations.
 * @param edge_t (output) Eigen::MatrixXd - edge translations.
 * @param num_nodes (output) int - number of nodes loaded.
 * @param num_edges (output) int - number of edges loaded.
 */
void readG2O(const char*   path,
             Ref<VectorXi> node_i,
             Ref<MatrixXd> node_r,
             Ref<MatrixXd> node_t,
             Ref<VectorXi> edge_i,
             Ref<VectorXi> edge_j,
             Ref<MatrixXd> edge_r,
             Ref<MatrixXd> edge_t,
             int&          num_nodes,
             int&          num_edges) {

    std::ifstream infile(path);
    if(!infile)
        return;

    string edge_tag = "EDGE_SE3:QUAT";
    string node_tag = "VERTEX_SE3:QUAT";

    num_nodes = 0;
    num_edges = 0;

    string line;
    while (std::getline(infile, line)) {
        vector<string> tokens = strSplit(line, " ");
        string tag = tokens[0];

        if (0 == tag.compare(edge_tag)) {
            int edge_i_id = stoi(tokens[1]);
            int edge_j_id = stoi(tokens[2]);

            // Edge translation
            Vector3d t( stod(tokens[3]), stod(tokens[4]), stod(tokens[5]) );

            // Edge quaternion
            Quaterniond q( stod(tokens[9]), stod(tokens[6]), stod(tokens[7]), stod(tokens[8]) );

            edge_i(num_edges) += edge_i_id;
            edge_j(num_edges) += edge_j_id;
            edge_r.block<3,3>(0,num_edges*3) += q.normalized().toRotationMatrix();
            edge_t.block<3,1>(0,num_edges) += t;
            num_edges++;

        } else if (0 == tag.compare(node_tag)) {
            int node_id = stoi(tokens[1]);

            // Node translation
            Vector3d t( stod(tokens[2]), stod(tokens[3]), stod(tokens[4]) );

            // Node quaternion
            Quaterniond q( stod(tokens[8]), stod(tokens[5]), stod(tokens[6]), stod(tokens[7]) );

            node_i(num_nodes) += node_id;
            node_r.block<3,3>(0,num_nodes*3) += q.normalized().toRotationMatrix();
            node_t.block<3,1>(0,num_nodes) += t;
            num_nodes++;
        };
    };
};

/**
 * Reads edge rotations from g2o file.
 *
 * @param path (input) char*  - g2o file name.
 * @param edge_i (output) Eigen::VectorXi - edge id (from).
 * @param edge_j (output) Eigen::VectorXi - edge id (to).
 * @param edge_r (output) Eigen::MatrixXd - Row block matrix containing SO(3) rotations.
 * @param num_edges (output) int - number of edges.
 */
void readG2O(const char*          path,
             Eigen::Ref<VectorXi> edge_i,
             Eigen::Ref<VectorXi> edge_j,
             Eigen::Ref<MatrixXd> edge_r,
             int&                 num_edges) {

    std::ifstream infile(path);
    if(!infile)
        return;

    string edge_tag = "EDGE_SE3:QUAT";

    num_edges = 0;

    string line;
    while (std::getline(infile, line)) {
        vector<string> tokens = strSplit(line, " ");
        string tag = tokens[0];

        if (0 == tag.compare(edge_tag)) {
            int edge_i_id = stoi(tokens[1]);
            int edge_j_id = stoi(tokens[2]);

            /* Edge quaternion */
            Quaterniond q( stod(tokens[9]), stod(tokens[6]), stod(tokens[7]), stod(tokens[8]) );

            edge_i(num_edges) += edge_i_id;
            edge_j(num_edges) += edge_j_id;
            edge_r.block<3,3>(0,num_edges*3) += q.normalized().toRotationMatrix();
            num_edges++;
        };
    };
};

void writeG2O(const char* filename,
              const vector<int>&         node_i,
              const vector<Quaterniond>& node_q,
              const vector<Vector3d>&    node_t,
              const vector<int>&         edge_i,
              const vector<int>&         edge_j,
              const vector<Quaterniond>& edge_q,
              const vector<Vector3d>&    edge_t) {

    int num_nodes = (int) node_i.size();
    int num_edges = (int) edge_i.size();

    string filepath = string(filename) + ".g2o";

    std::ostringstream buffer;
    buffer.clear();

    /* Write node data */
    for (int i = 0; i < num_nodes; ++i) {
        /* Write tag */
        buffer << "VERTEX_SE3:QUAT";

        /* Write node id */
        buffer << " " << std::to_string(node_i[i]);

        /* Write node translation */
        buffer << " " << std::fixed << std::setprecision(G2O_PRECISION) << node_t[i](0); // x
        buffer << " " << std::fixed << std::setprecision(G2O_PRECISION) << node_t[i](1); // y
        buffer << " " << std::fixed << std::setprecision(G2O_PRECISION) << node_t[i](2); // z

        /* Write node quaternion */
        buffer << " " << std::fixed << std::setprecision(G2O_PRECISION) << node_q[i].x();           // qx
        buffer << " " << std::fixed << std::setprecision(G2O_PRECISION) << node_q[i].y();           // qy
        buffer << " " << std::fixed << std::setprecision(G2O_PRECISION) << node_q[i].z();           // qz
        buffer << " " << std::fixed << std::setprecision(G2O_PRECISION) << node_q[i].w() << "\n";   // qw
    };

    /* Write edge data */
    for (int i = 0; i < num_edges; ++i) {
        /* Write tag */
        buffer << "EDGE_SE3:QUAT";

        /* Write edge ids */
        buffer << " " << std::to_string(edge_i[i]);
        buffer << " " << std::to_string(edge_j[i]);

        /* Write node translation */
        buffer << " " << std::fixed << std::setprecision(G2O_PRECISION) << edge_t[i](0); // x
        buffer << " " << std::fixed << std::setprecision(G2O_PRECISION) << edge_t[i](1); // y
        buffer << " " << std::fixed << std::setprecision(G2O_PRECISION) << edge_t[i](2); // z

        /* Write node quaternion */
        buffer << " " << std::fixed << std::setprecision(G2O_PRECISION) << edge_q[i].x();           // qx
        buffer << " " << std::fixed << std::setprecision(G2O_PRECISION) << edge_q[i].y();           // qy
        buffer << " " << std::fixed << std::setprecision(G2O_PRECISION) << edge_q[i].z();           // qz
        buffer << " " << std::fixed << std::setprecision(G2O_PRECISION) << edge_q[i].w() << "\n";   // qw
    };

    std::ofstream outfile;
    outfile.open(filepath);
    if(!outfile) {
        return;
    };

    outfile << buffer.str();
    outfile.close();
    printf("Pose graph saved to %s\n", filepath.c_str());
};