/*
# Copyright (c) 2014-2020 Reed A. Cartwright <reed@cartwright.ht>
#
# This file is part of the Ultimate Source Code Project.
#
# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the "Software"), to deal
# in the Software without restriction, including without limitation the rights
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
# copies of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:
#
# The above copyright notice and this permission notice shall be included in all
# copies or substantial portions of the Software.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
# SOFTWARE.
*/
#include <doctest/doctest.h>

#include <mutk/mutation.hpp>
#include <mutk/relationship_graph.hpp>

// Libraries needed for testing
#include <boost/numeric/ublas/matrix.hpp>

namespace {
// Structure useful for unit testing
struct kalleles_test_mat {
    using mat_t = boost::numeric::ublas::matrix<float,
        boost::numeric::ublas::column_major,std::vector<float>>;

    mat_t Q;

    kalleles_test_mat(size_t n, float k) : Q(n+1,n+1) {
        using namespace boost::numeric::ublas;
        // Build matrix
        vector<float> freqs(n+1, 1.0/k);
        freqs[n] = 1.0-n/k;

        for(int i=0; i<=n; ++i) {
            for(int j=0; j<=n; ++j) {
                Q(i,j) = freqs[j];
            }
            Q(i,i) = Q(i,i)-1.0;
        }

        // Scale matrix into substitution time
        Q *= k/(k-1);
    }
};

template<typename F>
void run_mutation_tests(F test) {
    test(2, 0.0,  4.0);
    test(4, 0.0,  4.0);

    test(1, 1e-8, 4.0);
    test(2, 1e-8, 4.0);
    test(3, 1e-8, 4.0);
    test(4, 1e-8, 4.0);

    test(4, 1e-8, 4.0);
    test(4, 1e-9, 5.0);
    test(4, 1e-6, 6.0);
    test(4, 1e-3, 7.0);
}
} // anon namespace

namespace {
constexpr int ALLELE[][2] = {
    {0,0},
    {0,1}, {1,1},
    {0,2}, {1,2}, {2,2},
    {0,3}, {1,3}, {2,3}, {3,3},
    {0,4}, {1,4}, {2,4}, {3,4}, {4,4},
};
}

using KAllelesModel = mutk::mutation::KAllelesModel;
using potential_t = mutk::RelationshipGraph::potential_t;

// LCOV_EXCL_START
TEST_CASE("KAllelesModel-Constructor") {
    CHECK_NOTHROW(KAllelesModel(4.0, 0.001, 0.0, 0.0, 0.0));
    CHECK_THROWS_AS(KAllelesModel(1.0, 0.001, 0.0, 0.0, 0.0), std::invalid_argument);
    CHECK_THROWS_AS(KAllelesModel(4.0, -0.1, 0.0, 0.0, 0.0), std::invalid_argument);
    CHECK_THROWS_AS(KAllelesModel(4.0, 0.001, 1.1, 0.0, 0.0), std::invalid_argument);
    CHECK_THROWS_AS(KAllelesModel(4.0, 0.001, -3000, 0.0, 0.0), std::invalid_argument);
    CHECK_THROWS_AS(KAllelesModel(4.0, 0.001, 0.0, 1.1, 0.0), std::invalid_argument);
    CHECK_THROWS_AS(KAllelesModel(4.0, 0.001, 0.0, -4000, 0.0), std::invalid_argument);
    CHECK_THROWS_AS(KAllelesModel(4.0, 0.001, 0.0, 0.0, 1.1), std::invalid_argument);
    CHECK_THROWS_AS(KAllelesModel(4.0, 0.001, 0.0, 0.0, -3000), std::invalid_argument);
}
// LCOV_EXCL_STOP

// ret(i,j) = P(i|j)
KAllelesModel::tensor_t KAllelesModel::CreateMatrix(size_t n, float t, any_t) const {
    assert(n > 0);
    assert(n <= 5);

    double beta = t*k_/(k_-1.0);
    double p_ji = -1.0/k_*expm1(-beta);
    double p_jj = exp(-beta) + p_ji;

    tensor_t ret = tensor_t::from_shape({n,n});

    for(size_t i = 0; i < n; ++i) {
        for(size_t j = 0; j < n; ++j) {
            ret(i,j) = (i == j) ? p_jj : p_ji;
        }
    }
    return ret;
}

// LCOV_EXCL_START
TEST_CASE("KAllelesModel-CreateMatrix with any_t") {
    using namespace boost::numeric::ublas;

    auto test = [&](size_t n, float u, float k) {
        CAPTURE(n);
        CAPTURE(k);
        CAPTURE(u);

        KAllelesModel model(k, 0.001, 0, 0, 0);

        using mat_t = kalleles_test_mat::mat_t;

        kalleles_test_mat mat(n,k);

        mat_t P = identity_matrix<float>(n+1);
        mat_t m = P;
        float t = 1.0;
        float f = 1.0;
        // Use the first 11 elements of the expm series
        for(int g=1;g <= 11; ++g) {
            mat_t a = prec_prod(m,mat.Q);
            m = a;
            t *= u;
            f *= g;
            P += m * (t/f);
        }
        auto obs = model.CreateMatrix(n, u, mutk::mutation::ANY);
        REQUIRE(obs.dimension() == 2);
        REQUIRE(obs.shape(0) == n);
        REQUIRE(obs.shape(1) == n);
        for(int i=0; i < n; ++i) {
            for(int j=0; j < n; ++j) {
                CAPTURE(i);
                CAPTURE(j);
                CHECK(obs(i,j) == doctest::Approx(P(i,j)));
            }
        }
    };
    run_mutation_tests(test);
}
// LCOV_EXCL_STOP

// ret(i,j) = P(i & x mutations | j)
KAllelesModel::tensor_t KAllelesModel::CreateMatrix(size_t n, float t, size_t x) const {
    assert(n > 0);
    assert(n <= 5);
    assert(x >= 0);

    double p_x;
    if(t == 0.0) {
        p_x = (x==0) ? 1.0 : 0.0;
    } else {
        p_x = exp(-t+x*log(t)-lgamma(x+1));
    }

    double p_ji = (1.0-pow(-1.0/(k_-1.0),x))/k_;
    double p_jj = (1.0+(k_-1.0)*pow(-1.0/(k_-1.0),x))/k_;

    tensor_t ret = tensor_t::from_shape({n,n});

    for(size_t i = 0; i < n; ++i) {
        for(size_t j = 0; j < n; ++j) {
            ret(i,j) = (i == j) ? p_x*p_jj : p_x*p_ji;
        }
    }
    return ret;
}

// LCOV_EXCL_START
TEST_CASE("KAllelesModel-CreateMatrix with int") {
    using namespace boost::numeric::ublas;

    auto test = [&](size_t n, float u, float k) {
        CAPTURE(n);
        CAPTURE(k);
        CAPTURE(u);

        KAllelesModel model(k, 0.001, 0, 0, 0);

        using mat_t = kalleles_test_mat::mat_t;

        kalleles_test_mat mat(n,k);

        mat_t P = identity_matrix<float>(n+1);
        mat_t m = P;
        mat_t J = mat.Q+m;
        float t;
        float f;
        for(int x=0; x <= 10; ++x) {
            CAPTURE(x);
            if(x == 0) {
                P = m*exp(-u);
                f = 1.0;
                t = 1.0;
            } else {
                mat_t a = prec_prod(m,J);
                m = a;
                t *= u;
                f *= x;
                P = m*(exp(-u)*t/f);
            }
            auto obs = model.CreateMatrix(n, u, x);
            REQUIRE(obs.dimension() == 2);
            REQUIRE(obs.shape(0) == n);
            REQUIRE(obs.shape(1) == n);
            for(int i=0; i < n; ++i) {
                for(int j=0; j < n; ++j) {
                    CAPTURE(i);
                    CAPTURE(j);
                    CHECK(obs(i,j) == doctest::Approx(P(i,j)));
                }
            }

        }
    };

    run_mutation_tests(test);
}
// LCOV_EXCL_STOP

// ret(i,j) = E[num of mutations | i,j]*P(i|j)
KAllelesModel::tensor_t KAllelesModel::CreateMatrix(size_t n, float t, mean_t) const {
    assert(n > 0);
    assert(n <= 5);

    double beta = k_*t/(k_-1.0);
    double p_jj = -t/k_*expm1(-beta);
    double p_ji = (t-p_jj)/(k_-1.0);

    tensor_t ret = tensor_t::from_shape({n,n});

    for(size_t i = 0; i < n; ++i) {
        for(size_t j = 0; j < n; ++j) {
            ret(i,j) = (i == j) ? p_jj : p_ji;
        }
    }
    return ret;
}

// LCOV_EXCL_START
TEST_CASE("KAllelesModel-CreateMatrix with mean_t") {
    using namespace boost::numeric::ublas;

    auto test = [&](size_t n, float u, float k) {
        CAPTURE(n);
        CAPTURE(k);
        CAPTURE(u);

        KAllelesModel model(k, 0.001, 0, 0, 0);

        using mat_t = kalleles_test_mat::mat_t;

        kalleles_test_mat mat(n,k);

        mat_t P = identity_matrix<float>(n+1);
        mat_t m = P;
        mat_t J = mat.Q+m;
        mat_t S;
        float t;
        float f;
        for(int x=0; x <= 10; ++x) {
            if(x == 0) {
                P = m*exp(-u);
                S = P*x;
                f = 1.0;
                t = 1.0;
            } else {
                mat_t a = prec_prod(m,J);
                m = a;
                t *= u;
                f *= x;
                P = m*(exp(-u)*t/f);
                S += P*x;
            }
        }
        auto obs = model.CreateMatrix(n, u, mutk::mutation::MEAN);
        REQUIRE(obs.dimension() == 2);
        REQUIRE(obs.shape(0) == n);
        REQUIRE(obs.shape(1) == n);
        for(int i=0; i < n; ++i) {
            for(int j=0; j < n; ++j) {
                CAPTURE(i);
                CAPTURE(j);
                CHECK(obs(i,j) == doctest::Approx(S(i,j)));
            }
        }
    };

    run_mutation_tests(test);
}
// LCOV_EXCL_STOP

KAllelesModel::tensor_t KAllelesModel::CreatePriorHaploid(size_t n) const {
    double k = k_;
    double e = theta_/(k-1.0);

    double p_R = (1.0+e+(k-1.0)*e*hap_bias_)/(1.0+k*e);
    double p_A = (e-e*hap_bias_)/(1.0+k*e);

    tensor_t ret = tensor_t::from_shape({n});

    for(size_t i = 0; i < n; ++i) {
        ret(i) = (i == 0) ? p_R : p_A;
    }
    return ret;
}

// LCOV_EXCL_START
TEST_CASE("KAllelesModel-CreatePriorHaploid") {
    auto test_haploid = [](size_t n, float theta, float hap_bias,
        float k) {
        CAPTURE(n);
        CAPTURE(theta);
        CAPTURE(hap_bias);
        CAPTURE(k);

        KAllelesModel model(k, theta, 0, 0, hap_bias);

        auto obs = model.CreatePriorHaploid(n);

        if(n == k) {
            float s = xt::sum(obs)();
            CHECK(s == doctest::Approx(1.0f));
        }

        const int sz = n;
        for(int a=0;a<sz;++a) {
            CAPTURE(a);
            float e = theta/(k-1.0);
            float r;
            if(a == 0) {
                // reference haploid
                r = (1.0+e+e*(k-1)*hap_bias)/(1.0+k*e);
            } else {
                r = (e-e*hap_bias)/(1.0+k*e);
            }
            float expected = r;
            CHECK(obs(a) == doctest::Approx(expected));
        }
    };

    auto test = [&](float theta, float hap_bias, float k) {
        for(int i=1;i<=5;++i) {
            test_haploid(i, theta, hap_bias, k);
        }
    };

    test(0.001, 0.0, 5.0);
    test(0.01, 1.0, 6.0);
    test(0.1, 0.0, 4.0);
    test(0.001, 1.0, 5.0);
    test(0.001, -1.0, 5.5);
}
// LCOV_EXCL_STOP

KAllelesModel::tensor_t KAllelesModel::CreatePriorDiploid(size_t n) const {
    double k = k_;
    double e = theta_/(k-1.0);

    double p_hom = (1.0+e)/(1.0+k*e);
    double p_hetk = e/(1.0+k*e);

    double p_RR = p_hom*(2.0+e+(k-1.0)*e*hom_bias_)/(2.0+k*e);
    double p_AA = p_hom*(e-e*hom_bias_)/(2.0+k*e);

    double p_RA = p_hetk*(2.0+2.0*e+(k-2.0)*e*het_bias_)/(2.0+k*e);
    double p_AB = p_hetk*(2.0*e-2.0*e*het_bias_)/(2.0+k*e);

    tensor_t ret = tensor_t::from_shape({mutk::dim_width<2>(n)});

    for(size_t i = 0; i < ret.size(); ++i) {
        auto a = ALLELE[i][0];
        auto b = ALLELE[i][1];
        if(a == b) {
            ret(i) = (a == 0) ? p_RR : p_AA;
        } else {
            ret(i) = (a == 0) ? p_RA : p_AB;
        }
    }
    return ret;
}

// LCOV_EXCL_START
TEST_CASE("KAllelesModel-CreatePriorDiploid") {
    auto test_diploid = [](size_t n, float theta, float hom_bias,
        float het_bias, float k) {
        CAPTURE(n);
        CAPTURE(theta);
        CAPTURE(hom_bias);
        CAPTURE(het_bias);
        CAPTURE(k);

        KAllelesModel model(k, theta, hom_bias, het_bias, 0);

        auto obs = model.CreatePriorDiploid(n);

        if(n == k) {
            float s = xt::sum(obs)();
            CHECK(s == doctest::Approx(1.0f));
        }

        const int sz = n;
        int y = 0;
        for(int a=0;a<sz;++a) {
            for(int b=0;b<=a;++b) {
                CAPTURE(a);
                CAPTURE(b);
                float e = theta/(k-1.0);

                double r = 0.0;
                if(a == 0 && b == 0) {
                    // Reference Homozygote
                    r = (1.0+e)/(1.0+k*e)*(2.0+e+(k-1.0)*e*hom_bias)/(2.0+k*e);
                } else if(a == b) {
                    r = (1.0+e)/(1.0+k*e)*(e-e*hom_bias)/(2.0+k*e);
                } else if(b == 0 || a == 0) {
                    // Reference Heterozygote
                    r = e/(1.0+k*e)*(2.0+2.0*e+(k-2.0)*e*het_bias)/(2.0+k*e);
                } else {
                    // Alt heterozygote
                    r = e/(1.0+k*e)*(2.0*e-2.0*e*het_bias)/(2.0+k*e);
                }
                float expected = r;
                CHECK(obs(y) == doctest::Approx(expected));
                ++y;
            }
        }
    };

    auto test = [&](float theta, float hom_bias, float het_bias, float k) {
        for(int i=1;i<=5;++i) {
            test_diploid(i, theta, hom_bias, het_bias, k);
        }
    };

    test(0.001, 0.0, 0.0, 5.0);
    test(0.01, 1.0, 0.0, 6.0);
    test(0.1, 0.0, 1.0, 4.0);
    test(0.001, 1.0, 1.0, 5.0);
    test(0.001, -1.0, -1.0, 5.5);
}
// LCOV_EXCL_STOP

template<typename Arg>
mutk::tensor_t create_transition_clone_haploid_impl(const mutk::mutation::Model &model,
    size_t n, float t, Arg arg) {

    return model.CreateMatrix(n, t, arg);
}

template<typename Arg>
mutk::tensor_t create_transition_gamete_diploid_impl(const mutk::mutation::Model &model,
    size_t n, float t, Arg arg) {

    using tensor_t = mutk::tensor_t;

    auto mat = create_transition_clone_haploid_impl(model, n, t, arg);
    tensor_t ret = tensor_t::from_shape({mutk::dim_width<1>(n),mutk::dim_width<2>(n)});

    // x/y -> a
    for(size_t u = 0; u < ret.size(); ++u) {
        auto s = xt::unravel_index(u, ret.shape());
        auto a = s[0];
        auto x = ALLELE[s[1]][0];
        auto y = ALLELE[s[1]][1];
        ret(u) = 0.5*(mat(a,x)+mat(a,y));
    }
    return ret;
}

template<int N, typename Arg>
inline
mutk::tensor_t create_transition_haploid_impl(const mutk::mutation::Model &model,
    size_t n, float t, Arg arg) {
    if constexpr(N == 1) {
        return create_transition_clone_haploid_impl(model,n,t,arg);
    } else if constexpr(N == 2) {
        return create_transition_gamete_diploid_impl(model,n,t,arg);
    } else {
        static_assert(N == 1 || N == 2);
    }
}

struct clone_diploid_tag {};

template<int N, int M>
struct child_diploid_tag {};

template<int N>
struct child_selfing_tag {};

template<typename A>
struct transition_traits;

// G1/G2 => a/b
template<>
struct transition_traits<clone_diploid_tag> {
    static constexpr int MAT1_SIZE = 1;
    static constexpr int MAT2_SIZE = 1;
    static constexpr int TENSOR_SIZE = 2;

    static mutk::shape_t Dims(size_t n) {
        return {mutk::dim_width<2>(n), mutk::dim_width<2>(n)};
    }
    static size_t G1(const mutk::strides_t &coords) {
        return ALLELE[coords[1]][0];
    }
    static size_t G2(const mutk::strides_t &coords) {
        return ALLELE[coords[1]][1];
    }
};

// G1 x G2 => a/b
template<int N, int M>
struct transition_traits<child_diploid_tag<N, M>> {
    static constexpr int MAT1_SIZE = N;
    static constexpr int MAT2_SIZE = M;
    static constexpr int TENSOR_SIZE = 3;

    static mutk::shape_t  Dims(size_t n) {
        return {mutk::dim_width<2>(n), mutk::dim_width<N>(n), mutk::dim_width<M>(n)};
    }
    static size_t G1(const mutk::strides_t &coords) {
        return coords[1];
    }
    static size_t G2(const mutk::strides_t &coords) {
        return coords[2];
    }
};

// G1 x G1 => a/b
template<int N>
struct transition_traits<child_selfing_tag<N>> {
    static constexpr int MAT1_SIZE = N;
    static constexpr int MAT2_SIZE = N;
    static constexpr int TENSOR_SIZE = 2;

    static mutk::shape_t Dims(size_t n) {
        return {mutk::dim_width<2>(n), mutk::dim_width<N>(n)};
    }
    static size_t G1(const mutk::strides_t &coords) {
        return coords[1];
    }
    static size_t G2(const mutk::strides_t &coords) {
        return coords[1];
    }
};

template<typename Tag>
struct generate_transition_diploid {
    using traits = transition_traits<Tag>;

    template<typename Arg1, typename Arg2>
    inline
    mutk::tensor_t
    operator()(const mutk::mutation::Model &model,
        size_t n, float t1, float t2, Arg1 arg1, Arg2 arg2) const {

        auto mat1 = create_transition_haploid_impl<traits::MAT1_SIZE>(model, n, t1, arg1);
        auto mat2 = create_transition_haploid_impl<traits::MAT2_SIZE>(model, n, t2, arg2);

        // G1/G2 => a/b OR G1 x G2 => a/b
        mutk::tensor_t ret = mutk::tensor_t::from_shape(traits::Dims(n));
        for(size_t u = 0; u < ret.size(); ++u) {
            auto s = xt::unravel_index(u, ret.shape());
            size_t a = ALLELE[s[0]][0];
            size_t b = ALLELE[s[0]][1];
            size_t g1 = traits::G1(s);
            size_t g2 = traits::G2(s);

            float f = mat1(a,g1) * mat2(b,g2);
            if(a != b) {
                f += mat1(b,g1) * mat2(a,g2);
            }
            ret(u) = f;
        }
        return ret;
    }
};

template<typename Tag, typename Arg>
auto create_transition_diploid_impl(const mutk::mutation::Model &model,
    size_t n, float t1, float t2, Arg arg) {
    generate_transition_diploid<Tag> gen;

    if constexpr(std::is_same_v<Arg, mutk::mutation::any_t>) {
        return gen(model, n, t1, t2, arg, arg);
    } else if constexpr(std::is_same_v<Arg, mutk::mutation::mean_t>) {
        auto ret = gen(model, n, t1, t2, arg, mutk::mutation::ANY);
        ret += gen(model, n, t1, t2, mutk::mutation::ANY, arg);
        return ret;
    } else if constexpr(std::is_integral_v<Arg>) {
        auto ret = gen(model, n, t1, t2, 0, arg);
        for(Arg a=1; a <= arg; ++a) {
            ret += gen(model, n, t1, t2, a, arg-a);
        }
        return ret;
    } else {
        static_assert(std::is_integral_v<Arg>
            || std::is_same_v<Arg, mutk::mutation::mean_t>
            || std::is_same_v<Arg, mutk::mutation::any_t> );
    }
}

template<int N, int M, typename Arg>
inline
mutk::tensor_t create_transition_child(const mutk::mutation::Model &model,
    size_t n, const potential_t &potential, Arg arg) {
    assert(potential.parents.size() == 2);
    assert(potential.axes.size() == 3);
    auto ret = create_transition_diploid_impl<child_diploid_tag<N,M>>(model, n, potential.parents[0].second,
        potential.parents[1].second, arg);
    return xt::transpose(ret, potential.axes);
}

template<int N, typename Arg>
inline
mutk::tensor_t create_transition_child_selfing(const mutk::mutation::Model &model,
    size_t n, const potential_t &potential, Arg arg) {
    assert(potential.parents.size() == 1);
    assert(potential.axes.size() == 2);
    auto ret = create_transition_diploid_impl<child_selfing_tag<N>>(model, n, potential.parents[0].second,
        potential.parents[0].second, arg);
    return xt::transpose(ret, potential.axes);
}

template<typename Arg>
mutk::tensor_t create_transition_clone_diploid(const mutk::mutation::Model &model,
    size_t n, const potential_t &potential, Arg arg) {
    assert(potential.parents.size() == 1);
    assert(potential.axes.size() == 2);
    auto ret = create_transition_diploid_impl<clone_diploid_tag>(model, n, potential.parents[0].second,
        potential.parents[0].second, arg);
    return xt::transpose(ret, potential.axes);
}

// LCOV_EXCL_START
TEST_CASE("create_transition_clone_diploid") {
    SUBCASE("Passing mutk::mutation::ANY as argument.") {
        auto test = [&](size_t n, float u, float k, std::vector<int> axes) {
            CAPTURE(n);
            CAPTURE(k);
            CAPTURE(u);

            potential_t pot{mutk::detail::Potential::CloneDiploid, 0, 1, u};
            pot.axes = axes;
            KAllelesModel model{k, 0.001, 0, 0, 0};
            auto obs = create_transition_clone_diploid(model, n, pot, mutk::mutation::ANY);

            auto mat = model.CreateMatrix(n, u, mutk::mutation::ANY);

            REQUIRE(obs.dimension() == 2);
            mutk::shape_t pot_shape = {n*(n+1)/2, n*(n+1)/2};
            REQUIRE(obs.shape(0) == pot_shape[pot.axes[0]]);
            REQUIRE(obs.shape(1) == pot_shape[pot.axes[1]]);
            for(int a1=0,a=0;a1<n;++a1) {
                for(int a2=0;a2<=a1;++a2,++a) {
                    for(int b1=0,b=0;b1<n;++b1) {
                        for(int b2=0;b2<=b1;++b2,++b) {
                            CAPTURE(a1);
                            CAPTURE(a2);
                            CAPTURE(a);
                            CAPTURE(b1);
                            CAPTURE(b2);
                            CAPTURE(b);
                            // b1/b2 -> a1/a2
                            float expected = mat(a1,b1)*mat(a2,b2);
                            if(a1 != a2) {
                                expected += mat(a2,b1)*mat(a1,b2);
                            }
                            std::array<int, 2> indexes = {a,b};
                            int o1 = indexes[pot.axes[0]];
                            int o2 = indexes[pot.axes[1]];
                            CHECK(obs(o1,o2) == doctest::Approx(expected));
                        }
                    }
                }
            }
        };
        SUBCASE("Axes are {0, 1}.") {
            auto subtest = [&](size_t n, float u, float k) {
                return test(n, u, k, {0,1});
            };
            run_mutation_tests(subtest);
        }
        SUBCASE("Axes are {1, 0}.") {
            auto subtest = [&](size_t n, float u, float k) {
                return test(n, u, k, {1,0});
            };
            run_mutation_tests(subtest);
        }
    }
    SUBCASE("Passing mutk::mutation::MEAN as argument.") {
        auto test = [&](size_t n, float u, float k, std::vector<int> axes) {
            CAPTURE(n);
            CAPTURE(k);
            CAPTURE(u);

            potential_t pot{mutk::detail::Potential::CloneDiploid, 0, 1, u};
            pot.axes = axes;
            KAllelesModel model{k, 0.001, 0, 0, 0};
            auto obs = create_transition_clone_diploid(model, n, pot, mutk::mutation::MEAN);

            auto mat1 = model.CreateMatrix(n, u, mutk::mutation::MEAN);
            auto mat2 = model.CreateMatrix(n, u, mutk::mutation::ANY);

            REQUIRE(obs.dimension() == 2);
            mutk::shape_t pot_shape = {n*(n+1)/2, n*(n+1)/2};
            REQUIRE(obs.shape(0) == pot_shape[pot.axes[0]]);
            REQUIRE(obs.shape(1) == pot_shape[pot.axes[1]]);
            for(int a1=0,a=0;a1<n;++a1) {
                for(int a2=0;a2<=a1;++a2,++a) {
                    for(int b1=0,b=0;b1<n;++b1) {
                        for(int b2=0;b2<=b1;++b2,++b) {
                            CAPTURE(a1);
                            CAPTURE(a2);
                            CAPTURE(a);
                            CAPTURE(b1);
                            CAPTURE(b2);
                            CAPTURE(b);
                            // b1/b2 -> a1/a2
                            float expected = mat1(a1,b1)*mat2(a2,b2) + mat2(a1,b1)*mat1(a2,b2);
                            if(a1 != a2) {
                                expected += mat1(a2,b1)*mat2(a1,b2) + mat2(a2,b1)*mat1(a1,b2);
                            }
                            std::array<int, 2> indexes = {a,b};
                            int o1 = indexes[pot.axes[0]];
                            int o2 = indexes[pot.axes[1]];
                            CHECK(obs(o1,o2) == doctest::Approx(expected));
                        }
                    }
                }
            }
        };
        SUBCASE("Axes are {0, 1}.") {
            auto subtest = [&](size_t n, float u, float k) {
                return test(n, u, k, {0,1});
            };
            run_mutation_tests(subtest);
        }
        SUBCASE("Axes are {1, 0}.") {
            auto subtest = [&](size_t n, float u, float k) {
                return test(n, u, k, {1,0});
            };
            run_mutation_tests(subtest);
        }
    }
    SUBCASE("Passing an integer as argument.") {
        auto test = [&](size_t n, float u, float k, std::vector<int> axes, int val) {
            CAPTURE(n);
            CAPTURE(k);
            CAPTURE(u);

            potential_t pot{mutk::detail::Potential::CloneDiploid, 0, 1, u};
            pot.axes = axes;
            KAllelesModel model{k, 0.001, 0, 0, 0};
            auto obs = create_transition_clone_diploid(model, n, pot, val);

            REQUIRE(obs.dimension() == 2);
            mutk::shape_t pot_shape = {n*(n+1)/2, n*(n+1)/2};
            REQUIRE(obs.shape(0) == pot_shape[pot.axes[0]]);
            REQUIRE(obs.shape(1) == pot_shape[pot.axes[1]]);
            for(int a1=0,a=0;a1<n;++a1) {
                for(int a2=0;a2<=a1;++a2,++a) {
                    for(int b1=0,b=0;b1<n;++b1) {
                        for(int b2=0;b2<=b1;++b2,++b) {
                            CAPTURE(a1);
                            CAPTURE(a2);
                            CAPTURE(a);
                            CAPTURE(b1);
                            CAPTURE(b2);
                            CAPTURE(b);
                            // b1/b2 -> a1/a2
                            float expected = 0.0;
                            for(int i=0;i<=val;++i) {
                                auto mat1 = model.CreateMatrix(n, u, i);
                                auto mat2 = model.CreateMatrix(n, u, val-i);
                                expected += mat1(a1,b1)*mat2(a2,b2);
                                if(a1 != a2) {
                                    expected += mat1(a2,b1)*mat2(a1,b2);
                                }
                            }
                            std::array<int, 2> indexes = {a,b};
                            int o1 = indexes[pot.axes[0]];
                            int o2 = indexes[pot.axes[1]];
                            CHECK(obs(o1,o2) == doctest::Approx(expected));
                        }
                    }
                }
            }
        };
        SUBCASE("Axes are {0, 1}. Integer is 0.") {
            auto subtest = [&](size_t n, float u, float k) {
                return test(n, u, k, {0,1}, 0);
            };
            run_mutation_tests(subtest);
        }
        SUBCASE("Axes are {1, 0}. Integer is 1.") {
            auto subtest = [&](size_t n, float u, float k) {
                return test(n, u, k, {1,0}, 1);
            };
            run_mutation_tests(subtest);
        }
    }
}
// LCOV_EXCL_STOP

template<typename Arg>
mutk::tensor_t create_transition_clone_haploid(const mutk::mutation::Model &model,
    size_t n, const potential_t &potential, Arg arg) {
    assert(potential.parents.size() == 1);
    assert(potential.axes.size() == 2);
    auto ret =  create_transition_haploid_impl<1>(model, n, potential.parents[0].second, arg);
    return xt::transpose(ret, potential.axes);
}

// LCOV_EXCL_START
TEST_CASE("create_transition_clone_haploid") {
    auto test = [&](size_t n, float u, float k, std::vector<int> axes) {
        CAPTURE(n);
        CAPTURE(k);
        CAPTURE(u);

        potential_t pot{mutk::detail::Potential::CloneHaploid, 0, 1, u};
        pot.axes = axes;
        KAllelesModel model{k, 0.001, 0, 0, 0};
        auto obs = create_transition_clone_haploid(model, n, pot, mutk::mutation::ANY);

        auto mat = model.CreateMatrix(n, u, mutk::mutation::ANY);
        REQUIRE(obs.dimension() == 2);
        mutk::shape_t pot_shape = {n,n};
        REQUIRE(obs.shape(0) == pot_shape[pot.axes[0]]);
        REQUIRE(obs.shape(1) == pot_shape[pot.axes[1]]);
        for(int a=0;a<n;++a) {
            for(int b=0;b<n;++b) {
                CAPTURE(a);
                CAPTURE(b);
                // b -> a
                float expected = mat(a,b);
                std::array<int, 2> indexes = {a,b};
                int o1 = indexes[pot.axes[0]];
                int o2 = indexes[pot.axes[1]];
                CHECK(obs(o1,o2) == doctest::Approx(expected));
            }
        }
    };
    SUBCASE("Axes are {0, 1}.") {
        auto subtest = [&](size_t n, float u, float k) {
            return test(n, u, k, {0,1});
        };
        run_mutation_tests(subtest);
    }
    SUBCASE("Axes are {1, 0}.") {
        auto subtest = [&](size_t n, float u, float k) {
            return test(n, u, k, {1,0});
        };
        run_mutation_tests(subtest);
    }
}
// LCOV_EXCL_STOP

template<typename Arg>
mutk::tensor_t create_transition_gamete_diploid(const mutk::mutation::Model &model,
    size_t n, const potential_t &potential, Arg arg) {
    assert(potential.parents.size() == 1);
    auto ret = create_transition_haploid_impl<2>(model, n, potential.parents[0].second, arg);
    return xt::transpose(ret, potential.axes);
}

// LCOV_EXCL_START
TEST_CASE("create_transition_gamete_diploid") {
    auto test = [&](size_t n, float u, float k, std::vector<int> axes) {
        CAPTURE(n);
        CAPTURE(k);
        CAPTURE(u);

        potential_t pot{mutk::detail::Potential::GameteDiploid, 0, 1, u};
        pot.axes = axes;
        KAllelesModel model(k, 0.001, 0, 0, 0);

        auto obs = create_transition_gamete_diploid(model, n, pot, mutk::mutation::ANY);
        auto mat = model.CreateMatrix(n, u, mutk::mutation::ANY);
        REQUIRE(obs.dimension() == 2);
        mutk::shape_t pot_shape = {n, n*(n+1)/2};
        REQUIRE(obs.shape(0) == pot_shape[pot.axes[0]]);
        REQUIRE(obs.shape(1) == pot_shape[pot.axes[1]]);
        for(int a=0;a<n;++a) {
            for(int b1=0,b=0;b1<n;++b1) {
                for(int b2=0;b2<=b1;++b2,++b) {
                    CAPTURE(a);
                    CAPTURE(b);
                    CAPTURE(b1);
                    CAPTURE(b2);
                    // b1/b2 -> a
                    float expected = 0.5*mat(a,b1) + 0.5*mat(a,b2);
                    std::array<int, 2> indexes = {a,b};
                    int o1 = indexes[pot.axes[0]];
                    int o2 = indexes[pot.axes[1]];
                    CHECK(obs(o1,o2) == doctest::Approx(expected));
                }
            }
        }
    };
    SUBCASE("Axes are {0, 1}.") {
        auto subtest = [&](size_t n, float u, float k) {
            return test(n, u, k, {0,1});
        };
        run_mutation_tests(subtest);
    }
    SUBCASE("Axes are {1, 0}.") {
        auto subtest = [&](size_t n, float u, float k) {
            return test(n, u, k, {1,0});
        };
        run_mutation_tests(subtest);
    }
}
// LCOV_EXCL_STOP

template<typename Arg>
mutk::tensor_t create_transition_child_diploid_diploid(const mutk::mutation::Model &model,
    size_t n, const potential_t &potential, Arg arg) {
    assert(potential.parents.size() == 2);
    // don't shuffle here because the subfunction does it for us.
    return create_transition_child<2, 2>(model, n, potential, arg);
}

// LCOV_EXCL_START
TEST_CASE("create_transition_child_diploid_diploid") {
    auto test = [&](size_t n, float u, float k, std::vector<int> axes) {
        CAPTURE(n);
        CAPTURE(k);
        CAPTURE(u);

        potential_t pot{mutk::detail::Potential:: ChildDiploidDiploid, 0, 1, u, 2, 1.1f*u};
        pot.axes = axes;
        KAllelesModel model(k, 0.001, 0, 0, 0);

        auto obs = create_transition_child_diploid_diploid(model, n, pot, mutk::mutation::ANY);
        auto mat1 = model.CreateMatrix(n, u, mutk::mutation::ANY);
        auto mat2 = model.CreateMatrix(n, 1.1f*u, mutk::mutation::ANY);

        REQUIRE(obs.dimension() == 3);
        mutk::shape_t pot_shape = {n*(n+1)/2, n*(n+1)/2, n*(n+1)/2};
        REQUIRE(obs.shape(0) == pot_shape[pot.axes[0]]);
        REQUIRE(obs.shape(1) == pot_shape[pot.axes[1]]);
        REQUIRE(obs.shape(2) == pot_shape[pot.axes[2]]);
        for(int a1=0,a=0;a1<n;++a1) {
            for(int a2=0;a2<=a1;++a2,++a) {
                for(int b1=0,b=0;b1<n;++b1) {
                    for(int b2=0;b2<=b1;++b2,++b) {
                        for(int c1=0,c=0;c1<n;++c1) {
                            for(int c2=0;c2<=c1;++c2,++c) {
                                CAPTURE(a1);
                                CAPTURE(a2);
                                CAPTURE(a);
                                CAPTURE(b1);
                                CAPTURE(b2);
                                CAPTURE(b);
                                CAPTURE(c1);
                                CAPTURE(c2);
                                CAPTURE(c);
                                // b1/b2 x c1/c2 -> a1/a2
                                float expected1 = 0.5f*mat1(a1,b1) + 0.5f*mat1(a1,b2);
                                float expected2 = 0.5f*mat2(a2,c1) + 0.5f*mat2(a2,c2);
                                float expected = expected1*expected2;
                                if(a1 != a2) {
                                    expected1 = 0.5f*mat1(a2,b1) + 0.5f*mat1(a2,b2);
                                    expected2 = 0.5f*mat2(a1,c1) + 0.5f*mat2(a1,c2);
                                    expected += expected1*expected2;
                                }

                                std::array<int, 3> indexes = {a,b,c};
                                int o1 = indexes[pot.axes[0]];
                                int o2 = indexes[pot.axes[1]];
                                int o3 = indexes[pot.axes[2]];
                                CHECK(obs(o1,o2,o3) == doctest::Approx(expected));
                            }
                        }
                    }
                }
            }
        }
    };
    SUBCASE("Axes are {0, 1, 2}.") {
        auto subtest = [&](size_t n, float u, float k) {
            return test(n, u, k, {0, 1, 2});
        };
        run_mutation_tests(subtest);
    }
    SUBCASE("Axes are {1, 0, 2}.") {
        auto subtest = [&](size_t n, float u, float k) {
            return test(n, u, k, {1, 0, 2});
        };
        run_mutation_tests(subtest);
    }
}
// LCOV_EXCL_STOP

template<typename Arg>
mutk::tensor_t create_transition_child_haploid_diploid(const mutk::mutation::Model &model,
    size_t n, const potential_t &potential, Arg arg) {
    assert(potential.parents.size() == 2);
    // no transpose here
    return create_transition_child<1, 2>(model, n, potential, arg);
}

// LCOV_EXCL_START
TEST_CASE("create_transition_child_haploid_diploid") {
    auto test = [&](size_t n, float u, float k, std::vector<int> axes) {
        CAPTURE(n);
        CAPTURE(k);
        CAPTURE(u);

        potential_t pot{mutk::detail::Potential:: ChildHaploidDiploid, 0, 1, u, 2, 1.1f*u};
        pot.axes = axes;
        KAllelesModel model(k, 0.001, 0, 0, 0);

        auto obs = create_transition_child_haploid_diploid(model, n, pot, mutk::mutation::ANY);
        auto mat1 = model.CreateMatrix(n, u, mutk::mutation::ANY);
        auto mat2 = model.CreateMatrix(n, 1.1f*u, mutk::mutation::ANY);

        REQUIRE(obs.dimension() == 3);
        mutk::shape_t pot_shape = {n*(n+1)/2, n, n*(n+1)/2};
        REQUIRE(obs.shape(0) == pot_shape[pot.axes[0]]);
        REQUIRE(obs.shape(1) == pot_shape[pot.axes[1]]);
        REQUIRE(obs.shape(2) == pot_shape[pot.axes[2]]);
        for(int a1=0,a=0;a1<n;++a1) {
            for(int a2=0;a2<=a1;++a2,++a) {
                for(int b=0;b<n;++b) {
                    for(int c1=0,c=0;c1<n;++c1) {
                        for(int c2=0;c2<=c1;++c2,++c) {
                            CAPTURE(a1);
                            CAPTURE(a2);
                            CAPTURE(a);
                            CAPTURE(b);
                            CAPTURE(c1);
                            CAPTURE(c2);
                            CAPTURE(c);
                            // b x c1/c2 -> a1/a2
                            float expected1 = mat1(a1,b);
                            float expected2 = 0.5f*mat2(a2,c1) + 0.5f*mat2(a2,c2);
                            float expected = expected1*expected2;
                            if(a1 != a2) {
                                expected1 = mat1(a2,b);
                                expected2 = 0.5f*mat2(a1,c1) + 0.5f*mat2(a1,c2);
                                expected += expected1*expected2;
                            }

                            std::array<int, 3> indexes = {a,b,c};
                            int o1 = indexes[pot.axes[0]];
                            int o2 = indexes[pot.axes[1]];
                            int o3 = indexes[pot.axes[2]];
                            CHECK(obs(o1,o2,o3) == doctest::Approx(expected));
                        }
                    }
                }
            }
        }
    };
    SUBCASE("Axes are {0, 1, 2}.") {
        auto subtest = [&](size_t n, float u, float k) {
            return test(n, u, k, {0, 1, 2});
        };
        run_mutation_tests(subtest);
    }
    SUBCASE("Axes are {1, 0, 2}.") {
        auto subtest = [&](size_t n, float u, float k) {
            return test(n, u, k, {1, 0, 2});
        };
        run_mutation_tests(subtest);
    }
}
// LCOV_EXCL_STOP

template<typename Arg>
mutk::tensor_t create_transition_child_diploid_haploid(const mutk::mutation::Model &model,
    size_t n, const potential_t &potential, Arg arg) {
    assert(potential.parents.size() == 2);
    // no transpose here
    return create_transition_child<2, 1>(model, n, potential, arg);
}

// LCOV_EXCL_START
TEST_CASE("create_transition_child_diploid_haploid") {
    auto test = [&](size_t n, float u, float k, std::vector<int> axes) {
        CAPTURE(n);
        CAPTURE(k);
        CAPTURE(u);

        potential_t pot{mutk::detail::Potential:: ChildDiploidHaploid, 0, 1, u, 2, 1.1f*u};
        pot.axes = axes;
        KAllelesModel model(k, 0.001, 0, 0, 0);

        auto obs = create_transition_child_diploid_haploid(model, n, pot, mutk::mutation::ANY);
        auto mat1 = model.CreateMatrix(n, u, mutk::mutation::ANY);
        auto mat2 = model.CreateMatrix(n, 1.1f*u, mutk::mutation::ANY);

        REQUIRE(obs.dimension() == 3);
        mutk::shape_t pot_shape = {n*(n+1)/2, n*(n+1)/2, n};
        REQUIRE(obs.shape(0) == pot_shape[pot.axes[0]]);
        REQUIRE(obs.shape(1) == pot_shape[pot.axes[1]]);
        REQUIRE(obs.shape(2) == pot_shape[pot.axes[2]]);
        for(int a1=0,a=0;a1<n;++a1) {
            for(int a2=0;a2<=a1;++a2,++a) {
                for(int b1=0,b=0;b1<n;++b1) {
                    for(int b2=0;b2<=b1;++b2,++b) {
                        for(int c=0;c<n;++c) {
                            CAPTURE(a1);
                            CAPTURE(a2);
                            CAPTURE(a);
                            CAPTURE(b1);
                            CAPTURE(b2);
                            CAPTURE(b);
                            CAPTURE(c);
                            // b1/b2 x c -> a1/a2
                            float expected1 = 0.5f*mat1(a1,b1) + 0.5f*mat1(a1,b2);
                            float expected2 = mat2(a2,c);
                            float expected = expected1*expected2;
                            if(a1 != a2) {
                                expected1 = 0.5f*mat1(a2,b1) + 0.5f*mat1(a2,b2);
                                expected2 = mat2(a1,c);
                                expected += expected1*expected2;
                            }

                            std::array<int, 3> indexes = {a,b,c};
                            int o1 = indexes[pot.axes[0]];
                            int o2 = indexes[pot.axes[1]];
                            int o3 = indexes[pot.axes[2]];
                            CHECK(obs(o1,o2,o3) == doctest::Approx(expected));
                        }
                    }
                }
            }
        }
    };
    SUBCASE("Axes are {0, 1, 2}.") {
        auto subtest = [&](size_t n, float u, float k) {
            return test(n, u, k, {0, 1, 2});
        };
        run_mutation_tests(subtest);
    }
    SUBCASE("Axes are {1, 0, 2}.") {
        auto subtest = [&](size_t n, float u, float k) {
            return test(n, u, k, {1, 0, 2});
        };
        run_mutation_tests(subtest);
    }
}
// LCOV_EXCL_STOP

template<typename Arg>
mutk::tensor_t create_transition_child_haploid_haploid(const mutk::mutation::Model &model,
    size_t n, const potential_t &potential, Arg arg) {
    assert(potential.parents.size() == 2);
    // no transpose here
    return create_transition_child<1, 1>(model, n, potential, arg);
}

// LCOV_EXCL_START
TEST_CASE("create_transition_child_haploid_haploid") {
    auto test = [&](size_t n, float u, float k, std::vector<int> axes) {
        CAPTURE(n);
        CAPTURE(k);
        CAPTURE(u);

        potential_t pot{mutk::detail::Potential:: ChildHaploidHaploid, 0, 1, u, 2, 1.1f*u};
        pot.axes = axes;
        KAllelesModel model(k, 0.001, 0, 0, 0);

        auto obs = create_transition_child_haploid_haploid(model, n, pot, mutk::mutation::ANY);
        auto mat1 = model.CreateMatrix(n, u, mutk::mutation::ANY);
        auto mat2 = model.CreateMatrix(n, 1.1f*u, mutk::mutation::ANY);

        REQUIRE(obs.dimension() == 3);
        mutk::shape_t pot_shape = {n*(n+1)/2, n, n};
        REQUIRE(obs.shape(0) == pot_shape[pot.axes[0]]);
        REQUIRE(obs.shape(1) == pot_shape[pot.axes[1]]);
        REQUIRE(obs.shape(2) == pot_shape[pot.axes[2]]);
        for(int a1=0,a=0;a1<n;++a1) {
            for(int a2=0;a2<=a1;++a2,++a) {
                for(int b=0;b<n;++b) {
                    for(int c=0;c<n;++c) {
                        CAPTURE(a1);
                        CAPTURE(a2);
                        CAPTURE(a);
                        CAPTURE(b);
                        CAPTURE(c);
                        // b x c -> a1/a2
                        float expected1 = mat1(a1,b);
                        float expected2 = mat2(a2,c);
                        float expected = expected1*expected2;
                        if(a1 != a2) {
                            expected1 = mat1(a2,b);
                            expected2 = mat2(a1,c);
                            expected += expected1*expected2;
                        }

                        std::array<int, 3> indexes = {a,b,c};
                        int o1 = indexes[pot.axes[0]];
                        int o2 = indexes[pot.axes[1]];
                        int o3 = indexes[pot.axes[2]];
                        CHECK(obs(o1,o2,o3) == doctest::Approx(expected));
                    }
                }
            }
        }
    };
    SUBCASE("Axes are {0, 1, 2}.") {
        auto subtest = [&](size_t n, float u, float k) {
            return test(n, u, k, {0, 1, 2});
        };
        run_mutation_tests(subtest);
    }
    SUBCASE("Axes are {1, 0, 2}.") {
        auto subtest = [&](size_t n, float u, float k) {
            return test(n, u, k, {1, 0, 2});
        };
        run_mutation_tests(subtest);
    }
}
// LCOV_EXCL_STOP

template<typename Arg>
mutk::tensor_t create_transition_child_selfing_diploid(const mutk::mutation::Model &model,
    size_t n, const potential_t &potential, Arg arg) {
    assert(potential.parents.size() == 1);
    // no transpose here
    return create_transition_child_selfing<2>(model, n, potential, arg);
}

// LCOV_EXCL_START
TEST_CASE("create_transition_child_selfing_diploid") {
    auto test = [&](size_t n, float u, float k, std::vector<int> axes) {
        CAPTURE(n);
        CAPTURE(k);
        CAPTURE(u);

        potential_t pot{mutk::detail::Potential::ChildSelfingDiploid, 0, 1, u};
        pot.axes = axes;
        KAllelesModel model(k, 0.001, 0, 0, 0);

        auto obs = create_transition_child_selfing_diploid(model, n, pot, mutk::mutation::ANY);
        auto mat1 = model.CreateMatrix(n, u, mutk::mutation::ANY);

        REQUIRE(obs.dimension() == 2);
        mutk::shape_t pot_shape = {n*(n+1)/2, n*(n+1)/2};
        REQUIRE(obs.shape(0) == pot_shape[pot.axes[0]]);
        REQUIRE(obs.shape(1) == pot_shape[pot.axes[1]]);
        for(int a1=0,a=0;a1<n;++a1) {
            for(int a2=0;a2<=a1;++a2,++a) {
                for(int b1=0,b=0;b1<n;++b1) {
                    for(int b2=0;b2<=b1;++b2,++b) {
                        CAPTURE(a1);
                        CAPTURE(a2);
                        CAPTURE(a);
                        CAPTURE(b1);
                        CAPTURE(b2);
                        CAPTURE(b);
                        // b1/b2 x b1/b2 -> a1/a2
                        float expected1 = 0.5f*mat1(a1,b1) + 0.5f*mat1(a1,b2);
                        float expected2 = 0.5f*mat1(a2,b1) + 0.5f*mat1(a2,b2);
                        float expected = expected1*expected2;
                        if(a1 != a2) {
                            expected1 = 0.5f*mat1(a2,b1) + 0.5f*mat1(a2,b2);
                            expected2 = 0.5f*mat1(a1,b1) + 0.5f*mat1(a1,b2);
                            expected += expected1*expected2;
                        }

                        std::array<int, 2> indexes = {a,b};
                        int o1 = indexes[pot.axes[0]];
                        int o2 = indexes[pot.axes[1]];
                        CHECK(obs(o1,o2) == doctest::Approx(expected));
                    }
                }
            }
        }
    };
    SUBCASE("Axes are {0, 1}.") {
        auto subtest = [&](size_t n, float u, float k) {
            return test(n, u, k, {0, 1});
        };
        run_mutation_tests(subtest);
    }
    SUBCASE("Axes are {1, 0}.") {
        auto subtest = [&](size_t n, float u, float k) {
            return test(n, u, k, {1, 0});
        };
        run_mutation_tests(subtest);
    }
}
// LCOV_EXCL_STOP

template<typename Arg>
mutk::tensor_t create_transition_child_selfing_haploid(const mutk::mutation::Model &model,
    size_t n, const potential_t &potential, Arg arg) {
    assert(potential.parents.size() == 1);
    // no transpose here
    return create_transition_child_selfing<1>(model, n, potential, arg);
}

// LCOV_EXCL_START
TEST_CASE("create_transition_child_selfing_haploid") {
    auto test = [&](size_t n, float u, float k, std::vector<int> axes) {
        CAPTURE(n);
        CAPTURE(k);
        CAPTURE(u);

        potential_t pot{mutk::detail::Potential::ChildSelfingHaploid, 0, 1, u};
        pot.axes = axes;
        KAllelesModel model(k, 0.001, 0, 0, 0);

        auto obs = create_transition_child_selfing_haploid(model, n, pot, mutk::mutation::ANY);
        auto mat1 = model.CreateMatrix(n, u, mutk::mutation::ANY);

        REQUIRE(obs.dimension() == 2);
        mutk::shape_t pot_shape = {n*(n+1)/2, n};
        REQUIRE(obs.shape(0) == pot_shape[pot.axes[0]]);
        REQUIRE(obs.shape(1) == pot_shape[pot.axes[1]]);
        for(int a1=0,a=0;a1<n;++a1) {
            for(int a2=0;a2<=a1;++a2,++a) {
                for(int b=0;b<n;++b) {
                    CAPTURE(a1);
                    CAPTURE(a2);
                    CAPTURE(a);
                    CAPTURE(b);
                    // b x b -> a1/a2
                    float expected1 = mat1(a1,b);
                    float expected2 = mat1(a2,b);
                    float expected = expected1*expected2;
                    if(a1 != a2) {
                        expected1 = mat1(a2,b);
                        expected2 = mat1(a1,b);
                        expected += expected1*expected2;
                    }

                    std::array<int, 2> indexes = {a,b};
                    int o1 = indexes[pot.axes[0]];
                    int o2 = indexes[pot.axes[1]];
                    CHECK(obs(o1,o2) == doctest::Approx(expected));
                }
            }
        }
    };
    SUBCASE("Axes are {0, 1}.") {
        auto subtest = [&](size_t n, float u, float k) {
            return test(n, u, k, {0, 1});
        };
        run_mutation_tests(subtest);
    }
    SUBCASE("Axes are {1, 0}.") {
        auto subtest = [&](size_t n, float u, float k) {
            return test(n, u, k, {1, 0});
        };
        run_mutation_tests(subtest);
    }
}
// LCOV_EXCL_STOP

mutk::tensor_t create_unit_potential(size_t n, const potential_t &potential) {
    assert(potential.type == mutk::detail::Potential::Unit);
    mutk::shape_t shape;
    for(auto &&a : potential.axes) {
        assert(a == 1 || a == 2);
        if(a == 1) {
            shape.push_back(mutk::dim_width<1>(n));
        } else {
            shape.push_back(mutk::dim_width<2>(n));
        }
    }
    return xt::ones<mutk::float_t>(shape);
}

namespace mutk {
template<typename Arg>
inline
mutk::tensor_t create_mutation_potential(const mutk::mutation::Model &model,
    size_t n, const potential_t &potential, Arg arg) {
    using P = mutk::detail::Potential;

    switch(potential.type) {
    case P::CloneDiploid:
        return create_transition_clone_diploid(model, n, potential, arg);
    case P::CloneHaploid:
        return create_transition_clone_haploid(model, n, potential, arg);
    case P::GameteDiploid:
        return create_transition_gamete_diploid(model, n, potential, arg);
    case P::ChildDiploidDiploid:
        return create_transition_child_diploid_diploid(model, n, potential, arg);
    case P::ChildHaploidDiploid:
        return create_transition_child_haploid_diploid(model, n, potential, arg);
    case P::ChildDiploidHaploid:
        return create_transition_child_diploid_haploid(model, n, potential, arg);
    case P::ChildHaploidHaploid:
        return create_transition_child_haploid_haploid(model, n, potential, arg);
    case P::ChildSelfingDiploid:
        return create_transition_child_selfing_diploid(model, n, potential, arg);
    case P::ChildSelfingHaploid:
        return create_transition_child_selfing_haploid(model, n, potential, arg);
    case P::Unit:
        return create_unit_potential(n, potential);
    case P::FounderDiploid:
        return model.CreatePriorDiploid(n);
    case P::FounderHaploid:
        return model.CreatePriorHaploid(n);
    default:
        break;
    };
    return {};
}
}

mutk::tensor_t mutk::mutation::Model::CreatePotential(size_t n, const potential_t &potential, any_t arg) {
    return create_mutation_potential(*this, n, potential, arg);
}

mutk::tensor_t mutk::mutation::Model::CreatePotential(size_t n, const potential_t &potential, mean_t arg) {
    return create_mutation_potential(*this, n, potential, arg);
}

mutk::tensor_t mutk::mutation::Model::CreatePotential(size_t n, const potential_t &potential, size_t arg) {
    return create_mutation_potential(*this, n, potential, arg);
}


#ifndef TEST_UNIT_MATH_REV_MAT_VECTORIZE_BUILD_REV_MATRIX_HPP
#define TEST_UNIT_MATH_REV_MAT_VECTORIZE_BUILD_REV_MATRIX_HPP

#include <stan/math/rev/core/var.hpp>
#include <vector>
#include <Eigen/Dense>

template <typename F, int R, int C>
static inline Eigen::Matrix<stan::math::var, R, C>
build_rev_matrix(const Eigen::Matrix<stan::math::var, R, C>& x) {
  using Eigen::Matrix;
  using std::vector;
  using stan::math::var;

  Matrix<var, R, C> var_matrix(x.rows(), x.cols());
  vector<double> inputs = F::valid_inputs();
  //Fills matrix with copies of valid_input values
  for (int i = 0; i < x.size(); ++i)
    var_matrix(i) = inputs[i % inputs.size()];
  return var_matrix;
}

#endif


#include <map>
#include <assert.h>
#include <sstream>
#include "Workspace.h"
#include "PrtUser.h"
#include "PrtExecution.h"
#include "goto_solver.h"
#include <Eigen/Dense>
#include <nav_msgs/Odometry.h>
#include <ros/ros.h>
#include "InitRos.h"
#include "Compat.h"
#include <pthread.h>

using namespace std;

extern "C" {
PRT_VALUE *P_FUN_RosInit_IMPL(PRT_MACHINEINST *context);
PRT_VALUE *P_FUN_StartExecutingPath_IMPL(PRT_MACHINEINST *context);
}

pthread_mutex_t publishers_map_lock = PTHREAD_MUTEX_INITIALIZER;
std::map<int, ros::Publisher> publishers;

double t_goto = 0.5;
double tscale = 1.0;

#ifndef USE_EMPTY
PRT_VALUE *P_FUN_RosInit_IMPL(PRT_MACHINEINST *context)
{
    PRT_MACHINEINST_PRIV *p_tmp_mach_priv = (PRT_MACHINEINST_PRIV *)context;
    PRT_VALUE *p_tmp_ret = NULL;
    PRT_FUNSTACK_INFO p_tmp_frame;
    PRT_VALUE *p_tmp_params;
    p_tmp_params = NULL;
    //remm to pop frame
    PrtPopFrame(p_tmp_mach_priv, &p_tmp_frame);
    //create a tuple value
    int robot_id = (int)PrtPrimGetInt(p_tmp_frame.locals[0]);
    int start_location = (int)PrtPrimGetInt(p_tmp_frame.locals[1]);
    std::ostringstream channel_name;
    channel_name << "odom" << robot_id;
    ros::Publisher pub = node_handle->advertise<nav_msgs::Odometry>(channel_name.str(), 10, true);
    pthread_mutex_lock(&publishers_map_lock);
    publishers[robot_id] = pub;
    pthread_mutex_unlock(&publishers_map_lock);
    SleepMs(1000); // Wait for publisher to register...
    printf("Registered robot %d at channel ", robot_id);
    std::cout << channel_name.str() << std::endl;

    WS_Coord start_coord = ExtractCoordFromGridLocation(start_location, WORKSPACE_INFO->dimension);
    nav_msgs::Odometry odom;
    odom.header.frame_id = "/simulator";
    odom.child_frame_id = "/quadrotor";
    double resolution;
    node_handle->param("/map/res", resolution, 1.0);
    odom.pose.pose.position.x = start_coord.x * resolution;
    odom.pose.pose.position.y = start_coord.y * resolution;
    odom.pose.pose.position.z = start_coord.z * resolution;
    pub.publish(odom);

    //remm to free the frame
    PrtFreeLocals(p_tmp_mach_priv, &p_tmp_frame);
    return p_tmp_ret;
}

static Eigen::Vector3d ReadVectorCoord(PRT_VALUE* trajSeq, PRT_UINT32 i)
{
    int loc = (int)PrtPrimGetInt(PrtSeqGetNCIntIndex(trajSeq, i));
    WS_Coord coord = ExtractCoordFromGridLocation(loc, WORKSPACE_INFO->dimension);
    return Eigen::Vector3d(coord.x, coord.y, coord.z);
}

static double cal_pos(double c[],double t)
{
    return c[0]*pow(t,7.0) +c[1]*pow(t,6.0) + c[2]*pow(t,5.0) + c[3]*pow(t,4.0) + c[4]*pow(t,3.0) + c[5]*pow(t,2.0) + c[6]*t + c[7];
}

static double cal_vel(double c[],double t)
{
    return (c[0]*7*pow(t,6.0) +c[1]*6*pow(t,5.0) + c[2]*5*pow(t,4.0) + c[3]*4*pow(t,3.0) + c[4]*3*pow(t,2.0) + c[5]*2*t + c[6])/tscale;
}

static double cal_acc(double c[],double t)
{
    return (c[0]*42*pow(t,5.0) +c[1]*30*pow(t,4.0) + c[2]*20*pow(t,3.0) + c[3]*12*pow(t,2.0) + c[4]*6*t + c[5]*2)/(tscale*tscale);
}

static void publish_straight_traj(int robot_id, Eigen::Vector3d start, Eigen::Vector3d end, double duration)
{
    nav_msgs::Odometry odom;
    odom.header.frame_id = "/simulator";
    odom.child_frame_id = "/quadrotor";

    double resolution;
    node_handle->param("/map/res", resolution, 1.0);
    start *= resolution;
    end *= resolution;

    ros::Rate odom_pub_duration(100);
    double t_start = ros::Time::now().toSec();
    double t_end = t_start + duration * tscale;

    Eigen::Vector3d vel = (end - start) / (duration / tscale);

    for(double cur_t = t_start; cur_t < t_end; cur_t = ros::Time::now().toSec()) {
        odom.header.stamp = ros::Time::now();

        double t = (cur_t - t_start) / tscale;
        Eigen::Vector3d loc = t * vel + start;
        odom.pose.pose.position.x = loc[0];
        odom.pose.pose.position.y = loc[1];
        odom.pose.pose.position.z = loc[2];

        odom.pose.pose.orientation.x = 0;
        odom.pose.pose.orientation.y = 0;
        odom.pose.pose.orientation.z = 0;
        odom.pose.pose.orientation.w = 0;

        odom.twist.twist.linear.x = vel[0];
        odom.twist.twist.linear.y = vel[1];
        odom.twist.twist.linear.z = vel[2];

        odom.twist.twist.angular.x = 0;
        odom.twist.twist.angular.y = 0;
        odom.twist.twist.angular.z = 0;

        publishers[robot_id].publish(odom);
        odom_pub_duration.sleep();
    }
}

static const double     _PI= 3.1415926535897932384626433832795028841971693993751058209749445923078164062862089986280348;
static const double _TWO_PI= 6.2831853071795864769252867665590057683943387987502116419498891846156328125724179972560696;

static void publish_turn_traj(int robot_id, Eigen::Vector3d start, Eigen::Vector3d turning, Eigen::Vector3d end, double duration)
{
    nav_msgs::Odometry odom;
    odom.header.frame_id = "/simulator";
    odom.child_frame_id = "/quadrotor";

    double resolution;
    node_handle->param("/map/res", resolution, 1.0);
    start *= resolution;
    end *= resolution;
    turning *= resolution;

    Eigen::Vector3d origin = end + (start - turning);

    ros::Rate odom_pub_duration(100);
    double t_start = ros::Time::now().toSec();
    double t_end = t_start + duration * tscale;

    double omega = (_PI / 2) / (duration / tscale);

    for(double cur_t = t_start; cur_t < t_end; cur_t = ros::Time::now().toSec()) {
        odom.header.stamp = ros::Time::now();

        double t = (cur_t - t_start) / tscale;
        double theta = (omega * t);
        Eigen::Vector3d loc = sin(theta) * (end - origin) + cos(theta) * (start - origin) + origin;
        odom.pose.pose.position.x = loc[0];
        odom.pose.pose.position.y = loc[1];
        odom.pose.pose.position.z = loc[2];

        odom.pose.pose.orientation.x = 0;
        odom.pose.pose.orientation.y = 0;
        odom.pose.pose.orientation.z = 0;
        odom.pose.pose.orientation.w = 0;

        odom.twist.twist.linear.x = 0;
        odom.twist.twist.linear.y = 0;
        odom.twist.twist.linear.z = 0;

        odom.twist.twist.angular.x = 0;
        odom.twist.twist.angular.y = 0;
        odom.twist.twist.angular.z = 0;

        publishers[robot_id].publish(odom);
        odom_pub_duration.sleep();
    }
}

bool time_initialized = false;
ros::Time initial_time;

PRT_VALUE *P_FUN_StartExecutingPath_IMPL(PRT_MACHINEINST *context)
{
    PRT_MACHINEINST_PRIV *p_tmp_mach_priv = (PRT_MACHINEINST_PRIV *)context;
    PRT_VALUE *p_tmp_ret = NULL;
    PRT_FUNSTACK_INFO p_tmp_frame;
    PRT_VALUE *p_tmp_params;
    p_tmp_params = NULL;
    //remm to pop frame
    PrtPopFrame(p_tmp_mach_priv, &p_tmp_frame);
    //create a tuple value
    PRT_VALUE* trajSeq = p_tmp_frame.locals[0];
    int start_time = (int)PrtPrimGetInt(p_tmp_frame.locals[1]);
    int robot_id = (int)PrtPrimGetInt(p_tmp_frame.locals[2]);

    if(!time_initialized) {
        initial_time = ros::Time::now() - ros::Duration(t_goto * start_time - 0.0001);
        time_initialized = true;
    }

    ros::Time sleep_until = initial_time + ros::Duration(start_time * t_goto);
#ifdef ASSERT_TIME_SYNC
    double went_by = (sleep_until- ros::Time::now()).toSec();
    printf("went by %f\n", went_by);
    assert(went_by >= 0);
#endif
    ros::Time::sleepUntil(sleep_until);

    PRT_UINT32 count = PrtSeqSizeOf(trajSeq);
    if(count >= 2) {
        Eigen::Vector3d diff;
        PRT_UINT32 i = 1;
        Eigen::Vector3d current_step = ReadVectorCoord(trajSeq, i) - ReadVectorCoord(trajSeq, i - 1);
        bool done = false;
        int traj_count = 0;
        while(!done) {
            int step_count = 0;
            Eigen::Vector3d start_coord = ReadVectorCoord(trajSeq, i - 1);
            do {
                i++;
                step_count++;
                if(i == count) {
                    done = true;
                    break;
                }
                diff = ReadVectorCoord(trajSeq, i) - ReadVectorCoord(trajSeq, i - 1);
            } while (diff == current_step);
            traj_count++;
            double duration = step_count * t_goto;
            Eigen::Vector3d straight_step = current_step * step_count;
#ifdef SIM_LOG_MOVE
            printf("robot %d move in (%f, %f, %f)\n", robot_id, straight_step[0], straight_step[1], straight_step[2]);
#endif
            Eigen::Vector3d end_coord = straight_step + start_coord;
            if(traj_count != 1) {
                start_coord += (current_step / 2);
                duration -= t_goto / 2;
            }
            if(i != count) {
                end_coord -= (current_step / 2);
                duration -= t_goto / 2;
            }
            if(step_count != 0) {
                publish_straight_traj(robot_id, start_coord, end_coord, duration);
            }
            if(i != count) {
                publish_turn_traj(robot_id, end_coord, end_coord + current_step / 2, end_coord + (current_step + diff) / 2, t_goto);
            }
            current_step = diff;
        }
    }
    //remm to free the frame
    PrtFreeLocals(p_tmp_mach_priv, &p_tmp_frame);
    return p_tmp_ret;
}
#endif


#include <iostream>
#include <unordered_map>
#include <boost/program_options.hpp>
#include "picpac.h"

using namespace std;
using namespace picpac;

int main(int argc, char const* argv[]) {
    fs::path db_path;

    namespace po = boost::program_options;
    po::options_description desc("Allowed options");
    desc.add_options()
        ("help,h", "produce help message.")
        ("db", po::value(&db_path), "")
        ;

    po::positional_options_description p;
    p.add("db", 1);

    po::variables_map vm;
    po::store(po::command_line_parser(argc, argv).
                     options(desc).positional(p).run(), vm);
    po::notify(vm);

    if (vm.count("help") || db_path.empty()) {
        cout << "Usage:" << endl;
        cout << "\tpicpac-stat ... <db>" << endl;
        cout << desc;
        cout << endl;
        return 0;
    }

    IndexedFileReader db(db_path);
    std::unordered_map<int, int> cnt;
    cout << "TOTAL: " << db.size() << endl;
    bool cat = true;
    for (unsigned i = 0; i < db.size(); ++i) {
        float l = db.group(i);
        if (int(l) != l) {
            cat = false;
            break;
        }
        cnt[l] += 1;
    }
    if (cat) {
        vector<std::pair<int, int>> all(cnt.begin(), cnt.end());
        sort(all.begin(), all.end());
        for (auto const &p: all) {
            cout << p.first << ": " << p.second << endl;
        }
    }

    return 0;
}


// (C) Copyright Jeremy Siek 2000. Permission to copy, use, modify, sell and
// distribute this software is granted provided this copyright notice appears
// in all copies. This software is provided "as is" without express or implied
// warranty, and with no claim as to its suitability for any purpose.

#include <functional>
#include <algorithm>
#include <iostream>
#include <boost/iterator_adaptors.hpp>

// What a bummer. We can't use std::binder1st with transform iterator
// because it does not have a default constructor. Here's a version
// that does.

namespace boost {

  template <class Operation>
  class binder1st
    : public std::unary_function<typename Operation::second_argument_type,
                                 typename Operation::result_type> {
  protected:
    Operation op;
    typename Operation::first_argument_type value;
  public:
    binder1st() { } // this had to be added!
    binder1st(const Operation& x,
              const typename Operation::first_argument_type& y)
        : op(x), value(y) {}
    typename Operation::result_type
    operator()(const typename Operation::second_argument_type& x) const {
      return op(value, x);
    }
  };

  template <class Operation, class T>
  inline binder1st<Operation> bind1st(const Operation& op, const T& x) {
    typedef typename Operation::first_argument_type arg1_type;
    return binder1st<Operation>(op, arg1_type(x));
  }

} // namespace boost

int
main(int, char*[])
{
  // This is a simple example of using the transform_iterators class to
  // generate iterators that multiply the value returned by dereferencing
  // the iterator. In this case we are multiplying by 2.
  // Would be cooler to use lambda library in this example.

  int x[] = { 1, 2, 3, 4, 5, 6, 7, 8 };
  const int N = sizeof(x)/sizeof(int);

  typedef boost::binder1st< std::multiplies<int> > Function;
  typedef boost::transform_iterator_generator<Function, int*>::type doubling_iterator;

  doubling_iterator i(x, boost::bind1st(std::multiplies<int>(), 2)),
    i_end(x + N, boost::bind1st(std::multiplies<int>(), 2));

  std::cout << "multiplying the array by 2:" << std::endl;
  while (i != i_end)
    std::cout << *i++ << " ";
  std::cout << std::endl;

  std::cout << "adding 4 to each element in the array:" << std::endl;

  std::copy(boost::make_transform_iterator(x, boost::bind1st(std::plus<int>(), 4)),
            boost::make_transform_iterator(x + N, boost::bind1st(std::plus<int>(), 4)),
            std::ostream_iterator<int>(std::cout, " "));
  std::cout << std::endl;

  return 0;
}


//        Copyright The Authors 2018.
//    Distributed under the 3-Clause BSD License.
//    (See accompanying file LICENSE or copy at
//   https://opensource.org/licenses/BSD-3-Clause)

#include <crypto/hash.h>

#include <boost/endian/conversion.hpp>

#include <common/platform.h>

namespace blocxxi {
namespace crypto {
namespace detail {

std::uint32_t HostToNetwork(std::uint32_t number) {
  return boost::endian::native_to_big(number);
}

std::uint32_t NetworkToHost(std::uint32_t number) {
  return boost::endian::big_to_native(number);
}

///@{
/// Find first set functions.
/// For details on the algorithms and options refer to
/// https://en.wikipedia.org/wiki/Find_first_set
///
/// These functions expect the range to be in big-endian byte order.

int CountLeadingZeroBits_SW(gsl::span<std::uint32_t const> buf) {
  auto const num = int(buf.size());
  std::uint32_t const* ptr = buf.data();

  for (int i = 0; i < num; i++) {
    if (ptr[i] == 0) continue;
    std::uint32_t v = NetworkToHost(ptr[i]);

    // http://graphics.stanford.edu/~seander/bithacks.html#IntegerLogObvious
    static const int MultiplyDeBruijnBitPosition[32] = {
        0, 9,  1,  10, 13, 21, 2,  29, 11, 14, 16, 18, 22, 25, 3, 30,
        8, 12, 20, 28, 15, 17, 24, 7,  19, 27, 23, 6,  26, 5,  4, 31};

    v |= v >> 1;  // first round down to one less than a power of 2
    v |= v >> 2;
    v |= v >> 4;
    v |= v >> 8;
    v |= v >> 16;

    // clang-format off
    return i * 32 + 31 - MultiplyDeBruijnBitPosition[static_cast<std::uint32_t>(v * 0x07C4ACDDU) >> 27];
    // clang-format on
  }

  return num * 32;
}

int CountLeadingZeroBits_HW(gsl::span<std::uint32_t const> buf) {
  auto const num = int(buf.size());
  std::uint32_t const* ptr = buf.data();

  for (int i = 0; i < num; i++) {
    if (ptr[i] == 0) continue;
    std::uint32_t v = NetworkToHost(ptr[i]);

#if ASAP_HAS_BUILTIN_CLZ
    return i * 32 + __builtin_clz(v);
#elif defined _MSC_VER
    unsigned long pos;
    _BitScanReverse(&pos, v);
    return i * 32 + 31 - pos;
#else
    ASAP_ASSERT_FAIL();
    return -1;
#endif
  }

  return num * 32;
}

int CountLeadingZeroBits(gsl::span<std::uint32_t const> buf) {
#if ASAP_HAS_BUILTIN_CLZ || defined _MSC_VER
  return CountLeadingZeroBits_HW(buf);
#else
  return CountLeadingZeroBits_SW(buf);
#endif
}

///@}

}  // namespace detail
}  // namespace crypto
}  // namespace blocxxi


#include <boost/test/unit_test.hpp>
#include "types/ds/fenwick_tree.hpp"

BOOST_AUTO_TEST_SUITE(TestFenwickTree)

    BOOST_AUTO_TEST_CASE(empty_ft)
    {
        Types::DS::FenwickTree<int> ft(1);
    }

    BOOST_AUTO_TEST_CASE(init_ft_with_values)
    {
        Types::DS::FenwickTree<int> ft({1, 2, 3, 4});
        BOOST_CHECK(ft.get(0) == 1);
        BOOST_CHECK(ft.get(1) == 2);
        BOOST_CHECK(ft.get(2) == 3);
        BOOST_CHECK(ft.get(3) == 4);
    }

    BOOST_AUTO_TEST_CASE(fill_ft)
    {
        Types::DS::FenwickTree<int> ft(6);
        ft.add(0, 3);
        ft.add(1, 4);
        ft.add(2, 5);
        ft.add(3, 6);
        ft.add(4, 7);
        ft.add(5, 8);

        BOOST_CHECK(ft.get(0) == 3);
        BOOST_CHECK(ft.get(1) == 4);
        BOOST_CHECK(ft.get(2) == 5);
        BOOST_CHECK(ft.get(3) == 6);
        BOOST_CHECK(ft.get(4) == 7);
        BOOST_CHECK(ft.get(5) == 8);
    }

    BOOST_AUTO_TEST_CASE(prefix_sum)
    {
        Types::DS::FenwickTree<int> ft(6);
        ft.add(0, 3);
        ft.add(1, 4);
        ft.add(2, 5);
        ft.add(3, 6);
        ft.add(4, 7);
        ft.add(5, 8);
        ft.add(3, 1);

        BOOST_CHECK(ft.prefix_sum(0) == 3);
        BOOST_CHECK(ft.prefix_sum(3) == 19);
        BOOST_CHECK(ft.prefix_sum(5) == 34);
    }

    BOOST_AUTO_TEST_CASE(set)
    {
        Types::DS::FenwickTree<int> ft({1, 2, 3, 4});
        ft.set(1, 10);

        BOOST_CHECK(ft.get(0) == 1);
        BOOST_CHECK(ft.get(1) == 10);
        BOOST_CHECK(ft.get(2) == 3);
        BOOST_CHECK(ft.get(3) == 4);
    }

BOOST_AUTO_TEST_SUITE_END()