/**
 * \file
 * \author Thomas Fischer
 * \date   2011-03-17
 * \brief  Implementation of the AngleSkewMetric class.
 *
 * \copyright
 * Copyright (c) 2012-2020, OpenGeoSys Community (http://www.opengeosys.org)
 *            Distributed under a Modified BSD License.
 *              See accompanying file LICENSE.txt or
 *              http://www.opengeosys.org/project/license
 *
 */

#include "AngleSkewMetric.h"

#include <cmath>
#include <boost/math/constants/constants.hpp>

#include "MeshLib/Node.h"

#include "MathLib/MathTools.h"

using namespace boost::math::double_constants;

namespace MeshLib
{
AngleSkewMetric::AngleSkewMetric(Mesh const& mesh) :
    ElementQualityMetric(mesh)
{}

void AngleSkewMetric::calculateQuality ()
{
    const std::vector<MeshLib::Element*>& elements(_mesh.getElements());
    const std::size_t nElements (_mesh.getNumberOfElements());

    for (std::size_t k(0); k < nElements; k++)
    {
        Element const& elem (*elements[k]);
        switch (elem.getGeomType())
        {
        case MeshElemType::LINE:
            _element_quality_metric[k] = -1.0;
            break;
        case MeshElemType::TRIANGLE:
            _element_quality_metric[k] = checkTriangle (elem);
            break;
        case MeshElemType::QUAD:
            _element_quality_metric[k] = checkQuad (elem);
            break;
        case MeshElemType::TETRAHEDRON:
            _element_quality_metric[k] = checkTetrahedron (elem);
            break;
        case MeshElemType::HEXAHEDRON:
            _element_quality_metric[k] = checkHexahedron (elem);
            break;
        case MeshElemType::PRISM:
            _element_quality_metric[k] = checkPrism (elem);
            break;
        default:
            break;
        }
    }
}

double AngleSkewMetric::checkTriangle (Element const& elem) const
{
    double const* const node0 (elem.getNode(0)->getCoords());
    double const* const node1 (elem.getNode(1)->getCoords());
    double const* const node2 (elem.getNode(2)->getCoords());

    double min_angle(two_pi);
    double max_angle(0.0);
    getMinMaxAngleFromTriangle (node0, node1, node2, min_angle, max_angle);

    return 1.0 -
           std::max((max_angle - third_pi) / two_thirds_pi,
                    (third_pi - min_angle) / third_pi);
}

double AngleSkewMetric::checkQuad (Element const& elem) const
{
    double const* const node0 (elem.getNode(0)->getCoords());
    double const* const node1 (elem.getNode(1)->getCoords());
    double const* const node2 (elem.getNode(2)->getCoords());
    double const* const node3 (elem.getNode(3)->getCoords());

    double min_angle (two_pi);
    double max_angle (0.0);

    getMinMaxAngleFromQuad (node0, node1, node2, node3, min_angle, max_angle);

    return 1.0 -
           std::max((max_angle - two_pi) / (-pi), (two_pi - min_angle) / (two_pi));
}

double AngleSkewMetric::checkTetrahedron (Element const& elem) const
{
    double const* const node0 (elem.getNode(0)->getCoords());
    double const* const node1 (elem.getNode(1)->getCoords());
    double const* const node2 (elem.getNode(2)->getCoords());
    double const* const node3 (elem.getNode(3)->getCoords());

    double min_angle (two_pi);
    double max_angle (0.0);

    // first triangle (0,1,2)
    getMinMaxAngleFromTriangle(node0, node1, node2, min_angle, max_angle);
    // second triangle (0,1,3)
    getMinMaxAngleFromTriangle(node0, node1, node3, min_angle, max_angle);
    // third triangle (0,2,3)
    getMinMaxAngleFromTriangle(node0, node2, node3, min_angle, max_angle);
    // fourth triangle (1,2,3)
    getMinMaxAngleFromTriangle(node1, node2, node3, min_angle, max_angle);

    return 1.0 - std::max((max_angle - two_pi) / two_thirds_pi,
                          (third_pi - min_angle) / third_pi);
}

double AngleSkewMetric::checkHexahedron (Element const& elem) const
{
    double const* const node0 (elem.getNode(0)->getCoords());
    double const* const node1 (elem.getNode(1)->getCoords());
    double const* const node2 (elem.getNode(2)->getCoords());
    double const* const node3 (elem.getNode(3)->getCoords());
    double const* const node4 (elem.getNode(4)->getCoords());
    double const* const node5 (elem.getNode(5)->getCoords());
    double const* const node6 (elem.getNode(6)->getCoords());
    double const* const node7 (elem.getNode(7)->getCoords());

    double min_angle (two_pi);
    double max_angle (0.0);

    // first surface (0,1,2,3)
    getMinMaxAngleFromQuad (node0, node1, node2, node3, min_angle, max_angle);
    // second surface (0,3,7,4)
    getMinMaxAngleFromQuad (node0, node3, node7, node4, min_angle, max_angle);
    // third surface (4,5,6,7)
    getMinMaxAngleFromQuad (node4, node5, node6, node7, min_angle, max_angle);
    // fourth surface (5,1,2,6)
    getMinMaxAngleFromQuad (node5, node1, node2, node6, min_angle, max_angle);
    // fifth surface (5,1,0,4)
    getMinMaxAngleFromQuad (node5, node1, node0, node4, min_angle, max_angle);
    // sixth surface (6,2,3,7)
    getMinMaxAngleFromQuad (node6, node2, node3, node7, min_angle, max_angle);

    return 1.0 -
           std::max((max_angle - two_pi) / (-pi), (two_pi - min_angle) / two_pi);
}

double AngleSkewMetric::checkPrism (Element const& elem) const
{
    double const* const node0 (elem.getNode(0)->getCoords());
    double const* const node1 (elem.getNode(1)->getCoords());
    double const* const node2 (elem.getNode(2)->getCoords());
    double const* const node3 (elem.getNode(3)->getCoords());
    double const* const node4 (elem.getNode(4)->getCoords());
    double const* const node5 (elem.getNode(5)->getCoords());

    double min_angle_tri (two_pi);
    double max_angle_tri (0.0);

    // first triangle (0,1,2)
    getMinMaxAngleFromTriangle (node0, node1, node2, min_angle_tri, max_angle_tri);
    // second surface (3,4,5)
    getMinMaxAngleFromTriangle (node3, node4, node5, min_angle_tri, max_angle_tri);

    double tri_criterion (1.0 - std::max((max_angle_tri - two_pi) / two_thirds_pi,
                                         (third_pi - min_angle_tri) / third_pi));

    double min_angle_quad (two_pi);
    double max_angle_quad (0.0);
    // surface (0,3,4,1)
    getMinMaxAngleFromQuad (node0, node3, node4, node1, min_angle_quad, max_angle_quad);
    // surface (2,5,3,0)
    getMinMaxAngleFromQuad (node2, node5, node3, node0, min_angle_quad, max_angle_quad);
    // surface (1,2,5,4)
    getMinMaxAngleFromQuad (node1, node2, node5, node4, min_angle_quad, max_angle_quad);

    double quad_criterion (1.0 - std::max((max_angle_quad - two_pi) / (-pi),
                                          (two_pi - min_angle_quad) / two_pi));

    return std::min (tri_criterion, quad_criterion);
}

void AngleSkewMetric::getMinMaxAngleFromQuad (
        double const* const n0, double const* const n1,
        double const* const n2, double const* const n3,
        double &min_angle, double &max_angle) const
{
    const double* nodes[4] = {n0, n1, n2, n3};
    for (unsigned i=0; i<4; ++i)
    {
        const double angle (MathLib::getAngle (nodes[i], nodes[(i+1)%4], nodes[(i+2)%4]));
        min_angle = std::min(angle, min_angle);
        max_angle = std::max(angle, max_angle);
    }
}

void AngleSkewMetric::getMinMaxAngleFromTriangle(double const* const n0,
                                                          double const* const n1,
                                                          double const* const n2,
                                                          double &min_angle,
                                                          double &max_angle) const
{
    const double* nodes[3] = {n0, n1, n2};
    for (unsigned i=0; i<3; ++i)
    {
        const double angle (MathLib::getAngle (nodes[i], nodes[(i+1)%3], nodes[(i+2)%3]));
        min_angle = std::min(angle, min_angle);
        max_angle = std::max(angle, max_angle);
    }
}

} // end namespace MeshLib


// Copyright (c)  Mingcheng Zuo, Dietmar Wolz.
//
// This source code is licensed under the MIT license found in the
// LICENSE file in the root directory.

// Eigen based implementation of differential evolution (GCL-DE) derived from
// "A case learning-based differential evolution algorithm for global optimization of interplanetary trajectory design,
//  Mingcheng Zuo, Guangming Dai, Lei Peng, Maocai Wang, Zhengquan Liu", https://doi.org/10.1016/j.asoc.2020.106451

#include <Eigen/Core>
#include <iostream>
#include <float.h>
#include <ctime>
#include <random>
#include "pcg_random.hpp"
#include "call_java.hpp"

using namespace std;

typedef Eigen::Matrix<double, Eigen::Dynamic, 1> vec;
typedef Eigen::Matrix<int, Eigen::Dynamic, 1> ivec;
typedef Eigen::Matrix<double, Eigen::Dynamic, Eigen::Dynamic> mat;

typedef void (*callback_parallel)(int, int, double[], double[]);

namespace cl_differential_evolution {

static uniform_real_distribution<> distr_01 = std::uniform_real_distribution<>(
        0, 1);
static normal_distribution<> gauss_01 = std::normal_distribution<>(0, 1);

static double normreal(pcg64 *rs, double mu, double sdev) {
    return gauss_01(*rs) * sdev + mu;
}

static vec zeros(int n) {
    return Eigen::MatrixXd::Zero(n, 1);
}

static vec constant(int n, double val) {
    return Eigen::MatrixXd::Constant(n, 1, val);
}

static Eigen::MatrixXd uniformVec(int dim, pcg64 &rs) {
    return Eigen::MatrixXd::NullaryExpr(dim, 1, [&]() {
        return distr_01(rs);
    });
}

static Eigen::MatrixXd uniform(int dx, int dy, pcg64 &rs) {
    return Eigen::MatrixXd::NullaryExpr(dx, dy, [&]() {
        return distr_01(rs);
    });
}

struct IndexVal {
    int index;
    double val;
};

static bool compareIndexVal(IndexVal i1, IndexVal i2) {
    return (i1.val < i2.val);
}

static ivec sort_index(const vec &x) {
    int size = x.size();
    IndexVal ivals[size];
    for (int i = 0; i < size; i++) {
        ivals[i].index = i;
        ivals[i].val = x[i];
    }
    std::sort(ivals, ivals + size, compareIndexVal);
    return Eigen::MatrixXi::NullaryExpr(size, 1, [&ivals](int i) {
        return ivals[i].index;
    });
}

static ivec sort_index(const vec &y, int *indices, int size) {
    IndexVal ivals[size];
    for (int i = 0; i < size; i++) {
        ivals[i].index = indices[i];
        ivals[i].val = y[indices[i]];
    }
    std::sort(ivals, ivals + size, compareIndexVal);
    return Eigen::MatrixXi::NullaryExpr(size, 1, [&ivals](int i) {
        return ivals[i].index;
    });
}

// wrapper around the fitness function, scales according to boundaries

class Fitness {

public:

    Fitness(CallJava *func_par_, const vec &lower_limit,
            const vec &upper_limit) {
        func_par = func_par_;
        lower = lower_limit;
        upper = upper_limit;
        evaluationCounter = 0;
        if (lower.size() > 0) // bounds defined
            scale = (upper - lower);
    }

    vec norm(vec &X) {
        return (X - lower).array() / scale.array();
    }

    vec getClosestFeasible(const vec &X) const {
        if (lower.size() > 0) {
            return X.cwiseMin(upper).cwiseMax(lower);
        }
        return X;
    }

    void values(const mat &popX, vec &ys) {
        int popsize = popX.cols();
        int n = popX.rows();
        double pargs[popsize * n];
        double res[popsize];
        for (int p = 0; p < popX.cols(); p++) {
            for (int i = 0; i < n; i++)
                pargs[p * n + i] = popX(i, p);
        }
        func_par->evalJava(popsize, n, pargs, res);
        for (int p = 0; p < popX.cols(); p++)
            ys[p] = res[p];
        evaluationCounter += popsize;
    }

    bool feasible(int i, double x) {
        return x >= lower[i] && x <= upper[i];
    }

    vec uniformX(pcg64 &rs) {
        vec rv = uniformVec(lower.size(), rs);
        return (rv.array() * scale.array()).matrix() + lower;
    }

    double uniformXi(int i, pcg64 &rs) {
        return lower[i] + scale[i] * distr_01(rs);
    }

    int getEvaluations() {
        return evaluationCounter;
    }

    vec lower;
    vec upper;

private:
    CallJava *func_par;
    long evaluationCounter;
    vec scale;
};

class ClDeOptimizer {

public:

    ClDeOptimizer(long runid_, Fitness *fitfun_, int dim_, int seed_,
            int popsize_, int maxEvaluations_, double pbest_,
            double stopfitness_, double K1_, double K2_) {
        // runid used to identify a specific run
        runid = runid_;
        // fitness function to minimize
        fitfun = fitfun_;
        // Number of objective variables/problem dimension
        dim = dim_;
        // Population size
        popsize0 = popsize_ > 0 ? popsize_ : int(dim * 8.5 + 150);
        // maximal number of evaluations allowed.
        maxEvaluations = maxEvaluations_;
        // use low value 0 < pbest <= 1 to narrow search.
        pbest0 = pbest_;
        // Limit for fitness value.
        stopfitness = stopfitness_;
        K1 = K1_;
        K2 = K2_;
        // stop criteria
        stop = 0;
        rs = new pcg64(seed_);
        init();
    }

    ~ClDeOptimizer() {
        delete rs;
    }

    double rnd01() {
        return distr_01(*rs);
    }

    int rndInt(int max) {
        return (int) (max * distr_01(*rs));
    }

    void doOptimize() {

        double CR, F;
        vector<vec> sp;

        double stage = 0.5;

        int popsize = popsize0;
        double pbest = pbest0;

        // -------------------- Generation Loop --------------------------------

        for (iterations = 1;; iterations++) {
            // sort population
            ivec sindex = sort_index(nextY);
            popY = nextY(sindex, Eigen::all);
            popX = nextX(Eigen::all, sindex);

            bestX = popX.col(0);
            bestY = popY[0];

            if (isfinite(stopfitness) && bestY < stopfitness) {
                stop = 1;
                return;
            }

            if (fitfun->getEvaluations() >= maxEvaluations)
                return;

            double evals = float(fitfun->getEvaluations()) / maxEvaluations;
            if (evals > stage)
                pbest = pbest0;
            else {
                for (double per = 1.0; per >= 0; per = per - 0.05)
                    if (evals * evals < per && evals * evals >= per - 0.05)
                        pbest = 1 - per;
            }
            popsize = min(popsize0,
                    max(7 * dim, int(popsize0 - (popsize0 - 50) * evals)));

            mat local_upper(dim, popsize);
            mat local_lower(dim, popsize);
            double bound_info[popsize];
            for (int p = 0; p < popsize; p++) {
                vec X = popX.col(p);
                local_upper.col(p) = fitfun->upper.cwiseMin(X);
                local_lower.col(p) = fitfun->lower.cwiseMax(X);
                vec norm = fitfun->norm(X);
                double avnorm = norm.sum() / popsize;
                bound_info[p] = 0;
                for (int j = 0; j < dim; j++)
                    bound_info[p] += abs(norm[j] - avnorm);
                bound_info[p] /= popsize;
            }

            for (int p = 0; p < popsize; p++) {
                int r1, r2, r3;
                do {
                    r1 = rndInt(popsize);
                } while (r1 == p);
                do {
                    r2 = rndInt(int(popsize * pbest));
                } while (r2 == p || r2 == r1);
                do {
                    r3 = rndInt(popsize + sp.size());
                } while (r3 == p || r3 == r2 || r3 == r1);
                int jr = rndInt(dim);

                if (iterations % 2 == 1)
//					CR = 1;
                    CR = normreal(rs, 0.95, 0.01);
                else
//					CR = 0;
                    CR = normreal(rs, 0.0, 0.01);

                mat mutationX;
                vec mutationY;
                int mutationI[3];
                if (evals > stage) {
                    mutationI[0] = r1;
                    mutationI[1] = r2;
                    mutationI[2] = r3 < popsize ? r3 : r3 - popsize;
                    ivec mindex = sort_index(nextY, mutationI, 3);
                    mutationY = nextY(mindex, Eigen::all);
                    mutationX = nextX(Eigen::all, mindex);
                    F = 2 * (mutationY[1] - mutationY[0])
                            / (mutationY[2] - mutationY[0]);
                } else {
                    if (rnd01() < 0.5)
                        F = normreal(rs, 0.1, 0.04);
                    else
                        F = normreal(rs, 1.0, 1.0);
                    if (F < 0 || F > 1)
                        F = rnd01();
                }
                vec ui = popX.col(p);
                vec ub = local_upper.col(p);
                vec lb = local_lower.col(p);
                for (int j = 0; j < dim; j++) {
                    if (j == jr || rnd01() < CR) {
                        if (bound_info[j] > 0.1 && evals < K1 && rnd01() > K2)
                            ui[j] = ub[j] + lb[j] - popX(j, r1);
                        else {
                            if (evals > stage) {
                                ui[j] = mutationX(j, 0)
                                        + F
                                                * (mutationX(j, 1)
                                                        - mutationX(j, 2));
                            } else {
                                if (r3 < popsize)
                                    ui[j] = popX(j, r1)
                                            + F * (popX(j, r2) - popX(j, r3));
                                else
                                    ui[j] =
                                            popX(j, r1)
                                                    + F
                                                            * ((popX)(j, r2)
                                                                    - sp[r3
                                                                            - popsize][j]);
                            }
                        }
                        if (!fitfun->feasible(j, ui[j]))
                            ui[j] = fitfun->uniformXi(j, *rs);
                    }
                }
                nextX.col(p) = ui;
            }
            fitfun->values(nextX, nextY);
            for (int p = 0; p < popsize; p++) {
                if (nextY[p] < popY[p]) {
                    if (sp.size() < popsize)
                        sp.push_back(popX.col(p));
                    else
                        sp[rndInt(popsize)] = popX.col(p);
                } else {    // no improvement, copy from parent
                    nextX.col(p) = popX.col(p);
                    nextY[p] = popY[p];
                }
            }
        }
    }

    void init() {
        popCR = zeros(popsize0);
        popF = zeros(popsize0);
        nextX = mat(dim, popsize0);
        for (int p = 0; p < popsize0; p++)
            nextX.col(p) = fitfun->uniformX(*rs);
        nextY = vec(popsize0);
        fitfun->values(nextX, nextY);
    }

    vec getBestX() {
        return bestX;
    }

    double getBestValue() {
        return bestY;
    }

    double getIterations() {
        return iterations;
    }

    double getStop() {
        return stop;
    }

private:
    long runid;
    Fitness *fitfun;
    int popsize0; // population size
    int dim;
    int maxEvaluations;
    double pbest0;
    double stopfitness;
    int iterations;
    double bestY;
    vec bestX;
    int stop;
    double K1;
    double K2;
    pcg64 *rs;
    mat popX;
    vec popY;
    mat nextX;
    vec nextY;
    vec popCR;
    vec popF;
};
}

using namespace cl_differential_evolution;

/*
 * Class:     fcmaes_core_Jni
 * Method:    optimizeCLDE
 * Signature: (Lfcmaes/core/Fitness;[D[D[DIDIDDDJI)I
 */
JNIEXPORT jint JNICALL Java_fcmaes_core_Jni_optimizeCLDE(JNIEnv *env,
        jclass cls, jobject func, jdoubleArray jlower, jdoubleArray jupper,
        jdoubleArray jinit, jint maxEvals, jdouble stopfitness, jint popsize,
        jdouble pbest, jdouble K1, jdouble K2, jlong seed, jint runid) {
    // init java function callback
    double *init = env->GetDoubleArrayElements(jinit, JNI_FALSE);
    double *lower = env->GetDoubleArrayElements(jlower, JNI_FALSE);
    double *upper = env->GetDoubleArrayElements(jupper, JNI_FALSE);
    int dim = env->GetArrayLength(jinit);

    vec lower_limit(dim), upper_limit(dim);
    bool useLimit = false;
    for (int i = 0; i < dim; i++) {
        lower_limit[i] = lower[i];
        upper_limit[i] = upper[i];
        useLimit |= (lower[i] != 0);
        useLimit |= (upper[i] != 0);
    }
    if (useLimit == false) {
        lower_limit.resize(0);
        upper_limit.resize(0);
    }
    CallJava callJava(func, env);
    Fitness fitfun(&callJava, lower_limit, upper_limit);
    ClDeOptimizer opt(runid, &fitfun, dim, seed, popsize, maxEvals, pbest,
            stopfitness, K1, K2);
    try {
        opt.doOptimize();
        vec bestX = opt.getBestX();
        double bestY = opt.getBestValue();

        for (int i = 0; i < dim; i++)
            init[i] = bestX[i];

        env->SetDoubleArrayRegion(jinit, 0, dim, (jdouble*) init);
        env->ReleaseDoubleArrayElements(jinit, init, 0);
        env->ReleaseDoubleArrayElements(jupper, upper, 0);
        env->ReleaseDoubleArrayElements(jlower, lower, 0);
        return fitfun.getEvaluations();

    } catch (std::exception &e) {
        cout << e.what() << endl;
        return fitfun.getEvaluations();
    }
    return 0;
}


//==================================================================================================
/**

  Copyright 2016 NumScale SAS

  Distributed under the Boost Software License, Version 1.0.
  (See accompanying file LICENSE.md or copy at http://boost.org/LICENSE_1_0.txt)
*/
//==================================================================================================
#ifndef BOOST_SIMD_ARCH_COMMON_SIMD_FUNCTION_SINHC_HPP_INCLUDED
#define BOOST_SIMD_ARCH_COMMON_SIMD_FUNCTION_SINHC_HPP_INCLUDED

#include <boost/simd/detail/overload.hpp>
#include <boost/simd/arch/common/detail/generic/sinhc_kernel.hpp>
#include <boost/simd/meta/as_logical.hpp>
#include <boost/simd/constant/half.hpp>
#include <boost/simd/constant/log_2.hpp>
#include <boost/simd/detail/constant/maxlog.hpp>
#include <boost/simd/constant/one.hpp>
#include <boost/simd/constant/zero.hpp>
#include <boost/simd/function/abs.hpp>
#include <boost/simd/function/average.hpp>
#include <boost/simd/function/divides.hpp>
#include <boost/simd/function/exp.hpp>
#include <boost/simd/function/if_else.hpp>
#include <boost/simd/function/nbtrue.hpp>
#include <boost/simd/function/is_greater.hpp>
#include <boost/simd/function/is_less.hpp>
#include <boost/simd/function/minus.hpp>
#include <boost/simd/function/multiplies.hpp>
#include <boost/simd/function/rec.hpp>
#include <boost/simd/function/sqr.hpp>
#include <boost/simd/function/unary_minus.hpp>

#ifndef BOOST_SIMD_NO_INFINITIES
#include <boost/simd/constant/inf.hpp>
#endif

namespace boost { namespace simd { namespace ext
{
   namespace bd = boost::dispatch;
   namespace bs = boost::simd;
   BOOST_DISPATCH_OVERLOAD_IF( sinhc_
                          , (typename A0, typename X)
                          , (detail::is_native<X>)
                          , bd::cpu_
                          , bs::pack_<bd::floating_<A0>, X>
                          )
   {
      BOOST_FORCEINLINE A0 operator()( const A0& a0) const BOOST_NOEXCEPT
      {
        //////////////////////////////////////////////////////////////////////////////
        // if x = abs(a0) is less than 1 sinhc is computed using a polynomial(float)
        // respectively rational(double) approx inspired from cephes sinh approx.
        // else according x < Threshold e =  exp(x) or exp(x/2) is respectively
        // computed
        // * in the first case sinh is ((e-rec(e))/2)/x
        // * in the second     sinh is (e/2/x)*e (avoiding undue overflow)
        // Threshold is Maxlog - Log_2
        //////////////////////////////////////////////////////////////////////////////
        A0 x = bs::abs(a0);
        auto lt1= is_less(x, One<A0>());
        std::size_t nb = nbtrue(lt1);
        A0 z = Zero<A0>();
        if( nb > 0)
        {
          z = detail::sinhc_kernel<A0>::compute(sqr(x));
          if(nb >=A0::static_size) return z;
        }
        auto test1 = is_greater(x, Maxlog<A0>()-Log_2<A0>());
        A0 fac = if_else(test1, Half<A0>(), One<A0>());
        A0 tmp = exp(x*fac);
        A0 tmp1 = (Half<A0>()*tmp)/x;
        A0 r =  if_else(test1, tmp1*tmp, average(tmp, -rec(tmp))/x);
        #ifndef BOOST_SIMD_NO_INFINITIES
        r = if_else(is_equal(x, Inf<A0>()), x, r);
        #endif
        return if_else(lt1, z, r);
      }
   };

} } }

#endif


#include <algorithm>
#include <unordered_map>

#include <Eigen/QR>
#include <descriptor-projection/build-projection-matrix.h>
#include <descriptor-projection/descriptor-projection.h>
#include <descriptor-projection/flags.h>

#include <vi-map/vertex.h>

namespace descriptor_projection {

// Compute the covariance of the descriptors, rows are states, columns are
// samples.
void ComputeCovariance(
    const Eigen::MatrixXf& data, Eigen::MatrixXf* covariance) {
  CHECK_NOTNULL(covariance);
  CHECK_GT(data.cols(), 0) << "Data must not be empty!";
  VLOG(4) << "Got " << data.cols()
          << " samples to compute the covariance from.";
  covariance->setZero(data.rows(), data.rows());
  constexpr int kBlockSize = 10000;
  const int num_blocks = data.cols() / kBlockSize + 1;
  for (int i = 0; i < num_blocks; ++i) {
    const int block_start = i * kBlockSize;
    const int block_size =
        std::min<int>((i + 1) * kBlockSize, data.cols()) - block_start;
    const Eigen::Block<const Eigen::MatrixXf>& data_block =
        data.block(0, block_start, data.rows(), block_size);

    const Eigen::MatrixXf centered =
        data_block.colwise() - data_block.rowwise().mean();
    double normalizer = std::max(static_cast<int>(data_block.cols() - 1), 1);
    covariance->noalias() += (centered * centered.adjoint()) / normalizer;
  }
  (*covariance) /= num_blocks;
}

void BuildListOfMatchesAndNonMatches(
    const Eigen::MatrixXf& all_descriptors, const std::vector<Track>& tracks,
    std::vector<descriptor_projection::DescriptorMatch>* matches,
    std::vector<descriptor_projection::DescriptorMatch>* non_matches) {
  CHECK_NOTNULL(matches);
  CHECK_NOTNULL(non_matches);
  for (const Track& track : tracks) {
    // Add pairs of matching descriptors to the list of matches.
    // TODO(slynen): Consider taking random pairs to not under-estimate the
    // variance.
    for (size_t i = 1; i < track.size(); ++i) {
      matches->emplace_back(track[i - 1], track[i]);
    }
  }

  std::random_device device;
  std::mt19937 generator(device());
  std::uniform_int_distribution<> distribution(0, all_descriptors.cols() - 1);

  for (size_t i = 1; i < static_cast<size_t>(all_descriptors.cols()) &&
                     i < matches->size() * 2;
       ++i) {
    unsigned int index_a = distribution(generator);
    unsigned int index_b = distribution(generator);
    if (index_a == index_b) {
      continue;
    }
    non_matches->emplace_back(index_a, index_b);
  }
}

void BuildCovarianceMatricesOfMatchesAndNonMatches(
    unsigned int descriptor_size, const Eigen::MatrixXf& all_descriptors,
    const std::vector<Track>& tracks, unsigned int* sample_size_matches,
    unsigned int* sample_size_non_matches, Eigen::MatrixXf* cov_matches,
    Eigen::MatrixXf* cov_non_matches) {
  CHECK_NOTNULL(sample_size_matches);
  CHECK_NOTNULL(sample_size_non_matches);
  CHECK_NOTNULL(cov_matches);
  CHECK_NOTNULL(cov_non_matches);

  {  // Scope to limit memory usage.
    unsigned int too_short_tracks = 0;
    unsigned int long_enough_tracks = 0;
    constexpr size_t kMinTrackLength = 5;
    size_t number_of_used_tracks = 0;
    Eigen::MatrixXf sumMuMu;
    sumMuMu.setZero(descriptor_size, descriptor_size);

    std::vector<size_t> descriptors_from_tracks;
    descriptors_from_tracks.reserve(500);

    // Centering.
    constexpr int kMaxNumSamples = 50000;
    for (const Track& track : tracks) {
      if (track.size() < kMinTrackLength) {
        ++too_short_tracks;
        continue;
      }
      if (number_of_used_tracks >= kMaxNumSamples) {
        LOG(WARNING) << "Truncated descriptors to " << kMaxNumSamples << ".";
        break;
      }
      ++long_enough_tracks;

      Eigen::Matrix<float, Eigen::Dynamic, 1> mean;
      mean.resize(descriptor_size, Eigen::NoChange);
      mean.setZero();

      for (const size_t& descriptor_idx : track) {
        descriptors_from_tracks.push_back(descriptor_idx);
        mean += all_descriptors.block(0, descriptor_idx, descriptor_size, 1);
      }

      mean /= track.size();
      CHECK_LE(mean.maxCoeff(), 1.0);
      CHECK_GE(mean.minCoeff(), 0.0);

      Eigen::MatrixXf mu_sq_current = (mean * mean.transpose()).eval();

      sumMuMu += mu_sq_current * track.size();
      ++number_of_used_tracks;
    }

    VLOG(3) << "Got " << long_enough_tracks << " tracks out of "
            << tracks.size() << " (dropped " << too_short_tracks
            << " tracks because they were too short)";

    CHECK(!descriptors_from_tracks.empty());

    VLOG(3) << "Computing matches covariance from "
            << descriptors_from_tracks.size()
            << " matches (descriptor size: " << descriptor_size << ")";

    *sample_size_matches = descriptors_from_tracks.size();

    Eigen::MatrixXf matches;
    matches.resize(descriptor_size, descriptors_from_tracks.size());
    matches.setZero();

    int matched_idx = 0;
    for (const size_t& descriptor_idx : descriptors_from_tracks) {
      matches.block(0, matched_idx, descriptor_size, 1) =
          all_descriptors.block(0, descriptor_idx, descriptor_size, 1);
      ++matched_idx;
    }

    CHECK_GT(descriptors_from_tracks.size(), number_of_used_tracks);

    // Covariance computation for matches.
    cov_matches->noalias() =
        (matches * matches.transpose() - sumMuMu) * 2.0 /
        static_cast<float>(
            descriptors_from_tracks.size() - number_of_used_tracks);
  }  // Scope to limit memory usage.

  // Use all descriptors to estimate the non-matching covariance.
  *sample_size_non_matches = all_descriptors.cols();
  // Compute sample covariances non matched descriptors.
  ComputeCovariance(all_descriptors, cov_non_matches);

  *cov_non_matches *= 2.0f;
}

void ComputeProjectionMatrix(
    const Eigen::MatrixXf& cov_matches, const Eigen::MatrixXf& cov_non_matches,
    Eigen::MatrixXf* A) {
  CHECK_NOTNULL(A);

  const int dimensionality = cov_matches.cols();
  A->resize(dimensionality, dimensionality);

  Eigen::JacobiSVD<Eigen::MatrixXf> svd(cov_matches, Eigen::ComputeFullV);

  CHECK_NE(svd.singularValues().minCoeff(), 0)
      << "Rank deficiency for matrix"
         " of samples detected. Probably too little matches.";

  Eigen::MatrixXf Av =
      svd.singularValues().cwiseSqrt().cwiseInverse().asDiagonal() *
      svd.matrixV().transpose();

  Eigen::JacobiSVD<Eigen::MatrixXf> svd_d(
      Av * cov_non_matches * Av.transpose(), Eigen::ComputeFullV);

  Eigen::MatrixXf singular_values_sqrt_inv =
      svd_d.singularValues().cwiseSqrt().cwiseInverse().asDiagonal();

  Eigen::MatrixXf eye;
  eye.resize(dimensionality, dimensionality);
  eye.setIdentity();

  *A = (eye - singular_values_sqrt_inv) * svd_d.matrixV().transpose() *
       svd.singularValues().cwiseInverse().asDiagonal() *
       svd.matrixV().transpose();
}
}  // namespace descriptor_projection


#include "utils.h"
#include "hard_clustering.h"
#include <boost/random/discrete_distribution.hpp>
#include <boost/random/mersenne_twister.hpp>

using namespace std;

CUGMMLite::CUGMMLite(AdaHardCluster& mdl_,
                     unsigned long int dim_)
    : copt::BoundedProblem<double>(2),
      mdl(mdl_),
      dim(dim_),
      k(mdl_.get_k()),
      mu(mdl_.get_mu()),
      m2(mdl_.get_m2()),
      m4(mdl_.get_m4()){
}

double CUGMMLite::value(const copt::Vector<double>& x) {
  // x = <alpha,kappa>
  double obj = 0.0;
  double alpha = x[0];
  double kappa = x[1];
  double var = 0.0;
  for (size_t c = 0; c < k; ++c) {
    var = mdl.variance(mu[c][dim],alpha,dim);
    obj += (pow(m2[c][dim] - kappa*var,2)
      )/(m4[c][dim] - 2*m2[c][dim]*kappa*var + pow(kappa*var,2));
  }
  return obj;
}

void CUGMMLite::gradient(const copt::Vector<double>& x,
                         copt::Vector<double>& grad) {
  double alpha = x[0];
  double kappa = x[1];
  double var = 0.0;
  double diff_var = 0.0;
  grad[0] = 0.0;
  grad[1] = 0.0;
  double coef = 0.0;
  for (size_t c = 0; c < k; ++c){
    var = mdl.variance(mu[c][dim],alpha,dim);
    diff_var = mdl.diff_variance(mu[c][dim],alpha,dim);
    coef = (2*(pow(m2[c][dim],2)-m4[c][dim])*(m2[c][dim] - kappa*var)
      )/pow(m4[c][dim] - 2*m2[c][dim]*kappa*var + pow(kappa*var,2),2);
    grad[0] += coef*kappa*diff_var;
    grad[1] += coef*var;
  }
}

AdaHardCluster::AdaHardCluster(const vector<vector<double>>& data_,
                       const vector<unsigned long int>& label_,
                       unsigned long int max_round_, unsigned long int k_)
    : data(data_), label(label_) {
  max_round = max_round_;
  k = k_;
  nmis.resize(max_round);
  fill(nmis.begin(), nmis.end(), 0);
  logliks.resize(max_round);
  fill(logliks.begin(), logliks.end(), -numeric_limits<double>::max());
  n_samples = data.size();
  n_dims = data[0].size();
  count.resize(k);
  kappa.resize(n_dims);
  kappa_a.resize(n_dims);
  kappa_b.resize(n_dims);
  mu.resize(k);
  mu_a.resize(k);
  mu_b.resize(k);
  m2.resize(k);
  m4.resize(k);
  v.resize(k);
  w_inv.resize(k);
  for (size_t c = 0; c < k; ++c){
    mu[c].resize(n_dims);
    mu_a[c].resize(n_dims);
    mu_b[c].resize(n_dims);
    m2[c].resize(n_dims);
    m4[c].resize(n_dims);
    v[c].resize(n_dims);
    w_inv[c].resize(n_dims);
  }
  asg.resize(n_samples);
  attr_discrete.resize(n_dims);
  attr_positive.resize(n_dims);
  attr_nonnegative.resize(n_dims);

  alpha.resize(n_dims);
  lb.resize(n_dims);
  ub.resize(n_dims);
  for (unsigned long int j = 0; j < n_dims; ++j) {
    attr_discrete[j] = is_discrete(data, j);
    attr_positive[j] = is_positive(data, j);
    attr_nonnegative[j] = attr_positive[j];
    if (!attr_positive[j]) attr_nonnegative[j] = is_nonnegative(data, j);
  }
  for (unsigned long int j = 0; j < n_dims; ++j) {
    if (attr_discrete[j]) {
      if (attr_nonnegative[j]) {
        alpha[j] = 1;
        lb[j] = 0;
        // ub[j] = numeric_limits<double>::max();
        ub[j] = 10;
      } else {
        alpha[j] = 1;
        lb[j] = 0;
        // ub[j] = numeric_limits<double>::max();
        ub[j] = 10;
      }
    } else {
      if (attr_positive[j]) {
        alpha[j] = 0;
        // lb[j] = -numeric_limits<double>::max();
        lb[j] = -10;
        ub[j] = 2;
      } else if (attr_nonnegative[j]) {
        alpha[j] = 0.5;
        lb[j] = 0;
        ub[j] = 1;
      } else {
        alpha[j] = 1;
        lb[j] = 0;
        // ub[j] = numeric_limits<double>::max();
        ub[j] = 10;
      }
    }
  }
}

double AdaHardCluster::distance(const vector<double>& x,
                                const unsigned long int c) {
  double dist = 0.0;
  for (unsigned long int j = 0; j < n_dims; ++j)
    dist += pow(pow(x[j],2) - kappa[j]*v[c][j],2)*w_inv[c][j];
  return dist;
}

double AdaHardCluster::variance(const double x,
                                const double alpha,
                                const unsigned long int dim){
  // return nnc_variance(x, alpha);
  if (attr_discrete[dim]){
    if (attr_nonnegative[dim])
      return nnd_variance(x, alpha);
    else
      return rc_variance(x, alpha);
  } else {
    if (attr_nonnegative[dim])
      return nnc_variance(x, alpha);
    else
      return rc_variance(x, alpha);
  }
}

double AdaHardCluster::diff_variance(const double x,
                                 const double alpha,
                                 const unsigned long int dim) {
  // return nnc_diff_variance(x, alpha);
  if (attr_discrete[dim]){
    if (attr_nonnegative[dim])
      return nnd_diff_variance(x, alpha);
    else
      return rc_diff_variance(x, alpha);
  } else {
    if (attr_nonnegative[dim])
      return nnc_diff_variance(x, alpha);
    else
      return rc_diff_variance(x, alpha);
  }
}

unsigned long int AdaHardCluster::get_k() { return k; }

unsigned long int AdaHardCluster::get_max_round() { return max_round; }

unsigned long int AdaHardCluster::get_n_dims() { return n_dims; }

unsigned long int AdaHardCluster::get_n_samples() { return n_samples; }

vector<double>& AdaHardCluster::get_kappa() { return kappa; }

vector<double>& AdaHardCluster::get_logliks() { return logliks; }

vector<double>& AdaHardCluster::get_nmis() { return nmis; }

copt::Vector<double>& AdaHardCluster::get_alpha() { return alpha; }

copt::Vector<double>& AdaHardCluster::get_lb() { return lb; }

copt::Vector<double>& AdaHardCluster::get_ub() { return ub; }

const vector<vector<double>>& AdaHardCluster::get_data() { return data; }

vector<vector<double>>& AdaHardCluster::get_mu() { return mu; }

vector<vector<double>>& AdaHardCluster::get_m2() { return m2; }

vector<vector<double>>& AdaHardCluster::get_m4() { return m4; }

void AdaHardCluster::initialize_random() {
  unsigned long int ri;
  for (size_t c = 0; c < k; ++c) {
    ri = rand() % n_samples;
    for (unsigned long int j = 0; j < n_dims; ++j){
      mu[c][j] = data[ri][j];
      mu_a[c][j] = 0;//mu[c][j];
      mu_b[c][j] = 1.0;
    }
  }
  for (unsigned long int j = 0; j < n_dims; ++j){
    kappa[j] = 1.0;
    kappa_a[j] = 1e-9;
    kappa_b[j] = 1.0;
  }
}

void AdaHardCluster::initialize_k_plus_plus() {
  boost::mt19937 gen;
  unsigned long int ri, k_eff;
  ri = rand() % n_samples;
  for (unsigned long int j = 0; j < n_dims; ++j) mu[0][j] = data[ri][j];
  k_eff = 1;
  vector<double> probs(n_samples, 0);
  double dist = 0;
  double best_dist = 0;
  unsigned long int best_asg = -1;
  unsigned long int cur_asg = -1;
  do {
    for (size_t i = 0; i < n_samples; ++i) {
      best_dist = numeric_limits<double>::max();
      for (size_t c = 0; c < k_eff; ++c) {
        dist = 0;
        for (unsigned long int j = 0; j < n_dims; ++j){
          // dist += distance(data[i][j], mu[c][j], alpha[j], j);
          dist += pow(data[i][j] - mu[c][j], 2);
        }
        if (dist < best_dist) best_dist = dist;
      }
      probs[i] = pow(best_dist, 2);
    }
    boost::random::discrete_distribution<> dist(probs.begin(), probs.end());
    ri = dist(gen);
    for (unsigned long int j = 0; j < n_dims; ++j) mu[k_eff][j] = data[ri][j];
    k_eff += 1;
  } while (k_eff != k);

  for (unsigned long int j = 0; j < n_dims; ++j)
    if (attr_nonnegative[j])
      for (size_t c = 0; c < k_eff; ++c)
        if (mu[c][j] == 0)
          mu[c][j] = 1e-9;
  for (unsigned long int j = 0; j < n_dims; ++j){
    kappa[j] = 1.0;
    kappa_a[j] = 1.0;
    kappa_b[j] = 1.0;
  }
  for (size_t c = 0; c < k_eff; ++c)
    for (unsigned long int j = 0; j < n_dims; ++j){
      mu_a[c][j] = mu[c][j];
      mu_b[c][j] = 1.0;
    }
  for (size_t i = 0; i < n_samples; ++i) {
    best_dist = numeric_limits<double>::max();
    best_asg = -1;
    for (size_t c = 0; c < k; ++c) {
      dist = 0.0;
      for (size_t j = 0; j < n_dims; ++j)
        dist += pow(data[i][j] - mu[c][j], 2);
      if (dist < best_dist) {
        best_dist = dist;
        best_asg = c;
      }
    }
    asg[i] = best_asg;
  }
  for (size_t c = 0; c < k; c++){
    fill(mu[c].begin(), mu[c].end(), 1e-9);
    fill(m2[c].begin(), m2[c].end(), 1e-9);
    fill(m4[c].begin(), m4[c].end(), 1e-9);
  }
  fill(count.begin(), count.end(), 0);
  for (size_t i = 0; i < n_samples; ++i) {
    cur_asg = asg[i];
    count[cur_asg] += 1;
    for (size_t j = 0; j < n_dims; ++j){
      mu[cur_asg][j] += data[i][j];
      m2[cur_asg][j] += pow(data[i][j],2);
      m4[cur_asg][j] += pow(data[i][j],4);
    }
  }
  for (size_t c = 0; c < k; c++)
    for (size_t j = 0; j < n_dims; ++j){
      mu[c][j] = (kappa[j]*mu_a[c][j] + mu[c][j]
                  )/(kappa[j]*mu_b[c][j] + count[c]);
      m2[c][j] = (kappa[j]*pow(mu_a[c][j],2) + m2[c][j]
                  )/(kappa[j]*mu_b[c][j] + count[c]);
      m4[c][j] = (kappa[j]*pow(mu_a[c][j],4) + m4[c][j]
                  )/(kappa[j]*mu_b[c][j] + count[c]);
      v[c][j] = variance(mu[c][j],alpha[j],j);
      w_inv[c][j] = (kappa[j]*mu_b[c][j] + count[c]
        ) / (m4[c][j] - 2*m2[c][j]*kappa[j]*v[c][j] + pow(kappa[j]*v[c][j],2));
    }
}

void AdaHardCluster::fit() {
  bool updated = 0;
  double inertia = 0;
  double nmi = 0;
  double dist = 0;
  double best_dist = 0;
  unsigned long int best_asg = -1;
  unsigned long int cur_asg = -1;
  for (size_t r = 0; r < max_round; r++) {
    updated = 0;
    inertia = 0.0;
    for (size_t i = 0; i < n_samples; ++i) {
      best_dist = numeric_limits<double>::max();
      best_asg = -1;
      for (size_t c = 0; c < k; ++c) {
        dist = distance(data[i], c);
        if (dist < best_dist) {
          best_dist = dist;
          best_asg = c;
        }
      }
      if (asg[i] != best_asg) {
        asg[i] = best_asg;
        updated = 1;
      }
      inertia += best_dist;
    }
    nmi = calc_nmi(label, asg);
    nmis[r] = nmi;
    if (!updated) break;

    for (size_t c = 0; c < k; c++){
      fill(mu[c].begin(), mu[c].end(), 1e-9);
      fill(m2[c].begin(), m2[c].end(), 1e-9);
      fill(m4[c].begin(), m4[c].end(), 1e-9);
    }
    fill(count.begin(), count.end(), 0);
    for (size_t i = 0; i < n_samples; ++i) {
      cur_asg = asg[i];
      count[cur_asg] += 1;
      for (size_t j = 0; j < n_dims; ++j){
        mu[cur_asg][j] += data[i][j];
        m2[cur_asg][j] += pow(data[i][j],2);
        m4[cur_asg][j] += pow(data[i][j],4);
      }
    }
    for (size_t c = 0; c < k; c++)
      for (size_t j = 0; j < n_dims; ++j){
        mu[c][j] = (kappa[j]*mu_a[c][j] + mu[c][j]
                    )/(kappa[j]*mu_b[c][j] + count[c]);
        m2[c][j] = (kappa[j]*pow(mu_a[c][j],2) + m2[c][j]
                    )/(kappa[j]*mu_b[c][j] + count[c]);
        m4[c][j] = (kappa[j]*pow(mu_a[c][j],4) + m4[c][j]
                    )/(kappa[j]*mu_b[c][j] + count[c]);
        v[c][j] = variance(mu[c][j],alpha[j],j);
        w_inv[c][j] = (kappa[j]*mu_b[c][j] + count[c]
          ) / (m4[c][j] - 2*m2[c][j]*kappa[j]*v[c][j] + pow(kappa[j]*v[c][j],2));
      }
    double obj = 0;
    for (size_t j = 0; j < n_dims; ++j){
      CUGMMLite prblm(*this, j);
      copt::Vector<double> lbj;
      lbj.resize(2);
      lbj[0] = lb[j];
      lbj[1] = 1e-9;
      prblm.setLowerBound(lbj);
      copt::Vector<double> ubj;
      ubj.resize(2);
      ubj[0] = ub[j];
      ubj[1] = numeric_limits<double>::max();
      prblm.setUpperBound(ubj);
      copt::Vector<double> xj;
      xj.resize(2);
      xj[0] = alpha[j];
      xj[1] = kappa[j];
      copt::LbfgsbSolver<CUGMMLite> solver;
      solver.minimize(prblm, xj);
      obj += prblm.value(xj);
      if (!isnan(xj[0]))
        alpha[j] = xj[0];
      if (!isnan(xj[1]))
        kappa[j] = xj[1];
    }
    cout << "round=" << r << " nmi=" << nmi << " inertia=" << inertia;
    cout << " obj=" << obj << endl;
  }
  cout << "Alpha=";
  for (unsigned long int j = 0; j < n_dims; ++j)
      cout << alpha[j] << " ";
  cout << endl;
  cout << "Kappa=";
  for (unsigned long int j = 0; j < n_dims; ++j)
      cout << kappa[j] << " ";
  cout << endl;
  for (size_t c = 0; c < k; ++c){
    cout << "mu[" << c << "]=";
    for (unsigned long int j = 0; j < n_dims; ++j)
        cout << mu[c][j] << " ";
    cout << endl;
  }
}


// Copyright 2021 RoboJackets
//
// 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.

// BEGIN STUDENT CODE

#ifndef KALMAN_FILTER_HPP_
#define KALMAN_FILTER_HPP_

#include <Eigen/Dense>

namespace mineral_deposit_tracking
{

template<int StateSize>
class KalmanFilter
{
public:
  using VectorType = Eigen::Matrix<double, StateSize, 1>;
  using MatrixType = Eigen::Matrix<double, StateSize, StateSize>;

  KalmanFilter(
    const MatrixType & transition_matrix,
    const MatrixType & process_covariance,
    const MatrixType & observation_matrix)
  : transition_matrix_(transition_matrix),
    process_covariance_(process_covariance),
    observation_matrix_(observation_matrix),
    estimate_(VectorType::Zero()),
    estimate_covariance_(MatrixType::Identity() * 500)
  {
  }

  void Reset(const VectorType & initial_state, const MatrixType & initial_covariance)
  {
    estimate_ = initial_state;
    estimate_covariance_ = initial_covariance;
  }

  void TimeUpdate()
  {
    estimate_ = transition_matrix_ * estimate_;
    estimate_covariance_ = transition_matrix_ * estimate_covariance_ *
      transition_matrix_.transpose() + process_covariance_;
  }

  void MeasurementUpdate(
    const VectorType & measurement,
    const MatrixType & measurement_covariance)
  {
    const MatrixType innovation_covariance =
      (observation_matrix_ * estimate_covariance_ * observation_matrix_.transpose()) +
      measurement_covariance;

    const MatrixType gain = estimate_covariance_ * observation_matrix_.transpose() *
      innovation_covariance.inverse();

    estimate_ = estimate_ + (gain * (measurement - (observation_matrix_ * estimate_)));

    const MatrixType tmp = MatrixType::Identity() - (gain * observation_matrix_);

    estimate_covariance_ = (tmp * estimate_covariance_ * tmp.transpose()) +
      (gain * measurement_covariance * gain.transpose());
  }

  const VectorType & GetEstimate() const
  {
    return estimate_;
  }

  const MatrixType & GetEstimateCovariance() const
  {
    return estimate_covariance_;
  }

private:
  const MatrixType transition_matrix_;
  const MatrixType process_covariance_;
  const MatrixType observation_matrix_;
  VectorType estimate_;
  MatrixType estimate_covariance_;
};

}  // namespace mineral_deposit_tracking

#endif  // KALMAN_FILTER_HPP_

// END STUDENT CODE


/******************************************************************************
 * Copyright (c) 2015 - 2016 Philipp Schubert.                                *
 * All rights reserved. This program and the accompanying materials are made  *
 * available under the terms of LICENSE.txt.                                  *
 *                                                                            *
 * Contributors:                                                              *
 *     Philipp Schubert                                                       *
 *****************************************************************************/

/**
 * @file pseudoinverse.cpp
 * @brief Enthält die Implementation der Datei pseudoinverse.hh.
 *
 * Hier wird der Prototyp zur Berechnung eines Pseudoinversen einer
 * Dreiecksmatrix implementiert.
 *
 * @author Philipp D. Schubert
 * @bug Keine Bugs bekannt.
 */

#include <armadillo>
#include <omp.h>
#include <stdlib.h>
#include <string.h>
#include <typeinfo>

extern "C" {
#include "m.h"
#include "tm.h"
#include "utils.h"
}

extern "C" tm_t *pseudoinverse(const tm_t *t) {
  m_t *m = initM(t->size, t->size);
  unsigned int i, j;
#pragma omp parallel for shared(t, m) private(i, j) schedule(dynamic)
  for (i = 0; i < t->size; ++i) {
    for (j = 0; j < t->size; ++j) {
      m->elems[i * t->size + j] = (j <= i) ? t->elems[(i * (i + 1)) / 2 + j]
                                           : t->elems[(j * (j + 1)) / 2 + i];
    }
  }
  // check if we should use single or double precision and use adequate
  // armadillo data structures
  if (typeid(real_t).name() == typeid(float).name()) {
    arma::fmat f_am =
        arma::fmat((float *)m->elems, t->size, t->size, false, false);
    arma::fmat f_inverse = arma::pinv(f_am);
    // caution: override data memory of m
    memcpy(m->elems, f_inverse.memptr(), t->size * t->size * sizeof(real_t));
  } else {
    arma::mat d_am =
        arma::mat((double *)m->elems, t->size, t->size, false, false);
    arma::mat d_inverse = arma::pinv(d_am);
    // caution: override data memory of m
    memcpy(m->elems, d_inverse.memptr(), t->size * t->size * sizeof(real_t));
  }
  tm_t *result = initTM(t->size);
#pragma omp parallel for shared(result, m) private(i, j) schedule(dynamic)
  for (i = 0; i < t->size; ++i) {
    for (j = 0; j <= i; ++j) {
      result->elems[(i * (i + 1)) / 2 + j] = m->elems[i * t->size + j];
    }
  }
  free(m);
  return result;
}


//  Created by jimmy on 2016-06-14.
//  Copyright © 2016 jimmy. All rights reserved.
//

#ifndef __Kabsch__
#define __Kabsch__

// This code is modifed from released in public domain from
// https://github.com/oleg-alexandrov/projects/blob/master/eigen/Kabsch.cpp

#include <Eigen/Geometry>

// This is general Kabsh algorithm,
// The input 3D points are stored as columns.
Eigen::Affine3d Find3DAffineTransform(Eigen::Matrix3Xd input_pts, Eigen::Matrix3Xd output_pts);

// This is Kabsh algorithm for camera pose estimation
// Assume the points in camera coordiantes and world coordinates have same scale
// The input 3D points are stored as columns. No scale effect
Eigen::Affine3d Find3DAffineTransformSameScale(Eigen::Matrix3Xd input_pts, Eigen::Matrix3Xd output_pts);

#endif /* __Kabsch__ */


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

// Copyright (c) 2007-2011 Barend Gehrels, Amsterdam, the Netherlands.
// Copyright (c) 2008-2011 Bruno Lalande, Paris, France.
// Copyright (c) 2009-2011 Mateusz Loskot, London, UK.

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

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

#ifndef BOOST_GEOMETRY_UTIL_MATH_HPP
#define BOOST_GEOMETRY_UTIL_MATH_HPP

#include <cmath>
#include <limits>

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

#include <boost/geometry/util/select_most_precise.hpp>

namespace boost { namespace geometry
{

namespace math
{

#ifndef DOXYGEN_NO_DETAIL
namespace detail
{

template <typename Type, bool IsFloatingPoint>
struct equals
{
    static inline bool apply(Type const& a, Type const& b)
    {
        return a == b;
    }
};

template <typename Type>
struct equals<Type, true>
{
    static inline bool apply(Type const& a, Type const& b)
    {
        // See http://www.parashift.com/c++-faq-lite/newbie.html#faq-29.17,
        // FUTURE: replace by some boost tool or boost::test::close_at_tolerance
        return std::abs(a - b) <= std::numeric_limits<Type>::epsilon() * std::abs(a);
    }
};

template <typename Type, bool IsFloatingPoint>
struct equals_with_epsilon : public equals<Type, IsFloatingPoint> {};

/*!
\brief Short construct to enable partial specialization for PI, currently not possible in Math.
*/
template <typename T>
struct define_pi
{
    static inline T apply()
    {
        // Default calls Boost.Math
        return boost::math::constants::pi<T>();
    }
};

} // namespace detail
#endif

template <typename T>
inline T pi() { return detail::define_pi<T>::apply(); }

// Maybe replace this by boost equals or boost ublas numeric equals or so

/*!
    \brief returns true if both arguments are equal.
    \ingroup utility
    \param a first argument
    \param b second argument
    \return true if a == b
    \note If both a and b are of an integral type, comparison is done by ==.
    If one of the types is floating point, comparison is done by abs and
    comparing with epsilon. If one of the types is non-fundamental, it might
    be a high-precision number and comparison is done using the == operator
    of that class.
*/

template <typename T1, typename T2>
inline bool equals(T1 const& a, T2 const& b)
{
    typedef typename select_most_precise<T1, T2>::type select_type;
    return detail::equals
        <
            select_type,
            boost::is_floating_point<select_type>::type::value
        >::apply(a, b);
}

template <typename T1, typename T2>
inline bool equals_with_epsilon(T1 const& a, T2 const& b)
{
    typedef typename select_most_precise<T1, T2>::type select_type;
    return detail::equals_with_epsilon
        <
            select_type,
            boost::is_floating_point<select_type>::type::value
        >::apply(a, b);
}

double const d2r = geometry::math::pi<double>() / 180.0;
double const r2d = 1.0 / d2r;

/*!
    \brief Calculates the haversine of an angle
    \ingroup utility
    \note See http://en.wikipedia.org/wiki/Haversine_formula
    haversin(alpha) = sin2(alpha/2)
*/
template <typename T>
inline T hav(T const& theta)
{
    T const half = T(0.5);
    T const sn = sin(half * theta);
    return sn * sn;
}

/*!
\brief Short utility to return the square
\ingroup utility
\param value Value to calculate the square from
\return The squared value
*/
template <typename T>
inline T sqr(T const& value)
{
    return value * value;
}

/*!
\brief Short utility to workaround gcc/clang problem that abs is converting to integer
\ingroup utility
*/
template<typename T>
inline T abs(const T& t)
{
    using std::abs;
    return abs(t);
}

} // namespace math

}} // namespace boost::geometry

#endif // BOOST_GEOMETRY_UTIL_MATH_HPP


#ifndef SKYLARK_WZT_HPP
#define SKYLARK_WZT_HPP

#ifndef SKYLARK_SKETCH_HPP
#error "Include top-level sketch.hpp instead of including individuals headers"
#endif

#include <boost/random.hpp>

namespace skylark { namespace sketch {

/**
 * Woodruff-Zhang Transform (data)
 *
 * Woodruff-Zhang Transform is very similar to the Clarkson-Woodruff Transform:
 * it replaces the +1/-1 diagonal with reciprocal exponentia random enteries.
 * It is sutiable for lp regression with 1 <= p <= 2.
 *
 * Reference:
 * D. Woodruff and Q. Zhang
 * Subspace Embeddings and L_p Regression Using Exponential Random
 * COLT 2013
 *
 * TODO current implementation is only one sketch index, when for 1 <= p <= 2
 *      you want more than one.
 */

template < typename InputMatrixType,
           typename OutputMatrixType = InputMatrixType >
struct WZT_t :
        public WZT_data_t,
        virtual public sketch_transform_t<InputMatrixType, OutputMatrixType > {

public:

    // We use composition to defer calls to hash_transform_t
    typedef hash_transform_t< InputMatrixType, OutputMatrixType,
                              boost::random::uniform_int_distribution,
                              boost::random::exponential_distribution > transform_t;

    typedef WZT_data_t data_type;
    typedef data_type::params_t params_t;

    WZT_t(int N, int S, double p, base::context_t& context)
        : data_type(N, S, p, context), _transform(*this) {

    }

    WZT_t(int N, int S, const params_t& params, base::context_t& context)
        : data_type(N, S, params, context),
          _transform(*this) {

    }

    WZT_t(const boost::property_tree::ptree &pt)
        : data_type(pt), _transform(*this) {

    }

    template< typename OtherInputMatrixType,
              typename OtherOutputMatrixType >
    WZT_t(const WZT_t<OtherInputMatrixType,OtherOutputMatrixType>& other)
        : data_type(other), _transform(*this) {

    }

    WZT_t(const data_type& other)
        : data_type(other), _transform(*this) {

    }

    /**
     * Apply columnwise the sketching transform that is described by the
     * the transform with output sketch_of_A.
     */
    void apply (const typename transform_t::matrix_type& A,
                typename transform_t::output_matrix_type& sketch_of_A,
                columnwise_tag dimension) const {
        _transform.apply(A, sketch_of_A, dimension);
    }

    /**
     * Apply rowwise the sketching transform that is described by the
     * the transform with output sketch_of_A.
     */
    void apply (const typename transform_t::matrix_type& A,
                typename transform_t::output_matrix_type& sketch_of_A,
                rowwise_tag dimension) const {
        _transform.apply(A, sketch_of_A, dimension);
    }

    int get_N() const { return this->_N; } /**< Get input dimesion. */
    int get_S() const { return this->_S; } /**< Get output dimesion. */

    const sketch_transform_data_t* get_data() const { return this; }

private:
    transform_t _transform;
};

template<>
class WZT_t<boost::any, boost::any> :
  public WZT_data_t,
  virtual public sketch_transform_t<boost::any, boost::any > {

public:

    typedef WZT_data_t data_type;
    typedef data_type::params_t params_t;

    WZT_t(int N, int S, double p, base::context_t& context)
        : data_type(N, S, p, context) {

    }

    WZT_t(int N, int S, const params_t& params, base::context_t& context)
        : data_type(N, S, params, context) {

    }

    WZT_t(const boost::property_tree::ptree &pt)
        : data_type(pt) {

    }

    /**
     * Copy constructor
     */
    template <typename OtherInputMatrixType,
              typename OtherOutputMatrixType>
    WZT_t (const WZT_t<OtherInputMatrixType, OtherOutputMatrixType>& other)
        : data_type(other) {

    }

    /**
     * Constructor from data
     */
    WZT_t (const data_type& other)
        : data_type(other) {

    }

    /**
     * Apply columnwise the sketching transform that is described by the
     * the transform with output sketch_of_A.
     */
    void apply(const boost::any &A, const boost::any &sketch_of_A,
                columnwise_tag dimension) const {

#if     !(defined SKYLARK_NO_ANY) || (defined SKYLARK_WITH_WZT_ANY)

        SKYLARK_SKETCH_ANY_APPLY_DISPATCH(mdtypes::matrix_t, mdtypes::matrix_t,
            WZT_t);
        SKYLARK_SKETCH_ANY_APPLY_DISPATCH(mdtypes::sparse_matrix_t,
            mdtypes::matrix_t, WZT_t);
        SKYLARK_SKETCH_ANY_APPLY_DISPATCH(mdtypes::sparse_matrix_t,
            mdtypes::sparse_matrix_t, WZT_t);
        SKYLARK_SKETCH_ANY_APPLY_DISPATCH(mdtypes::shared_matrix_t,
            mdtypes::shared_matrix_t, WZT_t);
        SKYLARK_SKETCH_ANY_APPLY_DISPATCH(mdtypes::root_matrix_t,
            mdtypes::root_matrix_t, WZT_t);
        SKYLARK_SKETCH_ANY_APPLY_DISPATCH(mdtypes::dist_matrix_t,
            mdtypes::shared_matrix_t, WZT_t);
        SKYLARK_SKETCH_ANY_APPLY_DISPATCH(mdtypes::dist_matrix_t,
            mdtypes::root_matrix_t, WZT_t);
        SKYLARK_SKETCH_ANY_APPLY_DISPATCH(mdtypes::dist_matrix_vc_star_t,
            mdtypes::root_matrix_t, WZT_t);
        SKYLARK_SKETCH_ANY_APPLY_DISPATCH(mdtypes::dist_matrix_vc_star_t,
            mdtypes::shared_matrix_t, WZT_t);
        SKYLARK_SKETCH_ANY_APPLY_DISPATCH(mdtypes::dist_matrix_vr_star_t,
            mdtypes::root_matrix_t, WZT_t);
        SKYLARK_SKETCH_ANY_APPLY_DISPATCH(mdtypes::dist_matrix_vr_star_t,
            mdtypes::shared_matrix_t, WZT_t);
        SKYLARK_SKETCH_ANY_APPLY_DISPATCH(mdtypes::dist_matrix_t,
            mdtypes::dist_matrix_t, WZT_t);
        SKYLARK_SKETCH_ANY_APPLY_DISPATCH(mdtypes::dist_matrix_vc_star_t,
            mdtypes::dist_matrix_vc_star_t, WZT_t);
        SKYLARK_SKETCH_ANY_APPLY_DISPATCH(mdtypes::dist_matrix_vr_star_t,
            mdtypes::dist_matrix_vr_star_t, WZT_t);
        SKYLARK_SKETCH_ANY_APPLY_DISPATCH(mdtypes::dist_matrix_star_vc_t,
            mdtypes::dist_matrix_star_vc_t, WZT_t);
        SKYLARK_SKETCH_ANY_APPLY_DISPATCH(mdtypes::dist_matrix_star_vr_t,
            mdtypes::dist_matrix_star_vr_t, WZT_t);

        SKYLARK_SKETCH_ANY_APPLY_DISPATCH(mftypes::matrix_t, mftypes::matrix_t,
            WZT_t);
        SKYLARK_SKETCH_ANY_APPLY_DISPATCH(mftypes::sparse_matrix_t,
            mftypes::matrix_t, WZT_t);
        SKYLARK_SKETCH_ANY_APPLY_DISPATCH(mftypes::sparse_matrix_t,
            mftypes::sparse_matrix_t, WZT_t);
        SKYLARK_SKETCH_ANY_APPLY_DISPATCH(mftypes::shared_matrix_t,
            mftypes::shared_matrix_t, WZT_t);
        SKYLARK_SKETCH_ANY_APPLY_DISPATCH(mftypes::root_matrix_t,
            mftypes::root_matrix_t, WZT_t);
        SKYLARK_SKETCH_ANY_APPLY_DISPATCH(mftypes::dist_matrix_t,
            mftypes::shared_matrix_t, WZT_t);
        SKYLARK_SKETCH_ANY_APPLY_DISPATCH(mftypes::dist_matrix_t,
            mftypes::root_matrix_t, WZT_t);
        SKYLARK_SKETCH_ANY_APPLY_DISPATCH(mftypes::dist_matrix_vc_star_t,
            mftypes::root_matrix_t, WZT_t);
        SKYLARK_SKETCH_ANY_APPLY_DISPATCH(mftypes::dist_matrix_vc_star_t,
            mftypes::shared_matrix_t, WZT_t);
        SKYLARK_SKETCH_ANY_APPLY_DISPATCH(mftypes::dist_matrix_vr_star_t,
            mftypes::root_matrix_t, WZT_t);
        SKYLARK_SKETCH_ANY_APPLY_DISPATCH(mftypes::dist_matrix_vr_star_t,
            mftypes::shared_matrix_t, WZT_t);
        SKYLARK_SKETCH_ANY_APPLY_DISPATCH(mftypes::dist_matrix_t,
            mftypes::dist_matrix_t, WZT_t);
        SKYLARK_SKETCH_ANY_APPLY_DISPATCH(mftypes::dist_matrix_vc_star_t,
            mftypes::dist_matrix_vc_star_t, WZT_t);
        SKYLARK_SKETCH_ANY_APPLY_DISPATCH(mftypes::dist_matrix_vr_star_t,
            mftypes::dist_matrix_vr_star_t, WZT_t);
        SKYLARK_SKETCH_ANY_APPLY_DISPATCH(mftypes::dist_matrix_star_vc_t,
            mftypes::dist_matrix_star_vc_t, WZT_t);
        SKYLARK_SKETCH_ANY_APPLY_DISPATCH(mftypes::dist_matrix_star_vr_t,
            mftypes::dist_matrix_star_vr_t, WZT_t);

#endif

        SKYLARK_THROW_EXCEPTION (
          base::sketch_exception()
              << base::error_msg(
                 "This combination has not yet been implemented for WZT"));

    }

    /**
     * Apply rowwise the sketching transform that is described by the
     * the transform with output sketch_of_A.
     */
    void apply (const boost::any &A, const boost::any &sketch_of_A,
        rowwise_tag dimension) const {

#if     !(defined SKYLARK_NO_ANY) || (defined SKYLARK_WITH_WZT_ANY)

        SKYLARK_SKETCH_ANY_APPLY_DISPATCH(mdtypes::matrix_t, mdtypes::matrix_t,
            WZT_t);
        SKYLARK_SKETCH_ANY_APPLY_DISPATCH(mdtypes::sparse_matrix_t,
            mdtypes::matrix_t, WZT_t);
        SKYLARK_SKETCH_ANY_APPLY_DISPATCH(mdtypes::sparse_matrix_t,
            mdtypes::sparse_matrix_t, WZT_t);
        SKYLARK_SKETCH_ANY_APPLY_DISPATCH(mdtypes::shared_matrix_t,
            mdtypes::shared_matrix_t, WZT_t);
        SKYLARK_SKETCH_ANY_APPLY_DISPATCH(mdtypes::root_matrix_t,
            mdtypes::root_matrix_t, WZT_t);
        SKYLARK_SKETCH_ANY_APPLY_DISPATCH(mdtypes::dist_matrix_t,
            mdtypes::shared_matrix_t, WZT_t);
        SKYLARK_SKETCH_ANY_APPLY_DISPATCH(mdtypes::dist_matrix_t,
            mdtypes::root_matrix_t, WZT_t);
        SKYLARK_SKETCH_ANY_APPLY_DISPATCH(mdtypes::dist_matrix_vc_star_t,
            mdtypes::root_matrix_t, WZT_t);
        SKYLARK_SKETCH_ANY_APPLY_DISPATCH(mdtypes::dist_matrix_vc_star_t,
            mdtypes::shared_matrix_t, WZT_t);
        SKYLARK_SKETCH_ANY_APPLY_DISPATCH(mdtypes::dist_matrix_vr_star_t,
            mdtypes::root_matrix_t, WZT_t);
        SKYLARK_SKETCH_ANY_APPLY_DISPATCH(mdtypes::dist_matrix_vr_star_t,
            mdtypes::shared_matrix_t, WZT_t);
        SKYLARK_SKETCH_ANY_APPLY_DISPATCH(mdtypes::dist_matrix_t,
            mdtypes::dist_matrix_t, WZT_t);
        SKYLARK_SKETCH_ANY_APPLY_DISPATCH(mdtypes::dist_matrix_vc_star_t,
            mdtypes::dist_matrix_vc_star_t, WZT_t);
        SKYLARK_SKETCH_ANY_APPLY_DISPATCH(mdtypes::dist_matrix_vr_star_t,
            mdtypes::dist_matrix_vr_star_t, WZT_t);
        SKYLARK_SKETCH_ANY_APPLY_DISPATCH(mdtypes::dist_matrix_star_vc_t,
            mdtypes::dist_matrix_star_vc_t, WZT_t);
        SKYLARK_SKETCH_ANY_APPLY_DISPATCH(mdtypes::dist_matrix_star_vr_t,
            mdtypes::dist_matrix_star_vr_t, WZT_t);

        SKYLARK_SKETCH_ANY_APPLY_DISPATCH(mftypes::matrix_t, mftypes::matrix_t,
            WZT_t);
        SKYLARK_SKETCH_ANY_APPLY_DISPATCH(mftypes::sparse_matrix_t,
            mftypes::matrix_t, WZT_t);
        SKYLARK_SKETCH_ANY_APPLY_DISPATCH(mftypes::sparse_matrix_t,
            mftypes::sparse_matrix_t, WZT_t);
        SKYLARK_SKETCH_ANY_APPLY_DISPATCH(mftypes::shared_matrix_t,
            mftypes::shared_matrix_t, WZT_t);
        SKYLARK_SKETCH_ANY_APPLY_DISPATCH(mftypes::root_matrix_t,
            mftypes::root_matrix_t, WZT_t);
        SKYLARK_SKETCH_ANY_APPLY_DISPATCH(mftypes::dist_matrix_t,
            mftypes::shared_matrix_t, WZT_t);
        SKYLARK_SKETCH_ANY_APPLY_DISPATCH(mftypes::dist_matrix_t,
            mftypes::root_matrix_t, WZT_t);
        SKYLARK_SKETCH_ANY_APPLY_DISPATCH(mftypes::dist_matrix_vc_star_t,
            mftypes::root_matrix_t, WZT_t);
        SKYLARK_SKETCH_ANY_APPLY_DISPATCH(mftypes::dist_matrix_vc_star_t,
            mftypes::shared_matrix_t, WZT_t);
        SKYLARK_SKETCH_ANY_APPLY_DISPATCH(mftypes::dist_matrix_vr_star_t,
            mftypes::root_matrix_t, WZT_t);
        SKYLARK_SKETCH_ANY_APPLY_DISPATCH(mftypes::dist_matrix_vr_star_t,
            mftypes::shared_matrix_t, WZT_t);
        SKYLARK_SKETCH_ANY_APPLY_DISPATCH(mftypes::dist_matrix_t,
            mftypes::dist_matrix_t, WZT_t);
        SKYLARK_SKETCH_ANY_APPLY_DISPATCH(mftypes::dist_matrix_vc_star_t,
            mftypes::dist_matrix_vc_star_t, WZT_t);
        SKYLARK_SKETCH_ANY_APPLY_DISPATCH(mftypes::dist_matrix_vr_star_t,
            mftypes::dist_matrix_vr_star_t, WZT_t);
        SKYLARK_SKETCH_ANY_APPLY_DISPATCH(mftypes::dist_matrix_star_vc_t,
            mftypes::dist_matrix_star_vc_t, WZT_t);
        SKYLARK_SKETCH_ANY_APPLY_DISPATCH(mftypes::dist_matrix_star_vr_t,
            mftypes::dist_matrix_star_vr_t, WZT_t);

#endif

        SKYLARK_THROW_EXCEPTION (
          base::sketch_exception()
              << base::error_msg(
                 "This combination has not yet been implemented for WZT"));

    }

    int get_N() const { return this->_N; } /**< Get input dimesion. */
    int get_S() const { return this->_S; } /**< Get output dimesion. */

    const sketch_transform_data_t* get_data() const { return this; }
};

} } /** namespace skylark::sketch */

#endif // SKYLARK_WZT_HPP


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

/*
 Copyright (C) 2001, 2002, 2003 Sadruddin Rejeb
 Copyright (C) 2003, 2004, 2005, 2006, 2007, 2008, 2012 Ferdinando Ametrano
 Copyright (C) 2006 Mark Joshi
 Copyright (C) 2006 StatPro Italia srl
 Copyright (C) 2007 Cristina Duminuco
 Copyright (C) 2007 Chiara Fornarola

 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.
*/

#include <ql/pricingengines/blackformula.hpp>
#include <ql/math/solvers1d/newtonsafe.hpp>
#include <ql/math/distributions/normaldistribution.hpp>
#include <boost/math/special_functions/atanh.hpp>

namespace {
    void checkParameters(QuantLib::Real strike,
                         QuantLib::Real forward,
                         QuantLib::Real displacement)
    {
        QL_REQUIRE(displacement >= 0.0, "displacement ("
                                            << displacement
                                            << ") must be non-negative");
        QL_REQUIRE(strike + displacement >= 0.0,
                   "strike + displacement (" << strike << " + " << displacement
                                             << ") must be non-negative");
        QL_REQUIRE(forward + displacement > 0.0, "forward + displacement ("
                                                     << forward << " + "
                                                     << displacement
                                                     << ") must be positive");
    }
}

namespace QuantLib {

    Real blackFormula(Option::Type optionType,
                      Real strike,
                      Real forward,
                      Real stdDev,
                      Real discount,
                      Real displacement)
    {
        checkParameters(strike, forward, displacement);
        QL_REQUIRE(stdDev>=0.0,
                   "stdDev (" << stdDev << ") must be non-negative");
        QL_REQUIRE(discount>0.0,
                   "discount (" << discount << ") must be positive");

        if (stdDev==0.0)
            return std::max((forward-strike)*optionType, Real(0.0))*discount;

        forward = forward + displacement;
        strike = strike + displacement;

        // since displacement is non-negative strike==0 iff displacement==0
        // so returning forward*discount is OK
        if (strike==0.0)
            return (optionType==Option::Call ? forward*discount : 0.0);

        Real d1 = std::log(forward/strike)/stdDev + 0.5*stdDev;
        Real d2 = d1 - stdDev;
        CumulativeNormalDistribution phi;
        Real nd1 = phi(optionType*d1);
        Real nd2 = phi(optionType*d2);
        Real result = discount * optionType * (forward*nd1 - strike*nd2);
        QL_ENSURE(result>=0.0,
                  "negative value (" << result << ") for " <<
                  stdDev << " stdDev, " <<
                  optionType << " option, " <<
                  strike << " strike , " <<
                  forward << " forward");
        return result;
    }

    Real blackFormula(const boost::shared_ptr<PlainVanillaPayoff>& payoff,
                      Real forward,
                      Real stdDev,
                      Real discount,
                      Real displacement) {
        return blackFormula(payoff->optionType(),
            payoff->strike(), forward, stdDev, discount, displacement);
    }

    Real blackFormulaImpliedStdDevApproximation(Option::Type optionType,
                                                Real strike,
                                                Real forward,
                                                Real blackPrice,
                                                Real discount,
                                                Real displacement)
    {
        checkParameters(strike, forward, displacement);
        QL_REQUIRE(blackPrice>=0.0,
                   "blackPrice (" << blackPrice << ") must be non-negative");
        QL_REQUIRE(discount>0.0,
                   "discount (" << discount << ") must be positive");

        Real stdDev;
        forward = forward + displacement;
        strike = strike + displacement;
        if (strike==forward)
            // Brenner-Subrahmanyan (1988) and Feinstein (1988) ATM approx.
            stdDev = blackPrice/discount*std::sqrt(2.0 * M_PI)/forward;
        else {
            // Corrado and Miller extended moneyness approximation
            Real moneynessDelta = optionType*(forward-strike);
            Real moneynessDelta_2 = moneynessDelta/2.0;
            Real temp = blackPrice/discount - moneynessDelta_2;
            Real moneynessDelta_PI = moneynessDelta*moneynessDelta/M_PI;
            Real temp2 = temp*temp-moneynessDelta_PI;
            if (temp2<0.0) // approximation breaks down, 2 alternatives:
                // 1. zero it
                temp2=0.0;
                // 2. Manaster-Koehler (1982) efficient Newton-Raphson seed
                //return std::fabs(std::log(forward/strike))*std::sqrt(2.0);
            temp2 = std::sqrt(temp2);
            temp += temp2;
            temp *= std::sqrt(2.0 * M_PI);
            stdDev = temp/(forward+strike);
        }
        QL_ENSURE(stdDev>=0.0,
                  "stdDev (" << stdDev << ") must be non-negative");
        return stdDev;
    }

    Real blackFormulaImpliedStdDevApproximation(
                      const boost::shared_ptr<PlainVanillaPayoff>& payoff,
                      Real forward,
                      Real blackPrice,
                      Real discount,
                      Real displacement) {
        return blackFormulaImpliedStdDevApproximation(payoff->optionType(),
            payoff->strike(), forward, blackPrice, discount, displacement);
    }

    class BlackImpliedStdDevHelper {
      public:
        BlackImpliedStdDevHelper(Option::Type optionType,
                                 Real strike,
                                 Real forward,
                                 Real undiscountedBlackPrice,
                                 Real displacement = 0.0)
        : halfOptionType_(0.5*optionType), signedStrike_(optionType*(strike+displacement)),
          signedForward_(optionType*(forward+displacement)),
          undiscountedBlackPrice_(undiscountedBlackPrice)
        {
            checkParameters(strike, forward, displacement);
            QL_REQUIRE(undiscountedBlackPrice>=0.0,
                       "undiscounted Black price (" <<
                       undiscountedBlackPrice << ") must be non-negative");
            signedMoneyness_ = optionType*std::log((forward+displacement)/(strike+displacement));
        }
        Real operator()(Real stdDev) const {
            #if defined(QL_EXTRA_SAFETY_CHECKS)
            QL_REQUIRE(stdDev>=0.0,
                       "stdDev (" << stdDev << ") must be non-negative");
            #endif
            if (stdDev==0.0)
                return std::max(signedForward_-signedStrike_, Real(0.0))
                                                   - undiscountedBlackPrice_;
            Real temp = halfOptionType_*stdDev;
            Real d = signedMoneyness_/stdDev;
            Real signedD1 = d + temp;
            Real signedD2 = d - temp;
            Real result = signedForward_ * N_(signedD1)
                - signedStrike_ * N_(signedD2);
            // numerical inaccuracies can yield a negative answer
            return std::max(Real(0.0), result) - undiscountedBlackPrice_;
        }
        Real derivative(Real stdDev) const {
            #if defined(QL_EXTRA_SAFETY_CHECKS)
            QL_REQUIRE(stdDev>=0.0,
                       "stdDev (" << stdDev << ") must be non-negative");
            #endif
            Real signedD1 = signedMoneyness_/stdDev + halfOptionType_*stdDev;
            return signedForward_*N_.derivative(signedD1);
        }
      private:
        Real halfOptionType_;
        Real signedStrike_, signedForward_;
        Real undiscountedBlackPrice_, signedMoneyness_;
        CumulativeNormalDistribution N_;
    };

    Real blackFormulaImpliedStdDev(Option::Type optionType,
                                   Real strike,
                                   Real forward,
                                   Real blackPrice,
                                   Real discount,
                                   Real displacement,
                                   Real guess,
                                   Real accuracy,
                                   Natural maxIterations)
    {
        checkParameters(strike, forward, displacement);

        QL_REQUIRE(discount>0.0,
                   "discount (" << discount << ") must be positive");

        QL_REQUIRE(blackPrice>=0.0,
                   "option price (" << blackPrice << ") must be non-negative");
        // check the price of the "other" option implied by put-call paity
        Real otherOptionPrice = blackPrice - optionType*(forward-strike)*discount;
        QL_REQUIRE(otherOptionPrice>=0.0,
                   "negative " << Option::Type(-1*optionType) <<
                   " price (" << otherOptionPrice <<
                   ") implied by put-call parity. No solution exists for " <<
                   optionType << " strike " << strike <<
                   ", forward " << forward <<
                   ", price " << blackPrice <<
                   ", deflator " << discount);

        // solve for the out-of-the-money option which has
        // greater vega/price ratio, i.e.
        // it is numerically more robust for implied vol calculations
        if (optionType==Option::Put && strike>forward) {
            optionType = Option::Call;
            blackPrice = otherOptionPrice;
        }
        if (optionType==Option::Call && strike<forward) {
            optionType = Option::Put;
            blackPrice = otherOptionPrice;
        }

        strike = strike + displacement;
        forward = forward + displacement;

        if (guess==Null<Real>())
            guess = blackFormulaImpliedStdDevApproximation(
                optionType, strike, forward, blackPrice, discount, displacement);
        else
            QL_REQUIRE(guess>=0.0,
                       "stdDev guess (" << guess << ") must be non-negative");
        BlackImpliedStdDevHelper f(optionType, strike, forward,
                                   blackPrice/discount);
        NewtonSafe solver;
        solver.setMaxEvaluations(maxIterations);
        Real minSdtDev = 0.0, maxStdDev = 24.0; // 24 = 300% * sqrt(60)
        Real stdDev = solver.solve(f, accuracy, guess, minSdtDev, maxStdDev);
        QL_ENSURE(stdDev>=0.0,
                  "stdDev (" << stdDev << ") must be non-negative");
        return stdDev;
    }

    Real blackFormulaImpliedStdDev(
                        const boost::shared_ptr<PlainVanillaPayoff>& payoff,
                        Real forward,
                        Real blackPrice,
                        Real discount,
                        Real displacement,
                        Real guess,
                        Real accuracy,
                        Natural maxIterations) {
        return blackFormulaImpliedStdDev(payoff->optionType(), payoff->strike(),
            forward, blackPrice, discount, displacement, guess, accuracy, maxIterations);
    }

    Real blackFormulaCashItmProbability(Option::Type optionType,
                                        Real strike,
                                        Real forward,
                                        Real stdDev,
                                        Real displacement) {
        checkParameters(strike, forward, displacement);
        if (stdDev==0.0)
            return (forward*optionType > strike*optionType ? 1.0 : 0.0);

        forward = forward + displacement;
        strike = strike + displacement;
        if (strike==0.0)
            return (optionType==Option::Call ? 1.0 : 0.0);
        Real d2 = std::log(forward/strike)/stdDev - 0.5*stdDev;
        CumulativeNormalDistribution phi;
        return phi(optionType*d2);
    }

    Real blackFormulaCashItmProbability(
                        const boost::shared_ptr<PlainVanillaPayoff>& payoff,
                        Real forward,
                        Real stdDev,
                        Real displacement) {
        return blackFormulaCashItmProbability(payoff->optionType(),
            payoff->strike(), forward, stdDev , displacement);
    }

    Real blackFormulaVolDerivative(Rate strike,
                                      Rate forward,
                                      Real stdDev,
                                      Real expiry,
                                      Real discount,
                                      Real displacement)
    {
        return  blackFormulaStdDevDerivative(strike,
                                     forward,
                                     stdDev,
                                     discount,
                                     displacement)*std::sqrt(expiry);
    }

    Real blackFormulaStdDevDerivative(Rate strike,
                                      Rate forward,
                                      Real stdDev,
                                      Real discount,
                                      Real displacement)
    {
        checkParameters(strike, forward, displacement);
        QL_REQUIRE(stdDev>=0.0,
                   "stdDev (" << stdDev << ") must be non-negative");
        QL_REQUIRE(discount>0.0,
                   "discount (" << discount << ") must be positive");

        forward = forward + displacement;
        strike = strike + displacement;

        if (stdDev==0.0 || strike==0.0)
            return 0.0;

        Real d1 = std::log(forward/strike)/stdDev + .5*stdDev;
        return discount * forward *
            CumulativeNormalDistribution().derivative(d1);
    }

    Real blackFormulaStdDevDerivative(
                        const boost::shared_ptr<PlainVanillaPayoff>& payoff,
                        Real forward,
                        Real stdDev,
                        Real discount,
                        Real displacement) {
        return blackFormulaStdDevDerivative(payoff->strike(), forward,
                                     stdDev, discount, displacement);
    }

    Real bachelierBlackFormula(Option::Type optionType,
                               Real strike,
                               Real forward,
                               Real stdDev,
                               Real discount)
    {
        QL_REQUIRE(stdDev>=0.0,
                   "stdDev (" << stdDev << ") must be non-negative");
        QL_REQUIRE(discount>0.0,
                   "discount (" << discount << ") must be positive");
        Real d = (forward-strike)*optionType, h = d/stdDev;
        if (stdDev==0.0)
            return discount*std::max(d, 0.0);
        CumulativeNormalDistribution phi;
        Real result = discount*(stdDev*phi.derivative(h) + d*phi(h));
        QL_ENSURE(result>=0.0,
                  "negative value (" << result << ") for " <<
                  stdDev << " stdDev, " <<
                  optionType << " option, " <<
                  strike << " strike , " <<
                  forward << " forward");
        return result;
    }

    Real bachelierBlackFormula(
                        const boost::shared_ptr<PlainVanillaPayoff>& payoff,
                        Real forward,
                        Real stdDev,
                        Real discount) {
        return bachelierBlackFormula(payoff->optionType(),
            payoff->strike(), forward, stdDev, discount);
    }

    static Real h(Real eta) {

        const static Real  A0          = 3.994961687345134e-1;
        const static Real  A1          = 2.100960795068497e+1;
        const static Real  A2          = 4.980340217855084e+1;
        const static Real  A3          = 5.988761102690991e+2;
        const static Real  A4          = 1.848489695437094e+3;
        const static Real  A5          = 6.106322407867059e+3;
        const static Real  A6          = 2.493415285349361e+4;
        const static Real  A7          = 1.266458051348246e+4;

        const static Real  B0          = 1.000000000000000e+0;
        const static Real  B1          = 4.990534153589422e+1;
        const static Real  B2          = 3.093573936743112e+1;
        const static Real  B3          = 1.495105008310999e+3;
        const static Real  B4          = 1.323614537899738e+3;
        const static Real  B5          = 1.598919697679745e+4;
        const static Real  B6          = 2.392008891720782e+4;
        const static Real  B7          = 3.608817108375034e+3;
        const static Real  B8          = -2.067719486400926e+2;
        const static Real  B9          = 1.174240599306013e+1;

        QL_REQUIRE(eta>=0.0,
                       "eta (" << eta << ") must be non-negative");

        const Real num = A0 + eta * (A1 + eta * (A2 + eta * (A3 + eta * (A4 + eta
                    * (A5 + eta * (A6 + eta * A7))))));

        const Real den = B0 + eta * (B1 + eta * (B2 + eta * (B3 + eta * (B4 + eta
                    * (B5 + eta * (B6 + eta * (B7 + eta * (B8 + eta * B9))))))));

        return std::sqrt(eta) * (num / den);

    }

    Real bachelierBlackFormulaImpliedVol(Option::Type optionType,
                                   Real strike,
                                   Real forward,
                                   Real tte,
                                   Real bachelierPrice,
                                   Real discount) {

        const static Real SQRT_QL_EPSILON = std::sqrt(QL_EPSILON);

        QL_REQUIRE(tte>0.0,
                   "tte (" << tte << ") must be positive");

        Real forwardPremium = bachelierPrice/discount;

        Real straddlePremium;
        if (optionType==Option::Call){
            straddlePremium = 2.0 * forwardPremium - (forward - strike);
        } else {
            straddlePremium = 2.0 * forwardPremium + (forward - strike);
        }

        Real nu = (forward - strike) / straddlePremium;
        QL_REQUIRE(nu<=1.0,
                   "nu (" << nu << ") must be <= 1.0");
        QL_REQUIRE(nu>=-1.0,
                     "nu (" << nu << ") must be >= -1.0");

        nu = std::max(-1.0 + QL_EPSILON, std::min(nu,1.0 - QL_EPSILON));

        // nu / arctanh(nu) -> 1 as nu -> 0
        Real eta = (std::fabs(nu) < SQRT_QL_EPSILON) ? 1.0 : nu / boost::math::atanh(nu);

        Real heta = h(eta);

        Real impliedBpvol = std::sqrt(M_PI / (2 * tte)) * straddlePremium * heta;

        return impliedBpvol;
    }
}


/**
 * @author     : Zhao Chonyyao (cyzhao@zju.edu.cn)
 * @date       : 2021-04-30
 * @description: some basic energy function.
 * @version    : 1.0
 */
#include <iostream>

#include <Eigen/SparseCore>

#include "basic_energy.h"

using namespace std;
using namespace Eigen;

namespace PhysIKA {
/******************************************momentum*******************************/

template <typename T, size_t dim_>
size_t momentum<T, dim_>::Nx() const
{
    return dim_ * dof_;
}

template <typename T, size_t dim_>
momentum<T, dim_>::momentum(const size_t dof, const Matrix<T, -1, 1>& mass_vec, const T& dt)
    : dof_(dof), dispk_(Matrix<T, -1, 1>::Zero(dim_ * dof)), vk_(Matrix<T, -1, 1>::Zero(dim_ * dof)), dt_(dt), d1dt_(1 / dt), d1dtdt_(1 / dt / dt), mass_vec_(dim_ * dof)
{

#pragma omp parallel for
    for (size_t i = 0; i < dof; ++i)
    {
        for (size_t j = 0; j < dim_; ++j)
        {
            mass_vec_(i * dim_ + j) = mass_vec(i);
        }
    }
}

template <typename T, size_t dim_>
momentum<T, dim_>::momentum(const T* rest, const size_t dof, const Matrix<T, -1, 1>& mass_vec, const T& dt)
    : dof_(dof), vk_(Matrix<T, -1, 1>::Zero(dim_ * dof)), dt_(dt), d1dt_(1 / dt), d1dtdt_(1 / dt / dt), mass_vec_(dim_ * dof)
{

    dispk_ = Eigen::Map<const Matrix<T, -1, 1>>(rest, dim_ * dof);
#pragma omp parallel for
    for (size_t i = 0; i < dof; ++i)
    {
        for (size_t j = 0; j < dim_; ++j)
        {
            mass_vec_(i * dim_ + j) = mass_vec(i);
        }
    }
}

template <typename T, size_t dim_>
int momentum<T, dim_>::Val(const T* x, data_ptr<T, dim_>& data) const
{
    Eigen::Map<const Matrix<T, -1, 1>> _x(x, dim_ * dof_);
    const Matrix<T, -1, 1>             acce = (_x - dispk_) * d1dt_ - vk_;
    data->save_val(0.5 * acce.dot(mass_vec_.cwiseProduct(acce)));

    return 0;
}
template <typename T, size_t dim_>
int momentum<T, dim_>::Gra(const T* x, data_ptr<T, dim_>& data) const
{
    Eigen::Map<const Matrix<T, -1, 1>> _x(x, dim_ * dof_);

    const Matrix<T, -1, 1> acce = (_x - dispk_) * d1dtdt_ - vk_ * d1dt_;
    data->save_gra(mass_vec_.cwiseProduct(acce));

    return 0;
}

template <typename T, size_t dim_>
int momentum<T, dim_>::Hes(const T* x, data_ptr<T, dim_>& data) const
{
#pragma omp parallel for
    for (size_t i = 0; i < dim_ * dof_; ++i)
        data->save_hes(i, i, d1dtdt_ * mass_vec_(i));

    return 0;
}

template <typename T, size_t dim_>
int momentum<T, dim_>::update_location_and_velocity(const T* new_dispk_ptr, const T* new_velo_ptr)
{
    Eigen::Map<const Matrix<T, -1, 1>> new_dispk(new_dispk_ptr, dim_ * dof_);
    if (new_velo_ptr == nullptr)
        vk_ = (new_dispk - dispk_) * d1dt_;
    else
        vk_ = Map<const Matrix<T, -1, 1>>(new_velo_ptr, dim_ * dof_);

    dispk_ = new_dispk;
    return 0;
}

template <typename T, size_t dim_>
int momentum<T, dim_>::set_initial_velocity(const Matrix<T, dim_, 1>& velo)
{
    Eigen::Map<Matrix<T, -1, -1>> myvelo(vk_.data(), dim_, dof_);
    for (size_t i = 0; i < dim_; ++i)
    {
        myvelo.row(i) = Matrix<T, 1, -1>::Ones(dof_) * velo(i);
    }
    return 0;
}

/******************************************momentum*******************************/
/******************************************position_constraint*******************************/
template <typename T, size_t dim_>
position_constraint<T, dim_>::position_constraint(const size_t dof, const T& w, const vector<size_t>& cons)
    : w_(w), cons_(cons), dof_(dof)
{
    rest_ = Matrix<T, -1, -1>::Zero(dim_, dof);
}

template <typename T, size_t dim_>
position_constraint<T, dim_>::position_constraint(const T* rest, const size_t dof, const T& w, const std::vector<size_t>& cons)
    : w_(w), cons_(cons), dof_(dof)
{
    rest_ = Eigen::Map<const Matrix<T, -1, -1>>(rest, dim_, dof_);
}

//TODO: simplify _x
template <typename T, size_t dim_>
int position_constraint<T, dim_>::Val(const T* x, data_ptr<T, dim_>& data) const
{
    Eigen::Map<const Matrix<T, -1, -1>> deformed(x, dim_, dof_);
    Matrix<T, -1, -1>                   _x = deformed - rest_;
    for (auto iter_c = cons_.begin(); iter_c != cons_.end(); ++iter_c)
    {
        data->save_val(w_ * _x.col(*iter_c).dot(_x.col(*iter_c)));
    }

    return 0;
}

template <typename T, size_t dim_>
int position_constraint<T, dim_>::Gra(const T* x, data_ptr<T, dim_>& data) const
{
    Eigen::Map<const Matrix<T, -1, -1>> deformed(x, dim_, dof_);
    Matrix<T, -1, -1>                   _x = deformed - rest_;

    for (auto iter_c = cons_.begin(); iter_c != cons_.end(); ++iter_c)
        data->save_gra(*iter_c, 2.0 * w_ * _x.col(*iter_c));
    return 0;
}

template <typename T, size_t dim_>
int position_constraint<T, dim_>::Hes(const T* x, data_ptr<T, dim_>& data) const
{
    for (auto iter_c = cons_.begin(); iter_c != cons_.end(); ++iter_c)
    {
        for (size_t j = 0; j < dim_; ++j)
        {
            data->save_hes(*iter_c * dim_ + j, *iter_c * dim_ + j, 2 * w_);
        }
    }
    return 0;
}

template <typename T, size_t dim_>
size_t position_constraint<T, dim_>::Nx() const
{
    return dim_ * dof_;
}
/******************************************position_constraint*******************************/

/******************************************gravity*******************************/
//dof here is not about dim_
template <typename T, size_t dim_>
gravity_energy<T, dim_>::gravity_energy(const size_t dof, const T& w_g, const T& gravity, const Matrix<T, -1, 1>& mass, const char& axis)
    : w_g_(w_g), dof_(dof), gravity_(gravity), mass_(mass), axis_(axis)
{
}

template <typename T, size_t dim_>
int gravity_energy<T, dim_>::Val(const T* x, data_ptr<T, dim_>& data) const
{

    Eigen::Map<const Matrix<T, -1, -1>> _x(x, dim_, dof_);
    size_t                              which_axis = size_t(axis_ - 'x');
    data->save_val((_x.row(which_axis).transpose().array() * mass_.array()).sum() * w_g_ * gravity_);
    return 0;
}

template <typename T, size_t dim_>
int gravity_energy<T, dim_>::Gra(const T* x, data_ptr<T, dim_>& data) const
{
    size_t which_axis = size_t(axis_ - 'x');

    Matrix<T, -1, -1> g(dim_, dof_);
    g.setZero();
    g.row(which_axis) = Matrix<T, -1, 1>::Constant(dof_, gravity_ * w_g_).cwiseProduct(mass_).transpose();
    Eigen::Map<const Matrix<T, -1, 1>> g_(g.data(), dim_ * dof_);
    data->save_gra(g_);
    return 0;
}
template <typename T, size_t dim_>
int gravity_energy<T, dim_>::Hes(const T* x, data_ptr<T, dim_>& data) const
{
    return 0;
}

template <typename T, size_t dim_>
size_t gravity_energy<T, dim_>::Nx() const
{
    return dim_ * dof_;
}
/******************************************gravity*******************************/

/*************************************collision*********************************/
template <typename T, size_t dim_>
collision<T, dim_>::collision(const size_t dof_, const T& w_coll, const char& ground_axis, const T& ground_pos, const size_t& num_surf_point, const shared_ptr<Matrix<T, -1, -1>>& init_points_ptr)
    : ground_axis_(ground_axis), ground_pos_(ground_pos), w_coll_(w_coll), num_surf_point_(num_surf_point), dof_(dof_), init_points_ptr_(init_points_ptr)
{
}

template <typename T, size_t dim_>
int collision<T, dim_>::Val(const T* x, data_ptr<T, dim_>& data) const
{
    const size_t which_axis = size_t(ground_axis_ - 'x');

    Eigen::Map<const Matrix<T, -1, -1>> _x(x, dim_, dof_);
#pragma omp parallel for
    for (size_t i = 0; i < dof_; ++i)
    {
        const T position_now = _x(which_axis, i) + (*init_points_ptr_)(which_axis, i);
        if ((position_now - ground_pos_) < 0)
        {

            data->save_val(w_coll_ * pow((ground_pos_ - position_now), 2));
        }
    }
    return 0;
}
template <typename T, size_t dim_>
int collision<T, dim_>::Gra(const T* x, data_ptr<T, dim_>& data) const
{
    const size_t which_axis = size_t(ground_axis_ - 'x');

    Eigen::Map<const Matrix<T, -1, -1>> _x(x, dim_, dof_);
#pragma omp parallel for
    for (size_t i = 0; i < dof_; ++i)
    {
        const T position_now = _x(which_axis, i) + (*init_points_ptr_)(which_axis, i);
        if ((position_now - ground_pos_) < 0)
        {
            data->save_gra(i * dim_ + which_axis, 2 * w_coll_ * (position_now - ground_pos_));
        }
    }
    return 0;
}
template <typename T, size_t dim_>
int collision<T, dim_>::Hes(const T* x, data_ptr<T, dim_>& data) const
{
    const size_t                        which_axis = size_t(ground_axis_ - 'x');
    Eigen::Map<const Matrix<T, -1, -1>> _x(x, dim_, dof_);
#pragma omp parallel for
    for (size_t i = 0; i < dof_; ++i)
    {
        const T position_now = _x(which_axis, i) + (*init_points_ptr_)(which_axis, i);
        if ((position_now - ground_pos_) < 0)
        {
            for (size_t j = 0; j < dim_; ++j)
            {
                data->save_hes(i * dim_ + j, i * dim_ + j, 2 * w_coll_);
            }
        }
    }

    return 0;
}
template <typename T, size_t dim_>
size_t collision<T, dim_>::Nx() const
{
    return dim_ * dof_;
}

/*************************************collision*********************************/

template class position_constraint<double, 3>;
template class position_constraint<float, 3>;
template class position_constraint<float, 1>;
template class position_constraint<double, 1>;
template class momentum<double, 3>;
template class momentum<float, 3>;
template class gravity_energy<double, 3>;
template class gravity_energy<float, 3>;
// template class collision<FLOAT_TYPE, 3>;

}  //namespace PhysIKA


/* +------------------------------------------------------------------------+
   |                     Mobile Robot Programming Toolkit (MRPT)            |
   |                          https://www.mrpt.org/                         |
   |                                                                        |
   | Copyright (c) 2005-2021, Individual contributors, see AUTHORS file     |
   | See: https://www.mrpt.org/Authors - All rights reserved.               |
   | Released under BSD License. See: https://www.mrpt.org/License          |
   +------------------------------------------------------------------------+ */

#include "math-precomp.h"  // Precompiled headers
//
#include <mrpt/math/TLine3D.h>
#include <mrpt/math/TSegment3D.h>
#include <mrpt/math/geometry.h>	 // distance()
#include <mrpt/serialization/CArchive.h>  // impl of << operator

#include <Eigen/Dense>

using namespace mrpt::math;

void TSegment3D::generate2DObject(TSegment2D& s) const
{
	s = TSegment2D(*this);
}

double TSegment3D::length() const { return math::distance(point1, point2); }
double TSegment3D::distance(const TPoint3D& point) const
{
	return std::min(
		std::min(math::distance(point, point1), math::distance(point, point2)),
		TLine3D(*this).distance(point));
}
double TSegment3D::distance(const TSegment3D& segment) const
{
	Eigen::Vector3d u, v, w;
	TPoint3D diff_vect = point2 - point1;
	diff_vect.asVector(u);
	diff_vect = segment.point2 - segment.point1;
	diff_vect.asVector(v);
	diff_vect = point1 - segment.point1;
	diff_vect.asVector(w);
	double a = u.dot(u);  // always >= 0
	double b = u.dot(v);
	double c = v.dot(v);  // always >= 0
	double d = u.dot(w);
	double e = v.dot(w);
	double D = a * c - b * b;  // always >= 0
	double sc, sN, sD = D;	// sc = sN / sD, default sD = D >= 0
	double tc, tN, tD = D;	// tc = tN / tD, default tD = D >= 0

	// compute the line parameters of the two closest points
	if (D < 0.00000001)
	{  // the lines are almost parallel
		sN = 0.0;  // force using point P0 on segment S1
		sD = 1.0;  // to prevent possible division by 0.0 later
		tN = e;
		tD = c;
	}
	else
	{  // get the closest points on the infinite lines
		sN = (b * e - c * d);
		tN = (a * e - b * d);
		if (sN < 0.0)
		{  // sc < 0 => the s=0 edge is visible
			sN = 0.0;
			tN = e;
			tD = c;
		}
		else if (sN > sD)
		{  // sc > 1 => the s=1 edge is visible
			sN = sD;
			tN = e + b;
			tD = c;
		}
	}

	if (tN < 0.0)
	{  // tc < 0 => the t=0 edge is visible
		tN = 0.0;
		// recompute sc for this edge
		if (-d < 0.0) sN = 0.0;
		else if (-d > a)
			sN = sD;
		else
		{
			sN = -d;
			sD = a;
		}
	}
	else if (tN > tD)
	{  // tc > 1 => the t=1 edge is visible
		tN = tD;
		// recompute sc for this edge
		if ((-d + b) < 0.0) sN = 0;
		else if ((-d + b) > a)
			sN = sD;
		else
		{
			sN = (-d + b);
			sD = a;
		}
	}
	// finally do the division to get sc and tc
	sc = (fabs(sN) < 0.00000001 ? 0.0 : sN / sD);
	tc = (fabs(tN) < 0.00000001 ? 0.0 : tN / tD);

	// get the difference of the two closest points
	const auto dP = (w + (sc * u) - (tc * v)).eval();
	return dP.norm();  // return the closest distance
}
bool TSegment3D::contains(const TPoint3D& point) const
{
	// Not very intuitive, but very fast, method.
	return std::abs(
			   math::distance(point1, point) + math::distance(point2, point) -
			   math::distance(point1, point2)) < getEpsilon();
}

bool TSegment3D::operator<(const TSegment3D& s) const
{
	if (point1 < s.point1) return true;
	else if (s.point1 < point1)
		return false;
	else
		return point2 < s.point2;
}

mrpt::serialization::CArchive& mrpt::math::operator>>(
	mrpt::serialization::CArchive& in, mrpt::math::TSegment3D& s)
{
	return in >> s.point1 >> s.point2;
}
mrpt::serialization::CArchive& mrpt::math::operator<<(
	mrpt::serialization::CArchive& out, const mrpt::math::TSegment3D& s)
{
	return out << s.point1 << s.point2;
}