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

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

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

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

#include <qle/models/crossassetanalytics.hpp>
#include <qle/models/crossassetmodel.hpp>

#include <ql/math/matrixutilities/pseudosqrt.hpp>
#include <ql/processes/eulerdiscretization.hpp>

#include <boost/make_shared.hpp>

using namespace QuantLib;

namespace {
static inline void setValue(Matrix& m, const Real& value, const QuantExt::CrossAssetModel* model,
                            const QuantExt::CrossAssetModelTypes::AssetType& t1, const Size& i1,
                            const QuantExt::CrossAssetModelTypes::AssetType& t2, const Size& i2,
                            const Size& offset1 = 0, const Size& offset2 = 0) {
    Size i = model->pIdx(t1, i1, offset1);
    Size j = model->pIdx(t2, i2, offset2);
    m[i][j] = m[j][i] = value;
}
} // anonymous namespace

namespace QuantExt {

using namespace CrossAssetAnalytics;

CrossAssetStateProcess::CrossAssetStateProcess(const CrossAssetModel* const model, discretization disc,
                                               SalvagingAlgorithm::Type salvaging)
    : StochasticProcess(), model_(model), salvaging_(salvaging) {

    if (disc == euler) {
        discretization_ = boost::make_shared<EulerDiscretization>();
    } else {
        discretization_ = boost::make_shared<CrossAssetStateProcess::ExactDiscretization>(model, salvaging);
    }
}

Size CrossAssetStateProcess::size() const { return model_->dimension(); }

void CrossAssetStateProcess::flushCache() const {
    cache_m_.clear();
    cache_v_.clear();
    cache_d_.clear();
    boost::shared_ptr<CrossAssetStateProcess::ExactDiscretization> tmp =
        boost::dynamic_pointer_cast<CrossAssetStateProcess::ExactDiscretization>(discretization_);
    if (tmp != NULL) {
        tmp->flushCache();
    }
}

Disposable<Array> CrossAssetStateProcess::initialValues() const {
    Array res(model_->dimension(), 0.0);
    /* irlgm1f processes have initial value 0 */
    for (Size i = 0; i < model_->components(FX); ++i) {
        /* fxbs processes are in log spot */
        res[model_->pIdx(FX, i, 0)] = std::log(model_->fxbs(i)->fxSpotToday()->value());
    }
    for (Size i = 0; i < model_->components(EQ); ++i) {
        /* eqbs processes are in log spot */
        res[model_->pIdx(EQ, i, 0)] = std::log(model_->eqbs(i)->eqSpotToday()->value());
    }
    /* infdk, crlgm1f processes have initial value 0 */
    return res;
}

Disposable<Array> CrossAssetStateProcess::drift(Time t, const Array& x) const {
    Array res(model_->dimension(), 0.0);
    Size n = model_->components(IR);
    Size n_eq = model_->components(EQ);
    Real H0 = model_->irlgm1f(0)->H(t);
    Real Hprime0 = model_->irlgm1f(0)->Hprime(t);
    Real alpha0 = model_->irlgm1f(0)->alpha(t);
    Real zeta0 = model_->irlgm1f(0)->zeta(t);
    boost::unordered_map<double, Array>::const_iterator i = cache_m_.find(t);
    if (i == cache_m_.end()) {
        /* z0 has drift 0 */
        for (Size i = 1; i < n; ++i) {
            Real Hi = model_->irlgm1f(i)->H(t);
            Real alphai = model_->irlgm1f(i)->alpha(t);
            Real sigmai = model_->fxbs(i - 1)->sigma(t);
            // ir-ir
            Real rhozz0i = model_->correlation(IR, 0, IR, i);
            // ir-fx
            Real rhozx0i = model_->correlation(IR, 0, FX, i - 1);
            Real rhozxii = model_->correlation(IR, i, FX, i - 1);
            // ir drifts
            res[model_->pIdx(IR, i, 0)] =
                -Hi * alphai * alphai + H0 * alpha0 * alphai * rhozz0i - sigmai * alphai * rhozxii;
            // log spot fx drifts (z0, zi independent parts)
            res[model_->pIdx(FX, i - 1, 0)] =
                H0 * alpha0 * sigmai * rhozx0i + model_->irlgm1f(0)->termStructure()->forwardRate(t, t, Continuous) -
                model_->irlgm1f(i)->termStructure()->forwardRate(t, t, Continuous) - 0.5 * sigmai * sigmai;
        }
        /* log equity spot drifts (the cache-able parts) */
        for (Size k = 0; k < n_eq; ++k) {
            Size i = model_->ccyIndex(model_->eqbs(k)->currency());
            // ir params (for equity currency)
            Real eps_ccy = (i == 0) ? 0.0 : 1.0;
            // Real Hi = model_->irlgm1f(i)->H(t);
            // Real alphai = model_->irlgm1f(i)->alpha(t);
            // eq vol
            Real sigmask = model_->eqbs(k)->sigma(t);
            // fx vol (eq ccy / base ccy)
            Real sigmaxi = (i == 0) ? 0.0 : model_->fxbs(i - 1)->sigma(t);
            // ir-eq corr
            // Real rhozsik = model_->correlation(EQ, k, IR, i); // eq cur
            Real rhozs0k = model_->correlation(EQ, k, IR, 0); // base cur
            // fx-eq corr
            Real rhoxsik = (i == 0) ? 0.0 : // no fx process for base-ccy
                               model_->correlation(FX, i - 1, EQ, k);
            // ir instantaneous forward rate (from curve used for eq forward projection)
            Real fr_i = model_->eqbs(k)->equityIrCurveToday()->forwardRate(t, t, Continuous);
            // div yield instantaneous forward rate
            Real fq_k = model_->eqbs(k)->equityDivYieldCurveToday()->forwardRate(t, t, Continuous);
            res[model_->pIdx(EQ, k, 0)] = fr_i - fq_k + (rhozs0k * H0 * alpha0 * sigmask) -
                                          (eps_ccy * rhoxsik * sigmaxi * sigmask) - (0.5 * sigmask * sigmask);
        }
        cache_m_.insert(std::make_pair(t, res));
    } else {
        res = i->second;
    }
    // non-cacheable sections of drifts
    for (Size i = 1; i < n; ++i) {
        // log spot fx drifts (z0, zi dependent parts)
        Real Hi = model_->irlgm1f(i)->H(t);
        Real Hprimei = model_->irlgm1f(i)->Hprime(t);
        Real zetai = model_->irlgm1f(i)->zeta(t);
        res[model_->pIdx(FX, i - 1, 0)] += x[model_->pIdx(IR, 0, 0)] * Hprime0 + zeta0 * Hprime0 * H0 -
                                           x[model_->pIdx(IR, i, 0)] * Hprimei - zetai * Hprimei * Hi;
    }
    for (Size k = 0; k < n_eq; ++k) {
        // log equity spot drifts (path-dependent parts)
        // notice the assumption in below that dividend yield curve is static
        Size i = model_->ccyIndex(model_->eqbs(k)->currency());
        // ir params (for equity currency)
        Real Hi = model_->irlgm1f(i)->H(t);
        Real Hprimei = model_->irlgm1f(i)->Hprime(t);
        Real zetai = model_->irlgm1f(i)->zeta(t);
        res[model_->pIdx(EQ, k, 0)] += (x[model_->pIdx(IR, i, 0)] * Hprimei) + (zetai * Hprimei * Hi);
    }
    /* no drift for infdk, crlgm1f components */
    return res;
}

Disposable<Matrix> CrossAssetStateProcess::diffusion(Time t, const Array& x) const {
    boost::unordered_map<double, Matrix>::const_iterator i = cache_d_.find(t);
    if (i == cache_d_.end()) {
        Matrix tmp = pseudoSqrt(diffusionImpl(t, x), salvaging_);
        cache_d_.insert(std::make_pair(t, tmp));
        return tmp;
    } else {
        // we have to make a copy, otherwise we destroy the map entry
        // since a disposable is returned
        Matrix tmp = i->second;
        return tmp;
    }
}

Disposable<Matrix> CrossAssetStateProcess::diffusionImpl(Time t, const Array&) const {
    Matrix res(model_->dimension(), model_->dimension());
    Size n = model_->components(IR);
    Size m = model_->components(FX);
    Size d = model_->components(INF);
    Size c = model_->components(CR);
    Size e = model_->components(EQ);
    // ir-ir
    for (Size i = 0; i < n; ++i) {
        for (Size j = 0; j <= i; ++j) {
            Real alphai = model_->irlgm1f(i)->alpha(t);
            Real alphaj = model_->irlgm1f(j)->alpha(t);
            Real rhozz = model_->correlation(IR, i, IR, j, 0, 0);
            setValue(res, alphai * alphaj * rhozz, model_, IR, i, IR, j, 0, 0);
        }
    }
    // ir-fx
    for (Size i = 0; i < n; ++i) {
        Real alphai = model_->irlgm1f(i)->alpha(t);
        for (Size j = 0; j < m; ++j) {
            Real sigmaj = model_->fxbs(j)->sigma(t);
            Real rhozx = model_->correlation(IR, i, FX, j, 0, 0);
            setValue(res, alphai * sigmaj * rhozx, model_, IR, i, FX, j, 0, 0);
        }
    }
    // fx-fx
    for (Size i = 0; i < m; ++i) {
        Real sigmai = model_->fxbs(i)->sigma(t);
        for (Size j = 0; j <= i; ++j) {
            Real sigmaj = model_->fxbs(j)->sigma(t);
            Real rhoxx = model_->correlation(FX, i, FX, j, 0, 0);
            setValue(res, sigmai * sigmaj * rhoxx, model_, FX, i, FX, j, 0, 0);
        }
    }
    // ir,fx,inf - inf
    for (Size j = 0; j < d; ++j) {
        Real alphaj = model_->infdk(j)->alpha(t);
        Real Hj = model_->infdk(j)->H(t);
        for (Size i = 0; i <= j; ++i) {
            Real alphai = model_->infdk(i)->alpha(t);
            Real Hi = model_->infdk(i)->H(t);
            Real rhoyy = model_->correlation(INF, i, INF, j, 0, 0);
            // infz-infz
            setValue(res, alphai * alphaj * rhoyy, model_, INF, i, INF, j, 0, 0);
            // infz-infy
            setValue(res, alphai * alphaj * Hj * rhoyy, model_, INF, i, INF, j, 0, 1);
            setValue(res, alphai * Hi * alphaj * rhoyy, model_, INF, i, INF, j, 1, 0);
            // infy-infy
            setValue(res, alphai * Hi * alphaj * Hj * rhoyy, model_, INF, i, INF, j, 1, 1);
        }
        for (Size i = 0; i < n; ++i) {
            Real alphai = model_->irlgm1f(i)->alpha(t);
            Real rhozy = model_->correlation(IR, i, INF, j, 0, 0);
            // ir-inf
            setValue(res, alphai * alphaj * rhozy, model_, IR, i, INF, j, 0, 0);
            setValue(res, alphai * alphaj * Hj * rhozy, model_, IR, i, INF, j, 0, 1);
        }
        for (Size i = 0; i < (n - 1); ++i) {
            Real sigmai = model_->fxbs(i)->sigma(t);
            Real rhoxy = model_->correlation(FX, i, INF, j, 0, 0);
            // fx-inf
            setValue(res, sigmai * alphaj * rhoxy, model_, FX, i, INF, j, 0, 0);
            setValue(res, sigmai * alphaj * Hj * rhoxy, model_, FX, i, INF, j, 0, 1);
        }
    }
    // ir,fx,inf,cr - cr
    for (Size j = 0; j < c; ++j) {
        Real alphaj = model_->crlgm1f(j)->alpha(t);
        Real Hj = model_->crlgm1f(j)->H(t);
        for (Size i = 0; i <= j; ++i) {
            Real alphai = model_->crlgm1f(i)->alpha(t);
            Real Hi = model_->crlgm1f(i)->H(t);
            Real rholl = model_->correlation(CR, i, CR, j, 0, 0);
            // crz-crz
            setValue(res, alphai * alphaj * rholl, model_, CR, i, CR, j, 0, 0);
            // crz-cry
            setValue(res, alphai * alphaj * Hj * rholl, model_, CR, i, CR, j, 0, 1);
            setValue(res, alphai * Hi * alphaj * rholl, model_, CR, i, CR, j, 1, 0);
            // cry-cry
            setValue(res, alphai * alphaj * Hi * Hj * rholl, model_, CR, i, CR, j, 1, 1);
        }
        for (Size i = 0; i < n; ++i) {
            Real alphai = model_->irlgm1f(i)->alpha(t);
            Real rhozl = model_->correlation(IR, i, CR, j, 0, 0);
            // ir-cr
            setValue(res, alphai * alphaj * rhozl, model_, IR, i, CR, j, 0, 0);
            setValue(res, alphai * alphaj * Hj * rhozl, model_, IR, i, CR, j, 0, 1);
        }
        for (Size i = 0; i < (n - 1); ++i) {
            Real sigmai = model_->fxbs(i)->sigma(t);
            Real rhoxl = model_->correlation(FX, i, CR, j, 0, 0);
            // fx-cr
            setValue(res, sigmai * alphaj * rhoxl, model_, FX, i, CR, j, 0, 0);
            setValue(res, sigmai * alphaj * Hj * rhoxl, model_, FX, i, CR, j, 0, 1);
        }
        for (Size i = 0; i < d; ++i) {
            Real alphai = model_->infdk(i)->alpha(t);
            Real Hi = model_->infdk(i)->H(t);
            Real rhoyl = model_->correlation(INF, i, CR, j, 0, 0);
            // inf-cr
            setValue(res, alphai * alphaj * rhoyl, model_, INF, i, CR, j, 0, 0);
            setValue(res, Hi * alphai * alphaj * rhoyl, model_, INF, i, CR, j, 1, 0);
            setValue(res, alphai * alphaj * Hj * rhoyl, model_, INF, i, CR, j, 0, 1);
            setValue(res, alphai * Hi * alphaj * Hj * rhoyl, model_, INF, i, CR, j, 1, 1);
        }
    }
    // ir,fx,inf,cr,eq - eq
    for (Size j = 0; j < e; ++j) {
        Real sigmaj = model_->eqbs(j)->sigma(t);
        // ir-eq
        for (Size i = 0; i < n; ++i) {
            Real alphai = model_->irlgm1f(i)->alpha(t);
            Real rhozs = model_->correlation(IR, i, EQ, j, 0, 0);
            Real value = alphai * sigmaj * rhozs;
            setValue(res, value, model_, IR, i, EQ, j, 0, 0);
        }
        // fx-eq
        for (Size i = 0; i < m; ++i) {
            Real sigmai = model_->fxbs(i)->sigma(t);
            Real rhoxs = model_->correlation(FX, i, EQ, j, 0, 0);
            Real value = sigmai * sigmaj * rhoxs;
            setValue(res, value, model_, FX, i, EQ, j, 0, 0);
        }
        // inf-eq
        for (Size i = 0; i < d; ++i) {
            Real alphai = model_->infdk(i)->alpha(t);
            Real Hi = model_->infdk(i)->H(t);
            Real rhoys = model_->correlation(INF, i, EQ, j, 0, 0);
            setValue(res, alphai * sigmaj * rhoys, model_, INF, i, EQ, j, 0, 0);
            setValue(res, Hi * alphai * sigmaj * rhoys, model_, INF, i, EQ, j, 1, 0);
        }
        // cr-eq
        for (Size i = 0; i < c; ++i) {
            Real alphai = model_->crlgm1f(i)->alpha(t);
            Real Hi = model_->crlgm1f(i)->H(t);
            Real rhols = model_->correlation(CR, i, EQ, j, 0, 0);
            setValue(res, alphai * sigmaj * rhols, model_, CR, i, EQ, j, 0, 0);
            setValue(res, Hi * alphai * sigmaj * rhols, model_, CR, i, EQ, j, 1, 0);
        }
        // eq-eq
        for (Size i = 0; i <= j; ++i) {
            Real sigmai = model_->eqbs(i)->sigma(t);
            Real rhoss = model_->correlation(EQ, i, EQ, j, 0, 0);
            Real value = sigmai * sigmaj * rhoss;
            setValue(res, value, model_, EQ, i, EQ, j, 0, 0);
        }
    }
    return res;
}

CrossAssetStateProcess::ExactDiscretization::ExactDiscretization(const CrossAssetModel* const model,
                                                                 SalvagingAlgorithm::Type salvaging)
    : model_(model), salvaging_(salvaging) {}

Disposable<Array> CrossAssetStateProcess::ExactDiscretization::drift(const StochasticProcess& p, Time t0,
                                                                     const Array& x0, Time dt) const {
    Array res;
    cache_key k = { t0, dt };
    boost::unordered_map<cache_key, Array>::const_iterator i = cache_m_.find(k);
    if (i == cache_m_.end()) {
        res = driftImpl1(p, t0, x0, dt);
        cache_m_.insert(std::make_pair(k, res));
    } else {
        res = i->second;
    }
    Array res2 = driftImpl2(p, t0, x0, dt);
    for (Size i = 0; i < res.size(); ++i) {
        res[i] += res2[i];
    }
    return res - x0;
}

Disposable<Matrix> CrossAssetStateProcess::ExactDiscretization::diffusion(const StochasticProcess& p, Time t0,
                                                                          const Array& x0, Time dt) const {
    cache_key k = { t0, dt };
    boost::unordered_map<cache_key, Matrix>::const_iterator i = cache_d_.find(k);
    if (i == cache_d_.end()) {
        Matrix res = pseudoSqrt(covariance(p, t0, x0, dt), salvaging_);
        // note that covariance actually does not depend on x0
        cache_d_.insert(std::make_pair(k, res));
        return res;
    } else {
        // see above about the copy
        Matrix tmp = i->second;
        return tmp;
    }
}

Disposable<Matrix> CrossAssetStateProcess::ExactDiscretization::covariance(const StochasticProcess& p, Time t0,
                                                                           const Array& x0, Time dt) const {
    cache_key k = { t0, dt };
    boost::unordered_map<cache_key, Matrix>::const_iterator i = cache_v_.find(k);
    if (i == cache_v_.end()) {
        Matrix res = covarianceImpl(p, t0, x0, dt);
        cache_v_.insert(std::make_pair(k, res));
        return res;
    } else {
        // see above about the copy
        Matrix tmp = i->second;
        return tmp;
    }
}

Disposable<Array> CrossAssetStateProcess::ExactDiscretization::driftImpl1(const StochasticProcess&, Time t0,
                                                                          const Array&, Time dt) const {
    Size n = model_->components(IR);
    Size m = model_->components(FX);
    Size e = model_->components(EQ);
    Array res(model_->dimension(), 0.0);
    for (Size i = 0; i < n; ++i) {
        res[model_->pIdx(IR, i, 0)] = ir_expectation_1(model_, i, t0, dt);
    }
    for (Size j = 0; j < m; ++j) {
        res[model_->pIdx(FX, j, 0)] = fx_expectation_1(model_, j, t0, dt);
    }
    for (Size k = 0; k < e; ++k) {
        res[model_->pIdx(EQ, k, 0)] = eq_expectation_1(model_, k, t0, dt);
    }
    return res;
}

Disposable<Array> CrossAssetStateProcess::ExactDiscretization::driftImpl2(const StochasticProcess&, Time t0,
                                                                          const Array& x0, Time dt) const {
    Size n = model_->components(IR);
    Size m = model_->components(FX);
    Size e = model_->components(EQ);
    Array res(model_->dimension(), 0.0);
    for (Size i = 0; i < n; ++i) {
        res[model_->pIdx(IR, i, 0)] += ir_expectation_2(model_, i, x0[model_->pIdx(IR, i, 0)]);
    }
    for (Size j = 0; j < m; ++j) {
        res[model_->pIdx(FX, j, 0)] += fx_expectation_2(model_, j, t0, x0[model_->pIdx(FX, j, 0)],
                                                        x0[model_->pIdx(IR, j + 1, 0)], x0[model_->pIdx(IR, 0, 0)], dt);
    }
    for (Size k = 0; k < e; ++k) {
        Size eqCcyIdx = model_->ccyIndex(model_->eqbs(k)->currency());
        res[model_->pIdx(EQ, k, 0)] +=
            eq_expectation_2(model_, k, t0, x0[model_->pIdx(EQ, k, 0)], x0[model_->pIdx(IR, eqCcyIdx, 0)], dt);
    }
    /*! inf, cr components have integrated drift 0, we have to return the conditional
        expected value though, since x0 is subtracted later */
    Size d = model_->components(INF);
    for (Size i = 0; i < d; ++i) {
        res[model_->pIdx(INF, i, 0)] = x0[model_->pIdx(INF, i, 0)];
        res[model_->pIdx(INF, i, 1)] = x0[model_->pIdx(INF, i, 1)];
    }
    Size c = model_->components(CR);
    for (Size i = 0; i < c; ++i) {
        res[model_->pIdx(CR, i, 0)] = x0[model_->pIdx(CR, i, 0)];
        res[model_->pIdx(CR, i, 1)] = x0[model_->pIdx(CR, i, 1)];
    }
    return res;
}

Disposable<Matrix> CrossAssetStateProcess::ExactDiscretization::covarianceImpl(const StochasticProcess&, Time t0,
                                                                               const Array&, Time dt) const {
    Matrix res(model_->dimension(), model_->dimension());
    Size n = model_->components(IR);
    Size m = model_->components(FX);
    Size d = model_->components(INF);
    Size c = model_->components(CR);
    Size e = model_->components(EQ);
    // ir-ir
    for (Size i = 0; i < n; ++i) {
        for (Size j = 0; j <= i; ++j) {
            setValue(res, ir_ir_covariance(model_, i, j, t0, dt), model_, IR, i, IR, j, 0, 0);
        }
    }
    // ir-fx
    for (Size i = 0; i < n; ++i) {
        for (Size j = 0; j < m; ++j) {
            setValue(res, ir_fx_covariance(model_, i, j, t0, dt), model_, IR, i, FX, j, 0, 0);
        }
    }
    // fx-fx
    for (Size i = 0; i < m; ++i) {
        for (Size j = 0; j <= i; ++j) {
            setValue(res, fx_fx_covariance(model_, i, j, t0, dt), model_, FX, i, FX, j);
        }
    }
    // ir,fx,inf - inf
    for (Size j = 0; j < d; ++j) {
        for (Size i = 0; i <= j; ++i) {
            // infz-infz
            setValue(res, infz_infz_covariance(model_, i, j, t0, dt), model_, INF, i, INF, j, 0, 0);
            // infz-infy
            setValue(res, infz_infy_covariance(model_, i, j, t0, dt), model_, INF, i, INF, j, 0, 1);
            setValue(res, infz_infy_covariance(model_, j, i, t0, dt), model_, INF, i, INF, j, 1, 0);
            // infy-infy
            setValue(res, infy_infy_covariance(model_, i, j, t0, dt), model_, INF, i, INF, j, 1, 1);
        }
        for (Size i = 0; i < n; ++i) {
            // ir-inf
            setValue(res, ir_infz_covariance(model_, i, j, t0, dt), model_, IR, i, INF, j, 0, 0);
            setValue(res, ir_infy_covariance(model_, i, j, t0, dt), model_, IR, i, INF, j, 0, 1);
        }
        for (Size i = 0; i < (n - 1); ++i) {
            // fx-inf
            setValue(res, fx_infz_covariance(model_, i, j, t0, dt), model_, FX, i, INF, j, 0, 0);
            setValue(res, fx_infy_covariance(model_, i, j, t0, dt), model_, FX, i, INF, j, 0, 1);
        }
    }
    // ir,fx,inf,cr - cr
    for (Size j = 0; j < c; ++j) {
        for (Size i = 0; i <= j; ++i) {
            // crz-crz
            setValue(res, crz_crz_covariance(model_, i, j, t0, dt), model_, CR, i, CR, j, 0, 0);
            // crz-cry
            setValue(res, crz_cry_covariance(model_, i, j, t0, dt), model_, CR, i, CR, j, 0, 1);
            setValue(res, crz_cry_covariance(model_, j, i, t0, dt), model_, CR, i, CR, j, 1, 0);
            // cry-cry
            setValue(res, cry_cry_covariance(model_, i, j, t0, dt), model_, CR, i, CR, j, 1, 1);
        }
        for (Size i = 0; i < n; ++i) {
            // ir-cr
            setValue(res, ir_crz_covariance(model_, i, j, t0, dt), model_, IR, i, CR, j, 0, 0);
            setValue(res, ir_cry_covariance(model_, i, j, t0, dt), model_, IR, i, CR, j, 0, 1);
        }
        for (Size i = 0; i < (n - 1); ++i) {
            // fx-cr
            setValue(res, fx_crz_covariance(model_, i, j, t0, dt), model_, FX, i, CR, j, 0, 0);
            setValue(res, fx_cry_covariance(model_, i, j, t0, dt), model_, FX, i, CR, j, 0, 1);
        }
        for (Size i = 0; i < d; ++i) {
            // inf-cr
            setValue(res, infz_crz_covariance(model_, i, j, t0, dt), model_, INF, i, CR, j, 0, 0);
            setValue(res, infy_crz_covariance(model_, i, j, t0, dt), model_, INF, i, CR, j, 1, 0);
            setValue(res, infz_cry_covariance(model_, i, j, t0, dt), model_, INF, i, CR, j, 0, 1);
            setValue(res, infy_cry_covariance(model_, i, j, t0, dt), model_, INF, i, CR, j, 1, 1);
        }
    }
    // ir,fx,inf,cr,eq - eq
    for (Size j = 0; j < e; ++j) {
        for (Size i = 0; i <= j; ++i) {
            // eq-eq
            setValue(res, eq_eq_covariance(model_, i, j, t0, dt), model_, EQ, i, EQ, j, 0, 0);
        }
        for (Size i = 0; i < n; ++i) {
            // ir-eq
            setValue(res, ir_eq_covariance(model_, i, j, t0, dt), model_, IR, i, EQ, j, 0, 0);
        }
        for (Size i = 0; i < (n - 1); ++i) {
            // fx-eq
            setValue(res, fx_eq_covariance(model_, i, j, t0, dt), model_, FX, i, EQ, j, 0, 0);
        }
        for (Size i = 0; i < d; ++i) {
            // inf-eq
            setValue(res, infz_eq_covariance(model_, i, j, t0, dt), model_, INF, i, EQ, j, 0, 0);
            setValue(res, infy_eq_covariance(model_, i, j, t0, dt), model_, INF, i, EQ, j, 1, 0);
        }
        for (Size i = 0; i < c; ++i) {
            // cr-eq
            setValue(res, crz_eq_covariance(model_, i, j, t0, dt), model_, CR, i, EQ, j, 0, 0);
            setValue(res, cry_eq_covariance(model_, i, j, t0, dt), model_, CR, i, EQ, j, 1, 0);
        }
    }
    return res;
}

void CrossAssetStateProcess::ExactDiscretization::flushCache() const {
    cache_m_.clear();
    cache_v_.clear();
    cache_d_.clear();
}

} // namespace QuantExt


#ifndef PYTHONIC_SCIPY_SPECIAL_KV_HPP
#define PYTHONIC_SCIPY_SPECIAL_KV_HPP

#include "pythonic/include/scipy/special/kv.hpp"

#include "pythonic/types/ndarray.hpp"
#include "pythonic/utils/functor.hpp"
#include "pythonic/utils/numpy_traits.hpp"

#define BOOST_MATH_THREAD_LOCAL thread_local
#include <boost/math/special_functions/bessel.hpp>

PYTHONIC_NS_BEGIN

namespace scipy
{
  namespace special
  {
    namespace details
    {
      template <class T0, class T1>
      double kv(T0 x, T1 y)
      {
        using namespace boost::math::policies;
        return boost::math::cyl_bessel_k(x, y,
                                         make_policy(promote_double<true>()));
      }
    }

#define NUMPY_NARY_FUNC_NAME kv
#define NUMPY_NARY_FUNC_SYM details::kv
#include "pythonic/types/numpy_nary_expr.hpp"
  }
}
PYTHONIC_NS_END

#endif


#include <cassert>
#include <iostream>
#include <fstream>
#include <sstream>
#include <cmath>
#include <iomanip>
#include <armadillo>

#include "constant.hpp"
#include "mt_random.hpp"

using namespace std;

namespace willow {

// global variables
const  int m_natom = 2;
const  int m_nnhc  = 4;
static int m_nbead;
static int m_nstep;
static int m_irst;
static int m_nref;
static double m_dt;
static double m_dt_ref;
static double m_gfree_nm;

static double m_ZPE;
static double m_omega;
static double m_omega_p2;
static double m_temp_nm;
static double m_beta_nm;
static double m_ekin_nm;

const  double delt_x = 4.e-3;

static arma::vec prob_dx;      // (0:399)
static arma::vec prob_dx_pos1; // (0:399)
static arma::vec prob_dx_pos2; // (0:399)

static arma::vec2 m_mass;     // (m_natom)
static double     m_rmass;    // reduced mass

static arma::mat  m_tmat_nm;  // (m_nbead, m_nbead)
static arma::mat  m_fict_mass;// (m_natom, m_nbead)

static arma::vec  m_rbath_nm; // (m_nnhc)
static arma::vec  m_vbath_nm; //
static arma::vec  m_qmass_nm; //

// one-dimensional system
// Cartesian Coordinate
static arma::mat  m_pos_qm;   // (m_natom, m_nbead)
static arma::mat  m_grd_qm;   // (m_natom, m_nbead)

// Normal Mode
static arma::mat  m_pos_nm;   // (m_natom, m_nbead)
static arma::mat  m_vel_nm;   // (m_natom, m_nbead)
static arma::mat  m_grd_nm;   // (m_natom, m_nbead)
static arma::mat  m_grd_nm_spr; // (m_natom, m_nbead)

static void sample_rho ()
{

  // QM (beads)

  for (auto ib = 0; ib < m_nbead; ++ib) {

    // (1) \rho(x) = \rho (x2-x1)
    double dx = m_pos_qm(1,ib) - m_pos_qm(0,ib);
    int  id1  = (int) round(dx/delt_x) + 200;

    if (id1 >= 0 && id1 < 400) prob_dx(id1) += 1;

    //
    // (2) distribution of beads around pos(0) and pos(1)
    // \varrho(x) = \varrho (x1 - X1)
    dx   = m_pos_qm(0,ib);
    id1  = (int) round(dx/delt_x) + 200;
    if (id1 >= 0 && id1 < 400) prob_dx_pos1(id1) += 1;

    //
    dx   = m_pos_qm(1,ib);
    id1  = (int) round(dx/delt_x) + 200;
    if (id1 >= 0 && id1 < 400) prob_dx_pos2(id1) += 1;

  }

}

void nm_nhc_integrate ()
{
  // Nose-Hoover Chain Method

  const double dt_ref  = m_dt_ref;
  const double dt_ref2 = 0.5*dt_ref;
  const double dt_ref4 = 0.5*dt_ref2;
  const double dt_ref8 = 0.5*dt_ref4;

  // m_nnhc = 4
  arma::vec4 rbath;
  arma::vec4 vbath;
  arma::vec4 fbath;
  arma::vec4 qmass;

  double ekin_nm = 0.0;

  for (size_t ib = 1; ib < m_nbead; ++ib) {

    for (size_t ia = 0; ia < m_natom; ++ia) {

      double mass = m_fict_mass (ia, ib);
      double v    = m_vel_nm(ia,ib);

      ekin_nm += mass * v * v;

    }

  }

  //---

  vbath = m_vbath_nm;
  rbath = m_rbath_nm;
  qmass = m_qmass_nm;

  // ihc = 0
  fbath(0) = (ekin_nm - m_gfree_nm*m_ekin_nm)/qmass(0);

  for (size_t ihc = 1; ihc < m_nnhc; ++ihc) {
    fbath(ihc) =
      (qmass(ihc-1)*vbath(ihc-1)*vbath(ihc-1) - m_ekin_nm) / qmass(ihc);
  }

  // Update Thermostat Velocities

  vbath(m_nnhc-1) = vbath(m_nnhc-1) + fbath(m_nnhc-1)*dt_ref4;

  for (auto ihc = 1; ihc < m_nnhc; ++ihc) {
    const auto jhc     = m_nnhc - ihc;
    const double vfact = exp (-vbath(jhc)*dt_ref8);
    const double vtmp  = vbath(jhc-1);
    vbath(jhc-1) = vtmp*vfact*vfact + fbath(jhc-1)*vfact*dt_ref4;
  }

  // Update atomic velocities
  const double pvfact = exp(-vbath(0)*dt_ref2);

  ekin_nm  = ekin_nm*pvfact*pvfact;

  // Update thermostat forces

  fbath(0) = (ekin_nm - m_gfree_nm*m_ekin_nm)/qmass(0);

  for (auto ihc = 0; ihc < m_nnhc; ++ihc) {
    rbath(ihc) += vbath(ihc)*dt_ref2;
  }

  // Update Thermostat velocities

  for (auto ihc = 0; ihc < m_nnhc-1; ++ihc) {
    const double vfact = exp(-vbath(ihc+1)*dt_ref8);
    const double vtmp  = vbath(ihc);

    vbath(ihc) = vtmp*vfact*vfact + fbath(ihc)*vfact*dt_ref4;
    fbath(ihc+1) =
      (qmass(ihc)*vbath(ihc)*vbath(ihc) - m_ekin_nm)/qmass(ihc+1);
  }

  vbath(m_nnhc-1) += fbath(m_nnhc-1)*dt_ref4;

  // backup
  m_vbath_nm = vbath;
  m_rbath_nm = rbath;

  // update velocities of normal modes
  for (auto ib = 1; ib < m_nbead; ++ib)
    for (auto ia = 0; ia < m_natom; ++ia) {
      m_vel_nm(ia,ib) *= pvfact;
    } // ia

}

void nm_pos_update ()
{

  m_pos_nm += m_dt_ref * m_vel_nm;

}

void nm_grad_spring ()
{

  //
  // centroid gradient is zero
  //
  m_grd_nm_spr.col(0).zeros();

  //
  // Gradients from the Springs between 'neighboring' beads
  //

  for (size_t ib = 1; ib < m_nbead; ++ib) {

    for (size_t ia = 0; ia < m_natom; ++ia) {

      const double fact = m_fict_mass(ia,ib)*m_omega_p2;

      m_grd_nm_spr(ia,ib) = fact * m_pos_nm(ia,ib);
    }
  }

}

void nm_pos_trans ()
{
  //
  // normal mode (nm) ---> Cartesian (qm)
  // pos_nm --->  pos_qm

  for (auto ib = 0; ib < m_nbead; ++ib) {

    arma::vec pos_x (m_natom, arma::fill::zeros);

    for (auto jb = 0; jb < m_nbead; ++jb) {
      pos_x += m_tmat_nm(ib,jb)*m_pos_nm.col(jb);
    }

    m_pos_qm.col(ib) = pos_x;
  }

}

void nm_grad_trans ()
{

  m_grd_nm.zeros();

  for (size_t ib = 0; ib < m_nbead; ++ib) {
    for (size_t jb = 0; jb < m_nbead; ++jb) {
      m_grd_nm.col(ib) += m_tmat_nm(jb,ib)*m_grd_qm.col(jb);
    }
  }

}

void nm_pos_init ()
{

  {// centroid particle
    m_pos_nm.col(0).zeros();
  }

  { // beads : normal mode
    const double dbead = m_nbead;
    const double usigma = 0.02*ang2bohr; // sigma_x = 0.02 A

    for (size_t ib = 1; ib < m_nbead; ++ib) {

      arma::vec2 pos_nm_ib;
      for (size_t ia = 0; ia < m_natom; ++ia) {

	double mass05 = sqrt(m_fict_mass(ia,ib));
	pos_nm_ib(ia) = usigma*rnd::gaus_dev()/mass05;
      }

      m_pos_nm.col(ib) = pos_nm_ib;

    }
  } // beads

}

void nm_vel_init ()
{

  // ---- centroid ----
  {
    // Here, vel_nm(ib = 0) zero
    m_vel_nm.col(0).zeros();
  }

  { // velocities for bead particles (or nm-mode particles)

    for (size_t ib = 1; ib < m_nbead; ++ib) {

      double sum_p    = 0.0;
      double sum_mass = 0.0;

      for (size_t ia = 0; ia < m_natom; ++ia) {

	double mass   = m_fict_mass(ia,ib);
	double vsigma = sqrt (m_ekin_nm/mass);

	double v = vsigma*rnd::gaus_dev();

	m_vel_nm(ia,ib) = v;

	sum_p    += mass*v;
	sum_mass += mass;
      }

      sum_p /= sum_mass;

      // translational motion is zero
      for (auto ia = 0; ia < m_natom; ++ia) {
	m_vel_nm(ia,ib) -= sum_p;
      }

    }

    //
    // one-dimensional diatomic molecule
    // does not have the rotational motion.
    //

    // Scale Velocity
    double ekin_nm = 0.0;

    for (size_t ib = 1; ib < m_nbead; ++ib) {

      for (size_t ia = 0; ia < m_natom; ++ia) {

	double mass  = m_fict_mass(ia,ib);
	double v     = m_vel_nm(ia,ib);
	ekin_nm += mass * v*v;

      }

    }

    double temp  = ekin_nm / (m_gfree_nm*boltz);
    double scale = sqrt(m_temp_nm / temp);

    for (size_t ib = 1; ib < m_nbead; ++ib) {

      for (size_t ia = 0; ia < m_natom; ++ia) {

	m_vel_nm(ia,ib) *= scale;

      }

    } // ib

  } // velocities for bead particles

}

void nm_vel_update ()
{

  // (ib = 0) belongs to the centroid velocity

  //---
  const double dt2 = 0.5*m_dt;

  for (size_t ib = 1; ib < m_nbead; ++ib) {
    for (size_t ia = 0; ia < m_natom; ++ia) {
      double mass   = m_fict_mass(ia,ib);
      double factor = dt2/mass;
      m_vel_nm(ia,ib) -= factor*m_grd_nm(ia,ib);
    }
  }

}

void nm_vel_spring_update ()
{

  double dt2 = 0.5*m_dt_ref;

  for (size_t ib = 1; ib < m_nbead; ++ib) {

    for (size_t ia = 0; ia < m_natom; ++ia) {
      double mass   = m_fict_mass(ia,ib);
      double factor = dt2/mass;
      m_vel_nm(ia,ib) -=
	factor*m_grd_nm_spr(ia,ib);
    }

  }

}

double nm_pot_grad ()
{

  // beads: normal mode ---> cartesian
  //        pos_nm ----> pos_qm
  nm_pos_trans ();

  double u_vib = 0.0;
  m_grd_qm.zeros();

  //
  // force constant: reduced_mass*omega*omega
  //

  double k_val = m_rmass*m_omega*m_omega;

  // Harmonic Oscilltor: Diatomic Molecule
  // U = 0.5 * k * (x2 - x1)**2
  //
  for (size_t ib = 0; ib < m_nbead; ib++) {

    // call your potential.
    double dx = m_pos_qm(1,ib) - m_pos_qm(0,ib);
    double en_harm = 0.5*k_val*dx*dx;

    m_grd_qm(1,ib) =  k_val*dx;
    m_grd_qm(0,ib) = -k_val*dx;

    u_vib += en_harm;
  } // ib

  // ---
  double d_nbead = m_nbead;
  u_vib /= d_nbead;

  m_grd_qm /= d_nbead;

  //
  // cartesian gradient ---> normal mode gradient
  //
  nm_grad_trans ();

  return u_vib;

}

void nm_report (const int& istep,
		const double& u_vib,
		double& E_eff)
{
  // kinetic energy for beads
  double ekin_nm = 0.0;

  for (auto ib = 1; ib < m_nbead; ++ib) {
    for (auto ia = 0; ia < m_natom; ++ia) {
      double mass = m_fict_mass(ia,ib);
      double v    = m_vel_nm(ia,ib);
      ekin_nm += mass*v*v;
    }
  }

  ekin_nm = 0.5*ekin_nm;
  double temp_nm = 2.0*ekin_nm/(m_gfree_nm*boltz);

  // Harmonic Potential of springs between neighboring beads

  double qkin_nm = 0.0;
  for (auto ib = 1; ib < m_nbead; ++ib)
    for (auto ia = 0; ia < m_natom; ++ia) {
      double fact  = 0.5*m_fict_mass(ia,ib)*m_omega_p2;
      double q     = m_pos_nm(ia,ib);
      qkin_nm += fact*q*q;
    }

  double ebath_nm = 0.0;

  arma::vec4 qmass = m_qmass_nm;
  ebath_nm += 0.5*qmass(0)*m_vbath_nm(0)*m_vbath_nm(0);
  ebath_nm += m_gfree_nm*m_ekin_nm*m_rbath_nm(0);

  for (auto ihc = 1; ihc < m_nnhc; ++ihc) {
    ebath_nm += 0.5*qmass(ihc)*m_vbath_nm(ihc)*m_vbath_nm(ihc);
    ebath_nm += m_ekin_nm*m_rbath_nm(ihc);
  }

  E_eff = qkin_nm + u_vib; // <E_eff> = E_ZPE
  double H_sys = ekin_nm + E_eff; // Hamiltonian of the system
  double H_tot = H_sys + ebath_nm; // Total H.

  // unit convert
  E_eff *= au_kcal;
  H_sys *= au_kcal;
  H_tot *= au_kcal;

  printf (" %8d %14.6f %14.6f %14.6f %14.6f %10.2f \n",
	  (istep+1), H_tot, H_sys, E_eff,
	  u_vib*au_kcal, temp_nm);

  fflush (stdout);

}

void read_restart_nm (int& istep0)
{

  // read a restart file
  std::string str_rst;

  std::ifstream ifs_rst ("pimdrr.sav");
  assert (ifs_rst.good());

  std::ostringstream oss;

  oss << ifs_rst.rdbuf();

  str_rst = oss.str();

  std::istringstream is (str_rst);

  std::string line;
  std::getline (is, line); // istep
  std::istringstream istep_ss (line);
  istep_ss >> istep0;

  // ib = 0 is for centroids
  for (auto ib = 1; ib < m_nbead; ++ib) {
    std::getline (is, line);
    std::istringstream iss_pos (line);

    iss_pos >> m_pos_nm(0,ib) >> m_pos_nm(1,ib);

    std::getline (is, line);
    std::istringstream iss_vel (line);

    iss_vel >> m_vel_nm(0,ib) >> m_vel_nm(1,ib);
  }

  for (auto ihc = 0; ihc < m_nnhc; ++ihc) {
    std::getline (is,line);
    std::istringstream iss (line);

    iss >> m_rbath_nm(ihc) >> m_vbath_nm(ihc);
  }

  m_pos_nm(0,0) = 0.0;
  m_pos_nm(1,0) = 0.0;
  m_vel_nm(0,0) = 0.0;
  m_vel_nm(1,0) = 0.0;

}

void write_restart_nm (const int& istep)
{

  FILE *ofs_rst = fopen ("pimdrr.sav", "w");

  fprintf( ofs_rst, " %10d \n", istep+1);

  // ib = 0 is for centroids
  for (auto ib = 1; ib < m_nbead; ++ib) {
    fprintf (ofs_rst, "  %E   %E   \n", m_pos_nm(0,ib), m_pos_nm(1,ib));
    fprintf (ofs_rst, "  %E   %E   \n", m_vel_nm(0,ib), m_vel_nm(1,ib));
  }

  for (auto ihc = 0; ihc < m_nnhc; ++ihc) {
    fprintf (ofs_rst, "  %E   %E   %E \n",
	     m_rbath_nm(ihc),
	     m_vbath_nm(ihc),
	     m_qmass_nm(ihc) );
  }

  fclose (ofs_rst);

}

void write_prob_bin (int& nsamp)
{

  double zeta2   = m_rmass*m_omega; // (unit 1/bohr**2)

  arma::vec p_dx = prob_dx/(m_nbead*(nsamp+1)*delt_x);

  std::ofstream ofs_qm ("prob_bin_qm.dat");
  // psi_0
  // zeta2 = mass*omega*(x*x)
  // debug
  const double zt    = sqrt(zeta2/M_PI);

  double ave = 0.0;

  for (auto ib = 0; ib < 400; ++ib) {
    double  x  = (ib - 200)*delt_x;
    double  x2 =  x* x;
    double rho =  zt*exp(-zeta2*x2);
    ave += x*p_dx(ib);
    ofs_qm << x << "   "  << p_dx(ib)
	   << "  " << rho << endl;
  }

  ofs_qm.close();

  //
  std::ofstream ofs_rho ("prob_bin_rho.dat");
  arma::vec p_dx1  = prob_dx_pos1/(m_nbead*(nsamp+1)*delt_x);
  arma::vec p_dx2  = prob_dx_pos2/(m_nbead*(nsamp+1)*delt_x);

  for (auto ib = 0; ib < 400; ++ib) {
    double x = (ib - 200)*delt_x;
    ofs_rho << x << "   "  << p_dx1(ib) << "  " << p_dx2(ib) << endl;
  }

  ofs_rho.close();
}

void wpimd_run ()
{

  //
  // This is a sample code,
  // in which WPIMD is running at T (thermal temperature) = 0 K.
  //

  int istep0 = 0;

  if (m_irst == 1) {
    read_restart_nm (istep0);
  }

  // initial gradients and potential energy
  double u_vib = nm_pot_grad ();

  double e_eff = 0.0;
  double sum_e_eff = 0.0;
  int    ncount  = 0;

  nm_grad_spring ();

  if (m_irst == 0)
    nm_report (istep0-1, u_vib, e_eff);

  for (auto istep = 0; istep < m_nstep; ++istep) {

    nm_vel_update ();

    for (auto iref = 0; iref < m_nref; ++iref) {
      nm_nhc_integrate ();
      nm_vel_spring_update();
      nm_pos_update ();
      nm_grad_spring();
      nm_vel_spring_update();
      nm_nhc_integrate ();
    }

    u_vib   = nm_pot_grad ();

    sample_rho ();
    nm_vel_update ();

    if ( (istep+1)%500 == 0) {
      nm_report (istep+istep0, u_vib, e_eff);
      sum_e_eff += e_eff;
      ncount++;
    }

    if ( (istep+1)%5000 == 0) {
      write_restart_nm   (istep+istep0);
      write_prob_bin (istep);
    }

  }

  cout << "AVE E_eff " << sum_e_eff / ncount << endl;

}

void wpimd_init ()
{

  // one-dimensional system
  m_gfree_nm = 1*m_nbead; // the vibrational degree of freedom

  // ZPE = 0.5 * hbar * omega
  // ZPE(au) = 0.5 * omega
  m_omega    = 2.0*m_ZPE;

  // Eq. (14)
  // temperature for the bead motions
  m_temp_nm  = m_omega/(m_gfree_nm*boltz);

  double dbead = m_nbead;
  double omega_p = sqrt(dbead)*boltz*m_temp_nm;
  m_omega_p2 = omega_p * omega_p;
  m_beta_nm  = 1.0 / (boltz*m_temp_nm);
  m_ekin_nm  = boltz*m_temp_nm;

  // mem alloc

  m_tmat_nm   = arma::mat (m_nbead, m_nbead, arma::fill::zeros);
  m_fict_mass = arma::mat (m_natom, m_nbead, arma::fill::zeros);

  m_rbath_nm  = arma::vec (m_nnhc, arma::fill::zeros);
  m_vbath_nm  = arma::vec (m_nnhc, arma::fill::zeros);
  m_qmass_nm  = arma::vec (m_nnhc, arma::fill::zeros);

  // one-dimensional system
  m_pos_qm    = arma::mat (m_natom, m_nbead, arma::fill::zeros);
  m_pos_nm    = arma::mat (m_natom, m_nbead, arma::fill::zeros);

  m_vel_nm    = arma::mat (m_natom, m_nbead, arma::fill::zeros);

  m_grd_qm    = arma::mat (m_natom, m_nbead, arma::fill::zeros);
  m_grd_nm    = arma::mat (m_natom, m_nbead, arma::fill::zeros);
  m_grd_nm_spr= arma::mat (m_natom, m_nbead, arma::fill::zeros);

  // --- initiate the normal mode matrix.
  for (size_t i = 0; i < m_nbead; ++i) {
    m_tmat_nm(i, 0) = 1.0;
  }

  for (size_t i = 0; i < m_nbead/2; ++i) {
    m_tmat_nm(2*i,   m_nbead-1) = -1.0;
    m_tmat_nm(2*i+1, m_nbead-1) =  1.0;
  }

  double dnorm = sqrt (2.0);

  for (size_t i = 0; i < m_nbead; ++i) {
    const double di    = i+1;
    const double phase = 2.0*di*(M_PI/dbead);
    for (size_t j = 0; j < (m_nbead-2)/2; ++j) {
      const double dj    = j+1;
      m_tmat_nm(i, 2*j+1) = dnorm*cos(phase*dj);
      m_tmat_nm(i, 2*j+2) = dnorm*sin(phase*dj);
    }
  }

  // --- mass init ---
  for (auto ia = 0; ia < m_natom; ++ia) {
    double mass = m_mass(ia);

    m_fict_mass(ia,0)         = mass;
    m_fict_mass(ia,m_nbead-1) = 4.0*dbead*mass;

    for (auto ib = 1; ib < m_nbead/2; ++ib) {
      double val = 2.0*(1.0 - cos (2.0*ib*(M_PI/dbead)))*dbead*mass;
      m_fict_mass(ia,2*ib-1) = val;
      m_fict_mass(ia,2*ib  ) = val;
    }
  }

  // reduced mass
  double rmass = 0.0;

  for (auto ia = 0; ia < m_natom; ++ia) {
    rmass += 1.0/m_mass(ia);
  }

  m_rmass = 1.0/rmass;

  // bath init for beads

  m_qmass_nm(0) = m_gfree_nm*m_ekin_nm/m_omega_p2;

  for (size_t ihc = 1; ihc < m_nnhc; ++ihc) {
    m_qmass_nm(ihc) = m_ekin_nm/m_omega_p2;
  }

  nm_pos_init ();
  nm_vel_init ();

  prob_dx      = arma::vec(400, arma::fill::zeros);
  prob_dx_pos1 = arma::vec(400, arma::fill::zeros);
  prob_dx_pos2 = arma::vec(400, arma::fill::zeros);

}

void read_input (const std::string& fname)
{

  std::ifstream is_input(fname);
  std::ostringstream oss;
  oss << is_input.rdbuf();

  std::istringstream ss (oss.str());

  m_nstep = 1000;
  m_dt   = 0.5; // fsec
  m_irst = 0;

  m_nbead = 8;
  m_nref  = 10;
  m_ZPE   = 0.0; // kcal/mol

  for (auto ia = 0; ia < m_natom; ++ia)
    m_mass(ia)  = 1.0*amu2au; // amu --> au

  std::string line;

  while(getline(ss, line)) {
    std::istringstream iss (line);
    std::string keyword;
    std::string val;
    iss >> keyword >> val;

    if (keyword == "nstep") {
      m_nstep = stoi (val);
    }
    else if (keyword == "dt") {
      m_dt = stod(val);
    }
    else if (keyword == "l_restart") {
      m_irst = stoi(val);
    }
    else if (keyword == "nbead") {
      m_nbead = stoi(val);
    }
    else if (keyword == "nref") {
      m_nref  = stoi(val);
    }
    else if (keyword == "ZPE") {
      m_ZPE  = stod(val)/au_kcal; // kcal/mol ---> au
    }
    else if (keyword == "mass1") {
      m_mass(0) = stod(val)*amu2au; // amu --> au
    }
    else if (keyword == "mass2") {
      m_mass(1) = stod(val)*amu2au; // amu --> au
    }

  }

  // -- time : [fs] ---> [au]
  m_dt  = m_dt * (1.0e-15/au_time);

  m_dt_ref = m_dt/m_nref;

}

} // namespace willow

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

  cout << std::setprecision (6);
  cout << std::fixed;

  //--- read an input file --
  const std::string fname = (argc > 1) ? argv[1] : "sample2.inp";

  willow::read_input (fname);

  willow::wpimd_init ();

  willow::wpimd_run ();

  return 0;

}


#include <boost/assert.hpp>
#include <cmath>
#include <limits>
#include <utility>

#include "mh_moves/mh_moves.hpp"
#include "mh_moves/mh_move_weights.hpp"
#include "sampling/simple.hpp"
#include "mh_moves/utility.hpp"

namespace biggles {

namespace mh_moves
{

namespace death_fun {
    float get_birth_obs_prob(const clutter_ptr& clutter, time_stamp first_time_stamp, time_stamp last_time_stamp,
        const shared_const_track_ptr& track_to_remove)
    {
        float noobs_weight(0.1f);
        float sample_weight(0.f);
        float total_weight(0.f);

        float new_obs_prob(0.f);
        observation last_obs;
        const observation_collection &track_to_remove_obs(track_to_remove->observations());
        observation_collection::const_iterator obs_it = track_to_remove_obs.begin();
        for (time_stamp new_obs_t = first_time_stamp; new_obs_t < last_time_stamp; ++new_obs_t)
        {

            std::deque<observation> candidates;
            if (noobs(last_obs)) {
                candidates = clutter->observations(new_obs_t);
                noobs_weight = 0.1f;
                total_weight = float(candidates.size()) + noobs_weight;
                if (new_obs_t == t(*obs_it)) {
                    last_obs = *obs_it++;
                    sample_weight = 1.f;
                } else {
                    sample_weight = noobs_weight;
                }
            } else {
                float radius = float(new_obs_t - t(last_obs)) * detail::speed_of_light_;
                candidates = clutter->locate_near(new_obs_t, last_obs, radius);
                std_normal_weight obs_weight(last_obs);
                noobs_weight = obs_weight.sigmas(2, new_obs_t);
                total_weight = fun_sum(candidates.begin(), candidates.end(), obs_weight, noobs_weight);
                if (obs_it != track_to_remove_obs.end() and new_obs_t == t(*obs_it)) {
                    last_obs = *obs_it++;
                    sample_weight = obs_weight(last_obs);
                } else {
                    sample_weight = noobs_weight;
                }
            }

            new_obs_prob += logf(sample_weight) - logf(total_weight);
        }

        return new_obs_prob;
    }
}

bool death(const partition_ptr_t& start_partition_ptr, partition_ptr_t& end_partition_ptr, float& proposal_density_ratio)
{
    // copy the start partition so that start_partition can equal end_partition and we can modify end_partition
    //partition original(start_partition);
    const clutter_t& clutter(start_partition_ptr->clutter());
    const track_collection& tracks(start_partition_ptr->tracks());

    capability_recorder_ptr cap_rec_ptr = start_partition_ptr->get_capability_recorder();
    BOOST_ASSERT(cap_rec_ptr.get() not_eq 0); // there actually is something assigned

    // early-out: there must be at least one track
    if(tracks.empty())
        return false;

    // choose a track to remove
    size_t n_to_sample = sampling::uniform_int(0, tracks.size());
    track_collection::const_iterator it(tracks.begin()); std::advance(it, n_to_sample);
    shared_const_track_ptr track_to_remove(*it);

    // create copies of the partition clutter and tracks
    clutter_ptr new_clutter(new clutter_t(clutter));
    boost::shared_ptr<track_collection> new_tracks(new track_collection(tracks));

    // put track obs into clutter
    BOOST_FOREACH(const observation& o, *track_to_remove)
    {
        new_clutter->insert(o);
    }

    cap_rec_ptr->add_erase_track(track_to_remove);
    cap_rec_ptr->set_editing_finished();

    // remove old track
    new_tracks->remove(track_to_remove);

    // update partition
    end_partition_ptr = partition_ptr_t(
        new partition(start_partition_ptr->pool(), new_tracks, new_clutter, start_partition_ptr->expansion())
    );
    end_partition_ptr->set_first_time_stamp(start_partition_ptr->first_time_stamp());
    end_partition_ptr->set_last_time_stamp(start_partition_ptr->last_time_stamp());

    // calculating the probs to sample the first and the last time stamp
    // see birth move for more details
    time_stamp last_clutter_time_stamp = new_clutter->last_time_stamp();
    time_stamp last_for_first_time_stamp = last_clutter_time_stamp - 1;
    time_stamp first_time_stamp = track_to_remove->first_time_stamp();
    float first_time_stamp_prob = -logf(last_for_first_time_stamp - end_partition_ptr->first_time_stamp());
    BOOST_ASSERT(std::isfinite(first_time_stamp_prob));
    time_stamp last_time_stamp = track_to_remove->last_time_stamp();
    float last_time_stamp_prob = -logf(end_partition_ptr->last_time_stamp() - first_time_stamp - 1);
    BOOST_ASSERT(std::isfinite(last_time_stamp_prob));

    float new_obs_prob =
        death_fun::get_birth_obs_prob(new_clutter, first_time_stamp, last_time_stamp, track_to_remove);

    // death prob == prob selecting correct track
    float death_log_prob = -logf(start_partition_ptr->tracks().size());

    // birth prob == prob of sampling correct observation
    float birth_log_prob = first_time_stamp_prob + last_time_stamp_prob + new_obs_prob;

    proposal_density_ratio = birth_log_prob - death_log_prob;
    return true;
}

}

}


<% render(instance) %>
<%= includes((0..x).size) %>

#include <boost/hana/integral.hpp>
#include <boost/hana/searchable/searchable.hpp>

int main() {
    auto go = boost::hana::any(
        <%= list(
            (0..x).map { |i| "decltype(boost::hana::int_<#{i}>)" },
            (0..x).map { |i| "boost::hana::int_<#{i}>" }
        ) %>,
        [](auto x) {
            return x == boost::hana::int_<  <%= x / 2 %>    >;
        }
    );
}


#ifndef IMAGE_REPROJECTION_PLUGINS_MESH_SURFACE_MODEL_HPP
#define IMAGE_REPROJECTION_PLUGINS_MESH_SURFACE_MODEL_HPP

#include <string>
#include <vector>

#include <image_reprojection/surface_model.hpp>
#include <image_reprojection_plugins/MeshStamped.h>
#include <topic_tools/shape_shifter.h>

#include <boost/foreach.hpp>
#include <boost/thread/locks.hpp>
#include <boost/thread/shared_mutex.hpp>

#include <opencv2/core/core.hpp>

namespace image_reprojection_plugins {

class MeshSurfaceModel : public image_reprojection::SurfaceModel {
public:
  MeshSurfaceModel() {}

  virtual ~MeshSurfaceModel() {}

  virtual void update(const topic_tools::ShapeShifter &surface) {
    const MeshStampedConstPtr mesh(surface.instantiate< MeshStamped >());
    CV_Assert(mesh);
    update(*mesh);
  }

  void update(const MeshStamped &mesh) {
    boost::unique_lock< boost::shared_mutex > write_lock(mutex_);

    frame_id_ = mesh.header.frame_id;

    vertices_.clear();
    BOOST_FOREACH (const geometry_msgs::Point &v, mesh.mesh.vertices) {
      vertices_.push_back(cv::Vec3f(v.x, v.y, v.z));
    }

    triangles_.clear();
    BOOST_FOREACH (const shape_msgs::MeshTriangle &t, mesh.mesh.triangles) {
      triangles_.push_back(
          cv::Vec3i(t.vertex_indices[0], t.vertex_indices[1], t.vertex_indices[2]));
    }
  }

  virtual std::string getFrameId() const {
    boost::shared_lock< boost::shared_mutex > read_lock(mutex_);
    return frame_id_;
  }

private:
  virtual void onInit() {}

  virtual void onIntersection(const cv::Vec3f &src_origin, const cv::Mat &src_direction,
                              cv::Mat &dst, cv::Mat &mask) const {
    boost::shared_lock< boost::shared_mutex > read_lock(mutex_);
    multirayMeshIntersection(src_origin, src_direction, dst, mask);
  }

  void multirayMeshIntersection(const cv::Vec3f &src_origin, const cv::Mat &src_direction,
                                cv::Mat &dst, cv::Mat &mask) const {
    dst.create(src_direction.size(), CV_32FC3);
    for (int x = 0; x < src_direction.size().width; ++x) {
      for (int y = 0; y < src_direction.size().height; ++y) {
        unsigned char &m(mask.at< unsigned char >(y, x));
        const cv::Vec3f &sd(src_direction.at< cv::Vec3f >(y, x));
        cv::Vec3f &d(dst.at< cv::Vec3f >(y, x));
        m = (m != 0 && rayMeshIntersection(src_origin, sd, d)) ? 1 : 0;
      }
    }
  }

  bool rayMeshIntersection(const cv::Vec3f &src_origin, const cv::Vec3f &src_direction,
                           cv::Vec3f &dst) const {
    bool intersects(false);
    double min_distance;
    BOOST_FOREACH (const cv::Vec3i &triangle, triangles_) {
      cv::Vec3f this_intersection;
      if (rayTriangleIntersection(src_origin, src_direction, triangle, this_intersection)) {
        const double this_distance(cv::norm(src_origin, this_intersection));
        if (!intersects || (intersects && this_distance < min_distance)) {
          intersects = true;
          min_distance = this_distance;
          dst = this_intersection;
        }
      }
    }

    return intersects;
  }

  bool rayTriangleIntersection(const cv::Vec3f &src_origin, const cv::Vec3f &src_direction,
                               const cv::Vec3i &triangle, cv::Vec3f &dst) const {
    // primitive vectors with respect to 0th vertex of triangle
    const cv::Vec3f r(src_origin - vertices_[triangle[0]]);              // position of ray origin
    const cv::Vec3f e1(vertices_[triangle[1]] - vertices_[triangle[0]]); // 1st edge of triangle
    const cv::Vec3f e2(vertices_[triangle[2]] - vertices_[triangle[0]]); // 2nd edge of triangle
    const cv::Vec3f n(e1.cross(e2));                                     // normal of triangle

    // intersection point (p) can be described as
    //    p = r + t * d = b1 * e1 + b2 * e2
    //    (t >= 0, b1 >= 0, b2 >= 0, b1 + b2 <= 1)
    // where d is ray direction

    // calculate dot(d, n) for convenience
    const double ddn(src_direction.dot(n));
    if (ddn == 0.) { // means ray and triangle are parallel
      return false;
    }
    const double ddn_sign(ddn > 0. ? 1. : -1.);
    const double ddn_norm(ddn > 0. ? ddn : -ddn);

    // t = -dot(r, n) / dot(d, n)
    const double srdn(ddn_sign * r.dot(n));
    if (srdn > 0.) { // means t < 0
      return false;
    }

    // b1 = dot(d, cross(r, e2)) / dot(d, n)
    const double sddrce2(ddn_sign * src_direction.dot(r.cross(e2)));
    if (sddrce2 < 0. || sddrce2 > ddn_norm) { // means b1 < 0 or b1 > 1
      return false;
    }

    // b2 = dot(d, cross(e1, r)) / dot(d, n)
    const double sdde1cr(ddn_sign * src_direction.dot(e1.cross(r)));
    if (sdde1cr < 0. || sdde1cr > ddn_norm) { // means b2 < 0 or b2 > 1
      return false;
    }

    if (sddrce2 + sdde1cr > ddn_norm) { // means b1 + b2 > 1
      return false;
    }

    // intersection point with respect to mesh coordinate frame
    dst = src_origin - srdn * src_direction / ddn_norm;
    return true;
  }

private:
  mutable boost::shared_mutex mutex_;
  std::string frame_id_;
  std::vector< cv::Vec3f > vertices_;
  std::vector< cv::Vec3i > triangles_;
};

} // namespace image_reprojection_plugins

#endif


/**
 * \file libsanm/sparse_solver.cpp
 * This file is part of SANM, a symbolic asymptotic numerical solver.
 */

#include "libsanm/sparse_solver.h"

#include <mkl_pardiso.h>
#include <mkl_service.h>
#include <mkl_spblas.h>
#include <mkl_types.h>

#include <cmath>
#include <cstring>

// set to 1 to construct a dense mat and print some statistics for debug
#define PRINT_MAT 0

#if PRINT_MAT
#include <Eigen/Dense>
#include <Eigen/SVD>
#include <iostream>
#endif

using namespace sanm;

namespace {

class ScopedMKLThreads {
public:
    explicit ScopedMKLThreads(int nr) {
        sanm_assert(nr > 0);
        mkl_set_num_threads_local(nr);
    }

    ~ScopedMKLThreads() { mkl_set_num_threads_local(0); }
};

}  // anonymous namespace

// use Intel PARDISO. See examples_core_c.tgz/solverc/source/pardiso_unsym_c.c

class SparseSolver::Impl final : public NonCopyable {
    const size_t m_nr_unknown;
    const int m_nr_threads;
    bool m_prepared = false;

    class MatBuilder;

    std::vector<MKL_INT> m_ia, m_ja;
    std::vector<double> m_a;

    //! pointers for A'*A + p*I
    sparse_matrix_t m_aTapI = nullptr;  //!< A'A
    MKL_INT *m_aTapI_ia0 = nullptr, *m_aTapI_ia1 = nullptr,
            *m_aTapI_ja = nullptr;
    double* m_aTapI_a = nullptr;
    TensorND m_aTapI_b;
    sparse_index_base_t m_aTapI_idx{};

    fp_t m_l2_penalty = 0;

    sparse_matrix_t m_a_csr = nullptr;  //!< cached CSR descriptor
    void* m_pt[64];
    MKL_INT m_iparm[64];
    MKL_INT m_mtype = -1;
    std::vector<std::unique_ptr<MatBuilder>> m_mat_builders;

    TensorND m_solution_cache;

    void call_pardiso(MKL_INT phase, const double* b, double* x) {
        if (m_mtype == -1) {
            // initialize at first call

            if (phase == -1) {
                // destroy before first init
                return;
            }

            init_pardiso();
        }
        sanm_assert(m_prepared);

        MKL_INT maxfct = 1;  // max number of numerical factorizations
        MKL_INT mnum = 1;    // 1 <= mnum <= maxfct
        MKL_INT n = m_nr_unknown;
        MKL_INT nrhs = 1;  // number of right hand sides
        MKL_INT msglvl = sm_verbosity;
        MKL_INT error = 0;
        double* pa;
        MKL_INT *ia, *ja;
        if (m_l2_penalty == 0) {
            pa = m_a.data();
            ia = m_ia.data();
            ja = m_ja.data();
        } else {
            pa = m_aTapI_a;
            ia = m_aTapI_ia0;
            ja = m_aTapI_ja;
        }
        pardiso(m_pt, &maxfct, &mnum, &m_mtype, &phase, &n, pa, ia, ja, nullptr,
                &nrhs, m_iparm, &msglvl, const_cast<double*>(b), x, &error);
        sanm_assert(error == 0, "pardiso phase=%d failed: error=%d",
                    static_cast<int>(phase), error);
    }

    void init_pardiso() {
        if (m_l2_penalty) {
            // real and symmetric positive definite
            m_mtype = 2;
        } else {
            // real and nonsymmetric matrix
            m_mtype = 11;
        }

        pardisoinit(m_pt, &m_mtype, m_iparm);

        m_iparm[17] = 0; /* I/O: Number of nonzeros in the factor LU */
        m_iparm[18] = 0; /* I/O: Mflops for LU factorization */
        m_iparm[34] = 1; /* Zero-based indexing */

        if (m_nr_threads > 1) {
            // The parallel (OpenMP) version of the nested dissection
            // algorithm
            m_iparm[1] = 3;
        }
    }

    //! get a descriptor for the coeff matrix
    sparse_matrix_t a_csr() {
        if (!m_a_csr) {
            auto err = mkl_sparse_d_create_csr(&m_a_csr, SPARSE_INDEX_BASE_ZERO,
                                               m_nr_unknown, m_nr_unknown,
                                               m_ia.data(), m_ia.data() + 1,
                                               m_ja.data(), m_a.data());
            sanm_assert(err == SPARSE_STATUS_SUCCESS,
                        "failed to create sparse matrix: err=%d",
                        static_cast<int>(err));
        }
        return m_a_csr;
    }

public:
    static int sm_verbosity;

    Impl(size_t nr_xs) : m_nr_unknown{nr_xs}, m_nr_threads{get_num_threads()} {
        sanm_assert(m_nr_threads >= 1);
    }

    ~Impl();

    void prepare(fp_t l2_penalty);

    TensorND solve(const TensorND& b) {
        SANM_SCOPED_PROFILER("sparse_solve");
        ScopedMKLThreads mkl_threads{m_nr_threads};
        sanm_assert(b.shape().total_nr_elems() == m_nr_unknown);
        auto bptr = b.ptr();
        for (size_t i = 0; i < m_nr_unknown; ++i) {
            sanm_assert(std::isfinite(bptr[i]), "b[%zu]=%g", i, bptr[i]);
        }
        if (m_l2_penalty) {
            m_aTapI_b.set_shape({m_nr_unknown});
            matrix_descr desc{
                    .type = SPARSE_MATRIX_TYPE_GENERAL,
                    .mode = SPARSE_FILL_MODE_UPPER,
                    .diag = SPARSE_DIAG_NON_UNIT,
            };
            auto status =
                    mkl_sparse_d_mv(SPARSE_OPERATION_TRANSPOSE, 1, a_csr(),
                                    desc, bptr, 0, m_aTapI_b.woptr());
            sanm_assert(status == SPARSE_STATUS_SUCCESS,
                        "failed to compute A'b: status=%d",
                        static_cast<int>(status));
            bptr = m_aTapI_b.ptr();
        }
        call_pardiso(33, bptr,
                     m_solution_cache.set_shape({m_nr_unknown}).woptr());
        return m_solution_cache;
    }

    void dump(const TensorND& b, FILE* fout) {
        sanm_assert(m_prepared);
        sanm_assert(b.empty() || (b.shape().total_nr_elems() == m_nr_unknown));
        auto bptr = b.empty() ? nullptr : b.ptr();
        fprintf(fout, "======== begin SparseSolver_%p dump ========\n", this);
        for (size_t i = 0; i < m_nr_unknown; ++i) {
            fprintf(fout, "<eqn%-3zu>: ", i);
            for (int j = m_ia[i]; j < m_ia[i + 1]; ++j) {
                fprintf(fout, "%3.3fx%-5zu", m_a[j],
                        static_cast<size_t>(m_ja[j]));
            }
            if (bptr) {
                fprintf(fout, " = %g\n", bptr[i]);
            } else {
                fprintf(fout, "\n");
            }
        }
        fprintf(fout, "======== end SparseSolver_%p dump ========\n", this);
    }

    TensorND apply(const TensorND& x) {
        ScopedMKLThreads mkl_threads{m_nr_threads};
        sanm_assert(m_prepared && x.rank() == 1 && x.shape(0) == m_nr_unknown);
        auto ret = x.make_same_shape();
        matrix_descr desc{.type = SPARSE_MATRIX_TYPE_GENERAL,
                          .mode = SPARSE_FILL_MODE_UPPER,
                          .diag = SPARSE_DIAG_NON_UNIT};
        auto err = mkl_sparse_d_mv(SPARSE_OPERATION_NON_TRANSPOSE, 1, a_csr(),
                                   desc, x.ptr(), 0, ret.woptr());
        sanm_assert(err == SPARSE_STATUS_SUCCESS,
                    "failed to compute sparse mv: err=%d",
                    static_cast<int>(err));
        return ret;
    }

    fp_t coeff_l2() const {
        fp_t s = 0;
        for (fp_t i : m_a) {
            s += i * i;
        }
        return std::sqrt(s);
    }

    SparseMatBuilder* make_builder(size_t cidx_offset);
};
int SparseSolver::Impl::sm_verbosity = 0;

class SparseSolver::Impl::MatBuilder final
        : public SparseSolver::SparseMatBuilder {
    struct CsrRowPair {
        leastsize_t col;
        double val;

        CsrRowPair() = default;
        CsrRowPair(size_t c, double v)
                : col{static_cast<leastsize_t>(c)}, val{v} {}

        bool operator<(const CsrRowPair& rhs) const { return col < rhs.col; }
    };

    Impl* const m_owner;
    const size_t m_offset;
    size_t m_prev_constraint_cidx;
    std::vector<CsrRowPair> m_csr_last_row;

    std::vector<MKL_INT> m_ia, m_ja;
    std::vector<double> m_a;

    void flush_csr_row() {
        sanm_assert(!m_csr_last_row.empty(), "empty row %zu",
                    m_prev_constraint_cidx);
        std::sort(m_csr_last_row.begin(), m_csr_last_row.end());
        size_t jr = 0, jsize = m_csr_last_row.size();
        m_csr_last_row.emplace_back(m_owner->m_nr_unknown,
                                    0.);  // sentinel value

        if (size_t tsize = m_a.size() + m_csr_last_row.size();
            m_a.capacity() < tsize) {
            tsize = std::max(tsize, m_a.capacity() * 3 / 2);
            m_a.reserve(tsize);
            m_ja.reserve(tsize);
        }

        m_ia.push_back(m_ja.size());
        while (jr < jsize) {
            leastsize_t c = m_csr_last_row[jr].col;
            double v = m_csr_last_row[jr].val;
            ++jr;
            while (m_csr_last_row[jr].col == c) {
                v += m_csr_last_row[jr].val;
                ++jr;
            }
            m_ja.push_back(c);
            m_a.push_back(v);
        }
        m_csr_last_row.clear();
    }

public:
    MatBuilder(Impl* owner, size_t offset)
            : m_owner{owner},
              m_offset{offset},
              m_prev_constraint_cidx{offset} {}

    void add_constraint(size_t cidx, size_t xidx, double coeff) override {
        sanm_assert(!m_owner->m_prepared);
        sanm_assert(std::isfinite(coeff), "coeff[%zu,%zu]=%g", cidx, xidx,
                    coeff);

        if (std::fabs(coeff) < 1e-9) {
            return;
        }

        sanm_assert(cidx < m_owner->m_nr_unknown &&
                    xidx < m_owner->m_nr_unknown);
        if (cidx != m_prev_constraint_cidx) {
            sanm_assert(cidx == m_prev_constraint_cidx + 1,
                        "constraints not added in order: prev=%zu cur=%zu",
                        m_prev_constraint_cidx, static_cast<size_t>(cidx));
            flush_csr_row();
            m_prev_constraint_cidx = cidx;
        }
        m_csr_last_row.emplace_back(xidx, coeff);
    }

    void prepare() {
        flush_csr_row();
        ++m_prev_constraint_cidx;
    }

    size_t ja_size() const { return m_ja.size(); }

    size_t first_cidx() const { return m_offset; }

    size_t last_cidx() const { return m_prev_constraint_cidx; }

    void copy_into(size_t ia_offset, MKL_INT* ia, MKL_INT* ja, double* a) {
        for (size_t i = m_offset; i < m_prev_constraint_cidx; ++i) {
            ia[i] = m_ia[i - m_offset] + ia_offset;
        }
        memcpy(ja + ia_offset, m_ja.data(), sizeof(MKL_INT) * m_ja.size());
        memcpy(a + ia_offset, m_a.data(), sizeof(double) * m_a.size());
    }
};

void SparseSolver::Impl::prepare(fp_t l2_penalty) {
    SANM_SCOPED_PROFILER("sparse_prep");
    sanm_assert(!m_prepared);
    m_prepared = true;
    m_l2_penalty = l2_penalty;

    ScopedMKLThreads mkl_threads{m_nr_threads};

    size_t tot_ja_size = 0;
    for (auto& i : m_mat_builders) {
        i->prepare();
        tot_ja_size += i->ja_size();
    }

    std::sort(m_mat_builders.begin(), m_mat_builders.end(),
              [](const std::unique_ptr<MatBuilder>& a,
                 const std::unique_ptr<MatBuilder>& b) {
                  return a->first_cidx() < b->first_cidx();
              });

    for (size_t i = 1; i < m_mat_builders.size(); ++i) {
        sanm_assert(m_mat_builders[i]->first_cidx() ==
                    m_mat_builders[i - 1]->last_cidx());
    }
    sanm_assert(m_mat_builders.back()->last_cidx() == m_nr_unknown);

    m_ia.resize(m_nr_unknown + 1);
    m_ja.resize(tot_ja_size);
    m_a.resize(tot_ja_size);
    m_ia.back() = tot_ja_size;
    {
        size_t offset = 0;
        for (auto& i : m_mat_builders) {
            i->copy_into(offset, m_ia.data(), m_ja.data(), m_a.data());
            offset += i->ja_size();
        }
    }
    m_mat_builders.clear();

    if (m_l2_penalty) {
        SANM_SCOPED_PROFILER("sparse_A'A");
        auto status =
                mkl_sparse_syrk(SPARSE_OPERATION_TRANSPOSE, a_csr(), &m_aTapI);
        sanm_assert(status == SPARSE_STATUS_SUCCESS,
                    "failed to compute A'A: status=%d",
                    static_cast<int>(status));
        MKL_INT r, c;
        status = mkl_sparse_d_export_csr(m_aTapI, &m_aTapI_idx, &r, &c,
                                         &m_aTapI_ia0, &m_aTapI_ia1,
                                         &m_aTapI_ja, &m_aTapI_a);
        sanm_assert(status == SPARSE_STATUS_SUCCESS,
                    "failed to export A'A: status=%d",
                    static_cast<int>(status));

        sanm_assert(r == c && static_cast<size_t>(r) == m_nr_unknown);

        sanm_assert(m_aTapI_idx == SPARSE_INDEX_BASE_ZERO);
        sanm_assert(m_aTapI_ia1 == m_aTapI_ia0 + 1);
        for (int i = 0; i < r; ++i) {
            int p = m_aTapI_ia0[i], c = m_aTapI_ja[p];
            sanm_assert(i == c, "row %d: first col is %d", i, c);
            m_aTapI_a[p] += m_l2_penalty;
        }
        if (sm_verbosity) {
            printf("Linear solve with L2: n=%zu size_A=%zu size_A'A=%zu\n",
                   m_nr_unknown, static_cast<size_t>(m_ia[m_nr_unknown]),
                   static_cast<size_t>(m_aTapI_ia1[m_nr_unknown - 1]));
        }
    }

#if PRINT_MAT
    {
        using Mat = Eigen::Matrix<fp_t, Eigen::Dynamic, Eigen::Dynamic>;
        Mat M(m_nr_unknown, m_nr_unknown);
        M.setZero();
        for (size_t i = 0; i < m_nr_unknown; ++i) {
            for (int j = m_ia[i]; j < m_ia[i + 1]; ++j) {
                M(i, m_ja[j]) = m_a[j];
            }
        }
        printf("======== dump sparse solver coeff: %zu\n", m_nr_unknown);
        if (m_nr_unknown <= 10) {
            Eigen::IOFormat fmt(3, 0, ", ", ";\n", "", "", "[", "]");
            std::cout << M.format(fmt) << std::endl;
        }
        Eigen::BDCSVD<Mat> svd(M);
        std::cout << "s: " << svd.singularValues().transpose() << std::endl;
        std::cout << "det=" << M.determinant() << std::endl;
        std::cout << "======================\n";
    }
#endif  // PRINT_MAT

    // analysis and factorization
    call_pardiso(12, nullptr, nullptr);
}

SparseSolver::SparseMatBuilder* SparseSolver::Impl::make_builder(
        size_t cidx_offset = 0) {
    auto& ptr = m_mat_builders.emplace_back();
    ptr.reset(new MatBuilder{this, cidx_offset});
    return ptr.get();
}

SparseSolver::Impl::~Impl() {
    call_pardiso(-1, nullptr, nullptr);
    if (m_a_csr) {
        mkl_sparse_destroy(m_a_csr);
        m_a_csr = nullptr;
    }
    if (m_aTapI) {
        mkl_sparse_destroy(m_aTapI);
        m_aTapI = nullptr;
    }
}

void SparseSolver::set_verbosity(int verbosity) {
    Impl::sm_verbosity = verbosity;
}

SparseSolver::SparseSolver(size_t nr_xs) {
    m_pimpl.reset(new Impl{nr_xs});
}

SparseSolver::~SparseSolver() = default;

SparseSolver::SparseMatBuilder* SparseSolver::make_builder(size_t cidx_offset) {
    return m_pimpl->make_builder(cidx_offset);
}

SparseSolver& SparseSolver::prepare(fp_t l2_penalty) {
    m_pimpl->prepare(l2_penalty);
    return *this;
}

TensorND SparseSolver::solve(const TensorND& b) const {
    return m_pimpl->solve(b);
}

void SparseSolver::dump(const TensorND& b, FILE* fout) const {
    m_pimpl->dump(b, fout);
}

TensorND SparseSolver::apply(const TensorND& x) const {
    return m_pimpl->apply(x);
}

fp_t SparseSolver::coeff_l2() const {
    return m_pimpl->coeff_l2();
}

static int g_solver_num_threads = 0;
void SparseSolver::set_num_threads(int nr) {
    sanm_assert(nr >= 0);
    g_solver_num_threads = nr;
}

int SparseSolver::get_num_threads() {
    if (g_solver_num_threads) {
        return g_solver_num_threads;
    }
    return sanm::get_num_threads();
}


#include <iostream>
#include "mod/vector.hpp"
#include <boost/test/unit_test.hpp>

BOOST_AUTO_TEST_SUITE( vector_suite )

BOOST_AUTO_TEST_CASE( mod_vector_eql ) {
  mod::vector<int> v1 = {1, 2, 3};
  mod::vector<int> v2 = {1, 2, 3};

  BOOST_TEST(v1 == v2);
}

BOOST_AUTO_TEST_CASE( mod_vector_drop ) {
  mod::vector<int> v1 = {1, 2, 3, 4, 5, 6};
  mod::vector<int> v2 = {4, 5, 6};

  BOOST_TEST(v1.drop(3) == v2);
}

BOOST_AUTO_TEST_CASE( mod_vector_take ) {
  mod::vector<int> v1 = {1, 2, 3, 4, 5, 6};
  mod::vector<int> v2 = {1, 2, 3};

  BOOST_TEST(v1.take(3) == v2);
}

BOOST_AUTO_TEST_CASE( mod_vector_fold ) {
  mod::vector<int> v1 = {1, 2, 3, 4, 5, 6};
  int val = 1 + 2 + 3 + 4 + 5 + 6;

  BOOST_TEST(val == v1.fold<int>([](int a, int b){ return a + b; }, 0));

  v1 = {};

  BOOST_TEST(0 == v1.fold<int>([](int a, int b){ return a + b; }, 0));
}

BOOST_AUTO_TEST_CASE( mod_vector_map ) {
  mod::vector<int> v1 = {1, 2, 3, 4, 5, 6};
  mod::vector<int> v2 = {2, 3, 4, 5, 6, 7};

  BOOST_TEST(v1.map<int>([](int x){ return x + 1; }) == v2);
}

BOOST_AUTO_TEST_CASE( mod_vector_zip ) {
  mod::vector<int> v1 = {1, 2, 3, 4, 5, 6};
  mod::vector<int> v2 = {2, 3, 4, 5, 6, 7};
  mod::vector<std::pair<int, int>> v3 = {
    std::pair<int, int>(1, 2),
    std::pair<int, int>(2, 3),
    std::pair<int, int>(3, 4),
    std::pair<int, int>(4, 5),
    std::pair<int, int>(5, 6),
    std::pair<int, int>(6, 7)};

  BOOST_TEST(v1.zip(v2) == v3);
}

// BOOST_AUTO_TEST_CASE( mod_vector_filter ) {
//   mod::vector<int> v1 = {1, 2, 3, 4, 5, 6, 7, 8, 9, 10};
//   mod::vector<int> v2 = {2, 4, 6, 8, 10};

//   mod::vector<int> out = v1.

//   BOOST_TEST(v1.zip(v2) == v3);
// }

BOOST_AUTO_TEST_SUITE_END()


/* boost random/random_number_generator.hpp header file
 *
 * Copyright Jens Maurer 2000-2001
 * Distributed under the Boost Software License, Version 1.0. (See
 * accompanying file LICENSE_1_0.txt or copy at
 * http://www.boost.org/LICENSE_1_0.txt)
 *
 * See http://www.boost.org for most recent version including documentation.
 *
 * $Id: random_number_generator.hpp,v 1.1 2007/02/12 18:25:54 irving Exp $
 *
 * Revision history
 *  2001-02-18  moved to individual header files
 */

#ifndef BOOST_RANDOM_RANDOM_NUMBER_GENERATOR_HPP
#define BOOST_RANDOM_RANDOM_NUMBER_GENERATOR_HPP

#include <boost/config.hpp>
#include <boost/limits.hpp>
#include <boost/static_assert.hpp>
#include <boost/random/uniform_int.hpp>
#include <boost/random/variate_generator.hpp>

namespace boost {

// a model for RandomNumberGenerator std:25.2.11 [lib.alg.random.shuffle]
template<class UniformRandomNumberGenerator, class IntType = long>
class random_number_generator
{
public:
  typedef UniformRandomNumberGenerator base_type;
  typedef IntType argument_type;
  typedef IntType result_type;
  random_number_generator(base_type& rng) : _rng(rng)
  {
#ifndef BOOST_NO_LIMITS_COMPILE_TIME_CONSTANTS
    BOOST_STATIC_ASSERT(std::numeric_limits<result_type>::is_integer);
#endif
  }
  // compiler-generated copy ctor is fine
  // assignment is disallowed because there is a reference member

  result_type operator()(argument_type n)
  {
    typedef uniform_int<IntType> dist_type;
    return variate_generator<base_type&, dist_type>(_rng, dist_type(0, n-1))();
  }

private:
  base_type& _rng;
};

} // namespace boost

#endif // BOOST_RANDOM_RANDOM_NUMBER_GENERATOR_HPP


#include "stdafx.h"

#include "problem.hpp"
#include "utility.hpp"

#include <fstream>
#include <iomanip>
#include <regex>
#include <utility>
#include <map>
#include <numeric>
#include <chrono>
#include <boost/algorithm/string.hpp>

struct advent_2017_17 : problem
{
	advent_2017_17() noexcept : problem(2017, 17) {
	}

protected:
	std::size_t steps = 0;

	void prepare_input(std::ifstream& fin) override {
		std::string line;
		std::getline(fin, line);
		steps = std::stoull(line);
	}

	std::string part_1() override {
		std::vector<std::size_t> buffer;
		buffer.push_back(0);
		std::size_t current_position = 0;
		for(std::size_t i = 1; i <= 2017; ++i) {
			current_position = ((current_position + steps) % i) + 1ui64;
			buffer.insert(std::begin(buffer) + gsl::narrow<std::ptrdiff_t>(current_position), i);
		}
		const auto position = std::find(std::begin(buffer), std::end(buffer), 2017ui64) + 1;
		return std::to_string((position != std::end(buffer)) ? *position : 0ui64);
	}

	std::string part_2() override {
		std::size_t current_position = 0;
		std::size_t after_zero = 0;
		for(std::size_t i = 1; i <= 50'000'000; ++i) {
			current_position = ((current_position + steps) % i) + 1ui64;
			if(current_position == 1) {
				after_zero = i;
			}
		}
		return std::to_string(after_zero);
	}
};

REGISTER_SOLVER(2017, 17);


/* boost random/discard_block.hpp header file
 *
 * Copyright Jens Maurer 2002
 * Copyright Steven Watanabe 2010
 * Distributed under the Boost Software License, Version 1.0. (See
 * accompanying file LICENSE_1_0.txt or copy at
 * http://www.boost.org/LICENSE_1_0.txt)
 *
 * See http://www.boost.org for most recent version including documentation.
 *
 * $Id$
 *
 * Revision history
 *  2001-03-02  created
 */

#ifndef BOOST_RANDOM_DISCARD_BLOCK_HPP
#define BOOST_RANDOM_DISCARD_BLOCK_HPP

#include <iostream>
#include <boost/config.hpp>
#include <boost/cstdint.hpp>
#include <boost/limits.hpp>
#include <boost/static_assert.hpp>
#include <boost/random/detail/config.hpp>
#include <boost/random/detail/seed.hpp>
#include <boost/random/detail/seed_impl.hpp>

namespace boost {
namespace random {

/**
 * The class template \discard_block_engine is a model of
 * \pseudo_random_number_generator.  It modifies
 * another generator by discarding parts of its output.
 * Out of every block of @c p results, the first @c r
 * will be returned and the rest discarded.
 *
 * Requires: 0 < p <= r
 */
template<class UniformRandomNumberGenerator, std::size_t p, std::size_t r>
class discard_block_engine
{
    typedef typename detail::seed_type<
        typename UniformRandomNumberGenerator::result_type>::type seed_type;
public:
    typedef UniformRandomNumberGenerator base_type;
    typedef typename base_type::result_type result_type;

    BOOST_STATIC_CONSTANT(std::size_t, block_size = p);
    BOOST_STATIC_CONSTANT(std::size_t, used_block = r);

    BOOST_STATIC_CONSTANT(bool, has_fixed_range = false);
    BOOST_STATIC_CONSTANT(std::size_t, total_block = p);
    BOOST_STATIC_CONSTANT(std::size_t, returned_block = r);

    BOOST_STATIC_ASSERT(total_block >= returned_block);

    /** Uses the default seed for the base generator. */
    discard_block_engine() : _rng(), _n(0) { }
    /** Constructs a new \discard_block_engine with a copy of rng. */
    explicit discard_block_engine(const base_type & rng) : _rng(rng), _n(0) { }

#ifndef BOOST_NO_CXX11_RVALUE_REFERENCES
    /** Constructs a new \discard_block_engine with rng. */
    explicit discard_block_engine(base_type && rng) : _rng(rng), _n(0) { }
#endif

    /**
     * Creates a new \discard_block_engine and seeds the underlying
     * generator with @c value
     */
    BOOST_RANDOM_DETAIL_ARITHMETIC_CONSTRUCTOR(discard_block_engine,
                                               seed_type, value)
    { _rng.seed(value); _n = 0; }

    /**
     * Creates a new \discard_block_engine and seeds the underlying
     * generator with @c seq
     */
    BOOST_RANDOM_DETAIL_SEED_SEQ_CONSTRUCTOR(discard_block_engine, SeedSeq, seq)
    { _rng.seed(seq); _n = 0; }

    /**
     * Creates a new \discard_block_engine and seeds the underlying
     * generator with first and last.
     */
    template<class It> discard_block_engine(It& first, It last)
      : _rng(first, last), _n(0) { }

    /** default seeds the underlying generator. */
    void seed() { _rng.seed(); _n = 0; }
    /** Seeds the underlying generator with s. */
    BOOST_RANDOM_DETAIL_ARITHMETIC_SEED(discard_block_engine, seed_type, s)
    { _rng.seed(s); _n = 0; }
    /** Seeds the underlying generator with seq. */
    BOOST_RANDOM_DETAIL_SEED_SEQ_SEED(discard_block_engine, SeedSeq, seq)
    { _rng.seed(seq); _n = 0; }
    /** Seeds the underlying generator with first and last. */
    template<class It> void seed(It& first, It last)
    { _rng.seed(first, last); _n = 0; }

    /** Returns the underlying engine. */
    const base_type& base() const { return _rng; }

    /** Returns the next value of the generator. */
    result_type operator()()
    {
        if(_n >= returned_block) {
            // discard values of random number generator
            // Don't use discard, since we still need to
            // be somewhat compatible with TR1.
            // _rng.discard(total_block - _n);
            for(std::size_t i = 0; i < total_block - _n; ++i) {
                _rng();
            }
            _n = 0;
        }
        ++_n;
        return _rng();
    }

    void discard(boost::uintmax_t z)
    {
        for(boost::uintmax_t j = 0; j < z; ++j) {
            (*this)();
        }
    }

    template<class It>
    void generate(It first, It last)
    { detail::generate(*this, first, last); }

    /**
     * Returns the smallest value that the generator can produce.
     * This is the same as the minimum of the underlying generator.
     */
    static result_type min BOOST_PREVENT_MACRO_SUBSTITUTION ()
    { return (base_type::min)(); }
    /**
     * Returns the largest value that the generator can produce.
     * This is the same as the maximum of the underlying generator.
     */
    static result_type max BOOST_PREVENT_MACRO_SUBSTITUTION ()
    { return (base_type::max)(); }

#ifndef BOOST_RANDOM_NO_STREAM_OPERATORS
    /** Writes a \discard_block_engine to a @c std::ostream. */
    template<class CharT, class Traits>
    friend std::basic_ostream<CharT,Traits>&
    operator<<(std::basic_ostream<CharT,Traits>& os,
               const discard_block_engine& s)
    {
        os << s._rng << ' ' << s._n;
        return os;
    }

    /** Reads a \discard_block_engine from a @c std::istream. */
    template<class CharT, class Traits>
    friend std::basic_istream<CharT,Traits>&
    operator>>(std::basic_istream<CharT,Traits>& is, discard_block_engine& s)
    {
        is >> s._rng >> std::ws >> s._n;
        return is;
    }
#endif

    /** Returns true if the two generators will produce identical sequences. */
    friend bool operator==(const discard_block_engine& x,
                           const discard_block_engine& y)
    { return x._rng == y._rng && x._n == y._n; }
    /** Returns true if the two generators will produce different sequences. */
    friend bool operator!=(const discard_block_engine& x,
                           const discard_block_engine& y)
    { return !(x == y); }

private:
    base_type _rng;
    std::size_t _n;
};

#ifndef BOOST_NO_INCLASS_MEMBER_INITIALIZATION
//  A definition is required even for integral static constants
template<class URNG, std::size_t p, std::size_t r>
const bool discard_block_engine<URNG, p, r>::has_fixed_range;
template<class URNG, std::size_t p, std::size_t r>
const std::size_t discard_block_engine<URNG, p, r>::total_block;
template<class URNG, std::size_t p, std::size_t r>
const std::size_t discard_block_engine<URNG, p, r>::returned_block;
template<class URNG, std::size_t p, std::size_t r>
const std::size_t discard_block_engine<URNG, p, r>::block_size;
template<class URNG, std::size_t p, std::size_t r>
const std::size_t discard_block_engine<URNG, p, r>::used_block;
#endif

/// \cond \show_deprecated

template<class URNG, int p, int r>
class discard_block : public discard_block_engine<URNG, p, r>
{
    typedef discard_block_engine<URNG, p, r> base_t;
public:
    typedef typename base_t::result_type result_type;
    discard_block() {}
    template<class T>
    discard_block(T& arg) : base_t(arg) {}
    template<class T>
    discard_block(const T& arg) : base_t(arg) {}
    template<class It>
    discard_block(It& first, It last) : base_t(first, last) {}
    result_type min BOOST_PREVENT_MACRO_SUBSTITUTION ()
    { return (this->base().min)(); }
    result_type max BOOST_PREVENT_MACRO_SUBSTITUTION ()
    { return (this->base().max)(); }
};

/// \endcond

namespace detail {

    template<class Engine>
    struct generator_bits;

    template<class URNG, std::size_t p, std::size_t r>
    struct generator_bits<discard_block_engine<URNG, p, r> > {
        static std::size_t value() { return generator_bits<URNG>::value(); }
    };

    template<class URNG, int p, int r>
    struct generator_bits<discard_block<URNG, p, r> > {
        static std::size_t value() { return generator_bits<URNG>::value(); }
    };

}

} // namespace random

} // namespace boost

#endif // BOOST_RANDOM_DISCARD_BLOCK_HPP


/*
	This file is part of the implementation for the technical paper

		Field-Aligned Online Surface Reconstruction
		Nico Schertler, Marco Tarini, Wenzel Jakob, Misha Kazhdan, Stefan Gumhold, Daniele Panozzo
		ACM TOG 36, 4, July 2017 (Proceedings of SIGGRAPH 2017)

	Use of this source code is granted via a BSD-style license, which can be found
	in License.txt in the repository root.

	@author Wenzel Jakob
	@author Nico Schertler
*/

#include "osr/field.h"
#include "osr/common.h"

#include <Eigen/Geometry>

namespace osr
{

	static const Float sqrt_3_over_4 = 0.866025403784439f;

	Vector3f rotate180(const Vector3f &q, const Vector3f &/* unused */) {
		return -q;
	}

	Vector3f rotate180_by(const Vector3f &q, const Vector3f &/* unused */, int amount) {
		return (amount & 1) ? Vector3f(-q) : q;
	}

	Vector2i rshift180(Vector2i shift, int amount) {
		if (amount & 1)
			shift = -shift;
		return shift;
	}

	Vector3f rotate90(const Vector3f &q, const Vector3f &n) {
		return n.cross(q);
	}

	Vector3f rotate90_by(const Vector3f &q, const Vector3f &n, int amount) {
		return ((amount & 1) ? (n.cross(q)) : q) * (amount < 2 ? 1.0f : -1.0f);
	}

	Vector2i rshift90(Vector2i shift, int amount) {
		if (amount & 1)
			shift = Vector2i(-shift.y(), shift.x());
		if (amount >= 2)
			shift = -shift;
		return shift;
	}

	Vector3f rotate60(const Vector3f &d, const Vector3f &n) {
		return sqrt_3_over_4 * n.cross(d) + 0.5f*(d + n * n.dot(d));
	}

	Vector2i rshift60(Vector2i shift, int amount) {
		for (int i = 0; i < amount; ++i)
			shift = Vector2i(-shift.y(), shift.x() + shift.y());
		return shift;
	}

	Vector3f rotate60_by(const Vector3f &d, const Vector3f &n, int amount) {
		switch (amount) {
		case 0: return d;
		case 1: return rotate60(d, n);
		case 2: return -rotate60(d, -n);
		case 3: return -d;
		case 4: return -rotate60(d, n);
		case 5: return rotate60(d, -n);
		}
		throw std::runtime_error("rotate60: invalid argument");
	}

	Vector3f rotate_vector_into_plane(Vector3f q, const Vector3f &source_normal, const Vector3f &target_normal) {
		const Float cosTheta = source_normal.dot(target_normal);
		if (cosTheta < 0.9999f) {
			Vector3f axis = source_normal.cross(target_normal);
			q = q * cosTheta + axis.cross(q) +
				axis * (axis.dot(q) * (1.0f - cosTheta) / axis.dot(axis));
		}
		return q;
	}

	inline Vector3f middle_point(const Vector3f &p0, const Vector3f &n0, const Vector3f &p1, const Vector3f &n1) {
		/* How was this derived?
		 *
		 * Minimize \|x-p0\|^2 + \|x-p1\|^2, where
		 * dot(n0, x) == dot(n0, p0)
		 * dot(n1, x) == dot(n1, p1)
		 *
		 * -> Lagrange multipliers, set derivative = 0
		 *  Use first 3 equalities to write x in terms of
		 *  lambda_1 and lambda_2. Substitute that into the last
		 *  two equations and solve for the lambdas. Finally,
		 *  add a small epsilon term to avoid issues when n1=n2.
		 */
		Float n0p0 = n0.dot(p0), n0p1 = n0.dot(p1),
			n1p0 = n1.dot(p0), n1p1 = n1.dot(p1),
			n0n1 = n0.dot(n1),
			denom = 1.0f / (1.0f - n0n1*n0n1 + 1e-4f),
			lambda_0 = 2.0f*(n0p1 - n0p0 - n0n1*(n1p0 - n1p1))*denom,
			lambda_1 = 2.0f*(n1p0 - n1p1 - n0n1*(n0p1 - n0p0))*denom;

		return 0.5f * (p0 + p1) - 0.25f * (n0 * lambda_0 + n1 * lambda_1);
	}

	IOrientationFieldTraits::IOrientationFieldTraits(int rosy)
		: _rosy(rosy)
	{ }

	template<int RoSy>
	OrientationFieldTraits<RoSy>::OrientationFieldTraits()
		: IOrientationFieldTraits(RoSy)
	{ }

	template <>
	void OrientationFieldTraits<2>::findCompatible(const Vector3f &q0, const Vector3f &n0, const Vector3f &q1, const Vector3f &n1, Vector3f& compat1, Vector3f& compat2)
	{
		compat1 = q0;
		compat2 = q1 * signum(q0.dot(q1));
	}

	template <>
	void OrientationFieldTraits<2>::findCompatible(const Vector3f &dirField0, const Vector3f &normal0, const Matrix3Xf& principalDirections, Matrix3Xf& compatOut, VectorXf& scoreOut)
	{
		float dot = dirField0.dot(principalDirections.col(0));
		compatOut.col(0) = dirField0 * signum(dot);
		scoreOut(0) = std::abs(dot);
	}

	template <>
	void OrientationFieldTraits<2>::getPrincipalDirections(const Vector3f& dirField, const Vector3f& normal, Matrix3Xf& outDirections)
	{
		outDirections.col(0) = dirField;
	}

	template <>
	void OrientationFieldTraits<4>::getPrincipalDirections(const Vector3f& dirField, const Vector3f& normal, Matrix3Xf& outDirections)
	{
		outDirections.col(0) = dirField;
		outDirections.col(1) = normal.cross(dirField);
	}

	template <>
	void OrientationFieldTraits<4>::findCompatible(const Vector3f &q0, const Vector3f &n0, const Vector3f &q1, const Vector3f &n1, Vector3f& compat1, Vector3f& compat2)
	{
		Matrix3Xf A(3, 2), B(3, 2);
		getPrincipalDirections(q0, n0, A);
		getPrincipalDirections(q1, n1, B);

		Float best_score = -std::numeric_limits<Float>::infinity();
		int best_a = 0, best_b = 0;

		for (int i = 0; i < 2; ++i) {
			for (int j = 0; j < 2; ++j) {
				Float score = std::abs(A.col(i).dot(B.col(j)));
				if (score > best_score) {
					best_a = i;
					best_b = j;
					best_score = score;
				}
			}
		}

		const Float dp = A.col(best_a).dot(B.col(best_b));
		compat1 = A.col(best_a);
		compat2 = B.col(best_b) * signum(dp);
	}

	template <>
	void OrientationFieldTraits<4>::findCompatible(const Vector3f &dirField0, const Vector3f &normal0, const Matrix3Xf& principalDirections, Matrix3Xf& compatOut, VectorXf& scoreOut)
	{
		Matrix3Xf B(3, 2);
		getPrincipalDirections(dirField0, normal0, B);

		for (int i = 0; i < 2; ++i)
		{
			Float best_score = -std::numeric_limits<Float>::infinity();
			int best = 0;

			for (int j = 0; j < 2; ++j)
			{
				Float score = std::abs(principalDirections.col(i).dot(B.col(j)));
				if (score > best_score)
				{
					best = j;
					best_score = score;
				}
			}

			const Float dp = principalDirections.col(i).dot(B.col(best));
			compatOut.col(i) = B.col(best) * signum(dp);
			scoreOut(i) = best_score;
		}
	}

	template <>
	void OrientationFieldTraits<6>::getPrincipalDirections(const Vector3f& dirField, const Vector3f& normal, Matrix3Xf& outDirections)
	{
		outDirections.col(0) = rotate60(dirField, -normal);
		outDirections.col(1) = dirField;
		outDirections.col(2) = rotate60(dirField, normal);
	}

	template <>
	void OrientationFieldTraits<6>::findCompatible(const Vector3f &q0, const Vector3f &n0, const Vector3f &q1, const Vector3f &n1, Vector3f& compat1, Vector3f& compat2)
	{
		Matrix3Xf A(3, 3), B(3, 3);
		getPrincipalDirections(q0, n0, A);
		getPrincipalDirections(q1, n1, B);

		Float best_score = -std::numeric_limits<Float>::infinity();
		int best_a = 0, best_b = 0;

		for (int i = 0; i < 3; ++i) {
			for (int j = 0; j < 3; ++j) {
				Float score = std::abs(A.col(i).dot(B.col(j)));
				if (score > best_score) {
					best_a = i;
					best_b = j;
					best_score = score;
				}
			}
		}

		const Float dp = A.col(best_a).dot(B.col(best_b));
		compat1 = A.col(best_a);
		compat2 = B.col(best_b) * signum(dp);
	}

	template <>
	void OrientationFieldTraits<6>::findCompatible(const Vector3f &dirField0, const Vector3f &normal0, const Matrix3Xf& principalDirections, Matrix3Xf& compatOut, VectorXf& scoreOut)
	{
		Matrix3Xf B(3, 3);
		getPrincipalDirections(dirField0, normal0, B);

		for (int i = 0; i < 3; ++i)
		{
			Float best_score = -std::numeric_limits<Float>::infinity();
			int best = 0;
			for (int j = 0; j < 3; ++j)
			{
				Float score = std::abs(principalDirections.col(i).dot(B.col(j)));
				if (score > best_score)
				{
					best = j;
					best_score = score;
				}
			}
			const Float dp = principalDirections.col(i).dot(B.col(best));
			compatOut.col(i) = B.col(best) * signum(dp);
			scoreOut(i) = best_score;
		}
	}

	IOrientationFieldTraits * getOrientationFieldTraits(int roSy)
	{
		if (roSy == 2) {
			return new OrientationFieldTraits<2>();
		}
		else if (roSy == 4) {
			return new OrientationFieldTraits<4>();
		}
		else if (roSy == 6) {
			return new OrientationFieldTraits<6>();
		}
		else {
			throw std::runtime_error("Invalid rotation symmetry type " + std::to_string(roSy) + "!");
		}
	}

	// ---- Position Field ----

	IPositionFieldTraits::IPositionFieldTraits(int posy)
		:_posy(posy)
	{
	}

	template<int PoSy>
	PositionFieldTraits<PoSy>::PositionFieldTraits()
		: IPositionFieldTraits(PoSy)
	{
	}

	IPositionFieldTraits * getPositionFieldTraits(int poSy)
	{
		if (poSy == 4) {
			return new PositionFieldTraits<4>();
		}
		else if (poSy == 6) {
			return new PositionFieldTraits<6>();
		}
		else {
			throw std::runtime_error("Invalid position symmetry type " + std::to_string(poSy) + "!");
		}
	}

	inline Vector2i position_round_index_3(const Vector3f &o, const Vector3f &q,
		const Vector3f &n, const Vector3f &p,
		Float scale, Float inv_scale)
	{
		Vector3f t = rotate60(q, n);
		Vector3f d = p - o;
		Float dpq = q.dot(d), dpt = t.dot(d);
		int u = (int)std::floor((4 * dpq - 2 * dpt) * (1.0f / 3.0f) * inv_scale);
		int v = (int)std::floor((-2 * dpq + 4 * dpt) * (1.0f / 3.0f) * inv_scale);

		Float best_cost = std::numeric_limits<Float>::infinity();
		int best_i = -1;

		for (int i = 0; i < 4; ++i) {
			Vector3f ot = o + (q*(u + (i & 1)) + t * (v + ((i & 2) >> 1))) * scale;
			Float cost = (ot - p).squaredNorm();
			if (cost < best_cost) {
				best_i = i;
				best_cost = cost;
			}
		}

		return Vector2i(
			u + (best_i & 1), v + ((best_i & 2) >> 1)
		);
	}

	inline Vector2i position_round_index_4(const Vector3f &o, const Vector3f &q,
		const Vector3f &n, const Vector3f &p,
		Float /* unused */, Float inv_scale)
	{
		Vector3f t = n.cross(q);
		Vector3f d = p - o;
		return Vector2i(
			(int)std::round(q.dot(d) * inv_scale),
			(int)std::round(t.dot(d) * inv_scale));
	}

	inline Vector3f position_floor_3(const Vector3f &o, const Vector3f &q,
		const Vector3f &n, const Vector3f &p,
		Float scale, Float inv_scale) {
		Vector3f t = rotate60(q, n);
		Vector3f d = p - o;
		Float dpq = q.dot(d), dpt = t.dot(d);
		Float u = std::floor((4 * dpq - 2 * dpt) * (1.0f / 3.0f) * inv_scale);
		Float v = std::floor((-2 * dpq + 4 * dpt) * (1.0f / 3.0f) * inv_scale);

		return o + (q*u + t*v) * scale;
	}

	inline Vector3f position_floor_4(const Vector3f &o, const Vector3f &q,
		const Vector3f &n, const Vector3f &p,
		Float scale, Float inv_scale) {
		Vector3f t = n.cross(q);
		Vector3f d = p - o;
		return o +
			q * std::floor(q.dot(d) * inv_scale) * scale +
			t * std::floor(t.dot(d) * inv_scale) * scale;
	}

	inline Vector2i position_floor_index_3(const Vector3f &o, const Vector3f &q,
		const Vector3f &n, const Vector3f &p,
		Float /* scale */, Float inv_scale) {
		Vector3f t = rotate60(q, n);
		Vector3f d = p - o;
		Float dpq = q.dot(d), dpt = t.dot(d);
		int u = (int)std::floor((4 * dpq - 2 * dpt) * (1.0f / 3.0f) * inv_scale);
		int v = (int)std::floor((-2 * dpq + 4 * dpt) * (1.0f / 3.0f) * inv_scale);

		return Vector2i(u, v);
	}

	inline Vector2i position_floor_index_4(const Vector3f &o, const Vector3f &q,
		const Vector3f &n, const Vector3f &p,
		Float /* unused */, Float inv_scale) {
		Vector3f t = n.cross(q);
		Vector3f d = p - o;
		return Vector2i(
			(int)std::floor(q.dot(d) * inv_scale),
			(int)std::floor(t.dot(d) * inv_scale));
	}

	template<>
	Vector3f PositionFieldTraits<4>::positionRound(const Vector3f &o, const Vector3f &q,
		const Vector3f &n, const Vector3f &p,
		Float scale, Float inv_scale)
	{
		Vector3f t = n.cross(q);
		Vector3f d = p - o;
		return o +
			q * std::round(q.dot(d) * inv_scale) * scale +
			t * std::round(t.dot(d) * inv_scale) * scale;
	}

	template<>
	Vector3f PositionFieldTraits<6>::positionRound(const Vector3f &o, const Vector3f &q,
		const Vector3f &n, const Vector3f &p,
		Float scale, Float inv_scale)
	{
		Vector3f t = rotate60(q, n);
		Vector3f d = p - o;

		Float dpq = q.dot(d), dpt = t.dot(d);
		Float u = std::floor((4 * dpq - 2 * dpt) * (1.0f / 3.0f) * inv_scale);
		Float v = std::floor((-2 * dpq + 4 * dpt) * (1.0f / 3.0f) * inv_scale);

		Float best_cost = std::numeric_limits<Float>::infinity();
		int best_i = -1;

		for (int i = 0; i < 4; ++i) {
			Vector3f ot = o + (q*(u + (i & 1)) + t*(v + ((i & 2) >> 1))) * scale;
			Float cost = (ot - p).squaredNorm();
			if (cost < best_cost) {
				best_i = i;
				best_cost = cost;
			}
		}

		return o + (q*(u + (best_i & 1)) + t*(v + ((best_i & 2) >> 1))) * scale;
	}

	template <>
	void PositionFieldTraits<4>::findCompatible(
		const Vector3f &p0, const Vector3f &n0, const Vector3f &q0, const Vector3f &o0,
		const Vector3f &p1, const Vector3f &n1, const Vector3f &q1, const Vector3f &o1,
		Float scale, Float inv_scale, Vector3f& compat1, Vector3f& compat2)
	{

		Vector3f t0 = n0.cross(q0), t1 = n1.cross(q1);
		Vector3f middle = middle_point(p0, n0, p1, n1);
		Vector3f o0p = position_floor_4(o0, q0, n0, middle, scale, inv_scale);
		Vector3f o1p = position_floor_4(o1, q1, n1, middle, scale, inv_scale);

		Float best_cost = std::numeric_limits<Float>::infinity();
		int best_i = -1, best_j = -1;

		for (int i = 0; i < 4; ++i)
		{
			Vector3f o0t = o0p + (q0 * (i & 1) + t0 * ((i & 2) >> 1)) * scale;
			for (int j = 0; j < 4; ++j)
			{
				Vector3f o1t = o1p + (q1 * (j & 1) + t1 * ((j & 2) >> 1)) * scale;
				Float cost = (o0t - o1t).squaredNorm();

				if (cost < best_cost)
				{
					best_i = i;
					best_j = j;
					best_cost = cost;
				}
			}
		}

		compat1 = o0p + (q0 * (best_i & 1) + t0 * ((best_i & 2) >> 1)) * scale;
		compat2 = o1p + (q1 * (best_j & 1) + t1 * ((best_j & 2) >> 1)) * scale;
	}

	template <>
	std::pair<Vector2i, Vector2i> PositionFieldTraits<4>::findCompatibleIndex(
		const Vector3f &p0, const Vector3f &n0, const Vector3f &q0, const Vector3f &o0,
		const Vector3f &p1, const Vector3f &n1, const Vector3f &q1, const Vector3f &o1,
		Float scale, Float inv_scale, Float* error) {
		Vector3f t0 = n0.cross(q0), t1 = n1.cross(q1);
		Vector3f middle = middle_point(p0, n0, p1, n1);
		Vector2i o0p = position_floor_index_4(o0, q0, n0, middle, scale, inv_scale);
		Vector2i o1p = position_floor_index_4(o1, q1, n1, middle, scale, inv_scale);

		Float best_cost = std::numeric_limits<Float>::infinity();
		int best_i = -1, best_j = -1;

		for (int i = 0; i < 4; ++i) {
			Vector3f o0t = o0 + (q0 * ((i & 1) + o0p[0]) + t0 * (((i & 2) >> 1) + o0p[1])) * scale;
			for (int j = 0; j < 4; ++j) {
				Vector3f o1t = o1 + (q1 * ((j & 1) + o1p[0]) + t1 * (((j & 2) >> 1) + o1p[1])) * scale;
				Float cost = (o0t - o1t).squaredNorm();

				if (cost < best_cost) {
					best_i = i;
					best_j = j;
					best_cost = cost;
				}
			}
		}

		if (error)
			*error = best_cost;

		return std::make_pair(
			Vector2i((best_i & 1) + o0p[0], ((best_i & 2) >> 1) + o0p[1]),
			Vector2i((best_j & 1) + o1p[0], ((best_j & 2) >> 1) + o1p[1]));
	}

	template <>
	void PositionFieldTraits<6>::findCompatible(
		const Vector3f &p0, const Vector3f &n0, const Vector3f &q0, const Vector3f &_o0,
		const Vector3f &p1, const Vector3f &n1, const Vector3f &q1, const Vector3f &_o1,
		Float scale, Float inv_scale, Vector3f& compat1, Vector3f& compat2) {
		Vector3f middle = middle_point(p0, n0, p1, n1);
		Vector3f o0 = position_floor_3(_o0, q0, n0, middle, scale, inv_scale);
		Vector3f o1 = position_floor_3(_o1, q1, n1, middle, scale, inv_scale);

		Vector3f t0 = rotate60(q0, n0), t1 = rotate60(q1, n1);
		Float best_cost = std::numeric_limits<Float>::infinity();
		int best_i = -1, best_j = -1;
		for (int i = 0; i < 4; ++i) {
			Vector3f o0t = o0 + (q0*(i & 1) + t0*((i & 2) >> 1)) * scale;
			for (int j = 0; j < 4; ++j) {
				Vector3f o1t = o1 + (q1*(j & 1) + t1*((j & 2) >> 1)) * scale;
				Float cost = (o0t - o1t).squaredNorm();

				if (cost < best_cost) {
					best_i = i;
					best_j = j;
					best_cost = cost;
				}
			}
		}

		compat1 = o0 + (q0*(best_i & 1) + t0*((best_i & 2) >> 1)) * scale;
		compat2 = o1 + (q1*(best_j & 1) + t1*((best_j & 2) >> 1)) * scale;
	}

	template <>
	std::pair<Vector2i, Vector2i> PositionFieldTraits<6>::findCompatibleIndex(
		const Vector3f &p0, const Vector3f &n0, const Vector3f &q0, const Vector3f &o0,
		const Vector3f &p1, const Vector3f &n1, const Vector3f &q1, const Vector3f &o1,
		Float scale, Float inv_scale, Float* error) {
		Vector3f t0 = rotate60(q0, n0), t1 = rotate60(q1, n1);
		Vector3f middle = middle_point(p0, n0, p1, n1);
		Vector2i o0i = position_floor_index_3(o0, q0, n0, middle, scale, inv_scale);
		Vector2i o1i = position_floor_index_3(o1, q1, n1, middle, scale, inv_scale);

		Float best_cost = std::numeric_limits<Float>::infinity();
		int best_i = -1, best_j = -1;
		for (int i = 0; i < 4; ++i) {
			Vector3f o0t = o0 + (q0*(o0i.x() + (i & 1)) + t0*(o0i.y() + ((i & 2) >> 1))) * scale;
			for (int j = 0; j < 4; ++j) {
				Vector3f o1t = o1 + (q1*(o1i.x() + (j & 1)) + t1*(o1i.y() + ((j & 2) >> 1))) * scale;
				Float cost = (o0t - o1t).squaredNorm();

				if (cost < best_cost) {
					best_i = i;
					best_j = j;
					best_cost = cost;
				}
			}
		}

		if (error)
			*error = best_cost;

		return std::make_pair(
			Vector2i(o0i.x() + (best_i & 1), o0i.y() + ((best_i & 2) >> 1)),
			Vector2i(o1i.x() + (best_j & 1), o1i.y() + ((best_j & 2) >> 1)));
	}
}


#include <iostream>

#include <boost/shared_ptr.hpp>

#include <boost/numeric/ublas/vector.hpp>

typedef boost::numeric::ublas::vector<float> Vector3;
typedef boost::shared_ptr<Vector3> SpVector3;

class Test
{
public:
	Test() : pos(3)
	{
		pos(0) = pos(1) = pos(2) = .0f;
	}

	void Add(Vector3 translation)
	{
		pos += translation;
	}

	Vector3 pos;
};

typedef boost::shared_ptr<Test> SpTest;

class Sample
{
public:
	Sample()
	{
		SpTest test(new Test());
		this->spTest = test;
	}

	SpTest getTest()    { return this->spTest; }

	SpTest spTest;
};

int main()
{
	Sample sample;

	Vector3 vec(3);
	vec(0) = vec(1) = vec(2) = 2.0f;

	sample.getTest()->Add(vec);

	std::cout << sample.spTest->pos(0) << std::endl;

	return 0;
}


#define BOOST_TEST_MODULE MyTest
#define BOOST_TEST_DYN_LINK

#include <boost/test/unit_test.hpp>
#include "pose.h"
#include <cmath>
#include <iostream>
#include "Criteria/criterion.h"
#include "mcdmfunction.h"
#include "Criteria/criterioncomparator.h"
#include <Criteria/sensingtimecriterion.h>
#include <Criteria/traveldistancecriterion.h>

using namespace std;

BOOST_AUTO_TEST_CASE( test_criterioncomparator )
{
    // creation of a test position
    Pose p = Pose(2,4,40,5,30) ;
    // creation of two candidate criteria
    SensingTimeCriterion sensingTime(0.5);
    TravelDistanceCriterion travelDistance(0.8);
    // insertion of two criteria's evaluation
    sensingTime.insertEvaluation(p,0.4);
    travelDistance.insertEvaluation(p,0.7);

    /*
    //to be finished...
    CriterionComparator cp(p);
    BOOST_CHECK_EQUAL( CriterionComparator(p), true) ;
    */

}