/*
 * H2L2.cpp
 *
 *  Created on: 13.03.2018
 *      Author: thies
 */

#include <base/Util.h>
#include <deal.II/base/exceptions.h>
#include <deal.II/base/logstream.h>
#include <deal.II/base/timer.h>
#include <deal.II/lac/sparse_direct.h>
#include <deal.II/lac/sparse_matrix.h>
#include <deal.II/lac/sparsity_pattern.h>
#include <deal.II/lac/vector.h>
#include <norms/H2L2.h>
#include <stddef.h>
#include <cmath>
#include <iostream>
#include <vector>

namespace wavepi {
namespace norms {

using namespace dealii;

inline double square(const double x) { return x * x; }
inline double pow4(const double x) { return x * x * x * x; }

template <int dim>
H2L2<dim>::H2L2(double alpha, double beta) : alpha_(alpha), beta_(beta) {}

template <int dim>
double H2L2<dim>::norm(const DiscretizedFunction<dim>& u) const {
  auto mesh = u.get_mesh();

  // we may be able to use v, but this might introduce inconsistencies in the adjoints
  // Note: this function works even for non-constant meshes.
  auto deriv = u.calculate_derivative();

  // using deriv.calculate_derivative feels wrong, better use a specialized formula.
  auto deriv2 = u.calculate_second_derivative();

  double result = 0;

  for (size_t i = 0; i < mesh->length(); i++) {
    double nrm2        = mesh->get_mass_matrix(i)->matrix_norm_square(u[i]);
    double nrm2_deriv  = mesh->get_mass_matrix(i)->matrix_norm_square(deriv[i]);
    double nrm2_deriv2 = mesh->get_mass_matrix(i)->matrix_norm_square(deriv2[i]);

    // + trapezoidal rule in time:
    if (i > 0)
      result += (nrm2 + alpha_ * nrm2_deriv + beta_ * nrm2_deriv2) / 2 *
                (std::abs(mesh->get_time(i) - mesh->get_time(i - 1)));

    if (i < mesh->length() - 1)
      result += (nrm2 + alpha_ * nrm2_deriv + beta_ * nrm2_deriv2) / 2 *
                (std::abs(mesh->get_time(i + 1) - mesh->get_time(i)));
  }

  return std::sqrt(result);
}

template <int dim>
double H2L2<dim>::dot(const DiscretizedFunction<dim>& u, const DiscretizedFunction<dim>& v) const {
  auto mesh     = u.get_mesh();
  double result = 0.0;

  // we may be able to use v, but this might introduce inconsistencies in the adjoints
  // Note: this function works even for non-constant meshes.
  auto deriv  = u.calculate_derivative();
  auto Vderiv = v.calculate_derivative();

  // using deriv.calculate_derivative feels wrong, better use a specialized formula.
  auto deriv2  = u.calculate_second_derivative();
  auto Vderiv2 = v.calculate_second_derivative();

  for (size_t i = 0; i < mesh->length(); i++) {
    double doti        = mesh->get_mass_matrix(i)->matrix_scalar_product(u[i], v[i]);
    double doti_deriv  = mesh->get_mass_matrix(i)->matrix_scalar_product(deriv[i], Vderiv[i]);
    double doti_deriv2 = mesh->get_mass_matrix(i)->matrix_scalar_product(deriv2[i], Vderiv2[i]);

    // + trapezoidal rule in time
    if (i > 0)
      result += (doti + alpha_ * doti_deriv + beta_ * doti_deriv2) / 2 *
                (std::abs(mesh->get_time(i) - mesh->get_time(i - 1)));

    if (i < mesh->length() - 1)
      result += (doti + alpha_ * doti_deriv + beta_ * doti_deriv2) / 2 *
                (std::abs(mesh->get_time(i + 1) - mesh->get_time(i)));
  }

  return result;
}

template <int dim>
void H2L2<dim>::dot_transform(DiscretizedFunction<dim>& u) {
  u.mult_mass();
  dot_solve_mass_and_transform(u);
}

template <int dim>
void H2L2<dim>::dot_transform_inverse(DiscretizedFunction<dim>& u) {
  u.solve_mass();
  dot_mult_mass_and_transform_inverse(u);
}

template <int dim>
void H2L2<dim>::dot_solve_mass_and_transform(DiscretizedFunction<dim>& u) {
  auto mesh = u.get_mesh();

  // X = (T + \alpha D^t T D + \beta D_2^t T D_2) * M,
  // M (blocks of mass matrices) is already taken care of, D = derivative, T = trapezoidal rule

  auto dx  = u.calculate_derivative();
  auto d2x = u.calculate_second_derivative();

  // trapezoidal rule
  // (has to happen between D and D^t for dx)
  for (size_t i = 0; i < mesh->length(); i++) {
    double factor = 0.0;

    if (i > 0) factor += std::abs(mesh->get_time(i) - mesh->get_time(i - 1)) / 2.0;
    if (i < mesh->length() - 1) factor += std::abs(mesh->get_time(i + 1) - mesh->get_time(i)) / 2.0;

    dx[i] *= factor;
    d2x[i] *= factor;
    u[i] *= factor;
  }

  auto dtdx   = dx.calculate_derivative_transpose();
  auto d2td2x = d2x.calculate_second_derivative_transpose();

  // add derivative terms
  u.add(alpha_, dtdx);
  u.add(beta_, d2td2x);
}

template <int dim>
void H2L2<dim>::factorize_matrix(std::shared_ptr<SpaceTimeMesh<dim>> mesh) {
  deallog << "factorizing matrix" << std::endl;

  SparsityPattern pattern(mesh->length(), mesh->length(), 5);

  for (size_t i = 0; i < 3; i++)
    for (size_t j = 0; j < 3; j++)
      pattern.add(i, j);

  for (size_t i = 3; i < mesh->length() - 3; i++) {
    // fill row i and column i
    for (int j = -2; j <= 2; j++) {
      pattern.add(i, i + j);
      pattern.add(i + j, i);
    }
  }

  for (size_t i = 0; i < 3; i++)
    for (size_t j = 0; j < 3; j++)
      pattern.add(mesh->length() - 1 - i, mesh->length() - 1 - j);

  pattern.compress();

  // coefficients of trapezoidal rule
  std::vector<double> lambdas(mesh->length(), 0.0);

  for (size_t i = 0; i < mesh->length(); i++) {
    if (i > 0) lambdas[i] += std::abs(mesh->get_time(i) - mesh->get_time(i - 1)) / 2.0;

    if (i < mesh->length() - 1) lambdas[i] += std::abs(mesh->get_time(i + 1) - mesh->get_time(i)) / 2.0;
  }

  SparseMatrix<double> matrix(pattern);

  double p20 = 1.0 * 16 / pow4(mesh->get_time(2) - mesh->get_time(0));
  double p10 = 1.0 / pow4(mesh->get_time(1) - mesh->get_time(0));
  double p31 = 1.0 * 16 / pow4(mesh->get_time(3) - mesh->get_time(1));
  double p42 = 1.0 * 16 / pow4(mesh->get_time(4) - mesh->get_time(2));

  matrix.set(0, 0, lambdas[0] * p10 + lambdas[1] * p20);
  matrix.set(1, 1, 4 * lambdas[0] * p10 + 4 * lambdas[1] * p20 + lambdas[2] * p31);
  matrix.set(2, 2, lambdas[0] * p10 + lambdas[1] * p20 + 4 * lambdas[2] * p31 + lambdas[3] * p42);

  matrix.set(0, 1, -2 * lambdas[0] * p10 - 2 * lambdas[1] * p20);
  matrix.set(1, 0, -2 * lambdas[0] * p10 - 2 * lambdas[1] * p20);

  matrix.set(0, 2, lambdas[0] * p10 + lambdas[1] * p20);
  matrix.set(2, 0, lambdas[0] * p10 + lambdas[1] * p20);

  matrix.set(1, 2, -2 * lambdas[0] * p10 - 2 * lambdas[1] * p20 - 2 * lambdas[2] * p31);
  matrix.set(2, 1, -2 * lambdas[0] * p10 - 2 * lambdas[1] * p20 - 2 * lambdas[2] * p31);

  for (size_t i = 3; i < mesh->length() - 3; i++) {
    // fill row i and column i

    double p20  = 1.0 * 16 / pow4(mesh->get_time(i + 2) - mesh->get_time(i));
    double p0m2 = 1.0 * 16 / pow4(mesh->get_time(i - 2) - mesh->get_time(i));
    double p1m1 = 1.0 * 16 / pow4(mesh->get_time(i + 1) - mesh->get_time(i - 1));

    matrix.set(i, i, lambdas[i + 1] * p20 + 4 * lambdas[i] * p1m1 + lambdas[i - 1] * p0m2);

    matrix.set(i, i - 1, -2 * lambdas[i - 1] * p0m2 - 2 * lambdas[i] * p1m1);
    matrix.set(i - 1, i, -2 * lambdas[i - 1] * p0m2 - 2 * lambdas[i] * p1m1);

    matrix.set(i, i + 1, -2 * lambdas[i + 1] * p20 - 2 * lambdas[i] * p1m1);
    matrix.set(i + 1, i, -2 * lambdas[i + 1] * p20 - 2 * lambdas[i] * p1m1);

    matrix.set(i, i + 2, lambdas[i + 1] * p20);
    matrix.set(i + 2, i, lambdas[i + 1] * p20);

    matrix.set(i, i - 2, lambdas[i - 1] * p0m2);
    matrix.set(i - 2, i, lambdas[i - 1] * p0m2);
  }

  // (symmetric to the first entries)
  size_t N = mesh->length() - 1;  // makes it easier to read

  p20 = 1.0 * 16 / pow4(mesh->get_time(N - 2) - mesh->get_time(N - 0));
  p10 = 1.0 / pow4(mesh->get_time(N - 1) - mesh->get_time(N - 0));
  p31 = 1.0 * 16 / pow4(mesh->get_time(N - 3) - mesh->get_time(N - 1));
  p42 = 1.0 * 16 / pow4(mesh->get_time(N - 4) - mesh->get_time(N - 2));

  matrix.set(N - 0, N - 0, lambdas[N - 0] * p10 + lambdas[N - 1] * p20);
  matrix.set(N - 1, N - 1, 4 * lambdas[N - 0] * p10 + 4 * lambdas[N - 1] * p20 + lambdas[N - 2] * p31);
  matrix.set(N - 2, N - 2,
             lambdas[N - 0] * p10 + lambdas[N - 1] * p20 + 4 * lambdas[N - 2] * p31 + lambdas[N - 3] * p42);

  matrix.set(N - 0, N - 1, -2 * lambdas[N - 0] * p10 - 2 * lambdas[N - 1] * p20);
  matrix.set(N - 1, N - 0, -2 * lambdas[N - 0] * p10 - 2 * lambdas[N - 1] * p20);

  matrix.set(N - 0, N - 2, lambdas[N - 0] * p10 + lambdas[N - 1] * p20);
  matrix.set(N - 2, N - 0, lambdas[N - 0] * p10 + lambdas[N - 1] * p20);

  matrix.set(N - 1, N - 2, -2 * lambdas[N - 0] * p10 - 2 * lambdas[N - 1] * p20 - 2 * lambdas[N - 2] * p31);
  matrix.set(N - 2, N - 1, -2 * lambdas[N - 0] * p10 - 2 * lambdas[N - 1] * p20 - 2 * lambdas[N - 2] * p31);

  matrix *= beta_;

  // H1 part (+ trapezoidal rule)
  SparseMatrix<double> matrixH1(pattern);

  double sq20 = 1.0 / square(mesh->get_time(2) - mesh->get_time(0));
  double sq10 = 1.0 / square(mesh->get_time(1) - mesh->get_time(0));
  double sq31 = 1.0 / square(mesh->get_time(3) - mesh->get_time(1));

  matrixH1.set(0, 0, lambdas[1] * sq20 + lambdas[0] * sq10);
  matrixH1.set(1, 1, lambdas[2] * sq31 + lambdas[0] * sq10);
  matrixH1.set(0, 1, -lambdas[0] * sq10);
  matrixH1.set(1, 0, -lambdas[0] * sq10);

  for (size_t i = 2; i < mesh->length() - 2; i++) {
    // fill row i and column i

    double sq20  = 1.0 / square(mesh->get_time(i + 2) - mesh->get_time(i));
    double sq0m2 = 1.0 / square(mesh->get_time(i) - mesh->get_time(i - 2));

    matrixH1.set(i, i, lambdas[i + 1] * sq20 + lambdas[i - 1] * sq0m2);

    matrixH1.set(i, i - 2, -lambdas[i - 1] * sq0m2);
    matrixH1.set(i - 2, i, -lambdas[i - 1] * sq0m2);

    matrixH1.set(i, i + 2, -lambdas[i + 1] * sq20);
    matrixH1.set(i + 2, i, -lambdas[i + 1] * sq20);
  }

  // (symmetric to the first entries)
  sq20 = 1.0 / square(mesh->get_time(N - 2) - mesh->get_time(N));
  sq10 = 1.0 / square(mesh->get_time(N - 1) - mesh->get_time(N));
  sq31 = 1.0 / square(mesh->get_time(N - 3) - mesh->get_time(N - 1));

  matrixH1.set(N, N - 0, lambdas[N - 1] * sq20 + lambdas[N] * sq10);
  matrixH1.set(N - 1, N - 1, lambdas[N - 2] * sq31 + lambdas[N] * sq10);
  matrixH1.set(N, N - 1, -lambdas[N] * sq10);
  matrixH1.set(N - 1, N, -lambdas[N] * sq10);

  matrix.add(alpha_, matrixH1);

  // L2 part (+ trapezoidal rule)
  for (size_t i = 0; i < mesh->length(); i++)
    matrix.add(i, i, lambdas[i]);

  umfpack.factorize(matrix);
}

template <int dim>
void H2L2<dim>::dot_mult_mass_and_transform_inverse(DiscretizedFunction<dim>& u) {
  auto mesh = u.get_mesh();

  LogStream::Prefix p("h2l2_transform");
  Timer timer;
  timer.start();

  if (umfpack.n() != mesh->length()) factorize_matrix(mesh);

  // just to be sure
  for (size_t i = 0; i < mesh->length(); i++)
    Assert(u[i].size() == u[0].size(), ExcInternalError());

  // solve for every DoF
  Vector<double> tmp(mesh->length());

  for (size_t i = 0; i < u[0].size(); i++) {
    for (size_t j = 0; j < mesh->length(); j++)
      tmp[j] = u[j][i];

    umfpack.solve(tmp);

    for (size_t j = 0; j < mesh->length(); j++)
      u[j][i] = tmp[j];
  }

  deallog << "solved in " << Util::format_duration(timer.wall_time()) << std::endl;
}

template <int dim>
std::string H2L2<dim>::name() const {
  return "H²([0,T], L²(Ω))";
}

template <int dim>
std::string H2L2<dim>::unique_id() const {
  return "H²([0,T], L²(Ω)) with α=" + std::to_string(alpha_) + ", β=" + std::to_string(beta_);
}

template class H2L2<1>;
template class H2L2<2>;
template class H2L2<3>;

} /* namespace norms */
} /* namespace wavepi */


// Software License for MTL
//
// Copyright (c) 2007 The Trustees of Indiana University.
//               2008 Dresden University of Technology and the Trustees of Indiana University.
//               2010 SimuNova UG (haftungsbeschränkt), www.simunova.com.
// All rights reserved.
// Authors: Peter Gottschling and Andrew Lumsdaine
//
// This file is part of the Matrix Template Library
//
// See also license.mtl.txt in the distribution.

// Written by Cornelius Steinhardt

#ifndef ITL_GMRES_INCLUDE
#define ITL_GMRES_INCLUDE

#include <algorithm>
#include <boost/numeric/mtl/concept/collection.hpp>
#include <boost/numeric/mtl/vector/dense_vector.hpp>
#include <boost/numeric/mtl/matrix/dense2D.hpp>
#include <boost/numeric/mtl/matrix/multi_vector.hpp>
#include <boost/numeric/mtl/operation/givens.hpp>
#include <boost/numeric/mtl/operation/two_norm.hpp>
#include <boost/numeric/mtl/utility/exception.hpp>
#include <boost/numeric/mtl/utility/irange.hpp>

namespace itl {

/// Generalized Minimal Residual method (without restart)
/** It computes at most kmax_in iterations (or size(x) depending on what is smaller)
    regardless on whether the termination criterion is reached or not.   **/
template < typename Matrix, typename Vector, typename LeftPreconditioner, typename RightPreconditioner, typename Iteration >
int gmres_full(const Matrix &A, Vector &x, const Vector &b,
               LeftPreconditioner &L, RightPreconditioner &R, Iteration& iter)
{
    using mtl::size; using mtl::irange; using mtl::iall; using std::abs; using std::sqrt;
    typedef typename mtl::Collection<Vector>::value_type Scalar;
    typedef typename mtl::Collection<Vector>::size_type  Size;

    if (size(b) == 0) throw mtl::logic_error("empty rhs vector");

    const Scalar                zero= math::zero(Scalar());
    Scalar                      rho, nu, hr;
    Size                        k, kmax(std::min(size(x), Size(iter.max_iterations() - iter.iterations())));
    Vector                      r0(b - A *x), r(solve(L,r0)), va(resource(x)), va0(resource(x)), va00(resource(x));
    mtl::matrix::multi_vector<Vector>   V(Vector(resource(x), zero), kmax+1);
    mtl::vector::dense_vector<Scalar>   s(kmax+1, zero), c(kmax+1, zero), g(kmax+1, zero), y(kmax, zero);  // replicated in distributed solvers
    mtl::matrix::dense2D<Scalar>        H(kmax+1, kmax);                                             // dito
    H= 0;

    rho= g[0]= two_norm(r);
    if (iter.finished(rho))
	return iter;
    V.vector(0)= r / rho;
    H= zero;

    // GMRES iteration
    for (k= 0; k < kmax ; ++k, ++iter) {
        va0= A * Vector(solve(R, V.vector(k)));
        V.vector(k+1)= va= solve(L,va0);
	// orth(V, V[k+1], false);
        // modified Gram Schmidt method
        for (Size j= 0; j < k+1; j++) {
	    H[j][k]= dot(V.vector(j), V.vector(k+1));
	    V.vector(k+1)-= H[j][k] * V.vector(j);
        }
        H[k+1][k]= two_norm(V.vector(k+1));
        //reorthogonalize
        for(Size j= 0; j < k+1; j++) {
	    hr= dot(V.vector(k+1), V.vector(j));
            H[j][k]+= hr;
            V.vector(k+1)-= hr * V.vector(j);
        }
        H[k+1][k]= two_norm(V.vector(k+1));
	if (H[k+1][k] != zero)                // watch for breakdown
            V.vector(k+1)*= 1. / H[k+1][k];

        // k Given's rotations
	for(Size i= 0; i < k; i++)
	    mtl::matrix::givens<mtl::matrix::dense2D<Scalar> >(H, H[i][k-1], H[i+1][k-1]).trafo(i);

       nu= sqrt(H[k][k]*H[k][k]+H[k+1][k]*H[k+1][k]);
       if(nu != zero){
            c[k]=  H[k][k]/nu;
            s[k]= -H[k+1][k]/nu;
            H[k][k]=c[k]*H[k][k]-s[k]*H[k+1][k];
            H[k+1][k]=0;
 	    mtl::vector::givens<mtl::vector::dense_vector<Scalar> >(g, c[k], s[k]).trafo(k);
        }
	rho= abs(g[k+1]);
    }

    //reduce k, to get regular matrix
    while (k > 0 && abs(g[k-1]<= iter.atol())) k--;

    // iteration is finished -> compute x: solve H*y=g as far as rank of H allows
    irange                  range(k);
    for (; !range.empty(); --range) {
	try {
	    y[range]= lu_solve(H[range][range], g[range]);
	} catch (mtl::matrix_singular) { continue; } // if singular then try with sub-matrix
	break;
    }

    if (range.finish() < k)
  	std::cerr << "GMRES orhogonalized with " << k << " vectors but matrix singular, can only use "
		  << range.finish() << " vectors!\n";
    if (range.empty())
        return iter.fail(2, "GMRES did not find any direction to correct x");
    x+= Vector(solve(R, Vector(V.vector(range)*y[range])));

    r= b - A*x;
    return iter.terminate(r);
}

/// Generalized Minimal Residual method with restart
template < typename Matrix, typename Vector, typename LeftPreconditioner,
           typename RightPreconditioner, typename Iteration >
int gmres(const Matrix &A, Vector &x, const Vector &b,
          LeftPreconditioner &L, RightPreconditioner &R,
	  Iteration& iter, typename mtl::Collection<Vector>::size_type restart)
{
     do {
	 Iteration inner(iter);
	 inner.set_max_iterations(std::min(int(iter.iterations()+restart), iter.max_iterations()));
	 inner.suppress_resume(true);
	 gmres_full(A, x, b, L, R, inner);
	 iter.update_progress(inner);
     } while (!iter.finished());

     return iter;
}

/// Solver class for GMRES; right preconditioner ignored (prints warning if not identity)
template < typename LinearOperator, typename Preconditioner= pc::identity<LinearOperator>,
	   typename RightPreconditioner= pc::identity<LinearOperator> >
class gmres_solver
{
  public:
    /// Construct solver from a linear operator; generate (left) preconditioner from it
    explicit gmres_solver(const LinearOperator& A, size_t restart= 8)
      : A(A), restart(restart), L(A), R(A) {}

    /// Construct solver from a linear operator and left preconditioner
    gmres_solver(const LinearOperator& A, size_t restart, const Preconditioner& L)
      : A(A), restart(restart), L(L), R(A) {}

    /// Construct solver from a linear operator and left preconditioner
    gmres_solver(const LinearOperator& A, size_t restart, const Preconditioner& L, const RightPreconditioner& R)
      : A(A), restart(restart), L(L), R(R) {}

    /// Solve linear system approximately as specified by \p iter
    template < typename HilbertSpaceB, typename HilbertSpaceX, typename Iteration >
    int solve(const HilbertSpaceB& b, HilbertSpaceX& x, Iteration& iter) const
    {
	return gmres(A, x, b, L, R, iter, restart);
    }

    /// Perform one GMRES iteration on linear system
    template < typename HilbertSpaceB, typename HilbertSpaceX >
    int solve(const HilbertSpaceB& b, HilbertSpaceX& x) const
    {
	itl::basic_iteration<double> iter(x, 1, 0, 0);
	return solve(b, x, iter);
    }

  private:
    const LinearOperator& A;
    size_t                restart;
    Preconditioner        L;
    RightPreconditioner   R;
};

} // namespace itl

#endif // ITL_GMRES_INCLUDE


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

/*
 Copyright (C) 2015 Peter Caspers

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

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

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

/*! \file qg1dlocalvolmodel.hpp
    \brief base class for one factor quasi gaussian models with local
           volatility
*/

#ifndef quantlib_quasigaussian1d_model_hpp
#define quantlib_quasigaussian1d_model_hpp

#include <ql/handle.hpp>
#include <ql/indexes/swapindex.hpp>
#include <ql/math/solvers1d/brent.hpp>
#include <ql/math/integrals/integral.hpp>
#include <ql/models/model.hpp>
#include <ql/termstructures/yieldtermstructure.hpp>

#include <boost/make_shared.hpp>
#include <boost/function.hpp>
#include <boost/bind.hpp>

// the integrator_ is used for h,G,yApprox,varSApprox => use different ones ?
// for the linear model h is overwritten

namespace QuantLib {

class Qg1dLocalVolModel : public TermStructureConsistentModel {
  public:
    /*! the model is specified by a function \kappa(t) and a function g(t,x,y),
        with \kappa(t) = -h'(t) / h(t), \sigma_f(t,T) = g(t)h(T), the HJM
        specification
        df(t,T) = \sigma_f(t,T) ( ( \int_t^T sigma_f(t,u) du ) du + dW(t) )
        and
        dx = (y - \kappa x) dt + \sigma_f(t,t) dW
        dy = (\sigma_f(t,t)^2 - 2 \kappa y) dt
        x(0) = y(0) = 0 */
    Qg1dLocalVolModel(const Handle<YieldTermStructure> &yts);

    /* core interface, these methods must be implemented by derived classes
       the other virtual methods may be overwritten by more efficient versions
       applicable to the particular model specification. */
    virtual Real kappa(const Real t) const = 0;
    virtual Real g(const Real t, const Real x, const Real y) const = 0;

    virtual Real h(const Real t) const;

    /*! \int_t^T h(s) ds / h(t) */
    virtual Real G(const Real t, const Real T) const;

    virtual Real sigma_f(const Real t, const Real T, const Real x,
                         const Real y) const;

    Real zerobond(const Real T, const Real t, const Real x, const Real y,
                  const Handle<YieldTermStructure> &yts =
                      Handle<YieldTermStructure>()) const;

    /*! swap rate is calculated with forward = discount, no indexed coupons
        T0 is the start date of the swap, fixedTimes the payment times of
        the fixed leg and taus the year fractions of the fixed leg */
    Real swapRate(const Real T0, const Real t,
                  const std::vector<Real> &fixedTimes,
                  const std::vector<Real> &taus, const Real x, const Real y,
                  const Handle<YieldTermStructure> &yts =
                      Handle<YieldTermStructure>()) const;

    Real dSwapRateDx(const Real T0, const Real t,
                     const std::vector<Real> &fixedTimes,
                     const std::vector<Real> &taus, const Real x, const Real y,
                     const Handle<YieldTermStructure> &yts =
                         Handle<YieldTermStructure>()) const;

    Real d2SwapRateDx2(const Real T0, const Real t,
                       const std::vector<Real> &fixedTimes,
                       const std::vector<Real> &taus, const Real x,
                       const Real y, const Handle<YieldTermStructure> &yts =
                                         Handle<YieldTermStructure>()) const;

    /*! local volatility using yApprox and sInvX (see Piterbarg, equation 13.19
        and what follows immediately after that), if numericalInversion is true,
        otherwise xi is used (Piterbarg, prop 13.1.8) */
    Real
    phi(const Real t, const Real s, const Real T0,
        const std::vector<Real> &fixedTimes, const std::vector<Real> &taus,
        const Handle<YieldTermStructure> &yts = Handle<YieldTermStructure>(),
        bool numericalInversion = false) const;

    Disposable<std::vector<Real> >
    phi(const Real t, const std::vector<Real> &s, const Real T0,
        const std::vector<Real> &fixedTimes, const std::vector<Real> &taus,
        const Handle<YieldTermStructure> &yts = Handle<YieldTermStructure>(),
        bool numericalInversion = false) const;

    /*! date based variants, only the forwarding curve from the swap index
        (if given) is used and no indexed coupons are used, see above */
    Real zerobond(const Date &maturiy, const Date &referenceDate, const Real x,
                  const Real y, const Handle<YieldTermStructure> &yts =
                                    Handle<YieldTermStructure>());

    Real swapRate(const Date &startDate, const Date &referenceDate,
                  const boost::shared_ptr<SwapIndex> &index,
                  const Period &tenor, const Real x, const Real y) const;

    Real dSwapRateDx(const Date &startDate, const Date &referenceDate,
                     const boost::shared_ptr<SwapIndex> &index,
                     const Period &tenor, const Real x, const Real y) const;

    Real d2SwapRateDx2(const Date &startDate, const Date &referenceDate,
                       const boost::shared_ptr<SwapIndex> &index,
                       const Period &tenor, const Real x, const Real y) const;

    /*! utilitiy function that fills T0, tau and a times vector based on
      a given swap index */
    void timesAndTaus(const Date &startDate,
                      const boost::shared_ptr<SwapIndex> &index,
                      const Period &tenor, Real &T0, std::vector<Real> &times,
                      std::vector<Real> &taus) const;

    virtual Real yApprox(const Real t) const;

    /*! this is \overline{x(t)} in Piterbarg */
    virtual Real xApprox(const Real t, const Real T0,
                         const std::vector<Real> &fixedTimes,
                         const std::vector<Real> &taus,
                         const Handle<YieldTermStructure> &yts) const;

    /*! approximate inversion of s in the sense of
        13.28 in Piterbarg */
    virtual Real xi(const Real t, const Real T0,
                    const std::vector<Real> &fixedTimes,
                    const std::vector<Real> &taus,
                    const Handle<YieldTermStructure> &yts, const Real s) const;

    virtual Disposable<std::vector<Real> >
    xi(const Real t, const Real T0, const std::vector<Real> &fixedTimes,
       const std::vector<Real> &taus, const Handle<YieldTermStructure> &yts,
       const std::vector<Real> &s) const;

    /*! numerical inversion of s with y = yApprox fixed,
        i.e. this is X(t,s) in Piterbarg's notation */
    Real sInvX(const Real t, const Real T0, const std::vector<Real> &fixedTimes,
               const std::vector<Real> &taus,
               const Handle<YieldTermStructure> &yts, const Real s) const;

    /*! Var(S(T)) approximation like in Piterbarg, remark 13.1.7 */
    Real varSApprox(const Real T, const Real T0,
                    const std::vector<Real> &fixedTimes,
                    const std::vector<Real> &taus,
                    const Handle<YieldTermStructure> &yts) const;

  protected:
    /*! compute swap rate, the first and second derivative w.r.t. x
        (since they share a lot of intermediate results this is more
        efficient than computing each single number) */
    void swapRate_d0_d1_d2(const Real T0, const Real t,
                           const std::vector<Real> &fixedTimes,
                           const std::vector<Real> &taus, const Real x,
                           const Real y, const Handle<YieldTermStructure> &yts,
                           Real &result_d0, Real &result_d1, Real &result_d2,
                           const bool compute_d0, const bool compute_d1,
                           const bool compute_d2) const;

    /*! sigma_f(t,t,0,0)^2*h(t)^{-2}, precondition (not checked) is t > 0 */
    virtual Real sigma_r_0_0_h_sqr(const Real t) const;

    /*! sigma_f(t,t,0.0)^2*dS/dx(s,0,0)^2 */
    virtual Real
    sigma_r_0_0_dSdx_sqr(const Real T0, const Real t,
                         const std::vector<Real> &fixedTimes,
                         const std::vector<Real> &taus,
                         const Handle<YieldTermStructure> &yts) const;

    boost::shared_ptr<Integrator> integrator_;

  private:
    Real sInvX_helper(const Real t, const Real T0,
                      const std::vector<Real> &fixedTimes,
                      const std::vector<Real> &taus,
                      const Handle<YieldTermStructure> &yts, const Real s,
                      const Real x) const;
};

// inline

inline Real Qg1dLocalVolModel::h(const Real t) const {
    return std::exp(-integrator_->operator()(
        boost::bind(&Qg1dLocalVolModel::kappa, this, _1), 0.0, t));
}

inline Real Qg1dLocalVolModel::G(const Real t, const Real T) const {
    return integrator_->operator()(boost::bind(&Qg1dLocalVolModel::h, this, _1),
                                   t, T);
}

inline Real Qg1dLocalVolModel::sigma_f(const Real t, const Real T, const Real x,
                                       const Real y) const {
    return g(t, x, y) * h(T);
}

inline Real Qg1dLocalVolModel::yApprox(const Real t) const {
    if (t < 1E-10)
        return 0.0;
    Real tmp = h(t);
    return tmp * tmp *
           integrator_->operator()(
               boost::bind(&Qg1dLocalVolModel::sigma_r_0_0_h_sqr, this, _1),
               0.0, t);
}

inline Real Qg1dLocalVolModel::sigma_r_0_0_h_sqr(const Real t) const {
    Real tmp = g(t, 0.0, 0.0); // this is sigma_f(t, t, 0.0, 0.0) / h(t);
    return tmp * tmp;
}

inline Real Qg1dLocalVolModel::sigma_r_0_0_dSdx_sqr(
    const Real T0, const Real t, const std::vector<Real> &fixedTimes,
    const std::vector<Real> &taus,
    const Handle<YieldTermStructure> &yts) const {
    Real tmp = sigma_f(t, t, 0.0, 0.0) *
               dSwapRateDx(T0, t, fixedTimes, taus, 0.0, 0.0, yts);
    return tmp * tmp;
}

inline Real Qg1dLocalVolModel::sInvX_helper(
    const Real t, const Real T0, const std::vector<Real> &fixedTimes,
    const std::vector<Real> &taus, const Handle<YieldTermStructure> &yts,
    const Real s, const Real x) const {
    Real y = yApprox(t);
    return swapRate(T0, t, fixedTimes, taus, x, y, yts) - s;
}

inline Real Qg1dLocalVolModel::sInvX(const Real t, const Real T0,
                                     const std::vector<Real> &fixedTimes,
                                     const std::vector<Real> &taus,
                                     const Handle<YieldTermStructure> &yts,
                                     const Real s) const {
    Brent b;
    boost::function<Real(Real)> f =
        boost::bind(&Qg1dLocalVolModel::sInvX_helper, this, t, T0, fixedTimes,
                    taus, yts, s, _1);
    return b.solve(f, 1E-7, 0.0, 0.01);
}

inline Real
Qg1dLocalVolModel::varSApprox(const Real T, const Real T0,
                              const std::vector<Real> &fixedTimes,
                              const std::vector<Real> &taus,
                              const Handle<YieldTermStructure> &yts) const {
    return integrator_->operator()(
        boost::bind(&Qg1dLocalVolModel::sigma_r_0_0_dSdx_sqr, this, T0, _1,
                    fixedTimes, taus, yts),
        0.0, T);
}

} // namespace QuantLib

#endif


#include <Engine/MeshEdit/Simulate.h>

#include <math.h>
#include <Eigen/Sparse>

using namespace Ubpa;

using namespace std;
using namespace Eigen;

void Simulate::Clear() {
	this->positions.clear();
	this->velocity.clear();
}

bool Simulate::Init() {
	//Clear();
	g = -9.8;
	stiff = 15000;
	//g = 0;
	isfast =true;
	x.resize(3 * positions.size());
	y_.resize(3 * positions.size());
	xx.resize(3 * positions.size());
	iter = 3;
	SetFix();

	this->velocity.resize(positions.size());
	for (int i = 0; i < positions.size(); i++)
	{
		for (int j = 0; j < 3; j++)
		{
			this->velocity[i][j] = 0;
			x[3 * i + j] = positions[i][j];
			y_[3 * i + j] = positions[i][j];
			xx[3 * i + j] = positions[i][j];
		}
	}

	mass.resize(3 * positions.size());
	M.resize(3 * positions.size(), 3 * positions.size());
	M.fill(0);

	force_ext.resize(3 * positions.size());
	force_int.resize(3 * positions.size());
	for (int i = 0; i < 3 * positions.size(); i++)
	{
		force_ext[i] = 0.0;
		force_int[i] = 0.0;
		mass[i] = 1;
		M(i, i) = mass[i];
	}

	//gx_m.resize(3 * positions.size(), 1);

	for (int i = 0; i < edgelist.size() / 2; i++)
	{
		size_t index1 = edgelist[2 * i];
		size_t index2 = edgelist[2 * i + 1];

		pointf3 v1 = positions[index1];
		pointf3 v2 = positions[index2];

		l.push_back((v1 - v2).norm());
		//l.push_back(0.5);
	}

	SetInitG();
	GetSet();

	CacK();

	CacA();

	return true;
}

void Simulate::SetFast()
{
	isfast = true;
}

bool Simulate::Run() {
	SimulateOnce();

	// half-edge structure -> triangle mesh

	return true;
}

void Ubpa::Simulate::SetLeftFix()
{
	//固定网格x坐标最小点
	fixed_id.clear();
	double x = 100000;
	for (int i = 0; i < positions.size(); i++)
	{
		if (positions[i][0] < x)
		{
			x = positions[i][0];
		}
	}

	for (int i = 0; i < positions.size(); i++)
	{
		if (abs(positions[i][0] - x) < 1e-5)
		{
			fixed_id.push_back(i);
		}
	}

	Init();
}

void Simulate::SimulateOnce() {
	// TODO
	//cout << "WARNING::Simulate::SimulateOnce:" << endl;
//		<< "\t" << "not implemented" << endl;
	//stiff = 50;
	//stiff = 100000;
//SetStiff(100000);

	if (!isfast)
	{
		CacX();
		//void UpdateX();
		CacV();
		UpdatePos();
	}
	else
	{
		std::vector<double> x_=x;
		for (int i = 0; i < x.size(); i++)
			y_ [i]= 2 * x[i]  - xx[i];
		x = y_;
		for (int i = 0; i < iter; i++)
		{
			Local_CacD();
			Global_CacX();
		}
		xx = x_;
		UpdatePos();
	}

}

void Simulate::CacForce()
{
	force_int.clear();
	force_int.resize(3 * positions.size());
	for (int i = 0; i < edgelist.size() / 2; i++)
	{
		size_t index1 = edgelist[2 * i];
		size_t index2 = edgelist[2 * i + 1];

		//pointf3 v1 = positions[index1];
		//pointf3 v2 = positions[index2];
		pointf3 v1 = pointf3(x[3 * index1 + 0], x[3 * index1 + 1], x[3 * index1 + 2]);
		pointf3 v2 = pointf3(x[3 * index2 + 0], x[3 * index2 + 1], x[3 * index2 + 2]);
		vecf3 r = v1 - v2;

		for (int j = 0; j < 3; j++)
		{
			force_int[3 * index1 + j] += -stiff * (r.norm() - l[i]) * r[j] / r.norm();
			force_int[3 * index2 + j] += stiff * (r.norm() - l[i]) * r[j] / r.norm();
		}

	}
}

void Simulate::SetInitG()
{
	for (int i = 0; i < positions.size(); i++)
		force_ext[i * 3 + 1] = g * mass[i];
}

void Simulate::GetSet()
{
	//fix.insert(10);
	//fix.insert(120);
	//fix.insert(440);
	//fix.insert(20);
	//fix.insert(1*positions.size() / 4);
	//fix.insert(2 * positions.size() / 4);
	fix = std::set<int>(fixed_id.begin(), fixed_id.end());
}

void Simulate::CacK()
{
	K.resize(x.size() - 3 * fix.size(), x.size());
	K.fill(0);

	for (int i = 0, j = 0; i < x.size(); i++)
	{
		if (fix.find(i / 3) == fix.end())
		{
			K(j, i) = 1;
			//std::cout << j << " " << i << endl;
			j++;
		}
	}

	Eigen::MatrixXd xt;
	xt.resize((x.size()), 1);
	b.resize(x.size());

	for (int i = 0; i < x.size(); i++)
		xt(i, 0) = x[i];
	Eigen::MatrixXd t = K.transpose() * K * xt;

	for (int i = 0; i < x.size(); i++)
		b[i] = x[i] - t(i, 0);
}

void Simulate::CacX()
{
	std::vector<double> y(x.size());
	for (int i = 0; i < x.size(); i++)
		y[i] = x[i] + h * velocity[i / 3][i % 3] + h * h / mass[i] * force_ext[i];

	xk = y;
	int i = 0;
	do
	{
		CacForce();
		CacGX();
		CacDiff();

		//Eigen::MatrixXd inverG = diff_.inverse();

		CacGxM();

		Eigen::MatrixXd t = inverG * gx_m;

		xk_1.clear();
		xk_1.resize(gx.size());
		for (int i = 0; i < gx.size(); i++)
		{
			xk_1[i] = xk[i] - t(i, 0);
			//cout << t(i, 0) << endl;
		}
		i++;
		//cout << i++ << endl;
		UpdateX();
		xk = x;
	} while (!isconv() && i <= 10);//sparse 50

}

void Simulate::CacGxM()
{
	gx_m.resize(gx.size(), 1);
	for (int i = 0; i < gx.size(); i++)
		gx_m(i, 0) = gx[i];
}

bool Simulate::isconv()
{
	bool flag = true;
	double delta = 0.01;
	std::vector<double> zero(gx.size(), delta);

	for (int i = 0; i < gx.size() && flag; i++)
	{
		if (abs(gx[i]) > zero[i])
			flag = false;
	}
	return flag;
}

void Simulate::SetFix()
{
	if (fixed_id.empty())
	{

		if (positions.size() > 440)
		{
			fixed_id.push_back(20);
			fixed_id.push_back(440);
		}
		else
		{
			fixed_id.push_back(10);
			fixed_id.push_back(120);
		}
	}
	//fixed_id.push_back(0);

}

void Simulate::CacGX()
{
	std::vector<double> y(xk.size());
	gx.resize(xk.size());

	for (int i = 0; i < xk.size(); i++)
	{
		y[i] = xk[i] + h * velocity[i / 3][i % 3] + h * h / mass[i] * force_ext[i];
	}

	for (int i = 0; i < xk.size(); i++)
	{
		gx[i] = mass[i] * (xk[i] - y[i]) - h * h * force_int[i];
	}

	Eigen::MatrixXd t(gx.size(), 1);
	for (int i = 0; i < gx.size(); i++)
		t(i, 0) = gx[i];

	t = K * t;

	gx.clear();
	gx.resize(t.rows());
	for (int i = 0; i < gx.size(); i++)
		gx[i] = t(i, 0);

	Eigen::MatrixXd xt;
	xt.resize((xk.size()), 1);

	for (int i = 0; i < xk.size(); i++)
		xt(i, 0) = xk[i];

	t = K * xt;
	xk.clear();
	xk.resize(t.rows());
	for (int i = 0; i < xk.size(); i++)
		xk[i] = t(i, 0);
}

void Simulate::CacDiff()
{
	std::vector<Eigen::Triplet<double> > triple;
	//Eigen::MatrixXd I = Eigen::DiagonalMatrix<double,3,3>::DiagonalMatrix(2);
	Eigen::MatrixXd I = MatrixXd::Identity(3, 3);
	for (int i = 0; i < edgelist.size() / 2; i++)
	{
		size_t index1 = edgelist[2 * i];
		size_t index2 = edgelist[2 * i + 1];

		pointf3 v1 = pointf3(x[3 * index1 + 0], x[3 * index1 + 1], x[3 * index1 + 2]);
		pointf3 v2 = pointf3(x[3 * index2 + 0], x[3 * index2 + 1], x[3 * index2 + 2]);
		vecf3 r = v1 - v2;
		MatrixXd t(3, 1);
		t(0, 0) = r[0]; t(1, 0) = r[1]; t(2, 0) = r[2];

		Eigen::MatrixXd dif;
		dif.resize(3, 3);
		dif = stiff * (l[i] / r.norm() - 1) * I - stiff * l[i]
			/ ((r.norm()) * (r.norm()) * (r.norm())) * t * t.transpose();
		//cout << index1 << " " << index2 << endl;
		//cout << dif << endl;
		for (int j = 0; j < 3; j++)
			for (int k = 0; k < 3; k++)
			{
				triple.push_back(Eigen::Triplet<double>(3 * index1 + j, 3 * index1 + k, -h * h * dif(j, k)));
				//triple.push_back(Eigen::Triplet<double>(3 * index1 + j, 3 * index2 + k, h*h*dif(j, k)));
				//triple.push_back(Eigen::Triplet<double>(3 * index2 + k, 3 * index1 + j, h*h*dif(j, k)));
				triple.push_back(Eigen::Triplet<double>(3 * index2 + k, 3 * index2 + j, -h * h * dif(j, k)));
			}
	}

	for (int i = 0; i < x.size(); i++)
	{
		triple.push_back(Eigen::Triplet<double>(i, i, mass[i]));
	}

	Eigen::MatrixXd diff_;
	Eigen::SparseLU<Eigen::SparseMatrix<double>> LU_;

	diff.setZero();
	diff.resize(x.size(), x.size());
	diff.setFromTriplets(triple.begin(), triple.end());

	diff_ = K * diff * K.transpose();
	//cout << diff_ << endl;
	//cout << K << endl;

	diff = diff_.sparseView();

	I = MatrixXd::Identity(gx.size(), gx.size());
	LU_.analyzePattern(diff);
	LU_.factorize(diff);
	//LU_.compute(diff);
	inverG = LU_.solve(I);
	//cout << inverG << endl;
}

void Simulate::UpdateX()
{
	Eigen::MatrixXd xt;
	xt.resize((xk_1.size()), 1);

	for (int i = 0; i < xk_1.size(); i++)
		xt(i, 0) = xk_1[i];

	Eigen::MatrixXd t = K.transpose() * xt;

	for (int i = 0; i < x.size(); i++)
		x[i] = t(i, 0) + b[i];

}

void Simulate::UpdatePos()
{
	for (int i = 0; i < x.size(); i++)
	{
		positions[i / 3][i % 3] = x[i];
	}
}

void Simulate::CacV()
{
	for (int i = 0; i < x.size(); i++)
	{
		velocity[i / 3][i % 3] = (x[i] - positions[i / 3][i % 3]) / h;
	}
}

void Simulate::CacL()
{
	L = MatrixXd::Zero(positions.size() * 3, positions.size() * 3);

	MatrixXd t = MatrixXd::Zero(positions.size(), positions.size());
	//cout << t << endl;
	for (int i = 0; i < edgelist.size() / 2; i++)
	{
		size_t index1 = edgelist[2 * i];
		size_t index2 = edgelist[2 * i + 1];

		std::vector<double> Ai(positions.size(),0);
		Ai[index1] = 1;
		Ai[index2] = -1;

		Eigen::MatrixXd At;
		At.resize((Ai.size()), 1);

		for (int j = 0; j < Ai.size(); j++)
			At(j, 0) = Ai[j];
		//cout << At << endl<<endl;
		t += stiff * At * At.transpose();
	}
	//cout << t << endl;
	//cout << M << endl;
	Matrix3d I3 = Matrix3d::Identity();
	for (int i = 0; i < positions.size(); i++)
	{
		for (int j = 0; j < positions.size(); j++)
			L.block(i * 3, j * 3, 3, 3) = t(i, j) * I3;
	}
}

void Simulate::CacJ()
{
	J = MatrixXd::Zero(positions.size() * 3, edgelist.size() / 2 * 3);

	MatrixXd t = MatrixXd::Zero(positions.size(), edgelist.size() / 2);

	for (int i = 0; i < edgelist.size() / 2; i++)
	{
		size_t index1 = edgelist[2 * i];
		size_t index2 = edgelist[2 * i + 1];

		std::vector<double> Ai(positions.size(), 0);
		Ai[index1] = 1;
		Ai[index2] = -1;

		std::vector<double> Si(edgelist.size() / 2,0);
		Si[i] = 1;

		Eigen::MatrixXd At;
		At.resize((Ai.size()), 1);

		for (int j = 0; j < Ai.size(); j++)
			At(j, 0) = Ai[j];

		Eigen::MatrixXd St;
		St.resize((Si.size()), 1);

		for (int j = 0; j < Si.size(); j++)
			St(j, 0) = Si[j];
		//cout << stiff * At * St.transpose() << endl;
		t += stiff * At * St.transpose();
	}
	//cout << t << endl;
	Matrix3d I3 = Matrix3d::Identity();
	for (int i = 0; i < positions.size(); i++)
	{
		for (int j = 0; j < edgelist.size() / 2; j++)
			J.block(i * 3, j * 3, 3, 3) = t(i, j) * I3;
	}
}

void Simulate::CacA()
{
	CacL();
	CacJ();
	MatrixXd A_;
	//cout << L << endl;
	//cout << J << endl;
	A_.resize(K.rows(), K.rows());
	A_ = K * (M + h * h * L) * K.transpose();
	A = A_.sparseView();
	LLT_.compute(A);
}

void Simulate::Global_CacX()
{
	Eigen::MatrixXd dt;
	dt.resize((d.size()), 1);

	for (int j = 0; j < d.size(); j++)
		dt(j, 0) = d[j];

	Eigen::MatrixXd yt;
	yt.resize((y_.size()), 1);

	for (int j = 0; j < y_.size(); j++)
		yt(j, 0) = y_[j];

	Eigen::MatrixXd ft;
	ft.resize((force_ext.size()), 1);

	for (int j = 0; j < force_ext.size(); j++)
		ft(j, 0) = force_ext[j];

	Eigen::MatrixXd bt;
	bt.resize((b.size()), 1);

	for (int j = 0; j < b.size(); j++)
		bt(j, 0) = b[j];

	VectorXd B = K * (h * h * J * dt + M * yt + h * h * ft - (M + h * h * L) * bt);
	//cout << J << endl;
	//cout << L << endl;
	VectorXd xf = LLT_.solve(B);
	//cout << xf << endl;
	xf = K.transpose() * xf + bt;

	for (int j = 0; j < x.size(); j++)
		x[j] = xf[j];
}

void Simulate::Local_CacD()
{
	d.resize(3 * edgelist.size() / 2);

	for (int i = 0; i < edgelist.size() / 2; i++)
	{
		size_t index1 = edgelist[2 * i];
		size_t index2 = edgelist[2 * i + 1];
		pointf3 v1 = pointf3(x[3 * index1 + 0], x[3 * index1 + 1], x[3 * index1 + 2]);
		pointf3 v2 = pointf3(x[3 * index2 + 0], x[3 * index2 + 1], x[3 * index2 + 2]);
		vecf3 r = v1 - v2;

		d[3 * i + 0] = l[i] * r[0] / r.norm();
		d[3 * i + 1] = l[i] * r[1] / r.norm();
		d[3 * i + 2] = l[i] * r[2] / r.norm();
	}
}


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

#ifndef SMOOTH__FEEDBACK__COLLOCATION_HPP_
#define SMOOTH__FEEDBACK__COLLOCATION_HPP_

#include <Eigen/Core>
#include <Eigen/LU>
#include <Eigen/Sparse>
#include <smooth/diff.hpp>
#include <smooth/internal/utils.hpp>
#include <smooth/polynomial/quadrature.hpp>

#include <cstddef>
#include <numeric>
#include <ranges>
#include <vector>

#include "traits.hpp"
#include "utils/sparse.hpp"

namespace smooth::feedback {

namespace detail {

/**
 * @brief Legendre-Gauss-Radau nodes including an extra node at +1.
 */
template<std::size_t K, std::size_t I = 8>
constexpr std::pair<std::array<double, K + 1>, std::array<double, K + 1>> lgr_plus_one()
{
  auto lgr_norm = ::smooth::lgr_nodes<K, I>();

  std::array<double, K + 1> ns, ws;
  for (auto i = 0u; i < K; ++i) {
    ns[i] = lgr_norm.first[i];
    ws[i] = lgr_norm.second[i];
  }
  ns[K] = 1;
  ws[K] = 0;
  return {ns, ws};
}

}  // namespace detail

/**
 * @brief Collocation mesh of interval [0, 1].
 * @tparam _Kmin minimal number of collocation points per interval
 * @tparam _Kmax maximal number of collocation points per interval
 *
 * [0, 1] is divided into non-overlapping intervals I_i, and each interval I_i has K_i LGR
 * collocation points.
 */
template<std::size_t _Kmin = 5, std::size_t _Kmax = 10>
  requires(_Kmin <= _Kmax)
class Mesh
{
  using MatMap = Eigen::Map<const Eigen::Matrix<double, -1, -1, Eigen::RowMajor>>;

public:
  /// @brief Minimal number of collocation points per interval
  static constexpr auto Kmin = _Kmin;
  /// @brief Maximal number of collocation points per interval
  static constexpr auto Kmax = _Kmax;

  /**
   * @brief Create a mesh consisting of a single interval [0, 1].
   *
   * @param Kmin minimal polynomial degree in mesh
   * @param Kmax maximal polynomial degree in mesh
   *
   * @note It must hold that kKmin <= Kmin <= Kmax <= kKmax, where kKmin and kKmax are compile-time
   * constants that define which LGR nodes to pre-compute.
   */
  inline Mesh() : intervals_(1, Interval{.K = Kmin, .tau0 = 0}) {}

  /**
   * @brief Number of intervals in mesh.
   */
  inline std::size_t N_ivals() const { return intervals_.size(); }

  /**
   * @brief Number of collocation points in mesh.
   */
  inline std::size_t N_colloc() const
  {
    return std::accumulate(
      intervals_.begin(), intervals_.end(), 0u, [](std::size_t curr, const auto & x) {
        return curr + x.K;
      });
  }

  /**
   * @brief Number of collocation points in interval i.
   *
   * @note This is also equal to the polynomial degree inside interval i, since the polynomial is
   * fitted with an "extra" point belonging to the subsequent interval.
   */
  inline std::size_t N_colloc_ival(std::size_t i) const { return intervals_[i].K; }

  /**
   * @breif Refine interval using the ph strategy.
   *
   * @param i index of interval to refine
   * @param D target number of collocation points in refined interval
   *
   * If D > Kmax, or current degree > Kmax    then the interval is divided into
   *                                          n = max(2, ceil(D / Kmin)) intervals with deg Kmin
   * If D < current degree,                   then nothing is done.
   * If D <= Kmax,                            then the polynomial degree is increased to D.
   */
  inline void refine_ph(std::size_t i, std::size_t D)
  {
    if (D > Kmax || intervals_[i].K > Kmax) {
      // refine by splitting interval into n intervals, each with degree Kmin_
      std::size_t n = std::max<std::size_t>(2u, (D + Kmin - 1) / Kmin);

      const double tau0 = intervals_[i].tau0;
      const double tauf = i + 1 < intervals_.size() ? intervals_[i + 1].tau0 : 1.;
      const double taum = (tauf - tau0) / n;

      while (n-- > 1) {
        intervals_.insert(intervals_.begin() + i + 1, Interval{.K = Kmin, .tau0 = tau0 + n * taum});
      }
    } else if (D < intervals_[i].K) {
      return;
    } else if (D <= Kmax) {
      // refine by increasing degree in interval
      intervals_[i].K = D;
    }
  }

  /**
   * @brief Set the number of collocation points in interval i to K
   * @param i interval index
   * @param K number of collocation points s.t. (Kmin <= K <= Kmax + 1)
   */
  inline void set_N_colloc_ival(std::size_t i, std::size_t K)
  {
    assert(Kmin <= K);
    assert(K <= Kmax + 1);
    intervals_[i].K = K;
  }

  /**
   * @brief Interval nodes and quadrature weights (DOES include extra point)
   */
  inline std::pair<Eigen::VectorXd, Eigen::VectorXd> interval_nodes_and_weights(std::size_t i) const
  {
    const std::size_t k = intervals_[i].K;

    Eigen::VectorXd ns, ws;
    utils::static_for<Kmax + 2 - Kmin>([&](auto i) {
      static constexpr auto K = Kmin + i;
      if (K == k) {
        static constexpr auto nw_ext_s = detail::lgr_plus_one<K>();
        ns = Eigen::Map<const Eigen::VectorXd>(nw_ext_s.first.data(), k + 1);
        ws = Eigen::Map<const Eigen::VectorXd>(nw_ext_s.second.data(), k + 1);
      }
    });

    const double tau0  = intervals_[i].tau0;
    const double tauf  = i + 1 < intervals_.size() ? intervals_[i + 1].tau0 : 1.;
    const double alpha = (tauf - tau0) / 2;

    return {
      Eigen::VectorXd::Constant(ns.size(), tau0) + alpha * (ns + Eigen::VectorXd::Ones(ns.size())),
      alpha * ws,
    };
  }

  /**
   * @brief All Mesh nodes and quadrature weights (DOES include extra point)
   */
  inline std::pair<Eigen::VectorXd, Eigen::VectorXd> all_nodes_and_weights() const
  {
    Eigen::VectorXd n(N_colloc() + 1), w(N_colloc() + 1);

    std::size_t cntr = 0;
    for (auto i = 0u; i < intervals_.size(); ++i) {
      auto [ni, wi] = interval_nodes_and_weights(i);

      const std::size_t Ni = ni.size();

      // exclude last point that belongs to next interval..
      n.segment(cntr, Ni - 1) = ni.head(Ni - 1);
      w.segment(cntr, Ni - 1) = wi.head(Ni - 1);

      cntr += Ni - 1;
    }

    n.tail(1).setConstant(1);
    w.tail(1).setConstant(0);

    return {n, w};
  }

  /**
   * @brief Interval differentiation matrix w.r.t. [0, 1] timescale.
   *
   * Returns a \f$ (K+1 \times K) \f$ matrix \f$ D \f$ s.t.
   * \f[
   *   \begin{bmatrix} y'(\tau_{i, 0}) & y'(\tau_{i, 1}) & \cdots & y'(\tau_{i, K-1}) \end{bmatrix}
   *  =
   *   \begin{bmatrix} y(\tau_{i, 0}) & y(\tau_{i, 1}) & \cdots & y(\tau_{i, K}) \end{bmatrix} D
   * \f],
   * where \f$ y(\cdot) \in \mathbb{R}^{d \times 1} \f$ is a Lagrange polynomial in interval i.
   */
  inline Eigen::MatrixXd interval_diffmat(std::size_t i) const
  {
    const std::size_t k = intervals_[i].K;

    const double tau0 = intervals_[i].tau0;
    const double tauf = i + 1 < intervals_.size() ? intervals_[i + 1].tau0 : 1.;

    Eigen::MatrixXd ret;
    utils::static_for<Kmax + 2 - Kmin>([&](auto i) {
      static constexpr auto K = Kmin + i;
      if (K == k) {
        static constexpr auto nw_ext_s = detail::lgr_plus_one<K>();
        static constexpr auto B_ext_s  = lagrange_basis<K>(nw_ext_s.first);
        static constexpr auto D_ext_s =
          polynomial_basis_derivatives<K, K + 1>(B_ext_s, nw_ext_s.first)
            .template block<K + 1, K>(0, 0);
        ret = MatMap(D_ext_s[0].data(), k + 1, k);
        ;
      }
    });

    return (2. / (tauf - tau0)) * ret;
  }

  /**
   * @brief Interval integration matrix w.r.t. [0, 1] timescale.
   *
   * Returns a \f$ (K \times K) \f$ matrix \f$ I \f$ s.t.
   * \f[
   *   \begin{bmatrix}
   *      y(\tau_{i, 1}) & y(\tau_{i, 2}) & \cdots & y(\tau_{i, K})
   *   \end{bmatrix}
   *  = y(\tau_{i, 0}) \begin{bmatrix} 1 & \ldots & 1 \end{bmatrix}
   *    + \begin{bmatrix}
   *        \dot y(\tau_{i, 0}) & \dot y(\tau_{i, 1}) & \cdots & \dot y(\tau_{i, K-1})
   *      \end{bmatrix} I
   * \f],
   * where \f$ y(\cdot) \in \mathbb{R}^{d \times 1} \f$ is a Lagrange
   * polynomial in interval i.
   */
  inline Eigen::MatrixXd interval_intmat(std::size_t i) const
  {
    const std::size_t k = intervals_[i].K;
    return interval_diffmat(i).block(1, 0, k, k).inverse();
  }

  /**
   * @brief Find interval index that contains t
   */
  inline std::size_t interval_find(double t) const
  {
    if (t < 0) { return 0; }
    if (t > 1) { return intervals_.size() - 1; }
    auto it = utils::binary_interval_search(
      intervals_, t, [](const auto & ival, double _t) { return ival.tau0 <=> _t; });
    if (it != intervals_.end()) { return std::distance(intervals_.begin(), it); }
    return 0;
  }

  /**
   * @brief Evaluate a function
   *
   * @tparam RetT return value type
   *
   * @param t time value in [0, 1]
   * @param r values for the collocation points (size N [extend=false] or N+1 [extend=true])
   * @param derivative to evaluate
   * @param extend set to true if a value is provided for t=+1
   */
  template<typename RetT, std::ranges::sized_range R>
  RetT eval(double t, const R & r, std::size_t p = 0, bool extend = true) const
  {
    [[maybe_unused]] const std::size_t N = N_colloc();

    if (extend) {
      assert(std::ranges::size(r) == N + 1);
    } else {
      assert(std::ranges::size(r) == N);
    }

    const std::size_t ival = interval_find(t);
    const std::size_t k    = intervals_[ival].K;

    const double tau0 = intervals_[ival].tau0;
    const double tauf = ival + 1 < intervals_.size() ? intervals_[ival + 1].tau0 : 1.;

    const double u = 2 * (t - tau0) / (tauf - tau0) - 1;

    Eigen::RowVectorXd W;

    utils::static_for<Kmax + 2 - Kmin>([&](auto i) {
      static constexpr auto K = Kmin + i;
      if (K == k) {
        if (extend || ival + 1 < intervals_.size()) {
          static constexpr auto nw_ext_s = detail::lgr_plus_one<K>();
          static constexpr auto B_ext_s  = lagrange_basis<K>(nw_ext_s.first);
          const auto U                   = monomial_derivative<K>(u, p);
          W = MatMap(U[0].data(), 1, k + 1) * MatMap(B_ext_s[0].data(), k + 1, k + 1);
          assert(std::size_t(W.size()) == k + 1);
        } else {
          static constexpr auto nw_s = lgr_nodes<K>();
          static constexpr auto B_s  = lagrange_basis<K - 1>(nw_s.first);
          const auto U               = monomial_derivative<K - 1>(u, p);
          W                          = MatMap(U[0].data(), 1, k) * MatMap(B_s[0].data(), k, k);
          assert(std::size_t(W.size()) == k);
        }
      }
    });

    using namespace std::views;

    std::size_t N_before = 0;
    for (auto i = 0u; i < ival; ++i) { N_before += intervals_[i].K; }
    const auto r_ival = r | drop(int64_t(N_before));
    RetT ret          = W(0) * *std::ranges::begin(r_ival);

    for (auto i = 1u; const auto & v : r_ival | drop(1) | take(W.size() - 1)) { ret += W(i++) * v; }
    return ret;
  }

private:
  struct Interval
  {
    /// @brief Polynomial degree in interval
    std::size_t K;
    /// @brief Start of interval on [0, 1] timescale
    double tau0;
  };

  /// @brief Mesh intervals
  std::vector<Interval> intervals_;
};

/// @brief MeshType is a specialization of Mesh
template<typename T>
concept MeshType = traits::is_specialization_of_sizet_v<T, Mesh>;

/**
 * @brief Evaluate a function on all collocation points.
 *
 * Returns a nf x N matrix
 *
 *  F= [ f(t_0, X_0, U_0)  f(t_1, X_1, u_1) ... f(t_{N-1}, X_{N-1}, U_{N-1})]
 *
 * with the function evaluated at all collocation points t_i in the Mesh m.
 *
 * @tparam Der return derivatives w.r.t variables
 *
 * @param nf dimensionality of f image
 * @param f function (t, X, U) -> R^nf
 * @param m Mesh of time
 * @param t0 initial time variable
 * @param tf final time variable
 * @param X state variables (size nx x N+1)
 * @param U input variables (size nu x N)
 *
 * @return If Deriv == false,
 * If Deriv == true, {F, dvecF_dt0, dvecF_dtf, dvecF_dvecX, dvecF_dvecU}, where vec(X) stacks the
 * columns of X into a single column vector.
 */
template<bool Deriv>
auto colloc_eval(
  const std::size_t nf,
  auto && f,
  const MeshType auto & m,
  const double t0,
  const double tf,
  const Eigen::MatrixXd & X,
  const Eigen::MatrixXd & U)
{
  assert(m.N_colloc() + 1 == static_cast<std::size_t>(X.cols()));  //  extra variable at the end
  assert(m.N_colloc() == static_cast<std::size_t>(U.cols()));  // one input per collocation point

  const std::size_t nx = X.rows();
  const std::size_t nu = U.rows();

  // all nodes in mesh
  const auto [tau_s, w_s] = m.all_nodes_and_weights();

  Eigen::MatrixXd Fval(nf, tau_s.size() - 1);

  Eigen::SparseMatrix<double> dvecF_dt0, dvecF_dtf, dvecF_dvecX, dvecF_dvecU;

  if constexpr (Deriv) {
    dvecF_dt0.resize(Fval.size(), 1);
    dvecF_dtf.resize(Fval.size(), 1);
    dvecF_dvecX.resize(Fval.size(), X.size());
    dvecF_dvecU.resize(Fval.size(), U.size());

    dvecF_dt0.reserve(Fval.size());
    dvecF_dtf.reserve(Fval.size());

    Eigen::VectorXi FX_pattern = Eigen::VectorXi::Constant(X.size(), nf);
    FX_pattern.tail(nx).setZero();
    dvecF_dvecX.reserve(FX_pattern);
    dvecF_dvecU.reserve(Eigen::VectorXi::Constant(U.size(), nf));
  }

  for (auto i = 0u; i + 1 < tau_s.size(); ++i) {
    const double T = t0 + (tf - t0) * tau_s(i);

    const Eigen::VectorXd x = X.col(i);
    const Eigen::VectorXd u = U.col(i);

    if constexpr (Deriv) {
      const auto [fval, dfval] = diff::dr(f, wrt(T, x, u));

      assert(fval.rows() == Eigen::Index(nf));
      assert(dfval.rows() == Eigen::Index(nf));
      assert(dfval.cols() == Eigen::Index(1 + nu + nx));

      Fval.col(i) = fval;

      for (auto row = 0u; row < nf; ++row) {
        dvecF_dt0.insert(nf * i + row, 0) = dfval(row, 0) * (1. - tau_s(i));
        dvecF_dtf.insert(nf * i + row, 0) = dfval(row, 0) * tau_s(i);
        for (auto col = 0u; col < nx; ++col) {
          dvecF_dvecX.insert(nf * i + row, i * nx + col) = dfval(row, 1 + col);
        }
        for (auto col = 0u; col < nu; ++col) {
          dvecF_dvecU.insert(nf * i + row, i * nu + col) = dfval(row, 1 + nx + col);
        }
      }
    } else {
      Fval.col(i) = f(T, x, u);
    }
  }

  if constexpr (Deriv) {
    dvecF_dt0.makeCompressed();
    dvecF_dtf.makeCompressed();
    dvecF_dvecX.makeCompressed();
    dvecF_dvecU.makeCompressed();

    return std::make_tuple(
      std::move(Fval),
      std::move(dvecF_dt0),
      std::move(dvecF_dtf),
      std::move(dvecF_dvecX),
      std::move(dvecF_dvecU));
  } else {
    return Fval;
  }
}

/**
 * @brief Evaluate a function at endpoints.
 *
 * Returns a nf vector
 *
 *  F = f(t_0, t_f, x_0, x_f, q)
 *
 * @tparam Der return derivatives w.r.t variables
 *
 * @param nf dimensionality of f image
 * @param nf state space degrees of freedom
 * @param f function (t0, tf, x0, xf, q) -> R^nf
 * @param t0 initial time
 * @param tf final time
 * @param X state variables (size nx x N+1)
 * @param q integrals
 *
 * @return If Deriv == false,
 * If Deriv == true, {F, dF_dt0, dF_dtf, dF_dvecX, dF_dQ},
 */
template<bool Deriv>
auto colloc_eval_endpt(
  const std::size_t nf,
  const std::size_t nx,
  auto && f,
  [[maybe_unused]] const double t0,
  const double tf,
  const Eigen::MatrixXd & X,
  const Eigen::VectorXd & Q)
{
  assert(static_cast<std::size_t>(X.rows()) == nx);

  // NOTE: for now t0 = 0 and we don't want t0 in signatures
  assert(t0 == 0);

  const Eigen::VectorXd x0 = X.leftCols(1);
  const Eigen::VectorXd xf = X.rightCols(1);

  if constexpr (!Deriv) {
    return f(tf, x0, xf, Q);
  } else {
    const auto [Fval, J] = diff::dr(f, wrt(tf, x0, xf, Q));

    assert(static_cast<std::size_t>(J.rows()) == nf);
    assert(static_cast<std::size_t>(J.cols()) == 1 + 2 * nx + Q.size());

    Eigen::SparseMatrix<double> dF_dt0, dF_dtf, dF_dvecX, dF_dQ;

    dF_dt0.resize(nf, 1);
    // dF_dt0.reserve(nf);
    // for (auto i = 0u; i < nf; ++i) { dF_dt0.insert(i, 0) = J(i, 0); }

    dF_dtf.resize(nf, 1);
    dF_dtf.reserve(nf);
    for (auto i = 0u; i < nf; ++i) { dF_dtf.insert(i, 0) = J(i, 0); }

    dF_dvecX.resize(nf, X.size());
    Eigen::VectorXi pattern = Eigen::VectorXi::Zero(X.size());
    pattern.head(nx).setConstant(nf);
    pattern.tail(nx).setConstant(nf);
    dF_dvecX.reserve(pattern);

    for (auto row = 0u; row < nf; ++row) {
      for (auto col = 0u; col < nx; ++col) {
        dF_dvecX.insert(row, col)                 = J(row, 1 + col);
        dF_dvecX.insert(row, X.size() - nx + col) = J(row, 1 + nx + col);
      }
    }

    dF_dQ.resize(nf, Q.size());
    dF_dQ.reserve(Eigen::VectorXi::Constant(Q.size(), nf));

    for (auto row = 0u; row < nf; ++row) {
      for (auto col = 0u; col < Q.size(); ++col) {
        dF_dQ.insert(row, col) = J(row, 1 + 2 * nx + col);
      }
    }

    dF_dt0.makeCompressed();
    dF_dtf.makeCompressed();
    dF_dvecX.makeCompressed();
    dF_dQ.makeCompressed();

    return std::make_tuple(Fval, dF_dt0, dF_dtf, dF_dvecX, dF_dQ);
  }
}

/**
 * @brief Evaluate dynamics constraint in all collocation points of a Mesh.
 *
 * @tparam Der return derivatives w.r.t variables
 *
 * @param nx state space degrees of freedom
 * @param f right-hand side of dynamics with signature (t, x, u) -> dx where x and dx are size nx
 * x 1 and u is size nu x 1
 * @param m mesh with a total of N collocation points
 * @param tf final time (variable of size 1)
 * @param x state values (variable of size nx x N+1)
 * @param u input values (variable of size nu x N)
 *
 * @return {F, dvecF_dt0, dvecF_dtf, dvecF_dvecX, dvecF_dvecU},
 * where vec(X) stacks the columns of X into a single column vector.
 */
template<bool Deriv>
auto colloc_dyn(
  const std::size_t nx,
  auto && f,
  const MeshType auto & m,
  const double t0,
  const double tf,
  const Eigen::MatrixXd & X,
  const Eigen::MatrixXd & U)
{
  assert(m.N_colloc() + 1 == static_cast<std::size_t>(X.cols()));  // extra at the end
  assert(m.N_colloc() == static_cast<std::size_t>(U.cols()));      // one per collocation point
  assert(nx == static_cast<std::size_t>(X.rows()));                // one per collocation point

  Eigen::MatrixXd Fval;
  Eigen::MatrixXd XD(nx, m.N_colloc());
  Eigen::SparseMatrix<double> dvecF_dt0, dvecF_dtf, dvecF_dvecX, dvecF_dvecU, dvecXD_dvecX;

  if constexpr (!Deriv) {
    Fval = colloc_eval<0>(nx, f, m, t0, tf, X, U);
  } else {
    std::tie(Fval, dvecF_dt0, dvecF_dtf, dvecF_dvecX, dvecF_dvecU) =
      colloc_eval<1>(nx, f, m, t0, tf, X, U);

    dvecXD_dvecX.resize(XD.size(), X.size());

    // reserve sparsity pattern
    Eigen::VectorXi pattern = Eigen::VectorXi::Zero(X.size());
    for (auto M = 0u, i = 0u; i < m.N_ivals(); ++i) {
      const std::size_t K = m.N_colloc_ival(i);
      pattern.segment(M, (K + 1) * nx) += Eigen::VectorXi::Constant((K + 1) * nx, K);
      M += K * nx;
    }
    dvecXD_dvecX.reserve(pattern);
  }

  for (auto i = 0u, M = 0u; i < m.N_ivals(); M += m.N_colloc_ival(i), ++i) {
    const std::size_t K     = m.N_colloc_ival(i);
    const Eigen::MatrixXd D = m.interval_diffmat(i);
    XD.block(0, M, nx, K)   = X.block(0, M, nx, K + 1) * D;

    if constexpr (Deriv) {
      // vec(X * D) = kron(D', I) * vec(X), so derivative w.r.t vec(X) = kron(D', I)
      for (auto i = 0u; i < K; ++i) {
        for (auto j = 0u; j < K + 1; ++j) {
          for (auto diag = 0u; diag < nx; ++diag) {
            dvecXD_dvecX.coeffRef(M * nx + i * nx + diag, M * nx + j * nx + diag) += D(j, i);
          }
        }
      }
    }
  }

  Eigen::VectorXd Fv = (XD - (tf - t0) * Fval).reshaped();

  // scale equalities by by quadrature weights
  const auto N      = m.N_colloc();
  const auto [n, w] = m.all_nodes_and_weights();

  // vec(A * W) = kron(W', I) * vec(A), so we apply kron(W', I) on the left

  Eigen::SparseMatrix<double> W(N, N);
  W.reserve(Eigen::VectorXi::Ones(N));
  for (auto i = 0u; i < N; ++i) { W.insert(i, i) = w(i); }

  const Eigen::SparseMatrix<double> W_kron_I = kron_identity(W, nx);

  Fv.applyOnTheLeft(W_kron_I);

  if constexpr (!Deriv) {
    return Fv;
  } else {
    dvecXD_dvecX.makeCompressed();

    Eigen::SparseMatrix<double> dF_dt0 = -(tf - t0) * dvecF_dt0;
    dF_dt0 += Fval.reshaped().sparseView();  // OK since dvecF_dtf is dense
    dF_dt0 = W_kron_I * dF_dt0;

    Eigen::SparseMatrix<double> dF_dtf = -(tf - t0) * dvecF_dtf;
    dF_dtf -= Fval.reshaped().sparseView();  // OK since dvecF_dtf is dense
    dF_dtf = W_kron_I * dF_dtf;

    Eigen::SparseMatrix<double> dF_dvecX = dvecXD_dvecX;
    dF_dvecX -= (tf - t0) * dvecF_dvecX;
    dF_dvecX = W_kron_I * dF_dvecX;

    Eigen::SparseMatrix<double> dF_dvecU = -(tf - t0) * W_kron_I * dvecF_dvecU;

    dF_dt0.makeCompressed();
    dF_dtf.makeCompressed();
    dF_dvecX.makeCompressed();
    dF_dvecU.makeCompressed();

    return std::make_tuple(
      std::move(Fv),
      std::move(dF_dt0),
      std::move(dF_dtf),
      std::move(dF_dvecX),
      std::move(dF_dvecU));
  }
}

/**
 * @brief Calculate relative dynamics errors for each interval in mesh.
 *
 * @param nx state space dimension
 * @param f dynamics function
 * @param m Mesh
 * @param t0 initial time variable
 * @param tf final time variable
 * @param x state trajectory
 * @param u input trajectory
 *
 * @return vector with relative errors for every interval in m
 */
Eigen::VectorXd mesh_dyn_error(
  const std::size_t nx,
  auto && f,
  const MeshType auto & m,
  const double t0,
  const double tf,
  const std::function<Eigen::VectorXd(double)> xfun,
  const std::function<Eigen::VectorXd(double)> ufun)
{
  const auto N = m.N_ivals();

  // create a new mesh where each interval is extended
  Mesh mext = m;
  for (auto i = 0u; i < N; ++i) {
    const std::size_t K = m.N_colloc_ival(i);
    mext.set_N_colloc_ival(i, K + 1);
  }

  Eigen::VectorXd ival_errs(N);

  // for each interval
  for (auto i = 0u, M = 0u; i < N; M += m.N_colloc_ival(i), ++i) {
    const std::size_t Kext = mext.N_colloc_ival(i);

    const auto [tau_s, weights] = mext.interval_nodes_and_weights(i);

    assert(std::size_t(tau_s.size()) == Kext + 1);

    // evaluate X and F at those points
    Eigen::MatrixXd Fval(nx, Kext + 1);
    Eigen::MatrixXd Xval(nx, Kext + 1);
    for (auto j = 0u; j < Kext + 1; ++j) {
      const double tj = t0 + (tf - t0) * tau_s(j);

      // evaluate x and u values at tj using current degree polynomials
      const auto Xj = xfun(tj);
      const auto Uj = ufun(tj);

      // evaluate right-hand side of dynamics at tj
      Fval.col(j) = f(tj, Xj, Uj);

      // store x values for later comparison
      Xval.col(j) = Xj;
    }

    // "integrate" system inside interval
    const Eigen::MatrixXd Xval_est =
      Xval.col(0).replicate(1, Kext) + (tf - t0) * Fval.leftCols(Kext) * mext.interval_intmat(i);

    // absolute error in interval
    Eigen::VectorXd e_abs = (Xval_est - Xval.rightCols(Kext)).colwise().norm();
    Eigen::VectorXd e_rel = e_abs / (1. + Xval.rightCols(Kext).colwise().norm().maxCoeff());

    // mex relative error on interval
    ival_errs(i) = e_rel.maxCoeff();
  }

  return ival_errs;
}

/**
 * @brief Refine intervals in mesh to satisfy a target error criterion.
 * @param[in, out] m mesh to refine
 * @param[in] errs relative errors for all intervals (@see mesh_dyn_error())
 * @param[in] target_err target relative error
 */
void mesh_refine(MeshType auto & m, const Eigen::VectorXd & errs, const double target_err)
{
  const auto N = m.N_ivals();

  assert(N == std::size_t(errs.size()));

  for (auto i = 0u; i < N; ++i) {
    const auto Nmi = N - 1 - i;
    const auto Ki  = m.N_colloc_ival(Nmi);

    if (errs(Nmi) > target_err) {
      const auto Ktarget = Ki + std::lround(std::log(errs(Nmi) / target_err) / std::log(Ki) + 1);
      m.refine_ph(Nmi, Ktarget);
    }
  }
}

/**
 * @brief Evaluate integral constraint on Mesh.
 *
 * @tparam Der return derivatives w.r.t variables
 *
 * @param nq number of integrals
 * @param g integrand with signature (t, x, u) -> R^{nq} where x is size nx x 1 and u is size nu x
 * 1
 * @param m mesh
 * @param t0 initial time (variable of size 1)
 * @param tf final time (variable of size 1)
 * @param I values (variable of size nq)
 * @param X state values (variable of size nx x N+1)
 * @param U input values (variable of size nu x N)
 *
 * @return {G, dvecG_dt0, dvecG_dtf, dvecG_dvecX, dvecG_dvecU},
 * where vec(X) stacks the columns of X into a single column vector.
 */
template<bool Deriv>
auto colloc_int(
  const std::size_t nq,
  auto && g,
  const MeshType auto & m,
  const double t0,
  const double tf,
  const Eigen::VectorXd & I,
  const Eigen::MatrixXd & X,
  const Eigen::MatrixXd & U)
{
  assert(static_cast<std::size_t>(I.size()) == nq);

  const std::size_t N = m.N_colloc();

  const auto [n, w] = m.all_nodes_and_weights();

  if constexpr (Deriv == false) {
    const auto Gv              = colloc_eval<Deriv>(nq, g, m, t0, tf, X, U);
    const Eigen::VectorXd Iest = Gv * w.head(N);
    Eigen::VectorXd Rv         = (tf - t0) * Iest - I;
    return Rv;
  } else {
    const auto [Gv, dvecG_dt0, dvecG_dtf, dvecG_dvecX, dvecG_dvecU] =
      colloc_eval<Deriv>(nq, g, m, t0, tf, X, U);
    const Eigen::VectorXd Iest = Gv * w.head(N);

    Eigen::VectorXd Rv = (tf - t0) * Iest - I;

    const Eigen::SparseMatrix<double> w_kron_I =
      (tf - t0) * kron_identity(w.head(N).transpose(), nq);

    Eigen::SparseMatrix<double> dR_dt0 = w_kron_I * dvecG_dt0;
    for (auto i = 0u; i < Iest.size(); ++i) { dR_dt0.coeffRef(i, 0) -= Iest(i); }

    Eigen::SparseMatrix<double> dR_dtf = w_kron_I * dvecG_dtf;
    for (auto i = 0u; i < Iest.size(); ++i) { dR_dtf.coeffRef(i, 0) += Iest(i); }

    Eigen::SparseMatrix<double> dR_dvecI = -sparse_identity(nq);

    Eigen::SparseMatrix<double> dR_dvecX = w_kron_I * dvecG_dvecX;

    Eigen::SparseMatrix<double> dR_dvecU = w_kron_I * dvecG_dvecU;

    dR_dt0.makeCompressed();
    dR_dtf.makeCompressed();
    dR_dvecI.makeCompressed();
    dR_dvecX.makeCompressed();
    dR_dvecU.makeCompressed();

    return std::make_tuple(
      std::move(Rv),
      std::move(dR_dt0),
      std::move(dR_dtf),
      std::move(dR_dvecI),
      std::move(dR_dvecX),
      std::move(dR_dvecU));
  }
}

}  // namespace smooth::feedback

#endif  // SMOOTH__FEEDBACK__COLLOCATION_HPP_


//==============================================================================
//         Copyright 2003 - 2011 LASMEA UMR 6602 CNRS/Univ. Clermont II
//         Copyright 2009 - 2011 LRI    UMR 8623 CNRS/Univ Paris Sud XI
//
//          Distributed under the Boost Software License, Version 1.0.
//                 See accompanying file LICENSE.txt or copy at
//                     http://www.boost.org/LICENSE_1_0.txt
//==============================================================================
#ifndef NT2_EXPONENTIAL_FUNCTIONS_COMPLEX_GENERIC_POW_HPP_INCLUDED
#define NT2_EXPONENTIAL_FUNCTIONS_COMPLEX_GENERIC_POW_HPP_INCLUDED
#include <nt2/exponential/functions/pow.hpp>
#include <nt2/include/functions/pow.hpp>
#include <nt2/include/functions/abs.hpp>
#include <nt2/include/functions/sincos.hpp>
#include <nt2/include/functions/log.hpp>
#include <nt2/include/functions/exp.hpp>
#include <nt2/include/functions/seladd.hpp>
#include <nt2/include/functions/if_else.hpp>
#include <nt2/include/functions/is_real.hpp>
#include <nt2/include/functions/real.hpp>
#include <nt2/include/functions/imag.hpp>
#include <nt2/include/functions/arg.hpp>
#include <nt2/include/functions/logical_not.hpp>
#include <boost/dispatch/meta/as_floating.hpp>
#include <nt2/sdk/complex/meta/as_complex.hpp>
#include <nt2/sdk/complex/meta/as_real.hpp>
#include <nt2/sdk/complex/meta/as_dry.hpp>
#include <nt2/sdk/simd/logical.hpp>

namespace nt2 { namespace ext
{
  NT2_FUNCTOR_IMPLEMENTATION( nt2::tag::pow_, tag::cpu_
                              , (A0)
                            , (generic_< complex_<floating_<A0> > >)
                              (generic_< complex_<floating_<A0> > >)
                            )
  {
    typedef typename meta::as_complex<A0>::type result_type;
    NT2_FUNCTOR_CALL_REPEAT(2)
    {
      return exp(a1*log(a0));
    }
  };

  NT2_FUNCTOR_IMPLEMENTATION( nt2::tag::pow_, tag::cpu_
                              , (A0)(A1)
                              , (generic_< complex_<floating_<A0> > >)
                                (generic_< floating_<A1> >)
                            )
  {
    typedef typename meta::as_complex<A0>::type result_type;
    NT2_FUNCTOR_CALL(2)
    {
      typedef typename meta::as_real<result_type>::type rtype;
      typedef typename meta::as_imaginary< rtype>::type itype;
      rtype t = nt2::arg(a0);
      rtype a = nt2::abs(a0);
      return nt2::pow(a, a1)*nt2::exp(itype(t*a1));
    }
  };

  NT2_FUNCTOR_IMPLEMENTATION( nt2::tag::pow_, tag::cpu_
                              , (A0)(A1)
                              , (generic_< complex_<floating_<A0> > >)
                                (generic_< dry_<floating_<A1> > >)
                            )
  {
    typedef typename meta::as_complex<A0>::type result_type;
    NT2_FUNCTOR_CALL(2)
    {
      return pow(a0, nt2::real(a1));
    }
  };

  NT2_FUNCTOR_IMPLEMENTATION( nt2::tag::pow_, tag::cpu_
                              , (A0)(A1)
                              , (generic_< floating_<A0> >)
                                (generic_< complex_<floating_<A1> > >)
                            )
  {
    typedef typename meta::as_complex<A0>::type result_type;
    NT2_FUNCTOR_CALL(2)
      {
        typedef typename meta::as_dry<A0>::type dtype;
        return nt2::exp(a1*nt2::log(dtype(a0)));
      }
  };

  NT2_FUNCTOR_IMPLEMENTATION( nt2::tag::pow_, tag::cpu_
                              , (A0)(A1)
                              , (generic_< dry_ < floating_<A0> > > )
                                (generic_< complex_<floating_<A1> > >)
                            )
  {
    typedef typename meta::as_complex<A0>::type result_type;
    NT2_FUNCTOR_CALL(2)
      {
        return nt2::exp(a1*nt2::log(a0));
      }
  };

  NT2_FUNCTOR_IMPLEMENTATION( nt2::tag::pow_, tag::cpu_
                              , (A0)(A1)
                              , (generic_< complex_<floating_<A0> > >)
                              (generic_< imaginary_<floating_<A1> > >)
                              )
  {
    typedef typename meta::as_complex<A0>::type result_type;
    NT2_FUNCTOR_CALL(2)
      {
        return nt2::exp(a1*nt2::log(a0));
      }
  };

  NT2_FUNCTOR_IMPLEMENTATION( nt2::tag::pow_, tag::cpu_
                              , (A0)(A1)
                              , (generic_< imaginary_<floating_<A0> > >)
                              (generic_< complex_<floating_<A1> > >)
                              )
  {
    typedef typename meta::as_complex<A0>::type result_type;
    NT2_FUNCTOR_CALL(2)
      {
        return nt2::exp(a1*nt2::log(a0));
      }
  };

  NT2_FUNCTOR_IMPLEMENTATION( nt2::tag::pow_, tag::cpu_
                              , (A0)(A1)
                              , (generic_< imaginary_<floating_<A0> > >)
                              (generic_< floating_<A1> >)
                              )
  {
    typedef typename meta::as_complex<A0>::type result_type;
    NT2_FUNCTOR_CALL(2)
      {
        return nt2::exp(a1*nt2::log(a0));
      }
  };

  NT2_FUNCTOR_IMPLEMENTATION( nt2::tag::pow_, tag::cpu_
                              , (A0)(A1)
                              , (generic_< imaginary_<floating_<A0> > >)
                              (generic_< dry_ < floating_<A1> > >)
                              )
  {
    typedef typename meta::as_complex<A0>::type result_type;
    NT2_FUNCTOR_CALL(2)
      {
        return nt2::exp(a1*nt2::log(a0));
      }
  };

  NT2_FUNCTOR_IMPLEMENTATION( nt2::tag::pow_, tag::cpu_
                              , (A0)(A1)
                              ,  (generic_< floating_<A0> >)
                              (generic_< imaginary_<floating_<A1> > >)
                              )
  {
    typedef typename meta::as_complex<A0>::type result_type;
    typedef typename meta::as_dry<A0>::type dtype;
    NT2_FUNCTOR_CALL(2)
      {
        return nt2::exp(a1*nt2::log(dtype(a0)));
      }
  };

  NT2_FUNCTOR_IMPLEMENTATION( nt2::tag::pow_, tag::cpu_
                              , (A0)(A1)
                              ,  (generic_< dry_ < floating_<A0> > > )
                              (generic_< imaginary_<floating_<A1> > >)
                              )
  {
    typedef typename meta::as_complex<A0>::type result_type;
    NT2_FUNCTOR_CALL(2)
      {
        return nt2::exp(a1*nt2::log(a0));
      }
  };

  NT2_FUNCTOR_IMPLEMENTATION( nt2::tag::pow_, tag::cpu_
                              , (A0)(A1)
                              ,  (generic_< dry_ < floating_<A0> > > )
                              (generic_< dry_<floating_<A1> > >)
                              )
  {
    typedef typename meta::as_complex<A0>::type result_type;
    NT2_FUNCTOR_CALL(2)
      {
        return nt2::exp(a1*nt2::log(a0));
      }
  };
} }

#endif


/******************************************************************************
 *
 * AMDiS - Adaptive multidimensional simulations
 *
 * Copyright (C) 2013 Dresden University of Technology. All Rights Reserved.
 * Web: https://fusionforge.zih.tu-dresden.de/projects/amdis
 *
 * Authors:
 * Simon Vey, Thomas Witkowski, Andreas Naumann, Simon Praetorius, et al.
 *
 * This file is provided AS IS with NO WARRANTY OF ANY KIND, INCLUDING THE
 * WARRANTY OF DESIGN, MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE.
 *
 *
 * This file is part of AMDiS
 *
 * See also license.opensource.txt in the distribution.
 *
 ******************************************************************************/

// Written by Simon Praetorius (adopted from previous implementation)

#ifndef ITL_GMRES_HOUSEHOLDER_INCLUDE
#define ITL_GMRES_HOUSEHOLDER_INCLUDE

#include <algorithm>
#include <boost/numeric/mtl/concept/collection.hpp>
#include <boost/numeric/mtl/vector/dense_vector.hpp>
#include <boost/numeric/mtl/matrix/dense2D.hpp>
#include <boost/numeric/mtl/matrix/multi_vector.hpp>
#include <boost/numeric/mtl/operation/givens.hpp>
#include <boost/numeric/mtl/operation/two_norm.hpp>
#include <boost/numeric/mtl/utility/exception.hpp>
#include <boost/numeric/mtl/utility/irange.hpp>

#include "solver/itl/details.hpp"

namespace itl
{

  /// Generalized Minimal Residual method (without restart) using householder othogonalization.
  /** It computes at most kmax_in iterations (or size(x) depending on what is smaller)
      regardless on whether the termination criterion is reached or not.   **/
  template <typename Matrix, typename Vector, typename LeftPreconditioner, typename Iteration>
  int gmres_householder_full(const Matrix& A, Vector& x, const Vector& b,
                             LeftPreconditioner& L, Iteration& iter)
  {
    using mtl::irange;
    using std::abs;
    using math::reciprocal;
    using mtl::iall;
    using mtl::imax;
    using mtl::signum;
    using mtl::vector::dot;
    using mtl::conj;
    typedef typename mtl::Collection<Vector>::value_type Scalar;
    typedef typename mtl::Collection<Vector>::size_type  Size;

    if (size(b) == 0) throw mtl::logic_error("empty rhs vector");

    const Scalar zero= math::zero(Scalar()), dbl_tol= 1.e-16;
    Scalar       rho, bnrm2, temp, beta;
    Size         k, kmax(std::min(size(x), Size(iter.max_iterations() - iter.iterations())));
    Vector       w(b - A *x), r(solve(L,w));
    mtl::matrix::multi_vector<Vector>   V(Vector(resource(x), zero), kmax+1);
    mtl::vector::dense_vector<Scalar>   sn(kmax, zero), cs(kmax, zero), s(kmax+1, zero), y(kmax, zero);  // replicated in distributed solvers
    mtl::matrix::dense2D<Scalar>        H(kmax, kmax);

    bnrm2 = two_norm(b);
    if (bnrm2 < dbl_tol)
      bnrm2 = 1.0;

    temp = two_norm(r);				// norm of preconditioned residual
    rho = temp * reciprocal(bnrm2);
    if (iter.finished(rho))			// initial guess is good enough solution
      return iter;

    // u = r + sign(r(0))*||r||*e0
    beta = signum(r[0])*temp;
    w = r;
    w[0] += beta;
    w *= reciprocal(two_norm(w));

    V.vector(0) = w;
    H = zero;
    s[0] = -beta;

    // GMRES iteration
    for (k= 0; k < kmax && !iter.finished(rho); ++k, ++iter)
    {

      w = (-2.0 * V.vector(k)[k])*V.vector(k);
      w[k] += 1.0;
      // v := P_0*...*P_{k-2}*(P_{k-1} * e_k)
      for (Size i= k; i > 0; i--)
      {
        temp = 2.0 * dot(V.vector(i-1), w);
        w -= temp * V.vector(i-1);
      }

      temp = two_norm(w);
      if (temp == zero)
        return iter.fail(2, "GMRES: breakdown");

      // Explicitly normalize v to reduce the effects of round-off.
      w *= reciprocal(temp);
      w = solve(L, Vector(A*w));

      // P_{k-1}*...*P_0*Av
      for (Size i = 0; i <= k; i++)
      {
        temp = 2.0 * dot(V.vector(i), w);
        w -= temp * V.vector(i);
      }

      temp = two_norm(w);
      if (temp == zero)
        return iter.fail(3, "GMRES: breakdown");

      irange range_to_end(k+1,imax);
      set_to_zero(V.vector(k+1));
      V.vector(k+1)[range_to_end] = w[range_to_end];
      beta = two_norm(V.vector(k+1));
      if (beta != 0.0)
      {
        beta *= signum(w[k+1]);
        V.vector(k+1)[k+1] += beta;
        V.vector(k+1) *= reciprocal(two_norm(V.vector(k+1)));

        w[k+1] = -beta;
      }

      for (Size i= 0; i < k; i++)
      {
        temp   =  conj(cs[i])*w[i] + conj(sn[i])*w[i+1];
        w[i+1] = -sn[i]*w[i] + cs[i]*w[i+1];
        w[i]   =  temp;
      }

      details::rotmat(w[k], w[k+1], cs[k], sn[k]);

      s[k+1] = -sn[k]*s[k];
      s[k]   = conj(cs[k])*s[k];
      w[k]   = cs[k]*w[k] + sn[k]*w[k+1];
      w[k+1] = 0.0;

      irange range(num_rows(H));
      H[iall][k] = w[range];

      rho = std::abs(s[k+1]) / bnrm2;
    }

    // reduce k, to get regular matrix
    //     while (k > 0 && std::abs(s[k-1]) <= iter.atol()) k--;

    // iteration is finished -> compute x: solve H*y=s as far as rank of H allows
    irange range(k);
    for (; !range.empty(); --range)
    {
      try
      {
        y[range] = upper_trisolve(H[range][range], s[range]);
      }
      catch (mtl::matrix_singular)
      {
        continue;    // if singular then try with sub-matrix
      }
      break;
    }

    if (range.finish() < k)
      std::cerr << "GMRES orhogonalized with " << k << " vectors but matrix singular, can only use "
                << range.finish() << " vectors!\n";
    if (range.empty())
      return iter.fail(3, "GMRES did not find any direction to correct x");

    kmax = k-1;

    w = V.vector(kmax) * (-2.0 * y[kmax] * conj(V.vector(kmax)[kmax]));
    w[kmax] += y[kmax];
    for (Size i= kmax; i > 0; i--)
    {
      w[i-1] += y[i-1];
      temp = 2.0 * dot(V.vector(i-1), w);
      w -= temp * V.vector(i-1);
    }
    x += w;

    r = b - A*x;
    return iter.terminate(r);
  }

  /// Generalized Minimal Residual method with restart
  template <typename Matrix, typename Vector, typename LeftPreconditioner,
            typename Iteration>
  int gmres_householder(const Matrix& A, Vector& x, const Vector& b,
                        LeftPreconditioner& L,
                        Iteration& iter, typename mtl::Collection<Vector>::size_type restart)
  {
    do
    {
      Iteration inner(iter);
      inner.set_max_iterations(std::min(int(iter.iterations()+restart), iter.max_iterations()));
      inner.suppress_resume(true);
      gmres_householder_full(A, x, b, L, inner);
      iter.update_progress(inner);
    }
    while (!iter.finished());

    return iter;
  }

} // namespace itl

#endif // ITL_GMRES_HOUSEHOLDER_INCLUDE


#pragma once
#include <stdint.h>
#include "settings.hpp"
#include "weight_initialization.hpp"
#include "activations.hpp"
#if USE_EIGEN == 1
#include <Eigen/Dense>
#endif

namespace nn
{
	enum class LayerType
	{
		kInput,
		kFC,
		kSoftmax
	};

	// any layer has some amount of units and activation function
#if USE_EIGEN == 1
	class layer
	{
	public:
		using MatrixType = Eigen::Matrix<real, Eigen::Dynamic, Eigen::Dynamic>;

		layer(LayerType type,
			uint32_t unitsInLayer,
			uint32_t unitsInPreviousLayer,
			ActivationType activationType,
			WeightInitializationType weightInitializationType) :
			m_type(type),
			m_unitsInLayer(unitsInLayer),
			m_unitsInPreviousLayer(unitsInPreviousLayer)
		{
			// set activation type
			if (activationType == ActivationType::kSigmoid)
			{
				m_activation = activation<ActivationType::kSigmoid>;
				m_activationDerivative = activation_derivative<ActivationType::kSigmoid>;
			}
			else if (activationType == ActivationType::kLinear)
			{
				m_activation = activation<ActivationType::kLinear>;
				m_activationDerivative = activation_derivative<ActivationType::kLinear>;
			}
			else if (activationType == ActivationType::kTanh)
			{
				m_activation = activation<ActivationType::kTanh>;
				m_activationDerivative = activation_derivative<ActivationType::kTanh>;
			}
			else if (activationType == ActivationType::kRelu)
			{
				m_activation = activation<ActivationType::kRelu>;
				m_activationDerivative = activation_derivative<ActivationType::kRelu>;
			}
			else if (activationType == ActivationType::kLRelu)
			{
				m_activation = activation<ActivationType::kLRelu>;
				m_activationDerivative = activation_derivative<ActivationType::kLRelu>;
			}

			if (type != LayerType::kInput)
			{
				if (weightInitializationType == WeightInitializationType::kGaussian)
				{
					m_weight = MatrixType::Zero(UnitsInLayer(), UnitsInPreviousLayer()).unaryExpr(weight_initalization<WeightInitializationType::kZeros>());
					m_nabla_w = MatrixType::Zero(UnitsInLayer(), UnitsInPreviousLayer());
					m_bias = MatrixType::Zero(UnitsInLayer(), 1).unaryExpr(weight_initalization<WeightInitializationType::kZeros>());
					m_nabla_b = MatrixType::Zero(UnitsInLayer(), 1);
				}
				else if (weightInitializationType == WeightInitializationType::kSequentialDebug)
				{
					m_weight = MatrixType::Zero(UnitsInLayer(), UnitsInPreviousLayer()).unaryExpr(weight_initalization<WeightInitializationType::kSequentialDebug>());
					m_nabla_w = MatrixType::Zero(UnitsInLayer(), UnitsInPreviousLayer());
					m_bias = MatrixType::Zero(UnitsInLayer(), 1).unaryExpr(weight_initalization<WeightInitializationType::kSequentialDebug>());
					m_nabla_b = MatrixType::Zero(UnitsInLayer(), 1);
				}
				else if (weightInitializationType == WeightInitializationType::kUniform)
				{
					m_weight = MatrixType::Zero(UnitsInLayer(), UnitsInPreviousLayer()).unaryExpr(weight_initalization<WeightInitializationType::kUniform>());
					m_nabla_w = MatrixType::Zero(UnitsInLayer(), UnitsInPreviousLayer());
					m_bias = MatrixType::Zero(UnitsInLayer(), 1).unaryExpr(weight_initalization<WeightInitializationType::kUniform>());
					m_nabla_b = MatrixType::Zero(UnitsInLayer(), 1);
				}
				else if (weightInitializationType == WeightInitializationType::kGaussian)
				{
					m_weight = MatrixType::Zero(UnitsInLayer(), UnitsInPreviousLayer()).unaryExpr(weight_initalization<WeightInitializationType::kGaussian>());
					m_nabla_w = MatrixType::Zero(UnitsInLayer(), UnitsInPreviousLayer());
					m_bias = MatrixType::Zero(UnitsInLayer(), 1).unaryExpr(weight_initalization<WeightInitializationType::kGaussian>());
					m_nabla_b = MatrixType::Zero(UnitsInLayer(), 1);
				}
				else if (weightInitializationType == WeightInitializationType::kWeightedGaussian)
				{
					m_weight = MatrixType::Zero(UnitsInLayer(), UnitsInPreviousLayer()).unaryExpr(weight_initalization<WeightInitializationType::kWeightedGaussian>(UnitsInLayer()));
					m_nabla_w = MatrixType::Zero(UnitsInLayer(), UnitsInPreviousLayer());
					m_bias = MatrixType::Zero(UnitsInLayer(), 1).unaryExpr(weight_initalization<WeightInitializationType::kWeightedGaussian>(UnitsInLayer()));
					m_nabla_b = MatrixType::Zero(UnitsInLayer(), 1);
				}
			}
		}

		~layer() {}

		void computeWeightedSum(const MatrixType& input)
		{
			m_z.noalias() = m_weight * input;
			for (int i = 0; i < m_z.cols(); ++i)
				m_z.col(i).noalias() += m_bias;

		}
		void setActivations(const MatrixType& input) { m_a.noalias() = input; }
		void computeActivations(const MatrixType& input)
		{
			if (m_type == LayerType::kFC)
				m_a.noalias() = input.unaryExpr(m_activation);
			else if (m_type == LayerType::kSoftmax)
			{
				if (m_a.rows() != input.rows() || m_a.cols() != input.cols())
					m_a = MatrixType::Zero(input.rows(), input.cols());
				MatrixType maxCol(m_a.rows(), 1);
				for (int i = 0; i < input.cols(); ++i)
				{
					maxCol.setConstant(input.maxCoeff()); // prevent softmax overflow
					m_a.col(i).noalias() = (input.col(i) - maxCol).unaryExpr(&expf);
					m_a.col(i) /= m_a.col(i).sum();
				}
			}
		}
		void computeActivationDerivatives(const MatrixType& input)
		{
			if (m_type == LayerType::kFC)
				m_da.noalias() = input.unaryExpr(m_activationDerivative);
			else if (m_type == LayerType::kSoftmax)
				m_da = MatrixType::Ones(m_a.rows(), m_a.cols());
		}

		MatrixType computeWeightedSumExplicit(const MatrixType& input)
		{
			MatrixType result = m_weight * input;
			for (int i = 0; i < result.cols(); ++i)
				result.col(i).noalias() += m_bias;
			return result;
		}

		MatrixType computeActivationsExplicit(const MatrixType& input)
		{
			MatrixType result;
			if (m_type == LayerType::kFC)
				result.noalias() = input.unaryExpr(m_activation);
			else if (m_type == LayerType::kSoftmax)
			{
				if (result.rows() != input.rows() || result.cols() != input.cols())
					result = MatrixType::Zero(input.rows(), input.cols());
				MatrixType maxCol(result.rows(), 1);
				for (int i = 0; i < input.cols(); ++i)
				{
					maxCol.setConstant(input.maxCoeff()); // prevent softmax overflow
					result.col(i).noalias() = (input.col(i) - maxCol).unaryExpr(&expf);
					result.col(i) /= result.col(i).sum();
				}
			}
			return result;
		}
		MatrixType computeActivationDerivativesExplicit(const MatrixType& input)
		{
			MatrixType result;
			if (m_type == LayerType::kFC)
				result.noalias() = input.unaryExpr(m_activationDerivative);
			else if (m_type == LayerType::kSoftmax)
				result = MatrixType::Ones(input.rows(), input.cols());
			return result;
		}

		const MatrixType& getWeightedSum() const { return m_z; }
		const MatrixType& getActivations() const { return m_a; }
		const MatrixType& getActivationDerivatives() const { return m_da; }
		const MatrixType& getWeights() const { return m_weight; }
		MatrixType& getWeights() { return m_weight; }
		MatrixType& getBias() { return m_bias; }
		MatrixType& getNablaB() { return m_nabla_b; }
		MatrixType& getNablaW() { return m_nabla_w; }

		uint32_t UnitsInLayer() const { return m_unitsInLayer; }
		uint32_t UnitsInPreviousLayer() const { return m_unitsInPreviousLayer; }
	private:
		LayerType m_type;
		uint32_t m_unitsInLayer;
		uint32_t m_unitsInPreviousLayer;
		MatrixType m_z;
		MatrixType m_a;
		MatrixType m_da;

		MatrixType m_weight;
		MatrixType m_nabla_w;
		MatrixType m_bias;
		MatrixType m_nabla_b;
		ActivationFunction m_activation = nullptr;
		ActivationFunction m_activationDerivative = nullptr;
	};
#endif
}