/* Author: Wolfgang Bangerth, University of Heidelberg, 2000 */

/*    $Id: step-8.cc 27657 2012-11-21 13:19:08Z bangerth $       */
/*                                                                */
/*    Copyright (C) 2000-2004, 2006-2008, 2010-2012 by the deal.II authors */
/*                                                                */
/*    This file is subject to QPL and may not be  distributed     */
/*    without copyright and license information. Please refer     */
/*    to the file deal.II/doc/license.html for the  text  and     */
/*    further information on this license.                        */

// @sect3{Include files}

// As usual, the first few include files are already known, so we will not
// comment on them further.
#include <deal.II/base/quadrature_lib.h>
#include <deal.II/base/function.h>
#include <deal.II/base/logstream.h>
#include <deal.II/lac/vector.h>
#include <deal.II/lac/full_matrix.h>
#include <deal.II/lac/sparse_matrix.h>
#include <deal.II/lac/solver_cg.h>
#include <deal.II/lac/precondition.h>
#include <deal.II/lac/constraint_matrix.h>
#include <deal.II/grid/tria.h>
#include <deal.II/grid/grid_generator.h>
#include <deal.II/grid/grid_refinement.h>
#include <deal.II/grid/tria_accessor.h>
#include <deal.II/grid/tria_iterator.h>
#include <deal.II/grid/tria_boundary_lib.h>
#include <deal.II/dofs/dof_handler.h>
#include <deal.II/dofs/dof_accessor.h>
#include <deal.II/dofs/dof_tools.h>
#include <deal.II/fe/fe_values.h>
#include <deal.II/numerics/vector_tools.h>
#include <deal.II/numerics/matrix_tools.h>
#include <deal.II/numerics/data_out.h>
#include <deal.II/numerics/error_estimator.h>

// In this example, we need vector-valued finite elements. The support for
// these can be found in the following include file:
#include <deal.II/fe/fe_system.h>
// We will compose the vector-valued finite elements from regular Q1 elements
// which can be found here, as usual:
#include <deal.II/fe/fe_q.h>

// This again is C++:
#include <fstream>
#include <iostream>

// The last step is as in previous programs. In particular, just like in
// step-7, we pack everything that's specific to this program into a namespace
// of its own.
namespace Step8
{
  using namespace dealii;

  // @sect3{The ElasticProblem class template}

  // The main class is, except for its name, almost unchanged with respect to
  // the step-6 example.
  //
  // The only change is the use of a different class for the fe
  // variable: Instead of a concrete finite element class such as
  // FE_Q, we now use a more generic one,
  // FESystem. In fact, FESystem is not really a
  // finite element itself in that it does not implement shape functions of
  // its own.  Rather, it is a class that can be used to stack several other
  // elements together to form one vector-valued finite element. In our case,
  // we will compose the vector-valued element of FE_Q(1)
  // objects, as shown below in the constructor of this class.
  template <int dim>
  class ElasticProblem
  {
  public:
    ElasticProblem ();
    ~ElasticProblem ();
    void run ();

  private:
    void setup_system ();
    void assemble_system ();
    void solve ();
    void refine_grid ();
    void output_results (const unsigned int cycle) const;

    Triangulation<dim>   triangulation;
    DoFHandler<dim>      dof_handler;

    FESystem<dim>        fe;

    ConstraintMatrix     hanging_node_constraints;

    SparsityPattern      sparsity_pattern;
    SparseMatrix<double> system_matrix;

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

  // @sect3{Right hand side values}

  // Before going over to the implementation of the main class, we declare and
  // define the class which describes the right hand side. This time, the
  // right hand side is vector-valued, as is the solution, so we will describe
  // the changes required for this in some more detail.
  //
  // The first thing is that vector-valued functions have to have a
  // constructor, since they need to pass down to the base class of how many
  // components the function consists. The default value in the constructor of
  // the base class is one (i.e.: a scalar function), which is why we did not
  // need not define a constructor for the scalar function used in previous
  // programs.
  template <int dim>
  class RightHandSide :  public Function<dim>
  {
  public:
    RightHandSide ();

    // The next change is that we want a replacement for the
    // value function of the previous examples. There, a second
    // parameter component was given, which denoted which
    // component was requested. Here, we implement a function that returns the
    // whole vector of values at the given place at once, in the second
    // argument of the function. The obvious name for such a replacement
    // function is vector_value.
    //
    // Secondly, in analogy to the value_list function, there is
    // a function vector_value_list, which returns the values of
    // the vector-valued function at several points at once:
    virtual void vector_value (const Point<dim> &p,
                               Vector<double>   &values) const;

    virtual void vector_value_list (const std::vector<Point<dim> > &points,
                                    std::vector<Vector<double> >   &value_list) const;
  };

  // This is the constructor of the right hand side class. As said above, it
  // only passes down to the base class the number of components, which is
  // dim in the present case (one force component in each of the
  // dim space directions).
  //
  // Some people would have moved the definition of such a short function
  // right into the class declaration. We do not do that, as a matter of
  // style: the deal.II style guides require that class declarations contain
  // only declarations, and that definitions are always to be found
  // outside. This is, obviously, as much as matter of taste as indentation,
  // but we try to be consistent in this direction.
  template <int dim>
  RightHandSide<dim>::RightHandSide ()
    :
    Function<dim> (dim)
  {}

  // Next the function that returns the whole vector of values at the point
  // p at once.
  //
  // To prevent cases where the return vector has not previously been set to
  // the right size we test for this case and otherwise throw an exception at
  // the beginning of the function. Note that enforcing that output arguments
  // already have the correct size is a convention in deal.II, and enforced
  // almost everywhere. The reason is that we would otherwise have to check at
  // the beginning of the function and possibly change the size of the output
  // vector. This is expensive, and would almost always be unnecessary (the
  // first call to the function would set the vector to the right size, and
  // subsequent calls would only have to do redundant checks). In addition,
  // checking and possibly resizing the vector is an operation that can not be
  // removed if we can't rely on the assumption that the vector already has
  // the correct size; this is in contract to the Assert call
  // that is completely removed if the program is compiled in optimized mode.
  //
  // Likewise, if by some accident someone tried to compile and run the
  // program in only one space dimension (in which the elastic equations do
  // not make much sense since they reduce to the ordinary Laplace equation),
  // we terminate the program in the second assertion. The program will work
  // just fine in 3d, however.
  template <int dim>
  inline
  void RightHandSide<dim>::vector_value (const Point<dim> &p,
                                         Vector<double>   &values) const
  {
    Assert (values.size() == dim,
            ExcDimensionMismatch (values.size(), dim));
    Assert (dim >= 2, ExcNotImplemented());

    // The rest of the function implements computing force values. We will use
    // a constant (unit) force in x-direction located in two little circles
    // (or spheres, in 3d) around points (0.5,0) and (-0.5,0), and y-force in
    // an area around the origin; in 3d, the z-component of these centers is
    // zero as well.
    //
    // For this, let us first define two objects that denote the centers of
    // these areas. Note that upon construction of the Point
    // objects, all components are set to zero.
    Point<dim> point_1, point_2;
    point_1(0) = 0.5;
    point_2(0) = -0.5;

    // If now the point p is in a circle (sphere) of radius 0.2
    // around one of these points, then set the force in x-direction to one,
    // otherwise to zero:
    if (((p-point_1).square() < 0.2*0.2) ||
        ((p-point_2).square() < 0.2*0.2))
      values(0) = 1;
    else
      values(0) = 0;

    // Likewise, if p is in the vicinity of the origin, then set
    // the y-force to 1, otherwise to zero:
    if (p.square() < 0.2*0.2)
      values(1) = 1;
    else
      values(1) = 0;
  }

  // Now, this is the function of the right hand side class that returns the
  // values at several points at once. The function starts out with checking
  // that the number of input and output arguments is equal (the sizes of the
  // individual output vectors will be checked in the function that we call
  // further down below). Next, we define an abbreviation for the number of
  // points which we shall work on, to make some things simpler below.
  template <int dim>
  void RightHandSide<dim>::vector_value_list (const std::vector<Point<dim> > &points,
                                              std::vector<Vector<double> >   &value_list) const
  {
    Assert (value_list.size() == points.size(),
            ExcDimensionMismatch (value_list.size(), points.size()));

    const unsigned int n_points = points.size();

    // Finally we treat each of the points. In one of the previous examples,
    // we have explained why the
    // value_list/vector_value_list function had
    // been introduced: to prevent us from calling virtual functions too
    // frequently. On the other hand, we now need to implement the same
    // function twice, which can lead to confusion if one function is changed
    // but the other is not.
    //
    // We can prevent this situation by calling
    // RightHandSide::vector_value on each point in the input
    // list. Note that by giving the full name of the function, including the
    // class name, we instruct the compiler to explicitly call this function,
    // and not to use the virtual function call mechanism that would be used
    // if we had just called vector_value. This is important,
    // since the compiler generally can't make any assumptions which function
    // is called when using virtual functions, and it therefore can't inline
    // the called function into the site of the call. On the contrary, here we
    // give the fully qualified name, which bypasses the virtual function
    // call, and consequently the compiler knows exactly which function is
    // called and will inline above function into the present location. (Note
    // that we have declared the vector_value function above
    // inline, though modern compilers are also able to inline
    // functions even if they have not been declared as inline).
    //
    // It is worth noting why we go to such length explaining what we
    // do. Using this construct, we manage to avoid any inconsistency: if we
    // want to change the right hand side function, it would be difficult to
    // always remember that we always have to change two functions in the same
    // way. Using this forwarding mechanism, we only have to change a single
    // place (the vector_value function), and the second place
    // (the vector_value_list function) will always be consistent
    // with it. At the same time, using virtual function call bypassing, the
    // code is no less efficient than if we had written it twice in the first
    // place:
    for (unsigned int p=0; p<n_points; ++p)
      RightHandSide<dim>::vector_value (points[p],
                                        value_list[p]);
  }

  // @sect3{The ElasticProblem class implementation}

  // @sect4{ElasticProblem::ElasticProblem}

  // Following is the constructor of the main class. As said before, we would
  // like to construct a vector-valued finite element that is composed of
  // several scalar finite elements (i.e., we want to build the vector-valued
  // element so that each of its vector components consists of the shape
  // functions of a scalar element). Of course, the number of scalar finite
  // elements we would like to stack together equals the number of components
  // the solution function has, which is dim since we consider
  // displacement in each space direction. The FESystem class can
  // handle this: we pass it the finite element of which we would like to
  // compose the system of, and how often it shall be repeated:

  template <int dim>
  ElasticProblem<dim>::ElasticProblem ()
    :
    dof_handler (triangulation),
    fe (FE_Q<dim>(1), dim)
  {}
  // In fact, the FESystem class has several more constructors
  // which can perform more complex operations than just stacking together
  // several scalar finite elements of the same type into one; we will get to
  // know these possibilities in later examples.

  // @sect4{ElasticProblem::~ElasticProblem}

  // The destructor, on the other hand, is exactly as in step-6:
  template <int dim>
  ElasticProblem<dim>::~ElasticProblem ()
  {
    dof_handler.clear ();
  }

  // @sect4{ElasticProblem::setup_system}

  // Setting up the system of equations is identitical to the function used in
  // the step-6 example. The DoFHandler class and all other
  // classes used here are fully aware that the finite element we want to use
  // is vector-valued, and take care of the vector-valuedness of the finite
  // element themselves. (In fact, they do not, but this does not need to
  // bother you: since they only need to know how many degrees of freedom
  // there are per vertex, line and cell, and they do not ask what they
  // represent, i.e. whether the finite element under consideration is
  // vector-valued or whether it is, for example, a scalar Hermite element
  // with several degrees of freedom on each vertex).
  template <int dim>
  void ElasticProblem<dim>::setup_system ()
  {
    dof_handler.distribute_dofs (fe);
    hanging_node_constraints.clear ();
    DoFTools::make_hanging_node_constraints (dof_handler,
                                             hanging_node_constraints);
    hanging_node_constraints.close ();
    sparsity_pattern.reinit (dof_handler.n_dofs(),
                             dof_handler.n_dofs(),
                             dof_handler.max_couplings_between_dofs());
    DoFTools::make_sparsity_pattern (dof_handler, sparsity_pattern);

    hanging_node_constraints.condense (sparsity_pattern);

    sparsity_pattern.compress();

    system_matrix.reinit (sparsity_pattern);

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

  // @sect4{ElasticProblem::assemble_system}

  // The big changes in this program are in the creation of matrix and right
  // hand side, since they are problem-dependent. We will go through that
  // process step-by-step, since it is a bit more complicated than in previous
  // examples.
  //
  // The first parts of this function are the same as before, however: setting
  // up a suitable quadrature formula, initializing an FEValues
  // object for the (vector-valued) finite element we use as well as the
  // quadrature object, and declaring a number of auxiliary arrays. In
  // addition, we declare the ever same two abbreviations:
  // n_q_points and dofs_per_cell. The number of
  // degrees of freedom per cell we now obviously ask from the composed finite
  // element rather than from the underlying scalar Q1 element. Here, it is
  // dim times the number of degrees of freedom per cell of the
  // Q1 element, though this is not explicit knowledge we need to care about:
  template <int dim>
  void ElasticProblem<dim>::assemble_system ()
  {
    QGauss<dim>  quadrature_formula(2);

    FEValues<dim> fe_values (fe, quadrature_formula,
                             update_values   | update_gradients |
                             update_quadrature_points | update_JxW_values);

    const unsigned int   dofs_per_cell = fe.dofs_per_cell;
    const unsigned int   n_q_points    = quadrature_formula.size();

    FullMatrix<double>   cell_matrix (dofs_per_cell, dofs_per_cell);
    Vector<double>       cell_rhs (dofs_per_cell);

    std::vector<unsigned int> local_dof_indices (dofs_per_cell);

    // As was shown in previous examples as well, we need a place where to
    // store the values of the coefficients at all the quadrature points on a
    // cell. In the present situation, we have two coefficients, lambda and
    // mu.
    std::vector<double>     lambda_values (n_q_points);
    std::vector<double>     mu_values (n_q_points);

    // Well, we could as well have omitted the above two arrays since we will
    // use constant coefficients for both lambda and mu, which can be declared
    // like this. They both represent functions always returning the constant
    // value 1.0. Although we could omit the respective factors in the
    // assemblage of the matrix, we use them here for purpose of
    // demonstration.
    ConstantFunction<dim> lambda(1.), mu(1.);

    // Then again, we need to have the same for the right hand side. This is
    // exactly as before in previous examples. However, we now have a
    // vector-valued right hand side, which is why the data type of the
    // rhs_values array is changed. We initialize it by
    // n_q_points elements, each of which is a
    // Vector@<double@> with dim elements.
    RightHandSide<dim>      right_hand_side;
    std::vector<Vector<double> > rhs_values (n_q_points,
                                             Vector<double>(dim));

    // Now we can begin with the loop over all cells:
    typename DoFHandler<dim>::active_cell_iterator cell = dof_handler.begin_active(),
                                                   endc = dof_handler.end();
    for (; cell!=endc; ++cell)
      {
        cell_matrix = 0;
        cell_rhs = 0;

        fe_values.reinit (cell);

        // Next we get the values of the coefficients at the quadrature
        // points. Likewise for the right hand side:
        lambda.value_list (fe_values.get_quadrature_points(), lambda_values);
        mu.value_list     (fe_values.get_quadrature_points(), mu_values);

        right_hand_side.vector_value_list (fe_values.get_quadrature_points(),
                                           rhs_values);

        // Then assemble the entries of the local stiffness matrix and right
        // hand side vector. This follows almost one-to-one the pattern
        // described in the introduction of this example.  One of the few
        // comments in place is that we can compute the number
        // comp(i), i.e. the index of the only nonzero vector
        // component of shape function i using the
        // fe.system_to_component_index(i).first function call
        // below.
        //
        // (By accessing the first variable of the return value
        // of the system_to_component_index function, you might
        // already have guessed that there is more in it. In fact, the
        // function returns a std::pair@<unsigned int, unsigned
        // int@>, of which the first element is comp(i)
        // and the second is the value base(i) also noted in the
        // introduction, i.e.  the index of this shape function within all the
        // shape functions that are nonzero in this component,
        // i.e. base(i) in the diction of the introduction. This
        // is not a number that we are usually interested in, however.)
        //
        // With this knowledge, we can assemble the local matrix
        // contributions:
        for (unsigned int i=0; i<dofs_per_cell; ++i)
          {
            const unsigned int
            component_i = fe.system_to_component_index(i).first;

            for (unsigned int j=0; j<dofs_per_cell; ++j)
              {
                const unsigned int
                component_j = fe.system_to_component_index(j).first;

                for (unsigned int q_point=0; q_point<n_q_points;
                     ++q_point)
                  {
                    cell_matrix(i,j)
                    +=
                      // The first term is (lambda d_i u_i, d_j v_j) + (mu d_i
                      // u_j, d_j v_i).  Note that
                      // shape_grad(i,q_point) returns the
                      // gradient of the only nonzero component of the i-th
                      // shape function at quadrature point q_point. The
                      // component comp(i) of the gradient, which
                      // is the derivative of this only nonzero vector
                      // component of the i-th shape function with respect to
                      // the comp(i)th coordinate is accessed by the appended
                      // brackets.
                      (
                        (fe_values.shape_grad(i,q_point)[component_i] *
                         fe_values.shape_grad(j,q_point)[component_j] *
                         lambda_values[q_point])
                        +
                        (fe_values.shape_grad(i,q_point)[component_j] *
                         fe_values.shape_grad(j,q_point)[component_i] *
                         mu_values[q_point])
                        +
                        // The second term is (mu nabla u_i, nabla v_j).  We
                        // need not access a specific component of the
                        // gradient, since we only have to compute the scalar
                        // product of the two gradients, of which an
                        // overloaded version of the operator* takes care, as
                        // in previous examples.
                        //
                        // Note that by using the ?: operator, we only do this
                        // if comp(i) equals comp(j), otherwise a zero is
                        // added (which will be optimized away by the
                        // compiler).
                        ((component_i == component_j) ?
                         (fe_values.shape_grad(i,q_point) *
                          fe_values.shape_grad(j,q_point) *
                          mu_values[q_point])  :
                         0)
                      )
                      *
                      fe_values.JxW(q_point);
                  }
              }
          }

        // Assembling the right hand side is also just as discussed in the
        // introduction:
        for (unsigned int i=0; i<dofs_per_cell; ++i)
          {
            const unsigned int
            component_i = fe.system_to_component_index(i).first;

            for (unsigned int q_point=0; q_point<n_q_points; ++q_point)
              cell_rhs(i) += fe_values.shape_value(i,q_point) *
                             rhs_values[q_point](component_i) *
                             fe_values.JxW(q_point);
          }

        // The transfer from local degrees of freedom into the global matrix
        // and right hand side vector does not depend on the equation under
        // consideration, and is thus the same as in all previous
        // examples. The same holds for the elimination of hanging nodes from
        // the matrix and right hand side, once we are done with assembling
        // the entire linear system:
        cell->get_dof_indices (local_dof_indices);
        for (unsigned int i=0; i<dofs_per_cell; ++i)
          {
            for (unsigned int j=0; j<dofs_per_cell; ++j)
              system_matrix.add (local_dof_indices[i],
                                 local_dof_indices[j],
                                 cell_matrix(i,j));

            system_rhs(local_dof_indices[i]) += cell_rhs(i);
          }
      }

    hanging_node_constraints.condense (system_matrix);
    hanging_node_constraints.condense (system_rhs);

    // The interpolation of the boundary values needs a small modification:
    // since the solution function is vector-valued, so need to be the
    // boundary values. The ZeroFunction constructor accepts a
    // parameter that tells it that it shall represent a vector valued,
    // constant zero function with that many components. By default, this
    // parameter is equal to one, in which case the ZeroFunction
    // object would represent a scalar function. Since the solution vector has
    // dim components, we need to pass dim as number
    // of components to the zero function as well.
    std::map<unsigned int,double> boundary_values;
    VectorTools::interpolate_boundary_values (dof_handler,
                                              0,
                                              ZeroFunction<dim>(dim),
                                              boundary_values);
    MatrixTools::apply_boundary_values (boundary_values,
                                        system_matrix,
                                        solution,
                                        system_rhs);
  }

  // @sect4{ElasticProblem::solve}

  // The solver does not care about where the system of equations comes, as
  // long as it stays positive definite and symmetric (which are the
  // requirements for the use of the CG solver), which the system indeed
  // is. Therefore, we need not change anything.
  template <int dim>
  void ElasticProblem<dim>::solve ()
  {
    SolverControl           solver_control (1000, 1e-12);
    SolverCG<>              cg (solver_control);

    PreconditionSSOR<> preconditioner;
    preconditioner.initialize(system_matrix, 1.2);

    cg.solve (system_matrix, solution, system_rhs,
              preconditioner);

    hanging_node_constraints.distribute (solution);
  }

  // @sect4{ElasticProblem::refine_grid}

  // The function that does the refinement of the grid is the same as in the
  // step-6 example. The quadrature formula is adapted to the linear elements
  // again. Note that the error estimator by default adds up the estimated
  // obtained from all components of the finite element solution, i.e., it
  // uses the displacement in all directions with the same weight. If we would
  // like the grid to be adapted to the x-displacement only, we could pass the
  // function an additional parameter which tells it to do so and do not
  // consider the displacements in all other directions for the error
  // indicators. However, for the current problem, it seems appropriate to
  // consider all displacement components with equal weight.
  template <int dim>
  void ElasticProblem<dim>::refine_grid ()
  {
    Vector<float> estimated_error_per_cell (triangulation.n_active_cells());

    typename FunctionMap<dim>::type neumann_boundary;
    KellyErrorEstimator<dim>::estimate (dof_handler,
                                        QGauss<dim-1>(2),
                                        neumann_boundary,
                                        solution,
                                        estimated_error_per_cell);

    GridRefinement::refine_and_coarsen_fixed_number (triangulation,
                                                     estimated_error_per_cell,
                                                     0.3, 0.03);

    triangulation.execute_coarsening_and_refinement ();
  }

  // @sect4{ElasticProblem::output_results}

  // The output happens mostly as has been shown in previous examples
  // already. The only difference is that the solution function is vector
  // valued. The DataOut class takes care of this automatically,
  // but we have to give each component of the solution vector a different
  // name.
  template <int dim>
  void ElasticProblem<dim>::output_results (const unsigned int cycle) const
  {
    std::string filename = "solution-";
    filename += ('0' + cycle);
    Assert (cycle < 10, ExcInternalError());

    filename += ".vtk";
    std::ofstream output (filename.c_str());

    DataOut<dim> data_out;
    data_out.attach_dof_handler (dof_handler);

    // As said above, we need a different name for each component of the
    // solution function. To pass one name for each component, a vector of
    // strings is used. Since the number of components is the same as the
    // number of dimensions we are working in, the following
    // switch statement is used.
    //
    // We note that some graphics programs have restriction as to what
    // characters are allowed in the names of variables. The library therefore
    // supports only the minimal subset of these characters that is supported
    // by all programs. Basically, these are letters, numbers, underscores,
    // and some other characters, but in particular no whitespace and
    // minus/hyphen. The library will throw an exception otherwise, at least
    // if in debug mode.
    //
    // After listing the 1d, 2d, and 3d case, it is good style to let the
    // program die if we run upon a case which we did not consider. Remember
    // that the Assert macro generates an exception if the
    // condition in the first parameter is not satisfied. Of course, the
    // condition false can never be satisfied, so the program
    // will always abort whenever it gets to the default statement:
    std::vector<std::string> solution_names;
    switch (dim)
      {
      case 1:
        solution_names.push_back ("displacement");
        break;
      case 2:
        solution_names.push_back ("x_displacement");
        solution_names.push_back ("y_displacement");
        break;
      case 3:
        solution_names.push_back ("x_displacement");
        solution_names.push_back ("y_displacement");
        solution_names.push_back ("z_displacement");
        break;
      default:
        Assert (false, ExcNotImplemented());
      }

    // After setting up the names for the different components of the solution
    // vector, we can add the solution vector to the list of data vectors
    // scheduled for output. Note that the following function takes a vector
    // of strings as second argument, whereas the one which we have used in
    // all previous examples accepted a string there. In fact, the latter
    // function is only a shortcut for the function which we call here: it
    // puts the single string that is passed to it into a vector of strings
    // with only one element and forwards that to the other function.
    data_out.add_data_vector (solution, solution_names);
    data_out.build_patches ();
    data_out.write_vtk (output);
  }

  // @sect4{ElasticProblem::run}

  // The run function does the same things as in step-6, for
  // example. This time, we use the square [-1,1]^d as domain, and we refine
  // it twice globally before starting the first iteration.
  //
  // The reason is the following: we use the Gauss quadrature
  // formula with two points in each direction for integration of the right
  // hand side; that means that there are four quadrature points on each cell
  // (in 2D). If we only refine the initial grid once globally, then there
  // will be only four quadrature points in each direction on the
  // domain. However, the right hand side function was chosen to be rather
  // localized and in that case all quadrature points lie outside the support
  // of the right hand side function. The right hand side vector will then
  // contain only zeroes and the solution of the system of equations is the
  // zero vector, i.e. a finite element function that it zero everywhere. We
  // should not be surprised about such things happening, since we have chosen
  // an initial grid that is totally unsuitable for the problem at hand.
  //
  // The unfortunate thing is that if the discrete solution is constant, then
  // the error indicators computed by the KellyErrorEstimator
  // class are zero for each cell as well, and the call to
  // refine_and_coarsen_fixed_number on the
  // triangulation object will not flag any cells for refinement
  // (why should it if the indicated error is zero for each cell?). The grid
  // in the next iteration will therefore consist of four cells only as well,
  // and the same problem occurs again.
  //
  // The conclusion needs to be: while of course we will not choose the
  // initial grid to be well-suited for the accurate solution of the problem,
  // we must at least choose it such that it has the chance to capture the
  // most striking features of the solution. In this case, it needs to be able
  // to see the right hand side. Thus, we refine twice globally. (Note that
  // the refine_global function is not part of the
  // GridRefinement class in which
  // refine_and_coarsen_fixed_number is declared, for
  // example. The reason is first that it is not an algorithm that computed
  // refinement flags from indicators, but more importantly that it actually
  // performs the refinement, in contrast to the functions in
  // GridRefinement that only flag cells without actually
  // refining the grid.)
  template <int dim>
  void ElasticProblem<dim>::run ()
  {
    for (unsigned int cycle=0; cycle<8; ++cycle)
      {
        std::cout << "Cycle " << cycle << ':' << std::endl;

        if (cycle == 0)
          {
            GridGenerator::hyper_cube (triangulation, -1, 1);
            triangulation.refine_global (2);
          }
        else
          refine_grid ();

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

        setup_system ();

        std::cout << "   Number of degrees of freedom: "
                  << dof_handler.n_dofs()
                  << std::endl;

        assemble_system ();
        solve ();
        output_results (cycle);
      }
  }
}

// @sect3{The main function}

// After closing the Step8 namespace in the last line above, the
// following is the main function of the program and is again exactly like in
// step-6 (apart from the changed class names, of course).
int main ()
{
  try
    {
      dealii::deallog.depth_console (0);

      Step8::ElasticProblem<2> elastic_problem_2d;
      elastic_problem_2d.run ();
    }
  catch (std::exception &exc)
    {
      std::cerr << std::endl << std::endl
                << "----------------------------------------------------"
                << std::endl;
      std::cerr << "Exception on processing: " << std::endl
                << exc.what() << std::endl
                << "Aborting!" << std::endl
                << "----------------------------------------------------"
                << std::endl;

      return 1;
    }
  catch (...)
    {
      std::cerr << std::endl << std::endl
                << "----------------------------------------------------"
                << std::endl;
      std::cerr << "Unknown exception!" << std::endl
                << "Aborting!" << std::endl
                << "----------------------------------------------------"
                << std::endl;
      return 1;
    }

  return 0;
}


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

#include <iostream>
#include <fstream>
#include <opencv2/highgui.hpp>
#include <opencv2/imgproc.hpp>
#include <opencv2/objdetect.hpp>
#include <boost/program_options.hpp>
#include <boost/numeric/ublas/lu.hpp>

#include <list>

template<typename T>
std::vector<T> polyfit( const std::vector<T>& oX, const std::vector<T>& oY, int nDegree ) {
	using namespace boost::numeric::ublas;

	if ( oX.size() != oY.size() )
		throw std::invalid_argument( "X and Y vector sizes do not match" );

	// more intuative this way
	nDegree++;

	size_t nCount =  oX.size();
	matrix<T> oXMatrix( nCount, nDegree );
	matrix<T> oYMatrix( nCount, 1 );

	// copy y matrix
	for ( size_t i = 0; i < nCount; i++ )
	{
		oYMatrix(i, 0) = oY[i];
	}

	// create the X matrix
	for ( size_t nRow = 0; nRow < nCount; nRow++ )
	{
		T nVal = 1.0f;
		for ( int nCol = 0; nCol < nDegree; nCol++ )
		{
			oXMatrix(nRow, nCol) = nVal;
			nVal *= oX[nRow];
		}
	}

	// transpose X matrix
	matrix<T> oXtMatrix( trans(oXMatrix) );
	// multiply transposed X matrix with X matrix
	matrix<T> oXtXMatrix( prec_prod(oXtMatrix, oXMatrix) );
	// multiply transposed X matrix with Y matrix
	matrix<T> oXtYMatrix( prec_prod(oXtMatrix, oYMatrix) );

	// lu decomposition
	permutation_matrix<T> pert(oXtXMatrix.size1());
	const std::size_t singular = lu_factorize(oXtXMatrix, pert);
	// must be singular
	if( singular != 0 ) {
		// then the regression cannot be done... return an empty vector
		return std::vector<T>();
	}

	// backsubstitution
	lu_substitute(oXtXMatrix, pert, oXtYMatrix);

	// copy the result to coeff
	return std::vector<T>( oXtYMatrix.data().begin(), oXtYMatrix.data().end() );
}

float simpleShape(cv::Mat &image) {
	if(image.empty()) return 0;

	cv::Mat gray;
	cv::cvtColor(image, gray, cv::COLOR_BGR2GRAY);

	cv::Mat edges;
	cv::Canny(gray, edges, 200, 400);

	cv::Mat result(image.size(), CV_8UC3, cv::Scalar(0,0,0));

	std::vector<std::vector<cv::Point> > contours0;
	std::vector<std::vector<cv::Point> > contours;
	std::vector<cv::Vec4i> hierarchy;
    cv::findContours( edges, contours0, hierarchy, cv::RETR_TREE, cv::CHAIN_APPROX_SIMPLE);

	std::cout << contours0.size() << " ";

	cv::drawContours( result, contours, 3, cv::Scalar(128,0,255),
                  3, cv::LINE_AA, hierarchy, std::abs(3) );

    contours.resize(contours0.size());
    for( size_t k = 0; k < contours0.size(); k++ )
        cv::approxPolyDP(cv::Mat(contours0[k]), contours[k], 1, true);

	cv::drawContours( result, contours, 3, cv::Scalar(128,255,255),
                  3, cv::LINE_AA, hierarchy, std::abs(3) );

	cv::imshow("result", edges);
	cv::waitKey();

	return 0;
}

float  horizonLine(cv::Mat &image) {
	if(image.empty()) return 0;

	std::vector<float> x, y;

	cv::Mat gray;
	cv::cvtColor(image, gray, cv::COLOR_BGR2GRAY);

	cv::Mat edges;
	cv::Canny(gray, edges, 110, 220);

	std::vector<cv::Vec4i> lines;
    cv::HoughLinesP( edges, lines, 1, CV_PI/2, 0, 30, 1 );
    float meanY = 0; int nbLines = 0;
    float meanY2 = 0;
    for( size_t i = 0; i < lines.size(); i++ )
    {
    	if(lines[i][0] == lines[i][2]) continue;

		if(lines[i][1] > (.7 * image.rows)) continue;
		if(lines[i][1] < (.3 * image.rows)) continue;

        cv::line( image, cv::Point(lines[i][0], lines[i][1]),
            cv::Point(lines[i][2], lines[i][3]), cv::Scalar(0,0,255), 3, 8 );

        x.push_back(lines[i][0]);
		x.push_back(lines[i][2]);
		y.push_back(lines[i][1]);
		y.push_back(lines[i][3]);

		meanY += lines[i][1]; ++nbLines;
		meanY += lines[i][3]; ++nbLines;

		meanY2 += lines[i][1] * lines[i][1];
		meanY2 += lines[i][3] * lines[i][3];

    }
    meanY /= nbLines;
    meanY2 /= nbLines;

    float stdev = std::sqrt(meanY2 - meanY*meanY);
    float meanValue = meanY;

    meanY = 0; nbLines = 0;
    for( size_t i = 0; i < lines.size(); i++ )
    {
    	if(lines[i][0] == lines[i][2]) continue;

		if(lines[i][1] > (.7 * image.rows)) continue;
		if(lines[i][1] < (.3 * image.rows)) continue;

        if(std::abs(lines[i][1] - meanValue) > 1.3 * stdev) continue;

		meanY += lines[i][1]; ++nbLines;
		meanY += lines[i][3]; ++nbLines;

    }
    meanY /= nbLines;

    if(nbLines == 0) {
    	meanY = image.rows/2;
    } else {
    	std::vector<float> model = polyfit<float>(x, y, 1);

	    if(std::abs(model[0]- image.rows/2) < std::abs(meanY - image.rows/2))
	    	meanY = model[0];
    }

	// cv::line(image, cv::Point_<int>(0, meanY), cv::Point_<int>(image.cols, meanY), cv::Scalar( 0, 255, 0 ), 2);

return meanY;

}

float faceLine(cv::Mat &image) {
	cv::Mat gray;
	cv::cvtColor(image, gray, cv::COLOR_BGR2GRAY);

	cv::CascadeClassifier face_cascade;
	cv::CascadeClassifier faceProfil_cascade;
	bool faceCascadeEnabled = false;
	bool faceProfilCascadeEnabled = false;

	// ------------------------------------------------------------------------------------------------
	// init haar cascades framework

	if(face_cascade.load("haarcascade_frontalface_alt.xml")) {
		faceCascadeEnabled = true;
	} else {
		std::cerr<< "[I] cannot open: haarcascade_frontalface_alt.xml" << std::endl;
	}

	if(faceProfil_cascade.load("haarcascade_profileface.xml")) {
		faceProfilCascadeEnabled = true;
	} else {
		std::cerr<< "[I] cannot open: haarcascade_profileface.xml" << std::endl;
	}

	// ------------------------------------------------------------------------------------------------
	// find features

	std::list<cv::Rect> allFeatures;
	std::vector<cv::Rect> faceFeatures;
	if(faceCascadeEnabled)
		face_cascade.detectMultiScale( gray, faceFeatures, 2, 2, 0| cv::CASCADE_SCALE_IMAGE, cv::Size(15, 15) );

	for(size_t i = 0 ; i < faceFeatures.size() ; ++i) {
		cv::rectangle(image, faceFeatures[i], cv::Scalar( 0, 0, 255 ));
		allFeatures.push_back(faceFeatures[i]);
	}

	faceFeatures.clear();
	if(faceProfilCascadeEnabled)
		faceProfil_cascade.detectMultiScale( gray, faceFeatures, 2, 2, 0|cv::CASCADE_SCALE_IMAGE, cv::Size(15, 15) );

	for(size_t i = 0 ; i < faceFeatures.size() ; ++i) {
		cv::rectangle(image, faceFeatures[i], cv::Scalar( 255, 255, 0 ));
		allFeatures.push_back(faceFeatures[i]);
	}

	if(allFeatures.size() < 2) return -1;

	float mnX = image.cols;
	float mxX = 0;
	for( std::list<cv::Rect>::iterator it = allFeatures.begin() ; it != allFeatures.end() ; ++it ) {
		float x = it->x + it->width /2;

		if(x > mxX) mxX = x;
		if(x < mnX) mnX = x;
	}

	if((mxX-mnX)/image.cols < 0.2) return -1;	// if the detected faces are collocated in one area, then skip it.

	// ------------------------------------------------------------------------------------------------
	// find equatorial line

	float meanY = 0; int nbLines = 0;
    float meanY2 = 0;
    for( std::list<cv::Rect>::iterator it = allFeatures.begin() ; it != allFeatures.end() ; ++it ) {
    	meanY += it->y + it->height /2;
    	meanY2 += (it->y + it->height /2)*(it->y + it->height /2);
    	++nbLines;
    }

    if(nbLines == 0) return -1;

    meanY /= nbLines;
    meanY2 /= nbLines;

    float stdev = std::sqrt(meanY2 - meanY*meanY);
    float meanValue = meanY;

    meanY = 0; nbLines = 0;

    for( std::list<cv::Rect>::iterator it = allFeatures.begin() ; it != allFeatures.end() ; ++it ) {
    	float y = it->y + it->height /2;

    	if(std::abs(y-meanValue) > 1.3 * stdev) continue;

    	meanY += y;
    	++nbLines;
    }

    if(nbLines == 0) return meanValue;

	meanY /= nbLines;

	if(meanY < (gray.rows * 1 / 3) && allFeatures.size() < 2) { // we need at least two faces to justify an equatorial line lower higher the first third of the image
		return -1;
	}

	return meanY;

}

float percentageSaliency(const cv::Mat& image, float slicePosition, float sliceHeightDegree = 10.f) {

	double sumImage = cv::sum( image )[0];
	float sumSlice = 0;

	float sliceHeight = sliceHeightDegree * (image.rows / 180);

	for(int i = static_cast<int>(std::max(slicePosition - sliceHeight, 0.f)) ; i < std::min(static_cast<int>(slicePosition + sliceHeight), image.rows) ; ++i) {
		for(int j = 0 ; j < image.cols ; ++j) {
			sumSlice += image.at<float>(i,j);
		}
	}

	return sumSlice / sumImage;

}

float salientCenter(const cv::Mat& image, int step = 5) {

	std::vector<float> histogram(180/step);
	int offset = 0; //static_cast<int>(30.f*(step/180.f));

	float sumHist = 0;
	for(int i = offset ; i < 180-offset ; i+=step) {
		float top    = image.rows*static_cast<float>(i)/180;
		float bottom = image.rows*static_cast<float>(i+step)/180;

		float sum = 0;
		for(int ii = top ; ii < bottom ; ++ii) {
			for(int jj = 0 ; jj < image.cols ; ++jj) {
				sum += image.at<float>(ii,jj);
			}
		}

		sum /= (bottom-top) * image.cols;
		histogram[i/step] = sum;
		sumHist += sum;
	}

	float loc = 0;
	for(size_t i = offset ; i < histogram.size() - offset ; ++i) {
		loc += i*step*histogram[i]/sumHist;
	}

	return image.rows * static_cast<float>(loc)/180.f;

}

void applyGaussianEquatorialPrior(cv::Mat& image, float gaussianM = 0.f, float gaussianSD = 500.f, bool central = false) {

	for(int i = 0 ; i < image.rows ; ++i) {
		float lat = (.5f - static_cast<float>(i)/image.rows) * 180.f;
		float prior = 0.f;

		if(central)
			prior =  (.01 + .40 * exp(-((lat-gaussianM)*(lat-gaussianM))/(gaussianSD))) / .41f;
		else
			prior =  (.01 + .40 * exp(-((lat-gaussianM)*(lat-gaussianM))/(gaussianSD)) + .75 * exp(-((lat-100)*(lat-100))/(200)) + .75 * exp(-((lat+90)*(lat+90))/(200))) / .41f;

		prior = std::min(1.f, prior);

		for(int j = 0 ; j < image.cols ; ++j) {
			image.at<float>(i,j) = prior * image.at<float>(i,j);
		}
	}

}

void applyEquatorialPrior(cv::Mat& image, cv::Mat& colorImage, bool central = false, bool print = false) {
	float scaling_factor = static_cast<float>(image.cols) / 1400.f;	// normalize the size of the images

	cv::resize(colorImage, colorImage, cv::Size(colorImage.size().width/scaling_factor, colorImage.size().height/scaling_factor));
	float hz = colorImage.rows / 2;
	float fc = faceLine(colorImage);

	float equatorialLine = hz;

	if(fc > 0) {
		if(fc > colorImage.rows / 3)	// prior: faces are not higher than the first third of the image
			equatorialLine = fc;
	}

	float slCenter = salientCenter(image);

	slCenter = std::max(std::min((slCenter - image.rows/2)/image.rows, 0.1f), -0.1f);
	slCenter = slCenter * image.rows + image.rows/2;
	slCenter = (slCenter/scaling_factor + colorImage.rows / 2) / 2;

	if(fc > 0) {
		if(fc > colorImage.rows / 3)	// prior: faces are not higher than the first third of the image
			slCenter = fc;
	}

	equatorialLine = (.5f - slCenter / colorImage.rows) * 180.f;

	if(print) std::cout << equatorialLine << std::endl;

	applyGaussianEquatorialPrior(image, equatorialLine, 700.f, central);

}

void bottomThresholdSaliency(cv::Mat &saliency) {
	cv::Mat uCharSalmap;
	uCharSalmap = 255 * saliency;
	uCharSalmap.convertTo(uCharSalmap, CV_8UC1);

	int histSize[] = {255};
	int channels[] = {0};
	float range[] = { 0, 256 } ;
 	const float* histRange = { range };
	cv::Mat hist;

	cv::calcHist(&uCharSalmap, 1, channels, cv::Mat(), hist, 1, histSize, &histRange, true, true);

	float sum = 0;
	float mid = static_cast<float>(saliency.cols*saliency.rows) * .7;
	int theshold = 0;
	for(int i = 0 ; i < hist.rows ; ++i) {
		sum += hist.at<float>(i,0);

		if(sum > mid) {
			theshold = i;
			break;
		}
	}

	for(int i = 0 ; i < uCharSalmap.rows ; ++i) {
		for(int j = 0 ; j < uCharSalmap.cols ; ++j) {
			if(uCharSalmap.at<unsigned char>(i,j) < theshold) {
				uCharSalmap.at<unsigned char>(i,j) = theshold;
			}
		}
	}

	uCharSalmap.convertTo(saliency, CV_32FC1);
	uCharSalmap /= 255;

}

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

	// ------------------------------------------------------------------------------------------------------------------------------------------
	// parse parameters

	namespace po = boost::program_options;

	po::positional_options_description p;
	p.add("input-file", -1);

	po::options_description desc("Allowed options");
	desc.add_options()
			("help", "produce help message")
			("input-file,i", po::value< std::string >(), "Saliency input image.")
			("color-file,c", po::value< std::string >(), "Color image corresponding to the saliency map.")
			("line,l", po::value< float >(), "Add the ground truth equatorial line (in degree [-90,+90]).")
			("output-file,o", po::value< std::string >(), "Saliency map with the equatorial prior.")
			("central", "Consider only the central part of the Gaussian mixture")
			("no-adapt", "No adaptation")
			("print-horizon", "Output in the terminal the position of the horizon line")
	;

	po::variables_map vm;
	try {
		po::store(po::command_line_parser(argc, argv).options(desc).positional(p).run(), vm);
		po::notify(vm);
	} catch(boost::exception &) {
		std::cerr << "Error incorect program options. See --help... \n";
		return 0;
	}

	if (vm.count("help")) {
		std::cout << "--------------------------------------------------------------------------------\n";
		std::cout << "\t\tApply the equatorial prior on saliency maps\n";
		std::cout << "--------------------------------------------------------------------------------\n";
		std::cout << "\n\n";
		std::cout << desc << "\n";
		return 1;
	}

	std::string inputImagePath;
	std::string inputColorImagePath;
	std::string outputPath;

	bool central = false;
	bool noAdapt = false;

	if (vm.count("input-file")) {
		inputImagePath = vm["input-file"].as< std::string >();
	} else {
		std::cerr << "It is required to provide the an input saliency map. See --help\n";
		return 0;
	}

	if (vm.count("color-file")) {
		inputColorImagePath = vm["color-file"].as< std::string >();
	} else {
		std::cerr << "It is required to provide the an input color image. See --help\n";
		return 0;
	}

	if (vm.count("output-file")) {
		outputPath = vm["output-file"].as< std::string >();
	}

	if(vm.count("central")) {
		central = true;
	}

	if(vm.count("no-adapt")) {
		noAdapt = true;
	}

	bool printHz = false;
	if(vm.count("print-horizon")) {
		printHz = true;
	}

	// ---------------------------------------------------------------------------------------------------
	// Get statistics

	cv::Mat image = cv::imread(inputImagePath);
	if(image.empty()) {
		std::cerr << "Cannot open: " << inputImagePath << std::endl;
		return 0;
	}

	cv::Mat colorImage = cv::imread(inputColorImagePath);
	if(colorImage.empty()) {
		std::cerr << "Cannot open: " << inputColorImagePath << std::endl;
		return 0;
	}

	cv::cvtColor(image, image, cv::COLOR_BGR2GRAY);

	image.convertTo(image, CV_32FC1);
	image /= 255;

	if(vm.count("second-file") && vm.count("weight")) {
		cv::Mat imageW = cv::imread(vm["second-file"].as< std::string >());
		if(imageW.empty()) {
			std::cerr << "cannot open: " << vm["second-file"].as< std::string >() << std::endl ;
		} else {
			cv::cvtColor(imageW, imageW, cv::COLOR_BGR2GRAY);
			imageW.convertTo(imageW, CV_32FC1);
			imageW /= 255;

			float w = vm["weight"].as< float >();
			image = image * w  + (1 - w) * imageW;
		}
	}

	// image = cv::Mat(image.size(), CV_32FC1, cv::Scalar(1.0f));

	if(printHz) std::cout << inputColorImagePath << ", ";

	if(!noAdapt)
		applyEquatorialPrior(image, colorImage, central, printHz);
	else
		applyGaussianEquatorialPrior(image, 0, 700, central);

	// bottomThresholdSaliency(image);

	// if (vm.count("line")) {
	// 	float line = vm["line"].as< float >();

	// 	line =  image.rows * (90 - line) / 180.f ;

	// 	// std::cout << line << std::endl;
	// applyGaussianEquatorialPrior(image, 0, 500);

	double mn, mx;
	cv::minMaxLoc(image, &mn, &mx);

	image = (image - mn) / (mx - mn);

	// 	// line =  colorImage.rows * (90 - line) / 180.f ;

	// 	// if(line > 0 && line < colorImage.rows)
	// 	// 	cv::line(colorImage, cv::Point_<int>(0, line), cv::Point_<int>(colorImage.cols, line), cv::Scalar( 255, 0, 255), 2);
	// }

	// cv::imshow("colorImage", colorImage);
	// cv::waitKey();

	if(!outputPath.empty()) {
		cv::Mat tmp;
		cv::cvtColor(image, tmp, cv::COLOR_GRAY2BGR);
		tmp *= 255;
		tmp.convertTo(tmp, CV_8UC3);
		cv::imwrite(outputPath, tmp);

		return 0;
	} else {
		cv::imshow("with equatorial prior", image);
		cv::waitKey();
	}

	return 0;
}


#pragma once
#include <Eigen/Core>

#include "smoothness_matrix.hpp"
#include "total_derivative.hpp"
#include "../fixed_size_container_type_trait.hpp"
#include "../value_type_trait.hpp"

namespace ubs {
/** @brief Internal namespace */
namespace internal {

/**
 * @brief Implementation to evaluate a spline at the specified position.
 *
 * This is implemented using template meta programming (TMP) as one needs to have one sum per dimension. A spline
 * can be evaluated using the following formula.
 * @f[
 * \mathbf{f}(\mathbf{t}) = \sum_{i_1} N_{i_1} \sum_{i_2} N_{i_2} \cdots \sum_{i_M} N_{i_M}
 * \mathbf{r}(\mathbf{t})
 * @f]
 *
 * @note As this class uses the internal of the spline, it must be a friend of the UniformBSpline class.
 * @tparam InputDims_ The total number of input dimensions.
 * @tparam CurDim_ The current dimension.
 */
template <typename Spline_, int InputDims_, int CurDim_>
struct EvaluateBSpline {
    /**
     * @brief Evaluates the spline at the specified position.
     *
     * @param[in] spline The spline.
     * @param[in] strides The strides of the control points.
     * @param[in] idx The current flattened index in the control points.
     * @param[in] basisVal The current basis value.
     * @param[in] pos The position.
     * @param[in] basisFunc The basis function.
     * @param[in] evalFunc The evaluation function.
     */
    template <typename T, typename EvalFunc>
    static void apply(const std::array<typename Spline_::BasisType, Spline_::InputDims>& fullBasis,
                      const std::array<int, Spline_::InputDims>& strides,
                      int idx,
                      T basisVal,
                      EvalFunc evalFunc) {
        const auto basis = fullBasis[CurDim_];
        for (int i = 0; i < Spline_::Order; ++i, idx += strides[CurDim_]) {
            const T newBasisVal = basisVal * basis[i];
            EvaluateBSpline<Spline_, InputDims_, CurDim_ + 1>::apply(fullBasis, strides, idx, newBasisVal, evalFunc);
        }
    }
};

/**
 * @brief 'Recursion' end, if the current dimension matches the input dimensions.
 * @tparam InputDims_ The input dimensions.
 */
template <typename Spline_, int InputDims_>
struct EvaluateBSpline<Spline_, InputDims_, InputDims_> {
    /**
     * @brief Call the evaluate function.
     */
    template <typename T, typename EvalFunc>
    static void apply(const std::array<typename Spline_::BasisType, Spline_::InputDims>& /*fullBasis*/,
                      const std::array<int, Spline_::InputDims>& /*strides*/,
                      int idx,
                      T basisVal,
                      EvalFunc evalFunc) {
        evalFunc(idx, basisVal);
    }
};

/**
 * @brief Implementation to evaluate the spline smoothness.
 * @note As this class uses the internal of the spline, it must be a friend of the UniformBSpline class.
 * @tparam TotalDerivative_ The order of the total derivative.
 * @tparam PartialDerivativeIdx_ The current partial derivative index.
 * @tparam InputDims_ The total number of input dimensions.
 * @tparam CurDim_ The current dimension.
 */
template <int TotalDerivative_, int PartialDerivativeIdx_, int InputDims_, int CurDim_>
class EvaluateBSplineSmoothness {
private:
    static constexpr auto FullPartialDerivatives_ =
        TotalDerivative<InputDims_, TotalDerivative_>::getPartialDerivatives();

    static constexpr int PartialDerivative_ =
        FullPartialDerivatives_[PartialDerivativeIdx_].partialDerivatives[CurDim_];

public:
    /**
     * @brief Evaluates the spline smoothness.
     * @param[in] spline The spline.
     * @param[in] idx1 The first index.
     * @param[in] idx2 The seconds index.
     * @param[in] strides The strides of the control points.
     * @param[in] factor The multiplication factor of the smoothness matrix.
     * @param[out] res The resulting smoothness.
     */
    template <typename Spline>
    static void apply(const Spline& spline,
                      int idx1,
                      int idx2,
                      const std::array<int, Spline::InputDims>& strides,
                      typename Spline::ValueType factor,
                      typename Spline::OutputType& res) {

        const auto smoothnessMatrix = UniformBSplineSmoothnessBasis<Spline::Order, PartialDerivative_>::matrix()
                                          .template cast<typename Spline::ValueType>()
                                          .eval();

        for (int p = 0; p < spline.getNumControlPoints(CurDim_) - Spline::Degree;
             ++p, idx1 += strides[CurDim_], idx2 += strides[CurDim_]) {
            int newIdx1 = idx1;
            for (int i = 0; i < Spline::Order; ++i, newIdx1 += strides[CurDim_]) {
                int newIdx2 = idx2;
                for (int j = 0; j < Spline::Order; ++j, newIdx2 += strides[CurDim_]) {
                    auto newFactor = factor * smoothnessMatrix(i, j);
                    EvaluateBSplineSmoothness<TotalDerivative_, PartialDerivativeIdx_, InputDims_, CurDim_ + 1>::apply(
                        spline, newIdx1, newIdx2, strides, newFactor, res);
                }
            }
        }
    }
};

/**
 * @brief 'Recursion' end, if the current dimension matches the input dimensions.
 * @tparam TotalDerivative_ The order of the total derivative.
 * @tparam PartialDerivativeIdx_ The current partial derivative index.
 * @tparam InputDims_ The input dimensions.
 */
template <int TotalDerivative_, int PartialDerivativeIdx_, int InputDims_>
class EvaluateBSplineSmoothness<TotalDerivative_, PartialDerivativeIdx_, InputDims_, InputDims_> {
public:
    template <typename Spline>
    static void apply(const Spline& spline,
                      int idx1,
                      int idx2,
                      const std::array<int, Spline::InputDims>& /*strides*/,
                      typename Spline::ValueType factor,
                      typename Spline::OutputType& res) {
        assert(idx1 < int(spline.controlPoints_.getNumElements()));
        assert(idx2 < int(spline.controlPoints_.getNumElements()));

        res += FixedSizeContainerTypeTrait<typename Spline::OutputType>::evalSmoothness(
            factor, spline.controlPoints_.at(idx1), spline.controlPoints_.at(idx2));
    }
};

/**
 * @brief Compute the smoothness with respect to the total derivative.
 *
 * To calculate the smoothness of the total derivative, the problem can be split up by calculating the smoothness for
 * each partial derivative, multiply that value by the multiplicity of that derivative and sum them up.
 *
 * @tparam InputDims_ The number of input dimensions.
 * @tparam TotalDerivative_ The order of the total derivative.
 * @tparam NumPartialDerivatives_ The total number of distinct partial derivatives.
 * @tparam PartialDerivativeIdx_ The current number of the partial derivative to compute the smoothness.
 */
template <int InputDims_, int TotalDerivative_, int NumPartialDerivatives_, int PartialDerivativeIdx_>
struct ComputeSmoothness {
private:
    static constexpr auto FullPartialDerivatives_ =
        TotalDerivative<InputDims_, TotalDerivative_>::getPartialDerivatives();
    static constexpr Array<int, InputDims_> PartialDerivatives_ =
        FullPartialDerivatives_[PartialDerivativeIdx_].partialDerivatives;
    static constexpr double PartialDerivativesMultiplicity_ =
        FullPartialDerivatives_[PartialDerivativeIdx_].multiplicity;

public:
    /**
     * @brief Compute the smoothness value of the partial derivatives needed to calculate the total derivative
     *        smoothness value
     * @param[in] spline The spline.
     * @param[in] strides The strides of the control points.
     * @param[out] res The smoothness value for each dimension.
     */
    template <typename Spline>
    static void apply(const Spline& spline,
                      const std::array<int, InputDims_>& strides,
                      typename Spline::OutputType& res) {
        static_assert(InputDims_ == Spline::InputDims, "Input dimension missmatch.");
        using T = typename Spline::ValueType;
        using OutputType = typename Spline::OutputType;

        // Calculate the scaling factor. The scaling factor depends on the multiplicity of the partial derivative,
        // the control points scale and the partial derivative.
        T scale = T(PartialDerivativesMultiplicity_);

        for (int dim = 0; dim < InputDims_; ++dim) {
            using ValueTypeTrait = ValueTypeTrait<T>;
            scale *= ValueTypeTrait::pow(spline.getScale(dim), 2 * PartialDerivatives_[dim] - 1);
        }

        // Calculate the smoothness value.
        OutputType smoothness = FixedSizeContainerTypeTrait<OutputType>::zero();
        EvaluateBSplineSmoothness<TotalDerivative_, PartialDerivativeIdx_, InputDims_, 0>::apply(
            spline, 0, 0, strides, T(1.0), smoothness);
        res += scale * smoothness;

        // Compute smoothness value for the next partial derivative.
        ComputeSmoothness<InputDims_, TotalDerivative_, NumPartialDerivatives_, PartialDerivativeIdx_ + 1>::apply(
            spline, strides, res);
    }
};

/**
 * @brief 'Loop' end, if the current dimension matches the input dimensions.
 * @tparam InputDims_ The input dimensions.
 */
template <int InputDims_, int TotalDerivative_, int NumPartialDerivatives_>
struct ComputeSmoothness<InputDims_, TotalDerivative_, NumPartialDerivatives_, NumPartialDerivatives_> {
    template <typename Spline>
    static void apply(const Spline& /*spline*/,
                      const std::array<int, InputDims_>& /*strides*/,
                      const typename Spline::OutputType& /*res*/) {
        // No more partial derivatives left. Nothing more to do.
    }
};

} // namespace internal
} // namespace ubs

#include "uniform_bspline_eval_impl.hpp"


/**
 * @file
 * @license BSD 3-clause
 * @copyright Copyright (c) 2020, New York University and Max Planck
 * Gesellschaft
 *
 * @brief Implements the "Walking Control Based on Step Timing Adaptation" step
 * planner. The pdf can be found in https://arxiv.org/abs/1704.01271, and in the
 * `doc/` folder in this repository.
 */

#pragma once

#include <eigen-quadprog/QuadProg.h>
#include <Eigen/Eigen>
#include <cmath>
#include <iostream>
#include <pinocchio/spatial/se3.hpp>
#include <sstream>
#include <stdexcept>
#define RESET "\033[0m"
#define RED "\033[31m"  /* Red */
#define BLUE "\033[34m" /* Blue */

namespace Eigen
{
/** @brief Column vector of size 9 which correspond to the number of
 * optimization variables. */
typedef Matrix<double, 9, 1> Vector9d;
}  // namespace Eigen

namespace reactive_planners
{
/**
 * @brief Implements the "Walking Control Based on Step Timing Adaptation" step
 * planner. The pdf can be found in https://arxiv.org/abs/1704.01271, and in the
 * `doc/` folder in this repository.
 *
 * All quantities are here expressed in the local frame. Which means for a robot
 * that the quantities are expressed in the "base" frame.
 *
 * @todo write here the formulation of the QP.
 */
class DcmVrpPlanner
{
public:
    /**
     * @brief Construct a new DcmVrpPlanner object and initialize it.
     *
     * @param l_min [in] Minimum step length in the x direction (in the
     * direction of forward motion).
     * @param l_max [in] Maximum step length in the x direction (in the
     * direction of forward motion).
     * @param w_min [in] Minimum step length in the y direction (in the lateral
     * direction).
     * @param w_max [in] Maximum step length in the y direction (in the lateral
     * direction).
     * @param t_min [in] Minimum step time.
     * @param t_max [in] Maximum step time.
     * @param v_des [in] Desired average velocity in the x and y ([v_x, v_y]) 2d
     * vector.
     * @param l_p [in] Default lateral step length. Typically useful for
     * humanoid robot where this value refer to the distance between the 2 feet
     * while in the half-sitting/neutral position.
     * @param ht [in] Average desired height of the com above the ground.
     * @param cost_weights_local [in] Weights of the QP cost expressed in the
     * local frame.
     */
    DcmVrpPlanner(const double& l_min,
                  const double& l_max,
                  const double& w_min,
                  const double& w_max,
                  const double& t_min,
                  const double& t_max,
                  const double& l_p,
                  const double& ht,
                  const Eigen::Ref<const Eigen::Vector9d>& cost_weights_local);

    /**
     * @brief Construct a new DcmVrpPlanner object with some default parameters.
     */
    DcmVrpPlanner()
    {
        initialize(
            0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, Eigen::Vector9d::Zero());
    }

    /**
     * @brief Initialize The inner variable that are not time varying.
     *
     * @copydoc DcmVrpPlanner::DcmVrpPlanner()
     */
    void initialize(
        const double& l_min,
        const double& l_max,
        const double& w_min,
        const double& w_max,
        const double& t_min,
        const double& t_max,
        const double& l_p,
        const double& ht,
        const Eigen::Ref<const Eigen::Vector9d>& cost_weights_local);

    /**
     * @brief Set the nominal steptime.
     */
    void set_steptime_nominal(double t_nom);

    /**
     * @brief Computes adapted step location.
     *
     * We use a QP formulation with the following notation:
     * \f{eqnarray*}{
     *    minimize   & \\
     *       x       & (1/2) x^T Q x + q^T x \\
     *    subject to & \\
     *               & A_{ineq} x \leq b_{ineq} \\
     *               & A_{eq} x = b_{eq} \\
     *               & x_{opt_{lb}} \leq x \leq x_{opt_{ub}}
     * \f}
     * We use the off-the-shelf QP solver eigen-quadprog in order to solve it.
     * @see DcmVrpPlanner for the full formulation.
     *
     * @param current_step_location is the location of the previous foot step
     * location (2d vector) [ux, uy].
     * @param time_from_last_step_touchdown is the time elapsed since the last
     * foot step landed.
     * @param is_left_leg_in_contact is true is the current foot is on the left
     * side, false if it is on the right side.
     * @param v_des
     * @param com is the CoM position.
     * @param com_vel is the CoM velocity.
     * @param world_M_base SE3 position of the robot base expressed in the world
     * frame.
     */
    void update(const Eigen::Ref<const Eigen::Vector3d>& current_step_location,
                const double& time_from_last_step_touchdown,
                const bool& is_left_leg_in_contact,
                const Eigen::Ref<const Eigen::Vector3d>& v_des,
                const Eigen::Ref<const Eigen::Vector3d>& com,
                const Eigen::Ref<const Eigen::Vector3d>& com_vel,
                const pinocchio::SE3& world_M_base,
                const double& new_t_min);
    /**
     * @brief Computes adapted step location for python3.
     *
     * We use a QP formulation with the following notation:
     * \f{eqnarray*}{
     *    minimize   & \\
     *       x       & (1/2) x^T Q x + q^T x \\
     *    subject to & \\
     *               & A_{ineq} x \leq b_{ineq} \\
     *               & A_{eq} x = b_{eq} \\
     *               & x_{opt_{lb}} \leq x \leq x_{opt_{ub}}
     * \f}
     * We use the off-the-shelf QP solver eigen-quadprog in order to solve it.
     * @see DcmVrpPlanner for the full formulation.
     *
     * @param current_step_location is the location of the previous foot step
     * location (2d vector) [ux, uy].
     * @param time_from_last_step_touchdown is the time elapsed since the last
     * foot step landed.
     * @param is_left_leg_in_contact is true is the current foot is on the left
     * side, false if it is on the right side.
     * @param v_des
     * @param com is the CoM position.
     * @param com_vel is the CoM velocity.
     * @param yaw
     */
    void update(const Eigen::Ref<const Eigen::Vector3d>& current_step_location,
                const double& time_from_last_step_touchdown,
                const bool& is_left_leg_in_contact,
                const Eigen::Ref<const Eigen::Vector3d>& v_des,
                const Eigen::Ref<const Eigen::Vector3d>& com,
                const Eigen::Ref<const Eigen::Vector3d>& com_vel,
                const double& yaw,
                const double& new_t_min);

    /**
     * @brief Solve the Quadratic program and extract the solution. Use
     * DcmVrpPlanner::get_next_step_location() and
     * DcmVrpPlanner::get_duration_before_step_landing() to access the results.
     */
    bool solve();

    /**
     * @brief Perform an internal checks on the solver matrices. A warning
     * message is displayed in case of problems.
     *
     * @return true is everything seems fine.
     * @return false is something wrong.
     */
    bool internal_checks();

    /**
     * @brief Display the matrices of the Problem.
     */
    void print_solver() const;

    /**
     * @brief Convert the inner data to a string format.
     */
    std::string to_string() const;

    /*
     * Getters
     */

    /**
     * @brief @copydoc DcmVrpPlanner::t_nom_
     *
     * @return const double&
     */
    const double& get_t_nom() const
    {
        return t_nom_;
    }

    /**
     * @brief @copydoc DcmVrpPlanner::tau_nom_
     *
     * @return const double&
     */
    const double& get_tau_nom() const
    {
        return tau_nom_;
    }

    /**
     * @brief @copydoc DcmVrpPlanner::l_nom_
     *
     * @return const double&
     */
    const double& get_l_nom() const
    {
        return l_nom_;
    }

    /**
     * @brief @copydoc DcmVrpPlanner::w_nom_
     *
     * @return const double&
     */
    const double& get_w_nom() const
    {
        return w_nom_;
    }

    /**
     * @brief @copydoc DcmVrpPlanner::bx_nom_
     *
     * @return const double&
     */
    const double& get_bx_nom() const
    {
        return bx_nom_;
    }

    /**
     * @brief @copydoc DcmVrpPlanner::by_nom_
     *
     * @return const double&
     */
    const double& get_by_nom() const
    {
        return by_nom_;
    }

    /**
     * @brief @copydoc DcmVrpPlanner::world_M_local_
     *
     * @return const pinocchio::SE3&
     */
    const pinocchio::SE3& get_world_M_local() const
    {
        return world_M_local_;
    }

    /**
     * @brief @copydoc DcmVrpPlanner::dcm_local_
     *
     * @return Eigen::Ref<const Eigen::Vector3d>
     */

    Eigen::Vector3d dcm_local_REF;
    Eigen::Ref<const Eigen::Vector3d> get_dcm_local()
    {
        dcm_local_REF = world_M_local_.act(dcm_local_);
        dcm_local_REF = dcm_local_;
        return dcm_local_REF;
    }

    /**
     * @brief @copydoc DcmVrpPlanner::current_step_location_local_
     *
     * @return Eigen::Ref<const Eigen::Vector3d>
     */
    Eigen::Ref<const Eigen::Vector3d> get_current_step_location_local() const
    {
        return current_step_location_local_;
    }

    /**
     * @brief @copydoc DcmVrpPlanner::v_des_local_
     *
     * @return Eigen::Ref<const Eigen::Vector3d>
     */
    Eigen::Ref<const Eigen::Vector3d> get_v_des_local() const
    {
        return v_des_local_;
    }

    /**
     * @brief @copydoc DcmVrpPlanner::dcm_nominal_
     *
     * @return Eigen::Ref<const Eigen::Vector3d>
     */
    Eigen::Ref<const Eigen::Vector3d> get_dcm_nominal() const
    {
        return dcm_nominal_;
    }

    /*
     * Output
     */

    /**
     * @brief @copydoc DcmVrpPlanner::next_step_location_
     *
     * @return Eigen::Ref<const Eigen::Vector3d>
     */
    Eigen::Ref<const Eigen::Vector3d> get_next_step_location() const
    {
        return next_step_location_;
    }

    /**
     * @brief Return the slack values from the last solution.
     **/
    const Eigen::Vector4d& get_slack_variables() const
    {
        return slack_variables_;
    }

    /**
     * @brief @copydoc DcmVrpPlanner::duration_before_step_landing_
     *
     * @return const double&
     */
    const double& get_duration_before_step_landing() const
    {
        return duration_before_step_landing_;
    }

    /**
     * @brief Get the desired com height.
     *
     * @return const double&
     */
    const double& get_com_height() const
    {
        return ht_;
    }

    /**
     * @brief return cost.
     *
     * @return double
     */
    double cost()
    {
        return (0.5 * x_opt_.transpose() * Q_ * x_opt_ +
                q_.transpose() * x_opt_)(0, 0);
    }

    /**
     * @brief add time equation for fixing tau.
     */
    void add_t_eq(double time)
    {
        A_eq_(2, 2) = 1;
        B_eq_(2) = exp(omega_ * time);
    }

    /*
     * Private methods
     */
private:
    /**
     * @brief Compute the nominal step location from the user input.
     *
     * @param is_left_leg_in_contact [in] is used notably to define which
     * surface to use for the next contact.
     * - 1 if left leg is in contact
     * - 2 if right leg is in contact
     * @param v_des_local in the local frame.
     */
    void compute_nominal_step_values(
        const bool& is_left_leg_in_contact,
        const Eigen::Ref<const Eigen::Vector3d>& v_des_local);

    /*
     * Attributes
     */
private:
    /** @brief Minimum step length in the x direction (in the direction of
     * forward motion). */
    double l_min_;

    /** @brief Maximum step length in the x direction (in the direction of
     * forward motion). */
    double l_max_;

    /** @brief Nominal step length in the x direction (in the direction of
     * forward motion). */
    double l_nom_;

    /** @brief Minimum step length in the y direction (in the lateral
     * direction).
     */
    double w_min_;

    /** @brief Maximum step length in the y direction (in the lateral
     * direction).
     */
    double w_max_;

    /** @brief Nominal step length in the y direction (in the lateral
     * direction).
     */
    double w_nom_;

    /** @brief Minimum step time. */
    double t_min_;

    /** @brief Maximum step time. */
    double t_max_;

    /** @brief Nominal step time. */
    double t_nom_ = 0.1;

    /** @brief Minimum step time in logarithmic scale:
     * \f$ e^{\omega t_{nom}} \f$*/
    double tau_min_;

    /** @brief Maximum step time in logarithmic scale:
     * \f$ e^{\omega t_{nom}} \f$*/
    double tau_max_;

    /** @brief Nominal step time in logarithmic scale:
     * \f$ e^{\omega t_{nom}} \f$*/
    double tau_nom_;

    /** @brief Default step width. */
    double l_p_;

    /** @brief Average desired height of the com above the ground. */
    double ht_;

    /** @brief Natural frequency of the pendulum: \f$ \omega =
     * \sqrt{\frac{g}{z_0}} \f$. */
    double omega_;

    /** @brief Maximum DCM offset along the X-axis. */
    double bx_max_;

    /** @brief Minimum DCM offset along the X-axis. */
    double bx_min_;

    /** @brief Nominal DCM offset along the Y-axis. */
    double bx_nom_;

    /** @brief Maximum DCM offset along the Y-axis. */
    double by_max_out_;

    /** @brief Minimum DCM offset along the Y-axis. */
    double by_max_in_;

    /** @brief Nominal DCM offset along the Y-axis. */
    double by_nom_;

    /** @brief SE3 position of the robot base in the world frame. */
    pinocchio::SE3 world_M_local_;

    /** @brief Current DCM computed from the CoM estimation. */
    Eigen::Vector3d dcm_local_;

    /** @brief Nominal DCM computed from the CoM estimation and nominal time. */
    Eigen::Vector3d dcm_nominal_;

    /** @brief Current DCM computed from the CoM estimation. */
    Eigen::Vector3d current_step_location_local_;

    /** @brief Current DCM computed from the CoM estimation. */
    Eigen::Vector3d v_des_local_;

    /** @brief Store the time from last step touchdown in order to stop
     * optimizing after t_min_ is passed. */
    double time_from_last_step_touchdown_;

    /*
     * Problem results
     */

    /** @brief Next step location expressed in the world frame. */
    Eigen::Vector3d next_step_location_;

    /** @brief Slack variable values corresponding to last solution. */
    Eigen::Vector4d slack_variables_;

    /** @brief Remaining time before the step landing. */
    double duration_before_step_landing_;

    /**
     * QP variables
     */

    /** @brief Number of variabes in the optimization problem. */
    int nb_var_;

    /** @brief Number of equality constraints in the optimization problem. */
    int nb_eq_;

    /** @brief Number of inequality in the optimization problem. */
    int nb_ineq_;

    /** @brief Quadratic program solver.
     *
     * This is an eigen wrapper around the quad_prog fortran solver.
     */
    Eigen::QuadProgDense qp_solver_;

    /** @brief Solution of the optimization problem.
     * @see DcmVrpPlanner::compute_adapted_step_locations. */
    Eigen::VectorXd x_opt_;

    /** @brief Lower Bound on the solution of the optimization problem.
     * @see DcmVrpPlanner::compute_adapted_step_locations. */
    Eigen::VectorXd x_opt_lb_;

    /** @brief Upper Bound on the solution of the optimization problem.
     * @see DcmVrpPlanner::compute_adapted_step_locations. */
    Eigen::VectorXd x_opt_ub_;

    /** @brief Quadratic term of the quadratic cost.
     * @see DcmVrpPlanner::compute_adapted_step_locations. */
    Eigen::MatrixXd Q_;

    /** @brief Cost weights expressed in the local frame. */
    Eigen::Vector9d cost_weights_local_;

    /** @brief Linear term of the quadratic cost.
     * @see DcmVrpPlanner::compute_adapted_step_locations. */
    Eigen::VectorXd q_;

    /** @brief Linear equality matrix.
     * @see DcmVrpPlanner::compute_adapted_step_locations. */
    Eigen::MatrixXd A_eq_;

    /** @brief Linear equality vector.
     * @see DcmVrpPlanner::compute_adapted_step_locations. */
    Eigen::VectorXd B_eq_;

    /** @brief Linear inequality matrix.
     * @see DcmVrpPlanner::compute_adapted_step_locations. */
    Eigen::MatrixXd A_ineq_;

    /** @brief Linear inequality vector.
     * @see DcmVrpPlanner::compute_adapted_step_locations. */
    Eigen::VectorXd B_ineq_;
};

}  // namespace reactive_planners


#pragma once

#include <crest/geometry/indexed_mesh.hpp>
#include <crest/geometry/mesh_algorithms.hpp>
#include <crest/geometry/biscale_mesh.hpp>
#include <crest/basis/quasi_interpolation.hpp>
#include <crest/basis/lagrange_basis2d.hpp>
#include <crest/basis/detail/homogenized_basis_detail.hpp>
#include <crest/util/eigen_extensions.hpp>
#include <crest/util/stat.hpp>
#include <crest/util/timer.hpp>

#include <Eigen/Sparse>
#include <Eigen/SparseLU>

#include <set>
#include <cassert>
#include <unordered_map>

#include <iostream>

namespace crest
{
    template <typename Scalar>
    class HomogenizedBasis : public Basis<Scalar, HomogenizedBasis<Scalar>>
    {
    public:
        explicit HomogenizedBasis(const BiscaleMesh<Scalar, int> & mesh,
                                  Eigen::SparseMatrix<double> weights);

        virtual std::vector<int> boundary_nodes() const override { return _mesh.coarse_mesh().boundary_vertices(); }

        virtual std::vector<int> interior_nodes() const override { return _mesh.coarse_mesh().compute_interior_vertices(); }

        virtual Assembly<Scalar> assemble() const override;

        virtual int num_dof() const override { return _mesh.coarse_mesh().num_vertices(); }

        template <typename Function2d>
        VectorX<Scalar> interpolate(const Function2d & f) const;

        template <typename Function2d>
        VectorX<Scalar> interpolate_boundary(const Function2d & f) const;

        template <int QuadStrength, typename Function2d>
        VectorX<Scalar> load(const Function2d & f) const;

        template <int QuadStrength, typename Function2d>
        Scalar error_l2(const Function2d & f, const VectorX<Scalar> & weights) const;

        template <int QuadStrength, typename Function2d_x, typename Function2d_y>
        Scalar error_h1_semi(const Function2d_x & f_x,
                             const Function2d_y & f_y,
                             const VectorX<Scalar> & weights) const;

        const Eigen::SparseMatrix<Scalar> & basis_weights() const { return _basis_weights; }

    private:
        Eigen::SparseMatrix<Scalar> _basis_weights;
        const BiscaleMesh<Scalar, int> & _mesh;
    };

    /**
     * Base class for solvers that compute localized correctors for a given BiscaleMesh.
     */
    template <typename Scalar>
    class CorrectorSolver
    {
    public:
        typedef Eigen::Matrix<Scalar, Eigen::Dynamic, Eigen::Dynamic> DenseMatrix;

        virtual ~CorrectorSolver() {}

        crest::HomogenizedBasis<Scalar> compute_basis(const BiscaleMesh<Scalar, int> & mesh,
                                                      unsigned int oversampling);

        Eigen::SparseMatrix<Scalar> compute_correctors(const BiscaleMesh<Scalar, int> & mesh,
                                                       unsigned int oversampling);

        virtual std::vector<Eigen::Triplet<Scalar>> resolve_element_correctors_for_patch(
                const BiscaleMesh<Scalar, int> & mesh,
                const std::vector<int> & coarse_patch_interior,
                const std::vector<int> & fine_patch_interior,
                const Eigen::SparseMatrix<Scalar> & global_coarse_stiffness,
                const Eigen::SparseMatrix<Scalar> & global_fine_stiffness,
                const Eigen::SparseMatrix<Scalar> & global_quasi_interpolator,
                int coarse_element) const = 0;

        const std::unordered_map<std::string, AccumulatedDensityHistogram> & stats() const;

        void set_iterative_tolerance(Scalar tol) { _tol = tol; }
        Scalar iterative_tolerance() const { return _tol; }

    protected:
        explicit CorrectorSolver<Scalar>() : _tol(Scalar(10) * std::numeric_limits<Scalar>::epsilon()) {}

        VectorX<Scalar> local_rhs(const BiscaleMesh<Scalar, int> & mesh,
                                  const std::vector<int> & fine_patch_interior,
                                  int coarse_element,
                                  int local_index) const;

    private:
        void set_stats(std::unordered_map<std::string, AccumulatedDensityHistogram> stats);
        std::unordered_map<std::string, AccumulatedDensityHistogram> _stats;
        Scalar _tol;
    };

    /**
     * The default corrector solver, which uses SparseLU to compute correctors. It is slow, but very robust.
     */
    template <typename Scalar>
    class SparseLuCorrectorSolver : public CorrectorSolver<Scalar>
    {
    public:
        virtual std::vector<Eigen::Triplet<Scalar>> resolve_element_correctors_for_patch(
                const BiscaleMesh<Scalar, int> & mesh,
                const std::vector<int> & coarse_patch_interior,
                const std::vector<int> & fine_patch_interior,
                const Eigen::SparseMatrix<Scalar> & global_coarse_stiffness,
                const Eigen::SparseMatrix<Scalar> & global_fine_stiffness,
                const Eigen::SparseMatrix<Scalar> & global_quasi_interpolator,
                int coarse_element) const override
        {
            (void) global_coarse_stiffness;

            const auto I_H = sparse_submatrix(global_quasi_interpolator,
                                              coarse_patch_interior,
                                              fine_patch_interior);
            const auto A = sparse_submatrix(global_fine_stiffness,
                                            fine_patch_interior,
                                            fine_patch_interior);

            std::vector<Eigen::Triplet<Scalar>> triplets;
            const auto C = detail::construct_saddle_point_problem(A, I_H);

            Eigen::SparseLU<Eigen::SparseMatrix<Scalar>> solver;
            solver.analyzePattern(C);
            solver.factorize(C);
            assert(solver.info() == Eigen::Success);

            for (int i = 0; i < 3; ++i)
            {
                const auto b_local = this->local_rhs(mesh, fine_patch_interior, coarse_element, i);

                VectorX<Scalar> c(C.rows());
                c << b_local, VectorX<Scalar>::Zero(I_H.rows());

                // Recall that the solution is of the form [x, kappa], where kappa is merely a Lagrange multiplier, so
                // we extract x as the corrector.
                const VectorX<Scalar> corrector = solver.solve(c).topRows(A.rows());

                assert(static_cast<size_t>(corrector.rows()) == fine_patch_interior.size());
                const auto global_index = mesh.coarse_mesh().elements()[coarse_element].vertex_indices[i];
                for (size_t k = 0; k < fine_patch_interior.size(); ++k)
                {
                    const auto component = corrector(k);
                    // Due to rounding issues, some components that should perhaps be zero in exact arithmetic
                    // may be non-zero, and as such we might end up with a denser matrix than we should actually
                    // have. To prevent this, we introduce a threshold which determines whether to keep the entry.
                    // TODO: Make this threshold configurable?
                    if (std::abs(component) > 1e-12)
                    {
                        triplets.emplace_back(Eigen::Triplet<Scalar>(global_index, fine_patch_interior[k], component));
                    }
                }
            }

            return triplets;
        }
    };

    template <typename Scalar>
    VectorX<Scalar> CorrectorSolver<Scalar>::local_rhs(const BiscaleMesh<Scalar, int> & mesh,
                                                       const std::vector<int> & fine_patch_interior,
                                                       int coarse_element,
                                                       int local_index) const
    {
        assert(local_index >= 0 && local_index < 3);
        // TODO: Simplify this function
        assert(std::is_sorted(fine_patch_interior.cbegin(), fine_patch_interior.cend()));

        const auto coarse_triangle = mesh.coarse_mesh().triangle_for(coarse_element);
        const Eigen::Matrix<Scalar, 3, 3> coarse_coeff = detail::basis_coefficients_for_triangle(coarse_triangle);

        const auto coarse_grad_x = coarse_coeff(0, local_index);
        const auto coarse_grad_y = coarse_coeff(1, local_index);

        VectorX<Scalar> rhs(fine_patch_interior.size());
        rhs.setZero();
        for (auto k : mesh.descendants_for(coarse_element))
        {
            const auto vertex_indices = mesh.fine_mesh().elements()[k].vertex_indices;
            const auto fine_triangle = mesh.fine_mesh().triangle_for(k);
            const Eigen::Matrix<Scalar, 3, 3> fine_coeff = detail::basis_coefficients_for_triangle(fine_triangle);

            for (size_t j = 0; j < 3; ++j)
            {
                const auto fine_grad_x = fine_coeff(0, j);
                const auto fine_grad_y = fine_coeff(1, j);

                // TODO: Fix this
                const auto product = [&] (auto  , auto  )
                {
                    return coarse_grad_x * fine_grad_x + coarse_grad_y * fine_grad_y;
                };

                // At this point, we only know the index of the vertex in the global mesh,
                // but we need the index of the vertex with respect to the patch interior. For now,
                // we just perform a binary search to recover it, though there may be much more efficient ways.
                // For example, we can construct a hashmap in the beginning of this function (which may
                // or may not be more efficient).
                const auto vertex_index = vertex_indices[j];
                const auto range = std::equal_range(fine_patch_interior.cbegin(),
                                                    fine_patch_interior.cend(),
                                                    vertex_index);

                if (range.first != range.second)
                {
                    const auto inner_product = triquad<2>(product,
                                                          fine_triangle.a,
                                                          fine_triangle.b,
                                                          fine_triangle.c);

                    const auto local_index = range.first - fine_patch_interior.cbegin();
                    rhs(local_index) += inner_product;
                }

            }
        }
        return rhs;
    }

    template <typename Scalar>
    Eigen::SparseMatrix<Scalar> CorrectorSolver<Scalar>::compute_correctors(
            const BiscaleMesh<Scalar, int> & mesh,
            unsigned int oversampling)
    {
        AccumulatedDensityHistogramBuilder timing_density;

        const auto I_H = quasi_interpolator(mesh);

        Eigen::SparseMatrix<Scalar> A_fine(mesh.fine_mesh().num_vertices(), mesh.fine_mesh().num_vertices());
        Eigen::SparseMatrix<Scalar> A_coarse(mesh.coarse_mesh().num_elements(), mesh.coarse_mesh().num_elements());
        {
            // Currently we needlessly construct the mass matrix here too. For now we put this in a block scope
            // so that it will be deallocated shortly thereafter, but in the long-term this should
            // be remedied so that we don't redundantly compute it. TODO
            const auto fine_assembly = LagrangeBasis2d<Scalar>(mesh.fine_mesh()).assemble();
            A_fine = std::move(fine_assembly.stiffness);

            const auto coarse_assembly = LagrangeBasis2d<Scalar>(mesh.coarse_mesh()).assemble();
            A_coarse = std::move(coarse_assembly.stiffness);
        }

        std::vector<Eigen::Triplet<Scalar>> basis_triplets;
        for (int coarse_element = 0; coarse_element < mesh.coarse_mesh().num_elements(); ++coarse_element)
        {
            Timer element_patch_timer;

            const auto coarse_patch = mesh.coarse_element_patch(coarse_element, oversampling);
            const auto coarse_patch_interior = coarse_patch.interior();
            const auto fine_patch = mesh.fine_patch_from_coarse(coarse_patch);
            const auto fine_patch_interior = fine_patch.interior();

            if (!fine_patch_interior.empty())
            {
                const auto corrector_contributions = resolve_element_correctors_for_patch(
                        mesh,
                        coarse_patch_interior,
                        fine_patch_interior,
                        A_coarse,
                        A_fine,
                        I_H,
                        coarse_element);

                std::copy(corrector_contributions.cbegin(),
                          corrector_contributions.cend(),
                          std::back_inserter(basis_triplets));
            }

            timing_density.add_sample(static_cast<double>(fine_patch_interior.size()), element_patch_timer.elapsed());
        }

        std::unordered_map<std::string, AccumulatedDensityHistogram> stats;
        stats["timing_distribution"] = timing_density.with_bin_count(100).build();
        set_stats(stats);

        Eigen::SparseMatrix<Scalar> basis(mesh.coarse_mesh().num_vertices(), mesh.fine_mesh().num_vertices());
        basis.setFromTriplets(basis_triplets.cbegin(), basis_triplets.cend());
        return basis;
    }

    template <typename Scalar>
    HomogenizedBasis<Scalar> CorrectorSolver<Scalar>::compute_basis(const BiscaleMesh<Scalar, int> & mesh,
                                                                    unsigned int oversampling)
    {
        const auto corrector_weights = compute_correctors(mesh, oversampling);
        const auto lagrange_basis_weights = detail::standard_coarse_basis_in_fine_space(mesh);
        const auto basis_weights = lagrange_basis_weights - corrector_weights;
        return HomogenizedBasis<Scalar>(mesh, basis_weights);
    }

    template <typename Scalar>
    const std::unordered_map<std::string, AccumulatedDensityHistogram> &
    CorrectorSolver<Scalar>::stats() const
    {
        return _stats;
    };

    template <typename Scalar>
    void CorrectorSolver<Scalar>::set_stats(std::unordered_map<std::string, AccumulatedDensityHistogram> stats)
    {
        _stats = stats;
    }

    template <typename Scalar>
    HomogenizedBasis<Scalar>::HomogenizedBasis(const BiscaleMesh<Scalar, int> & mesh,
                                               Eigen::SparseMatrix<double> weights)
            : _mesh(mesh)
    {
        if (weights.rows() == mesh.coarse_mesh().num_vertices() && weights.cols() == mesh.fine_mesh().num_vertices())
        {
            _basis_weights = std::move(weights);
        } else
        {
            throw std::invalid_argument("Dimensions of basis weights are not compatible with "
                                                "supplied coarse and fine meshes.");
        }
    }

    template <typename Scalar>
    Assembly<Scalar> HomogenizedBasis<Scalar>::assemble() const
    {
        LagrangeBasis2d<Scalar> fine_basis(_mesh.fine_mesh());
        const auto fine_assembly = fine_basis.assemble();

        const auto & M = fine_assembly.mass;
        const auto & A = fine_assembly.stiffness;
        const auto & W = _basis_weights;

        Assembly<Scalar> assembly;
        assembly.mass = W * M * W.transpose();
        assembly.stiffness = W * A * W.transpose();
        return assembly;
    }

    template <typename Scalar>
    template <typename Function2d>
    VectorX<Scalar> HomogenizedBasis<Scalar>::interpolate(const Function2d & f) const
    {
        // A slightly more accurate way of interpolating would perhaps be to interpolate the function
        // in the fine space, and then quasi-interpolate it to weights in the coarse space.
        // However, it seems this has some unfortunate effects on the implementation of inhomogeneous
        // Dirichlet boundary conditions?
        return LagrangeBasis2d<Scalar>(_mesh.coarse_mesh()).interpolate(f);
    }

    template <typename Scalar>
    template <typename Function2d>
    VectorX<Scalar> HomogenizedBasis<Scalar>::interpolate_boundary(const Function2d & f) const
    {
        return LagrangeBasis2d<Scalar>(_mesh.fine_mesh()).interpolate_boundary(f);
    }

    template <typename Scalar>
    template <int QuadStrength, typename Function2d>
    VectorX<Scalar> HomogenizedBasis<Scalar>::load(const Function2d & f) const
    {
        const LagrangeBasis2d<Scalar> fine_basis(_mesh.fine_mesh());
        const auto fine_load = fine_basis.template load<QuadStrength>(f);
        const auto & W = _basis_weights;
        return W * fine_load;
    };

    template <typename Scalar>
    template <int QuadStrength, typename Function2d>
    Scalar HomogenizedBasis<Scalar>::error_l2(const Function2d & f, const VectorX<Scalar> & weights) const
    {
        const LagrangeBasis2d<Scalar> fine_basis(_mesh.fine_mesh());
        const auto & W = _basis_weights;
        const VectorX<Scalar> fine_weights = W.transpose() * weights;
        return fine_basis.template error_l2<QuadStrength>(f, fine_weights);
    };

    template <typename Scalar>
    template <int QuadStrength, typename Function2d_x, typename Function2d_y>
    Scalar HomogenizedBasis<Scalar>::error_h1_semi(const Function2d_x & f_x,
                                                   const Function2d_y & f_y,
                                                   const VectorX<Scalar> & weights) const
    {
        const LagrangeBasis2d<Scalar> fine_basis(_mesh.fine_mesh());
        const auto & W = _basis_weights;
        const VectorX<Scalar> fine_weights = W.transpose() * weights;
        return fine_basis.template error_h1_semi<QuadStrength>(f_x, f_y, fine_weights);
    };
}