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

/*
 Copyright (C) 2017, 2018 Klaus Spanderen

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

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

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

#include "andreasenhugevolatilityinterpl.hpp"
#include "utilities.hpp"

#include <ql/math/comparison.hpp>
#include <ql/math/functional.hpp>
#include <ql/math/interpolations/sabrinterpolation.hpp>
#include <ql/math/optimization/bfgs.hpp>
#include <ql/math/optimization/simplex.hpp>
#include <ql/quotes/simplequote.hpp>
#include <ql/instruments/vanillaoption.hpp>
#include <ql/instruments/barrieroption.hpp>
#include <ql/time/daycounters/actual365fixed.hpp>
#include <ql/models/equity/hestonmodel.hpp>
#include <ql/pricingengines/blackformula.hpp>
#include <ql/processes/blackscholesprocess.hpp>
#include <ql/pricingengines/vanilla/analytichestonengine.hpp>
#include <ql/pricingengines/vanilla/fdblackscholesvanillaengine.hpp>
#include <ql/pricingengines/barrier/fdblackscholesbarrierengine.hpp>
#include <ql/termstructures/yield/zerocurve.hpp>
#include <ql/termstructures/volatility/equityfx/hestonblackvolsurface.hpp>
#include <ql/termstructures/volatility/equityfx/andreasenhugelocalvoladapter.hpp>
#include <ql/termstructures/volatility/equityfx/andreasenhugevolatilityinterpl.hpp>
#include <ql/termstructures/volatility/equityfx/andreasenhugevolatilityadapter.hpp>

#include <boost/make_shared.hpp>
#include <boost/assign/std/vector.hpp>
#include <boost/math/special_functions/fpclassify.hpp>

using namespace QuantLib;
using namespace boost::unit_test_framework;

namespace {

    struct CalibrationData {
        const Handle<Quote> spot;
        Handle<YieldTermStructure> rTS, qTS;
        AndreasenHugeVolatilityInterpl::CalibrationSet calibrationSet;
    };

    struct CalibrationResults {
        AndreasenHugeVolatilityInterpl::CalibrationType calibrationType;
        AndreasenHugeVolatilityInterpl::InterpolationType interpolationType;

        Real maxError, avgError;
        Real lvMaxError, lvAvgError;
    };

    CalibrationData AndreasenHugeExampleData() {
        // This is the example market data from the original paper
        // Andreasen J., Huge B., 2010. Volatility Interpolation
        // https://ssrn.com/abstract=1694972

        const Handle<Quote> spot(boost::make_shared<SimpleQuote>(2772.7));

        const Time maturityTimes[] = {
                  0.025, 0.101, 0.197, 0.274, 0.523, 0.772,
                  1.769, 2.267, 2.784, 3.781, 4.778, 5.774
            };

        const Real raw[][13] = {
            { 0.5131, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.3366, 0.3291, 0.0000, 0.0000 },
            { 0.5864, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.3178, 0.3129, 0.3008, 0.0000 },
            { 0.6597, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.3019, 0.2976, 0.2975, 0.0000 },
            { 0.7330, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.2863, 0.2848, 0.2848, 0.0000 },
            { 0.7697, 0.0000, 0.0000, 0.0000, 0.3262, 0.3079, 0.3001, 0.2843, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000 },
            { 0.8063, 0.0000, 0.0000, 0.0000, 0.3058, 0.2936, 0.2876, 0.2753, 0.2713, 0.2711, 0.2711, 0.2722, 0.2809 },
            { 0.8430, 0.0000, 0.0000, 0.0000, 0.2887, 0.2798, 0.2750, 0.2666, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000 },
            { 0.8613, 0.3365, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000 },
            { 0.8796, 0.3216, 0.2906, 0.2764, 0.2717, 0.2663, 0.2637, 0.2575, 0.2555, 0.2580, 0.2585, 0.2611, 0.2693 },
            { 0.8979, 0.3043, 0.2797, 0.2672, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000 },
            { 0.9163, 0.2880, 0.2690, 0.2578, 0.2557, 0.2531, 0.2519, 0.2497, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000 },
            { 0.9346, 0.2724, 0.2590, 0.2489, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000 },
            { 0.9529, 0.2586, 0.2488, 0.2405, 0.2407, 0.2404, 0.2411, 0.2418, 0.2410, 0.2448, 0.2469, 0.2501, 0.2584 },
            { 0.9712, 0.2466, 0.2390, 0.2329, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000 },
            { 0.9896, 0.2358, 0.2300, 0.2253, 0.2269, 0.2284, 0.2299, 0.2347, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000 },
            { 1.0079, 0.2247, 0.2213, 0.2184, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000 },
            { 1.0262, 0.2159, 0.2140, 0.2123, 0.2142, 0.2173, 0.2198, 0.2283, 0.2275, 0.2322, 0.2384, 0.2392, 0.2486 },
            { 1.0445, 0.2091, 0.2076, 0.2069, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000 },
            { 1.0629, 0.2056, 0.2024, 0.2025, 0.2039, 0.2074, 0.2104, 0.2213, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000 },
            { 1.0812, 0.2045, 0.1982, 0.1984, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000 },
            { 1.0995, 0.2025, 0.1959, 0.1944, 0.1962, 0.1988, 0.2022, 0.2151, 0.2161, 0.2219, 0.2269, 0.2305, 0.2399 },
            { 1.1178, 0.1933, 0.1929, 0.1920, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000 },
            { 1.1362, 0.0000, 0.0000, 0.0000, 0.1902, 0.1914, 0.1950, 0.2091, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000 },
            { 1.1728, 0.0000, 0.0000, 0.0000, 0.1885, 0.1854, 0.1888, 0.2039, 0.2058, 0.2122, 0.2186, 0.2223, 0.2321 },
            { 1.2095, 0.0000, 0.0000, 0.0000, 0.1867, 0.1811, 0.1839, 0.1990, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000 },
            { 1.2461, 0.0000, 0.0000, 0.0000, 0.1871, 0.1785, 0.1793, 0.1945, 0.0000, 0.2054, 0.2103, 0.2164, 0.2251 },
            { 1.3194, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.1988, 0.2054, 0.2105, 0.2190 },
            { 1.3927, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.1930, 0.2002, 0.2054, 0.2135 },
            { 1.4660, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.1849, 0.1964, 0.2012, 0.0000 }
        };

        const DayCounter dc = Actual365Fixed();
        const Date today = Date(1, March, 2010);

        const Handle<YieldTermStructure> rTS(flatRate(today, 0.0, dc));
        const Handle<YieldTermStructure> qTS(flatRate(today, 0.0, dc));

        const Size nStrikes = LENGTH(raw);
        const Size nMaturities = LENGTH(maturityTimes);

        QL_REQUIRE(nMaturities == LENGTH(raw[1])-1, "check raw data");

        AndreasenHugeVolatilityInterpl::CalibrationSet calibrationSet;

        calibrationSet.reserve(std::count_if(
            &raw[0][0], &raw[nStrikes-1][nMaturities]+1,
            std::bind2nd(std::not_equal_to<Real>(), 0.0)) - nStrikes);

        for (Size i=0; i < LENGTH(raw); ++i) {
            const Real strike = spot->value()*raw[i][0];

            for (Size j=1; j < LENGTH(raw[i]); ++j)
                if (raw[i][j] > QL_EPSILON) {
                    const Date maturity
                        = today + Period(Size(365*maturityTimes[j-1]), Days);

                    const Volatility impliedVol = raw[i][j];

                    calibrationSet.push_back(std::make_pair(
                        boost::make_shared<VanillaOption>(
                            boost::make_shared<PlainVanillaPayoff>(
                                (strike < spot->value())? Option::Put
                                                        : Option::Call,
                                strike),
                            boost::make_shared<EuropeanExercise>(maturity)),
                        boost::make_shared<SimpleQuote>(impliedVol))
                    );
                }
        }

        const CalibrationData data = { spot, rTS, qTS, calibrationSet};

        return data;
    }

    void testAndreasenHugeVolatilityInterpolation(
        const CalibrationData& data, const CalibrationResults& expected) {

        SavedSettings backup;

        const Handle<YieldTermStructure> rTS = data.rTS;
        const Handle<YieldTermStructure> qTS = data.qTS;

        const DayCounter dc = rTS->dayCounter();
        const Date today = rTS->referenceDate();
        Settings::instance().evaluationDate() = today;

        const Handle<Quote> spot = data.spot;

        AndreasenHugeVolatilityInterpl::CalibrationSet calibrationSet =
            data.calibrationSet;

        const boost::shared_ptr<AndreasenHugeVolatilityInterpl>
            andreasenHugeVolInterplation(
                boost::make_shared<AndreasenHugeVolatilityInterpl>(
                    calibrationSet, spot, rTS, qTS,
                    expected.interpolationType,
                    expected.calibrationType));

        const boost::tuple<Real, Real, Real> error =
            andreasenHugeVolInterplation->calibrationError();

        const Real maxError = error.get<1>();
        const Real avgError = error.get<2>();

        if (maxError > expected.maxError || avgError > expected.avgError) {
            BOOST_FAIL("Failed to reproduce calibration error"
                       << "\n    max calibration error:     " << maxError
                       << "\n    average calibration error: " << avgError
                       << "\n    expected max error:        " << expected.maxError
                       << "\n    expected average error:    " << expected.avgError);
        }

        const boost::shared_ptr<AndreasenHugeVolatilityAdapter> volatilityAdapter(
            boost::make_shared<AndreasenHugeVolatilityAdapter>(
                andreasenHugeVolInterplation));

        const boost::shared_ptr<AndreasenHugeLocalVolAdapter> localVolAdapter(
            boost::make_shared<AndreasenHugeLocalVolAdapter>(
                andreasenHugeVolInterplation));

        const boost::shared_ptr<GeneralizedBlackScholesProcess> localVolProcess(
            boost::make_shared<GeneralizedBlackScholesProcess>(
                spot, qTS, rTS,
                Handle<BlackVolTermStructure>(volatilityAdapter),
                Handle<LocalVolTermStructure>(localVolAdapter)));

        Real lvAvgError = 0.0, lvMaxError = 0.0;
        for (Size i=0, n=0; i < calibrationSet.size(); ++i) {

            const boost::shared_ptr<VanillaOption> option =
                calibrationSet[i].first;

            const boost::shared_ptr<PlainVanillaPayoff> payoff =
                boost::dynamic_pointer_cast<PlainVanillaPayoff>(
                    option->payoff());
            const Real strike = payoff->strike();
            const Option::Type optionType = payoff->optionType();

            const Time t = dc.yearFraction(today, option->exercise()->lastDate());

            const Volatility expectedVol = calibrationSet[i].second->value();
            const Volatility calculatedVol =
                volatilityAdapter->blackVol(t, strike, true);

            const Real diffVol = std::fabs(expectedVol - calculatedVol);
            const Real tol = std::max(1e-10, 1.01*maxError);

            if (diffVol > tol) {
                BOOST_FAIL("Failed to reproduce calibration option price"
                           << "\n    calculated: " << calculatedVol
                           << "\n    expected:   " << expectedVol
                           << "\n    difference: " << diffVol
                           << "\n    tolerance:  " << tol);
            }

            const boost::shared_ptr<PricingEngine> fdEngine(
                boost::make_shared<FdBlackScholesVanillaEngine>(
                    localVolProcess, std::max<Size>(30, Size(100*t)),
                    200, 0, FdmSchemeDesc::Douglas(), true));

            option->setPricingEngine(fdEngine);

            const DiscountFactor discount = rTS->discount(t);
            const Real fwd = spot->value()*qTS->discount(t)/discount;

            const Volatility lvImpliedVol = blackFormulaImpliedStdDevLiRS(
                optionType, strike, fwd, option->NPV(),
                discount, 0.0, Null<Real>(), 1.0, 1e-12)/std::sqrt(t);

            const Real lvError = std::fabs(lvImpliedVol - expectedVol);

            lvMaxError = std::max(lvError, lvMaxError);

            lvAvgError = (n*lvAvgError + lvError)/(n+1);

            ++n;
        }

        if (lvMaxError > expected.lvMaxError || avgError > expected.lvAvgError) {
            BOOST_FAIL("Failed to reproduce local volatility calibration error"
                       << "\n    max calibration error:     " << lvMaxError
                       << "\n    average calibration error: " << lvAvgError
                       << "\n    expected max error:        " << expected.lvMaxError
                       << "\n    expected average error:    " << expected.lvAvgError);
        }
    }

    CalibrationData BorovkovaExampleData() {
        // see Svetlana Borovkova, Ferry J. Permana
        // Implied volatility in oil markets
        // http://www.researchgate.net/publication/46493859_Implied_volatility_in_oil_markets

        const DayCounter dc = Actual365Fixed();
        const Date today = Date(4, January, 2018);

        const Handle<YieldTermStructure> rTS(flatRate(today, 0.025, dc));
        const Handle<YieldTermStructure> qTS(flatRate(today, 0.085, dc));

        Handle<Quote> spot(boost::make_shared<SimpleQuote>(100));

        const Real b1 = 0.35;
        const Real b2 = 0.03;
        const Real b3 = 0.005;
        const Real b4 = -0.02;
        const Real b5 = -0.005;

        const Real strikes[] = { 35, 50, 75, 100, 125, 150, 200, 300 };
        const Size maturityMonths[] = { 1, 3, 6, 9, 12, 15, 18, 24};

        AndreasenHugeVolatilityInterpl::CalibrationSet calibrationSet;

        for (Size i=0; i < LENGTH(strikes); ++i) {
            const Real strike = strikes[i];

            for (Size j=0; j < LENGTH(maturityMonths); ++j) {
                const Date maturityDate = today +
                    Period(maturityMonths[j], Months);
                const Time t = dc.yearFraction(today, maturityDate);

                const Real fwd = spot->value()*qTS->discount(t)/rTS->discount(t);
                const Real mn = std::log(fwd/strike)/std::sqrt(t);

                const Volatility vol = b1 + b2*mn + b3*mn*mn + b4*t + b5*mn*t;

                if (std::fabs(mn) < 3.71*vol) {

                    calibrationSet.push_back(std::make_pair(
                        boost::make_shared<VanillaOption>(
                            boost::make_shared<PlainVanillaPayoff>(
                                Option::Call, strike),
                            boost::make_shared<EuropeanExercise>(maturityDate)),
                        boost::make_shared<SimpleQuote>(vol)));
                }
            }
        }

        const CalibrationData data = { spot, rTS, qTS, calibrationSet};

        return data;
    }

    CalibrationData arbitrageData() {

        const DayCounter dc = Actual365Fixed();
        const Date today = Date(4, January, 2018);

        const Handle<YieldTermStructure> rTS(flatRate(today, 0.13, dc));
        const Handle<YieldTermStructure> qTS(flatRate(today, 0.03, dc));

        Handle<Quote> spot(boost::make_shared<SimpleQuote>(100));

        const Real strikes[] = { 100, 100, 100, 150 };
        const Size maturities[] = { 1, 3, 6, 6 };
        const Volatility vols[] = { 0.25, 0.35, 0.05, 0.35 };
        AndreasenHugeVolatilityInterpl::CalibrationSet calibrationSet;

        for (Size i=0; i < LENGTH(strikes); ++i) {
            const Real strike = strikes[i];
            const Date maturityDate = today + Period(maturities[i], Months);
            const Volatility vol = vols[i];

            calibrationSet.push_back(std::make_pair(
                boost::make_shared<VanillaOption>(
                    boost::make_shared<PlainVanillaPayoff>(
                        Option::Call, strike),
                    boost::make_shared<EuropeanExercise>(maturityDate)),
                boost::make_shared<SimpleQuote>(vol)));
        }

        const CalibrationData data = { spot, rTS, qTS, calibrationSet};

        return data;
    }
}

void AndreasenHugeVolatilityInterplTest::testAndreasenHugePut() {

    BOOST_TEST_MESSAGE(
        "Testing Andreasen-Huge example with Put calibration...");

    const CalibrationData data = AndreasenHugeExampleData();

    const CalibrationResults expected = {
        AndreasenHugeVolatilityInterpl::Put,
        AndreasenHugeVolatilityInterpl::CubicSpline,
        0.0015, 0.00035,
        0.0020, 0.00035
    };

    testAndreasenHugeVolatilityInterpolation(data, expected);
}

void AndreasenHugeVolatilityInterplTest::testAndreasenHugeCall() {

    BOOST_TEST_MESSAGE(
        "Testing Andreasen-Huge example with Call calibration...");

    const CalibrationData data = AndreasenHugeExampleData();

    const CalibrationResults expected = {
        AndreasenHugeVolatilityInterpl::Call,
        AndreasenHugeVolatilityInterpl::CubicSpline,
        0.0015, 0.00035,
        0.0015, 0.00035
    };

    testAndreasenHugeVolatilityInterpolation(data, expected);
}

void AndreasenHugeVolatilityInterplTest::testAndreasenHugeCallPut() {

    BOOST_TEST_MESSAGE(
        "Testing Andreasen-Huge example with instantaneous "
         "Call and Put calibration...");

    const CalibrationData data = AndreasenHugeExampleData();

    const CalibrationResults expected = {
        AndreasenHugeVolatilityInterpl::CallPut,
        AndreasenHugeVolatilityInterpl::CubicSpline,
        0.0015, 0.00035,
        0.0015, 0.00035
    };

    testAndreasenHugeVolatilityInterpolation(data, expected);
}

void AndreasenHugeVolatilityInterplTest::testLinearInterpolation() {
    BOOST_TEST_MESSAGE(
        "Testing Andreasen-Huge example with linear interpolation...");

    const CalibrationData data = AndreasenHugeExampleData();

    const CalibrationResults expected = {
        AndreasenHugeVolatilityInterpl::CallPut,
        AndreasenHugeVolatilityInterpl::Linear,
        0.0020, 0.00015,
        0.0040, 0.00035
    };

    testAndreasenHugeVolatilityInterpolation(data, expected);
}

void AndreasenHugeVolatilityInterplTest::testPiecewiseConstantInterpolation() {
    BOOST_TEST_MESSAGE(
        "Testing Andreasen-Huge example with piecewise constant interpolation...");

    const CalibrationData data = AndreasenHugeExampleData();

    const CalibrationResults expected = {
        AndreasenHugeVolatilityInterpl::CallPut,
        AndreasenHugeVolatilityInterpl::PiecewiseConstant,
        0.0025, 0.00025,
        0.0040, 0.00035
    };

    testAndreasenHugeVolatilityInterpolation(data, expected);
}

void AndreasenHugeVolatilityInterplTest::testTimeDependentInterestRates() {

    BOOST_TEST_MESSAGE(
        "Testing Andreasen-Huge volatility interpolation with "
        "time dependent interest rates and dividend yield...");

    SavedSettings backup;

    const CalibrationData data = AndreasenHugeExampleData();

    const DayCounter dc = data.rTS->dayCounter();
    const Date today = data.rTS->referenceDate();
    Settings::instance().evaluationDate() = today;

    using namespace boost::assign;
    std::vector<Real> r, q;
    r += 0.0167, 0.023, 0.03234, 0.034, 0.038, 0.042, 0.047, 0.053;
    q += 0.01, 0.011, 0.013, 0.014, 0.02, 0.025, 0.067, 0.072;

    std::vector<Date> dates;
    dates += today,                     today + Period(41, Days),
             today + Period(75, Days),  today + Period(165, Days),
             today + Period(256, Days), today + Period(345, Days),
             today + Period(524, Days), today + Period(2190, Days);

    const Handle<YieldTermStructure> rTS(
        boost::make_shared<ZeroCurve>(dates, r, dc));
    const Handle<YieldTermStructure> qTS(
        boost::make_shared<ZeroCurve>(dates, q, dc));

    const CalibrationData origData = AndreasenHugeExampleData();
    AndreasenHugeVolatilityInterpl::CalibrationSet calibrationSet =
        origData.calibrationSet;

    const Handle<Quote> spot = origData.spot;

    const boost::shared_ptr<HestonModel> hestonModel(
        boost::make_shared<HestonModel>(
            boost::make_shared<HestonProcess>(
                rTS, qTS, spot, 0.09, 2.0, 0.09, 0.4, -0.75)));

    const boost::shared_ptr<PricingEngine> hestonEngine(
        boost::make_shared<AnalyticHestonEngine>(
            hestonModel,
            AnalyticHestonEngine::AndersenPiterbarg,
            AnalyticHestonEngine::Integration::discreteTrapezoid(128)));

    for (Size i=0; i < calibrationSet.size(); ++i) {
        const boost::shared_ptr<VanillaOption> option =
            calibrationSet[i].first;

        const boost::shared_ptr<PlainVanillaPayoff> payoff =
            boost::dynamic_pointer_cast<PlainVanillaPayoff>(option->payoff());

        const Real strike = payoff->strike();
        const Option::Type optionType = payoff->optionType();

        const Time t = dc.yearFraction(today, option->exercise()->lastDate());

        const DiscountFactor discount = rTS->discount(t);
        const Real fwd = spot->value()*qTS->discount(t)/discount;

        option->setPricingEngine(hestonEngine);
        const Real npv = option->NPV();

        const Volatility impliedVol = blackFormulaImpliedStdDevLiRS(
            optionType, strike, fwd, npv,
            discount, 0.0, Null<Real>(), 1.0, 1e-12)/std::sqrt(t);

        calibrationSet[i].second = boost::make_shared<SimpleQuote>(impliedVol);
    }

    CalibrationData irData = { spot, rTS, qTS, calibrationSet };

    const CalibrationResults expected = {
        AndreasenHugeVolatilityInterpl::CallPut,
        AndreasenHugeVolatilityInterpl::CubicSpline,
        0.0020, 0.0003,
        0.0020, 0.0004
    };

    testAndreasenHugeVolatilityInterpolation(irData, expected);
}

void AndreasenHugeVolatilityInterplTest::testSingleOptionCalibration() {
    BOOST_TEST_MESSAGE(
        "Testing Andreasen-Huge volatility interpolation with "
        "a single option...");

    SavedSettings backup;

    const DayCounter dc = Actual365Fixed();
    const Date today = Date(4, January, 2018);

    const Handle<YieldTermStructure> rTS(flatRate(today, 0.025, dc));
    const Handle<YieldTermStructure> qTS(flatRate(today, 0.085, dc));

    AndreasenHugeVolatilityInterpl::CalibrationSet calibrationSet;

    const Real strike = 10.0;
    const Volatility vol = 0.3;
    const Date maturity = today + Period(1, Years);
    Handle<Quote> spot(boost::make_shared<SimpleQuote>(strike));

    calibrationSet.push_back(std::make_pair(
        boost::make_shared<VanillaOption>(
            boost::make_shared<PlainVanillaPayoff>(Option::Call, strike),
            boost::make_shared<EuropeanExercise>(maturity)),
        boost::make_shared<SimpleQuote>(vol)));

    const AndreasenHugeVolatilityInterpl::InterpolationType interpl[] = {
        AndreasenHugeVolatilityInterpl::Linear,
        AndreasenHugeVolatilityInterpl::CubicSpline,
        AndreasenHugeVolatilityInterpl::PiecewiseConstant
    };

    const AndreasenHugeVolatilityInterpl::CalibrationType calibrationType[] = {
        AndreasenHugeVolatilityInterpl::Put,
        AndreasenHugeVolatilityInterpl::Call,
        AndreasenHugeVolatilityInterpl::CallPut
    };

    for (Size i=0; i < LENGTH(interpl); ++i)
        for (Size j=0; j < LENGTH(calibrationType); ++j) {
            const boost::shared_ptr<AndreasenHugeVolatilityInterpl>
                andreasenHugeVolInterplation(
                    boost::make_shared<AndreasenHugeVolatilityInterpl>(
                        calibrationSet, spot, rTS, qTS,
                        interpl[i], calibrationType[j], 50));

            const boost::shared_ptr<AndreasenHugeVolatilityAdapter>
                volatilityAdapter =
                    boost::make_shared<AndreasenHugeVolatilityAdapter>(
                        andreasenHugeVolInterplation);

            const Volatility calculated =
                volatilityAdapter->blackVol(maturity, strike);
            const Volatility expected = vol;

            if (std::fabs(calculated - expected) > 1e-4) {
                BOOST_FAIL("Failed to reproduce single option calibration"
                           << "\n    calculated: " << calculated
                           << "\n    expected:   " << expected);
            }
        }
}

void AndreasenHugeVolatilityInterplTest::testArbitrageFree() {
    BOOST_TEST_MESSAGE(
        "Testing Andreasen-Huge volatility interpolation gives "
        "arbitrage free prices...");

    SavedSettings backup;

    CalibrationData data[] = { BorovkovaExampleData(), arbitrageData() };;

    for (Size i=0; i < LENGTH(data); ++i) {
        const Handle<Quote>& spot = data[i].spot;
        const AndreasenHugeVolatilityInterpl::CalibrationSet& calibrationSet =
            data[i].calibrationSet;

        const Handle<YieldTermStructure>& rTS = data[i].rTS;
        const Handle<YieldTermStructure>& qTS = data[i].qTS;

        const DayCounter dc = rTS->dayCounter();
        const Date today = rTS->referenceDate();

        const boost::shared_ptr<AndreasenHugeVolatilityInterpl>
            andreasenHugeVolInterplation(
                boost::make_shared<AndreasenHugeVolatilityInterpl>(
                    calibrationSet, spot, rTS, qTS,
                    AndreasenHugeVolatilityInterpl::CubicSpline,
                    AndreasenHugeVolatilityInterpl::CallPut, 5000));

        const boost::shared_ptr<AndreasenHugeVolatilityAdapter> volatilityAdapter(
            boost::make_shared<AndreasenHugeVolatilityAdapter>(
                andreasenHugeVolInterplation));

        for (Real m = -0.7; m < 0.7; m+=0.05) {

            for (Size weeks=6; weeks < 52; ++weeks) {
                const Date maturityDate = today + Period(weeks, Weeks);

                const Time t = dc.yearFraction(today, maturityDate);

                const Real fwd = spot->value()*qTS->discount(t)/rTS->discount(t);

                // J. Gatheral, Arbitrage-free SVI volatility surfaces
                // http://mfe.baruch.cuny.edu/wp-content/uploads/2013/01/OsakaSVI2012.pdf
                const Real eps = 0.025;
                const Real k  = fwd*std::exp(m);
                const Real km = fwd*std::exp(m - eps);
                const Real kp = fwd*std::exp(m + eps);

                const Real w =
                    volatilityAdapter->blackVariance(t, k, true);
                const Real w_p =
                    volatilityAdapter->blackVariance(t, kp, true);
                const Real w_m =
                    volatilityAdapter->blackVariance(t, km, true);

                const Real w1 = (w_p - w_m)/(2*eps);
                const Real w2 = (w_p + w_m - 2*w)/(eps*eps);

                const Real g_k = square<Real>()(1-m*w1/(2*w))
                    - w1*w1/4*(1/w + 0.25) + 0.5*w2;

                if (g_k < 0) {
                    BOOST_FAIL("No-arbitrage condition g_k >= 0 failed"
                               << "\n    strike:  " << k
                               << "\n    forward: " << fwd
                               << "\n    time:    " << t
                               << "\n    g_k:    " << g_k);
                }

                const Real deltaT = 1.0/365.;
                const Real fwdpt = spot->value()*
                    qTS->discount(t+deltaT)/rTS->discount(t+deltaT);

                const Real kpt = fwdpt*std::exp(m);
                const Real w_pt =
                    volatilityAdapter->blackVariance(t+deltaT, kpt, true);

                const Real w_t = (w_pt - w)/deltaT;
                if (w_t < -1e-8) {
                    BOOST_FAIL("No-arbitrage condition w_t >= 0 failed"
                               << "\n    strike:  " << k
                               << "\n    forward: " << fwd
                               << "\n    time:    " << t
                               << "\n    w        " << w
                               << "\n    w_t:     " << w_t);
                }
            }
        }
    }
}

void AndreasenHugeVolatilityInterplTest::testBarrierOptionPricing() {
    BOOST_TEST_MESSAGE(
        "Testing Barrier option pricing with Andreasen-Huge "
         "local volatility surface...");

    SavedSettings backup;

    const DayCounter dc = Actual365Fixed();
    const Date today = Date(4, January, 2018);

    const Handle<YieldTermStructure> rTS(flatRate(today, 0.01, dc));
    const Handle<YieldTermStructure> qTS(flatRate(today, 0.03, dc));

    Handle<Quote> spot(boost::make_shared<SimpleQuote>(100));
    const boost::shared_ptr<HestonModel> hestonModel(
        boost::make_shared<HestonModel>(
            boost::make_shared<HestonProcess>(
                rTS, qTS, spot, 0.04, 2.0, 0.04, 0.4, -0.75)));

    const boost::shared_ptr<BlackVolTermStructure> hestonVol =
        boost::make_shared<HestonBlackVolSurface>(
            Handle<HestonModel>(hestonModel));

    const boost::shared_ptr<GeneralizedBlackScholesProcess>
        dupireLocalVolProcess =
            boost::make_shared<GeneralizedBlackScholesProcess>(
                spot, qTS, rTS, Handle<BlackVolTermStructure>(hestonVol));

    const Real strikes[] = { 25, 50, 75, 90, 100, 110, 125, 150, 200, 400};
    const Size maturityMonths[] = { 1, 3, 6, 9, 12};

    AndreasenHugeVolatilityInterpl::CalibrationSet calibrationSet;

    for (Size i=0; i < LENGTH(strikes); ++i) {
        const Real strike = strikes[i];

        for (Size j=0; j < LENGTH(maturityMonths); ++j) {
            const Date maturityDate = today + Period(maturityMonths[j], Months);
            const Time t = dc.yearFraction(today, maturityDate);

            const Volatility vol = hestonVol->blackVol(t, strike);

            const Real mn = std::log(spot->value()/strike)/std::sqrt(t);

            if (std::fabs(mn) < 3.07*vol) {
                calibrationSet.push_back(std::make_pair(
                    boost::make_shared<VanillaOption>(
                        boost::make_shared<PlainVanillaPayoff>(
                            Option::Call, strike),
                        boost::make_shared<EuropeanExercise>(maturityDate)),
                    boost::make_shared<SimpleQuote>(vol)));
            }
        }
    }

    const boost::shared_ptr<AndreasenHugeVolatilityInterpl>
        andreasenHugeVolInterplation(
            boost::make_shared<AndreasenHugeVolatilityInterpl>(
                calibrationSet, spot, rTS, qTS));

    const boost::shared_ptr<AndreasenHugeLocalVolAdapter> localVolAdapter(
        boost::make_shared<AndreasenHugeLocalVolAdapter>(
            andreasenHugeVolInterplation));

    const boost::shared_ptr<GeneralizedBlackScholesProcess>
        andreasenHugeLocalVolProcess =
            boost::make_shared<GeneralizedBlackScholesProcess>(
                spot, qTS, rTS,
                Handle<BlackVolTermStructure>(hestonVol),
                Handle<LocalVolTermStructure>(localVolAdapter));

    const Real strike = 120.0;
    const Real barrier=  80.0;
    const Real rebate =   0.0;
    const Date maturity = today + Period(1, Years);
    const Barrier::Type barrierType = Barrier::DownOut;

    BarrierOption barrierOption(barrierType, barrier, rebate,
        boost::make_shared<PlainVanillaPayoff>(Option::Put, strike),
        boost::make_shared<EuropeanExercise>(maturity));

    barrierOption.setPricingEngine(
        boost::make_shared<FdBlackScholesBarrierEngine>(
            dupireLocalVolProcess, 50, 100, 0,
            FdmSchemeDesc::Douglas(), true, 0.2));

    const Real dupireNPV = barrierOption.NPV();

    barrierOption.setPricingEngine(
        boost::make_shared<FdBlackScholesBarrierEngine>(
            andreasenHugeLocalVolProcess, 200, 400, 0,
            FdmSchemeDesc::Douglas(), true, 0.25));

    const Real andreasenHugeNPV = barrierOption.NPV();

    const Real tol = 0.15;
    const Real diff = std::fabs(andreasenHugeNPV - dupireNPV);

    if (diff > tol) {
        BOOST_FAIL("failed to reproduce barrier prices with Andreasen-Huge "
                "local volatility surface"
                   << "\n    Andreasen-Huge price: " << andreasenHugeNPV
                   << "\n    Dupire formula price: " << dupireNPV
                   << "\n    diff:                 " << diff
                   << "\n    tolerance:            " << tol);
    }
}

namespace {
    std::pair<CalibrationData, std::vector<Real> > sabrData() {

        const DayCounter dc = Actual365Fixed();
        const Date today = Date(4, January, 2018);

        const Real alpha = 0.15;
        const Real beta = 0.8;
        const Real nu = 0.5;
        const Real rho = -0.48;
        const Real forward = 0.03;
        const Size maturityInYears = 20;

        const Date maturityDate = today + Period(maturityInYears, Years);
        const Time maturity = dc.yearFraction(today, maturityDate);

        AndreasenHugeVolatilityInterpl::CalibrationSet calibrationSet;

        const Real strikes[] = { 0.02, 0.025, 0.03, 0.035, 0.04, 0.05, 0.06 };

        for (Size i=0; i < LENGTH(strikes); ++i) {
            const Real strike = strikes[i];
            const Volatility vol = sabrVolatility(
               strike, forward, maturity, alpha, beta, nu, rho);

            calibrationSet.push_back(std::make_pair(
                boost::make_shared<VanillaOption>(
                    boost::make_shared<PlainVanillaPayoff>(
                        Option::Call, strike),
                    boost::make_shared<EuropeanExercise>(maturityDate)),
                boost::make_shared<SimpleQuote>(vol)));
        }

        const Handle<YieldTermStructure> rTS(flatRate(today, forward, dc));
        const Handle<YieldTermStructure> qTS(flatRate(today, forward, dc));

        Handle<Quote> spot(boost::make_shared<SimpleQuote>(forward));

        const CalibrationData data = { spot, rTS, qTS, calibrationSet};

        using namespace boost::assign;
        std::vector<Real> parameter;
        parameter += alpha, beta, nu, rho, forward, maturity;

        return std::make_pair(data, parameter);
    }
}

void AndreasenHugeVolatilityInterplTest::testPeterAndFabiensExample() {
    BOOST_TEST_MESSAGE(
        "Testing Peter's and Fabien's SABR example...");

    // http://chasethedevil.github.io/post/andreasen-huge-extrapolation/

    SavedSettings backup;

    const std::pair<CalibrationData, std::vector<Real> > sd = sabrData();
    const CalibrationData& data = sd.first;
    const std::vector<Real>& parameter = sd.second;

    const boost::shared_ptr<AndreasenHugeVolatilityInterpl>
        andreasenHugeVolInterplation(
            boost::make_shared<AndreasenHugeVolatilityInterpl>(
                data.calibrationSet, data.spot, data.rTS, data.qTS));

    const boost::shared_ptr<AndreasenHugeVolatilityAdapter> volAdapter(
        boost::make_shared<AndreasenHugeVolatilityAdapter>(
            andreasenHugeVolInterplation));

    const Real alpha    = parameter[0];
    const Real beta     = parameter[1];
    const Real nu       = parameter[2];
    const Real rho      = parameter[3];
    const Real forward  = parameter[4];
    const Time maturity = parameter[5];

    for (Real strike = 0.02; strike < 0.06; strike+=0.001) {
        const Volatility sabrVol = sabrVolatility(
           strike, forward, maturity, alpha, beta, nu, rho);

        const Volatility ahVol = volAdapter->blackVol(maturity, strike, true);

        const Real tol = 0.0005;
        const Real diff = std::fabs(sabrVol - ahVol);

        if (boost::math::isnan(ahVol) || diff > 0.005) {
            BOOST_FAIL("failed to reproduce SABR volatility with "
                    "Andreasen-Huge interpolation"
                   << "\n    Andreasen-Huge vol: " << ahVol
                   << "\n    SABR volatility:    " << sabrVol
                   << "\n    diff:               " << diff
                   << "\n    tolerance:          " << tol);
        }
    }
}

void AndreasenHugeVolatilityInterplTest::testDifferentOptimizers() {
    BOOST_TEST_MESSAGE(
        "Testing different optimizer for Andreasen-Huge "
        "volatility interpolation...");

    const CalibrationData& data = sabrData().first;

    const boost::shared_ptr<OptimizationMethod> optimizationMethods[] = {
        boost::make_shared<LevenbergMarquardt>(),
        boost::make_shared<BFGS>(),
        boost::make_shared<Simplex>(0.2)
    };

    for (Size i=0; i < LENGTH(optimizationMethods); ++i) {
        const boost::shared_ptr<OptimizationMethod> optimizationMethod =
            optimizationMethods[i];

        const Real avgError = AndreasenHugeVolatilityInterpl(
            data.calibrationSet,
            data.spot,
            data.rTS, data.qTS,
            AndreasenHugeVolatilityInterpl::CubicSpline,
            AndreasenHugeVolatilityInterpl::Call,
            400,
            Null<Real>(), Null<Real>(),
            optimizationMethod).calibrationError().get<2>();

        if (boost::math::isnan(avgError) || avgError > 0.0001) {
            BOOST_FAIL("failed to calibrate Andreasen-Huge "
                    "volatility interpolation with different optimizera"
                   << "\n    calibration error: " << avgError);
        }
    }
}

void AndreasenHugeVolatilityInterplTest::testMovingReferenceDate() {
    BOOST_TEST_MESSAGE(
        "Testing that reference date of adapter surface moves along with "
        "evaluation date...");

    SavedSettings backup;

    const Date today = Date(4, January, 2018);
    Settings::instance().evaluationDate() = today;

    const DayCounter dc = Actual365Fixed();
    const Date maturity = today + Period(1, Months);

    Handle<YieldTermStructure> ts(flatRate(0.04, dc));

    const Real s0 = 100.0;
    const Volatility impliedVol = 0.2;
    const Handle<Quote> spot(boost::make_shared<SimpleQuote>(s0));

    AndreasenHugeVolatilityInterpl::CalibrationSet calibrationSet(
        1,
        std::make_pair(
            boost::make_shared<VanillaOption>(
                boost::make_shared<PlainVanillaPayoff>(Option::Call, s0),
                boost::make_shared<EuropeanExercise>(maturity)),
            boost::make_shared<SimpleQuote>(impliedVol))
    );

    const boost::shared_ptr<AndreasenHugeVolatilityInterpl>
        andreasenHugeVolInterplation(
            boost::make_shared<AndreasenHugeVolatilityInterpl>(
                calibrationSet, spot, ts, ts));

    const Real tol = 1e-8;
    const boost::shared_ptr<AndreasenHugeVolatilityAdapter> volatilityAdapter(
        boost::make_shared<AndreasenHugeVolatilityAdapter>(
            andreasenHugeVolInterplation, tol));

    const boost::shared_ptr<AndreasenHugeLocalVolAdapter> localVolAdapter(
        boost::make_shared<AndreasenHugeLocalVolAdapter>(
            andreasenHugeVolInterplation));

    const Date volRefDate = volatilityAdapter->referenceDate();
    const Date localRefDate = localVolAdapter->referenceDate();

    if (volRefDate != today || localRefDate != today)
        BOOST_FAIL("reference dates should match today's date"
               << "\n    today                     : " << today
               << "\n    local vol reference date  : " << localRefDate
               << "\n    implied vol reference date: " << volRefDate);

    const Date modToday = Date(15, January, 2018);
    Settings::instance().evaluationDate() = modToday;

    const Date modVolRefDate = volatilityAdapter->referenceDate();
    const Date modLocalRefDate = localVolAdapter->referenceDate();

    if (modVolRefDate != modToday || modLocalRefDate != modToday)
        BOOST_FAIL("reference dates should match modified today's date"
               << "\n    today                     : " << modToday
               << "\n    local vol reference date  : " << modLocalRefDate
               << "\n    implied vol reference date: " << modVolRefDate);

    // test update method
    const Volatility modImpliedVol =
        volatilityAdapter->blackVol(maturity, s0, true);

    const Real diff = std::fabs(modImpliedVol - impliedVol);
    if (diff > 10*tol)
        BOOST_FAIL("modified implied vol should match direct calculation"
                << "\n    implied vol         : " << impliedVol
                << "\n    modified implied vol: " << modImpliedVol
                << "\n    difference          : " << diff
                << "\n    tolerance           : " << tol);
}

test_suite* AndreasenHugeVolatilityInterplTest::suite(SpeedLevel speed) {
    test_suite* suite =
        BOOST_TEST_SUITE("Andreasen-Huge volatility interpolation tests");

    suite->add(QUANTLIB_TEST_CASE(
        &AndreasenHugeVolatilityInterplTest::testSingleOptionCalibration));
    suite->add(QUANTLIB_TEST_CASE(
        &AndreasenHugeVolatilityInterplTest::testArbitrageFree));
    suite->add(QUANTLIB_TEST_CASE(
        &AndreasenHugeVolatilityInterplTest::testPeterAndFabiensExample));
    suite->add(QUANTLIB_TEST_CASE(
        &AndreasenHugeVolatilityInterplTest::testDifferentOptimizers));
    suite->add(QUANTLIB_TEST_CASE(
        &AndreasenHugeVolatilityInterplTest::testMovingReferenceDate));

    if (speed == Slow) {
        suite->add(QUANTLIB_TEST_CASE(
            &AndreasenHugeVolatilityInterplTest::testAndreasenHugePut));
        suite->add(QUANTLIB_TEST_CASE(
            &AndreasenHugeVolatilityInterplTest::testAndreasenHugeCall));
        suite->add(QUANTLIB_TEST_CASE(
            &AndreasenHugeVolatilityInterplTest::testAndreasenHugeCallPut));
        suite->add(QUANTLIB_TEST_CASE(
            &AndreasenHugeVolatilityInterplTest::testLinearInterpolation));
        suite->add(QUANTLIB_TEST_CASE(
            &AndreasenHugeVolatilityInterplTest::testPiecewiseConstantInterpolation));
        suite->add(QUANTLIB_TEST_CASE(
            &AndreasenHugeVolatilityInterplTest::testBarrierOptionPricing));
        suite->add(QUANTLIB_TEST_CASE(
            &AndreasenHugeVolatilityInterplTest::testTimeDependentInterestRates));
    }
    return suite;
}


//==================================================================================================
/**
  EVE - Expressive Vector Engine
  Copyright : EVE Contributors & Maintainers
  SPDX-License-Identifier: MIT
**/
//==================================================================================================
#include "test.hpp"
#include <eve/module/core.hpp>
#include <eve/module/core.hpp>
#include <eve/module/math.hpp>
#include <eve/module/elliptic.hpp>
#include <boost/math/special_functions/ellint_rg.hpp>

//==================================================================================================
// Types tests
//==================================================================================================
EVE_TEST_TYPES( "Check return types of ellint_rg"
            , eve::test::simd::ieee_reals
            )
<typename T>(eve::as<T>)
{
  using v_t = eve::element_type_t<T>;

  TTS_EXPR_IS( eve::ellint_rg(T(), T(), T())  , T);
  TTS_EXPR_IS( eve::ellint_rg(v_t(), v_t(), v_t()), v_t);
  TTS_EXPR_IS( eve::ellint_rg(T(), v_t(), v_t()), T);
  TTS_EXPR_IS( eve::ellint_rg(v_t(), T(), v_t()), T);
  TTS_EXPR_IS( eve::ellint_rg(v_t(), v_t(), T()), T);
  TTS_EXPR_IS( eve::ellint_rg(T(), T(), v_t()), T);
  TTS_EXPR_IS( eve::ellint_rg(v_t(), T(), T()), T);
  TTS_EXPR_IS( eve::ellint_rg(T(), v_t(), T()), T);
};

//==================================================================================================
// ellint_rg  tests
//==================================================================================================
EVE_TEST( "Check behavior of ellint_rg on wide"
        , eve::test::simd::ieee_reals
        , eve::test::generate( eve::test::randoms(0, 100.0)
                             , eve::test::randoms(0, 100.0)
                             , eve::test::randoms(0, 100.0))
        )
<typename T>(T const& x, T const& y, T const& z)
{
  using eve::detail::map;
  using v_t = eve::element_type_t<T>;

  TTS_ULP_EQUAL(eve::ellint_rg(x, y, z) , map([](auto e, auto f, auto g) -> v_t { return boost::math::ellint_rg(e, f, g); }, x, y, z), 11);
};


/*******************************************************************************
Copyright(c) 2015-2021 Parker Hannifin Corp. All rights reserved.

MIT Licensed. See the included LICENSE.txt for a copy of the full MIT License.
*******************************************************************************/
#include "mscl/MicroStrain/Wireless/SyncSamplingFormulas.h"
#include "mscl/MicroStrain/SampleRate.h"

#include <boost/test/unit_test.hpp>

using namespace mscl;

BOOST_AUTO_TEST_SUITE(SyncSamplingFormulas_Test)

BOOST_AUTO_TEST_CASE(SyncSamplingFormulas_bytesPerSecond)
{
    SampleRate hz_1 = SampleRate::Hertz(1);
    SampleRate hz_256 = SampleRate::Hertz(256);
    SampleRate sec_10 = SampleRate::Seconds(10);

    BOOST_CHECK_CLOSE(SyncSamplingFormulas::bytesPerSecond(hz_1, 4, 2), 8.0, 0.1);
    BOOST_CHECK_CLOSE(SyncSamplingFormulas::bytesPerSecond(hz_1, 4, 4), 16.0, 0.1);

    BOOST_CHECK_CLOSE(SyncSamplingFormulas::bytesPerSecond(hz_256, 4, 2), 2048.0, 0.1);

    BOOST_CHECK_CLOSE(SyncSamplingFormulas::bytesPerSecond(sec_10, 4, 2), 0.8, 0.1);
}

BOOST_AUTO_TEST_CASE(SyncSamplingFormulas_maxBytesPerPacket)
{
    SampleRate hz_256 = SampleRate::Hertz(256);
    SampleRate hz_512 = SampleRate::Hertz(512);
    SampleRate sec_10 = SampleRate::Seconds(10);
    WirelessTypes::CommProtocol lxrs = WirelessTypes::commProtocol_lxrs;

    BOOST_CHECK_EQUAL(SyncSamplingFormulas::maxBytesPerPacket(hz_256, true, true, 1, lxrs), 32);
    BOOST_CHECK_EQUAL(SyncSamplingFormulas::maxBytesPerPacket(hz_256, true, false, 1, lxrs), 64);
    BOOST_CHECK_EQUAL(SyncSamplingFormulas::maxBytesPerPacket(hz_256, false, false, 1, lxrs), 96);
    BOOST_CHECK_EQUAL(SyncSamplingFormulas::maxBytesPerPacket(hz_512, true, true, 1, lxrs), 32);
    BOOST_CHECK_EQUAL(SyncSamplingFormulas::maxBytesPerPacket(hz_512, true, false, 1, lxrs), 48);
    BOOST_CHECK_EQUAL(SyncSamplingFormulas::maxBytesPerPacket(sec_10, true, false, 1, lxrs), 64);
}

BOOST_AUTO_TEST_CASE(SyncSamplingFormulas_txPerGroup)
{
    BOOST_CHECK_EQUAL(SyncSamplingFormulas::txPerGroup(128.0, 200, 1), 1);
    BOOST_CHECK_EQUAL(SyncSamplingFormulas::txPerGroup(128.0, 400, 2), 1);
    BOOST_CHECK_EQUAL(SyncSamplingFormulas::txPerGroup(2.0, 400, 16), 1);
    BOOST_CHECK_EQUAL(SyncSamplingFormulas::txPerGroup(128.0, 2, 1), 64);

    BOOST_CHECK_EQUAL(SyncSamplingFormulas::txPerGroup(2048.0, 64, 1), 32);
}

BOOST_AUTO_TEST_CASE(SyncSamplingFormulas_txPerSecond)
{
    BOOST_CHECK_CLOSE(SyncSamplingFormulas::txPerSecond(1, 0), 0.0, 0.1);
    BOOST_CHECK_CLOSE(SyncSamplingFormulas::txPerSecond(4, 1), 4.0, 0.1);
    BOOST_CHECK_CLOSE(SyncSamplingFormulas::txPerSecond(16, 2), 8.0, 0.1);
}

BOOST_AUTO_TEST_CASE(SyncSamplingFormulas_slotSpacing)
{
    BOOST_CHECK_EQUAL(SyncSamplingFormulas::slotSpacing(WirelessTypes::commProtocol_lxrs), 8);
    BOOST_CHECK_EQUAL(SyncSamplingFormulas::slotSpacing(WirelessTypes::commProtocol_lxrsPlus), 4);
}

BOOST_AUTO_TEST_CASE(SyncSamplingFormulas_maxTdmaAddress)
{
    BOOST_CHECK_EQUAL(SyncSamplingFormulas::maxTdmaAddress(4, 1, true, WirelessTypes::commProtocol_lxrs), 248);
    BOOST_CHECK_EQUAL(SyncSamplingFormulas::maxTdmaAddress(16, 2, true, WirelessTypes::commProtocol_lxrs), 120);
    BOOST_CHECK_EQUAL(SyncSamplingFormulas::maxTdmaAddress(128, 1, true, WirelessTypes::commProtocol_lxrs), 1);
    BOOST_CHECK_EQUAL(SyncSamplingFormulas::maxTdmaAddress(128, 1, false, WirelessTypes::commProtocol_lxrs), 8);
}

BOOST_AUTO_TEST_CASE(SyncSamplingFormulas_percentBandwidth)
{
    BOOST_CHECK_CLOSE(SyncSamplingFormulas::percentBandwidth(0.0f, true, WirelessTypes::commProtocol_lxrs), 0.0, 0.1);
    BOOST_CHECK_CLOSE(SyncSamplingFormulas::percentBandwidth(4.0f, true, WirelessTypes::commProtocol_lxrs), 3.128, 0.1);
    BOOST_CHECK_CLOSE(SyncSamplingFormulas::percentBandwidth(8.0f, true, WirelessTypes::commProtocol_lxrs), 6.256, 0.1);
    BOOST_CHECK_CLOSE(SyncSamplingFormulas::percentBandwidth(128.0f, true, WirelessTypes::commProtocol_lxrs), 100.097, 0.1);
    BOOST_CHECK_CLOSE(SyncSamplingFormulas::percentBandwidth(128.0f, false, WirelessTypes::commProtocol_lxrs), 100.0, 0.1);
}

BOOST_AUTO_TEST_CASE(SyncSamplingFormulas_sampleDuration)
{
    SampleRate hz_1 = SampleRate::Hertz(1);
    SampleRate hz_256 = SampleRate::Hertz(256);
    SampleRate sec_10 = SampleRate::Seconds(10);

    BOOST_CHECK_CLOSE(SyncSamplingFormulas::sampleDuration(200, hz_1), 200, 0.1);
    BOOST_CHECK_CLOSE(SyncSamplingFormulas::sampleDuration(400, hz_256), 1.5625, 0.1);
    BOOST_CHECK_CLOSE(SyncSamplingFormulas::sampleDuration(800, sec_10), 8000, 0.1);
}

BOOST_AUTO_TEST_CASE(SyncSamplingFormulas_bytesPerSweep)
{
    BOOST_CHECK_EQUAL(SyncSamplingFormulas::bytesPerSweep(2, 4), 8);
    BOOST_CHECK_EQUAL(SyncSamplingFormulas::bytesPerSweep(4, 8), 32);
}

BOOST_AUTO_TEST_CASE(SyncSamplingFormulas_checkSamplingDelay)
{
    SampleRate hz_1 = SampleRate::Hertz(1);
    SampleRate hz_256 = SampleRate::Hertz(256);

    BOOST_CHECK_EQUAL(SyncSamplingFormulas::checkSamplingDelay(WirelessTypes::samplingMode_sync, hz_1, WirelessModels::node_gLink_10g), true);
    BOOST_CHECK_EQUAL(SyncSamplingFormulas::checkSamplingDelay(WirelessTypes::samplingMode_sync, hz_256, WirelessModels::node_tcLink_1ch), true);
    BOOST_CHECK_EQUAL(SyncSamplingFormulas::checkSamplingDelay(WirelessTypes::samplingMode_sync, hz_256, WirelessModels::node_shmLink), true);

    BOOST_CHECK_EQUAL(SyncSamplingFormulas::checkSamplingDelay(WirelessTypes::samplingMode_sync, hz_256, WirelessModels::node_gLink_10g), false);
}

BOOST_AUTO_TEST_CASE(SyncSamplingFormulas_slotsBetweenTx)
{
    BOOST_CHECK_EQUAL(SyncSamplingFormulas::slotsBetweenTx(2, 2), 1024);
    BOOST_CHECK_EQUAL(SyncSamplingFormulas::slotsBetweenTx(8, 2), 256);
    BOOST_CHECK_EQUAL(SyncSamplingFormulas::slotsBetweenTx(2, 8), 4096);
}

BOOST_AUTO_TEST_CASE(SyncSamplingFormulas_canHaveSlot1)
{
    BOOST_CHECK_EQUAL(SyncSamplingFormulas::canHaveFirstSlot(WirelessModels::node_sgLink, 1), true);
    BOOST_CHECK_EQUAL(SyncSamplingFormulas::canHaveFirstSlot(WirelessModels::node_tcLink_1ch, 1), false);
    BOOST_CHECK_EQUAL(SyncSamplingFormulas::canHaveFirstSlot(WirelessModels::node_tcLink_1ch, 2), true);
}

BOOST_AUTO_TEST_CASE(SyncSamplingFormulas_totalBytesPerBurst)
{
    BOOST_CHECK_EQUAL(SyncSamplingFormulas::totalBytesPerBurst(256, 2), 512);
    BOOST_CHECK_EQUAL(SyncSamplingFormulas::totalBytesPerBurst(20, 30), 600);
}

BOOST_AUTO_TEST_CASE(SyncSamplingFormulas_maxDataBytesPerPacket)
{
    WirelessTypes::CommProtocol lxrs = WirelessTypes::commProtocol_lxrs;

    BOOST_CHECK_EQUAL(SyncSamplingFormulas::maxBytesPerBurstPacket(20, true, lxrs), 80);
    BOOST_CHECK_EQUAL(SyncSamplingFormulas::maxBytesPerBurstPacket(256, true, lxrs), 0);
    BOOST_CHECK_EQUAL(SyncSamplingFormulas::maxBytesPerBurstPacket(20, false, lxrs), 80);
    BOOST_CHECK_EQUAL(SyncSamplingFormulas::maxBytesPerBurstPacket(256, false, lxrs), 0);
}

BOOST_AUTO_TEST_CASE(SyncSamplingFormulas_totalNeededBurstTx)
{
    BOOST_CHECK_EQUAL(SyncSamplingFormulas::totalNeededBurstTx(200, 400), 1);
    BOOST_CHECK_EQUAL(SyncSamplingFormulas::totalNeededBurstTx(400, 20), 20);
}

BOOST_AUTO_TEST_CASE(SyncSamplingFormulas_minTimeBetweenBursts)
{
    BOOST_CHECK_EQUAL(SyncSamplingFormulas::minTimeBetweenBursts(20, 2.0, true), 7);
    BOOST_CHECK_EQUAL(SyncSamplingFormulas::minTimeBetweenBursts(1, 0.5, true), 5);
    BOOST_CHECK_EQUAL(SyncSamplingFormulas::minTimeBetweenBursts(400, 2.0, false), 13);
}

BOOST_AUTO_TEST_CASE(SyncSamplingFormulas_burstTxPerSecond)
{
    BOOST_CHECK_EQUAL(SyncSamplingFormulas::burstTxPerSecond(4, 20, 2.0, true), 1);
    BOOST_CHECK_EQUAL(SyncSamplingFormulas::burstTxPerSecond(4, 2, 2.0, true), 2);
}

BOOST_AUTO_TEST_SUITE_END()


// Copyright (c) 2018 Graphcore Ltd. All rights reserved.
#define BOOST_TEST_MODULE MeshGridTest

#include <boost/test/unit_test.hpp>
#include <poplar/Engine.hpp>
#include <poplibs_support/TestDevice.hpp>
#include <poplin/MeshGrid.hpp>

#include <vector>
using namespace poplibs_support;

struct TestCase {
  poplar::Type type;
  float left;
  float right;
  size_t count;
  std::string name;
  std::vector<float> values;
};

BOOST_AUTO_TEST_CASE(LinSpace) {
  auto device = createTestDevice(TEST_TARGET);
  poplar::Graph g(device.getTarget());

  // Define some linspace arguments with expected result values:
  std::vector<TestCase> testCases = {
      {poplar::FLOAT, 11.f, 42.f, 2, "two", {11.f, 42.f}},
      {poplar::FLOAT, -1.f, 1.f, 3, "small", {-1.f, 0.f, 1.f}},
      {poplar::FLOAT,
       10.f,
       -10.f,
       5,
       "backwards",
       {10.f, 5.f, 0.f, -5.f, -10.f}}};

  poplar::program::Sequence prog;

  for (auto c : testCases) {
    auto var1 = poplin::linspace(g, c.type, c.left, c.right, c.count);
    auto var2 = g.clone(c.type, var1);
    g.setTileMapping(var2, 0);
    prog.add(poplar::program::Copy(var1, var2));
    g.createHostRead(c.name, var2);
  }

  poplar::Engine e(g, prog);
  device.bind([&](const poplar::Device &d) {
    e.load(d);
    e.run();

    for (auto c : testCases) {
      std::vector<float> result(c.values.size(), 0.f);
      e.readTensor(c.name, result.data(), result.data() + result.size());
      for (auto i = 0u; i < result.size(); ++i) {
        BOOST_CHECK_EQUAL(result.at(i), c.values.at(i));
      }
    }
  });
}

BOOST_AUTO_TEST_CASE(MeshGrid) {
  auto device = createTestDevice(TEST_TARGET);
  poplar::Graph g(device.getTarget());

  auto xCoords = poplin::linspace(g, poplar::FLOAT, -1.f, 1.f, 3);
  auto yCoords = poplin::linspace(g, poplar::FLOAT, -2.f, 2.f, 2);
  auto grids = poplin::meshgrid2d(g, xCoords, yCoords);

  auto gridXOut = g.clone(poplar::FLOAT, grids.at(0));
  auto gridYOut = g.clone(poplar::FLOAT, grids.at(1));
  g.setTileMapping(gridXOut, 0);
  g.setTileMapping(gridYOut, 0);

  const auto rowsOut = 2u;
  const auto colsOut = 3u;
  BOOST_CHECK_EQUAL(gridXOut.shape()[0], rowsOut);
  BOOST_CHECK_EQUAL(gridXOut.shape()[1], colsOut);
  BOOST_CHECK_EQUAL(gridYOut.shape()[0], rowsOut);
  BOOST_CHECK_EQUAL(gridYOut.shape()[1], colsOut);

  poplar::program::Sequence prog = {
      poplar::program::Copy(grids.at(0), gridXOut),
      poplar::program::Copy(grids.at(1), gridYOut)};

  g.createHostRead("xs", gridXOut);
  g.createHostRead("ys", gridYOut);

  poplar::Engine e(g, prog);
  device.bind([&](const poplar::Device &d) {
    e.load(d);
    e.run();

    // In Poplar, matrices will come back row major so these are
    // the expected flat results:
    const std::vector<float> correctXs = {-1.f, 0.f, 1.f, -1.f, 0.f, 1.f};
    const std::vector<float> correctYs = {-2.f, -2.f, -2.f, 2.f, 2.f, 2.f};

    std::vector<float> resultX(2 * 3);
    e.readTensor("xs", resultX.data(), resultX.data() + resultX.size());
    for (auto i = 0u; i < resultX.size(); ++i) {
      BOOST_CHECK_EQUAL(resultX.at(i), correctXs.at(i));
    }

    std::vector<float> resultY(2 * 3);
    e.readTensor("ys", resultY.data(), resultY.data() + resultY.size());
    for (auto i = 0u; i < resultY.size(); ++i) {
      BOOST_CHECK_EQUAL(resultY.at(i), correctYs.at(i));
    }
  });
}


// example_error_handling.cpp

// Copyright Paul A. Bristow 2007, 2010.
// Copyright John Maddock 2007.

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

// Note that this file contains quickbook markup as well as code
// and comments, don't change any of the special comment markups!

// Optional macro definitions described in text below:
//   #define BOOST_MATH_DOMAIN_ERROR_POLICY ignore_error
//   #define BOOST_MATH_DOMAIN_ERROR_POLICY errno_on_error
//   #define BOOST_MATH_DOMAIN_ERROR_POLICY is set to: throw_on_error

//[error_handling_example
/*`
The following example demonstrates the effect of
setting the macro BOOST_MATH_DOMAIN_ERROR_POLICY
when an invalid argument is encountered.  For the
purposes of this example, we'll pass a negative
degrees of freedom parameter to the student's t
distribution.

Since we know that this is a single file program we could
just add:

   #define BOOST_MATH_DOMAIN_ERROR_POLICY ignore_error

to the top of the source file to change the default policy
to one that simply returns a NaN when a domain error occurs.
Alternatively we could use:

   #define BOOST_MATH_DOMAIN_ERROR_POLICY errno_on_error

To ensure the `::errno` is set when a domain error occurs
as well as returning a NaN.

This is safe provided the program consists of a single
translation unit /and/ we place the define /before/ any
#includes.  Note that should we add the define after the includes
then it will have no effect!  A warning such as:

[pre warning C4005: 'BOOST_MATH_OVERFLOW_ERROR_POLICY' : macro redefinition]

is a certain sign that it will /not/ have the desired effect.

We'll begin our sample program with the needed includes:
*/

// Boost
#include <boost/math/distributions/students_t.hpp>
   using boost::math::students_t;  // Probability of students_t(df, t).

// std
#include <iostream>
   using std::cout;
   using std::endl;

#include <stdexcept>
   using std::exception;

#include <cstddef>
   // using ::errno

/*`
Next we'll define the program's main() to call the student's t
distribution with an invalid degrees of freedom parameter,
the program is set up to handle either an exception or a NaN:
*/

int main()
{
   cout << "Example error handling using Student's t function. " << endl;
   cout << "BOOST_MATH_DOMAIN_ERROR_POLICY is set to: "
      << BOOST_STRINGIZE(BOOST_MATH_DOMAIN_ERROR_POLICY) << endl;

   double degrees_of_freedom = -1; // A bad argument!
   double t = 10;

   try
   {
      errno = 0; // Clear/reset.
      students_t dist(degrees_of_freedom); // exception is thrown here if enabled.
      double p = cdf(dist, t);
      // Test for error reported by other means:
      if((boost::math::isnan)(p))
      {
         cout << "cdf returned a NaN!" << endl;
         if (errno != 0)
         { // So errno has been set.
           cout << "errno is set to: " << errno << endl;
         }
      }
      else
         cout << "Probability of Student's t is " << p << endl;
   }
   catch(const std::exception& e)
   {
      std::cout <<
         "\n""Message from thrown exception was:\n   " << e.what() << std::endl;
   }
   return 0;
} // int main()

/*`

Here's what the program output looks like with a default build
(one that *does throw exceptions*):

[pre
Example error handling using Student's t function.
BOOST_MATH_DOMAIN_ERROR_POLICY is set to: throw_on_error

Message from thrown exception was:
   Error in function boost::math::students_t_distribution<double>::students_t_distribution:
   Degrees of freedom argument is -1, but must be > 0 !
]

Alternatively let's build with:

   #define BOOST_MATH_DOMAIN_ERROR_POLICY ignore_error

Now the program output is:

[pre
Example error handling using Student's t function.
BOOST_MATH_DOMAIN_ERROR_POLICY is set to: ignore_error
cdf returned a NaN!
]

And finally let's build with:

   #define BOOST_MATH_DOMAIN_ERROR_POLICY errno_on_error

Which gives the output show errno:

[pre
Example error handling using Student's t function.
BOOST_MATH_DOMAIN_ERROR_POLICY is set to: errno_on_error
cdf returned a NaN!
errno is set to: 33
]

*/

//] [error_handling_eg end quickbook markup]


// Copyright 2014-2015 Josh Pieper, jjp@pobox.com.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

#include "base/quaternion.h"

#include <fmt/format.h>

#include <boost/test/auto_unit_test.hpp>

#include "base/euler.h"

namespace {
using namespace mjmech;
typedef base::Euler Euler;
typedef base::Quaternion Quaternion;
typedef base::Point3D Point3D;

void CheckVectorClose(const Point3D& lhs,
                      double x, double y, double z) {
  BOOST_CHECK_SMALL(lhs.x() - x, 1e-2);
  BOOST_CHECK_SMALL(lhs.y() - y, 1e-2);
  BOOST_CHECK_SMALL(lhs.z() - z, 1e-2);
}
}

BOOST_AUTO_TEST_CASE(BasicQuaternion) {
  Point3D v(10., 0., 0);

  v = Quaternion::FromEuler(0, 0, M_PI_2).Rotate(v);
  CheckVectorClose(v, 0, 10, 0);

  v = Quaternion::FromEuler(0, 0, -M_PI_2).Rotate(v);
  CheckVectorClose(v, 10, 0, 0);

  v = Quaternion::FromEuler(0, M_PI_2, 0).Rotate(v);
  CheckVectorClose(v, 0, 0, -10);

  v = Quaternion::FromEuler(M_PI_2, 0, 0).Rotate(v);
  CheckVectorClose(v, 0, 10, 0);

  v = Quaternion::FromEuler(0, 0, M_PI_2).Rotate(v);
  CheckVectorClose(v, -10, 0, 0);

  v = Quaternion::FromEuler(0, M_PI_2, 0).Rotate(v);
  CheckVectorClose(v, 0, 0, 10);

  v = Quaternion::FromEuler(M_PI_2, 0, 0).Rotate(v);
  CheckVectorClose(v, 0, -10, 0);

  v = Quaternion::FromEuler(0, 0, M_PI_2).Rotate(v);
  CheckVectorClose(v, 10, 0, 0);
}

namespace {
void CheckEuler(Euler euler_rad,
                double roll_rad,
                double pitch_rad,
                double yaw_rad) {
  BOOST_CHECK_SMALL(euler_rad.roll - roll_rad, 1e-5);
  BOOST_CHECK_SMALL(euler_rad.pitch - pitch_rad, 1e-5);
  BOOST_CHECK_SMALL(euler_rad.yaw - yaw_rad, 1e-5);
}
}

BOOST_AUTO_TEST_CASE(QuaternionEulerAndBack) {
  struct TestCase {
    double roll_deg;
    double pitch_deg;
    double yaw_deg;
  };

  TestCase tests[] = {
    {45, 0, 0},
    {0, 45, 0},
    {0, 0, 45},
    {0, 90, 0},
    {0, 90, 20},
    {0, -90, 0},
    {0, -90, -10},
    {0, -90, 30},
    {10, 20, 30},
    {-30, 10, 20},
  };

  for (const auto& x: tests) {
    BOOST_TEST_CONTEXT(fmt::format("{} {} {}",
                                   x.roll_deg, x.pitch_deg, x.yaw_deg)) {
      Euler result_rad = Quaternion::FromEuler(
          x.roll_deg / 180 * M_PI,
          x.pitch_deg / 180 * M_PI,
          x.yaw_deg / 180 * M_PI).euler_rad();
      CheckEuler(result_rad,
                 x.roll_deg / 180 * M_PI,
                 x.pitch_deg / 180 *M_PI,
                 x.yaw_deg / 180 * M_PI);
    }
  }
}

BOOST_AUTO_TEST_CASE(QuaternionMultiply1) {
  Quaternion x90 = Quaternion::FromEuler(0, M_PI_2, 0);
  Quaternion xn90 = Quaternion::FromEuler(0, -M_PI_2, 0);
  Quaternion y90 = Quaternion::FromEuler(M_PI_2, 0, 0);

  Quaternion result = xn90 * y90 * x90;
  Point3D vector(0, 1, 0);
  vector = result.Rotate(vector);
  CheckVectorClose(vector, -1, 0, 0);

  Quaternion initial = Quaternion::FromEuler(0, 0, 0.5 * M_PI_2);
  initial = Quaternion::FromEuler(0, 0, 0.5 * M_PI_2) * initial;
  CheckEuler(initial.euler_rad(), 0, 0, M_PI_2);

  initial = Quaternion::FromEuler(10 / 180. * M_PI, 0, 0) * initial;
  vector = initial.Rotate(vector);
  CheckVectorClose(vector, 0, -0.9848078, -0.17364818);
  CheckEuler(initial.euler_rad(), 0, -10 / 180.0 * M_PI, M_PI_2);
}


#ifndef QDOFS_HPP
#define QDOFS_HPP

#pragma once
#include <Eigen/Core>
#include <igl/edges.h>
#include <igl/vertex_triangle_adjacency.h>
#include <iostream>
#include <set>

class QDofs {
public:
	// list of edges connected to given vertex
	std::map<int, std::set<int>> vertex2edge;
	Eigen::MatrixXd nodes;

private:
	const Eigen::MatrixXd &vertices_;
	const Eigen::MatrixXi &triangles_;
	Eigen::MatrixXi        edges_;
	Eigen::MatrixXi        qdof_;

public:
	//! QDofs constructor
	//!
	//! @param[in] vertices a list of triangle vertices
	//! @param[in] triangles a list of triangles
	QDofs(const Eigen::MatrixXd &vertices,
	      const Eigen::MatrixXi &triangles)
	    : vertices_(vertices), triangles_(triangles) {
		set_quadraticDofs();
	}

	int get_dofs(Eigen::MatrixXi &qdof) const {
		qdof = qdof_;
		return edges_.rows() + vertices_.rows();
	}

	Eigen::MatrixXd get_vertices() const {
		return vertices_;
	}

private:
	//! Create matrix containing the global indices of the degrees of freedom
	//! associated to quadratic FEM. First numbers vertices and then edges.
	//! First three columns are vertices, the last three columns are the edges.
	void set_quadraticDofs() {
		// get list of mesh edges
		igl::edges(triangles_, edges_);

		// resize dof_ so it has right length and 6 columns
		qdof_.resize(triangles_.rows(), 6);
		// set frist 3 columns to match the vertices defining the triangles
		qdof_.block(0, 0, triangles_.rows(), 3) = triangles_;

		int Ne = edges_.rows();
		int Nv = vertices_.rows();

		// resize matrix with nodes coordinates
		nodes.resize(Ne + Nv, 3);
		// set first Nv rows to be the vertices
		nodes.block(0, 0, vertices_.rows(), 3) = vertices_;

		// get list of elements to which each vertex is connected
		std::vector<std::vector<int>> VF;
		std::vector<std::vector<int>> VFi;
		igl::vertex_triangle_adjacency(Nv, triangles_, VF, VFi);

		// traverse the edges
		for (int i = 0; i < Ne; i++) {
			// get vertices attached to current edge
			auto edgevert = edges_.row(i);

			// add midpoint to matrix of nodes coordinates
			nodes.row(i + Nv) = (vertices_.row(edgevert(0)) + vertices_.row(edgevert(1))) / 2.;

			// get list of elements connect to those vertices
			auto tria_vert0 = VF[edgevert(0)];
			auto tria_vert1 = VF[edgevert(1)];
			// get list of local index of those vertices in said elements
			auto tria_vertid0 = VFi[edgevert(0)];
			auto tria_vertid1 = VFi[edgevert(1)];
			// traverse elements connected to first vertex
			for (unsigned int k0 = 0; k0 < tria_vert0.size(); k0++) {
				auto tria0 = tria_vert0[k0];
				// traverse elements connected to second element
				for (unsigned int k1 = 0; k1 < tria_vert1.size(); k1++) {
					auto tria1 = tria_vert1[k1];
					// if the elements match, it means we have found a triangle
					// that contains the current edge
					if (tria0 == tria1) {
						int id = tria_vertid0[k0] + tria_vertid1[k1];
						// fill the matrix dof_ accordingly
						if (id == 1) {
							qdof_(tria0, 3) = Nv + i;
						} // edge 0, dof 3
						else if (id == 3) {
							qdof_(tria0, 4) = Nv + i;
						} // edge 1, dof 4
						else if (id == 2) {
							qdof_(tria0, 5) = Nv + i;
						} // edge 2, dof 5
						else {
							std::cout << " error!" << std::endl;
						}

					} // end found element

				} //looping on second vertex
			}   //looping on first vertex

			// add edge to the vertex2edge vector
			for (int k = 0; k < 2; k++) {
				// get set of each vertex
				auto auxset = vertex2edge.find(edgevert(k));
				// if vertex index is already there
				if (auxset != vertex2edge.end()) {
					// add it to its set
					(*auxset).second.insert((Nv + i));
				} else { // if not, create set
					std::set<int> vset;
					vset.insert((Nv + i));
					vertex2edge.insert(std::pair<int, std::set<int>>(edgevert(k), vset));
				}
			}

		} // looping on edges
	}
};
#endif


#include <string>
#include <iostream>
#include <memory>
#include <vector>

#include <boost/random.hpp>

#include <Eigen/Dense>

#include "Kernel.h"
#include "GaussianProcess.h"
#include "SparseGaussianProcess.h"

#include "Likelihood.h"
#include "SparseLikelihood.h"

#include "GaussianProcessInference.h"

typedef gpr::GaussianProcess<double> GaussianProcessType;
typedef GaussianProcessType::VectorType VectorType;
typedef GaussianProcessType::MatrixType MatrixType;
typedef GaussianProcessType::DiagMatrixType DiagMatrixType;
typedef GaussianProcessType::VectorListType VectorListType;

typedef gpr::GaussianProcessInference<double>    GaussianProcessInferenceType;
typedef GaussianProcessInferenceType::Pointer    GaussianProcessInferenceTypePointer;
typedef gpr::GaussianLogLikelihood<double>       LikelihoodType;
typedef LikelihoodType::Pointer                  LikelihoodTypePointer;
typedef gpr::GaussianExpKernel<double>           GaussianExpKernelType;
typedef GaussianExpKernelType::Pointer           GaussianExpKernelTypePointer;
typedef gpr::GaussianKernel<double>              GaussianKernelType;
typedef GaussianKernelType::Pointer              GaussianKernelTypePointer;

void Test1(){
    std::cout << "Test 1: maximum gaussian log likelihood with gradient descent test ..." << std::flush;
    std::cout.precision(8);

    // global parameters
    unsigned n = 200;
    double noise = 0.1;

    // construct training data
    auto f = [](double x)->double { return (0.5*std::sin(x+10*x) + std::sin(4*x))*x*x; };

    static boost::minstd_rand randgen(static_cast<unsigned>(time(0)));
    static boost::normal_distribution<> dist(0, noise);
    static boost::variate_generator<boost::minstd_rand, boost::normal_distribution<> > r(randgen, dist);

    double start = -5;
    double stop = 10;
    VectorType Xn = VectorType::Zero(n);
    VectorType Yn = VectorType::Zero(n);
    for(unsigned i=0; i<n; i++){
        Xn[i] = start + i*(stop-start)/n;
        Yn[i] = f(Xn[i])+r();
    }

    GaussianExpKernelTypePointer gk(new GaussianExpKernelType(1, 1));
    GaussianProcessType::Pointer gp(new GaussianProcessType(gk));
    //gp->DebugOn();
    gp->SetSigma(noise);
    for(unsigned i=0; i<n; i++){
        gp->AddSample(VectorType::Constant(1,Xn[i]), VectorType::Constant(1,Yn[i]));
    }

    // setup likelihood
    double step = 1e-1;
    unsigned iterations = 100;

    LikelihoodTypePointer lh(new LikelihoodType());
    GaussianProcessInferenceTypePointer gpi(new GaussianProcessInferenceType(lh, gp, step, iterations));

    bool exp_output = true;
    gpi->Optimize(false, exp_output);

    //std::cout << "Parameters are: ";
    GaussianProcessInferenceType::ParameterVectorType parameters = gpi->GetParameters();
    for(unsigned i=0; i<parameters.size(); i++){
        parameters[i] = std::exp(parameters[i]);
        //std::cout << parameters[i] << ", ";
    }
    //std::cout << std::endl;

    GaussianKernelTypePointer k(new GaussianKernelType(1,1));
    k->SetParameters(parameters);
    gp->SetKernel(k);

    // evaluate error
    double error = 0;
    unsigned gt_n = 1000;
    for(unsigned i=0; i<gt_n; i++){
        double x = start + i*(stop-start)/gt_n;
        double p = gp->Predict(VectorType::Constant(1,x))[0];
        error += std::fabs(p-f(x));
    }

    if(error/gt_n > 2){
        std::stringstream ss; ss<<error/gt_n; throw ss.str();
    }
    else{
        std::cout << "[passed]" << std::endl;
    }

}

typedef gpr::SparseGaussianProcess<double> SparseGaussianProcessType;

void Test2(){
    std::cout << "Test 2: maximum gaussian log likelihood on dense gp and prediction with sparse gp test ... " << std::flush;
    std::cout.precision(8);

    // global parameters
    unsigned n = 300;
    unsigned m = 50;
    double noise = 0.1;
    double jitter = 0.1;

    // construct training data
    auto f = [](double x)->double { return (0.5*std::sin(x+10*x) + std::sin(4*x))*x*x; };

    static boost::minstd_rand randgen(static_cast<unsigned>(time(0)));
    static boost::normal_distribution<> dist(0, noise);
    static boost::variate_generator<boost::minstd_rand, boost::normal_distribution<> > r(randgen, dist);

    double start = -5;
    double stop = 10;
    VectorType Xn = VectorType::Zero(n);
    VectorType Yn = VectorType::Zero(n);
    for(unsigned i=0; i<n; i++){
        Xn[i] = start + i*(stop-start)/n;
        Yn[i] = f(Xn[i])+r();
    }

    std::vector<unsigned> indices;
    for(unsigned i=0; i<n; i++){
        indices.push_back(i);
    }
    std::random_shuffle(indices.begin(), indices.end());

    // setup dense gaussian process
    GaussianExpKernelTypePointer gk(new GaussianExpKernelType(1, 1));
    GaussianProcessType::Pointer gp(new GaussianProcessType(gk));
    gp->SetSigma(noise);

    // setup sparse gaussian process
    SparseGaussianProcessType::Pointer sgp(new SparseGaussianProcessType(gk));
    sgp->SetSigma(noise);
    sgp->SetJitter(jitter);

    for(unsigned i=0; i<m; i++){
        gp->AddSample(VectorType::Constant(1,Xn[indices[i]]), VectorType::Constant(1,Yn[indices[i]]));
        sgp->AddInducingSample(VectorType::Constant(1,Xn[indices[i]]), VectorType::Constant(1,Yn[indices[i]]));
    }
    for(unsigned i=0; i<n; i++){
        sgp->AddSample(VectorType::Constant(1,Xn[i]), VectorType::Constant(1,Yn[i]));
    }

    // setup likelihood
    double step = 1e-1;
    unsigned iterations = 100;

    LikelihoodTypePointer lh(new LikelihoodType());
    GaussianProcessInferenceTypePointer gpi(new GaussianProcessInferenceType(lh, gp, step, iterations));

    bool exp_output = true;
    gpi->Optimize(false, exp_output);

    //std::cout << "Parameters are: ";
    GaussianProcessInferenceType::ParameterVectorType parameters = gpi->GetParameters();
    for(unsigned i=0; i<parameters.size(); i++){
        parameters[i] = std::exp(parameters[i]);
        //std::cout << parameters[i] << ", ";
    }
    //std::cout << std::endl;

    GaussianKernelTypePointer k(new GaussianKernelType(1,1));
    k->SetParameters(parameters);
    gp->SetKernel(k);
    sgp->SetKernel(k);

    // evaluate error
    double dense_error = 0;
    double sparse_error = 0;
    unsigned gt_n = 1000;
    for(unsigned i=0; i<gt_n; i++){
        double x = start + i*(stop-start)/gt_n;
        double pd = gp->Predict(VectorType::Constant(1,x))[0];
        double ps = sgp->Predict(VectorType::Constant(1,x))[0];
        dense_error += std::fabs(pd-f(x));
        sparse_error += std::fabs(ps-f(x));
    }

    if(dense_error < sparse_error){
        throw std::string("dense gp should result in a smaller error than the sparse gp");
    }
    else{
        std::cout << "[passed]" << std::endl;
    }

}

void Test3(){
    //std::cout << "Test 3: maximum likelihood of periodic signal ..." << std::flush;

    // ground truth periodic variable
    auto f = [](double x)->double { return 3.5*std::sin(1.3*x); };

    double start = 0;
    double stop = 5 * 2*M_PI; // full interval
    unsigned n = 350;

    //--------------------------------------------------------------------------------
    // generating ground truth
    VectorType Xn = VectorType::Zero(n);
    VectorType Yn = VectorType::Zero(n);
    for(unsigned i=0; i<n; i++){
        Xn[i] = start + i*(stop-start)/n;
        Yn[i] = f(Xn[i]);
    }

    //--------------------------------------------------------------------------------
    // perform training
    double noise = 0.01;
    static boost::minstd_rand randgen(static_cast<unsigned>(time(0)));
    static boost::normal_distribution<> dist(0, noise);
    static boost::variate_generator<boost::minstd_rand, boost::normal_distribution<> > r(randgen, dist);

    double interval_training_end = 2 * 2*M_PI; // interval to train
    unsigned number_of_samples = 200;

    typedef gpr::PeriodicKernel<double>		KernelType;
    typedef KernelType::Pointer KernelTypePointer;

    KernelTypePointer k(new KernelType(1.5, 0.5, 4)); // scale, period, smoothness
    GaussianProcessType::Pointer gp(new GaussianProcessType(k));
    gp->SetSigma(noise); // noise

    // add samples
    double training_step_size = (interval_training_end - start) / number_of_samples;
    for(unsigned i=0; i<number_of_samples; i++){
        VectorType x(1);
        x(0) = start + i*training_step_size;

        VectorType y(1);
        y(0) = f(x(0)) + r();

        gp->AddSample(x, y);
    }

    //--------------------------------------------------------------------------------
    // maximum likelihood
    // setup likelihood
    double step = 1e-1;
    unsigned iterations = 300;
    GaussianProcessInferenceType::BooleanVectorType bv;
    bv.push_back(false);
    bv.push_back(true);
    bv.push_back(false);

    LikelihoodTypePointer lh(new LikelihoodType());
    GaussianProcessInferenceTypePointer gpi(new GaussianProcessInferenceType(lh, gp, step, iterations));
    gpi->SetParametersToOptimize(bv);

    bool exp_output = false;
    gpi->Optimize2(true, exp_output);

    std::cout << "print \"Parameters are: ";
    GaussianProcessInferenceType::ParameterVectorType parameters = gpi->GetParameters();
    for(unsigned i=0; i<parameters.size(); i++){
        if(exp_output) parameters[i] = std::exp(parameters[i]);
        std::cout << parameters[i] << ", ";
    }
    std::cout << "\"" << std::endl;

//    KernelTypePointer k(new KernelType(1,1,1));
//    k->SetParameters(parameters);
//    gp->SetKernel(k);

    //--------------------------------------------------------------------------------
    // predict full intervall
    VectorType y_predict(n);
    VectorType y(n);
    for(unsigned i=0; i<n; i++){
        VectorType x(1);
        x(0) = Xn[i];
        y(i) = Yn[i];
        y_predict[i] = gp->Predict(x)(0);
    }

    double err = (y-y_predict).norm();
    std::cout << "print \"" << err/y.rows() << "\""<< std::endl;
//    if(err>0.4){
//        std::cout << " [failed] with an error of " << err << std::endl;
//    }
//    else{
//        std::cout << " [passed]." << std::endl;
//    }

    return;
    std::cout << "import numpy as np" << std::endl;
    std::cout << "import pylab as plt" << std::endl;

    // ground truth
    //unsigned gt_n = 1000;
    std::cout << "x = np.array([";
    for(unsigned i=0; i<n; i++){
        std::cout << Xn[i] << ", ";
    }
    std::cout << "])" << std::endl;
    std::cout << "y = np.array([";
    for(unsigned i=0; i<n; i++){
        std::cout << f(Xn[i]) << ", ";
    }
    std::cout << "])" << std::endl;
    std::cout << "plt.plot(x,y)" << std::endl;

    // training
    std::cout << "x_train = np.array([";
    for(unsigned i=0; i<number_of_samples; i++){
        std::cout << start + i*training_step_size << ", ";
    }
    std::cout << "])" << std::endl;
    std::cout << "y_train = np.array([";
    for(unsigned i=0; i<number_of_samples; i++){
        std::cout << f(start + i*training_step_size ) + r()<< ", ";
    }
    std::cout << "])" << std::endl;
    std::cout << "plt.plot(x_train, y_train, '.k')" << std::endl;

    // dense prediction
    double dense_error = 0;
    std::cout << "gp_y = np.array([";
    for(unsigned i=0; i<n; i++){
        double p = gp->Predict(VectorType::Constant(1,Xn[i]))[0];
        dense_error += std::fabs(p-f(Xn[i]));
        std::cout << p << ", ";
    }
    std::cout << "])" << std::endl;
    std::cout << "plt.plot(x, gp_y, '-r')" << std::endl;

    std::cout << "plt.show()" << std::endl;
}

int main (int argc, char *argv[]){
    //std::cout << "Maximum likelihood test 2: " << std::endl;
    try{
//        Test1();
//        Test2();
        Test3();
    }
    catch(std::string& s){
        std::cout << "[failed] Error: " << s << std::endl;
        return -1;
    }

    return 0;
}


#ifndef ESKF_IMPL
#define ESKF_IMPL

#include "ESKF.h"
#include "utilities.h"

#include <Eigen/LU>
#include <Wire.h>

#ifndef SGN
#define SGN(X) ((X > 0) - (X < 0))
#endif

#ifndef RAD_TO_DEG
#define RAD_TO_DEG (180.0 / M_PI)
#endif

#ifndef DEG_TO_RAD
#define DEG_TO_RAD (M_PI / 180.0)
#endif

#ifndef MS2_TO_G
#define MS2_TO_G (1.0 / 9.81)
#endif

#ifndef G_TO_MS2
#define G_TO_MS2 (9.81)
#endif

namespace IMU_EKF
{

template <typename precision>
ESKF<precision>::ESKF()
{
    init();
}

template <typename precision>
void ESKF<precision>::init()
{
    qref_ = Quaternion<precision>();
    x_.setZero();
    P_.setIdentity();
    Q_.setIdentity();
    // TODO don't hardcode this
    Q_(0, 0) = 0.05;
    Q_(1, 1) = 0.05;
    Q_(2, 2) = 0.05;
    Q_(3, 3) = 0.05;
    Q_(4, 4) = 0.05;
    Q_(5, 5) = 0.05;
    Q_(6, 6) = 0.025;
    Q_(7, 7) = 0.025;
    Q_(8, 8) = 0.025;
    Q_(9, 9) = 0.025;
    Q_(10, 10) = 0.025;
    Q_(11, 11) = 0.025;
    Q_(12, 12) = 0.01;
    Q_(13, 13) = 0.01;
    Q_(14, 14) = 0.01;

    R_Gyr_.setIdentity();
    R_Gyr_(0, 0) = 0.0000045494 * DEG_TO_RAD;
    R_Gyr_(1, 1) = 0.0000039704 * DEG_TO_RAD;
    R_Gyr_(2, 2) = 0.0000093844 * DEG_TO_RAD;

    R_Acc_.setIdentity();
    R_Acc_(0, 0) = 0.0141615383 * G_TO_MS2;
    R_Acc_(1, 1) = 0.0164647549 * G_TO_MS2;
    R_Acc_(2, 2) = 0.0100094303 * G_TO_MS2;

    R_Mag_.setIdentity();
    R_Mag_(0, 0) = 0.00004;
    R_Mag_(1, 1) = 0.00004;
    R_Mag_(2, 2) = 0.00004;
}

template <typename precision>
void ESKF<precision>::initWithAcc(const float ax, const float ay, const float az)
{
    init();

    // see https://cache.freescale.com/files/sensors/doc/app_note/AN3461.pdf
    // rotation sequence R = Rx * Ry * Rz
    // eq. 38
    precision roll = std::atan2(ay, SGN(-az) * std::sqrt(az * az + 0.01 * ax * ax));
    // eq. 37
    precision pitch = std::atan(-ax / std::sqrt(ay * ay + az * az));

    precision sr05 = std::sin(0.5 * roll);
    precision cr05 = std::cos(0.5 * roll);
    precision sp05 = std::sin(0.5 * pitch);
    precision cp05 = std::cos(0.5 * pitch);
    qref_ = Quaternion<precision>(sr05, 0, 0, cr05) * Quaternion<precision>(0, sp05, 0, cp05);
}

template <typename precision>
void ESKF<precision>::initWithAccAndMag(const float ax, const float ay, const float az, const float mx, const float my, const float mz, const Eigen::Matrix<precision, 3, 3> &Winv, const Eigen::Matrix<precision, 3, 1> &V)
{
    init();

    // see https://cache.freescale.com/files/sensors/doc/app_note/AN3461.pdf
    // rotation sequence R = Rx * Ry * Rz
    // eq. 38
    precision roll = std::atan2(ay, SGN(-az) * std::sqrt(az * az + 0.01 * ax * ax));
    // eq. 37
    precision pitch = std::atan(-ax / std::sqrt(ay * ay + az * az));

    // see https://www.nxp.com/docs/en/application-note/AN4246.pdf
    // eq. 6 - 10
    precision sr = std::sin(roll);
    precision cr = std::cos(roll);
    precision sp = std::sin(pitch);
    precision cp = std::cos(pitch);

    Eigen::Matrix<precision, 3, 3> RxT;
    RxT << 1, 0, 0,
        0, cr, sr,
        0, -sr, cr;
    Eigen::Matrix<precision, 3, 3> RyT;
    RyT << cp, 0, -sp,
        0, 1, 0,
        sp, 0, cp;

    Eigen::Matrix<precision, 3, 1> Bp;
    Bp << mx, my, mz;

    Eigen::Matrix<precision, 3, 1> Bf;
    Bf = RyT * RxT * Winv * (Bp - V);

    precision yaw = -atan2(-Bf(1), Bf(0));

    precision sr05 = std::sin(0.5 * roll);
    precision cr05 = std::cos(0.5 * roll);
    precision sp05 = std::sin(0.5 * pitch);
    precision cp05 = std::cos(0.5 * pitch);
    precision sy05 = std::sin(0.5 * yaw);
    precision cy05 = std::cos(0.5 * yaw);

    qref_ = Quaternion<precision>(sr05, 0, 0, cr05) * Quaternion<precision>(0, sp05, 0, cp05) * Quaternion<precision>(0, 0, sy05, cy05);
}

template <typename precision>
void ESKF<precision>::predict(precision dt)
{
    // eq. 23
    Quaternion<precision> angular_velocity_quat(x_[12], x_[13], x_[14], 0);
    qref_ += dt * ((0.5 * angular_velocity_quat) * qref_);
    qref_.normalize();

    Eigen::Matrix<precision, 9, 9> A;
    A.setIdentity();
    A(0, 3) = dt;
    A(0, 6) = dt * dt * 0.5;
    A(1, 4) = dt;
    A(1, 7) = dt * dt * 0.5;
    A(2, 5) = dt;
    A(2, 8) = dt * dt * 0.5;
    A(3, 6) = dt;
    A(4, 7) = dt;
    A(4, 8) = dt;

    x_.segment(0, 9) = A * x_.segment(0, 9);
    P_.topLeftCorner(9, 9) = A * P_.topLeftCorner(9, 9) * A.transpose() + dt * Q_.topLeftCorner(9, 9);

    // eq. 38
    Eigen::Matrix<precision, 3, 1> angular_velocity = x_.segment(12, 3);
    Eigen::Matrix<precision, 3, 1> error = x_.segment(9, 3);
    Eigen::Matrix<precision, 6, 6> Jac = Eigen::Matrix<precision, 6, 6>::Zero();
    Jac.topLeftCorner(3, 3) = toCrossMatrix<precision>(error);
    Jac.topRightCorner(3, 3) = -toCrossMatrix<precision>(angular_velocity);
    // eq. 39
    Eigen::Matrix<precision, 6, 6> G = Eigen::Matrix<precision, 6, 6>::Identity();
    G.bottomRightCorner(3, 3) *= -1;

    // eq. 33 adpated
    P_.bottomRightCorner(6, 6) = Jac * P_.bottomRightCorner(6, 6) * Jac.transpose() + G * dt * Q_.bottomRightCorner(6, 6) * G.transpose();
}

template <typename precision>
void ESKF<precision>::correctGyr(const float gx, const float gy, const float gz)
{
    // z
    Eigen::Matrix<precision, MEASSUREMENT_GYR_SIZE, 1> z;
    z << gx * DEG_TO_RAD, gy * DEG_TO_RAD, gz * DEG_TO_RAD;

    // Kalman Gain
    // K = P * H' (H * P * H' + V * R * V')^-
    // H = eye(3)
    Eigen::Matrix<precision, 3, MEASSUREMENT_GYR_SIZE> K;
    K = P_.bottomRightCorner(3, 3) * (P_.bottomRightCorner(3, 3) + R_Gyr_).inverse();

    // x = x + K * (z - H * x)
    x_.segment(12, 3) += K * (z - x_.segment(12, 3));

    // P = (I - KH)P
    Eigen::Matrix<precision, 3, 3> IKH = Eigen::Matrix<precision, 3, 3>::Identity();
    IKH -= K;

    P_.bottomRightCorner(3, 3) = IKH * P_.bottomRightCorner(3, 3);
}

template <typename precision>
void ESKF<precision>::correctAcc(const float ax, const float ay, const float az)
{
    // z
    Eigen::Matrix<precision, MEASSUREMENT_ACC_SIZE, 1> z;
    z << ax * G_TO_MS2, ay * G_TO_MS2, az * G_TO_MS2;

    Eigen::Matrix<precision, 3, 1> gravity;
    gravity << 0.0, 0.0, -9.81;

    if (std::abs(z.norm() - 9.81) < 0.7)
    {
        // set acc to zero
        // z
        Eigen::Matrix<precision, 3, 1> z_acc = Eigen::Matrix<precision, 3, 1>::Zero();

        // Kalman Gain
        // K = P * H' (H * P * H' + V * R * V')^-
        // H = eye(3)
        Eigen::Matrix<precision, 3, 3> K_acc;
        K_acc = P_.block(6, 6, 3, 3) * (P_.block(6, 6, 3, 3) + R_Acc_).inverse();

        // x = x + K * (z - H * x)
        x_.segment(6, 3) += K_acc * (z_acc - x_.segment(6, 3));

        // P = (I - KH)P
        Eigen::Matrix<precision, 3, 3> IKH = Eigen::Matrix<precision, 3, 3>::Identity();
        IKH -= K_acc;

        P_.block(6, 6, 3, 3) = IKH * P_.block(6, 6, 3, 3);

        // correct tilt
        Eigen::Matrix<precision, 3, 1> vi = gravity;

        // eq. 42
        Eigen::Matrix<precision, 3, 1> error = x_.segment(9, 3);
        // Eigen::Matrix<precision, 3, 3> Aa = toRotationMatrix<precision>(error);
        Eigen::Matrix<precision, 3, 3> Aq = qref_.toRotationMatrix();
        Eigen::Matrix<precision, 3, 1> vb_pred = Aq * vi;

        // H
        // eq. 44
        Eigen::Matrix<precision, 3, 3> Ha = toCrossMatrix<precision>(vb_pred);

        // eq. 46
        Eigen::Matrix<precision, 3, 3> K;
        K = P_.block(9, 9, 3, 3) * Ha.transpose() * (Ha * P_.block(9, 9, 3, 3) * Ha.transpose() + R_Acc_).inverse();

        // eq. 47
        // h = vb_pred
        x_.segment(9, 3) += K * (z - vb_pred - Ha * error);

        // eq. 48
        P_.block(9, 9, 3, 3) -= K * Ha * P_.block(9, 9, 3, 3);
    }
    else
    {
        Eigen::Matrix<precision, 3, 1> vi = gravity + x_.segment(6, 3);

        // eq. 42
        Eigen::Matrix<precision, 3, 1> error = x_.segment(9, 3);
        Eigen::Matrix<precision, 3, 3> Aa = toRotationMatrix<precision>(error);
        Eigen::Matrix<precision, 3, 3> Aq = qref_.toRotationMatrix();
        Eigen::Matrix<precision, 3, 1> vb_pred = Aq * vi;

        // H
        // eq. 44
        Eigen::Matrix<precision, 3, 6> H;
        Eigen::Matrix<precision, 3, 3> Ha = toCrossMatrix<precision>(vb_pred);
        H.topLeftCorner(3, 3) = Aa * Aq;
        H.bottomRightCorner(3, 3) = Ha;

        // eq. 46
        Eigen::Matrix<precision, 6, 3> K;
        K = P_.block(6, 6, 6, 6) * H.transpose() * (H * P_.block(6, 6, 6, 6) * H.transpose() + R_Acc_).inverse();

        // eq. 47
        // h = v_bred
        x_.segment(6, 6) += K * (z - vb_pred - Ha * error);

        // eq. 48
        P_.block(6, 6, 6, 6) -= K * H * P_.block(6, 6, 6, 6);
    }
}

template <typename precision>
void ESKF<precision>::correctMag(const float mx, const float my, const float mz, const float incl, const float B, const Eigen::Matrix<precision, 3, 3> &W, const Eigen::Matrix<precision, 3, 1> &V)
{
    // z
    Eigen::Matrix<precision, MEASSUREMENT_MAG_SIZE, 1> z;
    z << mx, my, mz;

    Eigen::Matrix<precision, 3, 1> vi;
    vi << std::cos(incl), 0, -std::sin(incl);
    vi = B * vi;

    // eq. 42
    Eigen::Matrix<precision, 3, 1> error = x_.segment(9, 3);
    // Eigen::Matrix<precision, 3, 3> Aa = toRotationMatrix<precision>(error);
    Eigen::Matrix<precision, 3, 3> Aq = qref_.toRotationMatrix();
    Eigen::Matrix<precision, 3, 1> vb_pred = Aq * vi;

    // h(v) = W * v + V
    Eigen::Matrix<precision, 3, 1> h = W * vb_pred + V;

    // H
    // eq. 44
    Eigen::Matrix<precision, 3, 3> Ha = W * toCrossMatrix<precision>(vb_pred);

    // eq. 46
    Eigen::Matrix<precision, 3, 3> K;
    K = P_.block(9, 9, 3, 3) * Ha.transpose() * (Ha * P_.block(9, 9, 3, 3) * Ha.transpose() + R_Mag_).inverse();

    // eq. 47
    x_.segment(9, 3) += K * (z - h - Ha * error);

    // eq. 48
    P_.block(9, 9, 3, 3) -= K * Ha * P_.block(9, 9, 3, 3);
}

template <typename precision>
void ESKF<precision>::reset()
{
    // eq. 21
    qref_ = Quaternion<precision>(x_(9), x_(10), x_(11), 2.0) * qref_;
    // eq. 22
    qref_.normalize();
    x_(9) = 0.0;
    x_(10) = 0.0;
    x_(11) = 0.0;
}

template <typename precision>
Eigen::Matrix<precision, STATE_SIZE, 1> ESKF<precision>::getState() const
{
    return x_;
}

template <typename precision>
void ESKF<precision>::getAttitude(float &roll, float &pitch, float &yaw) const
{
    Eigen::Matrix<precision, 3, 1> angles = qref_.toEulerAngles();
    roll = angles(0);
    pitch = angles(1);
    yaw = angles(2);
}

template <typename precision>
Quaternion<precision> ESKF<precision>::getAttitude() const
{
    return qref_;
}

template <typename precision>
void ESKF<precision>::getAcceleration(float &x, float &y, float &z) const
{
    x = x_(6);
    y = x_(7);
    z = x_(8);
}
} // namespace IMU_EKF

#endif // ESKF_IMPL


//
//  Family.hpp
//  test_rcpp
//
//  Created by Dongjie Wu on 18/11/2021.
//

#ifndef Family_hpp
#define Family_hpp

#include <stdio.h>
#include <RcppArmadillo.h>
//#include <boost/math/distributions/normal.hpp>
//#include <boost/multiprecision/cpp_bin_float.hpp>
//#include "distribution.hpp"
//using namespace boost::math;
using namespace Rcpp;

// [[Rcpp::depends(RcppArmadillo)]]

// !! need to decide provide the option of not log
template <class T>
class Family {
public:
  arma::mat logLik(const arma::vec& theta, const arma::mat& Y,
                   const arma::mat& X, const bool& lg = false) {
    return(static_cast<T*>(this) -> logLik(theta, Y, X, lg));
  }

};

class FamilyNormal : public Family<FamilyNormal> {
public:
  arma::mat logLik(const arma::vec& theta, const arma::mat& Y,
                   const arma::mat& X, const bool& lg = false)
  {
    arma::mat mean_mat = (X * theta.subvec(1,(theta.n_elem-1)));
    double sd = sqrt(theta[0] * theta[0]);
    arma::mat l = arma::zeros(Y.n_rows, Y.n_cols);
    for (int i=0; i<Y.n_rows;++i) {
      for (int j=0; j<Y.n_cols;++j) {
          //normal_distribution nd(mean_mat.at(i,j),sd);
          l.at(i,j) = R::dnorm4(Y.at(i,j), mean_mat.at(i,j), sd, lg);
      }
    }
    return l;
  }
};

class FamilyPoisson : public Family<FamilyPoisson> {
public:
  arma::mat logLik(const arma::vec& theta, const arma::mat& Y,
                   const arma::mat& X, const bool& lg = false)
  {
    arma::mat mean_mat = arma::exp(X * theta); // the precision problem
    arma::mat l = arma::zeros(Y.n_rows, Y.n_cols);

    for (int i=0; i<Y.n_rows;++i) {
      for (int j=0; j<Y.n_cols;++j) {
          l.at(i,j) = R::dpois(Y.at(i,j), mean_mat.at(i,j), lg);
      }
    }
    return l;
  }
};

class FamilyLogit : public Family<FamilyLogit> {
public:
  arma::mat sigmoid(arma::mat x) {
          return (1/(1+arma::exp(-x)));
      };
  arma::mat logLik(const arma::vec& theta, const arma::mat& Y,
                   const arma::mat& X, const bool& lg = false)
  {
//      arma::mat mean_t = (X * theta);
//      arma::mat mean_mat = log(sigmoid(mean_t));
//      arma::mat l = arma::zeros(Y.n_rows, Y.n_cols);
//      for (int i=0; i<Y.n_rows;i++) {
//        for (int j=0; j<Y.n_cols; j++) {
//           l.at(i,j) = R::dbinom(Y.at(i,j), 1, mean_mat.at(i,j), lg);
//        }
//      }
//      return l;
    arma::mat mean_t = (X * theta);
    arma::mat l = arma::zeros(Y.n_rows, Y.n_cols);
    l = (Y % mean_t) - arma::log1p(arma::exp(mean_t));
    return l;
  }
};

class FamilyMultiNomial : public Family<FamilyMultiNomial> {
public:
  arma::mat logLik(const arma::vec& theta, const arma::mat& Y,
                   const arma::mat& X, const bool& lg = false)
  {
    arma::mat theta_t = arma::mat(theta);
    theta_t.reshape(X.n_cols, Y.n_cols);
    arma::mat mean_mat = (X * theta_t);
    arma::vec r = arma::sum((Y % mean_mat), 1) - log(1 + arma::sum(arma::exp(mean_mat), 1));
    arma::mat l = arma::mat(r);
    return l;
  }
};

class FamilyConditionalLogit : public Family<FamilyConditionalLogit> {
public:
  arma::mat logLik(const arma::vec& theta, const arma::mat& Y,
                   const arma::mat& X, const bool& lg = false)
 // The group variable is the last variable of X.
  {
    //auto t1 = std::chrono::high_resolution_clock::now();
    arma::uvec gV = arma::conv_to<arma::uvec>::from(X.col(X.n_cols-1));
    arma::uvec g = unique(gV);
    arma::mat l = arma::zeros(g.n_elem, 1);
    arma::mat l1 = arma::zeros(g.n_elem, 1);
    arma::mat l2 = arma::zeros(g.n_elem, 1);
    //auto t2 = std::chrono::high_resolution_clock::now();
    for (arma::uword i=0, nr=Y.n_rows; i<nr; ++i) {
        arma::uword j = X.at(i,X.n_cols-1)-1;
        double mean_x = 0;
        for (arma::uword k=0, nc=X.n_cols; k<nc-1; ++k) {
            mean_x += theta.at(k) * X.at(i, k);
        }
        l1.at(j, 0) += Y.at(i, 0) * mean_x;
        l2.at(j, 0) += exp(mean_x);
    }
    //auto t3 = std::chrono::high_resolution_clock::now();
    l = l1 - log(l2);
    //auto t4 = std::chrono::high_resolution_clock::now();

    /*
    const arma::mat& Xd = X.cols(0, X.n_cols-2);
    auto t2 = std::chrono::high_resolution_clock::now();

    arma::cube Xc1(Y.n_rows, Y.n_cols, 1);
    Xc1.slice(0) = Y % (Xd * theta);
    Xc1.reshape(gV.n_elem/g.n_elem, Y.n_cols, g.n_elem);
    arma::cube Xc2(Y.n_rows, Y.n_cols, 1);
    Xc2.slice(0) = arma::exp(Xd * theta);
    Xc2.reshape(gV.n_elem/g.n_elem, Y.n_cols, g.n_elem);
    auto t3 = std::chrono::high_resolution_clock::now();

    for (arma::uword i=0; i<g.n_elem;++i) {
      l.at(i, 0) = arma::accu(Xc1.slice(i)) -     log(arma::accu(Xc2.slice(i)));
    }
    auto t4 = std::chrono::high_resolution_clock::now();

    std::chrono::duration<double, std::milli> d1 = t2 - t1;
    std::chrono::duration<double, std::milli> d2 = t3 - t2;
    std::chrono::duration<double, std::milli> d3 = t4 - t3;
    std::cout << d1.count() << "ms" << std::endl;
    std::cout << d2.count() << "ms" << std::endl;
    std::cout << d3.count() << "ms" << std::endl;
     */
    return l;
  }
};

class FamilyUnidiff : public Family<FamilyUnidiff> {
public:
  arma::mat logLik(const arma::vec& theta, const arma::mat& Y,
                   const arma::mat& X, const bool& lg = false)
 // The group variable is the last variable of X.
    {
        if (Y.n_cols != 1 ) {
            throw std::invalid_argument("Y should only have 1 column!");
        }
        if (X.n_cols != 2 ) {
            throw std::invalid_argument("X should have 2 columns!");
        }
        arma::vec uY = arma::unique(Y);
        arma::mat X_X = fast_dummy(X.col(0));
        arma::mat X_Z = fast_dummy(X.col(1));
        arma::mat Y_Y = fast_dummy(Y.col(0));
        arma::mat W = arma::ones(X_Z.n_rows, (X_Z.n_cols+1));
        W(arma::span::all,arma::span(1,W.n_cols-1)) = X_Z;
        arma::uword colX = X_X.n_cols;
        arma::uword colZ = X_Z.n_cols;
        arma::uword colY = Y_Y.n_cols;
        if (theta.n_elem != (colY*(colZ+1)+colY*colX+colZ)) {
            throw std::invalid_argument("Wrong size of theta!");
        }
        arma::uword end = colY*(colZ+1)-1;
        arma::uword end2 = colY*colX;
        arma::uword end3 = colZ;
        arma::mat theta_y = arma::mat(theta(arma::span(0, end)));
        theta_y.reshape(colZ+1, colY);
        arma::mat psi_y = arma::mat(theta(arma::span(end+1, end+end2)));
        psi_y.reshape(colX, colY);
        arma::mat phi = arma::mat(theta(arma::span(end+end2+1,end+end2+end3)));
        phi.reshape(colZ,1);
        arma::mat expZ = arma::exp(X_Z * phi);
        arma::mat phiX = X_X  * psi_y;
        arma::mat part2 = phiX.each_col() % expZ;
        arma::mat sum_part = W * theta_y + part2;
        arma::mat first = arma::sum(Y_Y % sum_part, 1);
        arma::mat second = log(1+arma::sum(arma::exp(sum_part), 1));
        arma::mat l = first - second;
        return l;
    }
    arma::mat fast_dummy(const arma::vec& x) {
        arma::vec vec_u = arma::unique(x);
        vec_u = vec_u(arma::span(1,vec_u.n_elem-1)); // Remove the baseline
        arma::mat mats = arma::zeros(x.n_elem, (vec_u.n_elem));
        for (arma::uword i=0; i<vec_u.n_elem; ++i) {
            mats.col(i) = x;
            double idx = vec_u.at(i);
            mats.col(i).for_each([idx](arma::vec::elem_type& val){
                (val == idx) ? (val = 1) : (val = 0);
            });
        }
        return mats;
    }
};

#endif /* Family_hpp */