/**
 * Copyright (c) 2009 Carnegie Mellon University.
 *     All rights reserved.
 *
 *  Licensed under the Apache License, Version 2.0 (the "License");
 *  you may not use this file except in compliance with the License.
 *  You may obtain a copy of the License at
 *
 *      http://www.apache.org/licenses/LICENSE-2.0
 *
 *  Unless required by applicable law or agreed to in writing,
 *  software distributed under the License is distributed on an "AS
 *  IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either
 *  express or implied.  See the License for the specific language
 *  governing permissions and limitations under the License.
 *
 * For more about this software visit:
 *
 *      http://www.graphlab.ml.cmu.edu
 *
 */

/**
 * \file
 *
 * This file contains an implementation of the weighted-ALS matrix factorization
 * algorithm. As described in:  Collaborative Filtering for Implicit Feedback Datasets Hu, Y.; Koren, Y.; Volinsky, C. IEEE International Conference on Data Mining (ICDM 2008), IEEE (2008).
 *
 * Code written By Danny Bickson, based on code by Joey Gonzalez
 */

#include <Eigen/Dense>

#include <boost/config/warning_disable.hpp>
#include <boost/spirit/include/qi.hpp>
#include <boost/spirit/include/phoenix_core.hpp>
#include <boost/spirit/include/phoenix_operator.hpp>
#include <boost/spirit/include/phoenix_stl.hpp>

// This file defines the serialization code for the eigen types.
#include "eigen_serialization.hpp"

#include <graphlab.hpp>
#include <graphlab/util/stl_util.hpp>
#include "stats.hpp"

#include <graphlab/macros_def.hpp>

const int SAFE_NEG_OFFSET = 2; //add 2 to negative node id
//to prevent -0 and -1 which arenot allowed

/**
 * \brief We use the eigen library's vector type to represent
 * mathematical vectors.
 */
typedef Eigen::VectorXd vec_type;

/**
 * \brief We use the eigen library's matrix type to represent
 * matrices.
 */
typedef Eigen::MatrixXd mat_type;

/**
 * \brief Remap the target id of each edge into a different id space
 * than the source id.
 */
bool REMAP_TARGET = true;

/**
 * \ingroup toolkit_matrix_factorization
 *
 * \brief the vertex data type which contains the latent factor.
 *
 * Each row and each column in the matrix corresponds to a different
 * vertex in the ALS graph.  Associated with each vertex is a factor
 * (vector) of latent parameters that represent that vertex.  The goal
 * of the ALS algorithm is to find the values for these latent
 * parameters such that the non-zero entries in the matrix can be
 * predicted by taking the dot product of the row and column factors.
 */
struct vertex_data {
  /**
   * \brief A shared "constant" that specifies the number of latent
   * values to use.
   */
  static size_t NLATENT;
  /** \brief The number of times this vertex has been updated. */
  uint32_t nupdates;
  /** \brief The most recent L1 change in the factor value */
  float residual; //! how much the latent value has changed
  /** \brief The latent factor for this vertex */
  vec_type factor;
  /**
   * \brief Simple default constructor which randomizes the vertex
   *  data
   */
  vertex_data() : nupdates(0), residual(1) { randomize(); }
  /** \brief Randomizes the latent factor */
  void randomize() { factor.resize(NLATENT); factor.setRandom(); }
  /** \brief Save the vertex data to a binary archive */
  void save(graphlab::oarchive& arc) const {
    arc << nupdates << residual << factor;
  }
  /** \brief Load the vertex data from a binary archive */
  void load(graphlab::iarchive& arc) {
    arc >> nupdates >> residual >> factor;
  }
}; // end of vertex data

size_t vertex_data::NLATENT = 20;

/**
 * \brief The edge data stores the entry in the matrix.
 *
 * In addition the edge data also stores the most recent error estimate.
 */
struct edge_data : public graphlab::IS_POD_TYPE {
  /**
   * \brief The type of data on the edge;
   *
   * \li *Train:* the observed value is correct and used in training
   * \li *Validate:* the observed value is correct but not used in training
   * \li *Predict:* The observed value is not correct and should not be
   *        used in training.
   */
  enum data_role_type { TRAIN, VALIDATE, PREDICT  };

  /** \brief the observed value for the edge */
  float obs;

  /** \brief the weight or time of the observation */
  float weight;

  /** \brief The train/validation/test designation of the edge */
  data_role_type role;

  /** \brief basic initialization */
  edge_data(float obs = 0, data_role_type role = TRAIN, float weight = 1) :
    obs(obs), weight(weight), role(role) { }

}; // end of edge data

/**
 * \brief The graph type is defined in terms of the vertex and edge
 * data.
 */
typedef graphlab::distributed_graph<vertex_data, edge_data> graph_type;

#include "implicit.hpp"

stats_info count_edges(const graph_type::edge_type & edge){
  stats_info ret;

  if (edge.data().role == edge_data::TRAIN)
     ret.training_edges = 1;
  else if (edge.data().role == edge_data::VALIDATE)
     ret.validation_edges = 1;
  ret.max_user = (size_t)edge.source().id();
  ret.max_item = (size_t)edge.target().id();
  return ret;
}

/**
 * \brief Given a vertex and an edge return the other vertex in the
 * edge.
 */
inline graph_type::vertex_type
get_other_vertex(graph_type::edge_type& edge,
                 const graph_type::vertex_type& vertex) {
  return vertex.id() == edge.source().id()? edge.target() : edge.source();
}; // end of get_other_vertex

/**
 * \brief The gather type used to construct XtX and Xty needed for the ALS
 * update
 *
 * To compute the ALS update we need to compute the sum of
 * \code
 *  sum: XtX = nbr.factor.transpose() * nbr.factor
 *  sum: Xy  = nbr.factor * edge.obs
 * \endcode
 * For each of the neighbors of a vertex.
 *
 * To do this in the Gather-Apply-Scatter model the gather function
 * computes and returns a pair consisting of XtX and Xy which are then
 * added. The gather type represents that tuple and provides the
 * necessary gather_type::operator+= operation.
 *
 */
class gather_type {
public:
  /**
   * \brief Stores the current sum of nbr.factor.transpose() *
   * nbr.factor
   */
  mat_type XtX;

  /**
   * \brief Stores the current sum of nbr.factor * edge.obs
   */
  vec_type Xy;

  /**
   * \brief Stores the weight of this edge
   */
  float weight;

  /** \brief basic default constructor */
  gather_type() { }

  /**
   * \brief This constructor computes XtX and Xy and stores the result
   * in XtX and Xy
   */
  gather_type(const vec_type& X, const double y, const float weight) :
    XtX(X.size(), X.size()), Xy(X.size()) {
    XtX.triangularView<Eigen::Upper>() = X * X.transpose() * weight;
    Xy = X * y * weight;
  } // end of constructor for gather type

  /** \brief Save the values to a binary archive */
  void save(graphlab::oarchive& arc) const { arc << XtX << Xy << weight; }

  /** \brief Read the values from a binary archive */
  void load(graphlab::iarchive& arc) { arc >> XtX >> Xy >> weight; }

  /**
   * \brief Computes XtX += other.XtX and Xy += other.Xy updating this
   * tuples value
   */
  gather_type& operator+=(const gather_type& other) {
    if(other.Xy.size() == 0) {
      ASSERT_EQ(other.XtX.rows(), 0);
      ASSERT_EQ(other.XtX.cols(), 0);
    } else {
      if(Xy.size() == 0) {
        ASSERT_EQ(XtX.rows(), 0);
        ASSERT_EQ(XtX.cols(), 0);
        XtX = other.XtX; Xy = other.Xy;
      } else {
        XtX.triangularView<Eigen::Upper>() += other.XtX;
        Xy += other.Xy;
      }
    }
    return *this;
  } // end of operator+=

}; // end of gather type

/**
 * \brief WALS vertex program implements the alternating least squares
 * algorithm in the Gather-Apply-Scatter abstraction.
 *
 * The ALS update treats adjacent vertices (rows or columns) as "X"
 * (independent) values and the edges (matrix entries) as observed "y"
 * (dependent) values and then updates the current vertex value as a
 * weight "w" such that:
 *
 *    y = X * w + noise
 *
 * This is accomplished using the following equation:
 *
 *    w = inv(X' * X) * (X * y)
 *
 * We implement this in the Gather-Apply-Scatter model by:
 *
 *  1) Gather: returns the tuple (X' * X, X * y)
 *     Sum:   (aX' * aX, aX * ay) + (bX' * bX, bX * by) =
 *                 (aX' * aX + bX' * bX, aX * ay + bX * by)
 *
 *  2) Apply: Solves  inv(X' * X) * (X * y)
 *
 *  3) Scatter: schedules the update of adjacent vertices if this
 *      vertex has changed sufficiently and the edge is not well
 *      predicted.
 *
 *
 */
class als_vertex_program :
  public graphlab::ivertex_program<graph_type, gather_type,
                                   graphlab::messages::sum_priority>,
  public graphlab::IS_POD_TYPE {
public:
  /** The convergence tolerance */
  static double TOLERANCE;
  static double LAMBDA;
  static size_t MAX_UPDATES;
  static double MAXVAL;
  static double MINVAL;

  /** The set of edges to gather along */
  edge_dir_type gather_edges(icontext_type& context,
                             const vertex_type& vertex) const {
    return graphlab::ALL_EDGES;
  }; // end of gather_edges

  /** The gather function computes XtX and Xy */
  gather_type gather(icontext_type& context, const vertex_type& vertex,
                     edge_type& edge) const {
    if(edge.data().role == edge_data::TRAIN) {
      const vertex_type other_vertex = get_other_vertex(edge, vertex);
      return gather_type(other_vertex.data().factor, edge.data().obs, edge.data().weight);
    } else return gather_type();
  } // end of gather function

  /** apply collects the sum of XtX and Xy */
  void apply(icontext_type& context, vertex_type& vertex,
             const gather_type& sum) {
    // Get and reset the vertex data
    vertex_data& vdata = vertex.data();
    // Determine the number of neighbors.  Each vertex has only in or
    // out edges depending on which side of the graph it is located
    if(sum.Xy.size() == 0) { vdata.residual = 0; ++vdata.nupdates; return; }
    mat_type XtX = sum.XtX;
    vec_type Xy = sum.Xy;
    // Add regularization
    for(int i = 0; i < XtX.rows(); ++i) XtX(i,i) += LAMBDA; // /nneighbors;
    // Solve the least squares problem using eigen ----------------------------
    const vec_type old_factor = vdata.factor;
    vdata.factor = XtX.selfadjointView<Eigen::Upper>().ldlt().solve(Xy);
    // Compute the residual change in the factor factor -----------------------
    vdata.residual = (vdata.factor - old_factor).cwiseAbs().sum() / XtX.rows();
    ++vdata.nupdates;
  } // end of apply

  /** The edges to scatter along */
  edge_dir_type scatter_edges(icontext_type& context,
                              const vertex_type& vertex) const {
    return graphlab::ALL_EDGES;
  }; // end of scatter edges

  /** Scatter reschedules neighbors */
  void scatter(icontext_type& context, const vertex_type& vertex,
               edge_type& edge) const {
    edge_data& edata = edge.data();
    if(edata.role == edge_data::TRAIN) {
      const vertex_type other_vertex = get_other_vertex(edge, vertex);
      const vertex_data& vdata = vertex.data();
      const vertex_data& other_vdata = other_vertex.data();
      const double pred = vdata.factor.dot(other_vdata.factor);
      const float error = std::fabs(edata.obs - pred);
      const double priority = (error * vdata.residual);
      // Reschedule neighbors ------------------------------------------------
      if( priority > TOLERANCE && other_vdata.nupdates < MAX_UPDATES)
        context.signal(other_vertex, priority);
    }
  } // end of scatter function

  /**
   * \brief Signal all vertices on one side of the bipartite graph
   */
  static graphlab::empty signal_left(icontext_type& context,
                                     const vertex_type& vertex) {
    if(vertex.num_out_edges() > 0) context.signal(vertex);
    return graphlab::empty();
  } // end of signal_left

}; // end of als vertex program

/**
 * \brief The graph loader function is a line parser used for
 * distributed graph construction.
 */
inline bool graph_loader(graph_type& graph,
                         const std::string& filename,
                         const std::string& line) {
  ASSERT_FALSE(line.empty());
  namespace qi = boost::spirit::qi;
  namespace ascii = boost::spirit::ascii;
  namespace phoenix = boost::phoenix;
  // Determine the role of the data
  edge_data::data_role_type role = edge_data::TRAIN;
  if(boost::ends_with(filename,".validate")) role = edge_data::VALIDATE;
  else if(boost::ends_with(filename, ".predict")) role = edge_data::PREDICT;
  // Parse the line
  std::stringstream strm(line);
  graph_type::vertex_id_type source_id(-1), target_id(-1);
  float obs(0), weight(1);
  strm >> source_id >> target_id;

  // for test files (.predict) no need to read the actual rating value.
  if(role == edge_data::TRAIN || role == edge_data::VALIDATE){
    strm >> obs >> weight;
    if (obs < als_vertex_program::MINVAL || obs > als_vertex_program::MAXVAL)
      logstream(LOG_FATAL)<<"Rating values should be between " << als_vertex_program::MINVAL << " and " << als_vertex_program::MAXVAL << ". Got value: " << obs << " [ user: " << source_id << " to item: " <<target_id << " ] " << std::endl;
  }
  target_id = -(graphlab::vertex_id_type(target_id + SAFE_NEG_OFFSET));

  // Create an edge and add it to the graph
  graph.add_edge(source_id, target_id, edge_data(obs, role, weight));
  return true; // successful load
}

// end of graph_loader

/**
 * \brief Given an edge compute the error associated with that edge
 */
double extract_l2_error(const graph_type::edge_type & edge) {
  double pred =
    edge.source().data().factor.dot(edge.target().data().factor);
  pred = std::min(als_vertex_program::MAXVAL, pred);
  pred = std::max(als_vertex_program::MINVAL, pred);
  return (edge.data().obs - pred) * (edge.data().obs - pred) * edge.data().weight;
} // end of extract_l2_error

double als_vertex_program::TOLERANCE = 1e-3;
double als_vertex_program::LAMBDA = 0.01;
size_t als_vertex_program::MAX_UPDATES = -1;
double als_vertex_program::MAXVAL = 1e+100;
double als_vertex_program::MINVAL = -1e+100;

/**
 * \brief The error aggregator is used to accumulate the overal
 * prediction error.
 *
 * The error aggregator is itself a "reduction type" and contains the
 * two static methods "map" and "finalize" which operate on
 * error_aggregators and are used by the engine.add_edge_aggregator
 * api.
 */
struct error_aggregator : public graphlab::IS_POD_TYPE {
  typedef als_vertex_program::icontext_type icontext_type;
  typedef graph_type::edge_type edge_type;
  double train_error, validation_error;
  error_aggregator() :
    train_error(0), validation_error(0){ }
  error_aggregator& operator+=(const error_aggregator& other) {
    train_error += other.train_error;
    validation_error += other.validation_error;
    return *this;
  }
  static error_aggregator map(icontext_type& context, const graph_type::edge_type& edge) {
    error_aggregator agg;
    if(edge.data().role == edge_data::TRAIN) {
      agg.train_error = extract_l2_error(edge);
    } else if(edge.data().role == edge_data::VALIDATE) {
      agg.validation_error = extract_l2_error(edge);
    }
    return agg;
  }
  static void finalize(icontext_type& context, const error_aggregator& agg) {
    const double train_error = std::sqrt(agg.train_error / info.training_edges);
    context.cout() << context.elapsed_seconds() << "\t" << train_error;
    if(info.validation_edges > 0) {
      const double validation_error =
        std::sqrt(agg.validation_error / info.validation_edges);
      context.cout() << "\t" << validation_error;
    }
    context.cout() << std::endl;
  }
}; // end of error aggregator

/**
 * \brief The prediction saver is used by the graph.save routine to
 * output the final predictions back to the filesystem.
 */
struct prediction_saver {
  typedef graph_type::vertex_type vertex_type;
  typedef graph_type::edge_type   edge_type;
  std::string save_vertex(const vertex_type& vertex) const {
    return ""; //nop
  }
  std::string save_edge(const edge_type& edge) const {
    if(edge.data().role == edge_data::PREDICT) {
      std::stringstream strm;
      const double prediction =
        edge.source().data().factor.dot(edge.target().data().factor);
      strm << edge.source().id() << '\t';
      if(REMAP_TARGET) strm << (-edge.target().id() - SAFE_NEG_OFFSET) << '\t';
      else strm << edge.target().id() << '\t';
      strm << prediction << '\n';
      return strm.str();
    } else return "";
  }
}; // end of prediction_saver

struct linear_model_saver_U {
  typedef graph_type::vertex_type vertex_type;
  typedef graph_type::edge_type   edge_type;
  /* save the linear model, using the format:
     nodeid) factor1 factor2 ... factorNLATENT \n
  */
  std::string save_vertex(const vertex_type& vertex) const {
    if (vertex.num_out_edges() > 0){
      std::string ret = boost::lexical_cast<std::string>(vertex.id()) + " ";
      for (uint i=0; i< vertex_data::NLATENT; i++)
        ret += boost::lexical_cast<std::string>(vertex.data().factor[i]) + " ";
        ret += "\n";
      return ret;
    }
    else return "";
  }
  std::string save_edge(const edge_type& edge) const {
    return "";
  }
};

struct linear_model_saver_V {
  typedef graph_type::vertex_type vertex_type;
  typedef graph_type::edge_type   edge_type;
  /* save the linear model, using the format:
     nodeid) factor1 factor2 ... factorNLATENT \n
  */
  std::string save_vertex(const vertex_type& vertex) const {
    if (vertex.num_out_edges() == 0){
      std::string ret = boost::lexical_cast<std::string>(-vertex.id()-SAFE_NEG_OFFSET) + ") ";
      for (uint i=0; i< vertex_data::NLATENT; i++)
        ret += boost::lexical_cast<std::string>(vertex.data().factor[i]) + " ";
        ret += "\n";
      return ret;
    }
    else return "";
  }
  std::string save_edge(const edge_type& edge) const {
    return "";
  }
};

/**
 * \brief The engine type used by the ALS matrix factorization
 * algorithm.
 *
 * The ALS matrix factorization algorithm currently uses the
 * synchronous engine.  However we plan to add support for alternative
 * engines in the future.
 */
typedef graphlab::omni_engine<als_vertex_program> engine_type;

int main(int argc, char** argv) {
  global_logger().set_log_level(LOG_INFO);
  global_logger().set_log_to_console(true);

  // Parse command line options -----------------------------------------------
  const std::string description =
    "Compute the Weighted-ALS factorization of a matrix.";
  graphlab::command_line_options clopts(description);
  std::string input_dir, output_dir;
  std::string predictions;
  size_t interval = 10;
  std::string exec_type = "synchronous";
  clopts.attach_option("matrix", input_dir,
                       "The directory containing the matrix file");
  clopts.add_positional("matrix");
  clopts.attach_option("D",  vertex_data::NLATENT,
                       "Number of latent parameters to use.");
  clopts.attach_option("max_iter", als_vertex_program::MAX_UPDATES,
                       "The maxumum number of udpates allowed for a vertex");
  clopts.attach_option("lambda", als_vertex_program::LAMBDA,
                       "wALS regularization weight");
  clopts.attach_option("tol", als_vertex_program::TOLERANCE,
                       "residual termination threshold");
  clopts.attach_option("maxval", als_vertex_program::MAXVAL, "max allowed value");
  clopts.attach_option("minval", als_vertex_program::MINVAL, "min allowed value");
  clopts.attach_option("interval", interval,
                       "The time in seconds between error reports");
  clopts.attach_option("predictions", predictions,
                       "The prefix (folder and filename) to save predictions.");
  clopts.attach_option("engine", exec_type,
                       "The engine type synchronous or asynchronous");
  // clopts.attach_option("remap", REMAP_TARGET,
  //                      "Renumber target vertex ids (internally) so that they\n"
  //                      "are in a different range allowing user 0 to connect to movie 0");
  clopts.attach_option("output", output_dir,
                       "Output results");
  if(!clopts.parse(argc, argv) || input_dir == "") {
    std::cout << "Error in parsing command line arguments." << std::endl;
    clopts.print_description();
    return EXIT_FAILURE;
  }

  ///! Initialize control plain using mpi
  graphlab::mpi_tools::init(argc, argv);
  graphlab::distributed_control dc;

  dc.cout() << "Loading graph." << std::endl;
  graphlab::timer timer;
  graph_type graph(dc, clopts);
  graph.load(input_dir, graph_loader);
  dc.cout() << "Loading graph. Finished in "
            << timer.current_time() << std::endl;

  if (dc.procid() == 0)
    add_implicit_edges4<edge_data>(implicitratingtype, graph, dc);

  dc.cout() << "Finalizing graph." << std::endl;
  timer.start();
  graph.finalize();
  dc.cout() << "Finalizing graph. Finished in "
            << timer.current_time() << std::endl;

  dc.cout()
      << "========== Graph statistics on proc " << dc.procid()
      << " ==============="
      << "\n Num vertices: " << graph.num_vertices()
      << "\n Num edges: " << graph.num_edges()
      << "\n Num replica: " << graph.num_replicas()
      << "\n Replica to vertex ratio: "
      << float(graph.num_replicas())/graph.num_vertices()
      << "\n --------------------------------------------"
      << "\n Num local own vertices: " << graph.num_local_own_vertices()
      << "\n Num local vertices: " << graph.num_local_vertices()
      << "\n Replica to own ratio: "
      << (float)graph.num_local_vertices()/graph.num_local_own_vertices()
      << "\n Num local edges: " << graph.num_local_edges()
      //<< "\n Begin edge id: " << graph.global_eid(0)
      << "\n Edge balance ratio: "
      << float(graph.num_local_edges())/graph.num_edges()
      << std::endl;

  dc.cout() << "Creating engine" << std::endl;
  engine_type engine(dc, graph, exec_type, clopts);

  // Add error reporting to the engine
  const bool success = engine.add_edge_aggregator<error_aggregator>
    ("error", error_aggregator::map, error_aggregator::finalize) &&
    engine.aggregate_periodic("error", interval);
  ASSERT_TRUE(success);

  // Signal all vertices on the vertices on the left (liberals)
  engine.map_reduce_vertices<graphlab::empty>(als_vertex_program::signal_left);
  info = graph.map_reduce_edges<stats_info>(count_edges);
  dc.cout()<<"Training edges: " << info.training_edges << " validation edges: " << info.validation_edges << std::endl;

  // Run the WALS ---------------------------------------------------------
  dc.cout() << "Running Weighted-ALS" << std::endl;
  timer.start();
  engine.start();

  const double runtime = timer.current_time();
  dc.cout() << "----------------------------------------------------------"
            << std::endl
            << "Final Runtime (seconds):   " << runtime
            << std::endl
            << "Updates executed: " << engine.num_updates() << std::endl
            << "Update Rate (updates/second): "
            << engine.num_updates() / runtime << std::endl;

  // Compute the final training error -----------------------------------------
  dc.cout() << "Final error: " << std::endl;
  engine.aggregate_now("error");

  // Make predictions ---------------------------------------------------------
  if(!predictions.empty()) {
    std::cout << "Saving predictions" << std::endl;
    const bool gzip_output = false;
    const bool save_vertices = false;
    const bool save_edges = true;
    const size_t threads_per_machine = 2;

    //save the predictions
    graph.save(predictions, prediction_saver(),
               gzip_output, save_vertices,
               save_edges, threads_per_machine);
    //save the linear model
    graph.save(predictions + ".U", linear_model_saver_U(),
		gzip_output, save_edges, save_vertices, threads_per_machine);
    graph.save(predictions + ".V", linear_model_saver_V(),
		gzip_output, save_edges, save_vertices, threads_per_machine);

  }

  graphlab::mpi_tools::finalize();
  return EXIT_SUCCESS;
} // end of main


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

/*!
 Copyright (C) 2007 Allen Kuo
 Copyright (C) 2015 Andres Hernandez

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

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

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

/*  This example shows how to fit a term structure to a set of bonds
    using four different fitting methodologies. Though fitting is most
    useful for large numbers of bonds with non-smooth yield tenor
    structures, for comparison purposes, relatively smooth bond yields
    are fit here and compared to known solutions (par coupons), or
    results generated from the bootstrap fitting method.
*/

#include <ql/qldefines.hpp>
#ifdef BOOST_MSVC
#  include <ql/auto_link.hpp>
#endif
#include <ql/termstructures/yield/fittedbonddiscountcurve.hpp>
#include <ql/termstructures/yield/piecewiseyieldcurve.hpp>
#include <ql/termstructures/yield/flatforward.hpp>
#include <ql/termstructures/yield/bondhelpers.hpp>
#include <ql/termstructures/yield/nonlinearfittingmethods.hpp>
#include <ql/pricingengines/bond/bondfunctions.hpp>
#include <ql/time/calendars/target.hpp>
#include <ql/time/daycounters/simpledaycounter.hpp>

#include <boost/timer.hpp>
#include <iostream>
#include <iomanip>
#include <boost/make_shared.hpp>

#define LENGTH(a) (sizeof(a)/sizeof(a[0]))

using namespace std;
using namespace QuantLib;

#if defined(QL_ENABLE_SESSIONS)
namespace QuantLib {
    Integer sessionId() { return 0; }
}
#endif

// par-rate approximation
Rate parRate(const YieldTermStructure& yts,
             const std::vector<Date>& dates,
             const DayCounter& resultDayCounter) {
    QL_REQUIRE(dates.size() >= 2, "at least two dates are required");
    Real sum = 0.0;
    Time dt;
    for (Size i=1; i<dates.size(); ++i) {
        dt = resultDayCounter.yearFraction(dates[i-1], dates[i]);
        QL_REQUIRE(dt>=0.0, "unsorted dates");
        sum += yts.discount(dates[i]) * dt;
    }
    Real result = yts.discount(dates.front()) - yts.discount(dates.back());
    return result/sum;
}

void printOutput(const std::string& tag,
                 const boost::shared_ptr<FittedBondDiscountCurve>& curve) {
    cout << tag << endl;
    cout << "reference date : "
         << curve->referenceDate()
         << endl;
    cout << "number of iterations : "
         << curve->fitResults().numberOfIterations()
         << endl
         << endl;
}

int main(int, char* []) {

    try {

        boost::timer timer;

        const Size numberOfBonds = 15;
        Real cleanPrice[numberOfBonds];

        for (Size i=0; i<numberOfBonds; i++) {
            cleanPrice[i]=100.0;
        }

        std::vector< boost::shared_ptr<SimpleQuote> > quote;
        for (Size i=0; i<numberOfBonds; i++) {
            boost::shared_ptr<SimpleQuote> cp(new SimpleQuote(cleanPrice[i]));
            quote.push_back(cp);
        }

        RelinkableHandle<Quote> quoteHandle[numberOfBonds];
        for (Size i=0; i<numberOfBonds; i++) {
            quoteHandle[i].linkTo(quote[i]);
        }

        Integer lengths[] = { 2, 4, 6, 8, 10, 12, 14, 16,
                              18, 20, 22, 24, 26, 28, 30 };
        Real coupons[] = { 0.0200, 0.0225, 0.0250, 0.0275, 0.0300,
                           0.0325, 0.0350, 0.0375, 0.0400, 0.0425,
                           0.0450, 0.0475, 0.0500, 0.0525, 0.0550 };

        Frequency frequency = Annual;
        DayCounter dc = SimpleDayCounter();
        BusinessDayConvention accrualConvention = ModifiedFollowing;
        BusinessDayConvention convention = ModifiedFollowing;
        Real redemption = 100.0;

        Calendar calendar = TARGET();
        Date today = calendar.adjust(Date::todaysDate());
        Date origToday = today;
        Settings::instance().evaluationDate() = today;

        // changing bondSettlementDays=3 increases calculation
        // time of exponentialsplines fitting method
        Natural bondSettlementDays = 0;
        Natural curveSettlementDays = 0;

        Date bondSettlementDate = calendar.advance(today, bondSettlementDays*Days);

        cout << endl;
        cout << "Today's date: " << today << endl;
        cout << "Bonds' settlement date: " << bondSettlementDate << endl;
        cout << "Calculating fit for 15 bonds....." << endl << endl;

        std::vector<boost::shared_ptr<BondHelper> > instrumentsA;
        std::vector<boost::shared_ptr<RateHelper> > instrumentsB;

        for (Size j=0; j<LENGTH(lengths); j++) {

            Date maturity = calendar.advance(bondSettlementDate, lengths[j]*Years);

            Schedule schedule(bondSettlementDate, maturity, Period(frequency),
                              calendar, accrualConvention, accrualConvention,
                              DateGeneration::Backward, false);

            boost::shared_ptr<BondHelper> helperA(
                     new FixedRateBondHelper(quoteHandle[j],
                                             bondSettlementDays,
                                             100.0,
                                             schedule,
                                             std::vector<Rate>(1,coupons[j]),
                                             dc,
                                             convention,
                                             redemption));

            boost::shared_ptr<RateHelper> helperB(
                     new FixedRateBondHelper(quoteHandle[j],
                                             bondSettlementDays,
                                             100.0,
                                             schedule,
                                             std::vector<Rate>(1, coupons[j]),
                                             dc,
                                             convention,
                                             redemption));
            instrumentsA.push_back(helperA);
            instrumentsB.push_back(helperB);
        }

        bool constrainAtZero = true;
        Real tolerance = 1.0e-10;
        Size max = 5000;

        boost::shared_ptr<YieldTermStructure> ts0 (
              new PiecewiseYieldCurve<Discount,LogLinear>(curveSettlementDays,
                                                          calendar,
                                                          instrumentsB,
                                                          dc));

        ExponentialSplinesFitting exponentialSplines(constrainAtZero);

        boost::shared_ptr<FittedBondDiscountCurve> ts1 (
                  new FittedBondDiscountCurve(curveSettlementDays,
                                              calendar,
                                              instrumentsA,
                                              dc,
                                              exponentialSplines,
                                              tolerance,
                                              max));

        printOutput("(a) exponential splines", ts1);

        SimplePolynomialFitting simplePolynomial(3, constrainAtZero);

        boost::shared_ptr<FittedBondDiscountCurve> ts2 (
                    new FittedBondDiscountCurve(curveSettlementDays,
                                                calendar,
                                                instrumentsA,
                                                dc,
                                                simplePolynomial,
                                                tolerance,
                                                max));

        printOutput("(b) simple polynomial", ts2);

        NelsonSiegelFitting nelsonSiegel;

        boost::shared_ptr<FittedBondDiscountCurve> ts3 (
                        new FittedBondDiscountCurve(curveSettlementDays,
                                                    calendar,
                                                    instrumentsA,
                                                    dc,
                                                    nelsonSiegel,
                                                    tolerance,
                                                    max));

        printOutput("(c) Nelson-Siegel", ts3);

        // a cubic bspline curve with 11 knot points, implies
        // n=6 (constrained problem) basis functions

        Time knots[] =  { -30.0, -20.0,  0.0,  5.0, 10.0, 15.0,
                           20.0,  25.0, 30.0, 40.0, 50.0 };

        std::vector<Time> knotVector;
        for (Size i=0; i< LENGTH(knots); i++) {
            knotVector.push_back(knots[i]);
        }

        CubicBSplinesFitting cubicBSplines(knotVector, constrainAtZero);

        boost::shared_ptr<FittedBondDiscountCurve> ts4 (
                       new FittedBondDiscountCurve(curveSettlementDays,
                                                   calendar,
                                                   instrumentsA,
                                                   dc,
                                                   cubicBSplines,
                                                   tolerance,
                                                   max));

        printOutput("(d) cubic B-splines", ts4);

        SvenssonFitting svensson;

        boost::shared_ptr<FittedBondDiscountCurve> ts5 (
                        new FittedBondDiscountCurve(curveSettlementDays,
                                                    calendar,
                                                    instrumentsA,
                                                    dc,
                                                    svensson,
                                                    tolerance,
                                                    max));

        printOutput("(e) Svensson", ts5);

        Handle<YieldTermStructure> discountCurve(
            boost::make_shared<FlatForward>(
                                    curveSettlementDays, calendar, 0.01, dc));
        SpreadFittingMethod nelsonSiegelSpread(
                                    boost::make_shared<NelsonSiegelFitting>(),
                                    discountCurve);

        boost::shared_ptr<FittedBondDiscountCurve> ts6 (
                        new FittedBondDiscountCurve(curveSettlementDays,
                                                    calendar,
                                                    instrumentsA,
                                                    dc,
                                                    nelsonSiegelSpread,
                                                    tolerance,
                                                    max));

        printOutput("(f) Nelson-Siegel spreaded", ts6);

        cout << "Output par rates for each curve. In this case, "
             << endl
             << "par rates should equal coupons for these par bonds."
             << endl
             << endl;

        cout << setw(6) << "tenor" << " | "
             << setw(6) << "coupon" << " | "
             << setw(6) << "bstrap" << " | "
             << setw(6) << "(a)" << " | "
             << setw(6) << "(b)" << " | "
             << setw(6) << "(c)" << " | "
             << setw(6) << "(d)" << " | "
             << setw(6) << "(e)" << " | "
             << setw(6) << "(f)" << endl;

        for (Size i=0; i<instrumentsA.size(); i++) {

            std::vector<boost::shared_ptr<CashFlow> > cfs =
                instrumentsA[i]->bond()->cashflows();

            Size cfSize = instrumentsA[i]->bond()->cashflows().size();
            std::vector<Date> keyDates;
            keyDates.push_back(bondSettlementDate);

            for (Size j=0; j<cfSize-1; j++) {
                if (!cfs[j]->hasOccurred(bondSettlementDate, false)) {
                    Date myDate =  cfs[j]->date();
                    keyDates.push_back(myDate);
                }
            }

            Real tenor = dc.yearFraction(today, cfs[cfSize-1]->date());

            cout << setw(6) << fixed << setprecision(3) << tenor << " | "
                 << setw(6) << fixed << setprecision(3)
                 << 100.*coupons[i] << " | "
                 // piecewise bootstrap
                 << setw(6) << fixed << setprecision(3)
                 << 100.*parRate(*ts0,keyDates,dc) << " | "
                 // exponential splines
                 << setw(6) << fixed << setprecision(3)
                 << 100.*parRate(*ts1,keyDates,dc) << " | "
                 // simple polynomial
                 << setw(6) << fixed << setprecision(3)
                 << 100.*parRate(*ts2,keyDates,dc) << " | "
                 // Nelson-Siegel
                 << setw(6) << fixed << setprecision(3)
                 << 100.*parRate(*ts3,keyDates,dc) << " | "
                 // cubic bsplines
                 << setw(6) << fixed << setprecision(3)
                 << 100.*parRate(*ts4,keyDates,dc) << " | "
                 // Svensson
                 << setw(6) << fixed << setprecision(3)
                 << 100.*parRate(*ts5,keyDates,dc)  << " | "
                 // Nelson-Siegel Spreaded
                 << setw(6) << fixed << setprecision(3)
                 << 100.*parRate(*ts6,keyDates,dc) << endl;
        }

        cout << endl << endl << endl;
        cout << "Now add 23 months to today. Par rates should be "  << endl
             << "automatically recalculated because today's date "  << endl
             << "changes.  Par rates will NOT equal coupons (YTM "  << endl
             << "will, with the correct compounding), but the "     << endl
             << "piecewise yield curve par rates can be used as "   << endl
             << "a benchmark for correct par rates."
             << endl
             << endl;

        today = calendar.advance(origToday,23,Months,convention);
        Settings::instance().evaluationDate() = today;
        bondSettlementDate = calendar.advance(today, bondSettlementDays*Days);

        printOutput("(a) exponential splines", ts1);

        printOutput("(b) simple polynomial", ts2);

        printOutput("(c) Nelson-Siegel", ts3);

        printOutput("(d) cubic B-splines", ts4);

        printOutput("(e) Svensson", ts5);

        printOutput("(f) Nelson-Siegel spreaded", ts6);

        cout << endl
             << endl;

        cout << setw(6) << "tenor" << " | "
             << setw(6) << "coupon" << " | "
             << setw(6) << "bstrap" << " | "
             << setw(6) << "(a)" << " | "
             << setw(6) << "(b)" << " | "
             << setw(6) << "(c)" << " | "
             << setw(6) << "(d)" << " | "
             << setw(6) << "(e)" << " | "
             << setw(6) << "(f)" << endl;

        for (Size i=0; i<instrumentsA.size(); i++) {

            std::vector<boost::shared_ptr<CashFlow> > cfs =
                instrumentsA[i]->bond()->cashflows();

            Size cfSize = instrumentsA[i]->bond()->cashflows().size();
            std::vector<Date> keyDates;
            keyDates.push_back(bondSettlementDate);

            for (Size j=0; j<cfSize-1; j++) {
                if (!cfs[j]->hasOccurred(bondSettlementDate, false)) {
                    Date myDate =  cfs[j]->date();
                    keyDates.push_back(myDate);
                }
            }

            Real tenor = dc.yearFraction(today, cfs[cfSize-1]->date());

            cout << setw(6) << fixed << setprecision(3) << tenor << " | "
                 << setw(6) << fixed << setprecision(3)
                 << 100.*coupons[i] << " | "
                 // piecewise bootstrap
                 << setw(6) << fixed << setprecision(3)
                 << 100.*parRate(*ts0,keyDates,dc) << " | "
                 // exponential splines
                 << setw(6) << fixed << setprecision(3)
                 << 100.*parRate(*ts1,keyDates,dc) << " | "
                 // simple polynomial
                 << setw(6) << fixed << setprecision(3)
                 << 100.*parRate(*ts2,keyDates,dc) << " | "
                 // Nelson-Siegel
                 << setw(6) << fixed << setprecision(3)
                 << 100.*parRate(*ts3,keyDates,dc) << " | "
                 // cubic bsplines
                 << setw(6) << fixed << setprecision(3)
                 << 100.*parRate(*ts4,keyDates,dc) << " | "
                 // Svensson
                 << setw(6) << fixed << setprecision(3)
                 << 100.*parRate(*ts5,keyDates,dc) << " | "
                 // Nelson-Siegel Spreaded
                 << setw(6) << fixed << setprecision(3)
                 << 100.*parRate(*ts6,keyDates,dc) << endl;
        }

        cout << endl << endl << endl;
        cout << "Now add one more month, for a total of two years " << endl
             << "from the original date. The first instrument is "  << endl
             << "now expired and par rates should again equal "     << endl
             << "coupon values, since clean prices did not change."
             << endl
             << endl;

        instrumentsA.erase(instrumentsA.begin(),
                           instrumentsA.begin()+1);
        instrumentsB.erase(instrumentsB.begin(),
                           instrumentsB.begin()+1);

        today = calendar.advance(origToday,24,Months,convention);
        Settings::instance().evaluationDate() = today;
        bondSettlementDate = calendar.advance(today, bondSettlementDays*Days);

        boost::shared_ptr<YieldTermStructure> ts00 (
              new PiecewiseYieldCurve<Discount,LogLinear>(curveSettlementDays,
                                                          calendar,
                                                          instrumentsB,
                                                          dc));

        boost::shared_ptr<FittedBondDiscountCurve> ts11 (
                  new FittedBondDiscountCurve(curveSettlementDays,
                                              calendar,
                                              instrumentsA,
                                              dc,
                                              exponentialSplines,
                                              tolerance,
                                              max));

        printOutput("(a) exponential splines", ts11);

        boost::shared_ptr<FittedBondDiscountCurve> ts22 (
                    new FittedBondDiscountCurve(curveSettlementDays,
                                                calendar,
                                                instrumentsA,
                                                dc,
                                                simplePolynomial,
                                                tolerance,
                                                max));

        printOutput("(b) simple polynomial", ts22);

        boost::shared_ptr<FittedBondDiscountCurve> ts33 (
                        new FittedBondDiscountCurve(curveSettlementDays,
                                                    calendar,
                                                    instrumentsA,
                                                    dc,
                                                    nelsonSiegel,
                                                    tolerance,
                                                    max));

        printOutput("(c) Nelson-Siegel", ts33);

        boost::shared_ptr<FittedBondDiscountCurve> ts44 (
                       new FittedBondDiscountCurve(curveSettlementDays,
                                                   calendar,
                                                   instrumentsA,
                                                   dc,
                                                   cubicBSplines,
                                                   tolerance,
                                                   max));

        printOutput("(d) cubic B-splines", ts44);

        boost::shared_ptr<FittedBondDiscountCurve> ts55 (
                       new FittedBondDiscountCurve(curveSettlementDays,
                                                   calendar,
                                                   instrumentsA,
                                                   dc,
                                                   svensson,
                                                   tolerance,
                                                   max));

        printOutput("(e) Svensson", ts55);

        boost::shared_ptr<FittedBondDiscountCurve> ts66 (
                        new FittedBondDiscountCurve(curveSettlementDays,
                                                    calendar,
                                                    instrumentsA,
                                                    dc,
                                                    nelsonSiegelSpread,
                                                    tolerance,
                                                    max));

        printOutput("(f) Nelson-Siegel spreaded", ts66);

        cout << setw(6) << "tenor" << " | "
             << setw(6) << "coupon" << " | "
             << setw(6) << "bstrap" << " | "
             << setw(6) << "(a)" << " | "
             << setw(6) << "(b)" << " | "
             << setw(6) << "(c)" << " | "
             << setw(6) << "(d)" << " | "
             << setw(6) << "(e)" << " | "
             << setw(6) << "(f)" << endl;

        for (Size i=0; i<instrumentsA.size(); i++) {

            std::vector<boost::shared_ptr<CashFlow> > cfs =
                instrumentsA[i]->bond()->cashflows();

            Size cfSize = instrumentsA[i]->bond()->cashflows().size();
            std::vector<Date> keyDates;
            keyDates.push_back(bondSettlementDate);

            for (Size j=0; j<cfSize-1; j++) {
                if (!cfs[j]->hasOccurred(bondSettlementDate, false)) {
                    Date myDate =  cfs[j]->date();
                    keyDates.push_back(myDate);
                }
            }

            Real tenor = dc.yearFraction(today, cfs[cfSize-1]->date());

            cout << setw(6) << fixed << setprecision(3) << tenor << " | "
                 << setw(6) << fixed << setprecision(3)
                 << 100.*coupons[i+1] << " | "
                 // piecewise bootstrap
                 << setw(6) << fixed << setprecision(3)
                 << 100.*parRate(*ts00,keyDates,dc) << " | "
                 // exponential splines
                 << setw(6) << fixed << setprecision(3)
                 << 100.*parRate(*ts11,keyDates,dc) << " | "
                 // simple polynomial
                 << setw(6) << fixed << setprecision(3)
                 << 100.*parRate(*ts22,keyDates,dc) << " | "
                 // Nelson-Siegel
                 << setw(6) << fixed << setprecision(3)
                 << 100.*parRate(*ts33,keyDates,dc) << " | "
                 // cubic bsplines
                 << setw(6) << fixed << setprecision(3)
                 << 100.*parRate(*ts44,keyDates,dc) << " | "
                 // Svensson
                 << setw(6) << fixed << setprecision(3)
                 << 100.*parRate(*ts55,keyDates,dc) << " | "
                 // Nelson-Siegel Spreaded
                 << setw(6) << fixed << setprecision(3)
                 << 100.*parRate(*ts66,keyDates,dc) << endl;
        }

        cout << endl << endl << endl;
        cout << "Now decrease prices by a small amount, corresponding"  << endl
             << "to a theoretical five basis point parallel + shift of" << endl
             << "the yield curve. Because bond quotes change, the new " << endl
             << "par rates should be recalculated automatically."
             << endl
             << endl;

        for (Size k=0; k<LENGTH(lengths)-1; k++) {

            Real P = instrumentsA[k]->quote()->value();
            const Bond& b = *instrumentsA[k]->bond();
            Rate ytm = BondFunctions::yield(b, P,
                                            dc, Compounded, frequency,
                                            today);
            Time dur = BondFunctions::duration(b, ytm,
                                               dc, Compounded, frequency,
                                               Duration::Modified,
                                               today);

            const Real bpsChange = 5.;
            // dP = -dur * P * dY
            Real deltaP = -dur * P * (bpsChange/10000.);
            quote[k+1]->setValue(P + deltaP);
        }

        cout << setw(6) << "tenor" << " | "
             << setw(6) << "coupon" << " | "
             << setw(6) << "bstrap" << " | "
             << setw(6) << "(a)" << " | "
             << setw(6) << "(b)" << " | "
             << setw(6) << "(c)" << " | "
             << setw(6) << "(d)" << " | "
             << setw(6) << "(e)" << " | "
             << setw(6) << "(f)" << endl;

        for (Size i=0; i<instrumentsA.size(); i++) {

            std::vector<boost::shared_ptr<CashFlow> > cfs =
                instrumentsA[i]->bond()->cashflows();

            Size cfSize = instrumentsA[i]->bond()->cashflows().size();
            std::vector<Date> keyDates;
            keyDates.push_back(bondSettlementDate);

            for (Size j=0; j<cfSize-1; j++) {
                if (!cfs[j]->hasOccurred(bondSettlementDate, false)) {
                    Date myDate =  cfs[j]->date();
                    keyDates.push_back(myDate);
                }
            }

            Real tenor = dc.yearFraction(today, cfs[cfSize-1]->date());

            cout << setw(6) << fixed << setprecision(3) << tenor << " | "
                 << setw(6) << fixed << setprecision(3)
                 << 100.*coupons[i+1] << " | "
                 // piecewise bootstrap
                 << setw(6) << fixed << setprecision(3)
                 << 100.*parRate(*ts00,keyDates,dc) << " | "
                 // exponential splines
                 << setw(6) << fixed << setprecision(3)
                 << 100.*parRate(*ts11,keyDates,dc) << " | "
                 // simple polynomial
                 << setw(6) << fixed << setprecision(3)
                 << 100.*parRate(*ts22,keyDates,dc) << " | "
                 // Nelson-Siegel
                 << setw(6) << fixed << setprecision(3)
                 << 100.*parRate(*ts33,keyDates,dc) << " | "
                 // cubic bsplines
                 << setw(6) << fixed << setprecision(3)
                 << 100.*parRate(*ts44,keyDates,dc) << " | "
                 // Svensson
                 << setw(6) << fixed << setprecision(3)
                 << 100.*parRate(*ts55,keyDates,dc) << " | "
                 // Nelson-Siegel Spreaded
                 << setw(6) << fixed << setprecision(3)
                 << 100.*parRate(*ts66,keyDates,dc) << endl;
        }

        double seconds = timer.elapsed();
        Integer hours = int(seconds/3600);
        seconds -= hours * 3600;
        Integer minutes = int(seconds/60);
        seconds -= minutes * 60;
        std::cout << " \nRun completed in ";
        if (hours > 0)
            std::cout << hours << " h ";
        if (hours > 0 || minutes > 0)
            std::cout << minutes << " m ";
        std::cout << std::fixed << std::setprecision(0)
                  << seconds << " s\n" << std::endl;

        return 0;

    } catch (std::exception& e) {
        cerr << e.what() << endl;
        return 1;
    } catch (...) {
        cerr << "unknown error" << endl;
        return 1;
    }

}


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

#ifndef SMOOTH__TN_HPP_
#define SMOOTH__TN_HPP_

#include <Eigen/Core>

#include "internal/lie_group_base.hpp"
#include "internal/macro.hpp"
#include "internal/tn.hpp"

namespace smooth {

/**
 * @brief Lie group \f$\mathbb{T}(n)\f$ of \f$n\f$-dimensional translations.
 *
 * Memory layout
 * -------------
 * - Group:    \f$ \mathbf{t} = [x_1, \ldots, x_n] \f$
 * - Tangent:  \f$ \mathbf{v} = [v_1, \ldots, v_n] \f$
 *
 * Lie group Matrix form
 * ---------------------
 *
 * \f[
 * \mathbf{X} =
 * \begin{bmatrix}
 *  I & \mathbf{t} \\
 *  0 & 1
 * \end{bmatrix} \in \mathbb{R}^{n+1 \times n+1}
 * \f]
 *
 * Lie algebra Matrix form
 * -----------------------
 *
 * \f[
 * \mathbf{v}^\wedge =
 * \begin{bmatrix}
 *  0 & \mathbf{v} \\
 *  0 & 0
 * \end{bmatrix} \in \mathbb{R}^{n+1 \times n+1}
 * \f]
 */
template<typename _Derived>
class TnBase : public LieGroupBase<_Derived>
{
  using Base = LieGroupBase<_Derived>;

protected:
  TnBase() = default;

public:
  SMOOTH_INHERIT_TYPEDEFS;

  /**
   * @brief Euclidean vector (Rn) representation.
   */
  Eigen::Map<Eigen::Matrix<Scalar, Dof, 1>> rn() requires is_mutable
  {
    return Eigen::Map<Eigen::Matrix<Scalar, Dof, 1>>(static_cast<_Derived &>(*this).data());
  }

  /**
   * @brief Euclidean vector (Rn) representation.
   */
  Eigen::Map<const Eigen::Matrix<Scalar, Dof, 1>> rn() const
  {
    return Eigen::Map<const Eigen::Matrix<Scalar, Dof, 1>>(
      static_cast<const _Derived &>(*this).data());
  }

  /**
   * @brief Translation action on Rn vector.
   */
  template<typename EigenDerived>
  Eigen::Matrix<Scalar, Dof, 1> operator*(const Eigen::MatrixBase<EigenDerived> & v) const
  {
    return rn() + v;
  }
};

// \cond
template<int N, typename _Scalar>
requires(N > 0) class Tn;
// \endcond

// \cond
template<int N, typename _Scalar>
struct lie_traits<Tn<N, _Scalar>>
{
  static constexpr bool is_mutable = true;

  using Impl   = TnImpl<N, _Scalar>;
  using Scalar = _Scalar;

  template<typename NewScalar>
  using PlainObject = Tn<N, NewScalar>;
};
// \endcond

/**
 * @brief Storage implementation of Tn Lie group.
 *
 * @see TnBase for memory layout.
 */
template<int N, typename _Scalar>
// \cond
requires(N > 0)
  // \endcond
  class Tn : public TnBase<Tn<N, _Scalar>>
{
  using Base = TnBase<Tn<N, _Scalar>>;
  SMOOTH_GROUP_API(Tn);

public:
  /**
   * @brief Construct from Eigen vector.
   *
   * @param rn Eigen vector.
   */
  template<typename Derived>
  Tn(const Eigen::MatrixBase<Derived> & rn) : coeffs_(rn)
  {}
};

}  // namespace smooth

// \cond
template<int N, typename _Scalar>
struct smooth::lie_traits<Eigen::Map<smooth::Tn<N, _Scalar>>>
    : public lie_traits<smooth::Tn<N, _Scalar>>
{};
// \endcond

/**
 * @brief Memory mapping of Tn Lie group.
 *
 * @see TnBase for memory layout.
 */
template<int N, typename _Scalar>
class Eigen::Map<smooth::Tn<N, _Scalar>> : public smooth::TnBase<Eigen::Map<smooth::Tn<N, _Scalar>>>
{
  using Base = smooth::TnBase<Eigen::Map<smooth::Tn<N, _Scalar>>>;

  SMOOTH_MAP_API(Map);
};

// \cond
template<int N, typename _Scalar>
struct smooth::lie_traits<Eigen::Map<const smooth::Tn<N, _Scalar>>>
    : public lie_traits<smooth::Tn<N, _Scalar>>
{
  static constexpr bool is_mutable = false;
};
// \endcond

/**
 * @brief Const memory mapping of Tn Lie group.
 *
 * @see TnBase for memory layout.
 */
template<int N, typename _Scalar>
class Eigen::Map<const smooth::Tn<N, _Scalar>>
    : public smooth::TnBase<Eigen::Map<const smooth::Tn<N, _Scalar>>>
{
  using Base = smooth::TnBase<Eigen::Map<const smooth::Tn<N, _Scalar>>>;

  SMOOTH_CONST_MAP_API(Map);
};

namespace smooth {
// \cond
template<typename Scalar>
using T1 = Tn<1, Scalar>;
template<typename Scalar>
using T2 = Tn<2, Scalar>;
template<typename Scalar>
using T3 = Tn<3, Scalar>;
template<typename Scalar>
using T4 = Tn<4, Scalar>;
template<typename Scalar>
using T5 = Tn<5, Scalar>;
template<typename Scalar>
using T6 = Tn<6, Scalar>;
template<typename Scalar>
using T7 = Tn<7, Scalar>;
template<typename Scalar>
using T8 = Tn<8, Scalar>;
template<typename Scalar>
using T9 = Tn<9, Scalar>;
template<typename Scalar>
using T10 = Tn<10, Scalar>;

using T1f  = T1<float>;
using T2f  = T2<float>;
using T3f  = T3<float>;
using T4f  = T4<float>;
using T5f  = T5<float>;
using T6f  = T6<float>;
using T7f  = T7<float>;
using T8f  = T8<float>;
using T9f  = T9<float>;
using T10f = T10<float>;

using T1d  = T1<double>;
using T2d  = T2<double>;
using T3d  = T3<double>;
using T4d  = T4<double>;
using T5d  = T5<double>;
using T6d  = T6<double>;
using T7d  = T7<double>;
using T8d  = T8<double>;
using T9d  = T9<double>;
using T10d = T10<double>;
// \endcond
}  // namespace smooth

#endif  // SMOOTH__TN_HPP_


// Boost.Geometry

// Copyright (c) 2018 Adeel Ahmad, Islamabad, Pakistan.

// Contributed and/or modified by Adeel Ahmad,
//   as part of Google Summer of Code 2018 program.

// This file was modified by Oracle on 2018.
// Modifications copyright (c) 2018 Oracle and/or its affiliates.
// Contributed and/or modified by Adam Wulkiewicz, on behalf of Oracle

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

// This file is converted from GeographicLib, https://geographiclib.sourceforge.io
// GeographicLib is originally written by Charles Karney.

// Author: Charles Karney (2008-2017)

// Last updated version of GeographicLib: 1.49

// Original copyright notice:

// Copyright (c) Charles Karney (2008-2017) <charles@karney.com> and licensed
// under the MIT/X11 License. For more information, see
// https://geographiclib.sourceforge.io

#ifndef BOOST_GEOMETRY_FORMULAS_KARNEY_DIRECT_HPP
#define BOOST_GEOMETRY_FORMULAS_KARNEY_DIRECT_HPP

#include <boost/array.hpp>

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

#include <boost/geometry/formulas/flattening.hpp>
#include <boost/geometry/formulas/result_direct.hpp>

#include <boost/geometry/util/condition.hpp>
#include <boost/geometry/util/math.hpp>
#include <boost/geometry/util/normalize_spheroidal_coordinates.hpp>
#include <boost/geometry/util/series_expansion.hpp>

namespace boost { namespace geometry { namespace formula
{

namespace se = series_expansion;

/*!
\brief The solution of the direct problem of geodesics on latlong coordinates,
       after Karney (2011).
\author See
- Charles F.F Karney, Algorithms for geodesics, 2011
https://arxiv.org/pdf/1109.4448.pdf
*/
template <
    typename CT,
    bool EnableCoordinates = true,
    bool EnableReverseAzimuth = false,
    bool EnableReducedLength = false,
    bool EnableGeodesicScale = false,
    size_t SeriesOrder = 8
>
class karney_direct
{
    static const bool CalcQuantities = EnableReducedLength || EnableGeodesicScale;
    static const bool CalcCoordinates = EnableCoordinates || CalcQuantities;
    static const bool CalcRevAzimuth = EnableReverseAzimuth || CalcCoordinates || CalcQuantities;

public:
    typedef result_direct<CT> result_type;

    template <typename T, typename Dist, typename Azi, typename Spheroid>
    static inline result_type apply(T const& lo1,
                                    T const& la1,
                                    Dist const& distance,
                                    Azi const& azimuth12,
                                    Spheroid const& spheroid)
    {
        result_type result;

        CT lon1 = lo1;
        CT const lat1 = la1;

        Azi azi12 = azimuth12;
        math::normalize_azimuth<degree, Azi>(azi12);

        Dist const dist_c0 = 0;

        if (math::equals(distance, dist_c0) || distance < dist_c0)
        {
            result.lon2 = lon1;
            result.lat2 = lat1;
            return result;
        }

        CT const c0 = 0;
        CT const c1 = 1;
        CT const c2 = 2;

        CT const b = CT(get_radius<2>(spheroid));
        CT const f = formula::flattening<CT>(spheroid);
        CT const one_minus_f = c1 - f;
        CT const two_minus_f = c2 - f;

        CT const n = f / two_minus_f;
        CT const e2 = f * two_minus_f;
        CT const ep2 = e2 / math::sqr(one_minus_f);

        CT sin_alpha1, cos_alpha1;
        math::sin_cos_degrees<CT>(math::round_angle<CT>(azi12), sin_alpha1, cos_alpha1);

        // Find the reduced latitude.
        CT sin_beta1, cos_beta1;
        math::sin_cos_degrees<CT>(math::round_angle<CT>(lat1), sin_beta1, cos_beta1);
        sin_beta1 *= one_minus_f;

        math::normalize_unit_vector<CT>(sin_beta1, cos_beta1);

        cos_beta1 = (std::max)(c0, cos_beta1);

        // Obtain alpha 0 by solving the spherical triangle.
        CT const sin_alpha0 = sin_alpha1 * cos_beta1;
        CT const cos_alpha0 = boost::math::hypot(cos_alpha1, sin_alpha1 * sin_beta1);

        CT const k2 = math::sqr(cos_alpha0) * ep2;

        CT const epsilon = k2 / (c2 * (c1 + math::sqrt(c1 + k2)) + k2);

        // Find the coefficients for A1 by computing the
        // series expansion using Horner scehme.
        CT const expansion_A1 = se::evaluate_A1<SeriesOrder>(epsilon);

        // Index zero element of coeffs_C1 is unused.
        se::coeffs_C1<SeriesOrder, CT> const coeffs_C1(epsilon);

        // Tau is an integration variable.
        CT const tau12 = distance / (b * (c1 + expansion_A1));

        CT const sin_tau12 = sin(tau12);
        CT const cos_tau12 = cos(tau12);

        CT sin_sigma1 = sin_beta1;
        CT sin_omega1 = sin_alpha0 * sin_beta1;

        CT cos_sigma1, cos_omega1;
        cos_sigma1 = cos_omega1 = sin_beta1 != c0 || cos_alpha1 != c0 ? cos_beta1 * cos_alpha1 : c1;
        math::normalize_unit_vector<CT>(sin_sigma1, cos_sigma1);

        CT const B11 = se::sin_cos_series(sin_sigma1, cos_sigma1, coeffs_C1);
        CT const sin_B11 = sin(B11);
        CT const cos_B11 = cos(B11);

        CT const sin_tau1 = sin_sigma1 * cos_B11 + cos_sigma1 * sin_B11;
        CT const cos_tau1 = cos_sigma1 * cos_B11 - sin_sigma1 * sin_B11;

        // Index zero element of coeffs_C1p is unused.
        se::coeffs_C1p<SeriesOrder, CT> const coeffs_C1p(epsilon);

        CT const B12 = - se::sin_cos_series
                             (sin_tau1 * cos_tau12 + cos_tau1 * sin_tau12,
                              cos_tau1 * cos_tau12 - sin_tau1 * sin_tau12,
                              coeffs_C1p);

        CT const sigma12 = tau12 - (B12 - B11);
        CT const sin_sigma12 = sin(sigma12);
        CT const cos_sigma12 = cos(sigma12);

        CT const sin_sigma2 = sin_sigma1 * cos_sigma12 + cos_sigma1 * sin_sigma12;
        CT const cos_sigma2 = cos_sigma1 * cos_sigma12 - sin_sigma1 * sin_sigma12;

        if (BOOST_GEOMETRY_CONDITION(CalcRevAzimuth))
        {
            CT const sin_alpha2 = sin_alpha0;
            CT const cos_alpha2 = cos_alpha0 * cos_sigma2;

            result.reverse_azimuth = atan2(sin_alpha2, cos_alpha2);

            // Convert the angle to radians.
            result.reverse_azimuth /= math::d2r<CT>();
        }

        if (BOOST_GEOMETRY_CONDITION(CalcCoordinates))
        {
            // Find the latitude at the second point.
            CT const sin_beta2 = cos_alpha0 * sin_sigma2;
            CT const cos_beta2 = boost::math::hypot(sin_alpha0, cos_alpha0 * cos_sigma2);

            result.lat2 = atan2(sin_beta2, one_minus_f * cos_beta2);

            // Convert the coordinate to radians.
            result.lat2 /= math::d2r<CT>();

            // Find the longitude at the second point.
            CT const sin_omega2 = sin_alpha0 * sin_sigma2;
            CT const cos_omega2 = cos_sigma2;

            CT const omega12 = atan2(sin_omega2 * cos_omega1 - cos_omega2 * sin_omega1,
                                     cos_omega2 * cos_omega1 + sin_omega2 * sin_omega1);

            se::coeffs_A3<SeriesOrder, CT> const coeffs_A3(n);

            CT const A3 = math::horner_evaluate(epsilon, coeffs_A3.begin(), coeffs_A3.end());
            CT const A3c = -f * sin_alpha0 * A3;

            se::coeffs_C3<SeriesOrder, CT> const coeffs_C3(n, epsilon);

            CT const B31 = se::sin_cos_series(sin_sigma1, cos_sigma1, coeffs_C3);

            CT const lam12 = omega12 + A3c *
                             (sigma12 + (se::sin_cos_series
                                             (sin_sigma2,
                                              cos_sigma2,
                                              coeffs_C3) - B31));

            // Convert to radians to get the
            // longitudinal difference.
            CT lon12 = lam12 / math::d2r<CT>();

            // Add the longitude at first point to the longitudinal
            // difference and normalize the result.
            math::normalize_longitude<degree, CT>(lon1);
            math::normalize_longitude<degree, CT>(lon12);

            result.lon2 = lon1 + lon12;
        }

        if (BOOST_GEOMETRY_CONDITION(CalcQuantities))
        {
            // Evaluate the coefficients for C2.
            // Index zero element of coeffs_C2 is unused.
            se::coeffs_C2<SeriesOrder, CT> const coeffs_C2(epsilon);

            CT const B21 = se::sin_cos_series(sin_sigma1, cos_sigma1, coeffs_C2);
            CT const B22 = se::sin_cos_series(sin_sigma2, cos_sigma2, coeffs_C2);

            // Find the coefficients for A2 by computing the
            // series expansion using Horner scehme.
            CT const expansion_A2 = se::evaluate_A2<SeriesOrder>(epsilon);

            CT const AB1 = (c1 + expansion_A1) * (B12 - B11);
            CT const AB2 = (c1 + expansion_A2) * (B22 - B21);
            CT const J12 = (expansion_A1 - expansion_A2) * sigma12 + (AB1 - AB2);

            CT const dn1 = math::sqrt(c1 + ep2 * math::sqr(sin_beta1));
            CT const dn2 = math::sqrt(c1 + k2 * math::sqr(sin_sigma2));

            // Find the reduced length.
            result.reduced_length = b * ((dn2 * (cos_sigma1 * sin_sigma2) -
                                          dn1 * (sin_sigma1 * cos_sigma2)) -
                                          cos_sigma1 * cos_sigma2 * J12);

            // Find the geodesic scale.
            CT const t = k2 * (sin_sigma2 - sin_sigma1) *
                              (sin_sigma2 + sin_sigma1) / (dn1 + dn2);

            result.geodesic_scale = cos_sigma12 +
                                    (t * sin_sigma2 - cos_sigma2 * J12) *
                                    sin_sigma1 / dn1;
        }

        return result;
    }
};

}}} // namespace boost::geometry::formula

#endif // BOOST_GEOMETRY_FORMULAS_KARNEY_DIRECT_HPP


#include <density_estimation/utility/stochastic_process/gaussian_process.h>

#include <Eigen/Cholesky>
#include <cgv/math/eig.h>

#include <density_estimation/utility/functions.h>
#include <density_estimation/utility/truncated_multivariate_normal_distribution.h>

using namespace Eigen;

gaussian_process::gaussian_process(
  stochastic_process::random_engine_ptr random_engine,
  unsigned index_set_sample_size,
  std::function<float(float)> mean,
  std::function<float(float, float)> covariance)
  : stochastic_process(random_engine, index_set_sample_size)
  , mean{ mean }
  , covariance{ covariance }
  , standard_normal_distribution{ 0.0f, 1.0f }
{
  bool success = determine_covariance_square_root_matrix();
  if (!success) {
    throw std::invalid_argument(
      __func__ + std::string(": covariance function generates a covariance "
                             "matrix that isn't positive definite"));
  }

  determine_mean_vector();
}

gaussian_process::~gaussian_process() {}

void
gaussian_process::marginal_densities(float arg_min,
                                     float arg_max,
                                     unsigned arg_count,
                                     float random_variable_index,
                                     float* densities)
{
  assert(arg_min <= arg_max);
  assert(arg_count > 0);
  float mean = this->mean(random_variable_index);
  float variance =
    this->covariance(random_variable_index, random_variable_index);
  float arg_dx = arg_count > 1 ? (arg_max - arg_min) / (arg_count - 1) : 0;
  for (unsigned i = 0; i < arg_count; ++i) {
    float arg = arg_min + i * arg_dx;
    float diff = arg - mean;
    float density =
      expf(-diff * diff / (2.0f * variance)) / sqrtf(2.0f * M_PIf32 * variance);
    densities[i] = density;
  }
}

void
gaussian_process::marginal_densities_truncated(
  float arg_min,
  float arg_max,
  unsigned arg_count,
  float random_variable_index,
  const std::vector<range_restriction>& restrictions,
  float* densities)
{
  assert(arg_min <= arg_max);
  assert(arg_count > 0);

  cgv::math::vec<float> density_vector;
  density_vector.set_extern_data(arg_count, densities);

  unsigned size = static_cast<unsigned>(restrictions.size()) + 1;

  cgv::math::vec<float> mean{ size };
  cgv::math::mat<float> sigma{ size, size };
  cgv::math::vec<float> lower{ size };
  cgv::math::vec<float> upper{ size };

  for (unsigned i = 0; i < size; ++i) {
    float curr_random_variable_index = random_variable_index;
    lower(i) = -std::numeric_limits<float>::infinity();
    upper(i) = +std::numeric_limits<float>::infinity();
    if (i < size - 1) {
      const range_restriction& restriction = restrictions.at(i);
      curr_random_variable_index = restriction.random_variable_index;
      lower(i) = restriction.value_min;
      upper(i) = restriction.value_max;
    }

    mean(i) = this->mean(random_variable_index);

    for (unsigned j = 0; j < size; ++j) {
      float other_random_variable_index = random_variable_index;
      if (j < size - 1) {
        const range_restriction& restriction = restrictions.at(j);
        other_random_variable_index = restriction.random_variable_index;
      }

      sigma(i, j) = this->covariance(curr_random_variable_index,
                                     other_random_variable_index);
    }
  }
  marginal_truncated_multivariate_normal_density(arg_min,
                                                 arg_max,
                                                 arg_count,
                                                 size - 1,
                                                 mean,
                                                 sigma,
                                                 lower,
                                                 upper,
                                                 density_vector);
}

void
gaussian_process::generate_sample_function(float* sample_function_values)
{
  VectorXf standard_normal_vector{ this->index_set_sample_size };
  for (unsigned i = 0; i < this->index_set_sample_size; ++i) {
    float value = this->standard_normal_distribution(*this->random_engine);
    standard_normal_vector(i) = value;
  }

  Map<VectorXf> sample_function_vector{ sample_function_values,
                                        this->index_set_sample_size };
  sample_function_vector =
    this->mean_vector +
    this->covariance_square_root_matrix * standard_normal_vector;
}

void
gaussian_process::determine_mean_vector()
{
  this->mean_vector = Eigen::VectorXf(this->index_set_sample_size);
  for (unsigned i = 0; i < this->index_set_sample_size; ++i) {
    float random_variable_index = i / (this->index_set_sample_size - 1.0f);
    float mean = this->mean(random_variable_index);
    this->mean_vector(i) = mean;
  }
}

bool
gaussian_process::determine_covariance_square_root_matrix()
{
  MatrixXf covariance_matrix{ this->index_set_sample_size,
                              this->index_set_sample_size };
  for (unsigned i = 0; i < this->index_set_sample_size; ++i) {
    for (unsigned j = 0; j < this->index_set_sample_size; ++j) {
      float random_variable_index_1 = i / (this->index_set_sample_size - 1.0f);
      float random_variable_index_2 = j / (this->index_set_sample_size - 1.0f);
      float covariance =
        this->covariance(random_variable_index_1, random_variable_index_2);
      if (i == j) {
        covariance += 1e-04f;
      }
      covariance_matrix(i, j) = covariance;
    }
  }

  LLT<MatrixXf> llt_solver{ covariance_matrix };
  this->covariance_square_root_matrix = llt_solver.matrixL();
  if (llt_solver.info() != Eigen::ComputationInfo::Success) {
    return false;
  }
  return true;
}


#include "mesh.hpp"
#include <Eigen/Core>
#include <Eigen/Geometry>
#include <Eigen/LU>
#include <elasty/alembic-manager.hpp>
#include <elasty/fem.hpp>
#include <mathtoolbox/l-bfgs.hpp>
#include <timer.hpp>
#include <vector>

namespace
{
    constexpr double k_youngs_modulus = 600.0;
    constexpr double k_poisson_ratio  = 0.45;

    constexpr double k_first_lame  = elasty::fem::calcFirstLame(k_youngs_modulus, k_poisson_ratio);
    constexpr double k_second_lame = elasty::fem::calcSecondLame(k_youngs_modulus, k_poisson_ratio);

    constexpr unsigned k_num_substeps = 2;
    constexpr double   k_delta_time   = 1.0 / 60.0;

    constexpr double k_damping_factor = 0.0;

    constexpr double k_spring_stiffness    = 1000.0;
    constexpr double k_collision_stiffness = 100.0;

    enum class Model
    {
        CoRotational,
        StVenantKirchhoff
    };

    constexpr Model k_model = Model::CoRotational;
} // namespace

struct Constraint
{
    std::size_t                              vert_index;
    std::function<Eigen::Vector3d(double t)> motion;
    double                                   stiffness;
};

struct Collision
{
    std::size_t     vert_index;
    Eigen::Vector3d surface_pos;
    Eigen::Vector3d surface_normal;
};

class HalfSpaceCollider
{
public:
    HalfSpaceCollider(const Eigen::Vector3d& representative_point, const Eigen::Vector3d& normal)
        : m_representative_point(representative_point), m_normal(normal.normalized())
    {
        assert(!normal.isZero());
    }

    bool testCollision(const Eigen::Vector3d& point, const double tol) const
    {
        return (point - m_representative_point).dot(m_normal) < tol;
    }

    void retrieveCollisionInfo(const Eigen::Vector3d& point,
                               Eigen::Vector3d&       surface_pos,
                               Eigen::Vector3d&       surface_normal) const
    {
        surface_pos    = point - (point - m_representative_point).dot(m_normal) * m_normal;
        surface_normal = m_normal;
    }

private:
    const Eigen::Vector3d m_representative_point;
    const Eigen::Vector3d m_normal;
};

template <typename Derived> typename Derived::Scalar calcEnergyDensity(const Eigen::MatrixBase<Derived>& deform_grad)
{
    switch (k_model)
    {
        case Model::StVenantKirchhoff:
            return elasty::fem::calcStVenantKirchhoffEnergyDensity(deform_grad, k_first_lame, k_second_lame);
        case Model::CoRotational:
            return elasty::fem::calcCoRotationalEnergyDensity(deform_grad, k_first_lame, k_second_lame);
    }
}

template <typename Derived>
Eigen::Matrix<typename Derived::Scalar, 3, 3> calcPiolaStress(const Eigen::MatrixBase<Derived>& deform_grad)
{
    switch (k_model)
    {
        case Model::StVenantKirchhoff:
            return elasty::fem::calcStVenantKirchhoffPiolaStress(deform_grad, k_first_lame, k_second_lame);
        case Model::CoRotational:
            return elasty::fem::calcCoRotationalPiolaStress(deform_grad, k_first_lame, k_second_lame);
    }
}

class VariationalImplicit3dEngine
{
public:
    VariationalImplicit3dEngine() {}

    void proceedFrame()
    {
        const std::size_t num_verts = m_mesh.x_rest.size() / 3;
        const std::size_t num_elems = m_mesh.elems.cols();

        // Reset forces
        m_mesh.f = Eigen::VectorXd::Zero(3 * num_verts);

        // Apply gravity force
        for (std::size_t i = 0; i < num_verts; ++i)
        {
            m_mesh.f[i * 3 + 1] += m_mesh.lumped_mass(i * 3 + 1) * (-9.80665);
        }

        // Calculate the inverse lumped mass matrix
        const auto W = m_mesh.lumped_mass.cwiseInverse().asDiagonal();

        // Calculate the "inertia" position
        const double&         h = m_delta_physics_time;
        const Eigen::VectorXd y = m_mesh.x + h * m_mesh.v + h * h * W * m_mesh.f;

        // Define colliders
        // TODO: Improve the scene setup
        const HalfSpaceCollider floor{Eigen::Vector3d{0.0, -1.0, 0.0}, Eigen::Vector3d{0.0, 1.0, 0.0}};

        // Detect collisions
        // TODO: Improve the collision detection algorithtm
        std::vector<Collision> collisions;
        for (std::size_t vert_index = 0; vert_index < num_verts; ++vert_index)
        {
            const Eigen::Vector3d predicted_pos = y.segment<3>(vert_index * 3);

            if (floor.testCollision(predicted_pos, 0.0))
            {
                Eigen::Vector3d surface_pos;
                Eigen::Vector3d surface_normal;
                floor.retrieveCollisionInfo(predicted_pos, surface_pos, surface_normal);

                const Collision collision{vert_index, surface_pos, surface_normal};
                collisions.push_back(collision);
            }
        }

        const auto calcInternalPotential = [&](const Eigen::VectorXd& x)
        {
            double sum = 0.0;

            // Elastic potential
            for (std::size_t elem_index = 0; elem_index < num_elems; ++elem_index)
            {
                const auto& indices = m_mesh.elems.col(elem_index);

                // Retrieve precomputed values
                const auto& D_m_inv = m_mesh.rest_shape_mat_inv_array[elem_index];
                const auto& vol     = m_mesh.volume_array[elem_index];

                // Calculate the deformation gradient $\mathbf{F}$
                const auto F = elasty::fem::calcTetrahedronDeformGrad(x.segment<3>(3 * indices[0]),
                                                                      x.segment<3>(3 * indices[1]),
                                                                      x.segment<3>(3 * indices[2]),
                                                                      x.segment<3>(3 * indices[3]),
                                                                      D_m_inv);

                sum += vol * calcEnergyDensity(F);
            }

            // Attach-spring potential
            for (const auto& constraint : m_constraints)
            {
                const std::size_t vert_index   = constraint.vert_index;
                const double&     k            = constraint.stiffness;
                const auto        p            = x.segment<3>(vert_index * 3);
                const auto        q            = constraint.motion(m_physics_time + m_delta_physics_time);
                const double      squared_dist = (p - q).squaredNorm();

                sum += 0.5 * k * squared_dist;
            }

            // Collision energy
            for (const auto& collision : collisions)
            {
                const auto&  vert_index = collision.vert_index;
                const auto   pos        = x.segment<3>(vert_index * 3);
                const auto&  normal     = collision.surface_normal;
                const double a          = (collision.surface_pos - pos).transpose() * normal;

                if (a > 0.0)
                {
                    sum += k_collision_stiffness * a * a;
                }
            }

            return sum;
        };

        const auto calcInternalPotentialGrad = [&](const Eigen::VectorXd& x) -> Eigen::VectorXd
        {
            Eigen::VectorXd sum = Eigen::VectorXd::Zero(x.size());
            for (std::size_t elem_index = 0; elem_index < num_elems; ++elem_index)
            {
                const auto& indices = m_mesh.elems.col(elem_index);

                // Retrieve precomputed values
                const auto& D_m_inv  = m_mesh.rest_shape_mat_inv_array[elem_index];
                const auto& vol      = m_mesh.volume_array[elem_index];
                const auto& vec_PFPx = m_mesh.vec_PFPx_array[elem_index];

                // Calculate the deformation gradient $\mathbf{F}$
                const auto F = elasty::fem::calcTetrahedronDeformGrad(x.segment<3>(3 * indices[0]),
                                                                      x.segment<3>(3 * indices[1]),
                                                                      x.segment<3>(3 * indices[2]),
                                                                      x.segment<3>(3 * indices[3]),
                                                                      D_m_inv);

                // Calculate $\frac{\partial \Phi}{\partial \mathbf{x}}$ and related values
                const auto P      = calcPiolaStress(F);
                const auto vec_P  = Eigen::Map<const Eigen::Matrix<double, 9, 1>>(P.data(), P.size());
                const auto PPsiPx = vec_PFPx.transpose() * vec_P;

                // Calculate $\frac{\partial E}{\partial \mathbf{x}}$
                const auto PEPx = vol * PPsiPx;

                sum.segment<3>(3 * indices[0]) += PEPx.segment<3>(0 * 3);
                sum.segment<3>(3 * indices[1]) += PEPx.segment<3>(1 * 3);
                sum.segment<3>(3 * indices[2]) += PEPx.segment<3>(2 * 3);
                sum.segment<3>(3 * indices[3]) += PEPx.segment<3>(3 * 3);
            }

            for (const auto& constraint : m_constraints)
            {
                const std::size_t vert_index = constraint.vert_index;
                const double&     k          = constraint.stiffness;
                const auto        p          = x.segment<3>(vert_index * 3);
                const auto        q          = constraint.motion(m_physics_time + m_delta_physics_time);
                const auto        r          = p - q;

                sum.segment<3>(3 * vert_index) += k * r;
            }

            // Collision energy
            for (const auto& collision : collisions)
            {
                const auto&  vert_index = collision.vert_index;
                const auto   pos        = x.segment<3>(vert_index * 3);
                const auto&  normal     = collision.surface_normal;
                const double a          = (collision.surface_pos - pos).transpose() * normal;

                if (a > 0.0)
                {
                    sum.segment<3>(3 * vert_index) += 2.0 * k_collision_stiffness * a * (-normal);
                }
            }

            return sum;
        };

        const auto calcMomentumPotential = [&](const Eigen::VectorXd& x)
        {
            return (0.5 / (h * h)) * (x - y).transpose() * m_mesh.lumped_mass.asDiagonal() * (x - y);
        };

        const auto calcMomentumPotentialGrad = [&](const Eigen::VectorXd& x) -> Eigen::VectorXd
        {
            return (1.0 / (h * h)) * m_mesh.lumped_mass.asDiagonal() * (x - y);
        };

        const auto calcObjective = [&](const Eigen::VectorXd& x)
        {
            const double momentum_potential = calcMomentumPotential(x);
            const double internal_potential = calcInternalPotential(x);

            assert(momentum_potential >= 0.0);
            assert(internal_potential >= 0.0);

            return momentum_potential + internal_potential;
        };

        const auto calcObjectiveGrad = [&](const Eigen::VectorXd& x) -> Eigen::VectorXd
        {
            return calcMomentumPotentialGrad(x) + calcInternalPotentialGrad(x);
        };

        // Solve the minimization problem
        unsigned        num_iters;
        Eigen::VectorXd x_opt;
        mathtoolbox::optimization::RunLBfgs(y, calcObjective, calcObjectiveGrad, 1e-06, 50, x_opt, num_iters);

        // Update the internal state
        m_mesh.v = (1.0 / h) * (x_opt - m_mesh.x);
        m_mesh.x = x_opt;

        // Apply velocity correction: a naive damping
        m_mesh.v *= std::exp(-k_damping_factor * m_delta_physics_time);

        // Apply velocity correction: a naive friction
        for (const auto& collision : collisions)
        {
            const auto v_i = m_mesh.v.segment<3>(collision.vert_index * 3);

            m_mesh.v.segment<3>(collision.vert_index * 3) =
                v_i.dot(collision.surface_normal) * collision.surface_normal;
        }

        // Update time counter
        m_physics_time += m_delta_physics_time;
    }

    void initializeScene()
    {
#ifdef USE_EXTERNAL_TETRA_MESH
        // Load a tetrahedral mesh
        m_mesh = ReadTetraMesh("./assets/torus.mesh");

        const std::size_t num_verts = m_mesh.x_rest.size() / 3;
        const std::size_t num_elems = m_mesh.elems.cols();

        // Set initial conditions
        constexpr double k_pi = 3.141592653589793238;

        const auto offset = Eigen::Vector3d{0.0, 1.5, 0.0};
        const auto rotate = Eigen::AngleAxisd(0.5 * k_pi, Eigen::Vector3d::UnitX());

        for (std::size_t vert_index = 0; vert_index < num_verts; ++vert_index)
        {
            m_mesh.x.segment<3>(vert_index * 3) = offset + rotate * m_mesh.x.segment<3>(vert_index * 3);
        }
#else
        constexpr std::size_t num_blocks_x = 16;
        constexpr std::size_t num_blocks_y = 4;
        constexpr std::size_t num_blocks_z = 4;
        constexpr std::size_t num_verts    = (num_blocks_x + 1) * (num_blocks_y + 1) * (num_blocks_z + 1);
        constexpr std::size_t num_elems    = num_blocks_x * num_blocks_y * num_blocks_z * 5;
        constexpr double      scale        = 0.5 / num_blocks_y;

        m_mesh.elems.resize(4, num_elems);
        m_mesh.x_rest.resize(num_verts * 3);

        for (std::size_t i_z = 0; i_z < num_blocks_z; ++i_z)
        {
            for (std::size_t i_y = 0; i_y < num_blocks_y; ++i_y)
            {
                for (std::size_t i_x = 0; i_x < num_blocks_x; ++i_x)
                {
                    const std::size_t block_index = i_x + num_blocks_x * i_y + num_blocks_x * num_blocks_y * i_z;
                    const std::size_t vert_base_index =
                        i_x + (num_blocks_x + 1) * i_y + (num_blocks_x + 1) * (num_blocks_y + 1) * i_z;
                    const std::size_t     elem_base_index = 5 * block_index;
                    const Eigen::Vector3d vert_base_pos{i_x, i_y, i_z};

                    const std::size_t indices[] = {vert_base_index,
                                                   vert_base_index + 1,
                                                   vert_base_index + (num_blocks_x + 1),
                                                   vert_base_index + (num_blocks_x + 1) + 1,
                                                   vert_base_index + (num_blocks_x + 1) * (num_blocks_y + 1),
                                                   vert_base_index + (num_blocks_x + 1) * (num_blocks_y + 1) + 1,
                                                   vert_base_index + (num_blocks_x + 1) * (num_blocks_y + 2),
                                                   vert_base_index + (num_blocks_x + 1) * (num_blocks_y + 2) + 1};

                    m_mesh.elems.col(elem_base_index + 0) << indices[0], indices[4], indices[1], indices[2];
                    m_mesh.elems.col(elem_base_index + 1) << indices[2], indices[7], indices[1], indices[3];
                    m_mesh.elems.col(elem_base_index + 2) << indices[6], indices[2], indices[7], indices[4];
                    m_mesh.elems.col(elem_base_index + 3) << indices[5], indices[1], indices[4], indices[7];
                    m_mesh.elems.col(elem_base_index + 4) << indices[1], indices[2], indices[4], indices[7];

                    m_mesh.x_rest.segment<3>(indices[0] * 3) << 0.0, 0.0, 0.0;
                    m_mesh.x_rest.segment<3>(indices[1] * 3) << 1.0, 0.0, 0.0;
                    m_mesh.x_rest.segment<3>(indices[2] * 3) << 0.0, 1.0, 0.0;
                    m_mesh.x_rest.segment<3>(indices[3] * 3) << 1.0, 1.0, 0.0;
                    m_mesh.x_rest.segment<3>(indices[4] * 3) << 0.0, 0.0, 1.0;
                    m_mesh.x_rest.segment<3>(indices[5] * 3) << 1.0, 0.0, 1.0;
                    m_mesh.x_rest.segment<3>(indices[6] * 3) << 0.0, 1.0, 1.0;
                    m_mesh.x_rest.segment<3>(indices[7] * 3) << 1.0, 1.0, 1.0;

                    for (std::size_t i = 0; i < 8; ++i)
                    {
                        m_mesh.x_rest.segment<3>(indices[i] * 3) += vert_base_pos;
                    }
                }
            }
        }

        // Set transform
        for (std::size_t vert = 0; vert < num_verts; ++vert)
        {
            m_mesh.x_rest[3 * vert + 1] -= 0.5 * num_blocks_y;
            m_mesh.x_rest[3 * vert + 2] -= 0.5 * num_blocks_z;
        }
        m_mesh.x_rest *= scale;

        // Initialize other values
        m_mesh.x = m_mesh.x_rest;
        m_mesh.v = Eigen::VectorXd::Zero(3 * num_verts);
        m_mesh.f = Eigen::VectorXd::Zero(3 * num_verts);

        // Set constraints
        for (std::size_t i_z = 0; i_z < num_blocks_z + 1; ++i_z)
        {
            for (std::size_t i_y = 0; i_y < num_blocks_y + 1; ++i_y)
            {
                const std::size_t i_x = 0;

                const std::size_t vert_index =
                    i_x + (num_blocks_x + 1) * i_y + (num_blocks_x + 1) * (num_blocks_y + 1) * i_z;

                const auto motion = [&, vert_index](double) -> Eigen::Vector3d
                {
                    return m_mesh.x_rest.segment<3>(vert_index * 3);
                };

                m_constraints.push_back(Constraint{vert_index, motion, k_spring_stiffness});
            }
        }
        for (std::size_t i_z = 0; i_z < num_blocks_z + 1; ++i_z)
        {
            for (std::size_t i_y = 0; i_y < num_blocks_y + 1; ++i_y)
            {
                const std::size_t i_x = num_blocks_x;

                const std::size_t vert_index =
                    i_x + (num_blocks_x + 1) * i_y + (num_blocks_x + 1) * (num_blocks_y + 1) * i_z;

                const auto motion = [&, vert_index](double t) -> Eigen::Vector3d
                {
                    constexpr double pi = 3.1415926535897932;

                    const auto ease = [](double x)
                    {
                        return x < 0.5 ? 4.0 * x * x * x : 1.0 - 0.5 * std::pow(-2.0 * x + 2.0, 3.0);
                    };

                    const auto   x_init    = m_mesh.x_rest.segment<3>(vert_index * 3);
                    const auto   axis      = Eigen::Vector3d::UnitX();
                    const double t_0       = 0.5;
                    const double t_1       = t_0 + 2.0;
                    const double a         = (t < t_0) ? 0.0 : ((t < t_1) ? t - t_0 : t_1 - t_0);
                    const double b         = ease(a / 2.0);
                    const double total_rot = 1.5 * pi;

                    return Eigen::AngleAxisd(total_rot * b, axis) * x_init;
                };

                m_constraints.push_back(Constraint{vert_index, motion, k_spring_stiffness});
            }
        }
#endif

        // Perform precomputation
        m_mesh.volume_array.resize(num_elems);
        m_mesh.rest_shape_mat_inv_array.resize(num_elems);
        m_mesh.vec_PFPx_array.resize(num_elems);
        for (std::size_t elem_index = 0; elem_index < num_elems; ++elem_index)
        {
            const auto& indices = m_mesh.elems.col(elem_index);

            const auto& x_0 = m_mesh.x_rest.segment<3>(3 * indices[0]);
            const auto& x_1 = m_mesh.x_rest.segment<3>(3 * indices[1]);
            const auto& x_2 = m_mesh.x_rest.segment<3>(3 * indices[2]);
            const auto& x_3 = m_mesh.x_rest.segment<3>(3 * indices[3]);

            m_mesh.volume_array[elem_index]             = elasty::fem::calcTetrahedronVolume(x_0, x_1, x_2, x_3);
            m_mesh.rest_shape_mat_inv_array[elem_index] = elasty::fem::calc3dShapeMatrix(x_0, x_1, x_2, x_3).inverse();
            m_mesh.vec_PFPx_array[elem_index] =
                elasty::fem::calcVecTetrahedronPartDeformGradPartPos(m_mesh.rest_shape_mat_inv_array[elem_index]);

            assert(!m_mesh.rest_shape_mat_inv_array[elem_index].hasNaN());
        }
        m_mesh.lumped_mass = elasty::fem::calcTetraMeshLumpedMass(m_mesh.x_rest, m_mesh.elems, m_mesh.mass);
    }

    /// \brief Getter of the delta physics time.
    ///
    /// \details The value equals to the delta frame time devided by the number of substeps.
    double getDeltaPhysicsTime() const { return m_delta_physics_time; }

    void setDeltaPhysicsTime(const double delta_physics_time) { m_delta_physics_time = delta_physics_time; }

    const TetraMesh* getMesh() const { return &m_mesh; }

private:
    double m_delta_physics_time = 1.0 / 60.0;
    double m_physics_time       = 0.0;

    std::vector<Constraint> m_constraints;

    TetraMesh m_mesh;
};

int main(int argc, char** argv)
{
    VariationalImplicit3dEngine engine;

    engine.initializeScene();
    engine.setDeltaPhysicsTime(k_delta_time / static_cast<double>(k_num_substeps));

    const auto        mesh      = engine.getMesh();
    const std::size_t num_verts = mesh->x_rest.size() / 3;
    const std::size_t num_elems = mesh->elems.cols();

    auto alembic_manager = elasty::createTetraMeshAlembicManager(
        "./out.abc", k_delta_time, num_verts, num_elems, mesh->x.data(), mesh->elems.data());

    for (unsigned int frame = 0; frame < 240; ++frame)
    {
        timer::Timer t(std::to_string(frame));

        alembic_manager->submitCurrentStatus();

        for (int i = 0; i < k_num_substeps; ++i)
        {
            engine.proceedFrame();
        }
    }

    return 0;
}