/**
 * Copyright (c) 2022 <Daumantas Kavolis>
 *
 * 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.
 */

#pragma once

#include "config.hpp"

MSVC_WARNING_DISABLE(4619)
#include <boost/math/constants/constants.hpp>
#include <boost/math/special_functions/chebyshev.hpp>
#include <boost/math/special_functions/hermite.hpp>
#include <boost/math/special_functions/laguerre.hpp>
#include <boost/math/special_functions/legendre.hpp>
#include <boost/math/special_functions/legendre_stieltjes.hpp>
MSVC_WARNING_POP()

#include "quadrature.hpp"

namespace poly {
namespace detail {

template <typename Integral>
struct BasicImpl {
  using OrderType = Integral;
  using Storage = Integral;

  constexpr static inline bool is_trivial = true;

  static auto make_storage(OrderType order = 0) noexcept -> Storage { return order; }
  static auto get_order(Storage const& item) noexcept -> OrderType { return item; }
  static void set_order(Storage& item, OrderType const& value) noexcept { item = value; }
};

}  // namespace detail

template <typename Real_, unsigned N = QuadraturePoints>
struct LegendreImpl : detail::BasicImpl<std::int16_t> {
  using Real = Real_;
  constexpr static inline bool is_orthonormal = true;

  static auto eval(Storage const& item, Real x) -> Real { return boost::math::legendre_p(item, x); }
  static auto prime(Storage const& item, Real x) -> Real {
    return boost::math::legendre_p_prime(item, x);
  }
  static auto zeros(Storage const& item) -> std::vector<Real> {
    using std::abs;
    auto z = boost::math::legendre_p_zeros<Real>(item);
    detail::reflect_in_place<detail::Reflection::Odd>(
        z, abs(z[0]) < 2 * std::numeric_limits<Real>::epsilon());
    return z;
  }
  static auto next(Storage const& item, Real x, Real Pl, Real Plm1) -> Real {
    return boost::math::legendre_next(item, x, Pl, Plm1);
  }

  template <bool check = false>
  static auto weights(Storage const& item, bounds_check<check> c = no_bounds_check)
      -> view<Real const> {
    return GaussQuadrature<Real, N>::weights(item + 1, c);
  }
  template <bool check = false>
  static auto abscissa(Storage const& item, bounds_check<check> c = no_bounds_check)
      -> view<Real const> {
    return GaussQuadrature<Real, N>::abscissa(item + 1, c);
  }
  static auto domain() -> std::pair<Real, Real> { return {-1, 1}; }
};

template <typename Real_>
struct LaguerreImpl : detail::BasicImpl<std::int16_t> {
  using Real = Real_;
  constexpr static inline bool is_orthonormal = true;

  static auto eval(Storage const& item, Real x) -> Real { return boost::math::laguerre(item, x); }
  static auto next(Storage const& item, Real x, Real Pl, Real Plm1) -> Real {
    return boost::math::laguerre_next(item, x, Pl, Plm1);
  }
  static auto domain() -> std::pair<Real, Real> { return {0, 1}; }
};

template <typename Real_>
struct HermiteImpl : detail::BasicImpl<std::uint16_t> {
  using Real = Real_;
  constexpr static inline bool is_orthonormal = true;

  static auto eval(Storage const& item, Real x) -> Real { return boost::math::hermite(item, x); }
  static auto next(Storage const& item, Real x, Real Pl, Real Plm1) -> Real {
    return boost::math::hermite_next(item, x, Pl, Plm1);
  }
  static auto domain() -> std::pair<Real, Real> {
    return {-std::numeric_limits<Real>::infinity(), std::numeric_limits<Real>::infinity()};
  }
};

template <typename, unsigned>
struct ChebyshevImpl;
namespace detail {
template <typename Real, unsigned N = QuadraturePoints>
struct ChebyshevQuadrature {
  static auto abscissa() -> view<Real const> {
    static std::vector<Real> v = ChebyshevImpl<Real, N>::zeros({N});
    return {v.data(), v.size()};
  }
  static auto weights() -> view<Real const> {
    static std::vector<Real> v(std::size_t(N), boost::math::constants::pi<Real>() / N);
    return {v.data(), v.size()};
  }
};
}  // namespace detail

template <typename Real, unsigned N = QuadraturePoints>
using ChebyshevQuadrature = detail::Quadrature<Real, detail::ChebyshevQuadrature, N, false>;

template <typename Real_, unsigned N = QuadraturePoints>
struct ChebyshevImpl : detail::BasicImpl<std::int16_t> {
  using Real = Real_;
  constexpr static inline bool is_orthonormal = true;

  static auto eval(Storage const& item, Real x) -> Real {
    return boost::math::chebyshev_t(item, x);
  }
  static auto prime(Storage const& item, Real x) -> Real {
    return boost::math::chebyshev_t_prime(item, x);
  }
  static auto next(Storage const& /* unused */, Real x, Real Pl, Real Plm1) -> Real {
    return boost::math::chebyshev_next(x, Pl, Plm1);
  }
  static auto zeros(Storage const& item) -> std::vector<Real> {
    using std::cos;
    std::vector<Real> values(item, 0);
    Real factor{boost::math::constants::pi<Real>() / 2.0 / item};
    auto half = item / 2;
    auto start = (item + 1) / 2;
    for (OrderType i = 0; i < half; ++i) {
      Real x = cos((2 * (start + i) + 1) * factor);
      values[start + i] = -x;
      values[half - i - 1] = x;
    }
    return values;
  }

  template <bool check = false>
  static auto weights(Storage const& item, bounds_check<check> c = no_bounds_check)
      -> view<Real const> {
    return ChebyshevQuadrature<Real, N>::weights(item + 1, c);
  }
  template <bool check = false>
  static auto abscissa(Storage const& item, bounds_check<check> c = no_bounds_check)
      -> view<Real const> {
    return ChebyshevQuadrature<Real, N>::abscissa(item + 1, c);
  }
  static auto domain() -> std::pair<Real, Real> { return {-1, 1}; }
};

template <typename Real_>
struct LegendreStieltjesImpl {
  using Real = Real_;
  using OrderType = std::int16_t;

  constexpr static inline bool is_orthonormal = true;

  class Storage {
   public:
    explicit Storage(OrderType m) : polynomial{m}, order{m} {};

   private:
    boost::math::legendre_stieltjes<Real> polynomial;
    OrderType order;

    template <class>
    friend struct LegendreStieltjesImpl;
  };

  static auto make_storage(OrderType order = 0) noexcept -> Storage { return Storage{order}; }
  static auto get_order(Storage const& item) noexcept -> OrderType { return item.order; }
  static void set_order(Storage& item, OrderType const& value) noexcept { item = Storage{value}; }

  static auto eval(Storage const& item, Real x) -> Real { return item.polynomial(x); }
  static auto prime(Storage const& item, Real x) -> Real { return item.polynomial.prime(x); }
  static auto zeros(Storage const& item) -> std::vector<Real> {
    auto vector = item.polynomial.zeros();
    detail::reflect_in_place<detail::Reflection::Odd>(vector, vector[0] == 0);
    return vector;
  }
  static auto domain() -> std::pair<Real, Real> { return {-1, 1}; }
};

template <typename Real_, unsigned N = QuadraturePoints>
struct GaussKronrodImpl : detail::BasicImpl<std::int16_t> {
  using Real = Real_;

  constexpr static inline bool is_orthonormal = true;

  template <bool check = false>
  static auto weights(Storage const& item, bounds_check<check> c = no_bounds_check)
      -> view<Real const> {
    return GaussKronrodQuadrature<Real, N>::weights(item, c);
  }
  template <bool check = false>
  static auto abscissa(Storage const& item, bounds_check<check> c = no_bounds_check)
      -> view<Real const> {
    return GaussKronrodQuadrature<Real, N>::abscissa(item, c);
  }
  static auto domain() -> std::pair<Real, Real> { return {-1, 1}; }
};

}  // namespace poly


#include "FiniteElement.h"

#include <Eigen/Core>
#include <iostream>

using namespace Eigen;

FiniteElement::FiniteElement(
        FEMObject* femObj,
        Cell cell,
        ElasticMaterial material)
    : mFemObj(femObj)
    , mMaterial(material)
{
    mDensity = 100.0;
    mCell = cell;
    mStiffnessMatrixDirty = false;
    mQ.setIdentity();
}

void FiniteElement::setMaterial(ElasticMaterial material)
{
    mMaterial = material;
    mStiffnessMatrixDirty = true;
}

void FiniteElement::initialize()
{
    updateDnx();

    // Initialize mQ with the correct rotation calculated with the slow singular
    // value decomposition. Use that result with the iterative approach later.
    updateRotation();
    mQ = Eigen::Quaterniond(mR);

    updateCauchyStressStrain();
    updateLinearStiffnessMatrix();
}

void FiniteElement::updateRotation()
{
    updateF();
    JacobiSVD<Matrix3d> svd(mF, ComputeFullU | ComputeFullV);
    double det = (svd.matrixU() * svd.matrixV().transpose()).determinant();
    Matrix3d C = Matrix3d::Identity();
    C(2,2) = det;
    mR = svd.matrixU() * C * svd.matrixV().transpose();
//    if (det < 0)
//        std::cout << "inversion detected! det = " << det << "\n";

}

void FiniteElement::updateRotationFast(size_t maxIterations, double tolerance)
{
    updateF();

    for (size_t i = 0; i < maxIterations; ++i)
    {
        Eigen::Matrix3d R = mQ.toRotationMatrix();
        double s = std::abs(R.col(0).dot(mF.col(0)) +
                            R.col(1).dot(mF.col(1)) +
                            R.col(2).dot(mF.col(2)));

        if (s < 1e-12)
            break;

        double sInv = 1.0 / s + tolerance;
        Eigen::Vector3d v = R.col(0).cross(mF.col(0)) +
                            R.col(1).cross(mF.col(1)) +
                            R.col(2).cross(mF.col(2));

        Eigen::Vector3d omega = sInv * v;
        double w = omega.norm();

        if (w < tolerance)
            break;

        Eigen::Quaterniond omegaQ =
                Eigen::Quaterniond(Eigen::AngleAxisd(w, 1.0 / w * omega));

        mQ = omegaQ * mQ;
    }

    mR = mQ.toRotationMatrix();
}

void FiniteElement::update(bool corotated)
{
    if (corotated)
    {
        updateCorotatedStiffnessMatrix();
        updateCorotatedForces();
    }
    else
    {
        updateLinearForces();
    }
}

void FiniteElement::updateStiffnessMatrix(bool corotated)
{
    if (corotated)
        updateCorotatedStiffnessMatrix();
    else if (mStiffnessMatrixDirty)
        updateLinearStiffnessMatrix();
}

void FiniteElement::updateLinearStiffnessMatrix()
{
    for (size_t a = 0; a < 4; ++a)
    {
        for (size_t b = 0; b < 4; ++b)
        {
            double second_part = mMaterial.getLameMu() * mDnx[a].dot(mDnx[b]);
            Matrix3d& K = mK[a][b];
            for (Index i = 0; i < 3; ++i)
            {
                for (Index k = 0; k < 3; ++k)
                {
                    //K(i,k) = ...
                    double value = 0.0;
                    if (i == k)
                        value = second_part;
                    value += mMaterial.getLameLambda() * mDnx[a](i) * mDnx[b](k)
                            + mMaterial.getLameMu() * mDnx[a](k) * mDnx[b](i);
                    K(i,k) = value * mVolume;
                }
            }
        }
    }

    updatePlasticityMatrix();
}

void FiniteElement::updateCorotatedStiffnessMatrix()
{
    for (size_t a = 0; a < 4; ++a)
    {
        for (size_t b = 0; b < 4; ++b)
        {
            Matrix3d& K = mK[a][b];
            Matrix3d& K_corot = mKCorot[a][b];
            K_corot = mR * K * mR.transpose();
        }
    }
}

void FiniteElement::updatePlasticityMatrix()
{
    Eigen::Matrix<double, 6, 6> D;
    double mu = mMaterial.getLameMu();
    double lambda = mMaterial.getLameLambda();

    double temp = lambda + 2 * mu;

    D << temp, lambda, lambda, 0, 0, 0,
            lambda, temp, lambda, 0, 0, 0,
            lambda, lambda, temp, 0, 0, 0,
            0, 0, 0, mu, 0, 0,
            0, 0, 0, 0, mu, 0,
            0, 0, 0, 0, 0, mu;

    mP = mVolume * mB.transpose() * D;
}

void FiniteElement::updateForces(bool corotated)
{
    if (corotated)
        updateCorotatedForces();
    else
        updateLinearForces();
}

void FiniteElement::updateCorotatedForces()
{
    // TODO: Optimizations
    // - Apply rotation only once
    // - reformulate equation to precalculate parts of it
    for (size_t a = 0; a < 4; ++a)
    {
        mForces[a].setZero();
        for (size_t b = 0; b < 4; ++b)
        {
            mForces[a] += mR * mK[a][b] *
                    (mR.transpose() * y(b) - x(b));
        }
    }

    updatePlasticForces(0.01);
}

void FiniteElement::updateLinearForces()
{
    for (size_t a = 0; a < 4; ++a)
    {
        mForces[a] = Vector::Zero();
        for (size_t b = 0; b < 4; ++b)
        {
            mForces[a] += mK[a][b] * u(b);

            // directly globally updates them this way
            // but this is not wanted here
            //f(a) += mK[a][b] * u(b);
        }
    }
}

void FiniteElement::updatePlasticForces(double stepSize)
{
    if (mMaterial.getPlasticMaxStrain() > 1e-12)
    {
        // Calculate total strain.
        Eigen::Matrix<double, 6, 1> strainTotal =
                Eigen::Matrix<double, 6, 1>::Zero();
        Eigen::Matrix3d RInv = mR.inverse();
        for (Eigen::Index a = 0; a < 4; ++a)
        {
            strainTotal += mB.block(0, a * 3, 6, 3) * (RInv * y(a) - x(a));
        }

        // Update plastic strain.
        Eigen::Matrix<double, 6, 1> strainElastic =
                strainTotal - mStrainPlastic;
        if (strainElastic.norm() > mMaterial.getPlasticYield())
            mStrainPlastic += stepSize * mMaterial.getPlasticCreep() * strainElastic;
        double strainPlasticNorm = mStrainPlastic.norm();
        if (strainPlasticNorm > mMaterial.getPlasticMaxStrain())
            mStrainPlastic *= mMaterial.getPlasticMaxStrain() / strainPlasticNorm;

        // Calculate plastic forces.
        for (Eigen::Index a = 0; a < 4; ++a)
        {
            mForces[a] -= mR * mP.block(a * 3, 0, 3, 6) * mStrainPlastic;
        }
    }
}

void FiniteElement::updateF()
{
//    mF = Matrix3d::Identity();
//    for (size_t i = 0; i < 4; ++i) {
//        mF += u(i) * mDnx[i].transpose();
//    }
    mF = Matrix3d::Zero();
    for (size_t i = 0; i < 4; ++i) {
        mF += y(i) * mDnx[i].transpose();
    }

    mB << mDnx[0](0), 0, 0, mDnx[1](0), 0, 0, mDnx[2](0), 0, 0, mDnx[3](0), 0, 0,
            0, mDnx[0](1), 0, 0, mDnx[1](1), 0, 0, mDnx[2](1), 0, 0, mDnx[3](1), 0,
            0, 0, mDnx[0](2), 0, 0, mDnx[1](2), 0, 0, mDnx[2](2), 0, 0, mDnx[3](2),
            mDnx[0](1), mDnx[0](0), 0, mDnx[1](1), mDnx[1](0), 0, mDnx[2](1), mDnx[2](0), 0, mDnx[3](1), mDnx[3](0), 0,
            mDnx[0](2), 0, mDnx[0](0), mDnx[1](2), 0, mDnx[1](0), mDnx[2](2), 0, mDnx[2](0), mDnx[3](2), 0, mDnx[3](0),
            0, mDnx[0](2), mDnx[0](1), 0, mDnx[1](2), mDnx[1](1), 0, mDnx[2](2), mDnx[2](1), 0, mDnx[3](2), mDnx[3](1);
}

void FiniteElement::updateCauchyStressStrain()
{
    // Cauchy strain
    mCauchyStrain = 0.5 * (mF + mF.transpose()) - Matrix3d::Identity();
    double trace = mCauchyStrain.trace();
    for (Index i = 0; i < 3; ++i)
         for (Index j = 0; j < 3; ++j)
            mCauchyStress(i, j) = (i == j) * mMaterial.getLameLambda() * trace
                    + 2 * mMaterial.getLameMu() * mCauchyStrain(i, j);
}

void FiniteElement::updateDnx()
{
    // As described in
    // http://www.iue.tuwien.ac.at/phd/nentchev/node30.html and
    // http://www.iue.tuwien.ac.at/phd/nentchev/node31.html

    Vector r1 = mFemObj->x(mCell[1]) - mFemObj->x(mCell[0]);
    Vector r2 = mFemObj->x(mCell[2]) - mFemObj->x(mCell[0]);
    Vector r3 = mFemObj->x(mCell[3]) - mFemObj->x(mCell[0]);

    Vector r4 = mFemObj->x(mCell[2]) - mFemObj->x(mCell[1]);
    Vector r5 = mFemObj->x(mCell[1]) - mFemObj->x(mCell[3]);

    mVolume = r1.cross(r2).dot(r3);

    // Sometimes tetrahedrons can be so degenerated that their volume is zero.
    if (mVolume < 1e-10)
        mVolume = 1e-5;

    mDnx[0] = r4.cross(r5) / mVolume;
    mDnx[1] = r2.cross(r3) / mVolume;
    mDnx[2] = r3.cross(r1) / mVolume;
    mDnx[3] = r1.cross(r2) / mVolume;

    mVolume /= 6.0;
}

void FiniteElement::updateDny()
{
    updateF(); // TODO: look if necessary
    Matrix3d F_inv = mF.inverse();
    for (size_t i = 0; i < 4; ++i) {
        mDny[i] = Vector(0,0,0);
        for (int a = 0; a < 3; ++a) {
            for (int k = 0; k < 3; ++k) {
                mDny[i](a) += mDnx[i](k) * F_inv(k,a);
            }
        }
    }
}


#include <iostream>
#include <math.h>
#include <Eigen/Eigen>
#include <Eigen/Dense>
#include <Eigen/Core>

#include "../include/UKFilter.h"
#include "../include/Robot.h"
#include "../include/Odom.h"
#include "../include/Imu.h"
#include "../include/Gps.h"

UnscentedKalmanFilter unscentedKalmanFilter;
UnscentedKalmanFilter *punscentedKalmanFilter;

// Constructor
UnscentedKalmanFilter::UnscentedKalmanFilter()
{
    is_initialized_ = false;
    lastTimeStamp_ = 0;
    nowTimeStamp_ = 0;
    deltaTime_ = 0;
}

//Destructor
UnscentedKalmanFilter::~UnscentedKalmanFilter()
{
    delete punscentedKalmanFilter;
}

void UnscentedKalmanFilter::Initialization()
{
    lastTimeStamp_ = nowTimeStamp_ = getSysTime();

    Eigen::VectorXf x_in = Eigen::VectorXf::Zero(5,1);
    SetX(x_in);

    // state covariance matrix the prediction error
    Eigen::MatrixXf P_in(5,5);
    P_in << 0.1, 0.0, 0.0, 0.0, 0.0,
            0.0, 0.1, 0.0, 0.0, 0.0,
            0.0, 0.0, 0.1, 0.0, 0.0,
            0.0, 0.0, 0.0, 0.1, 0.0,
            0.0, 0.0, 0.0, 0.0, 0.1;
    SetP(P_in);

    //predicted sigma
    Eigen::MatrixXf PreSigmaX_in = Eigen::MatrixXf::Zero(5, 15);
    SetPreSigmaX(PreSigmaX_in);

    //measurement covariance matrix
    // R is provided by Sensor supplier
    Eigen::MatrixXf R_in(3,3);
    R_in << 0.1, 0.0, 0.0,
            0.0, 0.1, 0.0,
            0.0, 0.0, 0.1;
    SetR(R_in);

    //Observation
    Eigen::VectorXf z_in(3,1);
    z_in << 0.0,
            0.0,
            0.0;
    SetZ(z_in);

    is_initialized_ = true;

}

//state prediction
//state defination: x, y, v, yaw, w
void UnscentedKalmanFilter::StatePrediction(float d_t)
{
    Eigen::MatrixXf augmented_sigma = Eigen::MatrixXf::Zero(7, 15);
    Eigen::VectorXf augmented_x = Eigen::VectorXf::Zero(7);
    Eigen::MatrixXf augmented_P = Eigen::MatrixXf::Zero(7, 7);
    augmented_x.head(5) = x_;
    augmented_P.topLeftCorner(5, 5) = P_;
    augmented_P(5, 5) = 0.1;
    augmented_P(6, 6) = 0.1;
    const Eigen::MatrixXf L = augmented_P.llt().matrixL();
    augmented_sigma.col(0) = augmented_x;
    for(int c = 0; c < 7; c++)
    {
        const int i = c + 1;
        augmented_sigma.col(i)  = augmented_x + 1.732 * L.col(c);
        augmented_sigma.col(i + 7)  = augmented_x - 1.732 * L.col(c);
    }

    Eigen::MatrixXf predicted_sigma = Eigen::MatrixXf(5, 15);
    for(int c = 0; c < 15; ++c)
    {
        /*************************************
        * Get the current state
        *************************************/
        const float px = augmented_sigma(0, c);
        const float py = augmented_sigma(1, c);
        const float speed = augmented_sigma(2, c);
        const float yaw = augmented_sigma(3, c);
        const float yawrate = augmented_sigma(4, c);
        const float speed_noise = augmented_sigma(5, c);
        const float yawrate_noise = augmented_sigma(6, c);

        /*************************************
        * predict the next state with noise
        * USING THE CTRV MODEL
        *************************************/
        const float cos_yaw = cos(yaw);
        const float sin_yaw = sin(yaw);
        const float d_t2 = d_t * d_t;

        // predicted position noise
        const float p_noise = 0.5 * speed_noise * d_t2;
        // predicted yaw noise
        const float y_noise = 0.5 * yawrate_noise * d_t2;
        const float dyaw = yawrate * d_t; //change in yaw
        const float dspeed = speed * d_t; //change in speed

        // predicted speed = assumed constant speed + noise
        const float p_speed = speed + speed_noise * d_t;

        // predicted yaw
        const float p_yaw = yaw + dyaw + y_noise;
        // predicted yaw rate = assumed constant yawrate + noise
        const float p_yawrate = yawrate + yawrate_noise * d_t;
        // where predicted positions will be stored
        float p_px, p_py;

        if(fabs(yawrate) <= 0.001)
        {
            // moving straight
            p_px = px + dspeed * cos_yaw + p_noise * cos_yaw;
            p_py = py + dspeed * sin_yaw + p_noise * sin_yaw;

        }
        else
        {
            const float k = speed / yawrate;
            const float theta = yaw + dyaw;
            p_px = px + k * (sin(theta) - sin_yaw) + p_noise * cos_yaw;
            p_py = py + k * (cos_yaw - cos(theta)) + p_noise * sin_yaw ;
        }

        /*************************************
        * Write the prediction to the appropriate column
        *************************************/
        predicted_sigma(0, c) = p_px;
        predicted_sigma(1, c) = p_py;
        predicted_sigma(2, c) = p_speed;
        predicted_sigma(3, c) = p_yaw;
        predicted_sigma(4, c) = p_yawrate;
    }

    Eigen::VectorXf predicted_x = Eigen::VectorXf::Zero(5);
    for(int c = 0; c < 15; c++)
    {
        predicted_x += WEIGHTS[c] * predicted_sigma.col(c);
    }
    Eigen::MatrixXf predicted_P = Eigen::MatrixXf::Zero(5, 5);
    Eigen::VectorXf dx = Eigen::VectorXf(5);
    for(int c = 0; c < 15; c++)
    {
        dx = predicted_sigma.col(c) - predicted_x;
        dx(3) = (fabs(dx(3)) > M_PI) ? remainder(dx(3), 2. * M_PI) : dx(3);
        predicted_P += WEIGHTS[c] * dx * dx.transpose();
    }
    SetPreSigmaX(predicted_sigma);
    SetX(predicted_x);
    SetP(predicted_P);
}

void UnscentedKalmanFilter::MeasurementPrediction()//measurement prediction
{
    Eigen::MatrixXf sigma = Eigen::MatrixXf::Zero(3, 15);
    for(int c = 0; c < 15; c++)
    {
        const float px = PreSigmaX_(0, c);
        const float py = PreSigmaX_(1, c);
        const float v = PreSigmaX_(2, c);
        const float yaw = PreSigmaX_(3, c);

        const float vx = cos(yaw) * v;
        const float vy = sin(yaw) * v;

        const float rho = sqrt(px * px + py * py);
        const float phi = atan2(py, px);
        const float rhodot = (rho > 0.0001) ? ((px * vx + py * vy) / rho) : 0.0;

        // avoid division by zero
        sigma(0, c) = rho;
        sigma(1, c) = phi;
        sigma(2, c) = rhodot;
    }

    Eigen::VectorXf z = Eigen::VectorXf::Zero(3);
    Eigen::VectorXf dz;
    Eigen::MatrixXf S = Eigen::MatrixXf::Zero(3, 3);
    for(int c = 0; c < 15; c++){
        z += WEIGHTS[c] * sigma.col(c);
    }
    for(int c = 0; c < 15; c++)
    {
        dz = sigma.col(c) - z;
        dz(1) =  (fabs(dz(1)) > M_PI) ? remainder(dz(1), 2. * M_PI) : dz(1);
        S += WEIGHTS[c] * dz * dz.transpose();
    }
    S += R_;

    SetPreSigmaZ(sigma);
    SetZ(z);
    SetS(S);
}

void UnscentedKalmanFilter::StateUpdate(Eigen::VectorXf z_new)
{
    Eigen::VectorXf dz;
    Eigen::VectorXf dx;
    Eigen::MatrixXf Tc = Eigen::MatrixXf::Zero(5, 3);

    for(int c = 0; c < 15; c++)
    {
        dx = PreSigmaX_.col(c) - x_;
        dx(3) = (fabs(dx(3)) > M_PI) ? remainder(dx(3), 2. * M_PI) : dx(3);
        dz = PreSigmaZ_.col(c) - z_;
        dz(1) = (fabs(dz(1)) > M_PI) ? remainder(dz(1), 2. * M_PI) : dz(1);
        Tc += WEIGHTS[c] * dx * dz.transpose();
    }
    Eigen::MatrixXf Si = S_.inverse();
    Eigen::MatrixXf K = Tc * Si;
    Eigen::VectorXf dz = z_new - z_;
    dz(1) = (fabs(dz(1)) > M_PI) ? remainder(dz(1), 2. * M_PI) : dz(1);

    SetX(x_ + K * dz);
    SetP(P_ - K * S_ * K.transpose());
}

bool UnscentedKalmanFilter::GetIsInitialized()
{
    return is_initialized_;
}

void UnscentedKalmanFilter::SetX(Eigen::VectorXf x_in)
{
    x_ = x_in;
}

void UnscentedKalmanFilter::SetP(Eigen::MatrixXf P_in)
{
    P_ = P_in;
}

void UnscentedKalmanFilter::SetPreSigmaX(Eigen::MatrixXf PreSigmaX_in)
{
    PreSigmaX_ = PreSigmaX_in;
}

void UnscentedKalmanFilter::SetPreSigmaZ(Eigen::MatrixXf PreSigmaZ_in)
{
    PreSigmaZ_ = PreSigmaZ_in;
}

void UnscentedKalmanFilter::SetZ(Eigen::VectorXf z_in)
{
    z_ = z_in;
}

void UnscentedKalmanFilter::SetR(Eigen::MatrixXf S_in)
{
    S_ = S_in;
}

void UnscentedKalmanFilter::SetR(Eigen::MatrixXf R_in)
{
    R_ = R_in;
}

void UnscentedKalmanFilter::setLastTimeStamp(const long int lastTimeStamp)
{
    this->lastTimeStamp_ = lastTimeStamp;
}

void UnscentedKalmanFilter::setNowTimeStamp(const long int nowTimeStamp )
{
    this->nowTimeStamp_ = nowTimeStamp;
}
void UnscentedKalmanFilter::setDeltaTime(const long int deltaTime)
{
    this->deltaTime_ = deltaTime;
}

void UnscentedKalmanFilter::getLastTimeStamp(long int& lastTimeStamp)
{
    lastTimeStamp = this->lastTimeStamp_;
}
void UnscentedKalmanFilter::getNowTimeStamp(long int& nowTimeStamp)
{
    nowTimeStamp = this->nowTimeStamp_;
}
void UnscentedKalmanFilter::getDeltaTime(long int& deltaTime)
{
    deltaTime = this->deltaTime_;
}


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

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

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

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

/*! \file qle/math/deltagammavar.hpp
    \brief functions to compute delta or delta-gamma VaR numbers
    \ingroup math
*/

#ifndef quantext_deltagammavar_hpp
#define quantext_deltagammavar_hpp

#include <qle/math/covariancesalvage.hpp>

#include <ql/math/array.hpp>
#include <ql/math/matrix.hpp>
#include <ql/math/matrixutilities/choleskydecomposition.hpp>
#include <ql/math/matrixutilities/pseudosqrt.hpp>
#include <ql/math/randomnumbers/rngtraits.hpp>
#include <ql/utilities/disposable.hpp>

// fix for boost 1.64, see https://lists.boost.org/Archives/boost/2016/11/231756.php
#if BOOST_VERSION >= 106400
#include <boost/serialization/array_wrapper.hpp>
#endif
#include <boost/accumulators/accumulators.hpp>
#include <boost/accumulators/statistics/stats.hpp>
#include <boost/accumulators/statistics/tail_quantile.hpp>
#include <boost/foreach.hpp>

namespace QuantExt {
using namespace QuantLib;

//! function that computes a delta VaR
/*! For a given covariance matrix and a delta vector this function computes a parametric var w.r.t. a given
 * confidence level for multivariate normal risk factors. */
Real deltaVar(const Matrix& omega, const Array& delta, const Real p,
              const CovarianceSalvage& sal = NoCovarianceSalvage());

//! function that computes a delta-gamma normal VaR
/*! For a given a covariance matrix, a delta vector and a gamma matrix this function computes a parametric var
 * w.r.t. a given confidence level. The gamma matrix is taken into account when computing the variance of the PL
 * distirbution, but the PL distribution is still assumed to be normal. */
Real deltaGammaVarNormal(const Matrix& omega, const Array& delta, const Matrix& gamma, const Real p,
                         const CovarianceSalvage& sal = NoCovarianceSalvage());

//! function that computes a delta-gamma VaR using Monte Carlo (single quantile)
/*! For a given a covariance matrix, a delta vector and a gamma matrix this function computes a parametric var
 * w.r.t. a given confidence level. The var quantile is estimated from Monte-Carlo realisations of a second order
 * sensitivity based PL. */
template <class RNG>
Real deltaGammaVarMc(const Matrix& omega, const Array& delta, const Matrix& gamma, const Real p, const Size paths,
                     const Size seed, const CovarianceSalvage& sal = NoCovarianceSalvage());

//! function that computes a delta-gamma VaR using Monte Carlo (multiple quantiles)
/*! For a given a covariance matrix, a delta vector and a gamma matrix this function computes a parametric var
 * w.r.t. a vector of given confidence levels. The var quantile is estimated from Monte-Carlo realisations of a second
 * order sensitivity based PL. */
template <class RNG>
Disposable<std::vector<Real> > deltaGammaVarMc(const Matrix& omega, const Array& delta, const Matrix& gamma,
                                               const std::vector<Real>& p, const Size paths, const Size seed,
                                               const CovarianceSalvage& sal = NoCovarianceSalvage());

namespace detail {
void check(const Real p);
void check(const Matrix& omega, const Array& delta);
void check(const Matrix& omega, const Array& delta, const Matrix& gamma);
template <typename A> Real absMax(const A& a) {
    Real tmp = 0.0;
    BOOST_FOREACH (Real x, a) {
        if (std::abs(x) > tmp)
            tmp = std::abs(x);
    }
    return tmp;
}
} // namespace detail

// implementation

template <class RNG>
Disposable<std::vector<Real> > deltaGammaVarMc(const Matrix& omega, const Array& delta, const Matrix& gamma,
                                               const std::vector<Real>& p, const Size paths, const Size seed,
                                               const CovarianceSalvage& sal) {
    BOOST_FOREACH (Real q, p) { detail::check(q); }
    detail::check(omega, delta, gamma);

    Real num = std::max(detail::absMax(delta), detail::absMax(gamma));
    if (close_enough(num, 0.0)) {
        std::vector<Real> res(p.size(), 0.0);
        return res;
    }

    Matrix L = sal.salvage(omega).second;
    if (L.rows() == 0) {
        L = CholeskyDecomposition(omega, true);
    }

    Real pmin = QL_MAX_REAL;
    BOOST_FOREACH (Real q, p) { pmin = std::min(pmin, q); }

    Size cache = Size(std::floor(static_cast<double>(paths) * (1.0 - pmin) + 0.5)) + 2;
    boost::accumulators::accumulator_set<
        double, boost::accumulators::stats<boost::accumulators::tag::tail_quantile<boost::accumulators::right> > >
        acc(boost::accumulators::tag::tail<boost::accumulators::right>::cache_size = cache);

    typename RNG::rsg_type rng = RNG::make_sequence_generator(delta.size(), seed);

    for (Size i = 0; i < paths; ++i) {
        std::vector<Real> seq = rng.nextSequence().value;
        Array z(seq.begin(), seq.end());
        Array u = L * z;
        acc(DotProduct(u, delta) + 0.5 * DotProduct(u, gamma * u));
    }

    std::vector<Real> res;
    BOOST_FOREACH (Real q, p) {
        res.push_back(boost::accumulators::quantile(acc, boost::accumulators::quantile_probability = q));
    }

    return res;
}

template <class RNG>
Real deltaGammaVarMc(const Matrix& omega, const Array& delta, const Matrix& gamma, const Real p, const Size paths,
                     const Size seed, const CovarianceSalvage& sal) {

    std::vector<Real> pv(1, p);
    return deltaGammaVarMc<RNG>(omega, delta, gamma, pv, paths, seed, sal).front();
}

} // namespace QuantExt

#endif


#include <mtao/types.hpp>
#include <iostream>
#include<Eigen/IterativeLinearSolvers>
#include <Eigen/Sparse>
#include <mtao/logging/timer.hpp>
#include <mtao/eigen/mic0_preconditioner.hpp>
using Mat = mtao::MatrixX<float>;

float permeability = 100.0;
float timestep = 1000.0;

Eigen::SparseMatrix<float> grid_boundary(int ni, int nj, bool dirichlet_boundary) {
    int system_size = ni*nj;
    int usize = (ni+1) * nj;
    int vsize = (nj+1) * ni;
    int offset = usize;

    auto u_ind = [ni,nj](int i, int j) {
        return i + j*(ni + 1);
    };

    auto v_ind = [ni,nj,offset](int i, int j) {
        return i + j*ni + offset;
    };

    std::vector< Eigen::Triplet<float> > triplets;

    Eigen::SparseMatrix<float> D(usize + vsize, system_size);
    for (int j = 0; j < nj ; ++j) {
        for (int i = 0; i < ni ; ++i) {
            int index = i + ni*j;
            if(!(dirichlet_boundary && i == 0)) {
                triplets.push_back(Eigen::Triplet<float>( u_ind(i,j), index, 1.0/ni));
            }
            if(!(dirichlet_boundary && i == ni-1)) {
            triplets.push_back(Eigen::Triplet<float>( u_ind(i+1,j), index, -1.0/ni));
            }

            if(!(dirichlet_boundary && j == 0)) {
            triplets.push_back(Eigen::Triplet<float>(  v_ind(i,j), index, 1.0/nj));
            }
            if(!(dirichlet_boundary && j == ni-1)) {
                triplets.push_back(Eigen::Triplet<float>(  v_ind(i,j+1), index, -1.0/nj));
            }

        }
    }
    D.setFromTriplets(triplets.begin(), triplets.end());
    return D;
}

template <typename SolverType>
void run_test(int N, SolverType& solver, const std::string& name, bool print = false) {
    Mat data = Mat::Random(N,N);
    Eigen::Map<Eigen::VectorXf> theta(data.data(),data.size());
    int size = N*N;
    auto B = grid_boundary(N,N,true);
    auto dirichlet_energy = [&]() -> double {
        return (B * theta).squaredNorm();
    };
    //size ~ dx * dy ~ dx^2
    Eigen::SparseMatrix<float> L = B.transpose() * B;
    Eigen::SparseMatrix<float> A(size,size);

    auto t = mtao::logging::timer(name + "-total");
    for(int i = 0; i < 1; ++i) {
        timestep = std::pow(10,i);
        A.setIdentity();
        A /= timestep;
        A += L;
        //auto t = mtao::logging::timer(name);
        solver.compute(A);
        for(int j = 0; j < 10; ++j) {
            data = Mat::Random(N,N);
            for(int i = 0; i < 10; ++i) {
                theta = solver.solve(theta);
                if(print) {
                    std::cout << dirichlet_energy() << std::endl;
                }
            }
        }
    }

        auto dms = std::chrono::duration_cast<std::chrono::milliseconds>(t.duration());
        std::cout << dms.count() << ",";
}

int main() {
    mtao::logging::make_logger("default",mtao::logging::Level::Info);
    for(int i = 50; i < 1000; i+=10) {
        Eigen::SimplicialLDLT<Eigen::SparseMatrix<float>> ldlt_solver;
        Eigen::ConjugateGradient<Eigen::SparseMatrix<float>, Eigen::Lower|Eigen::Upper> cg_solver;
        Eigen::ConjugateGradient<Eigen::SparseMatrix<float>, Eigen::Lower|Eigen::Upper, mtao::eigen::preconditioners::MIC0Preconditioner<float>> mic0pcg_solver;
        //run_test(ldlt_solver, "LDLT");
        std::cout << i << ",";
        run_test(i,cg_solver, "CG");
        run_test(i,mic0pcg_solver, "MIC0PCG");
        std::cout << std::endl;
    }
    return 0;

    }


///////////////////////////////////////////////////////////////////////////////
// median.hpp
//
//  Copyright 2006 Eric Niebler, Olivier Gygi. Distributed under the Boost
//  Software License, Version 1.0. (See accompanying file
//  LICENSE_1_0.txt or copy at http://www.boost.org/LICENSE_1_0.txt)

#ifndef BOOST_ACCUMULATORS_STATISTICS_MEDIAN_HPP_EAN_28_10_2005
#define BOOST_ACCUMULATORS_STATISTICS_MEDIAN_HPP_EAN_28_10_2005

#include <boost/mpl/placeholders.hpp>
#include <boost/range/iterator_range.hpp>
#include <boost/accumulators/framework/accumulator_base.hpp>
#include <boost/accumulators/framework/extractor.hpp>
#include <boost/accumulators/numeric/functional.hpp>
#include <boost/accumulators/framework/parameters/sample.hpp>
#include <boost/accumulators/framework/depends_on.hpp>
#include <boost/accumulators/statistics_fwd.hpp>
#include <boost/accumulators/statistics/count.hpp>
#include <boost/accumulators/statistics/p_square_quantile.hpp>
#include <boost/accumulators/statistics/density.hpp>
#include <boost/accumulators/statistics/p_square_cumul_dist.hpp>

namespace boost { namespace accumulators
{

namespace impl
{
    ///////////////////////////////////////////////////////////////////////////////
    // median_impl
    //
    /**
        @brief Median estimation based on the \f$P^2\f$ quantile estimator

        The \f$P^2\f$ algorithm is invoked with a quantile probability of 0.5.
    */
    template<typename Sample>
    struct median_impl
      : accumulator_base
    {
        // for boost::result_of
        typedef typename numeric::functional::average<Sample, std::size_t>::result_type result_type;

        median_impl(dont_care) {}

        template<typename Args>
        result_type result(Args const &args) const
        {
            return p_square_quantile_for_median(args);
        }
    };
    ///////////////////////////////////////////////////////////////////////////////
    // with_density_median_impl
    //
    /**
        @brief Median estimation based on the density estimator

        The algorithm determines the bin in which the \f$0.5*cnt\f$-th sample lies, \f$cnt\f$ being
        the total number of samples. It returns the approximate horizontal position of this sample,
        based on a linear interpolation inside the bin.
    */
    template<typename Sample>
    struct with_density_median_impl
      : accumulator_base
    {
        typedef typename numeric::functional::average<Sample, std::size_t>::result_type float_type;
        typedef std::vector<std::pair<float_type, float_type> > histogram_type;
        typedef iterator_range<typename histogram_type::iterator> range_type;
        // for boost::result_of
        typedef float_type result_type;

        template<typename Args>
        with_density_median_impl(Args const &args)
          : sum(numeric::average(args[sample | Sample()], (std::size_t)1))
          , is_dirty(true)
        {
        }

        void operator ()(dont_care)
        {
            this->is_dirty = true;
        }

        template<typename Args>
        result_type result(Args const &args) const
        {
            if (this->is_dirty)
            {
                this->is_dirty = false;

                std::size_t cnt = count(args);
                range_type histogram = density(args);
                typename range_type::iterator it = histogram.begin();
                while (this->sum < 0.5 * cnt)
                {
                    this->sum += it->second * cnt;
                    ++it;
                }
                --it;
                float_type over = numeric::average(this->sum - 0.5 * cnt, it->second * cnt);
                this->median = it->first * over + (it + 1)->first * (1. - over);
            }

            return this->median;
        }

    private:
        mutable float_type sum;
        mutable bool is_dirty;
        mutable float_type median;
    };

    ///////////////////////////////////////////////////////////////////////////////
    // with_p_square_cumulative_distribution_median_impl
    //
    /**
        @brief Median estimation based on the \f$P^2\f$ cumulative distribution estimator

        The algorithm determines the first (leftmost) bin with a height exceeding 0.5. It
        returns the approximate horizontal position of where the cumulative distribution
        equals 0.5, based on a linear interpolation inside the bin.
    */
    template<typename Sample>
    struct with_p_square_cumulative_distribution_median_impl
      : accumulator_base
    {
        typedef typename numeric::functional::average<Sample, std::size_t>::result_type float_type;
        typedef std::vector<std::pair<float_type, float_type> > histogram_type;
        typedef iterator_range<typename histogram_type::iterator> range_type;
        // for boost::result_of
        typedef float_type result_type;

        with_p_square_cumulative_distribution_median_impl(dont_care)
          : is_dirty(true)
        {
        }

        void operator ()(dont_care)
        {
            this->is_dirty = true;
        }

        template<typename Args>
        result_type result(Args const &args) const
        {
            if (this->is_dirty)
            {
                this->is_dirty = false;

                range_type histogram = p_square_cumulative_distribution(args);
                typename range_type::iterator it = histogram.begin();
                while (it->second < 0.5)
                {
                    ++it;
                }
                float_type over = numeric::average(it->second - 0.5, it->second - (it - 1)->second);
                this->median = it->first * over + (it + 1)->first * ( 1. - over );
            }

            return this->median;
        }
    private:

        mutable bool is_dirty;
        mutable float_type median;
    };

} // namespace impl

///////////////////////////////////////////////////////////////////////////////
// tag::median
// tag::with_densisty_median
// tag::with_p_square_cumulative_distribution_median
//
namespace tag
{
    struct median
      : depends_on<p_square_quantile_for_median>
    {
        /// INTERNAL ONLY
        ///
        typedef accumulators::impl::median_impl<mpl::_1> impl;
    };
    struct with_density_median
      : depends_on<count, density>
    {
        /// INTERNAL ONLY
        ///
        typedef accumulators::impl::with_density_median_impl<mpl::_1> impl;
    };
    struct with_p_square_cumulative_distribution_median
      : depends_on<p_square_cumulative_distribution>
    {
        /// INTERNAL ONLY
        ///
        typedef accumulators::impl::with_p_square_cumulative_distribution_median_impl<mpl::_1> impl;
    };
}

///////////////////////////////////////////////////////////////////////////////
// extract::median
// extract::with_density_median
// extract::with_p_square_cumulative_distribution_median
//
namespace extract
{
    extractor<tag::median> const median = {};
    extractor<tag::with_density_median> const with_density_median = {};
    extractor<tag::with_p_square_cumulative_distribution_median> const with_p_square_cumulative_distribution_median = {};

    BOOST_ACCUMULATORS_IGNORE_GLOBAL(median)
    BOOST_ACCUMULATORS_IGNORE_GLOBAL(with_density_median)
    BOOST_ACCUMULATORS_IGNORE_GLOBAL(with_p_square_cumulative_distribution_median)
}

using extract::median;
using extract::with_density_median;
using extract::with_p_square_cumulative_distribution_median;

// median(with_p_square_quantile) -> median
template<>
struct as_feature<tag::median(with_p_square_quantile)>
{
    typedef tag::median type;
};

// median(with_density) -> with_density_median
template<>
struct as_feature<tag::median(with_density)>
{
    typedef tag::with_density_median type;
};

// median(with_p_square_cumulative_distribution) -> with_p_square_cumulative_distribution_median
template<>
struct as_feature<tag::median(with_p_square_cumulative_distribution)>
{
    typedef tag::with_p_square_cumulative_distribution_median type;
};

// for the purposes of feature-based dependency resolution,
// with_density_median and with_p_square_cumulative_distribution_median
// provide the same feature as median
template<>
struct feature_of<tag::with_density_median>
  : feature_of<tag::median>
{
};

template<>
struct feature_of<tag::with_p_square_cumulative_distribution_median>
  : feature_of<tag::median>
{
};

// So that median can be automatically substituted with
// weighted_median when the weight parameter is non-void.
template<>
struct as_weighted_feature<tag::median>
{
    typedef tag::weighted_median type;
};

template<>
struct feature_of<tag::weighted_median>
  : feature_of<tag::median>
{
};

// So that with_density_median can be automatically substituted with
// with_density_weighted_median when the weight parameter is non-void.
template<>
struct as_weighted_feature<tag::with_density_median>
{
    typedef tag::with_density_weighted_median type;
};

template<>
struct feature_of<tag::with_density_weighted_median>
  : feature_of<tag::with_density_median>
{
};

// So that with_p_square_cumulative_distribution_median can be automatically substituted with
// with_p_square_cumulative_distribution_weighted_median when the weight parameter is non-void.
template<>
struct as_weighted_feature<tag::with_p_square_cumulative_distribution_median>
{
    typedef tag::with_p_square_cumulative_distribution_weighted_median type;
};

template<>
struct feature_of<tag::with_p_square_cumulative_distribution_weighted_median>
  : feature_of<tag::with_p_square_cumulative_distribution_median>
{
};

}} // namespace boost::accumulators

#endif


////////////////////////////////////////////////////////////////////////////////////
// The MIT License (MIT)                                                          //
//                                                                                //
// Copyright (c) 2015 Whit Armstrong                                              //
//                                                                                //
// 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.                                                                      //
////////////////////////////////////////////////////////////////////////////////////

#pragma once

#include <stdexcept>
#include <armadillo>
#include <armalogp/arma.extensions.hpp>
#include <armalogp/arma.math.hpp>

namespace armalogp {

  template<double LOGF(double), typename T, typename U, typename V>
  double normal_logp(const T& x, const U& mu, const V& tau) {
    return arma::accu(0.5*LOGF(0.5*tau/arma::datum::pi) - 0.5 * arma::schur_prod(tau, square(x - mu)));
  }

  template<double LOGF(double), typename T, typename U, typename V>
  double uniform_logp(const T& x, const U& lower, const V& upper) {
    return (arma::any(arma::vectorise(x < lower)) || arma::any(arma::vectorise(x > upper))) ? -std::numeric_limits<double>::infinity() : -arma::accu(LOGF(upper - lower));
  }

  template<double LOGF(double), typename T, typename U, typename V>
  double gamma_logp(const T& x, const U& alpha, const V& beta) {
    return arma::any(arma::vectorise(x < 0)) ?
      -std::numeric_limits<double>::infinity() :
      arma::accu(arma::schur_prod((alpha - 1.0),LOGF(x)) - arma::schur_prod(beta,x) - lgamma(alpha) + arma::schur_prod(alpha,LOGF(beta)));
  }

  template<double LOGF(double), typename T, typename U, typename V>
  double beta_logp(const T& x, const U& alpha, const V& beta) {
    const double one = 1.0;
    return arma::any(arma::vectorise(x <= 0)) || arma::any(arma::vectorise(x >= 1)) || arma::any(arma::vectorise(alpha <= 0)) || arma::any(arma::vectorise(beta <= 0)) ?
      -std::numeric_limits<double>::infinity() :
      arma::accu(lgamma(alpha+beta) - lgamma(alpha) - lgamma(beta) + arma::schur_prod((alpha-one),LOGF(x)) + arma::schur_prod((beta-one),LOGF(one-x)));
  }

  template<double LOGF(double)>
  double categorical_logp(const arma::ivec& x, const arma::mat& p) {
    if(arma::any(arma::vectorise(p <= 0)) || arma::any(arma::vectorise(p >= 1)) || arma::any(arma::vectorise(x < 0)) || arma::any(arma::vectorise(x >= p.n_cols))) {
      return -std::numeric_limits<double>::infinity();
    }
    // replace w/ call to p.elems later
    double ans(0);
    for(unsigned int i = 0; i < x.n_rows; i++) {
      ans += LOGF(p(i,x[i]));
    }
    return ans;
  }

  template<double LOGF(double)>
  double categorical_logp(const arma::ivec& x, const arma::vec& p) {
    if(arma::any(arma::vectorise(p <= 0)) || arma::any(arma::vectorise(p >= 1)) || arma::any(arma::vectorise(x < 0)) || arma::any(arma::vectorise(x >= p.n_elem))) {
      return -std::numeric_limits<double>::infinity();
    }
    // replace w/ call to p.elems later
    double ans(0);
    for(unsigned int i = 0; i < x.n_rows; i++) {
      ans += LOGF(p(x[i]));
    }
    return ans;
  }

  template<double LOGF(double)>
  double categorical_logp(const int x, const arma::vec& p) {
    return LOGF(p[x]);
  }

  template<double LOGF(double), typename T, typename U, typename V>
  double binomial_logp(const T& x, const U& n, const V& p) {
    if(arma::any(arma::vectorise(p <= 0)) || arma::any(arma::vectorise(p >= 1)) || arma::any(arma::vectorise(x < 0))  || arma::any(arma::vectorise(x > n))) {
      return -std::numeric_limits<double>::infinity();
    }
    return arma::accu(arma::schur_prod(x,LOGF(p)) + arma::schur_prod((n-x),LOGF(1-p)) + arma::factln(n) - arma::factln(x) - arma::factln(n-x));
  }

  template<double LOGF(double), typename T, typename U>
  double bernoulli_logp(const T& x, const U& p) {
    if( arma::any(arma::vectorise(p <= 0)) || arma::any(arma::vectorise(p >= 1)) || arma::any(arma::vectorise(x < 0))  || arma::any(arma::vectorise(x > 1)) ) {
      return -std::numeric_limits<double>::infinity();
    } else {
      return arma::accu(arma::schur_prod(x,LOGF(p)) + arma::schur_prod((1-x), LOGF(1-p)));
    }
  }

  template<double LOGF(double), typename T, typename U>
  double poisson_logp(const T& x, const U& mu) {
    if( arma::any(arma::vectorise(mu < 0)) || arma::any(arma::vectorise(x < 0))) {
      return -std::numeric_limits<double>::infinity();
    } else {
      return arma::accu(schur(x,LOGF(mu)) - mu - factln(x));
    }
  }

  template<double LOGF(double), typename T, typename U>
  double exponential_logp(const T& x, const U& lambda) {
    if(!arma::all(arma::vectorise(x > 0)) || !arma::all(arma::vectorise(lambda > 0))) {
        return -std::numeric_limits<double>::infinity();
    }
    return arma::accu(LOGF(lambda) - arma::schur_prod(lambda, x));
  }

  template<double LOGF(double), typename T, typename U>
  double multivariate_normal_chol_logp(const T& x, const U& mu, const arma::mat& R) {
    static double log_2pi = LOGF(2 * arma::datum::pi);
    double ldet = LOGF(cholesky_determinant(R));
    return -0.5 * (x.n_elem * log_2pi + ldet + mahalanobis_chol(x,mu,R));
  }

  // sigma denotes cov matrix rather than precision matrix
  template<double LOGF(double), typename T, typename U>
  double multivariate_normal_sigma_logp(const T& x, const U& mu, const arma::mat& sigma) {
    arma::mat R;
    bool chol_succeeded = chol(R,sigma);
    if(!chol_succeeded) { return -std::numeric_limits<double>::infinity(); }

    return multivariate_normal_chol_logp<LOGF>(x, mu, R);
  }

  // sigma denotes cov matrix rather than precision matrix
  template<double LOGF(double)>
  double multivariate_normal_sigma_logp(const arma::mat& x, const arma::vec& mu, const arma::mat& sigma) {
    arma::mat R;
    bool chol_succeeded = chol(R,sigma);
    if(!chol_succeeded) { return -std::numeric_limits<double>::infinity(); }
    const arma::rowvec mu_r = mu.t();
    double ans(0);
    for(size_t i = 0; i < x.n_rows; i++) {
      ans += multivariate_normal_chol_logp<LOGF>(x.row(i), mu_r, R);
    }
    return ans;
  }

  template<double LOGF(double)>
  double multivariate_normal_chol_logp(const arma::mat& x, const arma::vec& mu, const arma::mat& R) {
    const arma::rowvec mu_r = mu.t();
    double ans(0);
    for(size_t i = 0; i < x.n_rows; i++) {
      ans += multivariate_normal_chol_logp<LOGF>(x.row(i), mu_r, R);
    }
    return ans;
  }

  template<double LOGF(double)>
  double wishart_logp(const arma::mat& X, const arma::mat& tau, const unsigned int n) {
    if(X.n_cols != X.n_rows || tau.n_cols != tau.n_rows || X.n_cols != tau.n_rows || X.n_cols > n) { return -std::numeric_limits<double>::infinity(); }
    const double lg2 = LOGF(2.0);
    const int k = X.n_cols;
    const double dx(arma::det(X));
    const double db(arma::det(tau));
    if(dx <= 0 || db <= 0) { return -std::numeric_limits<double>::infinity(); }

    const double ldx(LOGF(dx));
    const double ldb(LOGF(db));
    const arma::mat bx(X * tau);
    const double tbx = arma::trace(bx);

    double cum_lgamma(0);
    for(size_t i = 0; i < X.n_rows; ++i) {
      cum_lgamma += lgamma((n + 1)/2.0);
    }
    return (n - k - 1)/2 * ldx + (n/2.0)*ldb - 0.5*tbx - (n*k/2.0)*lg2 - cum_lgamma;
  }

} // namespace armalogp


#include <boost/program_options.hpp>

#include <iostream>
#include <fstream>
#include <cassert>
#include <list>
#include <vector>

// https://doc.cgal.org/latest/Triangulation_3/
// https://doc.cgal.org/latest/Kernel_23/index.html#Chapter_2D_and_3D_Geometry_Kernel
// http://doc.cgal.org/latest/Triangulation_3/classTriangulationTraits__3.html
#include <CGAL/Exact_predicates_inexact_constructions_kernel.h>

// https://doc.cgal.org/latest/Triangulation_3/classCGAL_1_1Triangulation__3.html
// https://doc.cgal.org/latest/Triangulation_3/classCGAL_1_1Delaunay__triangulation__3.html
#include <CGAL/Delaunay_triangulation_3.h>
#include <CGAL/Triangulation_data_structure_3.h>

// The doc for Point is a little tricky to find. Delaunay and Triangulation pages
// don't link here:
// https://doc.cgal.org/latest/Kernel_23/classCGAL_1_1Point__3.html
// http://doc.cgal.org/latest/Kernel_23/classCGAL_1_1Cartesian.html
// http://www.cgal.org/FAQ.html

typedef CGAL::Exact_predicates_inexact_constructions_kernel K;
typedef CGAL::Delaunay_triangulation_3<K>                   Triangulation;
typedef Triangulation::Point                                Point;

// these two functions are for my readvtk function.
Triangulation::Vertex_handle
add_vertex(Triangulation &T, std::string& vname, double x, double y, double z)
{
  // vertex name doesn't have any meaning in this context.
  return T.insert(Point(x,y,z));
}

void add_edge(Triangulation &T,
               Triangulation::Vertex_handle s,
               Triangulation::Vertex_handle t)
{
  // no edges to add. CGAL does that.
}

#include <readvtk.hxx>

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

    namespace po = boost::program_options;
    po::options_description desc("Usage");

    std::string filename;
    desc.add_options()
        ("help", "produce help message")
        ("filename", po::value<std::string>(&filename)->default_value(""),
         "filename containing input points");

    po::variables_map opts;
    po::store(po::parse_command_line(argc, argv, desc), opts);

    try {
        po::notify(opts);
    } catch (std::exception& e) {
        std::cerr << "Error: " << e.what() << "\n";
        return 1;
    }

    if (filename == "") {
      std::cerr << "please provide a vtk file with the --filename <file> option" << std::endl;
      exit(-1);
    }

    std::ifstream input(filename.c_str());
    Triangulation T;

    /* if you just have a file with a list of points, you can do this:
       std::list<Point> L;
       Point p ;
       Triangulation T;
       while (input >> p) {
         T.insert(p);
         L.push_front(p);
       }
       // if you already have a list of points:
       // Triangulation T(L.begin(), L.end());

    */

    readvtk<Triangulation,Triangulation::Vertex_handle>(input, T);

    // at this point, the trianulation is already done. Just have to extract
    // the computed lines.
    Triangulation::size_type n = T.number_of_vertices();

    assert( T.is_valid() ); // checking validity of T

    std::cout << "# vtk DataFile Version 1.0\n";
    std::cout << "3D triangulation data\n";
    std::cout << "ASCII\n";
    std::cout << std::endl;
    std::cout << "DATASET POLYDATA\n";

    // according to http://doc.cgal.org/latest/TDS_3/classTriangulationDataStructure__3.html
    // Note that the triangulation data structure has one more
    // vertex than an associated geometric triangulation, if there
    // is one, since the infinite vertex is a standard vertex and
    // is thus also counted.
    // I am omitting this one by using the finite_vertices/edges iterators below
    int numvertices = T.number_of_vertices();
    std::cout << "POINTS " << numvertices << " float\n";

    // the APIs I'm using below are documented here:
    // http://doc.cgal.org/latest/Triangulation_3/classCGAL_1_1Triangulation__3.html

    // as I iterate over the vertices, I assign ids to each vertex. Seems like there'd
    // be an id method, but I haven't found it.
    std::map<Point, int> vertex_ids;
    Triangulation::Finite_vertices_iterator viter;
    for (viter =  T.finite_vertices_begin();
         viter != T.finite_vertices_end();
         viter++) {
      Triangulation::Triangulation_data_structure::Vertex v = *viter;
      // in the line below, looking up the point (the [] part) will add a new
      // element to the map. That's why I subtract 1.
      vertex_ids[v.point()] = vertex_ids.size() - 1;
      std::cout << v.point() << std::endl;
    }

    std::cout << "LINES " << T.number_of_finite_edges() << " "
              << 3*T.number_of_finite_edges() << std::endl;

    Triangulation::Finite_edges_iterator iter;
    for(iter =  T.finite_edges_begin();
        iter != T.finite_edges_end();
        iter++) {
      // edges are not represented as edges in CGAL triangulation graphs.
      // Instead, they are stored in faces/cells.

      Triangulation::Triangulation_data_structure::Edge e = *iter;
      Triangulation::Triangulation_data_structure::Cell_handle c = e.first;
      int i = e.second;
      int j = e.third;

      Triangulation::Triangulation_data_structure::Vertex_handle a = c->vertex(i);
      Point pa = a->point();
      Point pb = c->vertex(j)->point();
      int ida, idb;
      if (vertex_ids.find(pa) == vertex_ids.end()) {
        std::cout << "didn't find " << pa << std::endl;
        ida = 0;
      } else {
        ida = vertex_ids[pa];
      }
      if (vertex_ids.find(pb) == vertex_ids.end()) {
        std::cout << "didn't find " << pb << std::endl;
        idb = 0;
      } else {
        idb = vertex_ids[pb];
      }
      std::cout << "2 " << ida << " " << idb << std::endl;
    }

    return 0;
}


#include <cmath>
#include <cstring>
#include <Eigen/Eigenvalues>
#include <Eigen/SVD>
#include <esn/exceptions.hpp>
#include <esn/network_nsli.h>
#include <esn/network_nsli.hpp>
#include <network_nsli.h>

namespace ESN {

    std::unique_ptr< Network > CreateNetwork(
        const NetworkParamsNSLI & params )
    {
        return std::unique_ptr< NetworkNSLI >( new NetworkNSLI( params ) );
    }

    NetworkNSLI::NetworkNSLI( const NetworkParamsNSLI & params )
        : mParams( params )
        , mIn( params.inputCount )
        , mWIn( params.neuronCount, params.inputCount )
        , mWInScaling( params.inputCount )
        , mWInBias( params.inputCount )
        , mX( params.neuronCount )
        , mW( params.neuronCount, params.neuronCount )
        , mOut( params.outputCount )
        , mWOut( params.outputCount, params.neuronCount )
        , mWFB()
        , mWFBScaling()
        , mAdaptiveFilter( params.neuronCount,
            params.onlineTrainingForgettingFactor,
            params.onlineTrainingInitialCovariance )
    {
        if ( params.inputCount <= 0 )
            throw std::invalid_argument(
                "NetworkParamsNSLI::inputCount must be not null" );
        if ( params.neuronCount <= 0 )
            throw std::invalid_argument(
                "NetworkParamsNSLI::neuronCount must be not null" );
        if ( params.outputCount <= 0 )
            throw std::invalid_argument(
                "NetworkParamsNSLI::outputCount must be not null" );
        if ( !( params.leakingRateMin > 0.0 &&
                params.leakingRateMin <= 1.0 ) )
            throw std::invalid_argument(
                "NetworkParamsNSLI::leakingRateMin must be within "
                "interval (0,1]" );
        if ( !( params.leakingRateMax > 0.0 &&
                params.leakingRateMax <= 1.0 ) )
            throw std::invalid_argument(
                "NetworkParamsNSLI::leakingRateMax must be within "
                "interval (0,1]" );
        if ( params.leakingRateMin > params.leakingRateMax )
            throw std::invalid_argument(
                "NetworkParamsNSLI::leakingRateMin must be less then or "
                "equal to NetworkParamsNSLI::leakingRateMax" );
        if ( !( params.connectivity > 0.0f &&
                params.connectivity <= 1.0f ) )
            throw std::invalid_argument(
                "NetworkParamsNSLI::connectivity must be within "
                "interval (0,1]" );

        mWIn = Eigen::MatrixXf::Random(
            params.neuronCount, params.inputCount );

        Eigen::MatrixXf randomWeights =
            ( Eigen::MatrixXf::Random( params.neuronCount,
                params.neuronCount ).array().abs()
                    <= params.connectivity ).cast< float >() *
            Eigen::MatrixXf::Random( params.neuronCount,
                params.neuronCount ).array();
        if ( params.useOrthonormalMatrix )
        {
            auto svd = randomWeights.jacobiSvd(
                Eigen::ComputeFullU | Eigen::ComputeFullV );
            mW = ( svd.matrixU() * svd.matrixV() ).sparseView();
        }
        else
        {
            float spectralRadius =
                randomWeights.eigenvalues().cwiseAbs().maxCoeff();
            mW = ( randomWeights / spectralRadius *
                params.spectralRadius ).sparseView() ;
        }

        mWInScaling = Eigen::VectorXf::Constant( params.inputCount, 1.0f );
        mWInBias = Eigen::VectorXf::Zero( params.inputCount );

        mWOut = Eigen::MatrixXf::Zero(
            params.outputCount, params.neuronCount );

        if (params.hasOutputFeedback)
        {
            mWFB = Eigen::MatrixXf::Random(
                params.neuronCount, params.outputCount);
            mWFBScaling = Eigen::VectorXf::Constant(
                params.outputCount, 1.0f);
        }

        mLeakingRate = ( Eigen::ArrayXf::Random( params.neuronCount ) *
            ( mParams.leakingRateMax - mParams.leakingRateMin ) +
            ( mParams.leakingRateMin + mParams.leakingRateMax ) ) / 2.0f;
        mOneMinusLeakingRate = 1.0f - mLeakingRate.array();

        mIn = Eigen::VectorXf::Zero( params.inputCount );
        mX = Eigen::VectorXf::Random( params.neuronCount );
        mOut = Eigen::VectorXf::Zero( params.outputCount );
    }

    NetworkNSLI::~NetworkNSLI()
    {
    }

    void NetworkNSLI::SetInputs( const std::vector< float > & inputs )
    {
        if ( inputs.size() != mIn.rows() )
            throw std::invalid_argument( "Wrong size of the input vector" );
        mIn = (Eigen::Map<Eigen::VectorXf>(
            const_cast<float*>(inputs.data()), inputs.size()) +
            mWInBias).cwiseProduct(mWInScaling);
    }

    void NetworkNSLI::SetInputScalings(
        const std::vector< float > & scalings )
    {
        if ( scalings.size() != mParams.inputCount )
            throw std::invalid_argument(
                "Wrong size of the scalings vector" );
        mWInScaling = Eigen::Map< Eigen::VectorXf >(
            const_cast< float * >( scalings.data() ), scalings.size() );
    }

    void NetworkNSLI::SetInputBias(
        const std::vector< float > & bias )
    {
        if ( bias.size() != mParams.inputCount )
            throw std::invalid_argument(
                "Wrong size of the scalings vector" );
        mWInBias = Eigen::Map< Eigen::VectorXf >(
            const_cast< float * >( bias.data() ), bias.size() );
    }

    void NetworkNSLI::SetFeedbackScalings(
        const std::vector< float > & scalings )
    {
        if (!mParams.hasOutputFeedback)
            throw std::logic_error(
                "Trying to set up feedback scaling for a network "
                "which doesn't have an output feedback");
        if ( scalings.size() != mParams.outputCount )
            throw std::invalid_argument(
                "Wrong size of the scalings vector" );
        mWFBScaling = Eigen::Map< Eigen::VectorXf >(
            const_cast< float * >( scalings.data() ), scalings.size() );
    }

    void NetworkNSLI::Step( float step )
    {
        if ( step <= 0.0f )
            throw std::invalid_argument(
                "Step size must be positive value" );

        auto tanh = [] ( float x ) -> float { return std::tanh( x ); };

        #define TEMP mWIn * mIn + mW * mX

        #define CALC_X(val) \
            mX = mOneMinusLeakingRate.cwiseProduct(mX) + \
                mLeakingRate.cwiseProduct(val).unaryExpr(tanh)

        #define CALC_X_WITH_FB(val) \
            if (mParams.hasOutputFeedback) \
                CALC_X(TEMP + val); \
            else \
                CALC_X(TEMP)

        if ( mParams.linearOutput )
        {
            CALC_X_WITH_FB(
                mWFB * mOut.unaryExpr(tanh).cwiseProduct(mWFBScaling));
            mOut = mWOut * mX;
        }
        else
        {
            CALC_X_WITH_FB(
                mWFB * mOut.cwiseProduct(mWFBScaling));
            mOut = ( mWOut * mX ).unaryExpr( tanh );
        }

        #undef TEMP
        #undef CALC_X
        #undef CALC_X_WITH_FB

        auto isnotfinite =
            [] (float n) -> bool { return !std::isfinite(n); };
        if (mOut.unaryExpr(isnotfinite).any())
            throw OutputIsNotFinite();
    }

    void NetworkNSLI::CaptureTransformedInput(
        std::vector< float > & input )
    {
        if ( input.size() != mParams.inputCount )
            throw std::invalid_argument(
                "Size of the vector must be equal to "
                "the number of inputs" );
        for ( int i = 0; i < mParams.inputCount; ++ i )
            input[ i ] = mIn( i );
    }

    void NetworkNSLI::CaptureActivations(
        std::vector< float > & activations )
    {
        if ( activations.size() != mParams.neuronCount )
            throw std::invalid_argument(
                "Size of the vector must be equal "
                "actual number of neurons" );

        for ( int i = 0; i < mParams.neuronCount; ++ i )
            activations[ i ] = mX( i );
    }

    void NetworkNSLI::CaptureOutput( std::vector< float > & output )
    {
        if ( output.size() != mParams.outputCount )
            throw std::invalid_argument(
                "Size of the vector must be equal "
                "actual number of outputs" );

        for ( int i = 0; i < mParams.outputCount; ++ i )
            output[ i ] = mOut( i );
    }

    void NetworkNSLI::Train(
        const std::vector< std::vector< float > > & inputs,
        const std::vector< std::vector< float > > & outputs )
    {
        if ( inputs.size() == 0 )
            throw std::invalid_argument(
                "Number of samples must be not null" );
        if ( inputs.size() != outputs.size() )
            throw std::invalid_argument(
                "Number of input and output samples must be equal" );
        const unsigned kSampleCount = inputs.size();

        Eigen::MatrixXf matX( mParams.neuronCount, kSampleCount );
        Eigen::MatrixXf matY( mParams.outputCount, kSampleCount );
        for ( int i = 0; i < kSampleCount; ++ i )
        {
            SetInputs( inputs[i] );
            Step( 0.1f );
            matX.col( i ) = mX;
            matY.col( i ) = Eigen::Map< Eigen::VectorXf >(
                const_cast< float * >( outputs[i].data() ),
                    mParams.outputCount );
        }

        Eigen::MatrixXf matXT = matX.transpose();

        mWOut = ( matY * matXT * ( matX * matXT ).inverse() );
    }

    void NetworkNSLI::TrainOnline( const std::vector< float > & output,
        bool forceOutput )
    {
        for ( unsigned i = 0; i < mParams.outputCount; ++ i )
        {
            Eigen::VectorXf w = mWOut.row( i ).transpose();
            if ( mParams.linearOutput )
                mAdaptiveFilter.Train( w, mOut( i ), output[i], mX );
            else
                mAdaptiveFilter.Train( w, std::atanh( mOut( i ) ),
                    std::atanh( output[i] ), mX );
            mWOut.row( i ) = w.transpose();
        }

        if ( forceOutput )
            mOut = Eigen::Map< Eigen::VectorXf >(
                const_cast< float * >( output.data() ),
                mParams.outputCount );
    }

} // namespace ESN

#define SIZEOF_MEMBER( structure, member ) \
    sizeof( ( ( structure * ) 0 )->member )

void * esnCreateNetworkNSLI( esnNetworkParamsNSLI * params )
{
    static_assert( ( sizeof( esnNetworkParamsNSLI ) -
        SIZEOF_MEMBER( esnNetworkParamsNSLI, structSize ) ) ==
        sizeof( ESN::NetworkParamsNSLI ),
        "Wrong size of esnNetworkParamsNSLI" );

    if ( params->structSize != sizeof( esnNetworkParamsNSLI ) )
        throw std::invalid_argument(
            "esnNetworkParamsNSLI::structSize must be equal the "
            "sizeof( esnNetworkParamsNSLI )" );

    ESN::NetworkParamsNSLI p;
    std::memcpy( &p, reinterpret_cast< char * >( params ) +
        SIZEOF_MEMBER( esnNetworkParamsNSLI, structSize ),
        sizeof( ESN::NetworkParamsNSLI ) );

    return new ESN::NetworkNSLI( p );
}

#undef SIZEOF_MEMBER


#ifndef STACK_HH
#define STACK_HH

// 2020: Changes by Benjamin Huth
// - templated all classes

// StackT classes
//
// These are the low-level classes that implement feed-forward and
// recurrent neural networks. All the Eigen-dependant code in this
// library should live in this file.
//
// To keep the Eigen code out of the high-level interface, the STL ->
// Eigen ``preprocessor'' classes are also defined here.
//
// The ordering of classes is as follows:
//  - Feed-forward Stack class
//  - Feed-forward Layer classes
//  - RecurrentStackT class
//  - Recurrent layers
//  - Activation functions
//  - Various utility functions

#include "Exceptions.hh"
#include "NNLayerConfig.hh"

#include <Eigen/Dense>

#include <vector>
#include <functional>

namespace lwt {

  template<typename T>
  using VectorX = Eigen::Matrix<T, Eigen::Dynamic, 1>;

  template<typename T>
  using MatrixX = Eigen::Matrix<T, Eigen::Dynamic, Eigen::Dynamic>;

  template<typename T>
  using ArrayX = Eigen::Array<T, Eigen::Dynamic, Eigen::Dynamic>;

  template<typename T>
  class ILayerT;

  template<typename T>
  class IRecurrentLayerT;

  class FittableLWTNN;

  // ______________________________________________________________________
  // Feed forward Stack class

  template<typename T>
  class StackT
  {
    friend class FittableLWTNN;

  public:
    // constructor for dummy net
    StackT();
    // constructor for real net
    StackT(size_t n_inputs, const std::vector<LayerConfig>& layers,
          size_t skip_layers = 0);
    ~StackT();

    // make non-copyable for now
    StackT(StackT&) = delete;
    StackT& operator=(StackT&) = delete;

    VectorX<T> compute(VectorX<T>) const;
    size_t n_outputs() const;

  private:
    // returns the size of the next layer
    size_t add_layers(size_t n_inputs, const LayerConfig&);
    size_t add_dense_layers(size_t n_inputs, const LayerConfig&);
    size_t add_normalization_layers(size_t n_inputs, const LayerConfig&);
    size_t add_highway_layers(size_t n_inputs, const LayerConfig&);
    size_t add_maxout_layers(size_t n_inputs, const LayerConfig&);
    std::vector<ILayerT<T>*> m_layers;
    size_t m_n_outputs;
  };

  using Stack = StackT<double>;

  // _______________________________________________________________________
  // Feed-forward layers

  template<typename T>
  class ILayerT
  {
  public:
    virtual ~ILayerT() {}
    virtual VectorX<T> compute(const VectorX<T>&) const = 0;
  };

  using ILayer = ILayerT<double>;

  template<typename T>
  class DummyLayerT: public ILayerT<T>
  {
  public:
    virtual VectorX<T> compute(const VectorX<T>&) const override;
  };

  using DummyLayer = DummyLayerT<double>;

  template<typename T>
  class UnaryActivationLayerT: public ILayerT<T>
  {
  public:
    UnaryActivationLayerT(ActivationConfig);
    virtual VectorX<T> compute(const VectorX<T>&) const override;
  private:
    std::function<T(T)> m_func;
  };

  using UnaryActivationLayer = UnaryActivationLayerT<double>;

  template<typename T>
  class SoftmaxLayerT: public ILayerT<T>
  {
  public:
    virtual VectorX<T> compute(const VectorX<T>&) const override;
  };

  using SoftmaxLayer = SoftmaxLayerT<double>;

  template<typename T>
  class BiasLayerT: public ILayerT<T>
  {
    friend class FittableLWTNN;

  public:
    BiasLayerT(const VectorX<T>& bias);
    template<typename U> BiasLayerT(const std::vector<U>& bias);
    virtual VectorX<T> compute(const VectorX<T>&) const override;
  private:
    VectorX<T> m_bias;
  };

  using BiasLayer = BiasLayerT<double>;

  template<typename T>
  class MatrixLayerT: public ILayerT<T>
  {
    friend class FittableLWTNN;

  public:
    MatrixLayerT(const MatrixX<T>& matrix);
    virtual VectorX<T> compute(const VectorX<T>&) const override;
  private:
    MatrixX<T> m_matrix;
  };

  using MatrixLayer = MatrixLayerT<double>;

  template<typename T>
  class MaxoutLayerT: public ILayerT<T>
  {
    friend class FittableLWTNN;

  public:
    typedef std::pair<MatrixX<T>, VectorX<T>> InitUnit;
    MaxoutLayerT(const std::vector<InitUnit>& maxout_tensor);
    virtual VectorX<T> compute(const VectorX<T>&) const override;
  private:
    std::vector<MatrixX<T>> m_matrices;
    MatrixX<T> m_bias;
  };

  using MaxoutLayer = MaxoutLayerT<double>;

  /// Normalization layer ///
  /// https://arxiv.org/abs/1502.03167 ///
  template<typename T>
  class NormalizationLayerT : public ILayerT<T>
  {
    friend class FittableLWTNN;

  public:
    NormalizationLayerT(const VectorX<T>& W,const VectorX<T>& b);
    virtual VectorX<T> compute(const VectorX<T>&) const override;

  private:
    VectorX<T> _W;
    VectorX<T> _b;

  };

  using NormalizationLayer = NormalizationLayerT<double>;

  //http://arxiv.org/pdf/1505.00387v2.pdf
  template<typename T>
  class HighwayLayerT: public ILayerT<T>
  {
    friend class FittableLWTNN;

  public:
    HighwayLayerT(const MatrixX<T>& W,
                 const VectorX<T>& b,
                 const MatrixX<T>& W_carry,
                 const VectorX<T>& b_carry,
                 ActivationConfig activation);
    virtual VectorX<T> compute(const VectorX<T>&) const override;
  private:
    MatrixX<T> m_w_t;
    VectorX<T> m_b_t;
    MatrixX<T> m_w_c;
    VectorX<T> m_b_c;
    std::function<T(T)> m_act;
  };

  using HighwayLayer = HighwayLayerT<double>;

  // ______________________________________________________________________
  // Recurrent StackT

  template<typename T>
  class RecurrentStackT
  {
  public:
    RecurrentStackT(size_t n_inputs, const std::vector<LayerConfig>& layers);
    ~RecurrentStackT();
    RecurrentStackT(RecurrentStackT&) = delete;
    RecurrentStackT& operator=(RecurrentStackT&) = delete;
    MatrixX<T> scan(MatrixX<T> inputs) const;
    size_t n_outputs() const;
  private:
    std::vector<IRecurrentLayerT<T>*> m_layers;
    size_t add_lstm_layers(size_t n_inputs, const LayerConfig&);
    size_t add_gru_layers(size_t n_inputs, const LayerConfig&);
    size_t add_embedding_layers(size_t n_inputs, const LayerConfig&);
    size_t m_n_outputs;
  };

  using RecurrentStack = RecurrentStackT<double>;

  // This is the old RecurrentStack. Should probably absorb this into
  // the high-level interface in LightweightRNN, since all it does is
  // provide a slightly higher-level interface to a network which
  // combines recurrent + ff layers.
  template<typename T>
  class ReductionStackT
  {
  public:
    ReductionStackT(size_t n_in, const std::vector<LayerConfig>& layers);
    ~ReductionStackT();
    ReductionStackT(ReductionStackT&) = delete;
    ReductionStackT& operator=(ReductionStackT&) = delete;
    VectorX<T> reduce(MatrixX<T> inputs) const;
    size_t n_outputs() const;
  private:
    RecurrentStackT<T>* m_recurrent;
    StackT<T>* m_stack;
  };

  using ReductionStack = ReductionStackT<double>;

  // __________________________________________________________________
  // Recurrent layers

  template<typename T>
  class IRecurrentLayerT
  {
  public:
    virtual ~IRecurrentLayerT() {}
    virtual MatrixX<T> scan( const MatrixX<T>&) const = 0;
  };

  using IRecurrentLayer = IRecurrentLayerT<double>;

  template<typename T>
  class EmbeddingLayerT : public IRecurrentLayerT<T>
  {
  public:
    EmbeddingLayerT(int var_row_index, MatrixX<T> W);
    virtual ~EmbeddingLayerT() {};
    virtual MatrixX<T> scan( const MatrixX<T>&) const override;

  private:
    int m_var_row_index;
    MatrixX<T> m_W;
  };

  using EmbeddingLayer = EmbeddingLayerT<double>;

  /// long short term memory ///
  template<typename T> struct LSTMStateT;
  using LSTMState = LSTMStateT<double>;

  template<typename T>
  class LSTMLayerT : public IRecurrentLayerT<T>
  {
  public:
    LSTMLayerT(ActivationConfig activation,
              ActivationConfig inner_activation,
              MatrixX<T> W_i, MatrixX<T> U_i, VectorX<T> b_i,
              MatrixX<T> W_f, MatrixX<T> U_f, VectorX<T> b_f,
              MatrixX<T> W_o, MatrixX<T> U_o, VectorX<T> b_o,
              MatrixX<T> W_c, MatrixX<T> U_c, VectorX<T> b_c);

    virtual ~LSTMLayerT() {};
    virtual MatrixX<T> scan( const MatrixX<T>&) const override;
    void step( const VectorX<T>& input, LSTMState& ) const;

  private:
    std::function<T(T)> m_activation_fun;
    std::function<T(T)> m_inner_activation_fun;

    MatrixX<T> m_W_i;
    MatrixX<T> m_U_i;
    VectorX<T> m_b_i;

    MatrixX<T> m_W_f;
    MatrixX<T> m_U_f;
    VectorX<T> m_b_f;

    MatrixX<T> m_W_o;
    MatrixX<T> m_U_o;
    VectorX<T> m_b_o;

    MatrixX<T> m_W_c;
    MatrixX<T> m_U_c;
    VectorX<T> m_b_c;

    int m_n_outputs;
  };

  using LSTMLayer = LSTMLayerT<double>;

  /// gated recurrent unit ///
  template<typename T> struct GRUStateT;
  using GRUState = GRUStateT<double>;

  template<typename T>
  class GRULayerT : public IRecurrentLayerT<T>
  {
  public:
    GRULayerT(ActivationConfig activation,
             ActivationConfig inner_activation,
             MatrixX<T> W_z, MatrixX<T> U_z, VectorX<T> b_z,
             MatrixX<T> W_r, MatrixX<T> U_r, VectorX<T> b_r,
             MatrixX<T> W_h, MatrixX<T> U_h, VectorX<T> b_h);

    virtual ~GRULayerT() {};
    virtual MatrixX<T> scan( const MatrixX<T>&) const override;
    void step( const VectorX<T>& input, GRUState& ) const;

  private:
    std::function<T(T)> m_activation_fun;
    std::function<T(T)> m_inner_activation_fun;

    MatrixX<T> m_W_z;
    MatrixX<T> m_U_z;
    VectorX<T> m_b_z;

    MatrixX<T> m_W_r;
    MatrixX<T> m_U_r;
    VectorX<T> m_b_r;

    MatrixX<T> m_W_h;
    MatrixX<T> m_U_h;
    VectorX<T> m_b_h;

    int m_n_outputs;
  };

  using GRULayer = GRULayerT<double>;

  // ______________________________________________________________________
  // Activation functions

  // note that others are supported but are too simple to
  // require a special function
  template<typename T> T nn_sigmoidT( T x );
  template<typename T> T nn_hard_sigmoidT( T x );
  template<typename T> T nn_tanhT( T x );
  template<typename T> T nn_reluT( T x );

  double nn_sigmoid( double x );
  double nn_hard_sigmoid( double x );
  double nn_tanh( double x );
  double nn_relu( double x );

  template<typename T>
  class ELUT
  {
  public:
    ELUT(T alpha);
    T operator()(T) const;
  private:
    T m_alpha;
  };

  using ELU = ELUT<double>;

  template<typename T>
  class LeakyReLUT
  {
  public:
    LeakyReLUT(T alpha);
    T operator()(T) const;
  private:
    T m_alpha;
  };

  using LeakyReLU = LeakyReLUT<double>;

  template<typename T>
  class SwishT
  {
  public:
    SwishT(T alpha);
    T operator()(T) const;
  private:
    T m_alpha;
  };

  using Swish = SwishT<double>;

  template<typename T> std::function<T(T)> get_activationT(lwt::ActivationConfig);
  std::function<double(double)> get_activation(lwt::ActivationConfig);

  // WARNING: you own this pointer! Only call when assigning to member data!
  template<typename T> ILayerT<T>* get_raw_activation_layerT(ActivationConfig);
  ILayer* get_raw_activation_layer(ActivationConfig);

  // ______________________________________________________________________
  // utility functions

  // functions to build up basic units from vectors
  template<typename T1, typename T2> MatrixX<T1> build_matrixT(const std::vector<T2>& weights, size_t n_inputs);
  template<typename T1, typename T2> VectorX<T1> build_vectorT(const std::vector<T2>& bias);

  MatrixX<double> build_matrix(const std::vector<double>& weights, size_t n_inputs);
  VectorX<double> build_vector(const std::vector<double>& bias);

  // consistency checks
  void throw_if_not_maxout(const LayerConfig& layer);
  void throw_if_not_dense(const LayerConfig& layer);
  void throw_if_not_normalization(const LayerConfig& layer);

  // LSTM component for convenience in some layers
  template<typename T>
  struct DenseComponentsT
  {
    MatrixX<T> W;
    MatrixX<T> U;
    VectorX<T> b;
  };

  using DenseComponents = DenseComponentsT<double>;

  template<typename T> DenseComponentsT<T> get_componentT(const lwt::LayerConfig& layer, size_t n_in);
  DenseComponents get_component(const lwt::LayerConfig& layer, size_t n_in);
}

#include "Stack.txx"

#endif // STACK_HH


#include "quadrature/angular/level_symmetric_gaussian.h"

#include <deal.II/base/exceptions.h>
#include <deal.II/base/quadrature_lib.h>

namespace bart {

namespace quadrature {

namespace angular {

LevelSymmetricGaussian::LevelSymmetricGaussian(bart::quadrature::Order order)
    : order_(order.get()) {
  AssertThrow(order_ >= 2,
              dealii::ExcMessage("Error in constructor of "
                                 "LevelSymmetricGaussian order must be >= 2"))
  AssertThrow(order_ % 2 == 0,
              dealii::ExcMessage("Error in constructor of "
                                 "LevelSymmetricGaussian order must be even"))
  AssertThrow(order_ <= 16,
              dealii::ExcMessage("Error in constructor of "
                                 "LevelSymmetricGaussian order must be <= 16"))
}

std::vector<std::pair<CartesianPosition<3>, Weight>>
LevelSymmetricGaussian::GenerateSet() const {
  std::vector<std::pair<CartesianPosition<3>, Weight>> generated_set;
  const int n_points = order_/2;

  // Gaussian quadrature for theta, the lowest level has order_ points but we
  // will only use half of the quadrature
  dealii::QGauss<1> gaussian_quadrature(order_);

  // For each quadrature point in the lowest level, we will create equally
  // spaced points in phi
  for (int level = 0; level < n_points; ++level) {
    int n_points_this_level = level + 1;
    // mu is just the gaussian quadrature point
    double mu = 1 - gaussian_quadrature.point(level)[0] * 2;
    // points in phi will be equally spaced from 0 to pi/2
    double dphi = M_PI/(2*n_points_this_level);
    // weights on this level are equal to the gaussian times 2PI/total points in
    // 2PI. Here we have n_points_this_level*4 total points
    double weight =
        gaussian_quadrature.weight(level) * M_PI/(n_points_this_level);

    for (int j = 0; j < n_points_this_level; ++j) {
      double phi = (j + 0.5)*dphi;
      double x = std::sqrt(1 - mu*mu) * std::cos(phi);
      double y = std::sqrt(1 - mu*mu) * std::sin(phi);
      double z = mu;
      std::array<double, 3> position{x,y,z};
      generated_set.emplace_back(std::make_pair(CartesianPosition<3>(position),
                                                Weight(weight)));
    }
  }
  return generated_set;
}

} // namespace angular

} // namespace quadrature

} //namespace bart